Merge pull request #186 from browser-use/context-limits

feat(opencode): trim old tool output and media in one batch at a token mark

Author: exe-dev-github-integration[bot]

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)",

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<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[]
@@ -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,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)
         }
       }
@@ -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: {

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)
 }

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) {

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) =>