Skip to content

refactor: extract compaction prune planning #16

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jairad26 merged 1 commit into from
Mar 31, 2026
Merged

refactor: extract compaction prune planning #16

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
86 changes: 51 additions & 35 deletions packages/opencode/src/session/compaction.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,49 @@ export namespace SessionCompaction {
export const PRUNE_PROTECT = 40_000
const PRUNE_PROTECTED_TOOLS = ["skill"]

export function prunePlan(input: {
messages: MessageV2.WithParts[]
protect?: number
minimum?: number
turns?: number
protected?: readonly string[]
}) {
const protect = input.protect ?? PRUNE_PROTECT
const minimum = input.minimum ?? PRUNE_MINIMUM
const turnsToKeep = input.turns ?? 2
const protectedTools = input.protected ?? PRUNE_PROTECTED_TOOLS
let total = 0
let pruned = 0
let turns = 0
const parts: MessageV2.ToolPart[] = []

loop: for (let msgIndex = input.messages.length - 1; msgIndex >= 0; msgIndex--) {
const msg = input.messages[msgIndex]
if (msg.info.role === "user") turns++
if (turns < turnsToKeep) continue
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") continue
if (part.state.status !== "completed") continue
if (protectedTools.includes(part.tool)) continue
if (part.state.time.compacted) break loop
const estimate = Token.estimate(part.state.output)
total += estimate
if (total <= protect) continue
pruned += estimate
parts.push(part)
}
}

return {
total,
pruned,
parts,
shouldPrune: pruned > minimum,
}
}

export interface Interface {
readonly isOverflow: (input: {
tokens: MessageV2.Assistant["tokens"]
Expand Down Expand Up @@ -92,42 +135,15 @@ export namespace SessionCompaction {
.pipe(Effect.catchIf(NotFoundError.isInstance, () => Effect.succeed(undefined)))
if (!msgs) return

let total = 0
let pruned = 0
const toPrune: MessageV2.ToolPart[] = []
let turns = 0

loop: for (let msgIndex = msgs.length - 1; msgIndex >= 0; msgIndex--) {
const msg = msgs[msgIndex]
if (msg.info.role === "user") turns++
if (turns < 2) continue
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)
}
}
}
}

log.info("found", { pruned, total })
if (pruned > PRUNE_MINIMUM) {
for (const part of toPrune) {
if (part.state.status === "completed") {
part.state.time.compacted = Date.now()
yield* session.updatePart(part)
}
}
log.info("pruned", { count: toPrune.length })
const plan = prunePlan({ messages: msgs })
log.info("found", { pruned: plan.pruned, total: plan.total })
if (!plan.shouldPrune) return
for (const part of plan.parts) {
if (part.state.status !== "completed") continue
part.state.time.compacted = Date.now()
yield* session.updatePart(part)
}
log.info("pruned", { count: plan.parts.length })
})

const processCompaction = Effect.fn("SessionCompaction.process")(function* (input: {
Expand Down
70 changes: 70 additions & 0 deletions packages/opencode/test/session/compaction.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -458,6 +458,76 @@ describe("session.compaction.create", () => {
})

describe("session.compaction.prune", () => {
test("selects only eligible older completed tool outputs in deterministic order", async () => {
await using tmp = await tmpdir()
await Instance.provide({
directory: tmp.path,
fn: async () => {
const session = await Session.create({})
const u0 = await user(session.id, "u0")
const a0 = await assistant(session.id, u0.id, tmp.path)
const p0 = await tool(session.id, a0.id, "bash", "x".repeat(120_000))

const u1 = await user(session.id, "u1")
const a1 = await assistant(session.id, u1.id, tmp.path)
const p1 = await tool(session.id, a1.id, "bash", "x".repeat(120_000))

const u2 = await user(session.id, "u2")
const a2 = await assistant(session.id, u2.id, tmp.path)
const p2 = await tool(session.id, a2.id, "bash", "x".repeat(120_000))

const u3 = await user(session.id, "u3")
const a3 = await assistant(session.id, u3.id, tmp.path)
const p3 = await tool(session.id, a3.id, "bash", "x".repeat(120_000))

await user(session.id, "u4")

const msgs = await Session.messages({ sessionID: session.id })
const plan = SessionCompaction.prunePlan({ messages: msgs })

expect(plan.parts.map((part) => part.id)).toEqual([p1.id, p0.id])
expect(plan.parts.some((part) => part.id === p2.id)).toBe(false)
expect(plan.parts.some((part) => part.id === p3.id)).toBe(false)
expect(plan.shouldPrune).toBe(true)
},
})
})

test("does not prune when candidate total is at or below minimum", async () => {
await using tmp = await tmpdir()
await Instance.provide({
directory: tmp.path,
fn: async () => {
const session = await Session.create({})
const u1 = await user(session.id, "first")
const a1 = await assistant(session.id, u1.id, tmp.path)
await tool(session.id, a1.id, "bash", "x".repeat(80_000))

const u2 = await user(session.id, "second")
const a2 = await assistant(session.id, u2.id, tmp.path)
await tool(session.id, a2.id, "bash", "x".repeat(160_000))

await user(session.id, "third")
await user(session.id, "fourth")

const msgs = await Session.messages({ sessionID: session.id })
const plan = SessionCompaction.prunePlan({ messages: msgs })
expect(plan.pruned).toBe(20_000)
expect(plan.shouldPrune).toBe(false)

await SessionCompaction.prune({ sessionID: session.id })

const all = await Session.messages({ sessionID: session.id })
const compacted = all
.flatMap((msg) => msg.parts)
.flatMap((part) =>
part.type === "tool" && part.state.status === "completed" && part.state.time.compacted ? [part] : [],
)
expect(compacted).toHaveLength(0)
},
})
})

test("compacts old completed tool output", async () => {
await using tmp = await tmpdir()
await Instance.provide({
Expand Down