diff --git a/.changeset/restore-hot-inject.md b/.changeset/restore-hot-inject.md new file mode 100644 index 00000000000..48daa9a34a3 --- /dev/null +++ b/.changeset/restore-hot-inject.md @@ -0,0 +1,5 @@ +--- +"kilo-code": patch +--- + +Fix mid-turn message handling so a new prompt sent while the assistant is working no longer aborts the in-flight response. The current LLM reply streams to completion, any pending suggestion or question is automatically dismissed, and the new prompt runs immediately after the current step instead of waiting for the entire multi-step turn to finish. diff --git a/packages/kilo-vscode/src/KiloProvider.ts b/packages/kilo-vscode/src/KiloProvider.ts index c30299a5cdf..ba037497f4f 100644 --- a/packages/kilo-vscode/src/KiloProvider.ts +++ b/packages/kilo-vscode/src/KiloProvider.ts @@ -54,7 +54,7 @@ import { clearCommandsCache, loadCommands } from "./kilo-provider/commands" import { fetchMessagePage, MESSAGE_PAGE_LIMIT } from "./kilo-provider/message-page" import { childID } from "./kilo-provider/task-session" import { handleNetworkEvent, clearNetworkWaits } from "./kilo-provider/network" -import { abortSession, parseQueued } from "./kilo-provider/abort" +import { abortSession } from "./kilo-provider/abort" import { buildAutocompleteSettingsMessage, routeAutocompleteMessage, @@ -622,7 +622,7 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper } case "abort": this.cancelRetry(message.sessionID ?? "") - await this.handleAbort(message.sessionID, parseQueued(message.queuedMessageIDs)) + await this.handleAbort(message.sessionID) break case "revertSession": this.handleRevertSession(message.sessionID, message.messageID).catch((e) => @@ -2550,7 +2550,7 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper } } - private async handleAbort(sessionID?: string, queuedMessageIDs: string[] = []): Promise { + private async handleAbort(sessionID?: string): Promise { if (!this.client) { return } @@ -2565,7 +2565,6 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper client: this.client, sessionID: targetSessionID, dir: this.getWorkspaceDirectory(targetSessionID), - queuedMessageIDs, }) } catch (error) { console.error("[Kilo New] KiloProvider: Failed to abort session:", error) diff --git a/packages/kilo-vscode/src/agent-manager/types.ts b/packages/kilo-vscode/src/agent-manager/types.ts index 65e34a730b7..5f4846c96c7 100644 --- a/packages/kilo-vscode/src/agent-manager/types.ts +++ b/packages/kilo-vscode/src/agent-manager/types.ts @@ -623,7 +623,6 @@ interface ForkSessionIn { interface AbortIn { type: "abort" sessionID: string - queuedMessageIDs?: string[] } interface ContinueInWorktreeIn { diff --git a/packages/kilo-vscode/src/kilo-provider/abort.ts b/packages/kilo-vscode/src/kilo-provider/abort.ts index e295fea5454..8d3d00676e9 100644 --- a/packages/kilo-vscode/src/kilo-provider/abort.ts +++ b/packages/kilo-vscode/src/kilo-provider/abort.ts @@ -1,21 +1,5 @@ import type { KiloClient } from "@kilocode/sdk/v2/client" -export function parseQueued(value: unknown) { - if (!Array.isArray(value)) return [] - return value.filter((id): id is string => typeof id === "string") -} - -export async function abortSession(input: { - client: KiloClient - sessionID: string - dir: string - queuedMessageIDs: string[] -}) { +export async function abortSession(input: { client: KiloClient; sessionID: string; dir: string }) { await input.client.session.abort({ sessionID: input.sessionID, directory: input.dir }, { throwOnError: true }) - - for (const mid of new Set(input.queuedMessageIDs)) { - await input.client.session - .deleteMessage({ sessionID: input.sessionID, messageID: mid, directory: input.dir }, { throwOnError: true }) - .catch((err) => console.error("[Kilo New] KiloProvider: Failed to remove queued message:", err)) - } } diff --git a/packages/kilo-vscode/tests/unit/abort.test.ts b/packages/kilo-vscode/tests/unit/abort.test.ts index c9002a650c6..126333e6a4e 100644 --- a/packages/kilo-vscode/tests/unit/abort.test.ts +++ b/packages/kilo-vscode/tests/unit/abort.test.ts @@ -1,6 +1,6 @@ import { describe, expect, it } from "bun:test" import type { KiloClient } from "@kilocode/sdk/v2/client" -import { abortSession, parseQueued } from "../../src/kilo-provider/abort" +import { abortSession } from "../../src/kilo-provider/abort" function client(calls: unknown[], fail = false) { return { @@ -10,35 +10,15 @@ function client(calls: unknown[], fail = false) { if (fail) throw new Error("abort failed") return { data: true } }, - deleteMessage: async (params: unknown, opts: unknown) => { - calls.push({ type: "delete", params, opts }) - return { data: true } - }, }, } as unknown as KiloClient } -describe("parseQueued", () => { - it("keeps only string queued message ids", () => { - expect(parseQueued(["message_1", 2, null, "message_2", {}])).toEqual(["message_1", "message_2"]) - }) - - it("returns empty ids for invalid payloads", () => { - expect(parseQueued(undefined)).toEqual([]) - expect(parseQueued({ queuedMessageIDs: ["message_1"] })).toEqual([]) - }) -}) - describe("abortSession", () => { - it("aborts before removing queued follow-up messages", async () => { + it("calls session.abort with the session id and directory", async () => { const calls: unknown[] = [] - await abortSession({ - client: client(calls), - sessionID: "session_1", - dir: "/repo", - queuedMessageIDs: ["message_2", "message_3", "message_2"], - }) + await abortSession({ client: client(calls), sessionID: "session_1", dir: "/repo" }) expect(calls).toEqual([ { @@ -46,30 +26,15 @@ describe("abortSession", () => { params: { sessionID: "session_1", directory: "/repo" }, opts: { throwOnError: true }, }, - { - type: "delete", - params: { sessionID: "session_1", messageID: "message_2", directory: "/repo" }, - opts: { throwOnError: true }, - }, - { - type: "delete", - params: { sessionID: "session_1", messageID: "message_3", directory: "/repo" }, - opts: { throwOnError: true }, - }, ]) }) - it("does not remove queued messages when abort fails", async () => { + it("rejects when the abort request fails", async () => { const calls: unknown[] = [] - await expect( - abortSession({ - client: client(calls, true), - sessionID: "session_1", - dir: "/repo", - queuedMessageIDs: ["message_2"], - }), - ).rejects.toThrow("abort failed") + await expect(abortSession({ client: client(calls, true), sessionID: "session_1", dir: "/repo" })).rejects.toThrow( + "abort failed", + ) expect(calls).toEqual([ { diff --git a/packages/kilo-vscode/webview-ui/src/context/session.tsx b/packages/kilo-vscode/webview-ui/src/context/session.tsx index 485f4ba964a..0a1936d7232 100644 --- a/packages/kilo-vscode/webview-ui/src/context/session.tsx +++ b/packages/kilo-vscode/webview-ui/src/context/session.tsx @@ -46,7 +46,6 @@ import { import { Identifier } from "../utils/id" import { resolveModelSelection } from "./model-selection" import { resolveSessionAgent } from "./session-agent" -import { queuedUserMessageIDs } from "./session-queue" import { PartStash } from "./part-stash" import { KILO_AUTO, parseModelString } from "../../../src/shared/provider-model" @@ -1748,12 +1747,9 @@ export const SessionProvider: ParentComponent = (props) => { return } - const queuedMessageIDs = queuedUserMessageIDs(messages(), statusInfo()) - vscode.postMessage({ type: "abort", sessionID, - queuedMessageIDs, }) } diff --git a/packages/kilo-vscode/webview-ui/src/types/messages.ts b/packages/kilo-vscode/webview-ui/src/types/messages.ts index 473528888a6..fe803a43c1a 100644 --- a/packages/kilo-vscode/webview-ui/src/types/messages.ts +++ b/packages/kilo-vscode/webview-ui/src/types/messages.ts @@ -1710,7 +1710,6 @@ export interface SendMessageRequest { export interface AbortRequest { type: "abort" sessionID: string - queuedMessageIDs?: string[] } export interface RevertSessionRequest { diff --git a/packages/opencode/src/kilocode/question/index.ts b/packages/opencode/src/kilocode/question/index.ts new file mode 100644 index 00000000000..b827c866742 --- /dev/null +++ b/packages/opencode/src/kilocode/question/index.ts @@ -0,0 +1,62 @@ +import { Deferred, Effect } from "effect" +import { InstanceState } from "@/effect" +import { Log } from "@/util" +import { SessionID } from "@/session/schema" +import { KiloSessionPromptQueue } from "@/kilocode/session/prompt-queue" + +/** + * Kilo-specific helpers for the shared `@/question` module. + * + * Extracted here so the upstream file keeps just the import, an Interface entry + * for `dismissAll`, and one-liner calls at the use sites — minimising the + * surface area that conflicts with upstream. + */ +export namespace KiloQuestion { + const log = Log.create({ service: "question" }) + + /** Minimal entry shape both helpers need; matches `PendingEntry` in `@/question`. */ + type Entry = { + info: { id: unknown; sessionID: SessionID } + deferred: Deferred.Deferred + } + + /** + * Factory for `Question.dismissAll`: dismisses every pending question on a + * session so a new prompt can unblock an in-flight tool waiting on user + * input. Mirrors `Suggestion.dismissAll` so both read the same way at the + * callsite. + * + * The caller provides a `publishRejected` callback (closed over the already- + * resolved `Bus.Service` in the Question layer) and an error factory so this + * helper stays free of any `@/question` import and dodges a circular dep. + */ + export const makeDismissAll = + (args: { + state: InstanceState.InstanceState<{ pending: Map }> + publishRejected: (entry: PE) => Effect.Effect + makeError: () => PE["deferred"] extends Deferred.Deferred ? E : never + }) => + (sessionID: SessionID) => + Effect.gen(function* () { + const pending = (yield* InstanceState.get(args.state)).pending + for (const [id, entry] of Array.from(pending.entries())) { + if (entry.info.sessionID !== sessionID) continue + pending.delete(id) + log.info("dismissed", { requestID: id }) + yield* args.publishRejected(entry) + yield* Deferred.fail(entry.deferred, args.makeError()) + } + }) + + /** + * Auto-dismiss when a newer prompt is already queued on this session — a + * tool that calls `Question.ask` after the queue event would otherwise block + * the run while the user waits for their queued prompt to take over. + */ + export const guardFollowup = (sessionID: SessionID, makeError: () => E) => + Effect.gen(function* () { + if (!KiloSessionPromptQueue.hasFollowup(sessionID)) return + log.info("auto-dismissed — followup queued", { sessionID }) + return yield* Effect.fail(makeError()) + }) +} diff --git a/packages/opencode/src/kilocode/session/prompt-queue.ts b/packages/opencode/src/kilocode/session/prompt-queue.ts index fa520935b63..04df601c98e 100644 --- a/packages/opencode/src/kilocode/session/prompt-queue.ts +++ b/packages/opencode/src/kilocode/session/prompt-queue.ts @@ -3,6 +3,7 @@ import { MessageV2 } from "@/session/message-v2" import { MessageID, SessionID } from "@/session/schema" type Slot = { + readonly seq: number readonly version: number readonly previous: Promise readonly done: PromiseWithResolvers @@ -18,6 +19,13 @@ export namespace KiloSessionPromptQueue { const tails = new Map>() const versions = new Map() const targets = new Map() + // Monotonic arrival counter per session. latest holds the seq of the most + // recently enqueued slot; activeSince snapshots latest at the moment the + // currently running slot actually started. hasFollowup returns true only when + // a newer slot was enqueued after the active one began running. + const latest = new Map() + const activeSince = new Map() + let seq = 0 const version = (sessionID: SessionID) => versions.get(sessionID) ?? 0 const settle = (promise: Promise) => @@ -45,6 +53,18 @@ export namespace KiloSessionPromptQueue { targets.set(sessionID, { base: current.base, extras }) } + /** + * True when a newer prompt was enqueued after the currently running slot + * began. runLoop calls this between LLM steps to break out so the next + * queued prompt can take over without starting another LLM round-trip for + * the now-superseded turn. + */ + export function hasFollowup(sessionID: SessionID): boolean { + const l = latest.get(sessionID) ?? 0 + const a = activeSince.get(sessionID) ?? 0 + return l > a + } + export function scope(sessionID: SessionID, messages: MessageV2.WithParts[]) { const target = targets.get(sessionID) if (!target) return messages @@ -90,17 +110,23 @@ export namespace KiloSessionPromptQueue { ): Effect.Effect { return Effect.acquireUseRelease( Effect.sync(() => { + const mine = ++seq + latest.set(sessionID, mine) const previous = tails.get(sessionID) ?? Promise.resolve() const done = Promise.withResolvers() // Keep later queued prompts moving; each caller still observes its own failure. const tail = settle(previous).then(() => done.promise) tails.set(sessionID, tail) - return { version: version(sessionID), previous, done, tail } satisfies Slot + return { seq: mine, version: version(sessionID), previous, done, tail } satisfies Slot }), (slot) => Effect.promise(() => settle(slot.previous)).pipe( Effect.flatMap(() => { if (slot.version !== version(sessionID)) return cancelled + // Snapshot the latest seq at the moment this slot actually starts + // running. hasFollowup compares against this value so the slot only + // breaks when something newer than itself arrives. + activeSince.set(sessionID, latest.get(sessionID) ?? slot.seq) return Effect.acquireUseRelease( Effect.sync(() => { targets.set(sessionID, { base: target, extras: new Set() }) @@ -120,6 +146,8 @@ export namespace KiloSessionPromptQueue { tails.delete(sessionID) versions.delete(sessionID) targets.delete(sessionID) + latest.delete(sessionID) + activeSince.delete(sessionID) }), ) } diff --git a/packages/opencode/src/kilocode/suggestion/index.ts b/packages/opencode/src/kilocode/suggestion/index.ts index d5fefdbd85d..4607d09602c 100644 --- a/packages/opencode/src/kilocode/suggestion/index.ts +++ b/packages/opencode/src/kilocode/suggestion/index.ts @@ -1,8 +1,10 @@ import { Bus } from "../../bus" import { BusEvent } from "../../bus/bus-event" import { Identifier } from "../../id/id" +import { SessionID } from "../../session/schema" import { Log } from "../../util" import z from "zod" +import { KiloSessionPromptQueue } from "../session/prompt-queue" export namespace Suggestion { const log = Log.create({ service: "suggestion" }) @@ -95,6 +97,14 @@ export namespace Suggestion { blocking?: boolean tool?: { messageID: string; callID: string } }): Promise { + // Auto-dismiss if a newer prompt is already queued on this session. + // Synchronous check immediately before the pending set, so there's no + // interleaving with dismissAll called from SessionPrompt.prompt. + if (KiloSessionPromptQueue.hasFollowup(SessionID.make(input.sessionID))) { + log.info("auto-dismissed — followup queued", { sessionID: input.sessionID }) + throw new DismissedError() + } + const s = { pending } const id = Identifier.ascending("suggestion") diff --git a/packages/opencode/src/kilocode/tool/question.ts b/packages/opencode/src/kilocode/tool/question.ts new file mode 100644 index 00000000000..9e093e7c7ee --- /dev/null +++ b/packages/opencode/src/kilocode/tool/question.ts @@ -0,0 +1,27 @@ +import { Effect } from "effect" +import { Question } from "@/question" + +/** + * Helpers for the shared `@/tool/question` tool that surface a dismissed-question + * outcome (from `Question.dismissAll` when a new prompt arrives mid-question) as + * a normal tool result instead of letting `Effect.orDie` turn the + * `QuestionRejectedError` into a defect that kills the in-flight stream. + * + * Extracted here so the shared tool file keeps just a one-liner pipe plus an + * early return, minimising the surface area that conflicts with upstream. + */ +export namespace KiloQuestionTool { + const DISMISSED = "dismissed" as const + type Dismissed = typeof DISMISSED + + export const catchDismissed = (eff: Effect.Effect) => + eff.pipe(Effect.catchTag("QuestionRejectedError", () => Effect.succeed(DISMISSED))) + + export const isDismissed = (v: unknown): v is Dismissed => v === DISMISSED + + export const dismissedResult = () => ({ + title: "Question dismissed", + output: "User dismissed the question.", + metadata: { answers: [] as ReadonlyArray, dismissed: true as const }, + }) +} diff --git a/packages/opencode/src/question/index.ts b/packages/opencode/src/question/index.ts index 7315e0a96c7..4fd6501e90c 100644 --- a/packages/opencode/src/question/index.ts +++ b/packages/opencode/src/question/index.ts @@ -8,6 +8,7 @@ import { Log } from "@/util" import { withStatics } from "@/util/schema" import { QuestionID } from "./schema" import { makeRuntime } from "@/effect/run-service" // kilocode_change +import { KiloQuestion } from "@/kilocode/question" // kilocode_change const log = Log.create({ service: "question" }) @@ -148,6 +149,7 @@ export interface Interface { readonly reply: (input: { requestID: QuestionID; answers: ReadonlyArray }) => Effect.Effect readonly reject: (requestID: QuestionID) => Effect.Effect readonly list: () => Effect.Effect> + readonly dismissAll: (sessionID: SessionID) => Effect.Effect // kilocode_change } export class Service extends Context.Service()("@opencode/Question") {} @@ -193,6 +195,11 @@ export const layer = Layer.effect( blocking: input.blocking, // kilocode_change tool: input.tool, }) + + // kilocode_change start + yield* KiloQuestion.guardFollowup(input.sessionID, () => new RejectedError()) + // kilocode_change end + pending.set(id, { info, deferred }) yield* bus.publish(Event.Asked, info) @@ -245,7 +252,16 @@ export const layer = Layer.effect( return Array.from(pending.values(), (x) => x.info) }) - return Service.of({ ask, reply, reject, list }) + // kilocode_change start - body lives in @/kilocode/question/KiloQuestion.makeDismissAll + const dismissAll = KiloQuestion.makeDismissAll({ + state, + publishRejected: (entry) => + bus.publish(Event.Rejected, { sessionID: entry.info.sessionID, requestID: entry.info.id }), + makeError: () => new RejectedError(), + }) + // kilocode_change end + + return Service.of({ ask, reply, reject, list, dismissAll }) // kilocode_change }), ) @@ -257,6 +273,7 @@ export const list = () => runPromise((svc) => svc.list()) export const ask = (input: Parameters[0]) => runPromise((svc) => svc.ask(input)) export const reply = (input: Parameters[0]) => runPromise((svc) => svc.reply(input)) export const reject = (requestID: QuestionID) => runPromise((svc) => svc.reject(requestID)) +export const dismissAll = (sessionID: string) => runPromise((svc) => svc.dismissAll(SessionID.make(sessionID))) // kilocode_change end export * as Question from "." diff --git a/packages/opencode/src/session/prompt.ts b/packages/opencode/src/session/prompt.ts index c73e625ef0a..0281aaf2816 100644 --- a/packages/opencode/src/session/prompt.ts +++ b/packages/opencode/src/session/prompt.ts @@ -5,6 +5,7 @@ import { KiloSessionPrompt } from "@/kilocode/session/prompt" // kilocode_change import { KiloSessionPromptQueue } from "@/kilocode/session/prompt-queue" // kilocode_change import { KiloSession } from "@/kilocode/session" // kilocode_change import { Suggestion } from "@/kilocode/suggestion" // kilocode_change +import { Question } from "@/question" // kilocode_change import z from "zod" import { SessionID, MessageID, PartID } from "./schema" import { MessageV2 } from "./message-v2" @@ -1285,6 +1286,7 @@ NOTE: At any point in time through this workflow you should feel free to ask the function* (input: PromptInput) { const session = yield* sessions.get(input.sessionID) yield* revert.cleanup(session) + // kilocode_change start - persist queued prompts immediately while serializing each follow-up loop const message = yield* createUserMessage(input) yield* sessions.touch(input.sessionID) @@ -1297,12 +1299,17 @@ NOTE: At any point in time through this workflow you should feel free to ask the yield* sessions.setPermission({ sessionID: session.id, permission: permissions }) } - if (input.noReply === true) return message - // kilocode_change start — dismiss pending suggestions so a previous loop - // blocked on a suggestion can settle before the queue runs the next prompt + // kilocode_change start — unblock tools waiting on user input so any in-flight + // handle.process can return. Adding a new user message is the signal that any + // pending tool prompt is superseded, so we dismiss even on the noReply path. + // Critically we never cancel the in-flight fiber here — that would abort the + // streamText call mid-tokens and cut off the assistant reply. The enqueue call + // below serializes this prompt after the current turn's current LLM step, and + // runLoop checks hasFollowup between steps to break out once it has been + // enqueued during the turn. yield* Effect.promise(() => Suggestion.dismissAll(input.sessionID)) - // kilocode_change end - // kilocode_change start - serialize follow-up loops via queue + yield* Effect.promise(() => Question.dismissAll(input.sessionID)) + if (input.noReply === true) return message return yield* KiloSessionPromptQueue.enqueue( input.sessionID, message.info.id, @@ -1312,6 +1319,7 @@ NOTE: At any point in time through this workflow you should feel free to ask the // kilocode_change end }, ) + // kilocode_change end const lastAssistant = Effect.fnUntraced(function* (sessionID: SessionID) { // kilocode_change start - retry when cancel races before shellImpl writes messages @@ -1604,6 +1612,15 @@ NOTE: At any point in time through this workflow you should feel free to ask the overflow: !handle.message.finish, }) } + // kilocode_change start — break out so a newer queued prompt can take over + // instead of starting another LLM step for the now-superseded turn. The + // current handle.process has fully drained (tokens + inline tool calls) by + // the time we get here, so nothing is cut off. + if (KiloSessionPromptQueue.hasFollowup(sessionID)) { + closeReasons.set(sessionID, "interrupted") + return "break" as const + } + // kilocode_change end return "continue" as const }).pipe(Effect.ensuring(instruction.clear(handle.message.id))) if (outcome === "break") break diff --git a/packages/opencode/src/tool/question.ts b/packages/opencode/src/tool/question.ts index e5bb33aa69f..4b52ebd7cfd 100644 --- a/packages/opencode/src/tool/question.ts +++ b/packages/opencode/src/tool/question.ts @@ -3,6 +3,7 @@ import { Effect } from "effect" import * as Tool from "./tool" import { Question } from "../question" import DESCRIPTION from "./question.txt" +import { KiloQuestionTool } from "@/kilocode/tool/question" // kilocode_change const parameters = z.object({ questions: z.array(Question.Prompt.zod).describe("Questions to ask"), @@ -10,6 +11,7 @@ const parameters = z.object({ type Metadata = { answers: ReadonlyArray + dismissed?: boolean // kilocode_change } export const QuestionTool = Tool.define( @@ -22,11 +24,18 @@ export const QuestionTool = Tool.define, ctx: Tool.Context) => Effect.gen(function* () { - const answers = yield* question.ask({ - sessionID: ctx.sessionID, - questions: params.questions, - tool: ctx.callID ? { messageID: ctx.messageID, callID: ctx.callID } : undefined, - }) + // kilocode_change start - surface Question.dismissAll's RejectedError as a normal + // tool result via KiloQuestionTool helpers, so Effect.orDie below does not turn + // it into a defect and kill the in-flight stream. + const answers = yield* question + .ask({ + sessionID: ctx.sessionID, + questions: params.questions, + tool: ctx.callID ? { messageID: ctx.messageID, callID: ctx.callID } : undefined, + }) + .pipe(KiloQuestionTool.catchDismissed) + if (KiloQuestionTool.isDismissed(answers)) return KiloQuestionTool.dismissedResult() + // kilocode_change end const formatted = params.questions .map((q, i) => `"${q.question}"="${answers[i]?.length ? answers[i].join(", ") : "Unanswered"}"`) diff --git a/packages/opencode/test/kilocode/prompt-dismiss-contract.test.ts b/packages/opencode/test/kilocode/prompt-dismiss-contract.test.ts index 0ce893e44c1..c5981d9f6e7 100644 --- a/packages/opencode/test/kilocode/prompt-dismiss-contract.test.ts +++ b/packages/opencode/test/kilocode/prompt-dismiss-contract.test.ts @@ -1,9 +1,11 @@ /** * Contract test for prompt.ts Kilo-specific invariants. * - * prompt.ts is a shared upstream file. PR #8988 added Suggestion.dismissAll - * there with kilocode_change markers. An upstream merge that restructures - * the prompt handling could silently remove this call — this test catches that. + * prompt.ts is a shared upstream file. The Kilo-specific "new prompt unblocks + * pending suggestions/questions then enqueues without cancelling the in-flight + * stream" behaviour lives inside a kilocode_change block. An upstream merge + * that restructures the prompt handling could silently remove these calls — + * this test catches that. */ import { describe, test, expect } from "bun:test" @@ -18,18 +20,36 @@ describe("prompt.ts Kilo-specific invariants", () => { expect(content).toMatch(/import\s*\{[^}]*Suggestion[^}]*\}\s*from\s*["']@\/kilocode\/suggestion["']/) }) + test("imports Question from the question module", () => { + const content = fs.readFileSync(PROMPT_FILE, "utf-8") + expect(content).toMatch(/import\s*\{[^}]*Question[^}]*\}\s*from\s*["']@\/question["']/) + }) + test("calls Suggestion.dismissAll before restarting the session loop", () => { const content = fs.readFileSync(PROMPT_FILE, "utf-8") expect(content).toContain("Suggestion.dismissAll") }) - test("dismissAll runs before the prompt queue enqueues the new loop", () => { + test("dismissAll for suggestions and questions runs before enqueue, without cancelling the in-flight fiber", () => { const content = fs.readFileSync(PROMPT_FILE, "utf-8") - // dismissAll must precede KiloSessionPromptQueue.enqueue so a previous loop - // blocked on a suggestion can settle before the queue starts the next prompt. + // dismissAll for both suggestions and questions must precede the enqueue so + // an in-flight handle.process blocked on a pending tool prompt can return. + // Critically, the block must NOT call state.cancel or KiloSessionPromptQueue.reserve — + // either of those would abort the running streamText mid-tokens, which was + // the #9332 regression. Order: dismissAll(Suggestion) → dismissAll(Question) → enqueue. const block = content.match( - /kilocode_change start[^\n]*dismiss[\s\S]*?Suggestion\.dismissAll[\s\S]*?kilocode_change end[\s\S]*?KiloSessionPromptQueue\.enqueue/, + /kilocode_change start[^\n]*unblock tools[\s\S]*?Suggestion\.dismissAll[\s\S]*?Question\.dismissAll[\s\S]*?KiloSessionPromptQueue\.enqueue/, ) expect(block).not.toBeNull() + expect(content).not.toMatch(/state\.cancel\(input\.sessionID\)/) + expect(content).not.toMatch(/KiloSessionPromptQueue\.reserve/) + }) + + test("runLoop breaks out between LLM steps when a newer prompt was enqueued", () => { + const content = fs.readFileSync(PROMPT_FILE, "utf-8") + // hasFollowup has to be checked inside runLoop so the current handle.process + // finishes naturally (tokens + inline tool calls) and the next LLM step is + // skipped when a follow-up is already queued. + expect(content).toContain("KiloSessionPromptQueue.hasFollowup(sessionID)") }) }) diff --git a/packages/opencode/test/kilocode/question-dismiss-all.test.ts b/packages/opencode/test/kilocode/question-dismiss-all.test.ts new file mode 100644 index 00000000000..002765bd34c --- /dev/null +++ b/packages/opencode/test/kilocode/question-dismiss-all.test.ts @@ -0,0 +1,174 @@ +import { describe, expect, test } from "bun:test" +import { Effect } from "effect" +import { KiloSessionPromptQueue } from "../../src/kilocode/session/prompt-queue" +import { Instance } from "../../src/project/instance" +import { Question } from "../../src/question" +import { MessageID, SessionID } from "../../src/session/schema" +import { tmpdir } from "../fixture/fixture" + +describe("Question.dismissAll", () => { + test("rejects pending asks for the target session and clears them", async () => { + await using tmp = await tmpdir({ git: true }) + await Instance.provide({ + directory: tmp.path, + fn: async () => { + const sesA = SessionID.make("ses_a") + const sesB = SessionID.make("ses_b") + + const a1 = Question.ask({ + sessionID: sesA, + questions: [ + { + header: "Continue?", + question: "Should I continue?", + options: [ + { label: "Yes", description: "Go" }, + { label: "No", description: "Stop" }, + ], + }, + ], + }).catch((err) => { + if (err instanceof Question.RejectedError) return "rejected" + throw err + }) + + const a2 = Question.ask({ + sessionID: sesA, + questions: [ + { + header: "Retry?", + question: "Try again?", + options: [ + { label: "Retry", description: "Retry" }, + { label: "Cancel", description: "Cancel" }, + ], + }, + ], + }).catch((err) => { + if (err instanceof Question.RejectedError) return "rejected" + throw err + }) + + const b1 = Question.ask({ + sessionID: sesB, + questions: [ + { + header: "Deploy?", + question: "Deploy now?", + options: [ + { label: "Ship", description: "Ship" }, + { label: "Wait", description: "Wait" }, + ], + }, + ], + }).catch((err) => { + if (err instanceof Question.RejectedError) return "rejected-b" + throw err + }) + + // Wait for all three asks to register so we can dismiss them. + for (let i = 0; i < 50; i++) { + if ((await Question.list()).length >= 3) break + await Bun.sleep(10) + } + expect(await Question.list()).toHaveLength(3) + + // Track whether B's promise settles. + let settled = false + b1.then(() => { + settled = true + }) + + await Question.dismissAll("ses_a") + + expect(await a1).toBe("rejected") + expect(await a2).toBe("rejected") + + await new Promise((r) => setTimeout(r, 10)) + expect(settled).toBe(false) + + const remaining = await Question.list() + expect(remaining).toHaveLength(1) + expect(remaining[0]?.sessionID).toBe(sesB) + + await Question.reject(remaining[0]!.id) + expect(await b1).toBe("rejected-b") + }, + }) + }) + + test("is a no-op when no questions exist", async () => { + await using tmp = await tmpdir({ git: true }) + await Instance.provide({ + directory: tmp.path, + fn: async () => { + await Question.dismissAll("ses_missing") + expect(await Question.list()).toEqual([]) + }, + }) + }) + + test("ask rejects immediately when a followup is queued on the session", async () => { + // When a newer prompt has already been enqueued on the session, a tool + // that subsequently calls Question.ask would otherwise block the run until + // the user manually dismisses it. Verify the pre-emptive hasFollowup check + // rejects with RejectedError before any pending entry is registered. + await using tmp = await tmpdir({ git: true }) + await Instance.provide({ + directory: tmp.path, + fn: async () => { + const sessionID = SessionID.make("ses_auto_ask") + const started = Promise.withResolvers() + const release = Promise.withResolvers() + + // Slot 1 stays running so activeSince is pinned to its seq. + const first = Effect.runPromise( + KiloSessionPromptQueue.enqueue( + sessionID, + MessageID.make("message_ask_1"), + Effect.gen(function* () { + started.resolve() + yield* Effect.promise(() => release.promise) + return "first" as const + }), + Effect.succeed("first-cancelled" as const), + ), + ) + await started.promise + + // Slot 2 arrives while slot 1 is active — latest > activeSince. + const second = Effect.runPromise( + KiloSessionPromptQueue.enqueue( + sessionID, + MessageID.make("message_ask_2"), + Effect.succeed("second" as const), + Effect.succeed("second-cancelled" as const), + ), + ) + await Bun.sleep(10) + expect(KiloSessionPromptQueue.hasFollowup(sessionID)).toBe(true) + + await expect( + Question.ask({ + sessionID, + questions: [ + { + header: "Continue?", + question: "Should I continue?", + options: [ + { label: "Yes", description: "Go" }, + { label: "No", description: "Stop" }, + ], + }, + ], + }), + ).rejects.toBeInstanceOf(Question.RejectedError) + expect(await Question.list()).toEqual([]) + + release.resolve() + expect(await first).toBe("first") + expect(await second).toBe("second") + }, + }) + }) +}) diff --git a/packages/opencode/test/kilocode/session-prompt-queue.test.ts b/packages/opencode/test/kilocode/session-prompt-queue.test.ts index 37bc26cf1d5..954f77234df 100644 --- a/packages/opencode/test/kilocode/session-prompt-queue.test.ts +++ b/packages/opencode/test/kilocode/session-prompt-queue.test.ts @@ -1,7 +1,10 @@ import path from "path" import { describe, expect, test } from "bun:test" import { Effect } from "effect" +import { Bus } from "../../src/bus" import { KiloSessionPromptQueue } from "../../src/kilocode/session/prompt-queue" +import { Suggestion } from "../../src/kilocode/suggestion" +import { Question } from "../../src/question" import { ModelID, ProviderID } from "../../src/provider/schema" import { Instance } from "../../src/project/instance" import { Session } from "../../src/session" @@ -221,11 +224,85 @@ describe("session prompt queue", () => { expect(ids[ids.length - 1]).toBe(injected) }) - test("continues a queued prompt after the active run finishes", async () => { + test("hasFollowup reports true only for prompts enqueued after the active slot started", async () => { + const sessionID = SessionID.make("session_followup_semantics") + const observed: Array<{ where: string; value: boolean }> = [] + const firstStarted = Promise.withResolvers() + const firstReleased = Promise.withResolvers() + const secondStarted = Promise.withResolvers() + const secondReleased = Promise.withResolvers() + + const first = Effect.runPromise( + KiloSessionPromptQueue.enqueue( + sessionID, + MessageID.make("message_followup_1"), + Effect.gen(function* () { + observed.push({ where: "first:start", value: KiloSessionPromptQueue.hasFollowup(sessionID) }) + firstStarted.resolve() + yield* Effect.promise(() => firstReleased.promise) + observed.push({ where: "first:end", value: KiloSessionPromptQueue.hasFollowup(sessionID) }) + return "first" + }), + Effect.succeed("first-cancelled"), + ), + ) + + await firstStarted.promise + // msg1 is alone — nothing newer has arrived yet. + expect(observed[0]?.value).toBe(false) + + const second = Effect.runPromise( + KiloSessionPromptQueue.enqueue( + sessionID, + MessageID.make("message_followup_2"), + Effect.gen(function* () { + observed.push({ where: "second:start", value: KiloSessionPromptQueue.hasFollowup(sessionID) }) + secondStarted.resolve() + yield* Effect.promise(() => secondReleased.promise) + return "second" + }), + Effect.succeed("second-cancelled"), + ), + ) + + // Enqueueing msg2 while msg1 is still running must flip hasFollowup to true + // for msg1's running slot. + await new Promise((resolve) => setTimeout(resolve, 10)) + expect(KiloSessionPromptQueue.hasFollowup(sessionID)).toBe(true) + + const third = Effect.runPromise( + KiloSessionPromptQueue.enqueue( + sessionID, + MessageID.make("message_followup_3"), + Effect.sync(() => { + observed.push({ where: "third:start", value: KiloSessionPromptQueue.hasFollowup(sessionID) }) + return "third" + }), + Effect.succeed("third-cancelled"), + ), + ) + + // Let msg1 finish. + firstReleased.resolve() + await first + await secondStarted.promise + + // msg2 started after msg3 was enqueued, so hasFollowup should be false for + // msg2 — everything waiting is older than msg2's activeSince snapshot. + expect(KiloSessionPromptQueue.hasFollowup(sessionID)).toBe(false) + secondReleased.resolve() + + expect(await second).toBe("second") + expect(await third).toBe("third") + + const events = observed.map((item) => `${item.where}=${item.value}`) + expect(events).toEqual(["first:start=false", "first:end=true", "second:start=false", "third:start=false"]) + }) + + test("processes queued prompts without aborting the in-flight stream", async () => { const ready = Promise.withResolvers() - const release = Promise.withResolvers() + const injected = Promise.withResolvers() const calls: number[] = [] - const replies = ["first reply", "second reply", "third reply"] const server = Bun.serve({ port: 0, fetch(req) { @@ -235,8 +312,8 @@ describe("session prompt queue", () => { calls.push(Date.now()) const body = calls.length === 1 - ? reply({ text: replies[0], ready: ready.resolve, wait: release.promise }) - : reply({ text: replies[calls.length - 1] ?? "extra reply" }) + ? reply({ text: "first reply", ready: ready.resolve }) + : reply({ text: "second reply", ready: injected.resolve }) return new Response(body, { status: 200, headers: { "Content-Type": "text/event-stream" }, @@ -288,43 +365,49 @@ describe("session prompt queue", () => { agent: "code", parts: [{ type: "text", text: "second prompt" }], }) - const third = SessionPrompt.prompt({ - sessionID: session.id, - agent: "code", - parts: [{ type: "text", text: "third prompt" }], - }) - await Bun.sleep(20) - expect(calls).toHaveLength(1) - const queued = await Session.messages({ sessionID: session.id }) - expect(queued.filter((msg) => msg.info.role === "user")).toHaveLength(3) - expect(queued.filter((msg) => msg.info.role === "assistant")).toHaveLength(1) - - release.resolve() - await first + const one = await first + await injected.promise const two = await second - const three = await third + expect(calls).toHaveLength(2) + + // The in-flight stream must complete; no aborted error on msg1's reply. + expect(one.info.role).toBe("assistant") + if (one.info.role === "assistant") expect(one.info.error).toBeUndefined() + expect(hasText(one, "first reply")).toBe(true) expect(hasText(two, "second reply")).toBe(true) - expect(hasText(three, "third reply")).toBe(true) - expect(calls).toHaveLength(3) const msgs = await Session.messages({ sessionID: session.id }) const users = msgs.filter((msg) => msg.info.role === "user") const assistants = msgs.filter((msg) => msg.info.role === "assistant") + const prompts = users.flatMap((msg) => + msg.parts.filter((part) => part.type === "text").map((part) => part.text), + ) const text = assistants.flatMap((msg) => msg.parts.filter((part) => part.type === "text").map((part) => part.text), ) - expect(users).toHaveLength(3) - expect(assistants).toHaveLength(3) + expect(users).toHaveLength(2) + expect(assistants).toHaveLength(2) + expect(prompts).toContain("first prompt") + expect(prompts).toContain("second prompt") expect(text).toContain("first reply") expect(text).toContain("second reply") - expect(text).toContain("third reply") - for (const [index, item] of assistants.entries()) { - const user = users[index]?.info - if (item.info.role !== "assistant" || user?.role !== "user") throw new Error("missing turn") - expect(item.info.parentID).toBe(user.id) + + const firstUser = users.find((msg) => hasText(msg, "first prompt")) + const secondUser = users.find((msg) => hasText(msg, "second prompt")) + const firstReply = assistants.find((msg) => hasText(msg, "first reply")) + const secondReply = assistants.find((msg) => hasText(msg, "second reply")) + if ( + firstUser?.info.role !== "user" || + secondUser?.info.role !== "user" || + firstReply?.info.role !== "assistant" || + secondReply?.info.role !== "assistant" + ) { + throw new Error("missing expected messages") } + expect(firstReply.info.parentID).toBe(firstUser.info.id) + expect(secondReply.info.parentID).toBe(secondUser.info.id) }, }) } finally { @@ -392,12 +475,14 @@ describe("session prompt queue", () => { parts: [{ type: "text", text: "third prompt" }], }) + // Let msg2/msg3's enqueue capture the current version before cancel bumps it. await Bun.sleep(20) expect(calls).toHaveLength(1) await SessionPrompt.cancel(session.id) await Promise.all([first, second, third]) + // The queued prompts must never reach the LLM once cancel flushes the queue. expect(calls).toHaveLength(1) const msgs = await Session.messages({ sessionID: session.id }) const assistants = msgs.filter((msg) => msg.info.role === "assistant") @@ -414,10 +499,241 @@ describe("session prompt queue", () => { ), ) expect(ids).toEqual([]) + expect(KiloSessionPromptQueue.hasFollowup(session.id)).toBe(false) }, }) } finally { server.stop(true) } }) + + test("new prompt dismisses a pending suggestion", async () => { + const shown = Promise.withResolvers() + const dismissed = Promise.withResolvers() + await using tmp = await tmpdir({ git: true }) + + await Instance.provide({ + directory: tmp.path, + fn: async () => { + const session = await Session.create({ title: "Suggestion unblock regression" }) + const offShown = Bus.subscribe(Suggestion.Event.Shown, (event) => { + if (event.properties.sessionID === session.id) shown.resolve() + }) + const offDismissed = Bus.subscribe(Suggestion.Event.Dismissed, (event) => { + if (event.properties.sessionID === session.id) dismissed.resolve() + }) + + try { + const base = Suggestion.show({ + sessionID: session.id, + text: "Run review?", + actions: [{ label: "Review", prompt: "/local-review-uncommitted" }], + }).catch((err) => { + if (err instanceof Suggestion.DismissedError) return "dismissed" + throw err + }) + + await shown.promise + await SessionPrompt.prompt({ + sessionID: session.id, + agent: "code", + parts: [{ type: "text", text: "replacement prompt" }], + noReply: true, + }) + await dismissed.promise + + expect(await base).toBe("dismissed") + expect(await Suggestion.list()).toEqual([]) + } finally { + offShown() + offDismissed() + } + }, + }) + }) + + test("new prompt dismisses a pending question", async () => { + const asked = Promise.withResolvers() + const rejected = Promise.withResolvers() + await using tmp = await tmpdir({ git: true }) + + await Instance.provide({ + directory: tmp.path, + fn: async () => { + const session = await Session.create({ title: "Question unblock regression" }) + const offAsked = Bus.subscribe(Question.Event.Asked, (event) => { + if (event.properties.sessionID === session.id) asked.resolve() + }) + const offRejected = Bus.subscribe(Question.Event.Rejected, (event) => { + if (event.properties.sessionID === session.id) rejected.resolve() + }) + + try { + const pending = Question.ask({ + sessionID: session.id, + questions: [ + { + header: "Continue?", + question: "Should I continue?", + options: [ + { label: "Yes", description: "Go ahead" }, + { label: "No", description: "Stop" }, + ], + }, + ], + }).catch((err) => { + if (err instanceof Question.RejectedError) return "rejected" + throw err + }) + + await asked.promise + await SessionPrompt.prompt({ + sessionID: session.id, + agent: "code", + parts: [{ type: "text", text: "replacement prompt" }], + noReply: true, + }) + await rejected.promise + + expect(await pending).toBe("rejected") + expect(await Question.list()).toEqual([]) + } finally { + offAsked() + offRejected() + } + }, + }) + }) + + test("auto-dismisses a suggestion shown after a queued prompt", async () => { + // Reverse ordering of the "new prompt dismisses a pending suggestion" test: + // queue the follow-up first, then open the blocker. Suggestion.show must see + // hasFollowup=true and reject synchronously, before any pending entry or + // Shown event is published. + await using tmp = await tmpdir({ git: true }) + await Instance.provide({ + directory: tmp.path, + fn: async () => { + const sessionID = SessionID.make("ses_auto_suggestion") + const started = Promise.withResolvers() + const release = Promise.withResolvers() + + // Slot 1: active, activeSince snapshots latest=1. + const first = Effect.runPromise( + KiloSessionPromptQueue.enqueue( + sessionID, + MessageID.make("message_auto_sug_1"), + Effect.gen(function* () { + started.resolve() + yield* Effect.promise(() => release.promise) + return "first" as const + }), + Effect.succeed("first-cancelled" as const), + ), + ) + await started.promise + + // Slot 2: enqueued while slot 1 is active → latest=2 > activeSince=1. + const second = Effect.runPromise( + KiloSessionPromptQueue.enqueue( + sessionID, + MessageID.make("message_auto_sug_2"), + Effect.succeed("second" as const), + Effect.succeed("second-cancelled" as const), + ), + ) + await Bun.sleep(10) + expect(KiloSessionPromptQueue.hasFollowup(sessionID)).toBe(true) + + let shown = 0 + const offShown = Bus.subscribe(Suggestion.Event.Shown, (event) => { + if (event.properties.sessionID === sessionID) shown++ + }) + try { + await expect( + Suggestion.show({ + sessionID, + text: "Run review?", + actions: [{ label: "Review", prompt: "/local-review-uncommitted" }], + }), + ).rejects.toBeInstanceOf(Suggestion.DismissedError) + } finally { + offShown() + } + expect(shown).toBe(0) + expect(await Suggestion.list()).toEqual([]) + + release.resolve() + expect(await first).toBe("first") + expect(await second).toBe("second") + }, + }) + }) + + test("auto-dismisses a question shown after a queued prompt", async () => { + await using tmp = await tmpdir({ git: true }) + await Instance.provide({ + directory: tmp.path, + fn: async () => { + const sessionID = SessionID.make("ses_auto_question") + const started = Promise.withResolvers() + const release = Promise.withResolvers() + + const first = Effect.runPromise( + KiloSessionPromptQueue.enqueue( + sessionID, + MessageID.make("message_auto_q_1"), + Effect.gen(function* () { + started.resolve() + yield* Effect.promise(() => release.promise) + return "first" as const + }), + Effect.succeed("first-cancelled" as const), + ), + ) + await started.promise + + const second = Effect.runPromise( + KiloSessionPromptQueue.enqueue( + sessionID, + MessageID.make("message_auto_q_2"), + Effect.succeed("second" as const), + Effect.succeed("second-cancelled" as const), + ), + ) + await Bun.sleep(10) + expect(KiloSessionPromptQueue.hasFollowup(sessionID)).toBe(true) + + let asked = 0 + const offAsked = Bus.subscribe(Question.Event.Asked, (event) => { + if (event.properties.sessionID === sessionID) asked++ + }) + try { + await expect( + Question.ask({ + sessionID, + questions: [ + { + header: "Continue?", + question: "Should I continue?", + options: [ + { label: "Yes", description: "Go ahead" }, + { label: "No", description: "Stop" }, + ], + }, + ], + }), + ).rejects.toBeInstanceOf(Question.RejectedError) + } finally { + offAsked() + } + expect(asked).toBe(0) + expect(await Question.list()).toEqual([]) + + release.resolve() + expect(await first).toBe("first") + expect(await second).toBe("second") + }, + }) + }) }) diff --git a/packages/opencode/test/kilocode/suggestion/auto-dismiss.test.ts b/packages/opencode/test/kilocode/suggestion/auto-dismiss.test.ts new file mode 100644 index 00000000000..d8c6a19122f --- /dev/null +++ b/packages/opencode/test/kilocode/suggestion/auto-dismiss.test.ts @@ -0,0 +1,65 @@ +import { describe, expect, test } from "bun:test" +import { Effect } from "effect" +import { KiloSessionPromptQueue } from "../../../src/kilocode/session/prompt-queue" +import { Suggestion } from "../../../src/kilocode/suggestion" +import { Instance } from "../../../src/project/instance" +import { MessageID, SessionID } from "../../../src/session/schema" +import { tmpdir } from "../../fixture/fixture" + +describe("Suggestion.show auto-dismiss on queued followup", () => { + test("show rejects immediately when a followup is queued on the session", async () => { + // A tool that calls Suggestion.show after a queued prompt has arrived would + // otherwise block the turn on user input. Verify the pre-emptive + // hasFollowup check rejects with DismissedError before any pending entry + // is registered or a Shown event is published. + await using tmp = await tmpdir({ git: true }) + await Instance.provide({ + directory: tmp.path, + fn: async () => { + const sessionID = SessionID.make("ses_auto_show") + const started = Promise.withResolvers() + const release = Promise.withResolvers() + + // Slot 1 stays running so activeSince is pinned to its seq. + const first = Effect.runPromise( + KiloSessionPromptQueue.enqueue( + sessionID, + MessageID.make("message_show_1"), + Effect.gen(function* () { + started.resolve() + yield* Effect.promise(() => release.promise) + return "first" as const + }), + Effect.succeed("first-cancelled" as const), + ), + ) + await started.promise + + // Slot 2 arrives while slot 1 is active — latest > activeSince. + const second = Effect.runPromise( + KiloSessionPromptQueue.enqueue( + sessionID, + MessageID.make("message_show_2"), + Effect.succeed("second" as const), + Effect.succeed("second-cancelled" as const), + ), + ) + await Bun.sleep(10) + expect(KiloSessionPromptQueue.hasFollowup(sessionID)).toBe(true) + + await expect( + Suggestion.show({ + sessionID, + text: "Run review?", + actions: [{ label: "Review", prompt: "/local-review-uncommitted" }], + }), + ).rejects.toBeInstanceOf(Suggestion.DismissedError) + expect(await Suggestion.list()).toEqual([]) + + release.resolve() + expect(await first).toBe("first") + expect(await second).toBe("second") + }, + }) + }) +})