Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
8 changes: 8 additions & 0 deletions packages/core/src/v1/config/config.ts
Original file line number Diff line number Diff line change
Expand Up @@ -154,6 +154,14 @@ export const Info = Schema.Struct({
prune: Schema.optional(Schema.Boolean).annotate({
description: "Enable pruning of old tool outputs (default: false)",
}),
prune_at: Schema.optional(PositiveInt).annotate({
description:
"With prune, clear old tool output and media in one batch inside the loop once the context reaches this many tokens (default: prune only when a turn ends)",
}),
keep_media: Schema.optional(PositiveInt).annotate({
description:
"With prune_at, the latest tool results with screenshots or other media that stay whole; everything older is cleared (default: 2)",
}),
tail_turns: Schema.optional(NonNegativeInt).annotate({
description:
"Number of recent user turns, including their following assistant/tool responses, to keep verbatim during compaction (default: 2)",
Expand Down
30 changes: 23 additions & 7 deletions packages/opencode/src/session/compaction.ts
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,9 @@ export const Event = SessionCompactionEvent

export const PRUNE_MINIMUM = 20_000
export const PRUNE_PROTECT = 40_000
const DEFAULT_KEEP_MEDIA = 2
// Media is not part of the output text estimate, so a cleared screenshot counts as this many tokens.
const MEDIA_TOKEN_ESTIMATE = 1_500
const TOOL_OUTPUT_MAX_CHARS = 2_000
const PRUNE_PROTECTED_TOOLS = ["skill"]
const DEFAULT_TAIL_TURNS = 2
Expand Down Expand Up @@ -132,7 +135,7 @@ export interface Interface {
tokens: SessionV1.Assistant["tokens"]
model: Provider.Model
}) => Effect.Effect<boolean>
readonly prune: (input: { sessionID: SessionID }) => Effect.Effect<void>
readonly prune: (input: { sessionID: SessionID; tokens?: number }) => Effect.Effect<boolean>
readonly process: (input: {
parentID: MessageID
messages: SessionV1.WithParts[]
Expand Down Expand Up @@ -240,15 +243,21 @@ const layer = Layer.effect(

// goes backwards through parts until there are PRUNE_PROTECT tokens worth of tool
// calls, then erases output of older tool calls to free context space
const prune = Effect.fn("SessionCompaction.prune")(function* (input: { sessionID: SessionID }) {
// With `tokens` (the in-loop call) it only runs once prune_at is reached and clears everything older than the
// latest keep_media results with media, in one batch; the turn-end call keeps the token-window rules.
const prune = Effect.fn("SessionCompaction.prune")(function* (input: { sessionID: SessionID; tokens?: number }) {
const cfg = yield* config.get()
if (!cfg.compaction?.prune) return
if (!cfg.compaction?.prune) return false
const trim = input.tokens !== undefined
if (trim && (cfg.compaction.prune_at === undefined || input.tokens! < cfg.compaction.prune_at)) return false
const keepMedia = cfg.compaction.keep_media ?? DEFAULT_KEEP_MEDIA
let media = 0
yield* Effect.logInfo("pruning")

const msgs = yield* session
.messages({ sessionID: input.sessionID })
.pipe(Effect.catchIf(NotFoundError.isInstance, () => Effect.succeed(undefined)))
if (!msgs) return
if (!msgs) return false

let total = 0
let pruned = 0
Expand All @@ -258,18 +267,23 @@ const layer = Layer.effect(
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 (!trim && 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") continue
if (part.state.status !== "completed") continue
if (PRUNE_PROTECTED_TOOLS.includes(part.tool)) continue
if (part.state.time.compacted) break loop
const attached = part.state.attachments?.filter((a) => MessageV2.isMedia(a.mime)).length ?? 0
if (trim && media < keepMedia) {
if (attached) media++
continue
}
const estimate = Token.estimate(part.state.output)
total += estimate
if (total <= PRUNE_PROTECT) continue
pruned += estimate
if (!trim && total <= PRUNE_PROTECT) continue
pruned += estimate + (trim ? attached * MEDIA_TOKEN_ESTIMATE : 0)
toPrune.push(part)
}
}
Expand All @@ -283,7 +297,9 @@ const layer = Layer.effect(
}
}
yield* Effect.logInfo("pruned", { count: toPrune.length })
return true
}
return false
})

const processCompaction = Effect.fn("SessionCompaction.process")(function* (input: {
Expand Down
8 changes: 5 additions & 3 deletions packages/opencode/src/session/overflow.ts
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,10 @@ export function usable(input: { cfg: ConfigV1.Info; model: Provider.Model; outpu
: Math.max(0, context - ProviderTransform.maxOutputTokens(input.model, input.outputTokenMax))
}

export function tokenCount(tokens: SessionV1.Assistant["tokens"]) {
return tokens.total || tokens.input + tokens.output + tokens.cache.read + tokens.cache.write
}

export function isOverflow(input: {
cfg: ConfigV1.Info
tokens: SessionV1.Assistant["tokens"]
Expand All @@ -28,7 +32,5 @@ export function isOverflow(input: {
if (input.cfg.compaction?.auto === false) return false
if (input.model.limit.context === 0) return false

const count =
input.tokens.total || input.tokens.input + input.tokens.output + input.tokens.cache.read + input.tokens.cache.write
return count >= usable(input)
return tokenCount(input.tokens) >= usable(input)
}
15 changes: 15 additions & 0 deletions packages/opencode/src/session/prompt.ts
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ import { Provider } from "@/provider/provider"
import { type Tool as AITool, tool, jsonSchema } from "ai"
import type { JSONSchema7 } from "@ai-sdk/provider"
import { SessionCompaction } from "./compaction"
import { tokenCount } from "./overflow"
import { SystemPrompt } from "./system"
import { Instruction } from "./instruction"
import { Plugin } from "../plugin"
Expand Down Expand Up @@ -1083,6 +1084,7 @@ const layer = Layer.effect(
const ctx = yield* InstanceState.context
let structured: unknown
let step = 0
let afterTrim = false
const session = yield* sessions.get(sessionID).pipe(Effect.orDie)

while (true) {
Expand Down Expand Up @@ -1142,6 +1144,17 @@ const layer = Layer.effect(
break
}

// A trim frees context the last finished message's usage does not show yet, so the next pass skips the
// overflow check instead of compacting on stale numbers.
const trimmed =
lastFinished && lastFinished.summary !== true
? yield* compaction.prune({ sessionID, tokens: tokenCount(lastFinished.tokens) })
: false
if (trimmed) {
afterTrim = true
continue
}

step++
if (step === 1)
yield* title({
Expand Down Expand Up @@ -1174,11 +1187,13 @@ const layer = Layer.effect(
if (
lastFinished &&
lastFinished.summary !== true &&
!afterTrim &&
(yield* compaction.isOverflow({ tokens: lastFinished.tokens, model }))
) {
yield* compaction.create({ sessionID, agent: lastUser.agent, model: lastUser.model, auto: true })
continue
}
afterTrim = false

const agent = yield* agents.get(lastUser.agent)
if (!agent) {
Expand Down
94 changes: 94 additions & 0 deletions packages/opencode/test/session/compaction.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -707,6 +707,100 @@ describe("session.compaction.prune", () => {
),
)

const screenshotResults = (count: number) =>
Effect.gen(function* () {
const ssn = yield* SessionNs.Service
const info = yield* ssn.create({})
const user = yield* ssn.updateMessage({
id: MessageID.ascending(),
role: "user",
sessionID: info.id,
agent: "build",
model: ref,
time: { created: Date.now() },
})
yield* ssn.updatePart({ id: PartID.ascending(), messageID: user.id, sessionID: info.id, type: "text", text: "task" })
const assistant: SessionV1.Assistant = {
id: MessageID.ascending(),
role: "assistant",
sessionID: info.id,
mode: "build",
agent: "build",
path: { cwd: "/", root: "/" },
cost: 0,
tokens: { output: 0, input: 0, reasoning: 0, cache: { read: 0, write: 0 } },
modelID: ref.modelID,
providerID: ref.providerID,
parentID: user.id,
time: { created: Date.now() },
}
yield* ssn.updateMessage(assistant)
for (let i = 0; i < count; i++) {
yield* ssn.updatePart({
id: PartID.ascending(),
messageID: assistant.id,
sessionID: info.id,
type: "tool",
callID: crypto.randomUUID(),
tool: "bash",
state: {
status: "completed",
input: {},
output: "x".repeat(40_000),
title: "done",
metadata: {},
time: { start: Date.now(), end: Date.now() },
attachments: [
{
id: PartID.ascending(),
messageID: assistant.id,
sessionID: info.id,
type: "file",
mime: "image/png",
filename: "screenshot.png",
url: "data:image/png;base64,Zm9v",
},
],
},
})
}
return info.id
})
const clearedParts = (id: SessionID) =>
Effect.gen(function* () {
const ssn = yield* SessionNs.Service
const msgs = yield* ssn.messages({ sessionID: id })
return msgs
.flatMap((msg) => msg.parts)
.filter((part) => part.type === "tool")
.map((part) => part.type === "tool" && part.state.status === "completed" && !!part.state.time.compacted)
})

itCompaction.instance(
"with prune_at, waits for the mark, then clears everything older than the latest media results at once",
Effect.gen(function* () {
const id = yield* screenshotResults(6)

expect(yield* SessionCompaction.use.prune({ sessionID: id, tokens: 10_000 })).toBe(false)
expect(yield* clearedParts(id)).toEqual([false, false, false, false, false, false])

expect(yield* SessionCompaction.use.prune({ sessionID: id, tokens: 60_000 })).toBe(true)
expect(yield* clearedParts(id)).toEqual([true, true, true, true, false, false])

expect(yield* SessionCompaction.use.prune({ sessionID: id, tokens: 90_000 })).toBe(false)
}).pipe(withCompaction({ config: cfg({ prune: true, prune_at: 50_000, keep_media: 2 }) })),
)

itCompaction.instance(
"without prune_at, the in-loop call does nothing",
Effect.gen(function* () {
const id = yield* screenshotResults(6)

expect(yield* SessionCompaction.use.prune({ sessionID: id, tokens: 200_000 })).toBe(false)
expect(yield* clearedParts(id)).toEqual([false, false, false, false, false, false])
}).pipe(withCompaction({ config: cfg({ prune: true }) })),
)

it.live(
"skips protected skill tool output",
provideTmpdirInstance((dir) =>
Expand Down
Loading