Merge pull request #14816 from kr1shna-exe/fix/14454-single-load-prune

fix(cli): reuse the loaded transcript when pruning large payloads

.changeset/single-load-payload-prune.md

@@ -0,0 +1,5 @@
+---
+"@kilocode/cli": patch
+---
+
+Reduce the work done on every agent step in long sessions once the request passes the payload pruning limit. The session transcript is no longer reloaded and re-serialized after pruning, and pruning no longer reads the whole session history again. Plan mode also keeps its plan file instructions on those steps.

packages/opencode/src/session/compaction.ts

@@ -185,7 +185,13 @@ export interface Interface {
     tokens: SessionV1.Assistant["tokens"]
     model: Provider.Model
   }) => Effect.Effect<boolean>
-  readonly prune: (input: { sessionID: SessionID; reason?: PruneReason }) => Effect.Effect<void> // kilocode_change
+  // kilocode_change start - callers that already loaded the transcript pass it as messages; returns cleared outputs
+  readonly prune: (input: {
+    sessionID: SessionID
+    reason?: PruneReason
+    messages?: SessionV1.WithParts[]
+  }) => Effect.Effect<number>
+  // kilocode_change end
   readonly process: (input: {
     parentID: MessageID
     messages: SessionV1.WithParts[]
@@ -299,17 +305,20 @@ const layer = Layer.effect(
     const prune = Effect.fn("SessionCompaction.prune")(function* (input: {
       sessionID: SessionID
       reason?: PruneReason
+      messages?: SessionV1.WithParts[]
     }) {
       const cfg = yield* config.get()
       const reason = input.reason ?? "normal"
-      if (cfg.compaction?.prune === false) return
-      if (reason === "normal" && cfg.compaction?.prune !== true) return
+      if (cfg.compaction?.prune === false) return 0
+      if (reason === "normal" && cfg.compaction?.prune !== true) return 0
       yield* Effect.logInfo("pruning", { reason })
 
-      const msgs = yield* session
-        .messages({ sessionID: input.sessionID })
-        .pipe(Effect.catchIf(NotFoundError.isInstance, () => Effect.succeed(undefined)))
-      if (!msgs) return
+      const msgs =
+        input.messages ??
+        (yield* session
+          .messages({ sessionID: input.sessionID })
+          .pipe(Effect.catchIf(NotFoundError.isInstance, () => Effect.succeed(undefined))))
+      if (!msgs) return 0
 
       let total = 0
       let pruned = 0
@@ -342,15 +351,20 @@ const layer = Layer.effect(
       }
 
       yield* Effect.logInfo("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)
-          }
-        }
-        yield* Effect.logInfo("pruned", { reason, count: toPrune.length })
+      if (pruned <= PRUNE_MINIMUM) return 0
+      for (const part of toPrune) {
+        if (part.state.status !== "completed") continue
+        part.state.time.compacted = Date.now()
+        // Supplied messages may carry prompt-time transforms, so persist the stored part instead.
+        const stored = input.messages
+          ? yield* session.getPart({ sessionID: part.sessionID, messageID: part.messageID, partID: part.id })
+          : part
+        if (stored?.type !== "tool" || stored.state.status !== "completed") continue
+        stored.state.time.compacted = part.state.time.compacted
+        yield* session.updatePart(stored)
       }
+      yield* Effect.logInfo("pruned", { reason, count: toPrune.length })
+      return toPrune.length
     })
     // kilocode_change end
 

packages/opencode/src/session/prompt.ts

@@ -1824,19 +1824,13 @@ export const layer = Layer.effect(
           )
           const size = Buffer.byteLength(JSON.stringify(modelMsgs))
           if (size > REQUEST_PRUNE_BYTES) {
-            yield* compaction.prune({ sessionID, reason: "payload-limit" })
-            msgs = yield* MessageV2.filterCompactedEffect(sessionID).pipe(
-              Effect.provideService(Database.Service, database),
-            )
-            msgs = KiloSessionPromptQueue.scope(sessionID, msgs)
-            msgs = KiloSessionPrompt.trimBeforeLastSummary(msgs)
-            yield* plugin.trigger("experimental.chat.messages.transform", {}, { messages: msgs })
-            KiloSessionPrompt.injectEditorContext({ msgs, session, sessionID, cache: envCache })
-            msgs = KiloSessionPrompt.maybeStripHistoricalMedia(msgs)
-            modelMsgs = yield* MessageV2.toModelMessagesEffect(msgs, model).pipe(
-              Effect.provideService(Database.Service, database),
-            )
-            const nextSize = Buffer.byteLength(JSON.stringify(modelMsgs))
+            // prune marks cleared outputs on msgs, so re-convert only when it cleared something
+            const cleared = yield* compaction.prune({ sessionID, reason: "payload-limit", messages: msgs })
+            if (cleared > 0)
+              modelMsgs = yield* MessageV2.toModelMessagesEffect(msgs, model).pipe(
+                Effect.provideService(Database.Service, database),
+              )
+            const nextSize = cleared > 0 ? Buffer.byteLength(JSON.stringify(modelMsgs)) : size
             if (nextSize > REQUEST_PRUNE_BYTES)
               yield* Effect.logWarning("payload still large after pruning", { "session.id": sessionID, size: nextSize })
           }

packages/opencode/test/kilocode/session-prompt-compaction-safety.test.ts

@@ -229,14 +229,14 @@ function providerCfg(url: string) {
 const user = Effect.fn("prompt-safety.user")(function* (
   sessionID: SessionID,
   text: string,
-  input?: { synthetic?: boolean; editorContext?: MessageV2.User["editorContext"] },
+  input?: { synthetic?: boolean; editorContext?: MessageV2.User["editorContext"]; agent?: string },
 ) {
   const sessions = yield* Session.Service
   const msg = yield* sessions.updateMessage({
     id: MessageID.ascending(),
     role: "user",
     sessionID,
-    agent: "code",
+    agent: input?.agent ?? "code",
     model: ref,
     time: { created: Date.now() },
     tools: {},
@@ -376,6 +376,45 @@ describe("SessionPrompt compaction safety", () => {
     ),
   )
 
+  it.live("keeps plan reminders when pruning an oversized request", () =>
+    provideTmpdirServer(
+      Effect.fnUntraced(function* ({ llm }) {
+        const prompt = yield* SessionPrompt.Service
+        const sessions = yield* Session.Service
+        const chat = yield* sessions.create({ title: "Plan pruning" })
+        const request = yield* user(chat.id, "Plan the task", { agent: "plan" })
+        for (const output of ["stale-output".repeat(120_000), "recent-one", "recent-two"]) {
+          const response = yield* assistant(chat.id, request.id, { finish: "tool-calls" })
+          yield* sessions.updateMessage({ ...response, time: { ...response.time, completed: Date.now() } })
+          yield* sessions.updatePart({
+            id: PartID.ascending(),
+            sessionID: chat.id,
+            messageID: response.id,
+            type: "tool",
+            callID: crypto.randomUUID(),
+            tool: "bash",
+            state: {
+              status: "completed",
+              input: { command: "pwd" },
+              output,
+              title: "result",
+              metadata: {},
+              time: { start: Date.now(), end: Date.now() },
+            },
+          })
+        }
+        yield* llm.text("final answer")
+
+        yield* prompt.loop({ sessionID: chat.id })
+
+        const body = JSON.stringify((yield* llm.inputs).at(-1)?.messages)
+        expect(body).not.toContain("stale-output")
+        expect(body).toContain("## Plan File")
+      }),
+      { git: true, config: (url) => ({ ...providerCfg(url), compaction: { auto: false } }) },
+    ),
+  )
+
   it.live("compacts estimated outgoing context before the provider request", () =>
     provideTmpdirServer(
       Effect.fnUntraced(function* ({ llm }) {

packages/opencode/test/kilocode/session-pruning.test.ts

@@ -293,6 +293,50 @@ describe("single-turn tool-output pruning", () => {
     { config },
   )
 
+  it.instance(
+    "scans supplied messages instead of reloading the session",
+    Effect.gen(function* () {
+      const parent = yield* setup
+      const sessions = yield* Session.Service
+      const compact = yield* SessionCompaction.Service
+      const parts: MessageV2.ToolPart[] = []
+      for (const output of outputs) parts.push(yield* step(parent, output))
+      const old = new Set(parts.slice(0, 2).map((part) => part.messageID))
+      const messages = (yield* sessions.messages({ sessionID: parent.sessionID })).filter(
+        (message) => !old.has(message.info.id),
+      )
+
+      const cleared = yield* compact.prune({ sessionID: parent.sessionID, reason: "payload-limit", messages })
+
+      expect(cleared).toBe(0)
+      expect(marked(yield* tools(parent.sessionID))).toEqual([false, false, false, false, false, false])
+    }),
+  )
+
+  it.instance(
+    "marks supplied parts but persists the stored part",
+    Effect.gen(function* () {
+      const parent = yield* setup
+      const sessions = yield* Session.Service
+      const compact = yield* SessionCompaction.Service
+      for (const output of outputs) yield* step(parent, output)
+      const messages = yield* sessions.messages({ sessionID: parent.sessionID })
+      const supplied = messages.flatMap((message) => message.parts).filter((part) => part.type === "tool")
+      for (const part of supplied) if (part.state.status === "completed") part.state.title = "transformed"
+
+      const cleared = yield* compact.prune({ sessionID: parent.sessionID, reason: "payload-limit", messages })
+
+      expect(cleared).toBe(2)
+      expect(marked(supplied)).toEqual([true, true, false, false, false, false])
+      const stored = yield* tools(parent.sessionID)
+      expect(marked(stored)).toEqual([true, true, false, false, false, false])
+      expect(stored.map((part) => part.state.status === "completed" && part.state.title)).toEqual(
+        Array(6).fill("result"),
+      )
+      expect(stored.map((part) => part.state.status === "completed" && part.state.output)).toEqual(outputs)
+    }),
+  )
+
   for (const reason of ["normal", "payload-limit"] as const) {
     it.instance(
       reason === "normal" ? "keeps normal pruning opt-in" : "honors explicitly disabled payload pruning",