diff --git a/packages/core/src/v1/config/config.ts b/packages/core/src/v1/config/config.ts index cee9348166..f7be3ffddc 100644 --- a/packages/core/src/v1/config/config.ts +++ b/packages/core/src/v1/config/config.ts @@ -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)", diff --git a/packages/opencode/src/session/compaction.ts b/packages/opencode/src/session/compaction.ts index e61d1d1057..1d032cfa24 100644 --- a/packages/opencode/src/session/compaction.ts +++ b/packages/opencode/src/session/compaction.ts @@ -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 @@ -132,7 +135,7 @@ export interface Interface { tokens: SessionV1.Assistant["tokens"] model: Provider.Model }) => Effect.Effect - readonly prune: (input: { sessionID: SessionID }) => Effect.Effect + readonly prune: (input: { sessionID: SessionID; tokens?: number }) => Effect.Effect readonly process: (input: { parentID: MessageID messages: SessionV1.WithParts[] @@ -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 @@ -258,7 +267,7 @@ 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] @@ -266,10 +275,15 @@ const layer = Layer.effect( 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) } } @@ -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: { diff --git a/packages/opencode/src/session/overflow.ts b/packages/opencode/src/session/overflow.ts index 3374d1c9f8..a5d6780c8e 100644 --- a/packages/opencode/src/session/overflow.ts +++ b/packages/opencode/src/session/overflow.ts @@ -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"] @@ -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) } diff --git a/packages/opencode/src/session/prompt.ts b/packages/opencode/src/session/prompt.ts index 293abe3f45..7b4b021a6c 100644 --- a/packages/opencode/src/session/prompt.ts +++ b/packages/opencode/src/session/prompt.ts @@ -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" @@ -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) { @@ -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({ @@ -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) { diff --git a/packages/opencode/test/session/compaction.test.ts b/packages/opencode/test/session/compaction.test.ts index 4a4210cf08..b49bafb6e4 100644 --- a/packages/opencode/test/session/compaction.test.ts +++ b/packages/opencode/test/session/compaction.test.ts @@ -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) =>