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 a1db53ace7a..e90c343fa9e 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 4b2601eada0..07625eacefe 100644 --- a/packages/opencode/src/session/prompt.ts +++ b/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 }) } 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",