From 3f5d0893bbf02ceb1ef26ffa2b20e6b21c144856 Mon Sep 17 00:00:00 2001 From: Marius Date: Wed, 18 Mar 2026 10:11:42 +0100 Subject: [PATCH] perf(vscode): fix performance issues with write tool (#7211) * perf(vscode): fix performance issues with write tool * refactor: remove packages/ui changes, defer to follow-up issue * fix: preserve metadata.diagnostics in slimMeta --- packages/kilo-vscode/src/KiloProvider.ts | 72 ++++++++++++++++- .../src/agent-manager/AgentManagerProvider.ts | 4 +- packages/kilo-vscode/webview-ui/src/App.tsx | 23 ++---- .../src/components/chat/TaskToolExpanded.tsx | 3 +- .../webview-ui/src/context/session.tsx | 79 +++++++++++++++++-- .../webview-ui/src/stories/StoryProviders.tsx | 1 + .../src/stories/history.stories.tsx | 1 + 7 files changed, 158 insertions(+), 25 deletions(-) diff --git a/packages/kilo-vscode/src/KiloProvider.ts b/packages/kilo-vscode/src/KiloProvider.ts index 510481e8cc9..a53afedd854 100644 --- a/packages/kilo-vscode/src/KiloProvider.ts +++ b/packages/kilo-vscode/src/KiloProvider.ts @@ -40,6 +40,7 @@ import { resolveProjectDirectory } from "./project-directory" type KiloProviderOptions = { projectDirectory?: string | null + slimEditMetadata?: boolean } export class KiloProvider implements vscode.WebviewViewProvider, TelemetryPropertiesProvider { @@ -90,6 +91,7 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper private ignoreControllerDir: string | null = null private marketplace: MarketplaceService | null = null private projectDirectory: string | null | undefined + private slimEditMetadata = true /** Optional interceptor called before the standard message handler. * Return null to consume the message, or return a (possibly transformed) message. */ @@ -102,6 +104,7 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper options?: KiloProviderOptions, ) { this.projectDirectory = options?.projectDirectory + this.slimEditMetadata = options?.slimEditMetadata ?? true TelemetryProxy.getInstance().setProvider(this) } @@ -135,6 +138,63 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper } } + // Edit tool parts carry full file contents in metadata.filediff.before/after. + // A session with many edits can produce multi-MB payloads serialized through + // postMessage on every session switch. Stripping those strings down to just + // file path + addition/deletion counts eliminates the dominant cost. + + private slimMeta(meta: unknown) { + if (!this.slimEditMetadata) return meta + if (!meta || typeof meta !== "object") return undefined + + const obj = meta as Record + const filediff = obj.filediff + if (!filediff || typeof filediff !== "object") return undefined + + const diff = filediff as Record + const file = typeof diff.file === "string" ? diff.file : undefined + const additions = typeof diff.additions === "number" ? diff.additions : 0 + const deletions = typeof diff.deletions === "number" ? diff.deletions : 0 + + const result: Record = { + filediff: { + ...(file ? { file } : {}), + additions, + deletions, + }, + } + // Preserve diagnostics so post-edit LSP errors still render + if (obj.diagnostics) result.diagnostics = obj.diagnostics + return result + } + + /** Strip heavy metadata from a single edit tool part; pass-through for all other part types. */ + private slimPart(part: T): T { + if (!this.slimEditMetadata) return part + if (!part || typeof part !== "object") return part + + const obj = part as Record + if (obj.type !== "tool" || obj.tool !== "edit") return part + + const state = obj.state + if (!state || typeof state !== "object") return part + + const next = { ...(state as Record) } + const meta = this.slimMeta(next.metadata) + if (meta) next.metadata = meta + else delete next.metadata + + return { + ...obj, + state: next, + } as T + } + + private slimParts(parts: T[]) { + if (!this.slimEditMetadata) return parts + return parts.map((part) => this.slimPart(part)) + } + /** * Synchronize current extension-side state to the webview. * This is primarily used after a webview refresh where early postMessage calls @@ -939,7 +999,7 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper const messages = messagesData.map((m) => ({ ...m.info, - parts: m.parts, + parts: this.slimParts(m.parts), createdAt: new Date(m.info.time.created).toISOString(), })) @@ -988,7 +1048,7 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper const messages = messagesData.map((m) => ({ ...m.info, - parts: m.parts, + parts: this.slimParts(m.parts), createdAt: new Date(m.info.time.created).toISOString(), })) @@ -1002,6 +1062,7 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper messages, }) } catch (err) { + this.syncedChildSessions.delete(sessionID) console.error("[Kilo New] KiloProvider: Failed to sync child session:", err) } } @@ -2281,6 +2342,13 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper const msg = mapSSEEventToWebviewMessage(event, sessionID) if (msg) { + if (msg.type === "partUpdated") { + this.postMessage({ + ...msg, + part: this.slimPart(msg.part), + }) + return + } this.postMessage(msg) } } diff --git a/packages/kilo-vscode/src/agent-manager/AgentManagerProvider.ts b/packages/kilo-vscode/src/agent-manager/AgentManagerProvider.ts index ee235e930fa..162c6b87b8b 100644 --- a/packages/kilo-vscode/src/agent-manager/AgentManagerProvider.ts +++ b/packages/kilo-vscode/src/agent-manager/AgentManagerProvider.ts @@ -148,7 +148,9 @@ export class AgentManagerProvider implements vscode.Disposable { panel.webview.html = this.getHtml(panel.webview) - this.provider = new KiloProvider(this.extensionUri, this.connectionService, this.context) + this.provider = new KiloProvider(this.extensionUri, this.connectionService, this.context, { + slimEditMetadata: true, + }) this.provider.attachToWebview(panel.webview, { onBeforeMessage: (msg) => this.onMessage(msg), }) diff --git a/packages/kilo-vscode/webview-ui/src/App.tsx b/packages/kilo-vscode/webview-ui/src/App.tsx index 67d76e59c9f..7b52ba9f318 100644 --- a/packages/kilo-vscode/webview-ui/src/App.tsx +++ b/packages/kilo-vscode/webview-ui/src/App.tsx @@ -67,25 +67,16 @@ export const DataBridge: Component<{ children: any }> = (props) => { const data = createMemo(() => { const id = session.currentSessionID() - const allParts = session.allParts() - // Expose ALL session messages (including child sessions from sub-agents), - // not just the current session. This lets VscodeSessionTurn and - // TaskToolExpanded read child session data from the DataProvider store. - const allMessages = Object.fromEntries( - Object.entries(session.allMessages() as Record) - .filter(([, msgs]) => (msgs as SDKMessage[]).length > 0) - .map(([sid, msgs]) => [sid, msgs as SDKMessage[]]), - ) + const family = session.familyData(id) return { session: session.sessions().map((s) => ({ ...s, id: s.id, role: "user" as const })) as unknown as any[], - session_status: session.allStatusMap() as unknown as Record, + session_status: family.status as unknown as Record, session_diff: {} as Record, - message: allMessages, - part: Object.fromEntries( - Object.entries(allParts) - .filter(([, parts]) => (parts as SDKPart[]).length > 0) - .map(([msgId, parts]) => [msgId, parts as unknown as SDKPart[]]), - ), + // Restrict chat data to the selected session family (self + subagents). + // This keeps unrelated tracked sessions from invalidating the visible + // chat tree during streaming or background updates. + message: family.messages as Record, + part: family.parts as Record, permission: (() => { const grouped: Record = {} for (const p of session.permissions()) { 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 7c45b86265b..59ce576303c 100644 --- a/packages/kilo-vscode/webview-ui/src/components/chat/TaskToolExpanded.tsx +++ b/packages/kilo-vscode/webview-ui/src/components/chat/TaskToolExpanded.tsx @@ -49,7 +49,8 @@ const TaskToolRenderer: Component = (props) => { const running = createMemo(() => props.status === "pending" || props.status === "running") - // Sync child session into store whenever we have a sessionId + // Warm child session data immediately so completed task tools already have + // their compact child tool list available when the user expands them. createEffect(() => { const id = childSessionId() if (!id) return diff --git a/packages/kilo-vscode/webview-ui/src/context/session.tsx b/packages/kilo-vscode/webview-ui/src/context/session.tsx index 700d31cc41b..1a502552861 100644 --- a/packages/kilo-vscode/webview-ui/src/context/session.tsx +++ b/packages/kilo-vscode/webview-ui/src/context/session.tsx @@ -16,7 +16,7 @@ import { Accessor, batch, } from "solid-js" -import { createStore, produce } from "solid-js/store" +import { createStore, produce, reconcile } from "solid-js/store" import { useVSCode } from "./vscode" import { useServer } from "./server" import { useProvider } from "./provider" @@ -87,6 +87,13 @@ interface SessionContextValue { // All session statuses keyed by sessionID (for DataBridge) allStatusMap: () => Record + // Current session family data (self + subagents) for DataBridge + familyData: (sessionID: string | undefined) => { + messages: Record + parts: Record + status: Record + } + // Parts for a specific message getParts: (messageID: string) => Part[] @@ -197,6 +204,7 @@ export const SessionProvider: ParentComponent = (props) => { } const [loading, setLoading] = createSignal(false) + const [loaded, setLoaded] = createSignal>(new Set()) // Pending permissions const [permissions, setPermissions] = createSignal([]) @@ -566,6 +574,12 @@ export const SessionProvider: ParentComponent = (props) => { function handleMessagesLoaded(sessionID: string, messages: Message[]) { batch(() => { + setLoaded((prev) => { + if (prev.has(sessionID)) return prev + const next = new Set(prev) + next.add(sessionID) + return next + }) if (sessionID === currentSessionID()) setLoading(false) // Preserve optimistic messages that haven't been confirmed yet. @@ -576,15 +590,15 @@ export const SessionProvider: ParentComponent = (props) => { const loadedIds = new Set(messages.map((m) => m.id)) const current = store.messages[sessionID] ?? [] const orphans = current.filter((m) => pending.has(m.id) && !loadedIds.has(m.id)) - setStore("messages", sessionID, [...messages, ...orphans]) + setStore("messages", sessionID, reconcile([...messages, ...orphans], { key: "id" })) } else { - setStore("messages", sessionID, messages) + setStore("messages", sessionID, reconcile(messages, { key: "id" })) } // Also extract parts from messages for (const msg of messages) { if (msg.parts && msg.parts.length > 0) { - setStore("parts", msg.id, msg.parts) + setStore("parts", msg.id, reconcile(msg.parts, { key: "id" })) } } }) @@ -806,6 +820,42 @@ export const SessionProvider: ParentComponent = (props) => { return family } + function familyData(sessionID: string | undefined) { + if (!sessionID) { + return { + messages: {}, + parts: {}, + status: {}, + } + } + + const family = sessionFamily(sessionID) + const messages: Record = {} + const parts: Record = {} + const status: Record = {} + + for (const sid of family) { + const msgs = store.messages[sid] + if (msgs?.length) { + messages[sid] = msgs + for (const msg of msgs) { + const item = store.parts[msg.id] + if (!item?.length) continue + parts[msg.id] = item + } + } + + const info = statusMap[sid] + if (info) status[sid] = info + } + + return { + messages, + parts, + status, + } + } + /** Return permissions scoped to the given session's family (self + subagents). */ function scopedPermissions(sessionID: string | undefined): PermissionRequest[] { if (!sessionID) return [] @@ -930,6 +980,12 @@ export const SessionProvider: ParentComponent = (props) => { if (cloudPreviewId() !== cloudSessionId) return const key = `cloud:${cloudSessionId}` batch(() => { + setLoaded((prev) => { + if (prev.has(key)) return prev + const next = new Set(prev) + next.add(key) + return next + }) setStore("sessions", key, { id: key, title, @@ -951,6 +1007,12 @@ export const SessionProvider: ParentComponent = (props) => { const cloudKey = `cloud:${cloudSessionId}` const cloudMessages = store.messages[cloudKey] ?? [] batch(() => { + setLoaded((prev) => { + const next = new Set(prev) + next.add(session.id) + next.delete(cloudKey) + return next + }) setStore("sessions", session.id, session) const pendingAgent = pendingAgentSelection() @@ -1223,7 +1285,7 @@ export const SessionProvider: ParentComponent = (props) => { return } setCurrentSessionID(id) - setLoading(true) + setLoading(!loaded().has(id)) vscode.postMessage({ type: "loadMessages", sessionID: id }) } @@ -1251,6 +1313,12 @@ export const SessionProvider: ParentComponent = (props) => { delete sessions[id] }), ) + setLoaded((prev) => { + if (!prev.has(id)) return prev + const next = new Set(prev) + next.delete(id) + return next + }) vscode.postMessage({ type: "deleteSession", sessionID: id }) } @@ -1421,6 +1489,7 @@ export const SessionProvider: ParentComponent = (props) => { allMessages, allParts, allStatusMap, + familyData, variantList, currentVariant, selectVariant, diff --git a/packages/kilo-vscode/webview-ui/src/stories/StoryProviders.tsx b/packages/kilo-vscode/webview-ui/src/stories/StoryProviders.tsx index 3cdf84c201a..94d06670053 100644 --- a/packages/kilo-vscode/webview-ui/src/stories/StoryProviders.tsx +++ b/packages/kilo-vscode/webview-ui/src/stories/StoryProviders.tsx @@ -131,6 +131,7 @@ export function mockSessionValue(overrides?: { allMessages: () => ({}), allParts: () => ({}), allStatusMap: () => ({}), + familyData: () => ({ messages: {}, parts: {}, status: {} }), getParts: () => [], todos: () => [], permissions: () => permissions, diff --git a/packages/kilo-vscode/webview-ui/src/stories/history.stories.tsx b/packages/kilo-vscode/webview-ui/src/stories/history.stories.tsx index 1271b0dac3e..c0cffaea1cf 100644 --- a/packages/kilo-vscode/webview-ui/src/stories/history.stories.tsx +++ b/packages/kilo-vscode/webview-ui/src/stories/history.stories.tsx @@ -61,6 +61,7 @@ const WithSessions: ParentComponent<{ sessions?: typeof mockSessions }> = (props allMessages: () => ({}), allParts: () => ({}), allStatusMap: () => ({}), + familyData: () => ({ messages: {}, parts: {}, status: {} }), getParts: () => [], todos: () => [], permissions: () => [],