opensource/kilocode
3
0
Fork 0
mirror of https://github.com/Kilo-Org/kilocode.git synced 2026-08-28 19:11:03 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix: abort prompts during startup

Browse Source
This commit is contained in:
marius-kilocode
2026-06-12 09:57:49 +02:00
parent a09dafe61b
commit b23d3dfd75
15 changed files with 284 additions and 298 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.changeset/abort-moved-sessions.md
-5
View File
@@ -1,5 +0,0 @@
---
"kilo-code": patch
---
Stop active Agent Manager turns after moving or continuing their sessions in another worktree.
.changeset/abort-starting-sessions.md
+6
View File
@@ -0,0 +1,6 @@
---
"kilo-code": patch
"@kilocode/cli": patch
---
Allow Escape to stop Agent Manager prompts while their sessions are still starting.
packages/kilo-vscode/src/KiloProvider.ts
+12 -65
View File
@@ -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<string> = new Set()
private sessionStatusMap = new Map<string, SessionStatus["type"]>() // Latest status used for destructive config warnings.
private sessionDirectories = new Map<string, string>() // Per-session directory overrides, such as Agent Manager worktrees.
private activeSessionDirectories: ActiveSessionDirectories = new Map()
private permissionDirectories = new Map<string, string>()
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<void> {
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<string, unknown> | 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()
packages/kilo-vscode/src/kilo-provider/abort.ts
+1 -34
View File
@@ -1,37 +1,4 @@
import type { KiloClient, SessionStatus } from "@kilocode/sdk/v2/client"
import { sameDirectory } from "../kilo-provider-utils"
export type ActiveSessionDirectories = Map<string, Set<string>>
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 })
packages/kilo-vscode/tests/unit/abort-state.test.ts
+64
View File
@@ -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)
})
})
packages/kilo-vscode/tests/unit/abort.test.ts
+1 -58
View File
@@ -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[] = []
packages/kilo-vscode/tests/unit/kilo-provider-load-messages.test.ts
+2 -114
View File
@@ -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<unknown>
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<string, string>
activeSessionDirectories: Map<string, Set<string>>
trackedSessionIds: Set<string>
stopCurrentSessionProcesses: (next?: string) => void
handleEvent: (event: unknown, directory?: string) => void
handleAbort: (sid?: string) => Promise<void>
handleEvent: (event: unknown) => void
handleLoadMessages: (sid: string, opts?: { mode?: string; before?: string; limit?: number }) => Promise<void>
handleDeleteSession: (sid: string) => Promise<void>
}
@@ -145,108 +135,6 @@ function makeProvider(client: ReturnType<typeof createClient>) {
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({
packages/kilo-vscode/tests/unit/prompt-input-utils.test.ts
+13 -9
View File
@@ -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)
})
})
packages/kilo-vscode/webview-ui/src/components/chat/ChatView.tsx
+1 -1
View File
@@ -82,7 +82,7 @@ export const ChatView: Component<ChatViewProps> = (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()
}
packages/kilo-vscode/webview-ui/src/components/chat/PromptInput.tsx
+2 -1
View File
@@ -344,7 +344,8 @@ export const PromptInput: Component<PromptInputProps> = (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())
packages/kilo-vscode/webview-ui/src/components/chat/prompt-input-utils.ts
+2 -2
View File
@@ -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)
}
/**
packages/kilo-vscode/webview-ui/src/context/abort-state.ts
+50
View File
@@ -0,0 +1,50 @@
type Entry = {
phase: "pending" | "active"
messageID?: string
}
export function createAbortState() {
const entries = new Map<string, Entry>()
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)
},
}
}
packages/kilo-vscode/webview-ui/src/context/session.tsx
+16 -3
View File
@@ -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<Record<string, number>>({})
const [submissionMap, setSubmissionMap] = createStore<Record<string, number>>({})
const pendingSubmissions = new Map<string, string>()
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",
packages/opencode/src/effect/runner.ts
+19 -6
View File
@@ -83,11 +83,14 @@ export const make = <A, E = never>(
const startRun = (work: Effect.Effect<A, E>, done: Deferred.Deferred<A, E | Cancelled>) =>
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<A, E>
return { run: { id, done, fiber } satisfies RunHandle<A, E>, ready }
// kilocode_change end
})
const finishShell = (id: number) =>
@@ -98,8 +101,13 @@ export const make = <A, E = never>(
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 = <A, E = never>(
}
case "Idle": {
const done = yield* Deferred.make<A, E | Cancelled>()
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
}
}
}),
packages/opencode/test/kilocode/runner-start-order.test.ts
+95
View File
@@ -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<string>(scope)
const started = yield* Deferred.make<Runner.State<string, never>["_tag"]>()
const release = yield* Deferred.make<void>()
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<string>(scope)
const shell = yield* Deferred.make<void>()
const started = yield* Deferred.make<Runner.State<string, never>["_tag"]>()
const release = yield* Deferred.make<void>()
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<string>(scope)
const started = yield* Deferred.make<void>()
const release = yield* Deferred.make<void>()
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",
)
}),
)
})