enable session pruning and allow disabling with OPENCODE_DISABLE_PRUNE

4 files changed, 46 insertions, 30 deletions

diff --git a/packages/opencode/src/flag/flag.ts b/packages/opencode/src/flag/flag.ts
index a05379073..501279dd4 100644
--- a/packages/opencode/src/flag/flag.ts
+++ b/packages/opencode/src/flag/flag.ts
@@ -3,6 +3,7 @@ export namespace Flag {
   export const OPENCODE_CONFIG = process.env["OPENCODE_CONFIG"]
   export const OPENCODE_CONFIG_CONTENT = process.env["OPENCODE_CONFIG_CONTENT"]
   export const OPENCODE_DISABLE_AUTOUPDATE = truthy("OPENCODE_DISABLE_AUTOUPDATE")
+  export const OPENCODE_DISABLE_PRUNE = truthy("OPENCODE_DISABLE_PRUNE")
   export const OPENCODE_PERMISSION = process.env["OPENCODE_PERMISSION"]
   export const OPENCODE_DISABLE_DEFAULT_PLUGINS = truthy("OPENCODE_DISABLE_DEFAULT_PLUGINS")
   export const OPENCODE_DISABLE_LSP_DOWNLOAD = truthy("OPENCODE_DISABLE_LSP_DOWNLOAD")
diff --git a/packages/opencode/src/session/compaction.ts b/packages/opencode/src/session/compaction.ts
index 3f4d4b835..724517e9b 100644
--- a/packages/opencode/src/session/compaction.ts
+++ b/packages/opencode/src/session/compaction.ts
@@ -11,8 +11,12 @@ import z from "zod/v4"
 import type { ModelsDev } from "../provider/models"
 import { SessionPrompt } from "./prompt"
 import { Flag } from "../flag/flag"
+import { Token } from "../util/token"
+import { Log } from "../util/log"
 
 export namespace SessionCompaction {
+  const log = Log.create({ service: "session.compaction" })
+
   export const Event = {
     Compacted: Bus.event(
       "session.compacted",
@@ -32,6 +36,46 @@ export namespace SessionCompaction {
     return count > usable
   }
 
+  // goes backwards through parts until there are 40_000 tokens worth of tool
+  // calls. then erases output of previous tool calls. idea is to throw away old
+  // tool calls that are no longer relevant.
+  export async function prune(input: { sessionID: string }) {
+    if (Flag.OPENCODE_DISABLE_PRUNE) return
+    log.info("pruning")
+    const msgs = await Session.messages(input.sessionID)
+    let total = 0
+    let pruned = 0
+    const toPrune = []
+
+    loop: for (let msgIndex = msgs.length - 2; msgIndex >= 0; msgIndex--) {
+      const msg = msgs[msgIndex]
+      if (msg.info.role === "assistant" && msg.info.summary) return
+      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 (part.state.time.compacted) break loop
+            const estimate = Token.estimate(part.state.output)
+            total += estimate
+            if (total > 40_000) {
+              pruned += estimate
+              toPrune.push(part)
+            }
+          }
+      }
+    }
+    log.info("found", { pruned, total })
+    if (pruned > 20_000) {
+      for (const part of toPrune) {
+        if (part.state.status === "completed") {
+          part.state.time.compacted = Date.now()
+          await Session.updatePart(part)
+        }
+      }
+      log.info("pruned", { count: toPrune.length })
+    }
+  }
+
   export async function run(input: { sessionID: string; providerID: string; modelID: string }) {
     await Session.update(input.sessionID, (draft) => {
       draft.time.compacting = Date.now()
diff --git a/packages/opencode/src/session/index.ts b/packages/opencode/src/session/index.ts
index 18bb7aef6..cce1cf8c6 100644
--- a/packages/opencode/src/session/index.ts
+++ b/packages/opencode/src/session/index.ts
@@ -16,7 +16,6 @@ import { Log } from "../util/log"
 import { MessageV2 } from "./message-v2"
 import { Project } from "../project/project"
 import { Instance } from "../project/instance"
-import { Token } from "../util/token"
 import { SessionPrompt } from "./prompt"
 
 export namespace Session {
@@ -293,34 +292,6 @@ export namespace Session {
     return part
   }
 
-  // goes backwards through parts until there are 40_000 tokens worth of tool
-  // calls. then erases output of previous tool calls. idea is to throw away old
-  // tool calls that are no longer relevant.
-  export async function prune(input: { sessionID: string }) {
-    const msgs = await messages(input.sessionID)
-    let sum = 0
-    for (let msgIndex = msgs.length - 2; msgIndex >= 0; msgIndex--) {
-      const msg = msgs[msgIndex]
-      if (msg.info.role === "assistant" && msg.info.summary) return
-      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 (part.state.time.compacted) return
-            sum += Token.estimate(part.state.output)
-            if (sum > 40_000) {
-              log.info("pruning", {
-                sum,
-                id: part.id,
-              })
-              part.state.time.compacted = Date.now()
-              await updatePart(part)
-            }
-          }
-      }
-    }
-  }
-
   export function getUsage(model: ModelsDev.Model, usage: LanguageModelUsage, metadata?: ProviderMetadata) {
     const tokens = {
       input: usage.inputTokens ?? 0,
diff --git a/packages/opencode/src/session/prompt.ts b/packages/opencode/src/session/prompt.ts
index 90ffdec72..392413a41 100644
--- a/packages/opencode/src/session/prompt.ts
+++ b/packages/opencode/src/session/prompt.ts
@@ -321,7 +321,7 @@ export namespace SessionPrompt {
         item.callback(result)
       }
       state().queued.delete(input.sessionID)
-      // Session.prune(input)
+      SessionCompaction.prune(input)
       return result
     }
   }