diff --git a/.changeset/fresh-compaction-order.md b/.changeset/fresh-compaction-order.md new file mode 100644 index 00000000000..7d863475e6b --- /dev/null +++ b/.changeset/fresh-compaction-order.md @@ -0,0 +1,6 @@ +--- +"@kilocode/cli": patch +"kilo-code": patch +--- + +Keep post-compaction tool calls and follow-up messages ordered after the compaction summary in the CLI and VS Code transcript. diff --git a/packages/kilo-vscode/tests/unit/session-queue.test.ts b/packages/kilo-vscode/tests/unit/session-queue.test.ts index 40a9b282f85..8e8f6934c3a 100644 --- a/packages/kilo-vscode/tests/unit/session-queue.test.ts +++ b/packages/kilo-vscode/tests/unit/session-queue.test.ts @@ -17,6 +17,11 @@ const base = { const user = (id: string): Message => ({ ...base, id, role: "user" }) +const compact = (id: string): Message => ({ + ...user(id), + parts: [{ id: `part_${id}`, sessionID: base.sessionID, messageID: id, type: "compaction", auto: false }], +}) + const assistant = (id: string, parentID: string, opts: Partial = {}): Message => ({ ...base, id, @@ -316,6 +321,42 @@ describe("messageTurns", () => { ]) }) + it("keeps resumed replies after a persisted compaction turn", () => { + const messages = [ + user("message_1"), + assistant("message_2", "message_1"), + compact("message_3"), + assistant("message_4", "message_3", { summary: true, finish: "stop" }), + assistant("message_5", "message_1", { finish: "stop" }), + ] + + expect( + messageTurns(messages).map((turn) => ({ + user: turn.user.id, + assistant: turn.assistant.map((msg) => msg.id), + })), + ).toEqual([ + { user: "message_1", assistant: ["message_2"] }, + { user: "message_3", assistant: ["message_4", "message_5"] }, + ]) + }) + + it("detects persisted compaction parts through the lazy lookup", () => { + const messages = [ + user("message_1"), + assistant("message_2", "message_1"), + user("message_3"), + assistant("message_4", "message_3", { summary: true, finish: "stop" }), + assistant("message_5", "message_1", { finish: "stop" }), + ] + + expect( + visibleMessages(messages, undefined, (msg) => (msg.id === "message_3" ? compact(msg.id).parts : msg.parts)).map( + (msg) => msg.id, + ), + ).toEqual(["message_1", "message_2", "message_3", "message_4", "message_5"]) + }) + it("surfaces leading assistant output as partial turns grouped by parent", () => { const messages = [ assistant("message_2", "message_1"), @@ -437,6 +478,49 @@ describe("activeUserMessageID", () => { expect(activeUserMessageID(messages, { type: "busy" })).toBe("message_1") }) + it("maps resumed post-compaction tool calls to the compaction turn", () => { + const messages = [ + user("message_1"), + assistant("message_2", "message_1"), + compact("message_3"), + assistant("message_4", "message_3", { summary: true, finish: "stop" }), + assistant("message_5", "message_1", { finish: "tool-calls" }), + user("message_6"), + ] + + expect(activeUserMessageID(messages, { type: "busy" })).toBe("message_3") + expectLayout(messages, { type: "busy" }, { virtual: ["message_1"], direct: ["message_3"], queued: ["message_6"] }) + }) + + it("uses lazy compaction parts when mapping resumed tool calls", () => { + const messages = [ + user("message_1"), + assistant("message_2", "message_1"), + user("message_3"), + assistant("message_4", "message_3", { summary: true, finish: "stop" }), + assistant("message_5", "message_1", { finish: "tool-calls" }), + ] + + expect( + activeUserMessageID(messages, { type: "busy" }, (msg) => + msg.id === "message_3" ? compact(msg.id).parts : msg.parts, + ), + ).toBe("message_3") + }) + + it("advances beyond a completed post-compaction reply", () => { + const messages = [ + user("message_1"), + assistant("message_2", "message_1", { finish: "stop" }), + compact("message_3"), + assistant("message_4", "message_3", { summary: true, finish: "stop" }), + assistant("message_5", "message_1", { finish: "stop" }), + user("message_6"), + ] + + expect(activeUserMessageID(messages, { type: "busy" })).toBe("message_6") + }) + it("ignores completed tool-call assistants after the session becomes idle", () => { const messages = [ user("message_1"), diff --git a/packages/kilo-vscode/webview-ui/src/components/chat/MessageList.tsx b/packages/kilo-vscode/webview-ui/src/components/chat/MessageList.tsx index 702bd3a680e..7b2f1241592 100644 --- a/packages/kilo-vscode/webview-ui/src/components/chat/MessageList.tsx +++ b/packages/kilo-vscode/webview-ui/src/components/chat/MessageList.tsx @@ -92,14 +92,21 @@ export const MessageList: Component = (props) => { const boundary = () => session.revert()?.messageID const turns = createMemo((prev: MessageTurn[] | undefined) => - stableMessageTurns(messageTurns(session.messages(), boundary()), prev), + stableMessageTurns( + messageTurns(session.messages(), boundary(), (msg) => session.getParts(msg.id)), + prev, + ), ) const isEmpty = () => turns().length === 0 && !session.loading() && !boundary() const recent = createMemo(() => recentSessions(session.sessions())) - const activeUserID = createMemo(() => getActiveUserMessageID(session.messages(), session.statusInfo())) - const queuedIDs = createMemo(() => new Set(queuedUserMessageIDs(session.messages(), session.statusInfo()))) + const activeUserID = createMemo(() => + getActiveUserMessageID(session.messages(), session.statusInfo(), (msg) => session.getParts(msg.id)), + ) + const queuedIDs = createMemo( + () => new Set(queuedUserMessageIDs(session.messages(), session.statusInfo(), (msg) => session.getParts(msg.id))), + ) const [held, setHeld] = createSignal<{ sid: string; ids: Set }>() createEffect(() => { const id = activeUserID() diff --git a/packages/kilo-vscode/webview-ui/src/context/session-queue.ts b/packages/kilo-vscode/webview-ui/src/context/session-queue.ts index 18258bacb91..066df583295 100644 --- a/packages/kilo-vscode/webview-ui/src/context/session-queue.ts +++ b/packages/kilo-vscode/webview-ui/src/context/session-queue.ts @@ -42,27 +42,52 @@ function partials(messages: Message[]): MessageTurn[] { .map(partial) } -export function messageTurns(messages: Message[], boundary?: string): MessageTurn[] { +function isCompact(msg: Message, parts?: (msg: Message) => Message["parts"]) { + return msg.role === "user" && (parts?.(msg) ?? msg.parts)?.some((part) => part.type === "compaction") +} + +function target(messages: Message[], index: number, id: string, parts?: (msg: Message) => Message["parts"]) { + const parent = messages.findIndex((msg) => msg.id === id) + for (let i = index - 1; i > parent; i -= 1) { + const msg = messages[i] + if (msg && isCompact(msg, parts)) return msg.id + } + return id +} + +export function messageTurns( + messages: Message[], + boundary?: string, + parts?: (msg: Message) => Message["parts"], +): MessageTurn[] { const result: MessageTurn[] = [] const lead: Message[] = [] - const by = new Map() + const by = new Map() + let compact: { turn: MessageTurn; index: number } | undefined for (const msg of messages) { if (msg.role === "user") { if (boundary && msg.id >= boundary) break const turn = { id: msg.id, user: msg, assistant: [] } + const item = { turn, index: result.length } result.push(turn) - by.set(msg.id, turn) + by.set(msg.id, item) + if (isCompact(msg, parts)) compact = item continue } if (msg.role !== "assistant") continue - const turn = msg.parentID ? by.get(msg.parentID) : undefined - if (turn) { + const parent = msg.parentID ? by.get(msg.parentID) : undefined + if (parent) { + const turn = compact && parent.index < compact.index ? compact.turn : parent.turn turn.assistant.push(msg) continue } if (msg.parentID) { + if (compact) { + compact.turn.assistant.push(msg) + continue + } lead.push(msg) continue } @@ -78,8 +103,12 @@ export function messageTurns(messages: Message[], boundary?: string): MessageTur return [...partials(lead), ...result] } -export function visibleMessages(messages: Message[], boundary?: string): Message[] { - return messageTurns(messages, boundary).flatMap((turn) => +export function visibleMessages( + messages: Message[], + boundary?: string, + parts?: (msg: Message) => Message["parts"], +): Message[] { + return messageTurns(messages, boundary, parts).flatMap((turn) => turn.partial ? turn.assistant : [turn.user, ...turn.assistant], ) } @@ -106,7 +135,7 @@ export function stableMessageTurns(next: MessageTurn[], prev: MessageTurn[] = [] }) } -function active(messages: Message[], status: SessionStatusInfo) { +function active(messages: Message[], status: SessionStatusInfo, parts?: (msg: Message) => Message["parts"]) { let latest = true for (let i = messages.length - 1; i >= 0; i -= 1) { const msg = messages[i] @@ -118,8 +147,9 @@ function active(messages: Message[], status: SessionStatusInfo) { if (msg.error) continue if (msg.finish && !resumable) continue if (!msg.parentID) break - const parent = messages.find((item) => item.id === msg.parentID) - if (!parent) return msg.parentID + const id = target(messages, i, msg.parentID, parts) + const parent = messages.find((item) => item.id === id) + if (!parent) return id if (parent.role === "user") return parent.id break } @@ -127,20 +157,20 @@ function active(messages: Message[], status: SessionStatusInfo) { return undefined } -function done(messages: Message[]) { +function done(messages: Message[], parts?: (msg: Message) => Message["parts"]) { for (let i = messages.length - 1; i >= 0; i -= 1) { const msg = messages[i] - if (!msg || msg.role !== "assistant") continue - if (typeof msg.time?.completed === "number") return msg.parentID - if (msg.error) return msg.parentID - if (msg.finish && !["tool-calls", "unknown"].includes(msg.finish)) return msg.parentID + if (!msg || msg.role !== "assistant" || !msg.parentID) continue + if (typeof msg.time?.completed === "number") return target(messages, i, msg.parentID, parts) + if (msg.error) return target(messages, i, msg.parentID, parts) + if (msg.finish && !["tool-calls", "unknown"].includes(msg.finish)) return target(messages, i, msg.parentID, parts) } return undefined } -function pending(messages: Message[]) { +function pending(messages: Message[], parts?: (msg: Message) => Message["parts"]) { const users = messages.filter((msg) => msg.role === "user") - const id = done(messages) + const id = done(messages, parts) if (!id) return users[0]?.id const idx = users.findIndex((msg) => msg.id === id) @@ -149,23 +179,31 @@ function pending(messages: Message[]) { // Find the user message whose turn the server is actively processing. // Any user message after this one is "queued" (waiting for its turn). -export function activeUserMessageID(messages: Message[], status: SessionStatusInfo) { - const id = active(messages, status) +export function activeUserMessageID( + messages: Message[], + status: SessionStatusInfo, + parts?: (msg: Message) => Message["parts"], +) { + const id = active(messages, status, parts) if (id) return id if (status.type === "idle") return undefined - return pending(messages) + return pending(messages, parts) } -export function queuedUserMessageIDs(messages: Message[], status: SessionStatusInfo) { +export function queuedUserMessageIDs( + messages: Message[], + status: SessionStatusInfo, + parts?: (msg: Message) => Message["parts"], +) { if (status.type === "idle") return [] const users = messages.filter((msg) => msg.role === "user") - const running = active(messages, status) + const running = active(messages, status, parts) if (running) { const idx = users.findIndex((msg) => msg.id === running) if (idx < 0) return users.map((msg) => msg.id) return users.slice(idx + 1).map((msg) => msg.id) } - const id = pending(messages) + const id = pending(messages, parts) const idx = id ? users.findIndex((msg) => msg.id === id) : -1 if (idx < 0) return [] return users.slice(idx + 1).map((msg) => msg.id) diff --git a/packages/kilo-vscode/webview-ui/src/context/session.tsx b/packages/kilo-vscode/webview-ui/src/context/session.tsx index d93c05ffce5..0f564a07d01 100644 --- a/packages/kilo-vscode/webview-ui/src/context/session.tsx +++ b/packages/kilo-vscode/webview-ui/src/context/session.tsx @@ -2304,7 +2304,9 @@ export const SessionProvider: ParentComponent = (props) => { const userMessages = createMemo(() => messages().filter((m) => m.role === "user")) function visible(sessionID: string) { - return filterVisibleMessages(store.messages[sessionID] ?? [], store.sessions[sessionID]?.revert?.messageID) + return filterVisibleMessages(store.messages[sessionID] ?? [], store.sessions[sessionID]?.revert?.messageID, (msg) => + getParts(msg.id), + ) } const revert = createMemo(() => { diff --git a/packages/kilo-vscode/webview-ui/src/types/messages/parts.ts b/packages/kilo-vscode/webview-ui/src/types/messages/parts.ts index cdfb9efd3e6..336458f2cfc 100644 --- a/packages/kilo-vscode/webview-ui/src/types/messages/parts.ts +++ b/packages/kilo-vscode/webview-ui/src/types/messages/parts.ts @@ -73,7 +73,14 @@ export interface StepFinishPart extends BasePart { } } -export type Part = TextPart | FilePart | ToolPart | ReasoningPart | StepStartPart | StepFinishPart +export interface CompactionPart extends BasePart { + type: "compaction" + auto: boolean + overflow?: boolean + tail_start_id?: string +} + +export type Part = TextPart | FilePart | ToolPart | ReasoningPart | StepStartPart | StepFinishPart | CompactionPart // Part delta for streaming updates export interface PartDelta { diff --git a/packages/opencode/src/kilocode/session/message-order.ts b/packages/opencode/src/kilocode/session/message-order.ts new file mode 100644 index 00000000000..0d255eb22ce --- /dev/null +++ b/packages/opencode/src/kilocode/session/message-order.ts @@ -0,0 +1,66 @@ +import type { MessageV2 } from "@/session/message-v2" + +const chronology = new WeakMap() + +export namespace KiloSessionMessageOrder { + /** Preserve chronological order before model-facing projections rearrange messages. */ + export function annotate(msgs: MessageV2.WithParts[]) { + for (const [index, msg] of msgs.entries()) chronology.set(msg, index) + return msgs + } + + export function compare(a: MessageV2.WithParts, b: MessageV2.WithParts, indexA = -1, indexB = -1) { + if (a.info.time.created !== b.info.time.created) return a.info.time.created - b.info.time.created + const sequenceA = chronology.get(a) + const sequenceB = chronology.get(b) + if (sequenceA !== undefined && sequenceB !== undefined && sequenceA !== sequenceB) return sequenceA - sequenceB + return indexA - indexB + } + + /** Derive active messages by chronology while keeping queued tasks in model-facing projection order. */ + export function latest(msgs: MessageV2.WithParts[]) { + let user: MessageV2.WithParts | undefined + let assistant: MessageV2.WithParts | undefined + let finished: MessageV2.WithParts | undefined + let userIndex = -1 + let assistantIndex = -1 + let finishedIndex = -1 + + for (const [index, msg] of msgs.entries()) { + const info = msg.info + if (info.role === "user" && (!user || compare(msg, user, index, userIndex) > 0)) { + user = msg + userIndex = index + } + if (info.role === "assistant" && (!assistant || compare(msg, assistant, index, assistantIndex) > 0)) { + assistant = msg + assistantIndex = index + } + if (info.role === "assistant" && info.finish && (!finished || compare(msg, finished, index, finishedIndex) > 0)) { + finished = msg + finishedIndex = index + } + } + + const pivot = msgs.findLastIndex((msg) => msg.info.role === "assistant" && msg.info.finish) + const tasks = msgs + .slice(pivot + 1) + .reverse() + .flatMap((msg) => + msg.parts.filter( + (part): part is MessageV2.CompactionPart | MessageV2.SubtaskPart => + part.type === "compaction" || part.type === "subtask", + ), + ) + + return { + user: user?.info.role === "user" ? user.info : undefined, + assistant: assistant?.info.role === "assistant" ? assistant.info : undefined, + finished: finished?.info.role === "assistant" ? finished.info : undefined, + userMessage: user, + assistantMessage: assistant, + finishedMessage: finished, + tasks, + } + } +} diff --git a/packages/opencode/src/kilocode/session/prompt.ts b/packages/opencode/src/kilocode/session/prompt.ts index c151a8f7f45..7ff0b20d726 100644 --- a/packages/opencode/src/kilocode/session/prompt.ts +++ b/packages/opencode/src/kilocode/session/prompt.ts @@ -12,6 +12,7 @@ import type { SessionStatus } from "@/session/status" import { Flag } from "@opencode-ai/core/flag/flag" import { PlanFollowup } from "@/kilocode/plan-followup" import { KiloSession } from "@/kilocode/session" +import { KiloSessionMessageOrder } from "@/kilocode/session/message-order" import { Permission } from "@/permission" import { environmentDetails, type EditorContext } from "@/kilocode/editor-context" import { Identifier } from "@/id/id" @@ -342,14 +343,18 @@ export namespace KiloSessionPrompt { * `msgs`, `msgs` is returned unchanged. */ export function trimBeforeLastSummary(msgs: MessageV2.WithParts[]): MessageV2.WithParts[] { - for (let i = msgs.length - 1; i >= 0; i--) { - const info = msgs[i].info - if (info.role !== "assistant" || info.summary !== true || !info.finish || info.error) continue - const parentIdx = msgs.findIndex((m) => m.info.id === info.parentID) - if (parentIdx === -1) return msgs - return parentIdx === 0 ? msgs : msgs.slice(parentIdx) - } - return msgs + const summary = msgs.reduce<{ msg: MessageV2.WithParts; index: number } | undefined>((latest, msg, index) => { + const info = msg.info + if (info.role !== "assistant" || info.summary !== true || !info.finish || info.error) return latest + if (!latest || KiloSessionMessageOrder.compare(msg, latest.msg, index, latest.index) > 0) return { msg, index } + return latest + }, undefined) + if (!summary) return msgs + const info = summary.msg.info + if (info.role !== "assistant") return msgs + const parentIdx = msgs.findIndex((m) => m.info.id === info.parentID) + if (parentIdx === -1) return msgs + return parentIdx === 0 ? msgs : msgs.slice(parentIdx) } /** diff --git a/packages/opencode/src/session/message-v2.ts b/packages/opencode/src/session/message-v2.ts index 70abefb6909..38be168c2d3 100644 --- a/packages/opencode/src/session/message-v2.ts +++ b/packages/opencode/src/session/message-v2.ts @@ -24,6 +24,7 @@ import type { Provider } from "@/provider/provider" import { ModelID, ProviderID } from "@/provider/schema" import { SessionNetwork } from "./network" // kilocode_change import { CodexAuthExpiredError } from "@/kilocode/provider/codex-refresh" // kilocode_change +import { KiloSessionMessageOrder } from "@/kilocode/session/message-order" // kilocode_change import { Effect, Schema, Types } from "effect" import { zod, ZodOverride } from "@/util/effect-zod" import { NonNegativeInt, withStatics } from "@/util/schema" @@ -1218,6 +1219,7 @@ export function filterCompacted(msgs: Iterable) { completed.add(msg.info.parentID) } result.reverse() + KiloSessionMessageOrder.annotate(result) // kilocode_change - preserve chronology before retained-tail projection const compactionIndex = result.findLastIndex( (msg) => msg.info.role === "user" && diff --git a/packages/opencode/src/session/prompt.ts b/packages/opencode/src/session/prompt.ts index f1181bc9bdf..ecdd048bb41 100644 --- a/packages/opencode/src/session/prompt.ts +++ b/packages/opencode/src/session/prompt.ts @@ -2,6 +2,7 @@ import path from "path" import os from "os" import fs from "fs/promises" import { KiloSessionPrompt } from "@/kilocode/session/prompt" // kilocode_change +import { KiloSessionMessageOrder } from "@/kilocode/session/message-order" // kilocode_change import { KiloSessionPromptQueue } from "@/kilocode/session/prompt-queue" // kilocode_change import { KiloSession } from "@/kilocode/session" // kilocode_change import { KiloCostPropagation } from "@/kilocode/session/cost-propagation" // kilocode_change @@ -1496,19 +1497,9 @@ NOTE: At any point in time through this workflow you should feel free to ask the msgs = KiloSessionPromptQueue.scope(sessionID, msgs) // kilocode_change - hide later queued prompts msgs = KiloSessionPrompt.trimBeforeLastSummary(msgs) // kilocode_change - trim on any completed summary (e.g. manual /compact against a text user) - let lastUser: MessageV2.User | undefined - let lastAssistant: MessageV2.Assistant | undefined - let lastFinished: MessageV2.Assistant | undefined - let tasks: (MessageV2.CompactionPart | MessageV2.SubtaskPart)[] = [] - for (let i = msgs.length - 1; i >= 0; i--) { - const msg = msgs[i] - if (!lastUser && msg.info.role === "user") lastUser = msg.info - if (!lastAssistant && msg.info.role === "assistant") lastAssistant = msg.info - if (!lastFinished && msg.info.role === "assistant" && msg.info.finish) lastFinished = msg.info - if (lastUser && lastFinished) break - const task = msg.parts.filter((part) => part.type === "compaction" || part.type === "subtask") - if (task && !lastFinished) tasks.push(...task) - } + // kilocode_change start - select loop state by chronology after retained-tail projection + const latest = KiloSessionMessageOrder.latest(msgs) + const { user: lastUser, assistant: lastAssistant, finished: lastFinished, tasks } = latest // kilocode_change end if (!lastUser) throw new Error("No user message found in stream. This should never happen.") @@ -1516,6 +1507,12 @@ NOTE: At any point in time through this workflow you should feel free to ask the const lastAssistantMsg = msgs.findLast( (msg) => msg.info.role === "assistant" && msg.info.id === lastAssistant?.id, ) + // kilocode_change start - compare chronology, not generated IDs + const userBeforeAssistant = + latest.userMessage && + latest.assistantMessage && + KiloSessionMessageOrder.compare(latest.userMessage, latest.assistantMessage) < 0 + // kilocode_change end // kilocode_change start - carry local review command marker into LLM telemetry const telemetry = KiloSessionProcessor.extractReviewTelemetry( @@ -1537,7 +1534,7 @@ NOTE: At any point in time through this workflow you should feel free to ask the lastAssistant?.finish && hasToolCalls && lastAssistant.parentID === lastUser.id && - lastUser.id < lastAssistant.id && + userBeforeAssistant && KiloSessionPrompt.shouldAskPlanFollowup({ messages: msgs, abort: AbortSignal.any([]) }) ) { const action = yield* Effect.promise((signal) => @@ -1554,7 +1551,7 @@ NOTE: At any point in time through this workflow you should feel free to ask the !["tool-calls"].includes(lastAssistant.finish) && !hasToolCalls && lastAssistant.parentID === lastUser.id && // kilocode_change - unrelated later assistants do not answer this turn - lastUser.id < lastAssistant.id + userBeforeAssistant // kilocode_change - compare chronology, not generated IDs ) { // kilocode_change start - ask follow-up when plan_exit tool was called const action = yield* Effect.promise((signal) => @@ -1690,7 +1687,11 @@ NOTE: At any point in time through this workflow you should feel free to ask the if (step > 1 && lastFinished) { for (const m of msgs) { - if (m.info.role !== "user" || m.info.id <= lastFinished.id) continue + // kilocode_change start - compare chronology, not generated IDs + const finishedBeforeMessage = + latest.finishedMessage && KiloSessionMessageOrder.compare(latest.finishedMessage, m) < 0 + if (m.info.role !== "user" || !finishedBeforeMessage) continue + // kilocode_change end for (const p of m.parts) { if (p.type !== "text" || p.ignored || p.synthetic) continue if (!p.text.trim()) continue diff --git a/packages/opencode/test/kilocode/session-compaction-safety.test.ts b/packages/opencode/test/kilocode/session-compaction-safety.test.ts index 8254e748a68..9ba3694cbb8 100644 --- a/packages/opencode/test/kilocode/session-compaction-safety.test.ts +++ b/packages/opencode/test/kilocode/session-compaction-safety.test.ts @@ -4,6 +4,7 @@ import { describe, expect, test } from "bun:test" import { KiloSessionPrompt } from "../../src/kilocode/session/prompt" +import { KiloSessionMessageOrder } from "../../src/kilocode/session/message-order" import { MessageV2 } from "../../src/session/message-v2" import { ModelID, ProviderID } from "../../src/provider/schema" import { MessageID, PartID, SessionID } from "../../src/session/schema" @@ -68,6 +69,29 @@ function syntheticTextPart(messageID: string, text: string, partID = "p_syn_" + } } +function compactionPart(messageID: string, tailStartID: string): MessageV2.CompactionPart { + return { + id: PartID.make("p_compact_" + messageID), + sessionID, + messageID: MessageID.make(messageID), + type: "compaction", + auto: false, + tail_start_id: MessageID.make(tailStartID), + } +} + +function subtaskPart(messageID: string): MessageV2.SubtaskPart { + return { + id: PartID.make("p_subtask_" + messageID), + sessionID, + messageID: MessageID.make(messageID), + type: "subtask", + prompt: "continue", + description: "Continue queued task", + agent: "test", + } +} + function filePart( messageID: string, mime: string, @@ -149,6 +173,11 @@ function assistant( return { info: assistantInfo(id, parentID, opts), parts } } +function created(msg: MessageV2.WithParts, time: number) { + msg.info.time.created = time + return msg +} + const apiError = new MessageV2.APIError({ message: "boom", isRetryable: true, @@ -178,6 +207,54 @@ describe("KiloSessionPrompt.hasCompletedSummary", () => { }) }) +describe("MessageV2.latest", () => { + test("selects chronological state after retained pre-compaction tail", () => { + const msgs = MessageV2.filterCompacted([ + created(assistant("msg_summary", "msg_compact", [], { summary: true, finish: "end_turn" }), 4), + created(user("msg_compact", [compactionPart("msg_compact", "msg_tail")]), 3), + created(assistant("msg_tail_reply", "msg_tail", [], { finish: "end_turn" }), 2), + created(user("msg_tail", [textPart("msg_tail", "historical retained prompt")]), 1), + ]) + + expect(msgs.map((m) => m.info.id)).toEqual([ + MessageID.make("msg_compact"), + MessageID.make("msg_summary"), + MessageID.make("msg_tail"), + MessageID.make("msg_tail_reply"), + ]) + + const state = KiloSessionMessageOrder.latest(msgs) + expect(state.user?.id).toBe(MessageID.make("msg_compact")) + expect(state.assistant?.id).toBe(MessageID.make("msg_summary")) + expect(state.finished?.id).toBe(MessageID.make("msg_summary")) + expect(state.tasks).toEqual([]) + }) + + test("keeps queued subtasks moved after a chronologically later assistant", () => { + const part = subtaskPart("msg_queued") + const active = created(user("msg_active"), 1) + const queued = created(user("msg_queued", [part]), 2) + const done = created(assistant("msg_done", "msg_active", [], { finish: "end_turn" }), 3) + KiloSessionMessageOrder.annotate([active, queued, done]) + + const state = KiloSessionMessageOrder.latest([active, done, queued]) + expect(state.user?.id).toBe(MessageID.make("msg_queued")) + expect(state.finished?.id).toBe(MessageID.make("msg_done")) + expect(state.tasks).toEqual([part]) + }) + + test("processes projected tasks in queue order", () => { + const first = subtaskPart("msg_first") + const second = subtaskPart("msg_second") + const msgs = [created(user("msg_first", [first]), 1), created(user("msg_second", [second]), 2)] + + const state = KiloSessionMessageOrder.latest(msgs) + expect(state.user?.id).toBe(MessageID.make("msg_second")) + expect(state.tasks).toEqual([second, first]) + expect(state.tasks.pop()).toBe(first) + }) +}) + describe("KiloSessionPrompt.trimBeforeLastSummary", () => { test("returns input unchanged when no summary present", () => { const msgs = [user("msg_u1"), assistant("msg_a1", "msg_u1", [], { finish: "end_turn" })] @@ -228,6 +305,31 @@ describe("KiloSessionPrompt.trimBeforeLastSummary", () => { ]) }) + test("keeps the newest summary when retained history contains an older summary", () => { + const filtered = MessageV2.filterCompacted([ + created(assistant("msg_s8", "msg_c7", [], { summary: true, finish: "end_turn" }), 8), + created(user("msg_c7", [compactionPart("msg_c7", "msg_u1")]), 7), + created(assistant("msg_a6", "msg_u5", [], { finish: "end_turn" }), 6), + created(user("msg_u5"), 5), + created(assistant("msg_s4", "msg_c3", [], { summary: true, finish: "end_turn" }), 4), + created(user("msg_c3", [compactionPart("msg_c3", "msg_u1")]), 3), + created(assistant("msg_a2", "msg_u1", [], { finish: "end_turn" }), 2), + created(user("msg_u1"), 1), + ]) + + expect(filtered.map((m) => m.info.id)).toEqual([ + MessageID.make("msg_c7"), + MessageID.make("msg_s8"), + MessageID.make("msg_u1"), + MessageID.make("msg_a2"), + MessageID.make("msg_c3"), + MessageID.make("msg_s4"), + MessageID.make("msg_u5"), + MessageID.make("msg_a6"), + ]) + expect(KiloSessionPrompt.trimBeforeLastSummary(filtered)).toBe(filtered) + }) + test("ignores errored and unfinished summaries when choosing boundary", () => { const msgs = [ user("msg_u1"),