diff --git a/.changeset/soft-compaction-recovery.md b/.changeset/soft-compaction-recovery.md new file mode 100644 index 00000000000..7be1338ca62 --- /dev/null +++ b/.changeset/soft-compaction-recovery.md @@ -0,0 +1,5 @@ +--- +"@kilocode/cli": patch +--- + +Recover compaction when large tool results or media attachments exceed provider payload limits. diff --git a/packages/opencode/src/kilocode/session/compaction-payload-recovery.ts b/packages/opencode/src/kilocode/session/compaction-payload-recovery.ts new file mode 100644 index 00000000000..10865a3ee38 --- /dev/null +++ b/packages/opencode/src/kilocode/session/compaction-payload-recovery.ts @@ -0,0 +1,111 @@ +import { Effect } from "effect" +import type { Agent } from "@/agent/agent" +import type { Provider } from "@/provider/provider" +import type { LLM } from "@/session/llm" +import { MessageV2 } from "@/session/message-v2" +import type { SessionProcessor } from "@/session/processor" +import type { MessageID, SessionID } from "@/session/schema" + +type Update = (part: T) => Effect.Effect +type UpdateMessage = (msg: T) => Effect.Effect + +const pattern = /request entity too large|function_payload_too_large/i + +export namespace KiloCompactionPayloadRecovery { + export function matches(error: MessageV2.Assistant["error"]) { + if (!error) return false + if (error.name !== "ContextOverflowError" && error.name !== "APIError") return false + return pattern.test([error.data.message, error.data.responseBody].filter(Boolean).join("\n")) + } + + export function prompt(text: string) { + return [ + "The previous compaction request exceeded the provider's 4MB payload limit.", + "Older tool outputs and media attachments were removed from this compaction request.", + text, + ].join("\n\n") + } + + export function strip(input: { + messages: MessageV2.WithParts[] + update: Update + }) { + return Effect.forEach( + input.messages, + (msg) => + Effect.forEach(msg.parts, (part) => { + if (part.type === "tool" && part.state.status === "completed" && !part.state.time.compacted) { + part.state.time.compacted = Date.now() + return input.update(part) + } + if (part.type === "file" && MessageV2.isMedia(part.mime)) { + return input.update({ + id: part.id, + messageID: part.messageID, + sessionID: part.sessionID, + type: "text", + text: `[Attached ${part.mime}: ${part.filename ?? "file"}]`, + }) + } + return Effect.void + }), + { concurrency: 1 }, + ) + } + + export function process(input: { + processor: SessionProcessor.Handle + user: MessageV2.User + agent: Agent.Info + sessionID: SessionID + model: Provider.Model + messages: LLM.StreamInput["messages"] + prompt: string + recovery: MessageV2.WithParts[] + updateMessage: UpdateMessage + updatePart: Update + }) { + const run = Effect.fn("KiloCompactionPayloadRecovery.process")(function* ( + messages: LLM.StreamInput["messages"], + text: string, + ) { + return yield* input.processor.process({ + user: input.user, + agent: input.agent, + sessionID: input.sessionID, + tools: {}, + system: [], + messages: [ + ...messages, + { + role: "user", + content: [{ type: "text", text }], + }, + ], + model: input.model, + }) + }) + + return run(input.messages, input.prompt).pipe( + Effect.flatMap((result) => { + if (result !== "compact" && (result !== "stop" || !matches(input.processor.message.error))) { + return Effect.succeed(result) + } + if (result === "compact" && !matches(input.processor.compactError?.())) { + return Effect.succeed(result) + } + return Effect.gen(function* () { + input.processor.message.error = undefined + input.processor.message.finish = undefined + yield* input.updateMessage(input.processor.message) + yield* strip({ messages: input.recovery, update: input.updatePart }) + const stripped = yield* MessageV2.toModelMessagesEffect(input.recovery, input.model, { + stripMedia: true, + toolOutputMaxChars: 0, + }) + return yield* run(stripped, prompt(input.prompt)) + }) + }), + ) + } +} diff --git a/packages/opencode/src/session/compaction.ts b/packages/opencode/src/session/compaction.ts index b9a4294837a..91b0ef4b863 100644 --- a/packages/opencode/src/session/compaction.ts +++ b/packages/opencode/src/session/compaction.ts @@ -19,6 +19,7 @@ import { isOverflow as overflow, usable } from "./overflow" import { makeRuntime } from "@/effect/run-service" import { fn } from "@/util/fn" import { KiloSessionPromptQueue } from "@/kilocode/session/prompt-queue" // kilocode_change +import { KiloCompactionPayloadRecovery } from "@/kilocode/session/compaction-payload-recovery" // kilocode_change const log = Log.create({ service: "session.compaction" }) @@ -440,21 +441,20 @@ export const layer: Layer.Layer< sessionID: input.sessionID, model, }) - const result = yield* processor.process({ + // kilocode_change start + const result = yield* KiloCompactionPayloadRecovery.process({ + processor, user: userMessage, agent, sessionID: input.sessionID, - tools: {}, - system: [], - messages: [ - ...modelMessages, - { - role: "user", - content: [{ type: "text", text: nextPrompt }], - }, - ], model, + messages: modelMessages, + prompt: nextPrompt, + recovery: selected.head, + updateMessage: session.updateMessage, + updatePart: session.updatePart, }) + // kilocode_change end if (result === "compact") { processor.message.error = new MessageV2.ContextOverflowError({ diff --git a/packages/opencode/src/session/processor.ts b/packages/opencode/src/session/processor.ts index 4d418901ce5..122c44baf71 100644 --- a/packages/opencode/src/session/processor.ts +++ b/packages/opencode/src/session/processor.ts @@ -47,6 +47,7 @@ export interface Handle { }, ) => Effect.Effect readonly process: (streamInput: LLM.StreamInput) => Effect.Effect + readonly compactError?: () => ReturnType | undefined // kilocode_change } type Input = { @@ -73,6 +74,7 @@ interface ProcessorContext extends Input { snapshot: string | undefined blocked: boolean needsCompaction: boolean + compactionError: ReturnType | undefined // kilocode_change currentText: MessageV2.TextPart | undefined reasoningMap: Record stepStart: number // kilocode_change @@ -130,6 +132,7 @@ export const layer: Layer.Layer< snapshot: initialSnapshot, blocked: false, needsCompaction: false, + compactionError: undefined, // kilocode_change currentText: undefined, reasoningMap: {}, telemetry: input.telemetry, // kilocode_change @@ -505,6 +508,11 @@ export const layer: Layer.Layer< isOverflow({ cfg: yield* config.get(), tokens: usage.tokens, model: ctx.model }) ) { ctx.needsCompaction = true + // kilocode_change start + ctx.compactionError = new MessageV2.ContextOverflowError({ + message: "Input exceeds context window of this model", + }).toObject() + // kilocode_change end } return } @@ -635,6 +643,9 @@ export const layer: Layer.Layer< const halt = Effect.fn("SessionProcessor.halt")(function* (e: unknown) { slog.error("process", { error: errorMessage(e), stack: e instanceof Error ? e.stack : undefined }) const error = parse(e) + // kilocode_change start + ctx.compactionError = MessageV2.ContextOverflowError.isInstance(error) ? error : ctx.compactionError + // kilocode_change end if (MessageV2.ContextOverflowError.isInstance(error)) { ctx.needsCompaction = true yield* bus.publish(Session.Event.Error, { sessionID: ctx.sessionID, error }) @@ -648,9 +659,16 @@ export const layer: Layer.Layer< yield* status.set(ctx.sessionID, { type: "idle" }) }) + // kilocode_change start + const output = { + compactError: () => ctx.compactionError, + } + // kilocode_change end + const process = Effect.fn("SessionProcessor.process")(function* (streamInput: LLM.StreamInput) { slog.info("process") ctx.needsCompaction = false + ctx.compactionError = undefined // kilocode_change ctx.shouldBreak = (yield* config.get()).experimental?.continue_loop_on_deny !== true return yield* Effect.gen(function* () { @@ -700,7 +718,7 @@ export const layer: Layer.Layer< if (ctx.needsCompaction) return "compact" if (ctx.blocked || ctx.assistantMessage.error) return "stop" - return "continue" + return "continue" // kilocode_change - remove once compactError is no longer Kilo-specific }) }) @@ -710,6 +728,7 @@ export const layer: Layer.Layer< }, updateToolCall, completeToolCall, + ...output, // kilocode_change process, } satisfies Handle }) diff --git a/packages/opencode/test/kilocode/compaction-payload-recovery.test.ts b/packages/opencode/test/kilocode/compaction-payload-recovery.test.ts new file mode 100644 index 00000000000..2b415375df3 --- /dev/null +++ b/packages/opencode/test/kilocode/compaction-payload-recovery.test.ts @@ -0,0 +1,406 @@ +import { afterEach, describe, expect, mock, test } from "bun:test" +import { APICallError } from "ai" +import { Effect, Layer, ManagedRuntime } from "effect" +import * as Stream from "effect/Stream" +import { Agent } from "../../src/agent/agent" +import { Bus } from "../../src/bus" +import { Config } from "../../src/config/config" +import { KiloCompactionPayloadRecovery } from "../../src/kilocode/session/compaction-payload-recovery" +import { Permission } from "../../src/permission" +import { Plugin } from "../../src/plugin" +import { Instance } from "../../src/project/instance" +import { ModelID, ProviderID } from "../../src/provider/schema" +import { Snapshot } from "../../src/snapshot" +import { LLM } from "../../src/session/llm" +import { MessageV2 } from "../../src/session/message-v2" +import * as SessionProcessorModule from "../../src/session/processor" +import { Session as SessionNs } from "../../src/session/session" +import { MessageID, PartID, SessionID } from "../../src/session/schema" +import { SessionCompaction } from "../../src/session/compaction" +import { SessionStatus } from "../../src/session/status" +import { SessionSummary } from "../../src/session/summary" +import { ProviderTest } from "../fake/provider" +import { tmpdir } from "../fixture/fixture" + +const sessionID = SessionID.make("ses_payload_recovery") +const userID = MessageID.ascending() +const assistantID = MessageID.ascending() +const providerID = ProviderID.make("test") +const modelID = ModelID.make("test-model") + +const ref = { + providerID, + modelID, +} + +function run(fx: Effect.Effect) { + return Effect.runPromise(fx.pipe(Effect.provide(SessionNs.defaultLayer))) +} + +const svc = { + create(input?: SessionNs.CreateInput) { + return run(SessionNs.Service.use((svc) => svc.create(input))) + }, + messages(input: Parameters[0]) { + return run(SessionNs.Service.use((svc) => svc.messages(input))) + }, + updateMessage(msg: T) { + return run(SessionNs.Service.use((svc) => svc.updateMessage(msg))) + }, + updatePart(part: T) { + return run(SessionNs.Service.use((svc) => svc.updatePart(part))) + }, +} + +const summary = Layer.succeed( + SessionSummary.Service, + SessionSummary.Service.of({ + summarize: () => Effect.void, + diff: () => Effect.succeed([]), + computeDiff: () => Effect.succeed([]), + }), +) + +function base(id: MessageID) { + return { + id: PartID.ascending(), + messageID: id, + sessionID, + } +} + +async function user(sessionID: SessionID, text: string) { + const msg = await svc.updateMessage({ + id: MessageID.ascending(), + role: "user", + sessionID, + agent: "build", + model: ref, + time: { created: Date.now() }, + }) + await svc.updatePart({ + id: PartID.ascending(), + messageID: msg.id, + sessionID, + type: "text", + text, + }) + return msg +} + +async function assistant(sessionID: SessionID, parentID: MessageID, root: string) { + const msg: MessageV2.Assistant = { + id: MessageID.ascending(), + role: "assistant", + sessionID, + mode: "build", + agent: "build", + path: { cwd: root, root }, + cost: 0, + tokens: { output: 0, input: 0, reasoning: 0, cache: { read: 0, write: 0 } }, + modelID, + providerID, + parentID, + time: { created: Date.now() }, + finish: "end_turn", + } + await svc.updateMessage(msg) + return msg +} + +function llm() { + const queue: Array< + Stream.Stream | ((input: LLM.StreamInput) => Stream.Stream) + > = [] + + return { + push(stream: Stream.Stream | ((input: LLM.StreamInput) => Stream.Stream)) { + queue.push(stream) + }, + layer: Layer.succeed( + LLM.Service, + LLM.Service.of({ + stream: (input) => { + const item = queue.shift() ?? Stream.empty + const stream = typeof item === "function" ? item(input) : item + return stream.pipe(Stream.mapEffect((event) => Effect.succeed(event))) + }, + raw: () => Effect.die("raw not implemented in test LLM"), + }), + ), + } +} + +function reply( + text: string, + capture?: (input: LLM.StreamInput) => void, +): (input: LLM.StreamInput) => Stream.Stream { + return (input) => { + capture?.(input) + return Stream.make( + { type: "start" } as LLM.Event, + { type: "text-start", id: "txt-0" } as LLM.Event, + { type: "text-delta", id: "txt-0", delta: text, text } as LLM.Event, + { type: "text-end", id: "txt-0" } as LLM.Event, + { + type: "finish-step", + finishReason: "stop", + rawFinishReason: "stop", + response: { id: "res", modelId: "test-model", timestamp: new Date() }, + providerMetadata: undefined, + usage: { + inputTokens: 1, + outputTokens: 1, + totalTokens: 2, + inputTokenDetails: { noCacheTokens: undefined, cacheReadTokens: undefined, cacheWriteTokens: undefined }, + outputTokenDetails: { textTokens: undefined, reasoningTokens: undefined }, + }, + } as LLM.Event, + { + type: "finish", + finishReason: "stop", + rawFinishReason: "stop", + totalUsage: { + inputTokens: 1, + outputTokens: 1, + totalTokens: 2, + inputTokenDetails: { noCacheTokens: undefined, cacheReadTokens: undefined, cacheWriteTokens: undefined }, + outputTokenDetails: { textTokens: undefined, reasoningTokens: undefined }, + }, + } as LLM.Event, + ) + } +} + +function runtime(layer: Layer.Layer, config = Config.defaultLayer) { + const bus = Bus.layer + const status = SessionStatus.layer.pipe(Layer.provide(bus)) + const processor = SessionProcessorModule.SessionProcessor.layer.pipe(Layer.provide(summary)) + const model = ProviderTest.model({ providerID, id: modelID, limit: { context: 100_000, output: 32_000 } }) + return ManagedRuntime.make( + Layer.mergeAll(SessionCompaction.layer.pipe(Layer.provide(processor)), processor, bus, status).pipe( + Layer.provide(ProviderTest.fake({ model }).layer), + Layer.provide(SessionNs.defaultLayer), + Layer.provide(Snapshot.defaultLayer), + Layer.provide(layer), + Layer.provide(Permission.defaultLayer), + Layer.provide(Agent.defaultLayer), + Layer.provide(Plugin.defaultLayer), + Layer.provide(status), + Layer.provide(bus), + Layer.provide(config), + ), + ) +} + +afterEach(() => { + mock.restore() +}) + +describe("KiloCompactionPayloadRecovery", () => { + test("detects Kilo gateway payload-size errors", () => { + const error = new MessageV2.ContextOverflowError({ + message: "Request Entity Too Large", + responseBody: "Request Entity Too Large\n\nFUNCTION_PAYLOAD_TOO_LARGE", + }).toObject() + + expect(KiloCompactionPayloadRecovery.matches(error)).toBe(true) + }) + + test("strips media and marks completed tool outputs compacted", async () => { + const user: MessageV2.WithParts = { + info: { + id: userID, + role: "user", + sessionID, + agent: "build", + model: { providerID, modelID }, + time: { created: Date.now() }, + }, + parts: [ + { + ...base(userID), + type: "file", + mime: "image/png", + filename: "screen.png", + url: "data:image/png;base64,abc", + }, + ], + } + const tool: MessageV2.ToolPart = { + ...base(assistantID), + type: "tool", + callID: "call-1", + tool: "bash", + state: { + status: "completed", + input: {}, + output: "large output", + title: "Bash", + metadata: {}, + time: { start: Date.now(), end: Date.now() }, + }, + } + const assistant: MessageV2.WithParts = { + info: { + id: assistantID, + role: "assistant", + parentID: userID, + sessionID, + mode: "build", + agent: "build", + path: { cwd: "/tmp", root: "/tmp" }, + cost: 0, + tokens: { input: 0, output: 0, reasoning: 0, cache: { read: 0, write: 0 } }, + modelID, + providerID, + time: { created: Date.now() }, + }, + parts: [tool], + } + const updated: MessageV2.Part[] = [] + + await Effect.runPromise( + KiloCompactionPayloadRecovery.strip({ + messages: [user, assistant], + update: (part) => Effect.sync(() => updated.push(part)).pipe(Effect.as(part)), + }), + ) + + expect(updated).toHaveLength(2) + expect(updated[0]).toMatchObject({ + type: "text", + text: "[Attached image/png: screen.png]", + }) + expect(updated[1]?.type).toBe("tool") + if (updated[1]?.type === "tool" && updated[1].state.status === "completed") { + expect(updated[1].state.time.compacted).toBeNumber() + } + }) + + test("retries compaction without media and tool outputs after payload-size failure", async () => { + await using tmp = await tmpdir({ git: true }) + const stub = llm() + const captures: string[] = [] + stub.push((input) => { + captures.push(JSON.stringify(input.messages)) + return Stream.make({ type: "start" } as LLM.Event).pipe( + Stream.concat( + Stream.fail( + new APICallError({ + message: "Request Entity Too Large", + url: "https://api.kilo.ai/api/openrouter/responses", + requestBodyValues: {}, + statusCode: 413, + responseHeaders: { "content-type": "text/plain" }, + responseBody: "Request Entity Too Large\n\nFUNCTION_PAYLOAD_TOO_LARGE", + isRetryable: false, + }), + ), + ), + ) + }) + stub.push( + reply("summary", (input) => { + captures.push(JSON.stringify(input.messages)) + }), + ) + + await Instance.provide({ + directory: tmp.path, + fn: async () => { + const session = await svc.create({}) + const old = await user(session.id, "old image turn") + await svc.updatePart({ + id: PartID.ascending(), + messageID: old.id, + sessionID: session.id, + type: "file", + mime: "image/png", + filename: "old.png", + url: `data:image/png;base64,${"a".repeat(8_000)}`, + }) + const oldReply = await assistant(session.id, old.id, tmp.path) + await svc.updatePart({ + id: PartID.ascending(), + messageID: oldReply.id, + sessionID: session.id, + type: "tool", + callID: crypto.randomUUID(), + tool: "bash", + state: { + status: "completed", + input: {}, + output: "old output".repeat(10_000), + title: "old", + metadata: {}, + time: { start: Date.now(), end: Date.now() }, + }, + }) + await user(session.id, "latest turn") + const keep = await user(session.id, "preserved tail turn") + const keepReply = await assistant(session.id, keep.id, tmp.path) + await svc.updatePart({ + id: PartID.ascending(), + messageID: keepReply.id, + sessionID: session.id, + type: "tool", + callID: crypto.randomUUID(), + tool: "bash", + state: { + status: "completed", + input: {}, + output: "keep output", + title: "keep", + metadata: {}, + time: { start: Date.now(), end: Date.now() }, + }, + }) + await SessionCompaction.create({ + sessionID: session.id, + agent: "build", + model: ref, + auto: false, + }) + + const rt = runtime(stub.layer, Config.defaultLayer) + try { + const msgs = await svc.messages({ sessionID: session.id }) + const parent = msgs.at(-1)?.info.id + expect(parent).toBeTruthy() + const result = await rt.runPromise( + SessionCompaction.Service.use((svc) => + svc.process({ + parentID: parent!, + messages: msgs, + sessionID: session.id, + auto: false, + }), + ), + ) + + expect(result).toBe("continue") + expect(captures).toHaveLength(2) + expect(captures[0]).toContain("Attached image/png: old.png") + expect(captures[0]).toContain("old output") + expect(captures[0]).not.toContain("keep output") + expect(captures[1]).not.toContain("data:image/png;base64") + expect(captures[1]).not.toContain("old output") + expect(captures[1]).not.toContain("keep output") + expect(captures[1]).toContain("Attached image/png: old.png") + expect(captures[1]).toContain("Old tool result content cleared") + const tools = (await svc.messages({ sessionID: session.id })) + .flatMap((msg) => msg.parts) + .filter((part): part is MessageV2.ToolPart => part.type === "tool") + expect(tools).toHaveLength(2) + expect(tools[0]?.type).toBe("tool") + if (tools[0]?.state.status === "completed") expect(tools[0].state.time.compacted).toBeNumber() + expect(tools[1]?.type).toBe("tool") + if (tools[1]?.state.status === "completed") expect(tools[1].state.time.compacted).toBeUndefined() + } finally { + await rt.dispose() + } + }, + }) + }) +})