From d6c2cd9bf6fcd99c73b5cb401aa0762e59f4cc1f Mon Sep 17 00:00:00 2001 From: "kiloconnect[bot]" <240665456+kiloconnect[bot]@users.noreply.github.com> Date: Mon, 13 Apr 2026 08:37:44 +0000 Subject: [PATCH 1/3] refactor(vscode): extract session handlers from KiloProvider Extract session CRUD, message send/command, abort, revert/unrevert, compact, and retry-with-backoff into a dedicated session-handler module. Reduces KiloProvider from 3327 to 2819 lines, fixing the max-lines lint error. --- packages/kilo-vscode/src/KiloProvider.ts | 719 +++--------------- .../kilo-provider/handlers/session-handler.ts | 662 ++++++++++++++++ 2 files changed, 781 insertions(+), 600 deletions(-) create mode 100644 packages/kilo-vscode/src/kilo-provider/handlers/session-handler.ts diff --git a/packages/kilo-vscode/src/KiloProvider.ts b/packages/kilo-vscode/src/KiloProvider.ts index 0db78a29768..bc3546f0bb9 100644 --- a/packages/kilo-vscode/src/KiloProvider.ts +++ b/packages/kilo-vscode/src/KiloProvider.ts @@ -2,15 +2,7 @@ import * as path from "path" import * as vscode from "vscode" import { buildPreviewPath, getPreviewCommand, getPreviewDir, parseImage, trimEntries } from "./image-preview" import { isAbsolutePath } from "./path-utils" -import type { - KiloClient, - Session, - SessionStatus, - Event, - TextPartInput, - FilePartInput, - Config, -} from "@kilocode/sdk/v2/client" +import type { KiloClient, Session, SessionStatus, Event, Config } from "@kilocode/sdk/v2/client" import { type KiloConnectionService, type KilocodeNotification, ServerStartupError } from "./services/cli-backend" import type { EditorContext } from "./services/cli-backend/types" import { FileIgnoreController } from "./services/autocomplete/shims/FileIgnoreController" @@ -27,12 +19,9 @@ import { isEventFromForeignProject, MessageConfirmation, runWithMessageConfirmation, - loadSessions as loadSessionsUtil, - flushPendingSessionRefresh as flushPendingSessionRefreshUtil, resolveContextDirectory, resolveWorkspaceDirectory, mergeFileSearchResults, - type SessionRefreshContext, } from "./kilo-provider-utils" import { GitOps } from "./agent-manager/GitOps" import { GitStatsPoller, type LocalStats } from "./agent-manager/GitStatsPoller" @@ -48,7 +37,7 @@ import { parseMessageFiles } from "./kilo-provider/message-files" import { matchFollowup, recordFollowup, type Followup } from "./kilo-provider/followup-session" import { childID } from "./kilo-provider/task-session" import { handleNetworkEvent, clearNetworkWaits } from "./kilo-provider/network" -import { retryable, backoff, MAX_RETRIES } from "./util/retry" + // legacy-migration start import { checkAndShowMigrationWizard, @@ -78,6 +67,24 @@ import { fetchAndSendPendingPermissions, type PermissionContext, } from "./kilo-provider/handlers/permission-handler" +import { + handleCreateSession as handleCreateSessionHandler, + handleLoadMessages as handleLoadMessagesHandler, + handleSyncSession as handleSyncSessionHandler, + flushPendingSessionRefresh as flushPendingSessionRefreshHandler, + handleLoadSessions as handleLoadSessionsHandler, + handleDeleteSession as handleDeleteSessionHandler, + handleRenameSession as handleRenameSessionHandler, + handleSendMessage as handleSendMessageHandler, + handleSendCommand as handleSendCommandHandler, + handleAbort as handleAbortHandler, + handleRevertSession as handleRevertSessionHandler, + handleUnrevertSession as handleUnrevertSessionHandler, + handleCompact as handleCompactHandler, + withRetry as withRetryHandler, + cancelRetry as cancelRetryHandler, + type SessionContext, +} from "./kilo-provider/handlers/session-handler" import { handleQuestionReply, handleQuestionReject, @@ -163,8 +170,7 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper private sessionDirectories = new Map() /** Project ID for the current workspace, used to filter out sessions from other repositories. */ private projectID: string | undefined - /** Abort controller for the current loadMessages request; aborted when a new session is selected. */ - private loadMessagesAbort: AbortController | null = null + // loadMessagesAbort moved to kilo-provider/handlers/session-handler.ts /** Set when refreshSessions() is called before the client is ready. * Cleared and retried once the connection transitions to "connected". */ private pendingSessionRefresh = false @@ -1229,304 +1235,84 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper return sessionToWebview(session) } - /** - * Handle creating a new session. - */ - private async handleCreateSession(): Promise { - if (!this.client) { - this.postMessage({ - type: "error", - message: "Not connected to CLI backend", - }) - return - } + // Session CRUD, message send, abort, revert, compact — extracted to + // kilo-provider/handlers/session-handler.ts - try { - const workspaceDir = this.getContextDirectory() - const { data: session } = await this.client.session.create({ directory: workspaceDir }, { throwOnError: true }) - this.currentSession = session - this.contextSessionID = session.id - this.trackDirectory(session.id, workspaceDir) - this.trackedSessionIds.add(session.id) - - // Notify webview of the new session - this.postMessage({ - type: "sessionCreated", - session: this.sessionToWebview(this.currentSession!), - }) - } catch (error) { - console.error("[Kilo New] KiloProvider: Failed to create session:", error) - this.postMessage({ - type: "error", - message: getErrorMessage(error) || "Failed to create session", - }) + private get sessionCtx(): SessionContext { + const self = this + return { + get client() { + return self.client + }, + get currentSession() { + return self.currentSession + }, + set currentSession(session) { + self.currentSession = session + }, + get contextSessionID() { + return self.contextSessionID + }, + set contextSessionID(id) { + self.contextSessionID = id + }, + get connectionState() { + return self.connectionState + }, + trackedSessionIds: this.trackedSessionIds, + syncedChildSessions: this.syncedChildSessions, + sessionDirectories: this.sessionDirectories, + confirmations: this.confirmations, + slimEditMetadata: this.slimEditMetadata, + get pendingSessionRefresh() { + return self.pendingSessionRefresh + }, + set pendingSessionRefresh(v) { + self.pendingSessionRefresh = v + }, + get projectID() { + return self.projectID + }, + set projectID(v) { + self.projectID = v + }, + postMessage: (msg) => this.postMessage(msg), + getWorkspaceDirectory: (sid) => this.getWorkspaceDirectory(sid), + getContextDirectory: () => this.getContextDirectory(), + gatherEditorContext: () => this.gatherEditorContext(), + recordMessageSessionId: (mid, sid) => this.connectionService.recordMessageSessionId(mid, sid), + trackDirectory: (sid, dir) => this.trackDirectory(sid, dir), + recoverPendingPrompts: () => this.recoverPendingPrompts(), + focusSession: (id) => this.focusSession(id), } } - /** - * Handle loading messages for a session. - */ - private async handleLoadMessages(sessionID: string): Promise { - // Track the session so we receive its SSE events - this.trackedSessionIds.add(sessionID) - this.focusSession(sessionID) - this.contextSessionID = sessionID - - if (!this.client) { - this.postMessage({ - type: "error", - message: "Not connected to CLI backend", - sessionID, - }) - return - } - - // Abort any previous in-flight loadMessages request so the backend - // isn't overwhelmed when the user switches sessions rapidly. - this.loadMessagesAbort?.abort() - const abort = new AbortController() - this.loadMessagesAbort = abort - - try { - const workspaceDir = this.getWorkspaceDirectory(sessionID) - const { data: messagesData } = await retry(() => - this.client!.session.messages( - { sessionID, directory: workspaceDir }, - { throwOnError: true, signal: abort.signal }, - ), - ) - - // If this request was aborted while awaiting, skip posting stale results - if (abort.signal.aborted) return - - // Update currentSession so fallback logic in handleSendMessage/handleAbort - // references the correct session after switching. loadMessages is the - // canonical "user switched to this session" signal, so always update — - // the old guard `this.currentSession.id === sessionID` prevented updates - // when switching between different sessions. - // Non-blocking: don't let a failure here prevent messages from loading. - // 404s are expected for cross-worktree sessions — use silent to suppress HTTP error logs. - this.client.session - .get({ sessionID, directory: workspaceDir }) - .then((result) => { - if (result.data && !abort.signal.aborted) { - this.currentSession = result.data - this.contextSessionID = result.data.id - } - }) - .catch((err: unknown) => console.warn("[Kilo New] KiloProvider: getSession failed (non-critical):", err)) - - this.postMessage({ - type: "workspaceDirectoryChanged", - directory: this.getWorkspaceDirectory(sessionID), - }) - - // Fetch current session status so the webview has the correct busy/idle - // state after switching tabs (SSE events may have been missed). - this.client.session - .status({ directory: workspaceDir }) - .then((result) => { - if (!result.data) return - for (const [sid, info] of Object.entries(result.data) as [string, SessionStatus][]) { - if (!this.trackedSessionIds.has(sid)) continue - this.postMessage({ - type: "sessionStatus", - sessionID: sid, - status: info.type, - ...(info.type === "retry" ? { attempt: info.attempt, message: info.message, next: info.next } : {}), - }) - } - }) - .catch((err: unknown) => console.error("[Kilo New] KiloProvider: Failed to fetch session statuses:", err)) - - const messages = messagesData.map((m) => ({ - ...m.info, - parts: this.slimParts(m.parts), - createdAt: new Date(m.info.time.created).toISOString(), - })) - - for (const message of messages) { - this.connectionService.recordMessageSessionId(message.id, message.sessionID) - } - - this.postMessage({ - type: "messagesLoaded", - sessionID, - messages, - }) - - // Recover any prompts missed while the webview was loading or during an SSE reconnection. - this.recoverPendingPrompts() - } catch (error) { - // Silently ignore aborted requests — the user switched to a different session - if (abort.signal.aborted) return - console.error("[Kilo New] KiloProvider: Failed to load messages:", error) - this.postMessage({ - type: "error", - message: getErrorMessage(error) || "Failed to load messages", - sessionID, - }) - } + private handleCreateSession(): Promise { + return handleCreateSessionHandler(this.sessionCtx) } - /** - * Handle syncing a child session (e.g. spawned by the task tool). - * Tracks the session for SSE events and fetches its messages. - */ - private async handleSyncSession(sessionID: string, parentSessionID?: string): Promise { - if (!this.client) return - if (this.syncedChildSessions.has(sessionID)) return - - this.syncedChildSessions.add(sessionID) - this.trackedSessionIds.add(sessionID) - - // Inherit the parent's worktree directory so permission responses use - // the correct backend Instance. Without this, child sessions in Agent - // Manager worktrees fall back to workspace root and fail to find the - // pending permission request. - if (!this.sessionDirectories.has(sessionID) && parentSessionID) { - const dir = this.sessionDirectories.get(parentSessionID) - if (dir) { - this.sessionDirectories.set(sessionID, dir) - } - } - - try { - const workspaceDir = this.getWorkspaceDirectory(sessionID) - const { data: messagesData } = await retry(() => - this.client!.session.messages({ sessionID, directory: workspaceDir }, { throwOnError: true }), - ) - - const messages = messagesData.map((m) => ({ - ...m.info, - parts: this.slimParts(m.parts), - createdAt: new Date(m.info.time.created).toISOString(), - })) - - for (const message of messages) { - this.connectionService.recordMessageSessionId(message.id, message.sessionID) - } - - this.postMessage({ - type: "messagesLoaded", - sessionID, - messages, - }) - - // Recover any prompts emitted by the child before we started tracking it. - this.recoverPendingPrompts() - } catch (err) { - this.syncedChildSessions.delete(sessionID) - console.error("[Kilo New] KiloProvider: Failed to sync child session:", err) - } + private handleLoadMessages(sessionID: string): Promise { + return handleLoadMessagesHandler(this.sessionCtx, sessionID) } - /** - * Build the context object used by the extracted session-refresh helpers. - */ - private get sessionRefreshContext(): SessionRefreshContext { - const client = this.client - return { - pendingSessionRefresh: this.pendingSessionRefresh, - connectionState: this.connectionState, - listSessions: client - ? (dir: string) => - client.session.list({ directory: dir, roots: true }, { throwOnError: true }).then(({ data }) => data) - : null, - sessionDirectories: this.sessionDirectories, - workspaceDirectory: this.getWorkspaceDirectory(), - postMessage: (msg: unknown) => this.postMessage(msg), - } + private handleSyncSession(sessionID: string, parentSessionID?: string): Promise { + return handleSyncSessionHandler(this.sessionCtx, sessionID, parentSessionID) } - /** - * Retry a deferred sessions refresh once the client is ready. - */ - private async flushPendingSessionRefresh(reason: string): Promise { - if (!this.pendingSessionRefresh) return - console.log("[Kilo New] KiloProvider: 🔄 Flushing deferred sessions refresh", { reason }) - const ctx = this.sessionRefreshContext - try { - const resolved = await flushPendingSessionRefreshUtil(ctx) - if (resolved) this.projectID = resolved - } catch (error) { - console.error("[Kilo New] KiloProvider: Failed to flush session refresh:", error) - } - this.pendingSessionRefresh = ctx.pendingSessionRefresh + private flushPendingSessionRefresh(reason: string): Promise { + return flushPendingSessionRefreshHandler(this.sessionCtx, reason) } - /** - * Handle loading all sessions. - */ - private async handleLoadSessions(): Promise { - const ctx = this.sessionRefreshContext - try { - const resolved = await loadSessionsUtil(ctx) - if (resolved) this.projectID = resolved - } catch (error) { - console.error("[Kilo New] KiloProvider: Failed to load sessions:", error) - this.postMessage({ - type: "error", - message: getErrorMessage(error) || "Failed to load sessions", - }) - } - this.pendingSessionRefresh = ctx.pendingSessionRefresh + private handleLoadSessions(): Promise { + return handleLoadSessionsHandler(this.sessionCtx) } - /** - * Handle deleting a session. - */ - private async handleDeleteSession(sessionID: string): Promise { - if (!this.client) { - this.postMessage({ type: "error", message: "Not connected to CLI backend" }) - return - } - - try { - const workspaceDir = this.getWorkspaceDirectory(sessionID) - await this.client.session.delete({ sessionID, directory: workspaceDir }, { throwOnError: true }) - this.trackedSessionIds.delete(sessionID) - this.syncedChildSessions.delete(sessionID) - this.sessionDirectories.delete(sessionID) - if (this.currentSession?.id === sessionID) { - this.currentSession = null - } - this.postMessage({ type: "sessionDeleted", sessionID }) - } catch (error) { - console.error("[Kilo New] KiloProvider: Failed to delete session:", error) - this.postMessage({ - type: "error", - message: getErrorMessage(error) || "Failed to delete session", - }) - } + private handleDeleteSession(sessionID: string): Promise { + return handleDeleteSessionHandler(this.sessionCtx, sessionID) } - /** - * Handle renaming a session. - */ - private async handleRenameSession(sessionID: string, title: string): Promise { - if (!this.client) { - this.postMessage({ type: "error", message: "Not connected to CLI backend" }) - return - } - - try { - const workspaceDir = this.getWorkspaceDirectory(sessionID) - const { data: updated } = await this.client.session.update( - { sessionID, directory: workspaceDir, title }, - { throwOnError: true }, - ) - if (this.currentSession?.id === sessionID) { - this.currentSession = updated - } - this.postMessage({ type: "sessionUpdated", session: this.sessionToWebview(updated) }) - } catch (error) { - console.error("[Kilo New] KiloProvider: Failed to rename session:", error) - this.postMessage({ - type: "error", - message: getErrorMessage(error) || "Failed to rename session", - }) - } + private handleRenameSession(sessionID: string, title: string): Promise { + return handleRenameSessionHandler(this.sessionCtx, sessionID, title) } /** Fetch providers and send to webview. Coalesced: at most one in-flight + one queued. */ @@ -2262,115 +2048,11 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper } } - /** - * Ensure a session exists, creating one if needed. Returns the resolved - * session ID and workspace directory, or undefined when the client is - * disconnected. - */ - private async resolveSession( - sessionID?: string, - draftID?: string, - ): Promise<{ sid: string; dir: string } | undefined> { - if (!this.client) return undefined - - const dir = sessionID ? this.getWorkspaceDirectory(sessionID) : this.getContextDirectory() - - if (!sessionID && !this.currentSession) { - const { data: session } = await this.client.session.create({ directory: dir }, { throwOnError: true }) - this.currentSession = session - this.contextSessionID = session.id - this.trackDirectory(session.id, dir) - this.trackedSessionIds.add(session.id) - if (draftID) this.contextSessionID = session.id - this.postMessage({ - type: "sessionCreated", - session: this.sessionToWebview(session), - draftID, - }) - } - - const sid = sessionID || this.currentSession?.id - if (!sid) throw new Error("No session available") - this.trackedSessionIds.add(sid) - return { sid, dir } - } - - /** Abort controllers for active retry loops, keyed by session ID */ - private retryAbortControllers = new Map() - - /** Execute an SDK call with visible exponential backoff for retryable HTTP errors. */ - private async withRetry( - fn: () => Promise<{ error?: unknown; response?: Response }>, - sid: string, - messageID?: string, - ): Promise { - const abortController = new AbortController() - this.retryAbortControllers.set(sid, abortController) - - try { - for (let attempt = 1; ; attempt++) { - if (abortController.signal.aborted) { - // User cancelled — return normally without triggering sendMessageFailed - return - } - - const result = await fn() - if (!result.error) return - if (this.confirmations.has(messageID)) return - - const status = result.response?.status ?? 0 - - // Non-retryable status codes fail immediately without retry - if (!retryable(status)) { - this.postMessage({ type: "sessionStatus", sessionID: sid, status: "idle" }) - throw result.error - } - - // Stop retrying after MAX_RETRIES attempts - if (attempt >= MAX_RETRIES) { - this.postMessage({ type: "sessionStatus", sessionID: sid, status: "idle" }) - throw result.error - } - - const delay = backoff(attempt, result.response?.headers) - console.log(`[Kilo New] KiloProvider: Retry on ${status}, attempt ${attempt}/${MAX_RETRIES}, delay ${delay}ms`) - - this.postMessage({ - type: "sessionStatus", - sessionID: sid, - status: "retry", - attempt, - message: `Error (${status}). Retrying...`, - next: Date.now() + delay, - }) - - // Wait for delay or until aborted - await new Promise((resolve) => { - const done = () => { - clearTimeout(timer) - abortController.signal.removeEventListener("abort", done) - resolve() - } - const timer = setTimeout(done, delay) - abortController.signal.addEventListener("abort", done, { once: true }) - }) - if (this.confirmations.has(messageID)) return - } - } finally { - this.retryAbortControllers.delete(sid) - } - } - - /** Cancel an active retry loop for a session */ private cancelRetry(sid: string): void { - const controller = this.retryAbortControllers.get(sid) - if (controller) { - controller.abort() - this.postMessage({ type: "sessionStatus", sessionID: sid, status: "idle" }) - } + cancelRetryHandler(this.sessionCtx, sid) } - private async handleSendMessage( + private handleSendMessage( text: string, messageID?: string, sessionID?: string, @@ -2381,71 +2063,21 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper variant?: string, files?: Array<{ mime: string; url: string }>, ): Promise { - if (!this.client) { - this.postMessage({ - type: "sendMessageFailed", - error: "Not connected to CLI backend", - text, - sessionID, - draftID, - messageID, - files, - }) - return - } - - let resolved: { sid: string; dir: string } | undefined - try { - resolved = await this.resolveSession(sessionID, draftID) - - const parts: Array = [] - if (files) { - for (const f of files) { - parts.push({ type: "file", mime: f.mime, url: f.url }) - } - } - parts.push({ type: "text", text }) - - const editorContext = await this.gatherEditorContext() - - if (messageID) { - this.connectionService.recordMessageSessionId(messageID, resolved!.sid) - } - - const sid = resolved!.sid - const dir = resolved!.dir - await runWithMessageConfirmation(this.confirmations, messageID, "KiloProvider: Message request", () => - this.withRetry( - () => - this.client!.session.promptAsync({ - sessionID: sid, - directory: dir, - messageID, - parts, - model: providerID && modelID ? { providerID, modelID } : undefined, - agent, - variant, - editorContext, - }), - sid, - messageID, - ), - ) - } catch (error) { - console.error("[Kilo New] KiloProvider: Failed to send message:", error) - this.postMessage({ - type: "sendMessageFailed", - error: getErrorMessage(error) || "Failed to send message", - text, - sessionID: resolved?.sid ?? sessionID, - draftID, - messageID, - files, - }) - } + return handleSendMessageHandler( + this.sessionCtx, + text, + messageID, + sessionID, + draftID, + providerID, + modelID, + agent, + variant, + files, + ) } - private async handleSendCommand( + private handleSendCommand( command: string, args: string, messageID?: string, @@ -2457,148 +2089,35 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper variant?: string, files?: Array<{ mime: string; url: string }>, ): Promise { - if (!this.client) { - this.postMessage({ - type: "sendMessageFailed", - error: "Not connected to CLI backend", - text: `/${command} ${args}`.trim(), - sessionID, - draftID, - messageID, - files, - }) - return - } - - let resolved: { sid: string; dir: string } | undefined - try { - resolved = await this.resolveSession(sessionID, draftID) - - if (messageID) { - this.connectionService.recordMessageSessionId(messageID, resolved!.sid) - } - - const parts = files?.map((f) => ({ type: "file" as const, mime: f.mime, url: f.url })) - - const sid = resolved!.sid - const dir = resolved!.dir - await runWithMessageConfirmation(this.confirmations, messageID, "KiloProvider: Command request", () => - this.withRetry( - () => - this.client!.session.command({ - sessionID: sid, - directory: dir, - command, - arguments: args, - messageID, - model: providerID && modelID ? `${providerID}/${modelID}` : undefined, - agent, - variant, - parts, - }), - sid, - messageID, - ), - ) - } catch (error) { - console.error("[Kilo New] KiloProvider: Failed to send command:", error) - this.postMessage({ - type: "sendMessageFailed", - error: getErrorMessage(error) || "Failed to send command", - text: `/${command} ${args}`.trim(), - sessionID: resolved?.sid ?? sessionID, - draftID, - messageID, - files, - }) - } + return handleSendCommandHandler( + this.sessionCtx, + command, + args, + messageID, + sessionID, + draftID, + providerID, + modelID, + agent, + variant, + files, + ) } - /** - * Handle abort request from the webview. - */ - private async handleAbort(sessionID?: string): Promise { - if (!this.client) { - return - } - - const targetSessionID = sessionID || this.currentSession?.id - if (!targetSessionID) { - return - } - - try { - const workspaceDir = this.getWorkspaceDirectory(targetSessionID) - await this.client.session.abort({ sessionID: targetSessionID, directory: workspaceDir }, { throwOnError: true }) - } catch (error) { - console.error("[Kilo New] KiloProvider: Failed to abort session:", error) - } + private handleAbort(sessionID?: string): Promise { + return handleAbortHandler(this.sessionCtx, sessionID) } - private async handleRevertSession(sessionID: string, messageID: string): Promise { - if (!this.client) return - const dir = this.getWorkspaceDirectory(sessionID) - const { data, error } = await this.client.session.revert({ sessionID, messageID, directory: dir }) - if (error) { - console.error("[Kilo New] KiloProvider: Failed to revert session:", error) - this.postMessage({ type: "error", message: "Failed to revert session", sessionID }) - return - } - if (data) this.postMessage({ type: "sessionUpdated", session: sessionToWebview(data) }) + private handleRevertSession(sessionID: string, messageID: string): Promise { + return handleRevertSessionHandler(this.sessionCtx, sessionID, messageID) } - private async handleUnrevertSession(sessionID: string): Promise { - if (!this.client) return - const dir = this.getWorkspaceDirectory(sessionID) - const { data, error } = await this.client.session.unrevert({ sessionID, directory: dir }) - if (error) { - console.error("[Kilo New] KiloProvider: Failed to unrevert session:", error) - this.postMessage({ type: "error", message: "Failed to redo session", sessionID }) - return - } - if (data) this.postMessage({ type: "sessionUpdated", session: sessionToWebview(data) }) + private handleUnrevertSession(sessionID: string): Promise { + return handleUnrevertSessionHandler(this.sessionCtx, sessionID) } - /** - * Handle compact (context summarization) request from the webview. - */ - private async handleCompact(sessionID?: string, providerID?: string, modelID?: string): Promise { - if (!this.client) { - this.postMessage({ - type: "error", - message: "Not connected to CLI backend", - }) - return - } - - const target = sessionID || this.currentSession?.id - if (!target) { - console.error("[Kilo New] KiloProvider: No sessionID for compact") - return - } - - if (!providerID || !modelID) { - console.error("[Kilo New] KiloProvider: No model selected for compact") - this.postMessage({ - type: "error", - message: "No model selected. Connect a provider to compact this session.", - }) - return - } - - try { - const workspaceDir = this.getWorkspaceDirectory(target) - await this.client.session.summarize( - { sessionID: target, directory: workspaceDir, providerID, modelID }, - { throwOnError: true }, - ) - } catch (error) { - console.error("[Kilo New] KiloProvider: Failed to compact session:", error) - this.postMessage({ - type: "error", - message: getErrorMessage(error) || "Failed to compact session", - }) - } + private handleCompact(sessionID?: string, providerID?: string, modelID?: string): Promise { + return handleCompactHandler(this.sessionCtx, sessionID, providerID, modelID) } // Permission + question handlers extracted to kilo-provider/handlers/permission.ts and question.ts diff --git a/packages/kilo-vscode/src/kilo-provider/handlers/session-handler.ts b/packages/kilo-vscode/src/kilo-provider/handlers/session-handler.ts new file mode 100644 index 00000000000..d71c8b616af --- /dev/null +++ b/packages/kilo-vscode/src/kilo-provider/handlers/session-handler.ts @@ -0,0 +1,662 @@ +/** + * Session lifecycle handlers — extracted from KiloProvider. + * + * Covers CRUD (create, load, delete, rename, sync), message send/command, + * abort, revert/unrevert, compact, and the visible retry-with-backoff loop. + * No vscode dependency. + */ + +import type { KiloClient, Session, SessionStatus, TextPartInput, FilePartInput } from "@kilocode/sdk/v2/client" +import type { EditorContext } from "../../services/cli-backend/types" +import { + getErrorMessage, + sessionToWebview, + runWithMessageConfirmation, + loadSessions as loadSessionsUtil, + flushPendingSessionRefresh as flushPendingSessionRefreshUtil, + type SessionRefreshContext, + type MessageConfirmation, +} from "../../kilo-provider-utils" +import { retry } from "../../services/cli-backend/retry" +import { slimParts } from "../slim-metadata" +import { retryable, backoff, MAX_RETRIES } from "../../util/retry" + +export interface SessionContext { + readonly client: KiloClient | null + currentSession: Session | null + contextSessionID: string | undefined + connectionState: "connecting" | "connected" | "disconnected" | "error" + readonly trackedSessionIds: Set + readonly syncedChildSessions: Set + readonly sessionDirectories: Map + readonly confirmations: MessageConfirmation + readonly slimEditMetadata: boolean + pendingSessionRefresh: boolean + projectID: string | undefined + postMessage(msg: unknown): void + getWorkspaceDirectory(sessionId?: string): string + getContextDirectory(): string + gatherEditorContext(): Promise + recordMessageSessionId(messageId: string, sessionId: string): void + trackDirectory(sessionId: string, dir: string): void + recoverPendingPrompts(): void + focusSession(id?: string): void +} + +// --- Session CRUD --- + +export async function handleCreateSession(ctx: SessionContext): Promise { + if (!ctx.client) { + ctx.postMessage({ type: "error", message: "Not connected to CLI backend" }) + return + } + + try { + const dir = ctx.getContextDirectory() + const { data: session } = await ctx.client.session.create({ directory: dir }, { throwOnError: true }) + ctx.currentSession = session + ctx.contextSessionID = session.id + ctx.trackDirectory(session.id, dir) + ctx.trackedSessionIds.add(session.id) + + ctx.postMessage({ + type: "sessionCreated", + session: sessionToWebview(session), + }) + } catch (error) { + console.error("[Kilo New] KiloProvider: Failed to create session:", error) + ctx.postMessage({ + type: "error", + message: getErrorMessage(error) || "Failed to create session", + }) + } +} + +/** Abort controller for the current loadMessages request; aborted when a new session is selected. */ +let loadMessagesAbort: AbortController | null = null + +export async function handleLoadMessages(ctx: SessionContext, sessionID: string): Promise { + ctx.trackedSessionIds.add(sessionID) + ctx.focusSession(sessionID) + ctx.contextSessionID = sessionID + + if (!ctx.client) { + ctx.postMessage({ type: "error", message: "Not connected to CLI backend", sessionID }) + return + } + + // Abort any previous in-flight loadMessages request so the backend + // isn't overwhelmed when the user switches sessions rapidly. + loadMessagesAbort?.abort() + const abort = new AbortController() + loadMessagesAbort = abort + + try { + const dir = ctx.getWorkspaceDirectory(sessionID) + const { data: messages } = await retry(() => + ctx.client!.session.messages({ sessionID, directory: dir }, { throwOnError: true, signal: abort.signal }), + ) + + // If this request was aborted while awaiting, skip posting stale results + if (abort.signal.aborted) return + + // Update currentSession so fallback logic in handleSendMessage/handleAbort + // references the correct session after switching. loadMessages is the + // canonical "user switched to this session" signal, so always update — + // the old guard `this.currentSession.id === sessionID` prevented updates + // when switching between different sessions. + // Non-blocking: don't let a failure here prevent messages from loading. + // 404s are expected for cross-worktree sessions — use silent to suppress HTTP error logs. + ctx.client.session + .get({ sessionID, directory: dir }) + .then((result) => { + if (result.data && !abort.signal.aborted) { + ctx.currentSession = result.data + ctx.contextSessionID = result.data.id + } + }) + .catch((err: unknown) => console.warn("[Kilo New] KiloProvider: getSession failed (non-critical):", err)) + + ctx.postMessage({ + type: "workspaceDirectoryChanged", + directory: ctx.getWorkspaceDirectory(sessionID), + }) + + // Fetch current session status so the webview has the correct busy/idle + // state after switching tabs (SSE events may have been missed). + ctx.client.session + .status({ directory: dir }) + .then((result) => { + if (!result.data) return + for (const [sid, info] of Object.entries(result.data) as [string, SessionStatus][]) { + if (!ctx.trackedSessionIds.has(sid)) continue + ctx.postMessage({ + type: "sessionStatus", + sessionID: sid, + status: info.type, + ...(info.type === "retry" ? { attempt: info.attempt, message: info.message, next: info.next } : {}), + }) + } + }) + .catch((err: unknown) => console.error("[Kilo New] KiloProvider: Failed to fetch session statuses:", err)) + + const slim = ctx.slimEditMetadata + const mapped = messages.map((m) => ({ + ...m.info, + parts: slim ? slimParts(m.parts) : m.parts, + createdAt: new Date(m.info.time.created).toISOString(), + })) + + for (const message of mapped) { + ctx.recordMessageSessionId(message.id, message.sessionID) + } + + ctx.postMessage({ type: "messagesLoaded", sessionID, messages: mapped }) + + // Recover any prompts missed while the webview was loading or during an SSE reconnection. + ctx.recoverPendingPrompts() + } catch (error) { + // Silently ignore aborted requests — the user switched to a different session + if (abort.signal.aborted) return + console.error("[Kilo New] KiloProvider: Failed to load messages:", error) + ctx.postMessage({ + type: "error", + message: getErrorMessage(error) || "Failed to load messages", + sessionID, + }) + } +} + +/** + * Handle syncing a child session (e.g. spawned by the task tool). + * Tracks the session for SSE events and fetches its messages. + */ +export async function handleSyncSession( + ctx: SessionContext, + sessionID: string, + parentSessionID?: string, +): Promise { + if (!ctx.client) return + if (ctx.syncedChildSessions.has(sessionID)) return + + ctx.syncedChildSessions.add(sessionID) + ctx.trackedSessionIds.add(sessionID) + + // Inherit the parent's worktree directory so permission responses use + // the correct backend Instance. Without this, child sessions in Agent + // Manager worktrees fall back to workspace root and fail to find the + // pending permission request. + if (!ctx.sessionDirectories.has(sessionID) && parentSessionID) { + const dir = ctx.sessionDirectories.get(parentSessionID) + if (dir) { + ctx.sessionDirectories.set(sessionID, dir) + } + } + + try { + const dir = ctx.getWorkspaceDirectory(sessionID) + const { data: messages } = await retry(() => + ctx.client!.session.messages({ sessionID, directory: dir }, { throwOnError: true }), + ) + + const slim = ctx.slimEditMetadata + const mapped = messages.map((m) => ({ + ...m.info, + parts: slim ? slimParts(m.parts) : m.parts, + createdAt: new Date(m.info.time.created).toISOString(), + })) + + for (const message of mapped) { + ctx.recordMessageSessionId(message.id, message.sessionID) + } + + ctx.postMessage({ type: "messagesLoaded", sessionID, messages: mapped }) + + // Recover any prompts emitted by the child before we started tracking it. + ctx.recoverPendingPrompts() + } catch (err) { + ctx.syncedChildSessions.delete(sessionID) + console.error("[Kilo New] KiloProvider: Failed to sync child session:", err) + } +} + +/** + * Build the context object used by the extracted session-refresh helpers. + */ +export function buildSessionRefreshContext(ctx: SessionContext): SessionRefreshContext { + const client = ctx.client + return { + pendingSessionRefresh: ctx.pendingSessionRefresh, + connectionState: ctx.connectionState, + listSessions: client + ? (dir: string) => + client.session.list({ directory: dir, roots: true }, { throwOnError: true }).then(({ data }) => data) + : null, + sessionDirectories: ctx.sessionDirectories, + workspaceDirectory: ctx.getWorkspaceDirectory(), + postMessage: (msg: unknown) => ctx.postMessage(msg), + } +} + +/** + * Retry a deferred sessions refresh once the client is ready. + */ +export async function flushPendingSessionRefresh(ctx: SessionContext, reason: string): Promise { + if (!ctx.pendingSessionRefresh) return + console.log("[Kilo New] KiloProvider: Flushing deferred sessions refresh", { reason }) + const refresh = buildSessionRefreshContext(ctx) + try { + const resolved = await flushPendingSessionRefreshUtil(refresh) + if (resolved) ctx.projectID = resolved + } catch (error) { + console.error("[Kilo New] KiloProvider: Failed to flush session refresh:", error) + } + ctx.pendingSessionRefresh = refresh.pendingSessionRefresh +} + +/** + * Handle loading all sessions. + */ +export async function handleLoadSessions(ctx: SessionContext): Promise { + const refresh = buildSessionRefreshContext(ctx) + try { + const resolved = await loadSessionsUtil(refresh) + if (resolved) ctx.projectID = resolved + } catch (error) { + console.error("[Kilo New] KiloProvider: Failed to load sessions:", error) + ctx.postMessage({ + type: "error", + message: getErrorMessage(error) || "Failed to load sessions", + }) + } + ctx.pendingSessionRefresh = refresh.pendingSessionRefresh +} + +export async function handleDeleteSession(ctx: SessionContext, sessionID: string): Promise { + if (!ctx.client) { + ctx.postMessage({ type: "error", message: "Not connected to CLI backend" }) + return + } + + try { + const dir = ctx.getWorkspaceDirectory(sessionID) + await ctx.client.session.delete({ sessionID, directory: dir }, { throwOnError: true }) + ctx.trackedSessionIds.delete(sessionID) + ctx.syncedChildSessions.delete(sessionID) + ctx.sessionDirectories.delete(sessionID) + if (ctx.currentSession?.id === sessionID) { + ctx.currentSession = null + } + ctx.postMessage({ type: "sessionDeleted", sessionID }) + } catch (error) { + console.error("[Kilo New] KiloProvider: Failed to delete session:", error) + ctx.postMessage({ + type: "error", + message: getErrorMessage(error) || "Failed to delete session", + }) + } +} + +export async function handleRenameSession(ctx: SessionContext, sessionID: string, title: string): Promise { + if (!ctx.client) { + ctx.postMessage({ type: "error", message: "Not connected to CLI backend" }) + return + } + + try { + const dir = ctx.getWorkspaceDirectory(sessionID) + const { data: updated } = await ctx.client.session.update( + { sessionID, directory: dir, title }, + { throwOnError: true }, + ) + if (ctx.currentSession?.id === sessionID) { + ctx.currentSession = updated + } + ctx.postMessage({ type: "sessionUpdated", session: sessionToWebview(updated) }) + } catch (error) { + console.error("[Kilo New] KiloProvider: Failed to rename session:", error) + ctx.postMessage({ + type: "error", + message: getErrorMessage(error) || "Failed to rename session", + }) + } +} + +// --- Message send, abort, revert, compact --- + +/** Abort controllers for active retry loops, keyed by session ID */ +const retryControllers = new Map() + +/** Execute an SDK call with visible exponential backoff for retryable HTTP errors. */ +export async function withRetry( + ctx: SessionContext, + fn: () => Promise<{ error?: unknown; response?: Response }>, + sid: string, + messageID?: string, +): Promise { + const abort = new AbortController() + retryControllers.set(sid, abort) + + try { + for (let attempt = 1; ; attempt++) { + if (abort.signal.aborted) { + // User cancelled — return normally without triggering sendMessageFailed + return + } + + const result = await fn() + if (!result.error) return + if (ctx.confirmations.has(messageID)) return + + const status = result.response?.status ?? 0 + + // Non-retryable status codes fail immediately without retry + if (!retryable(status)) { + ctx.postMessage({ type: "sessionStatus", sessionID: sid, status: "idle" }) + throw result.error + } + + // Stop retrying after MAX_RETRIES attempts + if (attempt >= MAX_RETRIES) { + ctx.postMessage({ type: "sessionStatus", sessionID: sid, status: "idle" }) + throw result.error + } + + const delay = backoff(attempt, result.response?.headers) + console.log(`[Kilo New] KiloProvider: Retry on ${status}, attempt ${attempt}/${MAX_RETRIES}, delay ${delay}ms`) + + ctx.postMessage({ + type: "sessionStatus", + sessionID: sid, + status: "retry", + attempt, + message: `Error (${status}). Retrying...`, + next: Date.now() + delay, + }) + + // Wait for delay or until aborted + await new Promise((resolve) => { + const done = () => { + clearTimeout(timer) + abort.signal.removeEventListener("abort", done) + resolve() + } + const timer = setTimeout(done, delay) + abort.signal.addEventListener("abort", done, { once: true }) + }) + if (ctx.confirmations.has(messageID)) return + } + } finally { + retryControllers.delete(sid) + } +} + +/** Cancel an active retry loop for a session */ +export function cancelRetry(ctx: SessionContext, sid: string): void { + const controller = retryControllers.get(sid) + if (controller) { + controller.abort() + ctx.postMessage({ type: "sessionStatus", sessionID: sid, status: "idle" }) + } +} + +/** + * Ensure a session exists, creating one if needed. Returns the resolved + * session ID and workspace directory, or undefined when the client is + * disconnected. + */ +async function resolveSession( + ctx: SessionContext, + sessionID?: string, + draftID?: string, +): Promise<{ sid: string; dir: string } | undefined> { + if (!ctx.client) return undefined + + const dir = sessionID ? ctx.getWorkspaceDirectory(sessionID) : ctx.getContextDirectory() + + if (!sessionID && !ctx.currentSession) { + const { data: session } = await ctx.client.session.create({ directory: dir }, { throwOnError: true }) + ctx.currentSession = session + ctx.contextSessionID = session.id + ctx.trackDirectory(session.id, dir) + ctx.trackedSessionIds.add(session.id) + if (draftID) ctx.contextSessionID = session.id + ctx.postMessage({ + type: "sessionCreated", + session: sessionToWebview(session), + draftID, + }) + } + + const sid = sessionID || ctx.currentSession?.id + if (!sid) throw new Error("No session available") + ctx.trackedSessionIds.add(sid) + return { sid, dir } +} + +export async function handleSendMessage( + ctx: SessionContext, + text: string, + messageID?: string, + sessionID?: string, + draftID?: string, + providerID?: string, + modelID?: string, + agent?: string, + variant?: string, + files?: Array<{ mime: string; url: string }>, +): Promise { + if (!ctx.client) { + ctx.postMessage({ + type: "sendMessageFailed", + error: "Not connected to CLI backend", + text, + sessionID, + draftID, + messageID, + files, + }) + return + } + + let resolved: { sid: string; dir: string } | undefined + try { + resolved = await resolveSession(ctx, sessionID, draftID) + + const parts: Array = [] + if (files) { + for (const f of files) { + parts.push({ type: "file", mime: f.mime, url: f.url }) + } + } + parts.push({ type: "text", text }) + + const editor = await ctx.gatherEditorContext() + + if (messageID) { + ctx.recordMessageSessionId(messageID, resolved!.sid) + } + + const sid = resolved!.sid + const dir = resolved!.dir + await runWithMessageConfirmation(ctx.confirmations, messageID, "KiloProvider: Message request", () => + withRetry( + ctx, + () => + ctx.client!.session.promptAsync({ + sessionID: sid, + directory: dir, + messageID, + parts, + model: providerID && modelID ? { providerID, modelID } : undefined, + agent, + variant, + editorContext: editor, + }), + sid, + messageID, + ), + ) + } catch (error) { + console.error("[Kilo New] KiloProvider: Failed to send message:", error) + ctx.postMessage({ + type: "sendMessageFailed", + error: getErrorMessage(error) || "Failed to send message", + text, + sessionID: resolved?.sid ?? sessionID, + draftID, + messageID, + files, + }) + } +} + +export async function handleSendCommand( + ctx: SessionContext, + command: string, + args: string, + messageID?: string, + sessionID?: string, + draftID?: string, + providerID?: string, + modelID?: string, + agent?: string, + variant?: string, + files?: Array<{ mime: string; url: string }>, +): Promise { + if (!ctx.client) { + ctx.postMessage({ + type: "sendMessageFailed", + error: "Not connected to CLI backend", + text: `/${command} ${args}`.trim(), + sessionID, + draftID, + messageID, + files, + }) + return + } + + let resolved: { sid: string; dir: string } | undefined + try { + resolved = await resolveSession(ctx, sessionID, draftID) + + if (messageID) { + ctx.recordMessageSessionId(messageID, resolved!.sid) + } + + const parts = files?.map((f) => ({ type: "file" as const, mime: f.mime, url: f.url })) + + const sid = resolved!.sid + const dir = resolved!.dir + await runWithMessageConfirmation(ctx.confirmations, messageID, "KiloProvider: Command request", () => + withRetry( + ctx, + () => + ctx.client!.session.command({ + sessionID: sid, + directory: dir, + command, + arguments: args, + messageID, + model: providerID && modelID ? `${providerID}/${modelID}` : undefined, + agent, + variant, + parts, + }), + sid, + messageID, + ), + ) + } catch (error) { + console.error("[Kilo New] KiloProvider: Failed to send command:", error) + ctx.postMessage({ + type: "sendMessageFailed", + error: getErrorMessage(error) || "Failed to send command", + text: `/${command} ${args}`.trim(), + sessionID: resolved?.sid ?? sessionID, + draftID, + messageID, + files, + }) + } +} + +export async function handleAbort(ctx: SessionContext, sessionID?: string): Promise { + if (!ctx.client) return + + const target = sessionID || ctx.currentSession?.id + if (!target) return + + try { + const dir = ctx.getWorkspaceDirectory(target) + await ctx.client.session.abort({ sessionID: target, directory: dir }, { throwOnError: true }) + } catch (error) { + console.error("[Kilo New] KiloProvider: Failed to abort session:", error) + } +} + +export async function handleRevertSession(ctx: SessionContext, sessionID: string, messageID: string): Promise { + if (!ctx.client) return + const dir = ctx.getWorkspaceDirectory(sessionID) + const { data, error } = await ctx.client.session.revert({ sessionID, messageID, directory: dir }) + if (error) { + console.error("[Kilo New] KiloProvider: Failed to revert session:", error) + ctx.postMessage({ type: "error", message: "Failed to revert session", sessionID }) + return + } + if (data) ctx.postMessage({ type: "sessionUpdated", session: sessionToWebview(data) }) +} + +export async function handleUnrevertSession(ctx: SessionContext, sessionID: string): Promise { + if (!ctx.client) return + const dir = ctx.getWorkspaceDirectory(sessionID) + const { data, error } = await ctx.client.session.unrevert({ sessionID, directory: dir }) + if (error) { + console.error("[Kilo New] KiloProvider: Failed to unrevert session:", error) + ctx.postMessage({ type: "error", message: "Failed to redo session", sessionID }) + return + } + if (data) ctx.postMessage({ type: "sessionUpdated", session: sessionToWebview(data) }) +} + +export async function handleCompact( + ctx: SessionContext, + sessionID?: string, + providerID?: string, + modelID?: string, +): Promise { + if (!ctx.client) { + ctx.postMessage({ type: "error", message: "Not connected to CLI backend" }) + return + } + + const target = sessionID || ctx.currentSession?.id + if (!target) { + console.error("[Kilo New] KiloProvider: No sessionID for compact") + return + } + + if (!providerID || !modelID) { + console.error("[Kilo New] KiloProvider: No model selected for compact") + ctx.postMessage({ + type: "error", + message: "No model selected. Connect a provider to compact this session.", + }) + return + } + + try { + const dir = ctx.getWorkspaceDirectory(target) + await ctx.client.session.summarize( + { sessionID: target, directory: dir, providerID, modelID }, + { throwOnError: true }, + ) + } catch (error) { + console.error("[Kilo New] KiloProvider: Failed to compact session:", error) + ctx.postMessage({ + type: "error", + message: getErrorMessage(error) || "Failed to compact session", + }) + } +} From d864d4c61f0820ae52ac04e46ff5382332790ba4 Mon Sep 17 00:00:00 2001 From: "kiloconnect[bot]" <240665456+kiloconnect[bot]@users.noreply.github.com> Date: Mon, 13 Apr 2026 08:42:00 +0000 Subject: [PATCH 2/3] fix(vscode): update arch test and source-links after session-handler extraction Point the handleLoadSessions arch test at the new session-handler module and regenerate source-links.md. --- packages/kilo-docs/source-links.md | 4 ++-- packages/kilo-vscode/tests/unit/agent-manager-arch.test.ts | 7 ++++--- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/packages/kilo-docs/source-links.md b/packages/kilo-docs/source-links.md index 3a7ac81e6bf..70ade759cfa 100644 --- a/packages/kilo-docs/source-links.md +++ b/packages/kilo-docs/source-links.md @@ -82,7 +82,7 @@ - - + - - @@ -117,7 +117,7 @@ - - - + - - diff --git a/packages/kilo-vscode/tests/unit/agent-manager-arch.test.ts b/packages/kilo-vscode/tests/unit/agent-manager-arch.test.ts index 254423454d8..881ae998e08 100644 --- a/packages/kilo-vscode/tests/unit/agent-manager-arch.test.ts +++ b/packages/kilo-vscode/tests/unit/agent-manager-arch.test.ts @@ -409,9 +409,10 @@ describe("KiloProvider — pending session refresh on reconnect", () => { }) it("handleLoadSessions delegates to loadSessionsUtil", () => { - const start = provider.indexOf("private async handleLoadSessions()") - expect(start, "handleLoadSessions must exist").toBeGreaterThan(-1) - const snippet = provider.slice(start, start + 400) + const handler = fs.readFileSync(path.join(ROOT, "src/kilo-provider/handlers/session-handler.ts"), "utf-8") + const start = handler.indexOf("export async function handleLoadSessions") + expect(start, "handleLoadSessions must exist in session-handler").toBeGreaterThan(-1) + const snippet = handler.slice(start, start + 400) expect(snippet, "must call loadSessionsUtil").toContain("loadSessionsUtil") }) From 0c6f11e79cbf589f3f33d94510bd41438f94bde5 Mon Sep 17 00:00:00 2001 From: "kiloconnect[bot]" <240665456+kiloconnect[bot]@users.noreply.github.com> Date: Mon, 13 Apr 2026 08:47:39 +0000 Subject: [PATCH 3/3] fix(vscode): keep loadMessagesAbort and retryControllers per-instance Move loadMessagesAbort and retryControllers from module-scope variables into the SessionContext interface so they remain per-KiloProvider instance. Multiple providers (sidebar, editor tabs, Agent Manager) no longer interfere with each other's abort controllers. --- packages/kilo-vscode/src/KiloProvider.ts | 12 ++++++++++- .../kilo-provider/handlers/session-handler.ts | 20 +++++++++---------- 2 files changed, 20 insertions(+), 12 deletions(-) diff --git a/packages/kilo-vscode/src/KiloProvider.ts b/packages/kilo-vscode/src/KiloProvider.ts index bc3546f0bb9..ce9547367c1 100644 --- a/packages/kilo-vscode/src/KiloProvider.ts +++ b/packages/kilo-vscode/src/KiloProvider.ts @@ -170,7 +170,10 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper private sessionDirectories = new Map() /** Project ID for the current workspace, used to filter out sessions from other repositories. */ private projectID: string | undefined - // loadMessagesAbort moved to kilo-provider/handlers/session-handler.ts + /** Abort controller for the current loadMessages request; aborted when a new session is selected. */ + private loadMessagesAbort: AbortController | null = null + /** Abort controllers for active retry loops, keyed by session ID. */ + private retryAbortControllers = new Map() /** Set when refreshSessions() is called before the client is ready. * Cleared and retried once the connection transitions to "connected". */ private pendingSessionRefresh = false @@ -1276,6 +1279,13 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper set projectID(v) { self.projectID = v }, + get loadMessagesAbort() { + return self.loadMessagesAbort + }, + set loadMessagesAbort(v) { + self.loadMessagesAbort = v + }, + retryControllers: this.retryAbortControllers, postMessage: (msg) => this.postMessage(msg), getWorkspaceDirectory: (sid) => this.getWorkspaceDirectory(sid), getContextDirectory: () => this.getContextDirectory(), diff --git a/packages/kilo-vscode/src/kilo-provider/handlers/session-handler.ts b/packages/kilo-vscode/src/kilo-provider/handlers/session-handler.ts index d71c8b616af..fa75efba349 100644 --- a/packages/kilo-vscode/src/kilo-provider/handlers/session-handler.ts +++ b/packages/kilo-vscode/src/kilo-provider/handlers/session-handler.ts @@ -33,6 +33,10 @@ export interface SessionContext { readonly slimEditMetadata: boolean pendingSessionRefresh: boolean projectID: string | undefined + /** Per-instance abort controller for the current loadMessages request. */ + loadMessagesAbort: AbortController | null + /** Per-instance abort controllers for active retry loops, keyed by session ID. */ + readonly retryControllers: Map postMessage(msg: unknown): void getWorkspaceDirectory(sessionId?: string): string getContextDirectory(): string @@ -72,9 +76,6 @@ export async function handleCreateSession(ctx: SessionContext): Promise { } } -/** Abort controller for the current loadMessages request; aborted when a new session is selected. */ -let loadMessagesAbort: AbortController | null = null - export async function handleLoadMessages(ctx: SessionContext, sessionID: string): Promise { ctx.trackedSessionIds.add(sessionID) ctx.focusSession(sessionID) @@ -87,9 +88,9 @@ export async function handleLoadMessages(ctx: SessionContext, sessionID: string) // Abort any previous in-flight loadMessages request so the backend // isn't overwhelmed when the user switches sessions rapidly. - loadMessagesAbort?.abort() + ctx.loadMessagesAbort?.abort() const abort = new AbortController() - loadMessagesAbort = abort + ctx.loadMessagesAbort = abort try { const dir = ctx.getWorkspaceDirectory(sessionID) @@ -324,9 +325,6 @@ export async function handleRenameSession(ctx: SessionContext, sessionID: string // --- Message send, abort, revert, compact --- -/** Abort controllers for active retry loops, keyed by session ID */ -const retryControllers = new Map() - /** Execute an SDK call with visible exponential backoff for retryable HTTP errors. */ export async function withRetry( ctx: SessionContext, @@ -335,7 +333,7 @@ export async function withRetry( messageID?: string, ): Promise { const abort = new AbortController() - retryControllers.set(sid, abort) + ctx.retryControllers.set(sid, abort) try { for (let attempt = 1; ; attempt++) { @@ -387,13 +385,13 @@ export async function withRetry( if (ctx.confirmations.has(messageID)) return } } finally { - retryControllers.delete(sid) + ctx.retryControllers.delete(sid) } } /** Cancel an active retry loop for a session */ export function cancelRetry(ctx: SessionContext, sid: string): void { - const controller = retryControllers.get(sid) + const controller = ctx.retryControllers.get(sid) if (controller) { controller.abort() ctx.postMessage({ type: "sessionStatus", sessionID: sid, status: "idle" })