From 3ad45fdcfa7ed4bb5ef743c4568aedc7813cdf18 Mon Sep 17 00:00:00 2001 From: Krishna Date: Tue, 6 Oct 2026 03:05:11 +0530 Subject: [PATCH] fix(cli): reuse the loaded transcript when pruning large payloads Once a request passed the payload limit, every step handed prune only a session ID. Prune paged through the whole session history, then the loop reloaded the transcript, reran its transforms and serialized it again, even when nothing was cleared. The reload also skipped the mode reminders, so Plan sessions over the limit lost their plan instructions. Prune now takes the messages the loop already has and marks cleared outputs on them. It saves each cleared part from storage so in-memory transforms are never written back, then returns how many it cleared. The loop converts again only when that number is above zero. Refs #14454 --- .changeset/single-load-payload-prune.md | 5 +++ packages/opencode/src/session/compaction.ts | 44 ++++++++++++------- packages/opencode/src/session/prompt.ts | 20 +++------ .../session-prompt-compaction-safety.test.ts | 43 +++++++++++++++++- .../test/kilocode/session-pruning.test.ts | 44 +++++++++++++++++++ 5 files changed, 126 insertions(+), 30 deletions(-) create mode 100644 .changeset/single-load-payload-prune.md diff --git a/.changeset/single-load-payload-prune.md b/.changeset/single-load-payload-prune.md new file mode 100644 index 00000000000..6b52142d0d9 --- /dev/null +++ b/.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. diff --git a/packages/opencode/src/session/compaction.ts b/packages/opencode/src/session/compaction.ts index ee4fa7a8d46..42345f1ef14 100644 --- a/packages/opencode/src/session/compaction.ts +++ b/packages/opencode/src/session/compaction.ts @@ -185,7 +185,13 @@ export interface Interface { tokens: SessionV1.Assistant["tokens"] model: Provider.Model }) => Effect.Effect - readonly prune: (input: { sessionID: SessionID; reason?: PruneReason }) => Effect.Effect // 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 + // 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 diff --git a/packages/opencode/src/session/prompt.ts b/packages/opencode/src/session/prompt.ts index 353bc725d00..7e0ed82addb 100644 --- a/packages/opencode/src/session/prompt.ts +++ b/packages/opencode/src/session/prompt.ts @@ -1817,19 +1817,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 }) } diff --git a/packages/opencode/test/kilocode/session-prompt-compaction-safety.test.ts b/packages/opencode/test/kilocode/session-prompt-compaction-safety.test.ts index db73fa246c1..63c4768852f 100644 --- a/packages/opencode/test/kilocode/session-prompt-compaction-safety.test.ts +++ b/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 }) { diff --git a/packages/opencode/test/kilocode/session-pruning.test.ts b/packages/opencode/test/kilocode/session-pruning.test.ts index c7a3f90bcb4..1172e11797d 100644 --- a/packages/opencode/test/kilocode/session-pruning.test.ts +++ b/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",