diff --git a/.changeset/abort-moved-sessions.md b/.changeset/abort-moved-sessions.md deleted file mode 100644 index 287b8a8823..0000000000 --- a/.changeset/abort-moved-sessions.md +++ /dev/null @@ -1,5 +0,0 @@ ---- -"kilo-code": patch ---- - -Stop active Agent Manager turns after moving or continuing their sessions in another worktree. diff --git a/.changeset/abort-starting-sessions.md b/.changeset/abort-starting-sessions.md new file mode 100644 index 0000000000..99426d6997 --- /dev/null +++ b/.changeset/abort-starting-sessions.md @@ -0,0 +1,6 @@ +--- +"kilo-code": patch +"@kilocode/cli": patch +--- + +Allow Escape to stop Agent Manager prompts while their sessions are still starting. diff --git a/packages/kilo-vscode/src/KiloProvider.ts b/packages/kilo-vscode/src/KiloProvider.ts index 4a89ebc917..c393d27912 100644 --- a/packages/kilo-vscode/src/KiloProvider.ts +++ b/packages/kilo-vscode/src/KiloProvider.ts @@ -71,12 +71,7 @@ import { fetchMessagePage, MESSAGE_PAGE_LIMIT } from "./kilo-provider/message-pa import { childID } from "./kilo-provider/task-session" import { VisibleTaskStreams } from "./kilo-provider/visible-task-streams" import { handleNetworkEvent, clearNetworkWaits } from "./kilo-provider/network" -import { - abortSession, - resolveAbortDirectories, - updateActiveSessionDirectory, - type ActiveSessionDirectories, -} from "./kilo-provider/abort" +import { abortSession } from "./kilo-provider/abort" import { buildAutocompleteSettingsMessage, validAutocompleteSetting, @@ -286,7 +281,6 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper private syncedChildSessions: Set = new Set() private sessionStatusMap = new Map() // Latest status used for destructive config warnings. private sessionDirectories = new Map() // Per-session directory overrides, such as Agent Manager worktrees. - private activeSessionDirectories: ActiveSessionDirectories = new Map() private permissionDirectories = new Map() private projectID: string | undefined // Current workspace project ID used to filter sessions. private loadMessagesAbort: AbortController | null = null // Current load request cancellation. @@ -634,28 +628,10 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper * When set, all operations for this session use this directory instead of the workspace root. */ public setSessionDirectory(sessionId: string, directory: string): void { - const status = this.sessionStatusMap.get(sessionId) - if (status && status !== "idle" && !this.activeSessionDirectories.has(sessionId)) { - updateActiveSessionDirectory({ - active: this.activeSessionDirectories, - sessionID: sessionId, - status, - dir: this.getWorkspaceDirectory(sessionId), - }) - } this.sessionDirectories.set(sessionId, directory) } public clearSessionDirectory(sessionId: string): void { - const status = this.sessionStatusMap.get(sessionId) - if (status && status !== "idle" && !this.activeSessionDirectories.has(sessionId)) { - updateActiveSessionDirectory({ - active: this.activeSessionDirectories, - sessionID: sessionId, - status, - dir: this.getWorkspaceDirectory(sessionId), - }) - } this.sessionDirectories.delete(sessionId) } @@ -1732,7 +1708,6 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper this.visibleTaskStreams.delete(sessionID) this.syncedChildSessions.delete(sessionID) this.sessionDirectories.delete(sessionID) - this.activeSessionDirectories.delete(sessionID) this.lastReconciledAt.delete(sessionID) this.connectionService.pruneSession(sessionID) if (this.currentSession?.id === sessionID) { @@ -2730,33 +2705,23 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper } private async handleAbort(sessionID?: string): Promise { - const client = this.client - if (!client) return + if (!this.client) { + return + } const targetSessionID = sessionID || this.currentSession?.id if (!targetSessionID) { return } - const dirs = resolveAbortDirectories( - this.activeSessionDirectories, - targetSessionID, - this.getWorkspaceDirectory(targetSessionID), - ) - const results = await Promise.allSettled( - dirs.map((dir) => - abortSession({ - client, - sessionID: targetSessionID, - dir, - }), - ), - ) - const failures = results.flatMap((result, index) => - result.status === "rejected" ? [{ dir: dirs[index], error: result.reason }] : [], - ) - if (failures.length > 0) { - console.error("[Kilo New] KiloProvider: Failed to abort session in one or more directories:", failures) + try { + await abortSession({ + client: this.client, + sessionID: targetSessionID, + dir: this.getWorkspaceDirectory(targetSessionID), + }) + } catch (error) { + console.error("[Kilo New] KiloProvider: Failed to abort session:", error) } } @@ -3131,12 +3096,6 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper if (event.type === "session.status") { const sid = event.properties.sessionID this.sessionStatusMap.set(sid, event.properties.status.type) - updateActiveSessionDirectory({ - active: this.activeSessionDirectories, - sessionID: sid, - status: event.properties.status.type, - dir: directory, - }) const msg = mapSSEEventToWebviewMessage(event, sid) if (msg) { this.streams.flush(sid) @@ -3169,17 +3128,6 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper if (event.type === "server.instance.disposed") { const props = event.properties as Record | null const dir = typeof props?.directory === "string" ? props.directory : undefined - if (dir) { - for (const sessionID of [...this.activeSessionDirectories.keys()]) { - updateActiveSessionDirectory({ - active: this.activeSessionDirectories, - sessionID, - status: "idle", - dir, - }) - if (!this.activeSessionDirectories.has(sessionID)) this.sessionStatusMap.set(sessionID, "idle") - } - } if (dir && !sameDirectory(dir, this.getWorkspaceDirectory())) return void this.reloadAfterAuthChange() return @@ -3579,7 +3527,6 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper this.trackedSessionIds.clear() this.syncedChildSessions.clear() this.sessionDirectories.clear() - this.activeSessionDirectories.clear() this.permissionDirectories.clear() this.sessionStatusMap.clear() this.ignoreController?.dispose() diff --git a/packages/kilo-vscode/src/kilo-provider/abort.ts b/packages/kilo-vscode/src/kilo-provider/abort.ts index fa0eeb3edf..8d3d00676e 100644 --- a/packages/kilo-vscode/src/kilo-provider/abort.ts +++ b/packages/kilo-vscode/src/kilo-provider/abort.ts @@ -1,37 +1,4 @@ -import type { KiloClient, SessionStatus } from "@kilocode/sdk/v2/client" -import { sameDirectory } from "../kilo-provider-utils" - -export type ActiveSessionDirectories = Map> - -export function updateActiveSessionDirectory(input: { - active: ActiveSessionDirectories - sessionID: string - status: SessionStatus["type"] - dir?: string -}) { - const target = input.dir - if (!target) return - const dirs = input.active.get(input.sessionID) - if (input.status === "idle") { - if (!dirs) return - for (const dir of dirs) { - if (sameDirectory(dir, target)) dirs.delete(dir) - } - if (dirs.size === 0) input.active.delete(input.sessionID) - return - } - if (!dirs) { - input.active.set(input.sessionID, new Set([target])) - return - } - if (![...dirs].some((dir) => sameDirectory(dir, target))) dirs.add(target) -} - -export function resolveAbortDirectories(active: ActiveSessionDirectories, sessionID: string, fallback: string) { - const dirs = [...(active.get(sessionID) ?? [])] - if (!dirs.some((dir) => sameDirectory(dir, fallback))) dirs.push(fallback) - return dirs -} +import type { KiloClient } from "@kilocode/sdk/v2/client" export async function abortSession(input: { client: KiloClient; sessionID: string; dir: string }) { await input.client.session.abort({ sessionID: input.sessionID, directory: input.dir }, { throwOnError: true }) diff --git a/packages/kilo-vscode/tests/unit/abort-state.test.ts b/packages/kilo-vscode/tests/unit/abort-state.test.ts new file mode 100644 index 0000000000..1b37396b2d --- /dev/null +++ b/packages/kilo-vscode/tests/unit/abort-state.test.ts @@ -0,0 +1,64 @@ +import { describe, expect, it } from "bun:test" +import { createAbortState } from "../../webview-ui/src/context/abort-state" + +describe("pending prompt abort state", () => { + it("waits for the same submission to become cancellable", () => { + const aborts = createAbortState() + + expect(aborts.request("draft", "idle", "message")).toBe(false) + expect(aborts.update("draft", "busy")).toBe(true) + expect(aborts.update("draft", "busy")).toBe(false) + expect(aborts.update("draft", "idle")).toBe(false) + expect(aborts.update("draft", "busy")).toBe(false) + }) + + it("moves pending cancellation to the created session", () => { + const aborts = createAbortState() + + expect(aborts.request("draft", "idle", "message")).toBe(false) + aborts.move("draft", "session") + + expect(aborts.update("draft", "busy")).toBe(false) + expect(aborts.update("session", "busy")).toBe(true) + }) + + it("does not retain cancellation after an idle terminal status", () => { + const aborts = createAbortState() + + expect(aborts.request("session", "idle", "message")).toBe(false) + expect(aborts.update("session", "idle")).toBe(false) + expect(aborts.update("session", "busy")).toBe(false) + }) + + it("allows retrying an abort while the session remains active", () => { + const aborts = createAbortState() + + expect(aborts.request("session", "busy")).toBe(true) + expect(aborts.request("session", "busy")).toBe(true) + expect(aborts.update("session", "idle")).toBe(false) + expect(aborts.request("session", "busy")).toBe(true) + }) + + it("clears cancellation when the matching submission finishes", () => { + const aborts = createAbortState() + + expect(aborts.request("session", "idle", "message")).toBe(false) + aborts.finish("other") + expect(aborts.update("session", "busy")).toBe(true) + + expect(aborts.update("session", "idle")).toBe(false) + expect(aborts.request("session", "idle", "message")).toBe(false) + aborts.finish("message") + expect(aborts.update("session", "busy")).toBe(false) + }) + + it("preserves active destination state during duplicate draft migration", () => { + const aborts = createAbortState() + + expect(aborts.request("draft", "idle", "message")).toBe(false) + expect(aborts.request("session", "busy")).toBe(true) + aborts.move("draft", "session") + + expect(aborts.request("session", "busy")).toBe(true) + }) +}) diff --git a/packages/kilo-vscode/tests/unit/abort.test.ts b/packages/kilo-vscode/tests/unit/abort.test.ts index 2c182cbfe8..126333e6a4 100644 --- a/packages/kilo-vscode/tests/unit/abort.test.ts +++ b/packages/kilo-vscode/tests/unit/abort.test.ts @@ -1,11 +1,6 @@ import { describe, expect, it } from "bun:test" import type { KiloClient } from "@kilocode/sdk/v2/client" -import { - abortSession, - resolveAbortDirectories, - updateActiveSessionDirectory, - type ActiveSessionDirectories, -} from "../../src/kilo-provider/abort" +import { abortSession } from "../../src/kilo-provider/abort" function client(calls: unknown[], fail = false) { return { @@ -19,58 +14,6 @@ function client(calls: unknown[], fail = false) { } as unknown as KiloClient } -describe("active session directories", () => { - function update(active: ActiveSessionDirectories, status: "busy" | "retry" | "idle", dir?: string) { - updateActiveSessionDirectory({ active, sessionID: "session_1", status, dir }) - } - - it("includes the active owner and current session directory", () => { - const active: ActiveSessionDirectories = new Map() - update(active, "busy", "/repo/source") - - expect(resolveAbortDirectories(active, "session_1", "/repo/worktree")).toEqual(["/repo/source", "/repo/worktree"]) - }) - - it("falls back to the current session directory after the active turn becomes idle", () => { - const active: ActiveSessionDirectories = new Map() - update(active, "busy", "/repo/source") - update(active, "idle", "/repo/source") - - expect(resolveAbortDirectories(active, "session_1", "/repo/worktree")).toEqual(["/repo/worktree"]) - }) - - it("retains active directories when an unrelated instance reports idle", () => { - const active: ActiveSessionDirectories = new Map() - update(active, "retry", "/repo/source") - update(active, "idle", "/repo/worktree") - update(active, "idle") - - expect(resolveAbortDirectories(active, "session_1", "/repo/worktree")).toEqual(["/repo/source", "/repo/worktree"]) - }) - - it("tracks concurrent instances and ignores delayed idle events from the old one", () => { - const active: ActiveSessionDirectories = new Map() - update(active, "busy", "/repo/source") - update(active, "busy", "/repo/worktree") - - expect(resolveAbortDirectories(active, "session_1", "/repo/fallback")).toEqual([ - "/repo/source", - "/repo/worktree", - "/repo/fallback", - ]) - - update(active, "idle", "/repo/source") - expect(resolveAbortDirectories(active, "session_1", "/repo/fallback")).toEqual(["/repo/worktree", "/repo/fallback"]) - }) - - it("deduplicates the current directory when it is already active", () => { - const active: ActiveSessionDirectories = new Map() - update(active, "busy", "/repo/worktree") - - expect(resolveAbortDirectories(active, "session_1", "/repo/worktree/.")).toEqual(["/repo/worktree"]) - }) -}) - describe("abortSession", () => { it("calls session.abort with the session id and directory", async () => { const calls: unknown[] = [] diff --git a/packages/kilo-vscode/tests/unit/kilo-provider-load-messages.test.ts b/packages/kilo-vscode/tests/unit/kilo-provider-load-messages.test.ts index f8f44116d8..a63808d556 100644 --- a/packages/kilo-vscode/tests/unit/kilo-provider-load-messages.test.ts +++ b/packages/kilo-vscode/tests/unit/kilo-provider-load-messages.test.ts @@ -1,4 +1,4 @@ -import { describe, it, expect, spyOn } from "bun:test" +import { describe, it, expect } from "bun:test" // vscode mock is provided by the shared preload (tests/setup/vscode-mock.ts) const { KiloProvider } = await import("../../src/KiloProvider") @@ -43,15 +43,12 @@ function createClient(options?: { deleteDeferred?: Deferred sessionData?: unknown sessionGet?: (params: { sessionID: string; directory?: string }) => Promise<{ data: unknown }> - abortFailures?: string[] }) { const calls: { before?: string; limit?: number }[] = [] const stopped: { sessionID: string; directory?: string }[] = [] - const aborted: { sessionID: string; directory?: string }[] = [] return { calls, stopped, - aborted, session: { list: async () => ({ data: [] }), get: async (params: { sessionID: string; directory?: string }) => { @@ -59,11 +56,6 @@ function createClient(options?: { return { data: options?.sessionData ?? null } }, status: async () => ({ data: {} }), - abort: async (params: { sessionID: string; directory?: string }) => { - aborted.push(params) - if (params.directory && options?.abortFailures?.includes(params.directory)) throw new Error("abort failed") - return { data: true } - }, messages: async (params: { before?: string; limit?: number }) => { calls.push({ before: params.before, limit: params.limit }) if (options?.messagesDeferred) return options.messagesDeferred.promise @@ -122,11 +114,9 @@ type ProviderInternals = { currentSession: { id: string; directory?: string } | null contextSessionID: string | undefined sessionDirectories: Map - activeSessionDirectories: Map> trackedSessionIds: Set stopCurrentSessionProcesses: (next?: string) => void - handleEvent: (event: unknown, directory?: string) => void - handleAbort: (sid?: string) => Promise + handleEvent: (event: unknown) => void handleLoadMessages: (sid: string, opts?: { mode?: string; before?: string; limit?: number }) => Promise handleDeleteSession: (sid: string) => Promise } @@ -145,108 +135,6 @@ function makeProvider(client: ReturnType) { return { provider, internal, sent } } -describe("KiloProvider.handleAbort", () => { - it("keeps the active owner when a running session moves to a worktree", async () => { - const client = createClient() - const { provider, internal } = makeProvider(client) - internal.handleEvent({ - type: "session.status", - properties: { sessionID: "s1", status: { type: "busy" } }, - }) - provider.setSessionDirectory("s1", "/repo/worktree") - - await internal.handleAbort("s1") - - expect(client.aborted).toEqual([ - { sessionID: "s1", directory: "/repo" }, - { sessionID: "s1", directory: "/repo/worktree" }, - ]) - }) - - it("preserves a directory-qualified owner instead of inferring the current directory", async () => { - const client = createClient() - const { provider, internal } = makeProvider(client) - internal.handleEvent( - { - type: "session.status", - properties: { sessionID: "s1", status: { type: "busy" } }, - }, - "/repo/source", - ) - provider.setSessionDirectory("s1", "/repo/worktree") - - await internal.handleAbort("s1") - - expect(client.aborted).toEqual([ - { sessionID: "s1", directory: "/repo/source" }, - { sessionID: "s1", directory: "/repo/worktree" }, - ]) - }) - - it("falls back to the current directory after the old owner becomes idle", async () => { - const client = createClient() - const { provider, internal } = makeProvider(client) - internal.handleEvent({ - type: "session.status", - properties: { sessionID: "s1", status: { type: "busy" } }, - }) - provider.setSessionDirectory("s1", "/repo/worktree") - internal.handleEvent( - { - type: "session.status", - properties: { sessionID: "s1", status: { type: "idle" } }, - }, - "/repo", - ) - - await internal.handleAbort("s1") - - expect(client.aborted).toEqual([{ sessionID: "s1", directory: "/repo/worktree" }]) - }) - - it("deduplicates the current directory when it is also an active owner", async () => { - const client = createClient() - const { provider, internal } = makeProvider(client) - internal.trackedSessionIds.add("s1") - provider.setSessionDirectory("s1", "/repo/worktree") - internal.handleEvent( - { - type: "session.status", - properties: { sessionID: "s1", status: { type: "busy" } }, - }, - "/repo/worktree", - ) - - await internal.handleAbort("s1") - - expect(client.aborted).toEqual([{ sessionID: "s1", directory: "/repo/worktree" }]) - }) - - it("attempts every owner when one abort request fails", async () => { - const error = spyOn(console, "error").mockImplementation(() => {}) - const client = createClient({ abortFailures: ["/repo"] }) - const { provider, internal } = makeProvider(client) - internal.trackedSessionIds.add("s1") - internal.handleEvent( - { - type: "session.status", - properties: { sessionID: "s1", status: { type: "busy" } }, - }, - "/repo", - ) - provider.setSessionDirectory("s1", "/repo/worktree") - - await internal.handleAbort("s1") - - expect(client.aborted).toEqual([ - { sessionID: "s1", directory: "/repo" }, - { sessionID: "s1", directory: "/repo/worktree" }, - ]) - expect(error).toHaveBeenCalledTimes(1) - error.mockRestore() - }) -}) - describe("KiloProvider.handleLoadMessages / focus mode freshness", () => { it("stops background processes for the previous session when switching sessions", async () => { const client = createClient({ diff --git a/packages/kilo-vscode/tests/unit/prompt-input-utils.test.ts b/packages/kilo-vscode/tests/unit/prompt-input-utils.test.ts index 9ccd644809..f38a45a179 100644 --- a/packages/kilo-vscode/tests/unit/prompt-input-utils.test.ts +++ b/packages/kilo-vscode/tests/unit/prompt-input-utils.test.ts @@ -171,30 +171,34 @@ describe("isPromptBlocked", () => { describe("isPromptBusy", () => { it("returns true when busy and neither suggesting nor questioning", () => { - expect(isPromptBusy("busy", false, false)).toBe(true) + expect(isPromptBusy("busy", false, false, false)).toBe(true) + }) + + it("returns true while submitting before the backend reports busy", () => { + expect(isPromptBusy("idle", false, false, true)).toBe(true) }) it("returns false when idle regardless of suggesting/questioning", () => { - expect(isPromptBusy("idle", false, false)).toBe(false) - expect(isPromptBusy("idle", true, false)).toBe(false) - expect(isPromptBusy("idle", false, true)).toBe(false) - expect(isPromptBusy("idle", true, true)).toBe(false) + expect(isPromptBusy("idle", false, false, false)).toBe(false) + expect(isPromptBusy("idle", true, false, false)).toBe(false) + expect(isPromptBusy("idle", false, true, false)).toBe(false) + expect(isPromptBusy("idle", true, true, false)).toBe(false) }) it("returns false when busy but suggesting is true (suggestion decoupling)", () => { - expect(isPromptBusy("busy", true, false)).toBe(false) + expect(isPromptBusy("busy", true, false, false)).toBe(false) }) it("returns false when busy but questioning is true (question decoupling)", () => { - expect(isPromptBusy("busy", false, true)).toBe(false) + expect(isPromptBusy("busy", false, true, false)).toBe(false) }) it("returns false when busy and both suggesting and questioning", () => { - expect(isPromptBusy("busy", true, true)).toBe(false) + expect(isPromptBusy("busy", true, true, false)).toBe(false) }) it("returns true for non-idle non-busy status when not suggesting/questioning", () => { - expect(isPromptBusy("retry", false, false)).toBe(true) + expect(isPromptBusy("retry", false, false, false)).toBe(true) }) }) diff --git a/packages/kilo-vscode/webview-ui/src/components/chat/ChatView.tsx b/packages/kilo-vscode/webview-ui/src/components/chat/ChatView.tsx index a0d09a9a62..3cb25864e6 100644 --- a/packages/kilo-vscode/webview-ui/src/components/chat/ChatView.tsx +++ b/packages/kilo-vscode/webview-ui/src/components/chat/ChatView.tsx @@ -82,7 +82,7 @@ export const ChatView: Component = (props) => { onMount(() => { if (props.readonly) return const handler = (e: KeyboardEvent) => { - if (e.key !== "Escape" || session.status() === "idle" || e.defaultPrevented) return + if (e.key !== "Escape" || (!session.submitting() && session.status() === "idle") || e.defaultPrevented) return e.preventDefault() session.abort() } diff --git a/packages/kilo-vscode/webview-ui/src/components/chat/PromptInput.tsx b/packages/kilo-vscode/webview-ui/src/components/chat/PromptInput.tsx index 898301bc4e..2de7d0eb8c 100644 --- a/packages/kilo-vscode/webview-ui/src/components/chat/PromptInput.tsx +++ b/packages/kilo-vscode/webview-ui/src/components/chat/PromptInput.tsx @@ -344,7 +344,8 @@ export const PromptInput: Component = (props) => { window.addEventListener("exportSessionTranscript", onExport) onCleanup(() => window.removeEventListener("exportSessionTranscript", onExport)) - const isBusy = () => isPromptBusy(session.status(), !!props.suggesting?.(), !!props.questioning?.()) + const isBusy = () => + isPromptBusy(session.status(), !!props.suggesting?.(), !!props.questioning?.(), session.submitting()) const isDisabled = () => !server.isConnected() const canUseSpeech = () => canUseSpeechToText(config(), provider.connected(), server.profileData()) const speechModel = () => selectedSpeechToTextModel(config()) diff --git a/packages/kilo-vscode/webview-ui/src/components/chat/prompt-input-utils.ts b/packages/kilo-vscode/webview-ui/src/components/chat/prompt-input-utils.ts index 22fb8db90a..b78b97624c 100644 --- a/packages/kilo-vscode/webview-ui/src/components/chat/prompt-input-utils.ts +++ b/packages/kilo-vscode/webview-ui/src/components/chat/prompt-input-utils.ts @@ -87,8 +87,8 @@ export function isPromptBlocked(permissions: number): boolean { * Returns false (idle-like) when the session is busy only because * a suggestion or question tool call is pending. */ -export function isPromptBusy(status: string, suggesting: boolean, questioning: boolean): boolean { - return status !== "idle" && !suggesting && !questioning +export function isPromptBusy(status: string, suggesting: boolean, questioning: boolean, submitting: boolean): boolean { + return submitting || (status !== "idle" && !suggesting && !questioning) } /** diff --git a/packages/kilo-vscode/webview-ui/src/context/abort-state.ts b/packages/kilo-vscode/webview-ui/src/context/abort-state.ts new file mode 100644 index 0000000000..ca0c12887e --- /dev/null +++ b/packages/kilo-vscode/webview-ui/src/context/abort-state.ts @@ -0,0 +1,50 @@ +type Entry = { + phase: "pending" | "active" + messageID?: string +} + +export function createAbortState() { + const entries = new Map() + + return { + request(id: string, status: string, messageID?: string) { + const entry = entries.get(id) + if (entry?.phase === "active") return status !== "idle" + if (entry) return false + if (status === "idle") { + if (!messageID) return false + entries.set(id, { phase: "pending", messageID }) + return false + } + entries.set(id, { phase: "active" }) + return true + }, + move(from: string, to: string) { + const source = entries.get(from) + if (!source) return + entries.delete(from) + const target = entries.get(to) + if (target?.phase === "active") return + entries.set(to, source) + }, + update(id: string, status: string) { + const entry = entries.get(id) + if (!entry) return false + if (status === "idle") { + entries.delete(id) + return false + } + if (entry.phase === "active") return false + entries.set(id, { phase: "active" }) + return true + }, + finish(messageID: string) { + for (const [id, entry] of entries) { + if (entry.phase === "pending" && entry.messageID === messageID) entries.delete(id) + } + }, + clear(id: string) { + entries.delete(id) + }, + } +} diff --git a/packages/kilo-vscode/webview-ui/src/context/session.tsx b/packages/kilo-vscode/webview-ui/src/context/session.tsx index 8d296b2fac..f4a3a6c5c2 100644 --- a/packages/kilo-vscode/webview-ui/src/context/session.tsx +++ b/packages/kilo-vscode/webview-ui/src/context/session.tsx @@ -71,6 +71,7 @@ import { state as todoState } from "./todo-revert" import { getVariant, sessionVariantKeys, transferVariants, variantKey } from "./session-variant-store" import { KILO_AUTO, parseModelString } from "../../../src/shared/provider-model" import { visibleMessages as filterVisibleMessages } from "./session-queue" +import { createAbortState } from "./abort-state" const RECENT_LIMIT = 5 const MESSAGE_PAGE_LIMIT = 80 @@ -303,6 +304,7 @@ export const SessionProvider: ParentComponent = (props) => { const [busySinceMap, setBusySinceMap] = createStore>({}) const [submissionMap, setSubmissionMap] = createStore>({}) const pendingSubmissions = new Map() + const aborts = createAbortState() const idle: SessionStatusInfo = { type: "idle" } @@ -442,6 +444,7 @@ export const SessionProvider: ParentComponent = (props) => { if (!busySinceMap[sid]) setBusySinceMap(sid, Date.now()) } const finishSubmission = (messageID: string) => { + aborts.finish(messageID) const sid = pendingSubmissions.get(messageID) if (!sid) return pendingSubmissions.delete(messageID) @@ -464,7 +467,9 @@ export const SessionProvider: ParentComponent = (props) => { } const confirmSubmissions = (sid: string) => { for (const [id, scope] of pendingSubmissions) { - if (scope === sid) pendingSubmissions.delete(id) + if (scope !== sid) continue + aborts.finish(id) + pendingSubmissions.delete(id) } setSubmissionMap( produce((map) => { @@ -1070,6 +1075,7 @@ export const SessionProvider: ParentComponent = (props) => { // Event handlers function handleSessionCreated(session: SessionInfo, draftID?: string) { + if (draftID) aborts.move(draftID, session.id) batch(() => { setStore("sessions", session.id, session) @@ -1496,6 +1502,7 @@ export const SessionProvider: ParentComponent = (props) => { message?: string, next?: number, ) { + const shouldAbort = aborts.update(sessionID, newStatus) confirmSubmissions(sessionID) const prev = statusMap[sessionID] ?? { type: "idle" } const info: SessionStatusInfo = @@ -1522,6 +1529,7 @@ export const SessionProvider: ParentComponent = (props) => { // messages themselves will be reconciled on the next messagesLoaded. pendingOptimistic.delete(sessionID) } + if (shouldAbort) vscode.postMessage({ type: "abort", sessionID }) } function handlePermissionRequest(permission: PermissionRequest) { @@ -1640,6 +1648,7 @@ export const SessionProvider: ParentComponent = (props) => { function handleSendMessageFailed(message: SendMessageFailedMessage) { const sid = message.sessionID ?? message.draftID if (message.messageID) finishSubmission(message.messageID) + if (!message.messageID && sid) aborts.clear(sid) if (sid && message.messageID) { pendingOptimistic.get(sid)?.delete(message.messageID) stash.remove(message.messageID) @@ -1777,6 +1786,7 @@ export const SessionProvider: ParentComponent = (props) => { function handleSessionDeleted(sessionID: string) { pendingOptimistic.delete(sessionID) + aborts.clear(sessionID) confirmSubmissions(sessionID) batch(() => { // Collect message IDs so we can clean up their parts (store + stash) @@ -2197,10 +2207,13 @@ export const SessionProvider: ParentComponent = (props) => { function abort() { const sessionID = currentSessionID() - if (!sessionID) { - console.warn("[Kilo New] Cannot abort: no current session") + const scope = sessionID ?? draftSessionID() + if (!scope) { + console.warn("[Kilo New] Cannot abort: no current or pending session") return } + const messageID = [...pendingSubmissions].reverse().find(([, sid]) => sid === scope)?.[0] + if (!aborts.request(scope, status(), messageID) || !sessionID) return vscode.postMessage({ type: "abort", diff --git a/packages/opencode/src/effect/runner.ts b/packages/opencode/src/effect/runner.ts index f21a61c97e..e0e525b362 100644 --- a/packages/opencode/src/effect/runner.ts +++ b/packages/opencode/src/effect/runner.ts @@ -83,11 +83,14 @@ export const make = ( const startRun = (work: Effect.Effect, done: Deferred.Deferred) => Effect.gen(function* () { const id = next() - const fiber = yield* work.pipe( + // kilocode_change start - do not let work publish busy before the Running state is committed + const ready = yield* Latch.make() + const fiber = yield* ready.whenOpen(work).pipe( Effect.onExit((exit) => finishRun(id, done, exit)), Effect.forkIn(scope), ) - return { id, done, fiber } satisfies RunHandle + return { run: { id, done, fiber } satisfies RunHandle, ready } + // kilocode_change end }) const finishShell = (id: number) => @@ -98,8 +101,13 @@ export const make = ( return [idle, { _tag: "Idle" }] as const } if (st._tag === "ShellThenRun" && st.shell.id === id) { - const run = yield* startRun(st.run.work, st.run.done) - return [Effect.void, { _tag: "Running", run }] as const + // kilocode_change start - open work only after the Running state is committed + const started = yield* startRun(st.run.work, st.run.done) + return [ + started.ready.open.pipe(Effect.uninterruptible, Effect.asVoid), + { _tag: "Running", run: started.run }, + ] as const + // kilocode_change end } return [Effect.void, st] as const }), @@ -130,8 +138,13 @@ export const make = ( } case "Idle": { const done = yield* Deferred.make() - const run = yield* startRun(work, done) - return [awaitDone(done), { _tag: "Running", run }] as const + // kilocode_change start - open work only after the Running state is committed + const started = yield* startRun(work, done) + return [ + started.ready.open.pipe(Effect.uninterruptible, Effect.andThen(awaitDone(done))), + { _tag: "Running", run: started.run }, + ] as const + // kilocode_change end } } }), diff --git a/packages/opencode/test/kilocode/runner-start-order.test.ts b/packages/opencode/test/kilocode/runner-start-order.test.ts new file mode 100644 index 0000000000..1faedf2162 --- /dev/null +++ b/packages/opencode/test/kilocode/runner-start-order.test.ts @@ -0,0 +1,95 @@ +import { describe, expect } from "bun:test" +import { Deferred, Effect, Fiber, Scope } from "effect" +import { Runner } from "@/effect/runner" +import { awaitWithTimeout, it, pollWithTimeout } from "../lib/effect" + +describe("Runner start ordering", () => { + it.live( + "commits Running before work begins", + Effect.gen(function* () { + const scope = yield* Scope.Scope + const runner = Runner.make(scope) + const started = yield* Deferred.make["_tag"]>() + const release = yield* Deferred.make() + const fiber = yield* runner + .ensureRunning( + Effect.gen(function* () { + yield* Deferred.succeed(started, runner.state._tag) + yield* Deferred.await(release) + return "done" + }), + ) + .pipe(Effect.forkChild) + + expect(yield* Deferred.await(started)).toBe("Running") + yield* Deferred.succeed(release, undefined) + expect(yield* Fiber.join(fiber)).toBe("done") + }), + ) + + it.live( + "commits Running before queued work begins after a shell", + Effect.gen(function* () { + const scope = yield* Scope.Scope + const runner = Runner.make(scope) + const shell = yield* Deferred.make() + const started = yield* Deferred.make["_tag"]>() + const release = yield* Deferred.make() + const shellFiber = yield* runner.startShell(Deferred.await(shell).pipe(Effect.as("shell"))).pipe(Effect.forkChild) + yield* pollWithTimeout( + Effect.sync(() => (runner.state._tag === "Shell" ? true : undefined)), + "runner did not enter Shell", + ) + const runFiber = yield* runner + .ensureRunning( + Effect.gen(function* () { + yield* Deferred.succeed(started, runner.state._tag) + yield* Deferred.await(release) + return "done" + }), + ) + .pipe(Effect.forkChild) + yield* pollWithTimeout( + Effect.sync(() => (runner.state._tag === "ShellThenRun" ? true : undefined)), + "runner did not queue work", + ) + + yield* Deferred.succeed(shell, undefined) + yield* Fiber.join(shellFiber) + expect(yield* awaitWithTimeout(Deferred.await(started), "queued work did not start")).toBe("Running") + yield* Deferred.succeed(release, undefined) + expect(yield* Fiber.join(runFiber)).toBe("done") + }), + ) + + it.live( + "opens committed work when its first caller is interrupted", + Effect.gen(function* () { + const scope = yield* Scope.Scope + const runner = Runner.make(scope) + const started = yield* Deferred.make() + const release = yield* Deferred.make() + const fiber = yield* runner + .ensureRunning( + Effect.gen(function* () { + yield* Deferred.succeed(started, undefined) + yield* Deferred.await(release) + return "done" + }), + ) + .pipe(Effect.forkChild) + + yield* pollWithTimeout( + Effect.sync(() => (runner.state._tag === "Running" ? true : undefined)), + "runner did not commit Running", + ) + yield* Fiber.interrupt(fiber) + yield* awaitWithTimeout(Deferred.await(started), "committed work did not start") + yield* Deferred.succeed(release, undefined) + yield* pollWithTimeout( + Effect.sync(() => (runner.state._tag === "Idle" ? true : undefined)), + "runner did not return to Idle", + ) + }), + ) +})