fix(session): improve session compaction (#23870)

This commit is contained in:
Shoubhit Dash
2026-04-22 20:37:32 +05:30
committed by Aiden Cline
parent fa8b7bc4d2
commit 574b2c2170
6 changed files with 561 additions and 83 deletions

View File

@@ -143,6 +143,43 @@ async function assistant(sessionID: SessionID, parentID: MessageID, root: string
return msg
}
async function summaryAssistant(sessionID: SessionID, parentID: MessageID, root: string, text: string) {
const msg: MessageV2.Assistant = {
id: MessageID.ascending(),
role: "assistant",
sessionID,
mode: "compaction",
agent: "compaction",
path: { cwd: root, root },
cost: 0,
tokens: {
output: 0,
input: 0,
reasoning: 0,
cache: { read: 0, write: 0 },
},
modelID: ref.modelID,
providerID: ref.providerID,
parentID,
summary: true,
time: { created: Date.now() },
finish: "end_turn",
}
await svc.updateMessage(msg)
await svc.updatePart({
id: PartID.ascending(),
messageID: msg.id,
sessionID,
type: "text",
text,
})
return msg
}
async function lastCompactionPart(sessionID: SessionID) {
return (await svc.messages({ sessionID })).at(-2)?.parts.find((item): item is MessageV2.CompactionPart => item.type === "compaction")
}
function fake(
input: Parameters<SessionProcessorModule.SessionProcessor.Interface["create"]>[0],
result: "continue" | "compact",
@@ -946,12 +983,9 @@ describe("session.compaction.process", () => {
),
)
const part = (await svc.messages({ sessionID: session.id }))
.at(-2)
?.parts.find((item) => item.type === "compaction")
const part = await lastCompactionPart(session.id)
expect(part?.type).toBe("compaction")
if (part?.type === "compaction") expect(part.tail_start_id).toBe(keep.id)
expect(part?.tail_start_id).toBe(keep.id)
} finally {
await rt.dispose()
}
@@ -991,12 +1025,9 @@ describe("session.compaction.process", () => {
),
)
const part = (await svc.messages({ sessionID: session.id }))
.at(-2)
?.parts.find((item) => item.type === "compaction")
const part = await lastCompactionPart(session.id)
expect(part?.type).toBe("compaction")
if (part?.type === "compaction") expect(part.tail_start_id).toBe(keep.id)
expect(part?.tail_start_id).toBe(keep.id)
} finally {
await rt.dispose()
}
@@ -1042,12 +1073,9 @@ describe("session.compaction.process", () => {
),
)
const part = (await svc.messages({ sessionID: session.id }))
.at(-2)
?.parts.find((item) => item.type === "compaction")
const part = await lastCompactionPart(session.id)
expect(part?.type).toBe("compaction")
if (part?.type === "compaction") expect(part.tail_start_id).toBeUndefined()
expect(part?.tail_start_id).toBeUndefined()
expect(captured).toContain("yyyy")
} finally {
await rt.dispose()
@@ -1103,12 +1131,9 @@ describe("session.compaction.process", () => {
),
)
const part = (await svc.messages({ sessionID: session.id }))
.at(-2)
?.parts.find((item) => item.type === "compaction")
const part = await lastCompactionPart(session.id)
expect(part?.type).toBe("compaction")
if (part?.type === "compaction") expect(part.tail_start_id).toBeUndefined()
expect(part?.tail_start_id).toBeUndefined()
expect(captured).toContain("recent image turn")
expect(captured).toContain("Attached image/png: big.png")
} finally {
@@ -1118,6 +1143,76 @@ describe("session.compaction.process", () => {
})
})
test("retains a split turn suffix when a later message fits the preserve token budget", async () => {
await using tmp = await tmpdir({ git: true })
const stub = llm()
let captured = ""
stub.push(
reply("summary", (input) => {
captured = JSON.stringify(input.messages)
}),
)
await Instance.provide({
directory: tmp.path,
fn: async () => {
const session = await svc.create({})
await user(session.id, "older")
const recent = await user(session.id, "recent turn")
const large = await assistant(session.id, recent.id, tmp.path)
await svc.updatePart({
id: PartID.ascending(),
messageID: large.id,
sessionID: session.id,
type: "text",
text: "z".repeat(2_000),
})
const keep = await assistant(session.id, recent.id, tmp.path)
await svc.updatePart({
id: PartID.ascending(),
messageID: keep.id,
sessionID: session.id,
type: "text",
text: "keep tail",
})
await SessionCompaction.create({
sessionID: session.id,
agent: "build",
model: ref,
auto: false,
})
const rt = liveRuntime(stub.layer, wide(), cfg({ tail_turns: 1, preserve_recent_tokens: 100 }))
try {
const msgs = await svc.messages({ sessionID: session.id })
const parent = msgs.at(-1)?.info.id
expect(parent).toBeTruthy()
await rt.runPromise(
SessionCompaction.Service.use((svc) =>
svc.process({
parentID: parent!,
messages: msgs,
sessionID: session.id,
auto: false,
}),
),
)
const part = await lastCompactionPart(session.id)
expect(part?.type).toBe("compaction")
expect(part?.tail_start_id).toBe(keep.id)
expect(captured).toContain("zzzz")
expect(captured).not.toContain("keep tail")
const filtered = MessageV2.filterCompacted(MessageV2.stream(session.id))
expect(filtered[0]?.info.id).toBe(keep.id)
expect(filtered.map((msg) => msg.info.id)).not.toContain(large.id)
} finally {
await rt.dispose()
}
},
})
})
test("allows plugins to disable synthetic continue prompt", async () => {
await using tmp = await tmpdir()
await Instance.provide({
@@ -1530,6 +1625,80 @@ describe("session.compaction.process", () => {
})
})
test("anchors repeated compactions with the previous summary", async () => {
const stub = llm()
let captured = ""
stub.push(reply("summary one"))
stub.push(
reply("summary two", (input) => {
captured = JSON.stringify(input.messages)
}),
)
await using tmp = await tmpdir({ git: true })
await Instance.provide({
directory: tmp.path,
fn: async () => {
const session = await svc.create({})
await user(session.id, "older context")
await user(session.id, "keep this turn")
await SessionCompaction.create({
sessionID: session.id,
agent: "build",
model: ref,
auto: false,
})
const rt = liveRuntime(stub.layer, wide())
try {
let msgs = await svc.messages({ sessionID: session.id })
let parent = msgs.at(-1)?.info.id
expect(parent).toBeTruthy()
await rt.runPromise(
SessionCompaction.Service.use((svc) =>
svc.process({
parentID: parent!,
messages: msgs,
sessionID: session.id,
auto: false,
}),
),
)
await user(session.id, "latest turn")
await SessionCompaction.create({
sessionID: session.id,
agent: "build",
model: ref,
auto: false,
})
msgs = MessageV2.filterCompacted(MessageV2.stream(session.id))
parent = msgs.at(-1)?.info.id
expect(parent).toBeTruthy()
await rt.runPromise(
SessionCompaction.Service.use((svc) =>
svc.process({
parentID: parent!,
messages: msgs,
sessionID: session.id,
auto: false,
}),
),
)
expect(captured).toContain("<previous-summary>")
expect(captured).toContain("summary one")
expect(captured.match(/summary one/g)?.length).toBe(1)
expect(captured).toContain("## Constraints & Preferences")
expect(captured).toContain("## Progress")
} finally {
await rt.dispose()
}
},
})
})
test("keeps recent pre-compaction turns across repeated compactions", async () => {
const stub = llm()
stub.push(reply("summary one"))
@@ -1604,6 +1773,76 @@ describe("session.compaction.process", () => {
},
})
})
test("ignores previous summaries when sizing the retained tail", async () => {
await using tmp = await tmpdir()
await Instance.provide({
directory: tmp.path,
fn: async () => {
const session = await svc.create({})
await user(session.id, "older")
const keep = await user(session.id, "keep this turn")
const keepReply = await assistant(session.id, keep.id, tmp.path)
await svc.updatePart({
id: PartID.ascending(),
messageID: keepReply.id,
sessionID: session.id,
type: "text",
text: "keep reply",
})
await SessionCompaction.create({
sessionID: session.id,
agent: "build",
model: ref,
auto: false,
})
const firstCompaction = (await svc.messages({ sessionID: session.id })).at(-1)?.info.id
expect(firstCompaction).toBeTruthy()
await summaryAssistant(session.id, firstCompaction!, tmp.path, "summary ".repeat(800))
const recent = await user(session.id, "recent turn")
const recentReply = await assistant(session.id, recent.id, tmp.path)
await svc.updatePart({
id: PartID.ascending(),
messageID: recentReply.id,
sessionID: session.id,
type: "text",
text: "recent reply",
})
await SessionCompaction.create({
sessionID: session.id,
agent: "build",
model: ref,
auto: false,
})
const rt = runtime("continue", Plugin.defaultLayer, wide(), cfg({ tail_turns: 2, preserve_recent_tokens: 500 }))
try {
const msgs = await svc.messages({ sessionID: session.id })
const parent = msgs.at(-1)?.info.id
expect(parent).toBeTruthy()
await rt.runPromise(
SessionCompaction.Service.use((svc) =>
svc.process({
parentID: parent!,
messages: msgs,
sessionID: session.id,
auto: false,
}),
),
)
const part = await lastCompactionPart(session.id)
expect(part?.type).toBe("compaction")
expect(part?.tail_start_id).toBe(keep.id)
} finally {
await rt.dispose()
}
},
})
})
})
describe("util.token.estimate", () => {

View File

@@ -585,6 +585,76 @@ describe("session.message-v2.toModelMessage", () => {
])
})
test("truncates tool output when requested", async () => {
const userID = "m-user"
const assistantID = "m-assistant"
const input: MessageV2.WithParts[] = [
{
info: userInfo(userID),
parts: [
{
...basePart(userID, "u1"),
type: "text",
text: "run tool",
},
] as MessageV2.Part[],
},
{
info: assistantInfo(assistantID, userID),
parts: [
{
...basePart(assistantID, "a1"),
type: "tool",
callID: "call-1",
tool: "bash",
state: {
status: "completed",
input: { cmd: "ls" },
output: "abcdefghij",
title: "Bash",
metadata: {},
time: { start: 0, end: 1 },
},
},
] as MessageV2.Part[],
},
]
expect(await MessageV2.toModelMessages(input, model, { toolOutputMaxChars: 4 })).toStrictEqual([
{
role: "user",
content: [{ type: "text", text: "run tool" }],
},
{
role: "assistant",
content: [
{
type: "tool-call",
toolCallId: "call-1",
toolName: "bash",
input: { cmd: "ls" },
providerExecuted: undefined,
},
],
},
{
role: "tool",
content: [
{
type: "tool-result",
toolCallId: "call-1",
toolName: "bash",
output: {
type: "text",
value: "abcd\n[Tool output truncated for compaction: omitted 6 chars]",
},
},
],
},
])
})
test("converts assistant tool error into error-text tool result", async () => {
const userID = "m-user"
const assistantID = "m-assistant"

View File

@@ -837,6 +837,70 @@ describe("MessageV2.filterCompacted", () => {
})
})
test("retains an assistant tail when compaction starts inside a turn", async () => {
await Instance.provide({
directory: root,
fn: async () => {
const session = await svc.create({})
const u1 = await addUser(session.id, "first")
const a1 = await addAssistant(session.id, u1, { finish: "end_turn" })
await svc.updatePart({
id: PartID.ascending(),
sessionID: session.id,
messageID: a1,
type: "text",
text: "first reply",
})
const u2 = await addUser(session.id, "second")
const a2 = await addAssistant(session.id, u2, { finish: "end_turn" })
await svc.updatePart({
id: PartID.ascending(),
sessionID: session.id,
messageID: a2,
type: "text",
text: "second reply",
})
const a3 = await addAssistant(session.id, u2, { finish: "end_turn" })
await svc.updatePart({
id: PartID.ascending(),
sessionID: session.id,
messageID: a3,
type: "text",
text: "tail reply",
})
const c1 = await addUser(session.id)
await addCompactionPart(session.id, c1, a3)
const s1 = await addAssistant(session.id, c1, { summary: true, finish: "end_turn" })
await svc.updatePart({
id: PartID.ascending(),
sessionID: session.id,
messageID: s1,
type: "text",
text: "summary",
})
const u3 = await addUser(session.id, "third")
const a4 = await addAssistant(session.id, u3, { finish: "end_turn" })
await svc.updatePart({
id: PartID.ascending(),
sessionID: session.id,
messageID: a4,
type: "text",
text: "third reply",
})
const result = MessageV2.filterCompacted(MessageV2.stream(session.id))
expect(result.map((item) => item.info.id)).toEqual([a3, c1, s1, u3, a4])
await svc.remove(session.id)
},
})
})
test("prefers latest compaction boundary when repeated compactions exist", async () => {
await Instance.provide({
directory: root,