Merge pull request #7617 from Kilo-Org/fix/session-diff-memory-leak

fix(cli): prevent CLI memory leaks from oversized diffs and session accumulation
This commit is contained in:
Marian Alexandru Alecu
2026-03-26 14:24:32 +02:00
committed by GitHub
7 changed files with 114 additions and 53 deletions
+11
View File
@@ -268,6 +268,17 @@ function App() {
})
// kilocode_change end
// kilocode_change start — evict per-session data from store when navigating away
createEffect(
on(
() => (route.data.type === "session" ? route.data.sessionID : undefined),
(current, prev) => {
if (prev && prev !== current) sync.session.evict(prev)
},
),
)
// kilocode_change end
// Update terminal window title based on current route and session
createEffect(() => {
if (!terminalTitleEnabled() || Flag.KILO_DISABLE_TERMINAL_TITLE) return
@@ -53,7 +53,7 @@ export const { use: useSync, provider: SyncProvider } = createSimpleContext({
[sessionID: string]: SessionStatus
}
session_diff: {
[sessionID: string]: Snapshot.FileDiff[]
[sessionID: string]: Omit<Snapshot.FileDiff, "before" | "after">[] // kilocode_change
}
todo: {
[sessionID: string]: Todo[]
@@ -107,12 +107,42 @@ export const { use: useSync, provider: SyncProvider } = createSimpleContext({
const sdk = useSDK()
const fullSyncedSessions = new Set<string>() // kilocode_change
async function syncWorkspaces() {
const result = await sdk.client.experimental.workspace.list().catch(() => undefined)
if (!result?.data) return
setStore("workspaceList", reconcile(result.data))
}
// kilocode_change start
function evict(sessionID: string) {
// Collect child session IDs so we can evict them too.
const children = store.session.filter((s) => s.parentID === sessionID).map((s) => s.id)
setStore(
produce((draft) => {
const messages = draft.message[sessionID]
if (messages) {
for (const msg of messages) delete draft.part[msg.id]
}
delete draft.message[sessionID]
delete draft.session_diff[sessionID]
delete draft.session_status[sessionID]
delete draft.todo[sessionID]
}),
)
fullSyncedSessions.delete(sessionID)
for (const child of children) evict(child)
}
// Strip summary.diffs from user messages — the TUI never reads them
// and they can carry multi-MB before/after file content strings.
function strip(msg: Message): Message {
if (msg.role !== "user" || !msg.summary?.diffs) return msg
return { ...msg, summary: { ...msg.summary, diffs: [] } } as Message
}
// kilocode_change end
sdk.event.listen((e) => {
const event = e.details
switch (event.type) {
@@ -199,21 +229,29 @@ export const { use: useSync, provider: SyncProvider } = createSimpleContext({
break
case "session.diff":
setStore("session_diff", event.properties.sessionID, event.properties.diff)
setStore(
"session_diff",
event.properties.sessionID,
event.properties.diff.map(({ before: _, after: __, ...rest }) => rest),
)
break
// kilocode_change start
case "session.deleted": {
const result = Binary.search(store.session, event.properties.info.id, (s) => s.id)
if (result.found) {
const sid = event.properties.info.id
const match = Binary.search(store.session, sid, (s) => s.id)
if (match.found) {
setStore(
"session",
produce((draft) => {
draft.splice(result.index, 1)
draft.splice(match.index, 1)
}),
)
}
evict(sid)
break
}
// kilocode_change end
case "session.updated": {
const result = Binary.search(store.session, event.properties.info.id, (s) => s.id)
if (result.found) {
@@ -234,31 +272,33 @@ export const { use: useSync, provider: SyncProvider } = createSimpleContext({
break
}
// kilocode_change start
case "message.updated": {
const messages = store.message[event.properties.info.sessionID]
const info = strip(event.properties.info)
const messages = store.message[info.sessionID]
if (!messages) {
setStore("message", event.properties.info.sessionID, [event.properties.info])
setStore("message", info.sessionID, [info])
break
}
const result = Binary.search(messages, event.properties.info.id, (m) => m.id)
const result = Binary.search(messages, info.id, (m) => m.id)
if (result.found) {
setStore("message", event.properties.info.sessionID, result.index, reconcile(event.properties.info))
setStore("message", info.sessionID, result.index, reconcile(info))
break
}
setStore(
"message",
event.properties.info.sessionID,
info.sessionID,
produce((draft) => {
draft.splice(result.index, 0, event.properties.info)
draft.splice(result.index, 0, info)
}),
)
const updated = store.message[event.properties.info.sessionID]
const updated = store.message[info.sessionID]
if (updated.length > 100) {
const oldest = updated[0]
batch(() => {
setStore(
"message",
event.properties.info.sessionID,
info.sessionID,
produce((draft) => {
draft.shift()
}),
@@ -273,6 +313,7 @@ export const { use: useSync, provider: SyncProvider } = createSimpleContext({
}
break
}
// kilocode_change end
case "message.removed": {
const messages = store.message[event.properties.sessionID]
const result = Binary.search(messages, event.properties.messageID, (m) => m.id)
@@ -441,7 +482,6 @@ export const { use: useSync, provider: SyncProvider } = createSimpleContext({
bootstrap()
})
const fullSyncedSessions = new Set<string>()
const result = {
data: store,
set: setStore,
@@ -481,15 +521,16 @@ export const { use: useSync, provider: SyncProvider } = createSimpleContext({
if (match.found) draft.session[match.index] = session.data!
if (!match.found) draft.session.splice(match.index, 0, session.data!)
draft.todo[sessionID] = todo.data ?? []
draft.message[sessionID] = messages.data!.map((x) => x.info)
draft.message[sessionID] = messages.data!.map((x) => strip(x.info)) // kilocode_change
for (const message of messages.data!) {
draft.part[message.info.id] = message.parts
}
draft.session_diff[sessionID] = diff.data ?? []
draft.session_diff[sessionID] = (diff.data ?? []).map(({ before: _, after: __, ...rest }) => rest)
}),
)
fullSyncedSessions.add(sessionID)
},
evict, // kilocode_change
},
workspace: {
get(workspaceID: string) {
+9 -9
View File
@@ -1342,16 +1342,15 @@ export namespace Config {
const files = GLOBAL_CONFIG_FILES.map((file) => path.join(Global.Path.config, file))
// also check legacy TOML config — its presence means existing user
const legacy = path.join(Global.Path.config, "config")
const existing: string[] = []
for (const file of files) {
if (existsSync(file)) existing.push(file)
}
const existing = files.filter((file) => existsSync(file))
const hasLegacy = existsSync(legacy)
// no global config → new user, they'll get the new bash:ask default
if (existing.length === 0 && !hasLegacy) return
// check if any config file already has an explicit bash permission
for (const file of existing) {
const text = await Bun.file(file).text()
const text = await Bun.file(file)
.text()
.catch(() => "")
const data = parseJsonc(text) ?? {}
if (data.permission?.bash) return
}
@@ -1372,11 +1371,12 @@ export namespace Config {
formattingOptions: { insertSpaces: true, tabSize: 2 },
})
await Bun.write(target, applyEdits(text, edits))
} else {
const data = parseJsonc(text) ?? {}
const merged = { ...data, permission: { ...data.permission, bash: "allow" } }
await Bun.write(target, JSON.stringify(merged, null, 2))
log.info("migrated bash permission to allow for existing user", { path: target })
return
}
const data = parseJsonc(text) ?? {}
const merged = { ...data, permission: { ...data.permission, bash: "allow" } }
await Bun.write(target, JSON.stringify(merged, null, 2))
log.info("migrated bash permission to allow for existing user", { path: target })
}
// kilocode_change end
+10 -19
View File
@@ -12,7 +12,6 @@ import { LLM } from "@/session/llm"
import { MessageV2 } from "@/session/message-v2"
import { Todo } from "@/session/todo"
import { Log } from "@/util/log"
import fs from "fs/promises"
import path from "path"
function toText(item: MessageV2.WithParts): string {
@@ -115,20 +114,20 @@ export namespace PlanFollowup {
export const ANSWER_NEW_SESSION = "Start new session"
export const ANSWER_CONTINUE = "Continue here"
function resolveVariant(input: { value: string | undefined; model: Provider.Model | undefined }) {
if (!input.value) return undefined
if (!input.model?.variants?.[input.value]) return undefined
return input.value
function resolveVariant(value: string | undefined, model: Provider.Model | undefined) {
if (!value) return undefined
if (!model?.variants?.[value]) return undefined
return value
}
async function resolveCodeModel(input: Pick<MessageV2.User, "model" | "variant">) {
const state =
Flag.KILO_CLIENT === "cli"
? await fs
.readFile(path.join(Global.Path.state, "model.json"), "utf-8")
? await Bun.file(path.join(Global.Path.state, "model.json"))
.text()
.then(
(item) =>
JSON.parse(item) as {
(raw) =>
JSON.parse(raw) as {
model?: Record<string, MessageV2.User["model"]>
variant?: Record<string, string | undefined>
},
@@ -142,10 +141,7 @@ export namespace PlanFollowup {
const key = `${saved.providerID}/${saved.modelID}`
return {
model: saved,
variant: resolveVariant({
value: state?.variant?.[key],
model: full,
}),
variant: resolveVariant(state?.variant?.[key], full),
}
}
}
@@ -156,12 +152,7 @@ export namespace PlanFollowup {
if (full) {
return {
model: agent.model,
variant: agent.variant
? resolveVariant({
value: agent.variant,
model: full,
})
: undefined,
variant: resolveVariant(agent.variant, full),
}
}
}
+9 -2
View File
@@ -121,16 +121,23 @@ export namespace SessionSummary {
}),
async (input) => {
const diffs = await Storage.read<Snapshot.FileDiff[]>(["session_diff", input.sessionID]).catch(() => [])
// kilocode_change start — scrub oversized diffs from stored session_diff
const next = diffs.map((item) => {
const file = unquoteGitPath(item.file)
if (file === item.file) return item
const oversized =
Buffer.byteLength(item.before) > Snapshot.MAX_DIFF_SIZE ||
Buffer.byteLength(item.after) > Snapshot.MAX_DIFF_SIZE
if (file === item.file && !oversized) return item
return {
...item,
file,
before: oversized ? "" : item.before,
after: oversized ? "" : item.after,
}
})
const changed = next.some((item, i) => item.file !== diffs[i]?.file)
const changed = next.some((item, i) => item !== diffs[i])
if (changed) Storage.write(["session_diff", input.sessionID], next).catch(() => {})
// kilocode_change end
return next
},
)
+17 -6
View File
@@ -14,6 +14,7 @@ export namespace Snapshot {
const log = Log.create({ service: "snapshot" })
const hour = 60 * 60 * 1000
const prune = "7.days"
export const MAX_DIFF_SIZE = 256 * 1024 // kilocode_change
export function init() {
Scheduler.register({
@@ -34,10 +35,11 @@ export namespace Snapshot {
.then(() => true)
.catch(() => false)
if (!exists) return
const result = await $`git -c core.autocrlf=false -c core.longpaths=true -c core.symlinks=true --git-dir ${git} --work-tree ${Instance.worktree} gc --prune=${prune}`
.quiet()
.cwd(Instance.directory)
.nothrow()
const result =
await $`git -c core.autocrlf=false -c core.longpaths=true -c core.symlinks=true --git-dir ${git} --work-tree ${Instance.worktree} gc --prune=${prune}`
.quiet()
.cwd(Instance.directory)
.nothrow()
if (result.exitCode !== 0) {
log.warn("cleanup failed", {
exitCode: result.exitCode,
@@ -228,13 +230,22 @@ export namespace Snapshot {
if (!line) continue
const [additions, deletions, file] = line.split("\t")
const isBinaryFile = additions === "-" && deletions === "-"
const before = isBinaryFile
// kilocode_change start
const oversized =
!isBinaryFile &&
((parseInt(await $`git --git-dir ${git} cat-file -s ${from}:${file}`.quiet().nothrow().text()) || 0) >
MAX_DIFF_SIZE ||
(parseInt(await $`git --git-dir ${git} cat-file -s ${to}:${file}`.quiet().nothrow().text()) || 0) >
MAX_DIFF_SIZE)
const skip = isBinaryFile || oversized
// kilocode_change end
const before = skip
? ""
: await $`git -c core.autocrlf=false -c core.longpaths=true -c core.symlinks=true --git-dir ${git} --work-tree ${Instance.worktree} show ${from}:${file}`
.quiet()
.nothrow()
.text()
const after = isBinaryFile
const after = skip
? ""
: await $`git -c core.autocrlf=false -c core.longpaths=true -c core.symlinks=true --git-dir ${git} --work-tree ${Instance.worktree} show ${to}:${file}`
.quiet()
+1 -1
View File
@@ -61,7 +61,7 @@ export const TaskTool = Tool.define("task", async (ctx) => {
const agent = await Agent.get(params.subagent_type)
if (!agent) throw new Error(`Unknown agent type: ${params.subagent_type} is not a valid agent type`)
const allowsTask = agent.permission.some((rule) => rule.permission === "task" && rule.action === "allow")
const allowsTask = agent.permission.some((rule) => rule.permission === "task" && rule.action === "allow") // kilocode_change
const session = await iife(async () => {
if (params.task_id) {