diff --git a/.changeset/brisk-subagent-streaming.md b/.changeset/brisk-subagent-streaming.md new file mode 100644 index 00000000000..5f4524618ff --- /dev/null +++ b/.changeset/brisk-subagent-streaming.md @@ -0,0 +1,5 @@ +--- +"kilo-code": patch +--- + +Keep inline subagent streaming responsive during tool-heavy sessions. diff --git a/packages/kilo-vscode/src/KiloProvider.ts b/packages/kilo-vscode/src/KiloProvider.ts index 4556b4fcbc6..38cb9d9cef9 100644 --- a/packages/kilo-vscode/src/KiloProvider.ts +++ b/packages/kilo-vscode/src/KiloProvider.ts @@ -64,7 +64,7 @@ import { interceptMessage } from "./kilo-provider/git-changes-request" import { matchFollowup, recordFollowup, type Followup } from "./kilo-provider/followup-session" import { clearCommandsCache, loadCommands } from "./kilo-provider/commands" import { fetchMessagePage, MESSAGE_PAGE_LIMIT } from "./kilo-provider/message-page" -import { childID } from "./kilo-provider/task-session" +import { childID, VisibleTaskStreams } from "./kilo-provider/task-session" import { handleNetworkEvent, clearNetworkWaits } from "./kilo-provider/network" import { abortSession } from "./kilo-provider/abort" import { @@ -215,6 +215,7 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper * Cleared and retried once the connection transitions to "connected". */ private pendingSessionRefresh = false private readonly streams = new SessionStreamScheduler((msg) => this.postMessage(msg)) + private readonly visibleTaskStreams = new VisibleTaskStreams((id, visible) => this.streams.setVisible(id, visible)) private readonly confirmations = new MessageConfirmation() private unsubscribeEvent: (() => void) | null = null private unsubscribeState: (() => void) | null = null @@ -302,6 +303,10 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper else this.connectionService.unregisterFocused(this.instanceId) } + public setStreamVisibility(active: boolean): void { + this.visibleTaskStreams.setActive(active) + } + public setProjectDirectory(directory: string | null): void { if (this.projectDirectory === directory) return this.projectDirectory = directory @@ -452,9 +457,11 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper this.setupWebviewMessageHandler(webviewView.webview) this.setSidebarVisible(webviewView.visible) + this.setStreamVisibility(webviewView.visible) this.visibilityDisposable?.dispose() this.visibilityDisposable = webviewView.onDidChangeVisibility(() => { this.setSidebarVisible(webviewView.visible) + this.setStreamVisibility(webviewView.visible) if (this.statsPoller) { this.statsPoller.setEnabled(webviewView.visible) this.statsPoller.setVisible(webviewView.visible) @@ -485,10 +492,12 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper panel.webview.html = this._getHtmlForWebview(panel.webview) this.setupWebviewMessageHandler(panel.webview) + this.setStreamVisibility(panel.active) this.viewStateDisposable?.dispose() - this.viewStateDisposable = panel.onDidChangeViewState(() => - this.focusSession(panel.active ? this.currentSession?.id : undefined), - ) + this.viewStateDisposable = panel.onDidChangeViewState(() => { + this.setStreamVisibility(panel.active) + this.focusSession(panel.active ? this.currentSession?.id : undefined) + }) this.initializeConnection() } @@ -645,10 +654,12 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper ) { return } + this.visibleTaskStreams.handle(message) switch (message.type) { case "webviewReady": console.log("[Kilo New] KiloProvider: ✅ webviewReady received") this.isWebviewReady = true + this.visibleTaskStreams.clear() await this.syncWebviewState("webviewReady") this.flushPendingReviewComments() this.recoverPendingPrompts() @@ -1624,6 +1635,7 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper await this.client.session.delete({ sessionID, directory: workspaceDir }, { throwOnError: true }) this.trackedSessionIds.delete(sessionID) this.streams.drop(sessionID) + this.visibleTaskStreams.delete(sessionID) this.syncedChildSessions.delete(sessionID) this.sessionDirectories.delete(sessionID) this.lastReconciledAt.delete(sessionID) @@ -3563,7 +3575,7 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper this.speechToTextDisposable?.dispose() this.telemetryStateDisposable?.dispose() this.autoApproveBridge?.dispose() - this.streams.dispose() + ;(this.visibleTaskStreams.clear(), this.streams.dispose()) this.isWebviewReady = false this.promptRecoveryQueued = false clearNetworkWaits(this.trackedSessionIds) diff --git a/packages/kilo-vscode/src/agent-manager/vscode-host.ts b/packages/kilo-vscode/src/agent-manager/vscode-host.ts index 54e2b6ce5b3..94f28cba3b1 100644 --- a/packages/kilo-vscode/src/agent-manager/vscode-host.ts +++ b/packages/kilo-vscode/src/agent-manager/vscode-host.ts @@ -99,6 +99,10 @@ export class VscodeHost implements Host { provider.attachToWebview(panel.webview, { onBeforeMessage: opts.onBeforeMessage, }) + provider.setStreamVisibility(panel.active && panel.visible) + const streams = panel.onDidChangeViewState((event) => + provider.setStreamVisibility(event.webviewPanel.active && event.webviewPanel.visible), + ) if (this.autoApprove) provider.setAutoApproveController(this.autoApprove) const sessions: SessionProvider = { @@ -147,6 +151,7 @@ export class VscodeHost implements Host { return panel.onDidDispose(cb) }, dispose() { + streams.dispose() provider.dispose() panel.dispose() }, diff --git a/packages/kilo-vscode/src/kilo-provider/session-stream-scheduler.ts b/packages/kilo-vscode/src/kilo-provider/session-stream-scheduler.ts index 1f3c91e8bae..3209d3c6e9d 100644 --- a/packages/kilo-vscode/src/kilo-provider/session-stream-scheduler.ts +++ b/packages/kilo-vscode/src/kilo-provider/session-stream-scheduler.ts @@ -17,6 +17,7 @@ export type StreamSchedulerStats = { emitted: number batches: number active: number + visible: number background: number } @@ -32,6 +33,8 @@ export type StreamSchedulerOptions = { backgroundStepMs?: number /** Hard cap for the background cadence. Defaults to 400ms. */ backgroundMaxMs?: number + /** Flush cadence for visible inline child sessions. Defaults to 50ms. */ + visibleMs?: number } // Scheduler tuning — rationale: @@ -71,6 +74,7 @@ export type StreamSchedulerOptions = { // ~500ms the UI starts feeling disconnected; below 300ms the many-agent // backoff stops providing meaningful throttling. const DEFAULT_ACTIVE_MS = 16 +const DEFAULT_VISIBLE_MS = 50 const DEFAULT_BG_BASE_MS = 150 const DEFAULT_BG_STEP_MS = 20 const DEFAULT_BG_MAX_MS = 400 @@ -110,10 +114,14 @@ function mergePartUpdate(prev: PartUpdate | undefined, msg: PartUpdate): PartUpd export class SessionStreamScheduler { private active: string | undefined private atimer: ReturnType | null = null + private vtimer: ReturnType | null = null private btimer: ReturnType | null = null private bgFirstQueuedAt = 0 + private visibleFirstQueuedAt = 0 private readonly queues = new Map>() + private readonly visible = new Set() private readonly activeMs: number + private readonly visibleMs: number private readonly bgBase: number private readonly bgStep: number private readonly bgMax: number @@ -122,6 +130,7 @@ export class SessionStreamScheduler { emitted: 0, batches: 0, active: 0, + visible: 0, background: 0, } @@ -130,6 +139,7 @@ export class SessionStreamScheduler { opts?: StreamSchedulerOptions, ) { this.activeMs = opts?.activeMs ?? DEFAULT_ACTIVE_MS + this.visibleMs = opts?.visibleMs ?? DEFAULT_VISIBLE_MS this.bgBase = opts?.backgroundBaseMs ?? DEFAULT_BG_BASE_MS this.bgStep = opts?.backgroundStepMs ?? DEFAULT_BG_STEP_MS this.bgMax = opts?.backgroundMaxMs ?? DEFAULT_BG_MAX_MS @@ -143,10 +153,26 @@ export class SessionStreamScheduler { this.atimer = null } this.active = sessionID - if (prev && this.queues.get(prev)?.size) this.scheduleBackground() + if (prev && this.queues.get(prev)?.size) this.schedule(prev) if (sessionID) this.flush(sessionID) } + setVisible(sessionID: string, visible: boolean): void { + const changed = visible ? !this.visible.has(sessionID) : this.visible.has(sessionID) + if (!changed) return + if (visible) this.visible.add(sessionID) + else this.visible.delete(sessionID) + if (this.queues.get(sessionID)?.size) this.schedule(sessionID) + if (this.vtimer && !this.hasVisible()) { + clearTimeout(this.vtimer) + this.vtimer = null + } + if (this.btimer && !this.hasBackground()) { + clearTimeout(this.btimer) + this.btimer = null + } + } + push(msg: PartUpdate): void { this.counters.received++ const key = partUpdateKey(msg) @@ -183,6 +209,11 @@ export class SessionStreamScheduler { this.emit(this.take(sessionID)) + if (this.vtimer && !this.hasVisible()) { + clearTimeout(this.vtimer) + this.vtimer = null + } + if (this.btimer && !this.hasBackground()) { clearTimeout(this.btimer) this.btimer = null @@ -200,6 +231,10 @@ export class SessionStreamScheduler { */ drop(sessionID: string): void { this.queues.delete(sessionID) + if (this.vtimer && !this.hasVisible()) { + clearTimeout(this.vtimer) + this.vtimer = null + } if (this.btimer && !this.hasBackground()) { clearTimeout(this.btimer) this.btimer = null @@ -209,6 +244,7 @@ export class SessionStreamScheduler { dispose(): void { this.clearTimers() this.queues.clear() + this.visible.clear() } stats(): Readonly { @@ -230,9 +266,23 @@ export class SessionStreamScheduler { this.atimer = setTimeout(() => this.flushActive(), this.activeMs) return } + if (this.visible.has(sessionID)) { + this.scheduleVisible() + return + } this.scheduleBackground() } + private scheduleVisible(): void { + if (!this.hasVisible()) return + const now = Date.now() + if (!this.vtimer) this.visibleFirstQueuedAt = now + const elapsed = now - this.visibleFirstQueuedAt + const remaining = Math.max(0, this.visibleMs - elapsed) + if (this.vtimer) clearTimeout(this.vtimer) + this.vtimer = setTimeout(() => this.flushVisible(), remaining) + } + private scheduleBackground(): void { const count = this.backgroundCount() if (count === 0) return @@ -259,6 +309,12 @@ export class SessionStreamScheduler { this.emit(this.takeBackground()) } + private flushVisible(): void { + this.vtimer = null + this.visibleFirstQueuedAt = 0 + this.emit(this.takeVisible()) + } + private take(sessionID: string): PartUpdate[] { const queue = this.queues.get(sessionID) if (!queue) return [] @@ -276,20 +332,43 @@ export class SessionStreamScheduler { const updates: PartUpdate[] = [] for (const [sid, queue] of this.queues) { if (sid === this.active) continue + if (this.visible.has(sid)) continue updates.push(...queue.values()) this.queues.delete(sid) } return updates } + private takeVisible(): PartUpdate[] { + const updates: PartUpdate[] = [] + for (const [sid, queue] of this.queues) { + if (sid === this.active || !this.visible.has(sid)) continue + updates.push(...queue.values()) + this.queues.delete(sid) + } + return updates + } + + private visibleCount(): number { + let n = 0 + for (const [sid, queue] of this.queues) { + if (sid !== this.active && this.visible.has(sid) && queue.size > 0) n++ + } + return n + } + private backgroundCount(): number { let n = 0 for (const [sid, queue] of this.queues) { - if (sid !== this.active && queue.size > 0) n++ + if (sid !== this.active && !this.visible.has(sid) && queue.size > 0) n++ } return n } + private hasVisible(): boolean { + return this.visibleCount() > 0 + } + private hasBackground(): boolean { return this.backgroundCount() > 0 } @@ -315,12 +394,15 @@ export class SessionStreamScheduler { private countLane(sessionID: string): void { if (sessionID === this.active) this.counters.active++ + else if (this.visible.has(sessionID)) this.counters.visible++ else this.counters.background++ } private clearTimers(): void { if (this.atimer) clearTimeout(this.atimer) this.atimer = null + if (this.vtimer) clearTimeout(this.vtimer) + this.vtimer = null if (this.btimer) clearTimeout(this.btimer) this.btimer = null } diff --git a/packages/kilo-vscode/src/kilo-provider/task-session.ts b/packages/kilo-vscode/src/kilo-provider/task-session.ts index cff8f6661d7..d716d0a55cc 100644 --- a/packages/kilo-vscode/src/kilo-provider/task-session.ts +++ b/packages/kilo-vscode/src/kilo-provider/task-session.ts @@ -15,3 +15,47 @@ export function childID(part: Part): string | undefined { if (part.type !== "tool" || part.tool !== "task") return undefined return part.metadata?.sessionId ?? part.state?.metadata?.sessionId } + +type VisibleStreamMessage = { type?: unknown; sessionID?: unknown; visible?: unknown } + +function parse(message: unknown): VisibleStreamMessage | undefined { + if (!message || typeof message !== "object") return undefined + const msg = message as VisibleStreamMessage + if (msg.type !== "streamSessionVisible") return undefined + return msg +} + +export class VisibleTaskStreams { + private readonly refs = new Map() + private active = true + + constructor(private readonly set: (id: string, visible: boolean) => void) {} + + clear(): void { + for (const id of this.refs.keys()) this.set(id, false) + this.refs.clear() + } + + delete(id: string): void { + this.set(id, false) + this.refs.delete(id) + } + + setActive(active: boolean): void { + if (this.active === active) return + this.active = active + for (const id of this.refs.keys()) this.set(id, active) + } + + handle(message: unknown): boolean { + const msg = parse(message) + if (!msg) return false + if (typeof msg.sessionID !== "string" || typeof msg.visible !== "boolean") return true + const count = this.refs.get(msg.sessionID) ?? 0 + const next = msg.visible ? count + 1 : Math.max(0, count - 1) + if (next === 0) this.refs.delete(msg.sessionID) + else this.refs.set(msg.sessionID, next) + this.set(msg.sessionID, this.active && next > 0) + return true + } +} diff --git a/packages/kilo-vscode/tests/unit/part-stash.test.ts b/packages/kilo-vscode/tests/unit/part-stash.test.ts index 58d25f6aff8..8e51214c19c 100644 --- a/packages/kilo-vscode/tests/unit/part-stash.test.ts +++ b/packages/kilo-vscode/tests/unit/part-stash.test.ts @@ -26,6 +26,15 @@ describe("PartStash", () => { expect(stash.size()).toBe(0) }) + it("removePart() keeps off-screen removals authoritative", () => { + const stash = new PartStash() + stash.put("m1", [text("p1", "m1", "drop"), text("p2", "m1", "keep")]) + stash.removePart("m1", "p1") + expect(stash.peek("m1")?.map((part) => part.id)).toEqual(["p2"]) + stash.removePart("m1", "p2") + expect(stash.peek("m1")).toEqual([]) + }) + it("take() consumes stashed parts atomically", () => { const stash = new PartStash() stash.put("m1", [text("p1", "m1", "a")]) diff --git a/packages/kilo-vscode/tests/unit/session-stream-scheduler.test.ts b/packages/kilo-vscode/tests/unit/session-stream-scheduler.test.ts index 209d8c77079..6248b89067b 100644 --- a/packages/kilo-vscode/tests/unit/session-stream-scheduler.test.ts +++ b/packages/kilo-vscode/tests/unit/session-stream-scheduler.test.ts @@ -329,6 +329,66 @@ describe("SessionStreamScheduler / adaptive scheduling", () => { queue.dispose() }) + it("visible sessions flush faster than background sessions", async () => { + const sent: Sent[] = [] + const queue = new SessionStreamScheduler((msg) => sent.push(msg), { + activeMs: 5, + visibleMs: 10, + backgroundBaseMs: 100, + backgroundStepMs: 0, + backgroundMaxMs: 100, + }) + queue.focus("active") + queue.setVisible("visible", true) + queue.push(update("v", "v", "visible")) + queue.push(update("b", "b", "background")) + await sleep(25) + const first = items(sent) + expect(first).toHaveLength(1) + expect(first[0]!.sessionID).toBe("visible") + expect(queue.stats().visible).toBe(1) + expect(queue.stats().background).toBe(0) + await sleep(100) + expect(items(sent).some((msg) => msg.sessionID === "background")).toBe(true) + queue.dispose() + }) + + it("active sessions still win over visible sessions", async () => { + const sent: Sent[] = [] + const queue = new SessionStreamScheduler((msg) => sent.push(msg), { + activeMs: 5, + visibleMs: 40, + backgroundBaseMs: 100, + backgroundStepMs: 0, + backgroundMaxMs: 100, + }) + queue.focus("sess-1") + queue.setVisible("sess-1", true) + queue.push(update("a", "a", "sess-1")) + await sleep(15) + expect(items(sent).map((msg) => msg.sessionID)).toEqual(["sess-1"]) + expect(queue.stats().active).toBe(1) + expect(queue.stats().visible).toBe(0) + queue.dispose() + }) + + it("focus flushes queued visible updates immediately", () => { + const sent: Sent[] = [] + const queue = new SessionStreamScheduler((msg) => sent.push(msg), { + activeMs: 50, + visibleMs: 100, + backgroundBaseMs: 100, + backgroundStepMs: 0, + backgroundMaxMs: 100, + }) + queue.setVisible("sess-2", true) + queue.push(update("v", "v", "sess-2")) + queue.focus("sess-2") + expect(items(sent).map((msg) => msg.sessionID)).toEqual(["sess-2"]) + expect(queue.stats().active).toBe(1) + queue.dispose() + }) + it("focus(A→B) flushes B immediately and schedules A on background", async () => { const sent: Sent[] = [] const queue = new SessionStreamScheduler((msg) => sent.push(msg), { @@ -401,6 +461,6 @@ describe("SessionStreamScheduler / stats", () => { // Single flush produces one batch message. expect(stats.batches).toBe(1) // Lane counters are incremented once per emission (batch or single). - expect(stats.active + stats.background).toBe(stats.batches) + expect(stats.active + stats.visible + stats.background).toBe(stats.batches) }) }) diff --git a/packages/kilo-vscode/tests/unit/session-utils.test.ts b/packages/kilo-vscode/tests/unit/session-utils.test.ts index e308a576e3d..c3499539245 100644 --- a/packages/kilo-vscode/tests/unit/session-utils.test.ts +++ b/packages/kilo-vscode/tests/unit/session-utils.test.ts @@ -5,12 +5,18 @@ import { calcContextUsage, buildFamilyCosts, buildFamilyParents, + buildFamilyParentsFromTools, buildFamilyLabels, + buildFamilyLabelsFromTools, buildCostBreakdown, + buildSessionToolParts, collapseCostBreakdown, childID, + removeSessionToolPart, + removeSessionToolPartsForMessage, + upsertSessionToolPart, } from "../../webview-ui/src/context/session-utils" -import type { Part } from "../../webview-ui/src/types/messages" +import type { Message, Part, ToolPart } from "../../webview-ui/src/types/messages" const t = (key: string) => key @@ -149,6 +155,7 @@ function msg(id: string, role: string, cost?: number) { function toolPart(tool: string, sessionId?: string, input?: { subagent_type?: string; description?: string }) { return { + id: `part-${tool}-${sessionId ?? "none"}`, type: "tool" as const, tool, state: { @@ -158,6 +165,14 @@ function toolPart(tool: string, sessionId?: string, input?: { subagent_type?: st } } +function textPart(id: string, text = "text"): Part { + return { id, type: "text" as const, text } +} + +function indexMsg(id: string, role: "user" | "assistant", parts?: Part[]): Message { + return { id, role, sessionID: "s1", createdAt: id, parts } +} + describe("childID", () => { it("reads session ID from top-level metadata", () => { expect(childID({ type: "tool", tool: "task", metadata: { sessionId: "child1" } })).toBe("child1") @@ -367,6 +382,61 @@ describe("buildFamilyParents", () => { }) }) +describe("session tool indexes", () => { + it("builds tool parts in assistant message order", () => { + const messages = [ + indexMsg("u1", "user", [toolPart("read") as ToolPart]), + indexMsg("a1", "assistant", [toolPart("read") as ToolPart, textPart("t1")]), + indexMsg("a2", "assistant", [toolPart("grep") as ToolPart]), + ] + const tools = buildSessionToolParts(messages) + expect(tools.map((part) => part.tool)).toEqual(["read", "grep"]) + expect(tools.map((part) => part.messageID)).toEqual(["a1", "a2"]) + }) + + it("uses a lookup so stashed loaded parts can feed the index", () => { + const messages = [indexMsg("a1", "assistant")] + const parts: Record = { a1: [toolPart("websearch") as ToolPart, textPart("t1")] } + const tools = buildSessionToolParts(messages, (item) => parts[item.id]) + expect(tools.map((part) => part.tool)).toEqual(["websearch"]) + }) + + it("upserts tool parts without duplicating and ignores text deltas", () => { + const first = { ...toolPart("bash"), id: "p1", state: { status: "running", input: {}, title: "old" } } + const next = { ...toolPart("bash"), id: "p1", state: { status: "running", input: {}, title: "new" } } + const indexed = upsertSessionToolPart([], first as ToolPart, { id: "m1", sessionID: "s1" }) + const updated = upsertSessionToolPart(indexed, next as ToolPart, { id: "m1", sessionID: "s1" }) + const text = upsertSessionToolPart(updated, textPart("t1"), { id: "m1", sessionID: "s1" }) + expect(text).toHaveLength(1) + expect((text[0]!.state as { title?: string }).title).toBe("new") + }) + + it("removes indexed tools by part or message", () => { + const first = upsertSessionToolPart([], { ...toolPart("read"), id: "p1" } as ToolPart, { + id: "m1", + sessionID: "s1", + }) + const second = upsertSessionToolPart(first, { ...toolPart("grep"), id: "p2" } as ToolPart, { + id: "m2", + sessionID: "s1", + }) + expect(removeSessionToolPart(second, "p1").map((part) => part.id)).toEqual(["p2"]) + expect(removeSessionToolPartsForMessage(second, "m2").map((part) => part.id)).toEqual(["p1"]) + }) + + it("derives parents and labels from indexed tool parts", () => { + const family = new Set(["root", "child"]) + const tools = new Map([ + ["root", [toolPart("task", "child", { subagent_type: "explore" })]], + ["child", []], + ]) + const parents = buildFamilyParentsFromTools(family, (sid) => tools.get(sid) ?? []) + const labels = buildFamilyLabelsFromTools(family, (sid) => tools.get(sid) ?? []) + expect(parents.get("child")).toBe("root") + expect(labels.get("child")).toBe("explore") + }) +}) + describe("buildFamilyLabels", () => { it("returns empty map when no task tool parts exist", () => { const family = new Set(["s1"]) diff --git a/packages/kilo-vscode/tests/unit/task-session.test.ts b/packages/kilo-vscode/tests/unit/task-session.test.ts index d35d7e87830..87be97fc4e6 100644 --- a/packages/kilo-vscode/tests/unit/task-session.test.ts +++ b/packages/kilo-vscode/tests/unit/task-session.test.ts @@ -1,5 +1,5 @@ import { describe, expect, it } from "bun:test" -import { childID } from "../../src/kilo-provider/task-session" +import { childID, VisibleTaskStreams } from "../../src/kilo-provider/task-session" describe("childID", () => { it("reads session ID from top-level metadata", () => { @@ -25,3 +25,33 @@ describe("childID", () => { expect(childID({ type: "tool", tool: "read", state: { metadata: { sessionId: "child3" } } })).toBeUndefined() }) }) + +describe("VisibleTaskStreams", () => { + it("suspends and restores visible child sessions with provider activity", () => { + const sent: Array<[string, boolean]> = [] + const streams = new VisibleTaskStreams((id, visible) => sent.push([id, visible])) + streams.handle({ type: "streamSessionVisible", sessionID: "child", visible: true }) + streams.setActive(false) + streams.setActive(true) + streams.handle({ type: "streamSessionVisible", sessionID: "child", visible: false }) + expect(sent).toEqual([ + ["child", true], + ["child", false], + ["child", true], + ["child", false], + ]) + }) + + it("clears visible refs after a webview reload", () => { + const sent: Array<[string, boolean]> = [] + const streams = new VisibleTaskStreams((id, visible) => sent.push([id, visible])) + streams.handle({ type: "streamSessionVisible", sessionID: "child", visible: true }) + streams.clear() + streams.handle({ type: "streamSessionVisible", sessionID: "child", visible: false }) + expect(sent).toEqual([ + ["child", true], + ["child", false], + ["child", false], + ]) + }) +}) diff --git a/packages/kilo-vscode/webview-ui/src/components/chat/TaskToolExpanded.tsx b/packages/kilo-vscode/webview-ui/src/components/chat/TaskToolExpanded.tsx index 3b0ff7e6734..e672cbae562 100644 --- a/packages/kilo-vscode/webview-ui/src/components/chat/TaskToolExpanded.tsx +++ b/packages/kilo-vscode/webview-ui/src/components/chat/TaskToolExpanded.tsx @@ -7,38 +7,19 @@ * Call registerExpandedTaskTool() once at app startup to activate. */ -import { Component, createEffect, createMemo, For, Show } from "solid-js" +import { Component, createEffect, createMemo, For, Show, onCleanup } from "solid-js" import { ToolRegistry, ToolProps, getToolInfo } from "@kilocode/kilo-ui/message-part" import { BasicTool } from "@kilocode/kilo-ui/basic-tool" import { Icon } from "@kilocode/kilo-ui/icon" import { IconButton } from "@kilocode/kilo-ui/icon-button" -import { useData } from "@kilocode/kilo-ui/context/data" import { useLanguage } from "../../context/language" import { useI18n } from "@kilocode/kilo-ui/context/i18n" import { createAutoScroll } from "@kilocode/kilo-ui/hooks" import { useSession } from "../../context/session" import { useVSCode } from "../../context/vscode" import { childID } from "../../context/session-utils" -import type { ToolPart, Message as SDKMessage } from "@kilocode/sdk/v2" - -/** Collect all tool parts from all assistant messages in a given session. */ -function getSessionToolParts(store: ReturnType["store"], sessionId: string): ToolPart[] { - const messages = (store.message?.[sessionId] as SDKMessage[] | undefined)?.filter((m) => m.role === "assistant") - if (!messages) return [] - const parts: ToolPart[] = [] - for (const m of messages) { - const msgParts = store.part?.[m.id] - if (msgParts) { - for (const p of msgParts) { - if (p && p.type === "tool") parts.push(p as ToolPart) - } - } - } - return parts -} const TaskToolRenderer: Component = (props) => { - const data = useData() const i18n = useI18n() const language = useLanguage() const session = useSession() @@ -73,7 +54,24 @@ const TaskToolRenderer: Component = (props) => { const childToolParts = createMemo(() => { const id = childSessionId() if (!id) return [] - return getSessionToolParts(data.store, id) + return session.getSessionToolParts(id) + }) + + const childToolCount = createMemo(() => { + const id = childSessionId() + return id ? session.getSessionToolCount(id) : 0 + }) + + createEffect((prev: string | undefined) => { + const id = childSessionId() + if (prev && prev !== id) vscode.postMessage({ type: "streamSessionVisible", sessionID: prev, visible: false }) + if (id && id !== prev) vscode.postMessage({ type: "streamSessionVisible", sessionID: id, visible: true }) + return id + }) + + onCleanup(() => { + const id = childSessionId() + if (id) vscode.postMessage({ type: "streamSessionVisible", sessionID: id, visible: false }) }) const autoScroll = createAutoScroll({ @@ -93,11 +91,11 @@ const TaskToolRenderer: Component = (props) => { {title()} - 0}> + 0}> {description()} - 0}> - {description() ? " " : ""}({childToolParts().length}) + 0}> + {description() ? " " : ""}({childToolCount()}) @@ -127,7 +125,7 @@ const TaskToolRenderer: Component = (props) => { >
- +
{language.t("session.messages.taskStarting")}
diff --git a/packages/kilo-vscode/webview-ui/src/context/part-stash.ts b/packages/kilo-vscode/webview-ui/src/context/part-stash.ts index 9d622ce9fc4..0d01b69ca5c 100644 --- a/packages/kilo-vscode/webview-ui/src/context/part-stash.ts +++ b/packages/kilo-vscode/webview-ui/src/context/part-stash.ts @@ -36,6 +36,16 @@ export class PartStash { this.map.delete(messageID) } + /** Remove one stale part while keeping the message stash authoritative. */ + removePart(messageID: string, partID: string): void { + const parts = this.map.get(messageID) + if (!parts) return + this.map.set( + messageID, + parts.filter((part) => part.id !== partID), + ) + } + /** * Collect parts for the given IDs, consuming the stash. Used by the * virtualizer when a turn is about to render: the returned parts should diff --git a/packages/kilo-vscode/webview-ui/src/context/session-utils.ts b/packages/kilo-vscode/webview-ui/src/context/session-utils.ts index 2544246964a..e39f90bc12c 100644 --- a/packages/kilo-vscode/webview-ui/src/context/session-utils.ts +++ b/packages/kilo-vscode/webview-ui/src/context/session-utils.ts @@ -1,4 +1,4 @@ -import type { Part } from "../types/messages" +import type { Message, Part, ToolPart } from "../types/messages" export const SNAPSHOT_PROGRESS_TEXT = "Initializing snapshot..." @@ -19,7 +19,7 @@ export type CostMessage = { id: string; role: string; cost?: number } /** Minimal tool part shape for label extraction. */ type ToolState = { - input?: { description?: string; subagent_type?: string } + input?: Record metadata?: { sessionId?: string } } @@ -35,6 +35,63 @@ export function childID(part: TaskPart): string | undefined { return part.metadata?.sessionId ?? part.state?.metadata?.sessionId } +function stringField(value: unknown): string | undefined { + return typeof value === "string" ? value : undefined +} + +function withMessage(part: ToolPart, msg: { id: string; sessionID?: string }): ToolPart { + return { + ...part, + messageID: part.messageID ?? msg.id, + sessionID: part.sessionID ?? msg.sessionID, + } +} + +export type ToolIndexMessage = Pick + +/** + * Build the per-session compact tool index in assistant-message order. + * Text/reasoning deltas should not touch this index, keeping streaming cheap. + */ +export function buildSessionToolParts( + messages: ToolIndexMessage[], + lookup?: (message: ToolIndexMessage) => Part[] | undefined, +): ToolPart[] { + const tools: ToolPart[] = [] + for (const msg of messages) { + if (msg.role !== "assistant") continue + const parts = lookup?.(msg) ?? msg.parts + if (!parts) continue + for (const part of parts) { + if (part.type !== "tool") continue + tools.push(withMessage(part, msg)) + } + } + return tools +} + +export function upsertSessionToolPart( + current: readonly ToolPart[], + part: Part, + msg: { id: string; sessionID?: string }, +): ToolPart[] { + if (part.type !== "tool") return current.slice() + const next = withMessage(part, msg) + const index = current.findIndex((item) => item.id === part.id) + if (index < 0) return [...current, next] + const tools = current.slice() + tools[index] = next + return tools +} + +export function removeSessionToolPart(current: readonly ToolPart[], partID: string): ToolPart[] { + return current.filter((part) => part.id !== partID) +} + +export function removeSessionToolPartsForMessage(current: readonly ToolPart[], messageID: string): ToolPart[] { + return current.filter((part) => part.messageID !== messageID) +} + /** * Derive a human-readable status string from the last streaming part. * Returns undefined for part types that don't map to a status. @@ -148,19 +205,24 @@ export function buildFamilyParents( family: Set, messages: Record, parts: Record, +): Map { + return buildFamilyParentsFromTools(family, (sid) => { + const msgs = messages[sid] + if (!msgs) return [] + return msgs.flatMap((msg) => parts[msg.id] ?? []) + }) +} + +export function buildFamilyParentsFromTools( + family: Set, + tools: (sessionID: string) => readonly TaskPart[], ): Map { const parents = new Map() for (const sid of family) { - const msgs = messages[sid] - if (!msgs) continue - for (const msg of msgs) { - const list = parts[msg.id] - if (!list) continue - for (const p of list) { - const child = childID(p) - if (!child || !family.has(child) || parents.has(child)) continue - parents.set(child, sid) - } + for (const p of tools(sid)) { + const child = childID(p) + if (!child || !family.has(child) || parents.has(child)) continue + parents.set(child, sid) } } return parents @@ -176,22 +238,27 @@ export function buildFamilyLabels( family: Set, messages: Record, parts: Record, +): Map { + return buildFamilyLabelsFromTools(family, (sid) => { + const msgs = messages[sid] + if (!msgs) return [] + return msgs.flatMap((msg) => parts[msg.id] ?? []) + }) +} + +export function buildFamilyLabelsFromTools( + family: Set, + tools: (sessionID: string) => readonly TaskPart[], ): Map { const labels = new Map() for (const sid of family) { - const msgs = messages[sid] - if (!msgs) continue - for (const msg of msgs) { - const list = parts[msg.id] - if (!list) continue - for (const p of list) { - if (p.type !== "tool") continue - const child = childID(p) - if (!child || !family.has(child)) continue - const raw = p.state?.input?.subagent_type || p.state?.input?.description || p.tool || "task" - const desc = raw.length > LABEL_CAP ? raw.slice(0, LABEL_CAP - 2) + "…" : raw - if (!labels.has(child)) labels.set(child, desc) - } + for (const p of tools(sid)) { + if (p.type !== "tool") continue + const child = childID(p) + if (!child || !family.has(child)) continue + const raw = stringField(p.state?.input?.subagent_type) ?? stringField(p.state?.input?.description) ?? p.tool ?? "task" + const desc = raw.length > LABEL_CAP ? raw.slice(0, LABEL_CAP - 2) + "…" : raw + if (!labels.has(child)) labels.set(child, desc) } } return labels diff --git a/packages/kilo-vscode/webview-ui/src/context/session.tsx b/packages/kilo-vscode/webview-ui/src/context/session.tsx index cc0e4db56ad..8e4fa91eae7 100644 --- a/packages/kilo-vscode/webview-ui/src/context/session.tsx +++ b/packages/kilo-vscode/webview-ui/src/context/session.tsx @@ -33,16 +33,21 @@ import type { SendMessageFailedMessage, McpStatusEntry, MessageLoadMode, + ToolPart, } from "../types/messages" import { removeSessionPermissions, upsertPermission } from "./permission-queue" import { computeStatus, calcContextUsage, buildFamilyCosts, - buildFamilyParents, - buildFamilyLabels, + buildFamilyParentsFromTools, + buildFamilyLabelsFromTools, buildCostBreakdown, + buildSessionToolParts, childID, + removeSessionToolPart, + removeSessionToolPartsForMessage, + upsertSessionToolPart, } from "./session-utils" import { Identifier } from "../utils/id" import { resolveModelSelection } from "./model-selection" @@ -80,6 +85,7 @@ interface SessionStore { sessions: Record messages: Record // sessionID -> messages parts: Record // messageID -> parts + toolParts: Record // sessionID -> compact per-session tool index todos: Record // sessionID -> todos modelSelections: Record // agentName -> model (global, extension-lifetime) sessionOverrides: Record // sessionID -> per-session model override (compare mode) @@ -129,6 +135,10 @@ interface SessionContextValue { // Parts for a specific message getParts: (messageID: string) => Part[] + // Tool parts for a specific session, maintained incrementally for streaming views + getSessionToolParts: (sessionID: string) => ToolPart[] + getSessionToolCount: (sessionID: string) => number + // Hidden after model changes so switching models can clear stale provider errors // without removing messages and their checkpoint restore actions. isErrorHidden: (messageID: string) => boolean @@ -398,6 +408,7 @@ export const SessionProvider: ParentComponent = (props) => { sessions: {}, messages: {}, parts: {}, + toolParts: {}, todos: {}, modelSelections: {}, sessionOverrides: {}, @@ -937,6 +948,7 @@ export const SessionProvider: ParentComponent = (props) => { if (!store.messages[session.id]?.length) { setStore("messages", session.id, []) } + if (!store.toolParts[session.id]) setStore("toolParts", session.id, []) const pendingAgent = draftID ? store.agentSelections[draftID] : pendingAgentSelection() const pendingModel = draftID ? store.sessionOverrides[draftID] : undefined @@ -1045,6 +1057,38 @@ export const SessionProvider: ParentComponent = (props) => { return true } + function rebuildToolParts(sessionID: string, messages: Message[], parts?: Record) { + const tools = buildSessionToolParts( + messages, + (msg) => parts?.[msg.id] ?? store.parts[msg.id] ?? stash.peek(msg.id) ?? msg.parts, + ) + setStore("toolParts", sessionID, tools) + } + + function messageParts(messages: Message[]): Record { + const parts: Record = {} + for (const msg of messages) { + if (msg.parts && msg.parts.length > 0) parts[msg.id] = msg.parts + } + return parts + } + + function patchToolPart(sessionID: string | undefined, messageID: string, part: Part) { + const sid = sessionID ?? part.sessionID + if (!sid) return + if (part.type !== "tool") return + setStore("toolParts", sid, (tools = []) => upsertSessionToolPart(tools, part, { id: messageID, sessionID: sid })) + } + + function dropToolPart(sessionID: string | undefined, partID: string) { + if (!sessionID) return + setStore("toolParts", sessionID, (tools = []) => removeSessionToolPart(tools, partID)) + } + + function dropMessageTools(sessionID: string, messageID: string) { + setStore("toolParts", sessionID, (tools = []) => removeSessionToolPartsForMessage(tools, messageID)) + } + function handleMessagesLoaded( sessionID: string, messages: Message[], @@ -1057,6 +1101,11 @@ export const SessionProvider: ParentComponent = (props) => { // message+part-count already agrees with the server. Skip the reactive // store churn entirely — virtualizer and rendering stay untouched. if (mode === "reconcile" && sameReconcileShape(store.messages[sessionID] ?? [], messages)) { + const parts = messageParts(messages) + for (const msg of messages) { + if (store.parts[msg.id]) delete parts[msg.id] + } + rebuildToolParts(sessionID, messages, parts) patchPage(sessionID, { initialLoaded: true, lastMutation: "update" }) return } @@ -1075,6 +1124,7 @@ export const SessionProvider: ParentComponent = (props) => { mode === "prepend" || mode === "reconcile" ? mergeMessages(current, messages, mode) : withPending(sessionID, messages) + const loadedParts: Record = {} // "replace" mode (session switch): assign directly — reconcile's O(n) // diff is unnecessary when the entire list is new, and its reactive // proxy creation for each message object dominated the trace (~900ms). @@ -1087,6 +1137,7 @@ export const SessionProvider: ParentComponent = (props) => { for (const msg of messages) { if (!msg.parts || msg.parts.length === 0) continue + loadedParts[msg.id] = msg.parts if (mode === "reconcile" && store.parts[msg.id]) { // Reconcile on a message already hydrated into the reactive store: // write parts directly so visible turns pick up the server- @@ -1101,6 +1152,8 @@ export const SessionProvider: ParentComponent = (props) => { } } + rebuildToolParts(sessionID, merged, loadedParts) + // "reconcile" is a background tail refresh, not a page navigation — // preserve the existing pagination cursor/hasMore so "load earlier" // keeps working. @@ -1163,6 +1216,7 @@ export const SessionProvider: ParentComponent = (props) => { stash.remove(message.id) setStore("parts", message.id, message.parts) } + rebuildToolParts(message.sessionID, store.messages[message.sessionID] ?? []) } function handlePartUpdated( @@ -1180,6 +1234,7 @@ export const SessionProvider: ParentComponent = (props) => { } if (sessionID) patchPage(sessionID, { lastMutation: "update" }) + patchToolPart(sessionID, effectiveMessageID, part) // If the stash has parts for this message, hydrate them first so the // SSE update merges into the full part list rather than an empty array. @@ -1227,6 +1282,8 @@ export const SessionProvider: ParentComponent = (props) => { function handlePartRemoved(sessionID: string | undefined, messageID: string, partID: string) { if (sessionID) patchPage(sessionID, { lastMutation: "update" }) + dropToolPart(sessionID, partID) + stash.removePart(messageID, partID) setStore( "parts", @@ -1392,6 +1449,7 @@ export const SessionProvider: ParentComponent = (props) => { stash.remove(message.messageID) batch(() => { setStore("messages", message.sessionID!, (msgs = []) => msgs.filter((m) => m.id !== message.messageID)) + dropMessageTools(message.sessionID!, message.messageID!) setStore( "parts", produce((parts) => { @@ -1422,25 +1480,21 @@ export const SessionProvider: ParentComponent = (props) => { const queue = [rootID] while (queue.length > 0) { const sid = queue.pop()! - const msgs = source(sid) - for (const msg of msgs) { - const parts = store.parts[msg.id] - if (!parts) continue - for (const p of parts) { - if (p.type !== "tool") continue - // Webview ToolState omits runtime metadata; task parts still carry it from the backend. - const child = childID( - p as { - type: string - tool?: string - metadata?: { sessionId?: string } - state?: { metadata?: { sessionId?: string } } - }, - ) - if (child && !ids.has(child)) { - ids.add(child) - queue.push(child) - } + const visible = new Set(source(sid).map((msg) => msg.id)) + for (const p of store.toolParts[sid] ?? []) { + if (p.messageID && !visible.has(p.messageID)) continue + // Webview ToolState omits runtime metadata; task parts still carry it from the backend. + const child = childID( + p as { + type: string + tool?: string + metadata?: { sessionId?: string } + state?: { metadata?: { sessionId?: string } } + }, + ) + if (child && !ids.has(child)) { + ids.add(child) + queue.push(child) } } } @@ -1549,6 +1603,12 @@ export const SessionProvider: ParentComponent = (props) => { } }), ) + setStore( + "toolParts", + produce((parts) => { + delete parts[sessionID] + }), + ) setStore( "todos", produce((todos) => { @@ -1618,6 +1678,7 @@ export const SessionProvider: ParentComponent = (props) => { // Splices the message from the store and deletes its parts. function handleMessageRemoved(sessionID: string, messageID: string) { setStore("messages", sessionID, (msgs = []) => msgs.filter((m) => m.id !== messageID)) + dropMessageTools(sessionID, messageID) clearHiddenErrors([messageID]) setStore( "parts", @@ -1654,6 +1715,7 @@ export const SessionProvider: ParentComponent = (props) => { setStore("parts", msg.id, msg.parts) } } + rebuildToolParts(key, messages) setCurrentSessionID(key) setLoading(false) }) @@ -1678,6 +1740,7 @@ export const SessionProvider: ParentComponent = (props) => { // Carry over cloud messages so there's no loading flash setStore("messages", session.id, cloudMessages) + rebuildToolParts(session.id, cloudMessages) setCloudPreviewId(null) setCurrentSessionID(session.id) @@ -1718,6 +1781,12 @@ export const SessionProvider: ParentComponent = (props) => { delete messages[cloudKey] }), ) + setStore( + "toolParts", + produce((parts) => { + delete parts[cloudKey] + }), + ) }) // Load real messages in the background (picks up server-assigned IDs // and the new user message once the send completes via SSE) @@ -2178,6 +2247,10 @@ export const SessionProvider: ParentComponent = (props) => { return store.parts[messageID] || stash.peek(messageID) || [] } + const getSessionToolParts = (sessionID: string) => store.toolParts[sessionID] ?? [] + + const getSessionToolCount = (sessionID: string) => store.toolParts[sessionID]?.length ?? 0 + function hydrateParts(ids: string[]) { const pending = stash.take(ids, (id) => Boolean(store.parts[id])) if (Object.keys(pending).length === 0) return @@ -2274,7 +2347,7 @@ export const SessionProvider: ParentComponent = (props) => { const family = visibleFamily(id) const msgs: Record = {} for (const sid of family) msgs[sid] = visible(sid) - const parents = buildFamilyParents(family, msgs, store.parts) + const parents = buildFamilyParentsFromTools(family, (sid) => store.toolParts[sid] ?? []) return buildFamilyCosts(family, msgs, store.sessions, parents) }) @@ -2285,7 +2358,7 @@ export const SessionProvider: ParentComponent = (props) => { const family = visibleFamily(id) const msgs: Record = {} for (const sid of family) msgs[sid] = visible(sid) - return buildFamilyLabels(family, msgs, store.parts) + return buildFamilyLabelsFromTools(family, (sid) => store.toolParts[sid] ?? []) }) /** Combined cost breakdown with labels. */ @@ -2353,6 +2426,8 @@ export const SessionProvider: ParentComponent = (props) => { visibleMessages, userMessages, getParts, + getSessionToolParts, + getSessionToolCount, isErrorHidden: (messageID: string) => hiddenErrors().has(messageID), hydrateParts, todos, diff --git a/packages/kilo-vscode/webview-ui/src/types/messages/webview-messages.ts b/packages/kilo-vscode/webview-ui/src/types/messages/webview-messages.ts index 9c70feeb759..e1cd58af660 100644 --- a/packages/kilo-vscode/webview-ui/src/types/messages/webview-messages.ts +++ b/packages/kilo-vscode/webview-ui/src/types/messages/webview-messages.ts @@ -390,6 +390,12 @@ export interface RequestTimelineSettingMessage { type: "requestTimelineSetting" } +export interface StreamSessionVisibleMessage { + type: "streamSessionVisible" + sessionID: string + visible: boolean +} + export interface RequestBrowserSettingsMessage { type: "requestBrowserSettings" } @@ -1118,6 +1124,7 @@ export type WebviewMessage = | ChatCompletionAcceptedMessage | UpdateSettingRequest | RequestTimelineSettingMessage + | StreamSessionVisibleMessage | RequestBrowserSettingsMessage | RequestClaudeCompatSettingMessage | RequestConfigMessage