fix(vscode): speed up inline subagent streaming

.changeset/brisk-subagent-streaming.md

@@ -0,0 +1,5 @@
+---
+"kilo-code": patch
+---
+
+Keep inline subagent streaming responsive during tool-heavy sessions.

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)

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()
       },

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<typeof setTimeout> | null = null
+  private vtimer: ReturnType<typeof setTimeout> | null = null
   private btimer: ReturnType<typeof setTimeout> | null = null
   private bgFirstQueuedAt = 0
+  private visibleFirstQueuedAt = 0
   private readonly queues = new Map<string, Map<string, PartUpdate>>()
+  private readonly visible = new Set<string>()
   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<StreamSchedulerStats> {
@@ -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
   }

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<string, number>()
+  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
+  }
+}

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")])

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)
   })
 })

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<string, Part[]> = { 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"])

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],
+    ])
+  })
+})

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<typeof useData>["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<ToolProps> = (props) => {
-  const data = useData()
   const i18n = useI18n()
   const language = useLanguage()
   const session = useSession()
@@ -73,7 +54,24 @@ const TaskToolRenderer: Component<ToolProps> = (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<ToolProps> = (props) => {
         <span data-slot="basic-tool-tool-title" class="capitalize">
           {title()}
         </span>
-        <Show when={description() || childToolParts().length > 0}>
+        <Show when={description() || childToolCount() > 0}>
           <span data-slot="basic-tool-tool-subtitle">
             {description()}
-            <Show when={childToolParts().length > 0}>
-              {description() ? " " : ""}({childToolParts().length})
+            <Show when={childToolCount() > 0}>
+              {description() ? " " : ""}({childToolCount()})
             </Show>
           </span>
         </Show>
@@ -127,7 +125,7 @@ const TaskToolRenderer: Component<ToolProps> = (props) => {
       >
         <div ref={autoScroll.scrollRef} onScroll={autoScroll.handleScroll} data-component="tool-output" data-scrollable>
           <div ref={autoScroll.contentRef} data-component="task-tools">
-            <Show when={running() && childToolParts().length === 0}>
+            <Show when={running() && childToolCount() === 0}>
               <div data-slot="task-tool-item" data-state="starting">
                 <span data-slot="task-tool-title">{language.t("session.messages.taskStarting")}</span>
               </div>

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

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<string, unknown>
   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<Message, "id" | "sessionID" | "role" | "parts">
+
+/**
+ * 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<string>,
   messages: Record<string, CostMessage[]>,
   parts: Record<string, TaskPart[]>,
+): Map<string, string> {
+  return buildFamilyParentsFromTools(family, (sid) => {
+    const msgs = messages[sid]
+    if (!msgs) return []
+    return msgs.flatMap((msg) => parts[msg.id] ?? [])
+  })
+}
+
+export function buildFamilyParentsFromTools(
+  family: Set<string>,
+  tools: (sessionID: string) => readonly TaskPart[],
 ): Map<string, string> {
   const parents = new Map<string, string>()
   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<string>,
   messages: Record<string, CostMessage[]>,
   parts: Record<string, TaskPart[]>,
+): Map<string, string> {
+  return buildFamilyLabelsFromTools(family, (sid) => {
+    const msgs = messages[sid]
+    if (!msgs) return []
+    return msgs.flatMap((msg) => parts[msg.id] ?? [])
+  })
+}
+
+export function buildFamilyLabelsFromTools(
+  family: Set<string>,
+  tools: (sessionID: string) => readonly TaskPart[],
 ): Map<string, string> {
   const labels = new Map<string, string>()
   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

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<string, SessionInfo>
   messages: Record<string, Message[]> // sessionID -> messages
   parts: Record<string, Part[]> // messageID -> parts
+  toolParts: Record<string, ToolPart[]> // sessionID -> compact per-session tool index
   todos: Record<string, TodoItem[]> // sessionID -> todos
   modelSelections: Record<string, ModelSelection | null> // agentName -> model (global, extension-lifetime)
   sessionOverrides: Record<string, ModelSelection> // 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<string, Part[]>) {
+    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<string, Part[]> {
+    const parts: Record<string, Part[]> = {}
+    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<string, Part[]> = {}
       // "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<string, Message[]> = {}
     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<string, Message[]> = {}
     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,

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