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

@@ -32,16 +32,105 @@ export const Event = {
export const PRUNE_MINIMUM = 20_000
export const PRUNE_PROTECT = 40_000
const TOOL_OUTPUT_MAX_CHARS = 2_000
const PRUNE_PROTECTED_TOOLS = ["skill"]
const DEFAULT_TAIL_TURNS = 2
const MIN_PRESERVE_RECENT_TOKENS = 2_000
const MAX_PRESERVE_RECENT_TOKENS = 8_000
const SUMMARY_TEMPLATE = `Output exactly this Markdown structure and keep the section order unchanged:
---
## Goal
- [single-sentence task summary]
## Constraints & Preferences
- [user constraints, preferences, specs, or "(none)"]
## Progress
### Done
- [completed work or "(none)"]
### In Progress
- [current work or "(none)"]
### Blocked
- [blockers or "(none)"]
## Key Decisions
- [decision and why, or "(none)"]
## Next Steps
- [ordered next actions or "(none)"]
## Critical Context
- [important technical facts, errors, open questions, or "(none)"]
## Relevant Files
- [file or directory path: why it matters, or "(none)"]
---
Rules:
- Keep every section, even when empty.
- Use terse bullets, not prose paragraphs.
- Preserve exact file paths, commands, error strings, and identifiers when known.
- Do not mention the summary process or that context was compacted.`
type Turn = {
start: number
end: number
id: MessageID
}
type Tail = {
start: number
id: MessageID
}
type CompletedCompaction = {
userIndex: number
assistantIndex: number
summary: string | undefined
}
function summaryText(message: MessageV2.WithParts) {
const text = message.parts
.filter((part): part is MessageV2.TextPart => part.type === "text")
.map((part) => part.text.trim())
.filter(Boolean)
.join("\n\n")
.trim()
return text || undefined
}
function completedCompactions(messages: MessageV2.WithParts[]) {
const users = new Map<MessageID, number>()
for (let i = 0; i < messages.length; i++) {
const msg = messages[i]
if (msg.info.role !== "user") continue
if (!msg.parts.some((part) => part.type === "compaction")) continue
users.set(msg.info.id, i)
}
return messages.flatMap((msg, assistantIndex): CompletedCompaction[] => {
if (msg.info.role !== "assistant") return []
if (!msg.info.summary || !msg.info.finish || msg.info.error) return []
const userIndex = users.get(msg.info.parentID)
if (userIndex === undefined) return []
return [{ userIndex, assistantIndex, summary: summaryText(msg) }]
})
}
function buildPrompt(input: { previousSummary?: string; context: string[] }) {
const anchor = input.previousSummary
? [
"Update the anchored summary below using the conversation history above.",
"Preserve still-true details, remove stale details, and merge in the new facts.",
"<previous-summary>",
input.previousSummary,
"</previous-summary>",
].join("\n")
: "Create a new anchored summary from the conversation history above."
return [anchor, SUMMARY_TEMPLATE, ...input.context].join("\n\n")
}
function preserveRecentBudget(input: { cfg: Config.Info; model: Provider.Model }) {
return (
input.cfg.compaction?.preserve_recent_tokens ??
@@ -67,6 +156,31 @@ function turns(messages: MessageV2.WithParts[]) {
return result
}
function splitTurn(input: {
messages: MessageV2.WithParts[]
turn: Turn
model: Provider.Model
budget: number
estimate: (input: { messages: MessageV2.WithParts[]; model: Provider.Model }) => Effect.Effect<number>
}) {
return Effect.gen(function* () {
if (input.budget <= 0) return undefined
if (input.turn.end - input.turn.start <= 1) return undefined
for (let start = input.turn.start + 1; start < input.turn.end; start++) {
const size = yield* input.estimate({
messages: input.messages.slice(start, input.turn.end),
model: input.model,
})
if (size > input.budget) continue
return {
start,
id: input.messages[start]!.info.id,
} satisfies Tail
}
return undefined
})
}
export interface Interface {
readonly isOverflow: (input: {
tokens: MessageV2.Assistant["tokens"]
@@ -147,18 +261,28 @@ export const layer: Layer.Layer<
}),
{ concurrency: 1 },
)
if (sizes.at(-1)! > budget) {
log.info("tail fallback", { budget, size: sizes.at(-1) })
return { head: input.messages, tail_start_id: undefined }
}
let total = 0
let keep: Turn | undefined
let keep: Tail | undefined
for (let i = recent.length - 1; i >= 0; i--) {
const turn = recent[i]!
const size = sizes[i]
if (total + size > budget) break
total += size
keep = recent[i]
if (total + size <= budget) {
total += size
keep = { start: turn.start, id: turn.id }
continue
}
const remaining = budget - total
const split = yield* splitTurn({
messages: input.messages,
turn,
model: input.model,
budget: remaining,
estimate,
})
if (split) keep = split
else if (!keep) log.info("tail fallback", { budget, size, total })
break
}
if (!keep || keep.start === 0) return { head: input.messages, tail_start_id: undefined }
@@ -192,17 +316,15 @@ export const layer: Layer.Layer<
if (msg.info.role === "assistant" && msg.info.summary) break loop
for (let partIndex = msg.parts.length - 1; partIndex >= 0; partIndex--) {
const part = msg.parts[partIndex]
if (part.type === "tool")
if (part.state.status === "completed") {
if (PRUNE_PROTECTED_TOOLS.includes(part.tool)) continue
if (part.state.time.compacted) break loop
const estimate = Token.estimate(part.state.output)
total += estimate
if (total > PRUNE_PROTECT) {
pruned += estimate
toPrune.push(part)
}
}
if (part.type !== "tool") continue
if (part.state.status !== "completed") continue
if (PRUNE_PROTECTED_TOOLS.includes(part.tool)) continue
if (part.state.time.compacted) break loop
const estimate = Token.estimate(part.state.output)
total += estimate
if (total <= PRUNE_PROTECT) continue
pruned += estimate
toPrune.push(part)
}
}
@@ -263,8 +385,11 @@ export const layer: Layer.Layer<
: yield* provider.getModel(userMessage.model.providerID, userMessage.model.modelID)
const cfg = yield* config.get()
const history = compactionPart && messages.at(-1)?.info.id === input.parentID ? messages.slice(0, -1) : messages
const prior = completedCompactions(history)
const hidden = new Set(prior.flatMap((item) => [item.userIndex, item.assistantIndex]))
const previousSummary = prior.at(-1)?.summary
const selected = yield* select({
messages: history,
messages: history.filter((_, index) => !hidden.has(index)),
cfg,
model,
})
@@ -274,34 +399,13 @@ export const layer: Layer.Layer<
{ sessionID: input.sessionID },
{ context: [], prompt: undefined },
)
const defaultPrompt = `When constructing the summary, try to stick to this template:
---
## Goal
[What goal(s) is the user trying to accomplish?]
## Instructions
- [What important instructions did the user give you that are relevant]
- [If there is a plan or spec, include information about it so next agent can continue using it]
## Discoveries
[What notable things were learned during this conversation that would be useful for the next agent to know when continuing the work]
## Accomplished
[What work has been completed, what work is still in progress, and what work is left?]
## Relevant files / directories
[Construct a structured list of relevant files that have been read, edited, or created that pertain to the task at hand. If all the files in a directory are relevant, include the path to the directory.]
---`
const prompt = compacting.prompt ?? [defaultPrompt, ...compacting.context].join("\n\n")
const nextPrompt = compacting.prompt ?? buildPrompt({ previousSummary, context: compacting.context })
const msgs = structuredClone(selected.head)
yield* plugin.trigger("experimental.chat.messages.transform", {}, { messages: msgs })
const modelMessages = yield* MessageV2.toModelMessagesEffect(msgs, model, { stripMedia: true })
const modelMessages = yield* MessageV2.toModelMessagesEffect(msgs, model, {
stripMedia: true,
toolOutputMaxChars: TOOL_OUTPUT_MAX_CHARS,
})
const ctx = yield* InstanceState.context
const msg: MessageV2.Assistant = {
id: MessageID.ascending(),
@@ -345,7 +449,7 @@ export const layer: Layer.Layer<
...modelMessages,
{
role: "user",
content: [{ type: "text", text: prompt }],
content: [{ type: "text", text: nextPrompt }],
},
],
model,

View File

@@ -319,6 +319,12 @@ export const ToolStateCompleted = Schema.Struct({
.pipe(withStatics((s) => ({ zod: zod(s) })))
export type ToolStateCompleted = Types.DeepMutable<Schema.Schema.Type<typeof ToolStateCompleted>>
function truncateToolOutput(text: string, maxChars?: number) {
if (!maxChars || text.length <= maxChars) return text
const omitted = text.length - maxChars
return `${text.slice(0, maxChars)}\n[Tool output truncated for compaction: omitted ${omitted} chars]`
}
export const ToolStateError = Schema.Struct({
status: Schema.Literal("error"),
input: Schema.Record(Schema.String, Schema.Any),
@@ -700,7 +706,7 @@ function providerMeta(metadata: Record<string, any> | undefined) {
export const toModelMessagesEffect = Effect.fnUntraced(function* (
input: WithParts[],
model: Provider.Model,
options?: { stripMedia?: boolean },
options?: { stripMedia?: boolean; toolOutputMaxChars?: number },
) {
const result: UIMessage[] = []
const toolNames = new Set<string>()
@@ -839,7 +845,9 @@ export const toModelMessagesEffect = Effect.fnUntraced(function* (
if (part.type === "tool") {
toolNames.add(part.tool)
if (part.state.status === "completed") {
const outputText = part.state.time.compacted ? "[Old tool result content cleared]" : part.state.output
const outputText = part.state.time.compacted
? "[Old tool result content cleared]"
: truncateToolOutput(part.state.output, options?.toolOutputMaxChars)
const attachments = part.state.time.compacted || options?.stripMedia ? [] : (part.state.attachments ?? [])
// For providers that don't support media in tool results, extract media files
@@ -955,7 +963,7 @@ export const toModelMessagesEffect = Effect.fnUntraced(function* (
export function toModelMessages(
input: WithParts[],
model: Provider.Model,
options?: { stripMedia?: boolean },
options?: { stripMedia?: boolean; toolOutputMaxChars?: number },
): Promise<ModelMessage[]> {
return Effect.runPromise(toModelMessagesEffect(input, model, options).pipe(Effect.provide(EffectLogger.layer)))
}