opensource/kilocode
3
0
Fork 0
mirror of https://github.com/Kilo-Org/kilocode.git synced 2026-08-29 03:32:08 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(vscode): stabilize long session restores (#9444)

Browse Source 
* fix(vscode): stabilize long session restores

* fix(vscode): render partial long-session turns

* fix(vscode): split partial session turns
This commit is contained in:
Marius
2026-04-24 10:33:16 +02:00
committed by GitHub
parent b457920301
commit 35f13cbc86
7 changed files with 206 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.changeset/fix-long-session-loading.md
+5
View File
@@ -0,0 +1,5 @@
---
"kilo-code": patch
---
Fix restoring and paginating very long VS Code sessions.
packages/kilo-vscode/src/kilo-provider/message-page.ts
+6 -2
View File
@@ -3,6 +3,9 @@ import { retry } from "../services/cli-backend/retry"
export const MESSAGE_PAGE_LIMIT = 80
// Bound assistant-boundary backfill so corrupt histories cannot load an entire session.
const FILL_LIMIT = 2
/**
* Build the same base64url-encoded cursor format the server emits so a
* synthesized cursor round-trips through `session.messages({ before })`.
@@ -46,12 +49,13 @@ export async function fetchMessagePage(
return { items, cursor }
}
const fill = async (page: Awaited<ReturnType<typeof read>>): Promise<Awaited<ReturnType<typeof read>>> => {
const fill = async (page: Awaited<ReturnType<typeof read>>, depth = 0): Promise<Awaited<ReturnType<typeof read>>> => {
if (page.items[0]?.info.role !== "assistant") return page
if (depth >= FILL_LIMIT) return page
if (!page.cursor || input.signal?.aborted) return page
const next = await read(page.cursor)
const items = [...next.items, ...page.items]
return fill({ items, cursor: next.cursor })
return fill({ items, cursor: next.cursor }, depth + 1)
}
return fill(await read(input.before))
packages/kilo-vscode/tests/unit/message-page.test.ts
+30
View File
@@ -164,4 +164,34 @@ describe("fetchMessagePage / cursor fallback", () => {
expect(page.items.map((item) => item.info.id)).toEqual(["m1", "m2", "m3", "m4", "m5"])
expect(page.cursor).toBeUndefined()
})
it("bounds assistant turn filling when older pages never reach a user message", async () => {
const { client, calls } = mockClient([
{
items: [message("m5", "assistant", 50), message("m6", "assistant", 60)],
cursor: "c1",
},
{
items: [message("m3", "assistant", 30), message("m4", "assistant", 40)],
cursor: "c2",
},
{
items: [message("m1", "assistant", 10), message("m2", "assistant", 20)],
cursor: "c3",
},
{
items: [message("m0", "user", 0)],
},
])
const page = await fetchMessagePage(client as never, {
sessionID: "s1",
workspaceDir: "/repo",
limit: 2,
})
expect(calls.map((call) => call.before)).toEqual([undefined, "c1", "c2"])
expect(page.items.map((item) => item.info.id)).toEqual(["m1", "m2", "m3", "m4", "m5", "m6"])
expect(page.cursor).toBe("c3")
})
})
packages/kilo-vscode/tests/unit/session-queue.test.ts
+56 -1
View File
@@ -1,5 +1,10 @@
import { describe, expect, it } from "bun:test"
import { activeUserMessageID, messageTurns, queuedUserMessageIDs } from "../../webview-ui/src/context/session-queue"
import {
activeUserMessageID,
messageTurns,
queuedUserMessageIDs,
stableMessageTurns,
} from "../../webview-ui/src/context/session-queue"
import type { Message } from "../../webview-ui/src/types/messages"
const base = {
@@ -84,6 +89,56 @@ describe("messageTurns", () => {
{ user: "message_4", assistant: [] },
])
})
it("surfaces leading assistant output as partial turns grouped by parent", () => {
const messages = [
assistant("message_2", "message_1"),
assistant("message_4", "message_3"),
assistant("message_5", "message_3"),
user("message_6"),
]
const turns = messageTurns(messages)
expect(
turns.map((turn) => ({ id: turn.id, partial: turn.partial, assistant: turn.assistant.map((msg) => msg.id) })),
).toEqual([
{ id: "message_1", partial: true, assistant: ["message_2"] },
{ id: "message_3", partial: true, assistant: ["message_4", "message_5"] },
{ id: "message_6", partial: undefined, assistant: [] },
])
})
})
describe("stableMessageTurns", () => {
it("keeps existing turn identities stable when older turns are prepended", () => {
const u1 = user("message_1")
const a2 = assistant("message_2", "message_1")
const u3 = user("message_3")
const prev = messageTurns([u1, a2, u3])
const next = stableMessageTurns(messageTurns([user("message_0"), u1, a2, u3]), prev)
expect(next[1]).toBe(prev[0])
expect(next[2]).toBe(prev[1])
})
it("replaces a turn identity when its assistant messages change", () => {
const u1 = user("message_1")
const a2 = assistant("message_2", "message_1")
const prev = messageTurns([u1, a2])
const next = stableMessageTurns(messageTurns([u1, a2, assistant("message_3", "message_1")]), prev)
expect(next[0]).not.toBe(prev[0])
expect(next[0]?.assistant.map((msg) => msg.id)).toEqual(["message_2", "message_3"])
})
it("keeps partial turn identities stable while their assistant messages are unchanged", () => {
const a2 = assistant("message_2", "message_1")
const a3 = assistant("message_3", "message_1")
const prev = messageTurns([a2, a3])
const next = stableMessageTurns(messageTurns([a2, a3, user("message_4")]), prev)
expect(next[0]).toBe(prev[0])
})
})
describe("activeUserMessageID", () => {
packages/kilo-vscode/webview-ui/src/components/chat/MessageList.tsx
+6 -2
View File
@@ -29,6 +29,8 @@ import {
activeUserMessageID as getActiveUserMessageID,
messageTurns,
queuedUserMessageIDs,
stableMessageTurns,
type MessageTurn,
} from "../../context/session-queue"
import type { QuestionRequest, SuggestionRequest } from "../../types/messages"
@@ -84,7 +86,9 @@ export const MessageList: Component<MessageListProps> = (props) => {
const positions = new Map<string, { top: number; userScrolled: boolean }>()
const boundary = () => session.revert()?.messageID
const turns = createMemo(() => messageTurns(session.messages(), boundary()))
const turns = createMemo((prev: MessageTurn[] | undefined) =>
stableMessageTurns(messageTurns(session.messages(), boundary()), prev),
)
const isEmpty = () => turns().length === 0 && !session.loading() && !boundary()
const recent = createMemo(() =>
@@ -151,11 +155,11 @@ export const MessageList: Component<MessageListProps> = (props) => {
if (pos?.userScrolled) {
el.scrollTop = pos.top
autoScroll.pause()
maybeLoadOlder()
} else {
autoScroll.forceScrollToBottom()
}
setPendingRestore(undefined)
maybeLoadOlder()
})
})
})
packages/kilo-vscode/webview-ui/src/components/chat/VscodeSessionTurn.tsx
+30 -26
View File
@@ -50,6 +50,7 @@ export interface VscodeTurn {
id: string
user: WebMessage
assistant: WebMessage[]
partial?: boolean
}
interface VscodeSessionTurnProps {
@@ -71,7 +72,8 @@ export const VscodeSessionTurn: Component<VscodeSessionTurnProps> = (props) => {
createEffect(() => {
const turn = props.turn
session.hydrateParts([turn.user.id, ...turn.assistant.map((m) => m.id)])
const ids = turn.partial ? turn.assistant.map((m) => m.id) : [turn.user.id, ...turn.assistant.map((m) => m.id)]
session.hydrateParts(ids)
})
const message = createMemo(() => props.turn.user as SDKMessage & { role: "user" })
@@ -138,33 +140,35 @@ export const VscodeSessionTurn: Component<VscodeSessionTurnProps> = (props) => {
{(msg) => (
<div class="vscode-session-turn" data-message={msg().id}>
{/* User message */}
<div
class="vscode-session-turn-user"
data-revert-disabled={
assistantMessages().length > 0 && !session.revert() && session.status() !== "idle" ? "" : undefined
}
title={
assistantMessages().length > 0 && !session.revert() && session.status() !== "idle"
? language.t("revert.disabled.agentBusy")
: undefined
}
>
<UserMessageDisplay
message={msg() as unknown as Parameters<typeof UserMessageDisplay>[0]["message"]}
parts={parts() as unknown as Parameters<typeof UserMessageDisplay>[0]["parts"]}
interrupted={interrupted()}
queued={props.queued}
onFork={props.onForkMessage ? () => props.onForkMessage?.(msg().sessionID, msg().id) : undefined}
onRevert={
assistantMessages().length > 0 && !session.revert()
? () => {
if (session.status() !== "idle") return
session.revertSession(msg().id)
}
<Show when={!props.turn.partial}>
<div
class="vscode-session-turn-user"
data-revert-disabled={
assistantMessages().length > 0 && !session.revert() && session.status() !== "idle" ? "" : undefined
}
title={
assistantMessages().length > 0 && !session.revert() && session.status() !== "idle"
? language.t("revert.disabled.agentBusy")
: undefined
}
/>
</div>
>
<UserMessageDisplay
message={msg() as unknown as Parameters<typeof UserMessageDisplay>[0]["message"]}
parts={parts() as unknown as Parameters<typeof UserMessageDisplay>[0]["parts"]}
interrupted={interrupted()}
queued={props.queued}
onFork={props.onForkMessage ? () => props.onForkMessage?.(msg().sessionID, msg().id) : undefined}
onRevert={
assistantMessages().length > 0 && !session.revert()
? () => {
if (session.status() !== "idle") return
session.revertSession(msg().id)
}
: undefined
}
/>
</div>
</Show>
{/* Assistant parts — flat list, no context grouping */}
<Show when={assistantMessages().length > 0}>
packages/kilo-vscode/webview-ui/src/context/session-queue.ts
+73 -4
View File
@@ -4,10 +4,47 @@ export interface MessageTurn {
id: string
user: Message
assistant: Message[]
partial?: boolean
}
export function messageTurns(messages: Message[], boundary?: string) {
function key(msg: Message) {
return msg.parentID ?? msg.id
}
function partial(messages: Message[]): MessageTurn {
const first = messages[0]!
const id = first.parentID ?? `${first.id}:partial`
return {
id,
user: {
id,
sessionID: first.sessionID,
role: "user",
createdAt: first.createdAt,
time: first.time,
},
assistant: messages,
partial: true,
}
}
function partials(messages: Message[]): MessageTurn[] {
return messages
.reduce<Message[][]>((groups, msg) => {
const prev = groups[groups.length - 1]
if (!prev || key(prev[0]!) !== key(msg)) {
groups.push([msg])
return groups
}
prev.push(msg)
return groups
}, [])
.map(partial)
}
export function messageTurns(messages: Message[], boundary?: string): MessageTurn[] {
const result: MessageTurn[] = []
const lead: Message[] = []
const by = new Map<string, MessageTurn>()
for (const msg of messages) {
@@ -20,11 +57,43 @@ export function messageTurns(messages: Message[], boundary?: string) {
}
if (msg.role !== "assistant") continue
const turn = (msg.parentID ? by.get(msg.parentID) : undefined) ?? result[result.length - 1]
if (turn) turn.assistant.push(msg)
const turn = msg.parentID ? by.get(msg.parentID) : undefined
if (turn) {
turn.assistant.push(msg)
continue
}
const last = result[result.length - 1]
if (last) {
last.assistant.push(msg)
continue
}
lead.push(msg)
}
return result
if (lead.length === 0) return result
return [...partials(lead), ...result]
}
function sameMessages(a: Message[], b: Message[]) {
if (a.length !== b.length) return false
for (let i = 0; i < a.length; i++) {
if (a[i] !== b[i]) return false
}
return true
}
// Keep virtua's item keys stable across prepends; Solid's adapter keys by data object identity.
export function stableMessageTurns(next: MessageTurn[], prev: MessageTurn[] = []): MessageTurn[] {
if (prev.length === 0) return next
const by = new Map(prev.map((turn) => [turn.user.id, turn]))
return next.map((turn) => {
const old = by.get(turn.user.id)
if (!old) return turn
if (old.partial !== turn.partial) return turn
if (!turn.partial && old.user !== turn.user) return turn
if (!sameMessages(old.assistant, turn.assistant)) return turn
return old
})
}
function active(messages: Message[]) {