From a2cd5f4960bb51ba16d737d8cacbfb13ccb78ae4 Mon Sep 17 00:00:00 2001 From: "weijia2.chen" Date: Sun, 28 Jun 2026 07:26:28 +0800 Subject: [PATCH] fix(core): guard non-reducing compaction --- packages/core/src/session/compaction.ts | 40 ++-- packages/core/src/session/runner/llm.ts | 104 ++++++--- packages/core/test/session-runner.test.ts | 243 +++++++++++++++++----- 3 files changed, 295 insertions(+), 92 deletions(-) diff --git a/packages/core/src/session/compaction.ts b/packages/core/src/session/compaction.ts index 5229949cb958..4d3bd0d1127e 100644 --- a/packages/core/src/session/compaction.ts +++ b/packages/core/src/session/compaction.ts @@ -77,6 +77,14 @@ type Input = { } const estimate = (value: unknown) => Token.estimate(JSON.stringify(value)) +export const estimateRequest = (request: LLMRequest) => + estimate({ system: request.system, messages: request.messages, tools: request.tools }) + +export type OverflowCompactionResult = + | { readonly compacted: false } + | { readonly compacted: true; readonly beforeTokens: number; readonly budget: number } + +const notCompacted: OverflowCompactionResult = { compacted: false } const truncate = (value: string) => value.length <= TOOL_OUTPUT_MAX_CHARS ? value : `${value.slice(0, TOOL_OUTPUT_MAX_CHARS)}\n[truncated]` @@ -174,19 +182,27 @@ export const buildPrompt = (input: { readonly previousSummary?: string; readonly export const make = (dependencies: Dependencies) => { const config = settings(dependencies.config) + const budget = (input: Input) => { + const context = input.model.route.defaults.limits?.context + if (context === undefined || context <= 0) return + const output = input.request.generation?.maxTokens ?? input.model.route.defaults.limits?.output ?? 0 + return context - Math.max(output, config.buffer) + } const compactAfterOverflow = Effect.fn("SessionCompaction.compactAfterOverflow")(function* (input: Input) { const context = input.model.route.defaults.limits?.context - if (context === undefined || context <= 0) return false + const requestBudget = budget(input) + if (context === undefined || context <= 0 || requestBudget === undefined) return notCompacted const output = input.request.generation?.maxTokens ?? input.model.route.defaults.limits?.output ?? 0 + const beforeTokens = estimateRequest(input.request) const selected = select(input.entries, config.tokens) const previousSummary = input.entries.find((entry) => entry.message.type === "compaction")?.message - if (!selected || (selected.head.length === 0 && previousSummary?.type !== "compaction")) return false + if (!selected || (selected.head.length === 0 && previousSummary?.type !== "compaction")) return notCompacted const summaryPrompt = buildPrompt({ previousSummary: previousSummary?.type === "compaction" ? previousSummary.summary : undefined, context: [previousSummary?.type === "compaction" ? previousSummary.recent : "", selected.head].filter(Boolean), }) const summaryOutput = Math.min(output || SUMMARY_OUTPUT_TOKENS, SUMMARY_OUTPUT_TOKENS) - if (Token.estimate(summaryPrompt) > context - summaryOutput) return false + if (Token.estimate(summaryPrompt) > context - summaryOutput) return notCompacted const messageID = SessionMessage.ID.create() yield* dependencies.events.publish(SessionEvent.Compaction.Started, { sessionID: input.sessionID, @@ -216,7 +232,7 @@ export const make = (dependencies: Dependencies) => { Effect.catchTag("LLM.Error", () => Effect.succeed(false)), ) const summary = chunks.join("") - if (!summarized || failed || !summary.trim()) return false + if (!summarized || failed || !summary.trim()) return notCompacted yield* dependencies.events.publish(SessionEvent.Compaction.Ended, { sessionID: input.sessionID, messageID, @@ -225,19 +241,15 @@ export const make = (dependencies: Dependencies) => { text: summary, recent: selected.recent, }) - return true + const result: OverflowCompactionResult = { compacted: true, beforeTokens, budget: requestBudget } + return result }) const compactIfNeeded = Effect.fn("SessionCompaction.compactIfNeeded")(function* (input: Input) { if (!config.auto) return false - const context = input.model.route.defaults.limits?.context - if (context === undefined || context <= 0) return false - const output = input.request.generation?.maxTokens ?? input.model.route.defaults.limits?.output ?? 0 - if ( - estimate({ system: input.request.system, messages: input.request.messages, tools: input.request.tools }) <= - context - Math.max(output, config.buffer) - ) - return false - return yield* compactAfterOverflow(input) + const requestBudget = budget(input) + if (requestBudget === undefined) return false + if (estimateRequest(input.request) <= requestBudget) return false + return (yield* compactAfterOverflow(input)).compacted }) return { compactIfNeeded, diff --git a/packages/core/src/session/runner/llm.ts b/packages/core/src/session/runner/llm.ts index 7983c5e9c9b6..ce041d305ccf 100644 --- a/packages/core/src/session/runner/llm.ts +++ b/packages/core/src/session/runner/llm.ts @@ -144,11 +144,18 @@ export const layer = Layer.effect( const isQuestionRejected = (cause: Cause.Cause) => cause.reasons.some((reason) => Cause.isDieReason(reason) && reason.defect instanceof QuestionV2.RejectedError) + type OverflowRecoveryError = ProviderErrorEvent | LLMError + type OverflowRecoveryGuard = { + readonly error: OverflowRecoveryError + readonly beforeTokens: number + readonly budget: number + } + type TurnTransition = // Automatic compaction completed; rebuild the request from compacted history. | { readonly _tag: "ContinueAfterCompaction"; readonly step: number } // Overflow compaction completed; rebuild once through the path without overflow recovery. - | { readonly _tag: "ContinueAfterOverflowCompaction"; readonly step: number } + | { readonly _tag: "ContinueAfterOverflowCompaction"; readonly step: number; readonly guard: OverflowRecoveryGuard } class TurnTransitionError extends Error { constructor(readonly transition: TurnTransition) { @@ -157,8 +164,8 @@ export const layer = Layer.effect( } const continueAfterCompaction = (step: number) => new TurnTransitionError({ _tag: "ContinueAfterCompaction", step }) - const continueAfterOverflowCompaction = (step: number) => - new TurnTransitionError({ _tag: "ContinueAfterOverflowCompaction", step }) + const continueAfterOverflowCompaction = (step: number, guard: OverflowRecoveryGuard) => + new TurnTransitionError({ _tag: "ContinueAfterOverflowCompaction", step, guard }) const loadSystemContext = (agent: AgentV2.Selection) => Effect.all([systemContext.load(), skillGuidance.load(agent), referenceGuidance.load()], { @@ -170,6 +177,7 @@ export const layer = Layer.effect( promotion: SessionInput.Delivery | undefined, step: number, recoverOverflow?: typeof compaction.compactAfterOverflow, + overflowGuard?: OverflowRecoveryGuard, ) { const session = yield* getSession(sessionID) if (session.location.directory !== location.directory || session.location.workspaceID !== location.workspaceID) @@ -207,22 +215,43 @@ export const layer = Layer.effect( tools: toolMaterialization?.definitions ?? [], toolChoice: isLastStep ? "none" : undefined, }) + const makePublisher = Effect.fnUntraced(function* () { + const startSnapshot = yield* snapshots.capture() + const publisher = createLLMEventPublisher(events, { + sessionID: session.id, + agent: agent.id, + model: { + id: ModelV2.ID.make(model.id), + providerID: ProviderV2.ID.make(model.provider), + ...(session.model?.variant === undefined ? {} : { variant: session.model.variant }), + }, + snapshot: startSnapshot, + }) + const withPublication = Semaphore.makeUnsafe(1).withPermit + const publish = (event: LLMEvent, outputPaths: ReadonlyArray = []) => + withPublication(publisher.publish(event, outputPaths)) + return { startSnapshot, publisher, withPublication, publish } + }) + if (overflowGuard) { + const afterTokens = SessionCompaction.estimateRequest(request) + if (afterTokens >= overflowGuard.beforeTokens || afterTokens > overflowGuard.budget) { + const { publisher, withPublication, publish } = yield* makePublisher() + if (overflowGuard.error instanceof LLMError) { + yield* withPublication(publisher.failUnsettledTools("Provider did not return a tool result", true)) + yield* withPublication(publisher.failAssistant(overflowGuard.error.reason.message)) + yield* withPublication(publisher.flush()) + return yield* Effect.fail(overflowGuard.error) + } else { + yield* publish(overflowGuard.error) + yield* withPublication(publisher.failUnsettledTools("Tool execution interrupted")) + } + yield* withPublication(publisher.flush()) + return { needsContinuation: false, step: currentStep } + } + } if (yield* compaction.compactIfNeeded({ sessionID: session.id, entries, model, request })) return yield* Effect.die(continueAfterCompaction(currentStep)) - const startSnapshot = yield* snapshots.capture() - const publisher = createLLMEventPublisher(events, { - sessionID: session.id, - agent: agent.id, - model: { - id: ModelV2.ID.make(model.id), - providerID: ProviderV2.ID.make(model.provider), - ...(session.model?.variant === undefined ? {} : { variant: session.model.variant }), - }, - snapshot: startSnapshot, - }) - const withPublication = Semaphore.makeUnsafe(1).withPermit - const publish = (event: LLMEvent, outputPaths: ReadonlyArray = []) => - withPublication(publisher.publish(event, outputPaths)) + const { startSnapshot, publisher, withPublication, publish } = yield* makePublisher() let overflowFailure: ProviderErrorEvent | undefined const providerStream = llm.stream(request).pipe( Stream.runForEach((event) => @@ -277,10 +306,19 @@ export const layer = Layer.effect( if ( recoverOverflow && !publisher.hasAssistantStarted() && - isContextOverflowFailure(overflowFailure ?? failure) && - (yield* restore(recoverOverflow({ sessionID: session.id, entries, model, request }))) - ) - return yield* Effect.die(continueAfterOverflowCompaction(currentStep)) + isContextOverflowFailure(overflowFailure ?? failure) + ) { + const recovery = yield* restore(recoverOverflow({ sessionID: session.id, entries, model, request })) + const overflowError = overflowFailure ?? failure + if (recovery.compacted && (overflowError instanceof LLMError || LLMEvent.is.providerError(overflowError))) + return yield* Effect.die( + continueAfterOverflowCompaction(currentStep, { + error: overflowError, + beforeTokens: recovery.beforeTokens, + budget: recovery.budget, + }), + ) + } if (overflowFailure) yield* publish(overflowFailure) const llmFailure = failure instanceof LLMError ? failure : undefined if (llmFailure && !publisher.hasProviderError()) { @@ -346,16 +384,27 @@ export const layer = Layer.effect( promotion: SessionInput.Delivery | undefined, step: number, ) => Effect.Effect<{ readonly needsContinuation: boolean; readonly step: number }, RunError> + type RunTurnAfterOverflow = ( + sessionID: SessionSchema.ID, + promotion: SessionInput.Delivery | undefined, + step: number, + overflowGuard: OverflowRecoveryGuard, + ) => Effect.Effect<{ readonly needsContinuation: boolean; readonly step: number }, RunError> - const runAfterOverflowCompaction: RunTurn = Effect.fnUntraced(function* (sessionID, promotion, step) { - return yield* runTurnAttempt(sessionID, promotion, step).pipe( + const runAfterOverflowCompaction: RunTurnAfterOverflow = Effect.fnUntraced(function* ( + sessionID: SessionSchema.ID, + promotion: SessionInput.Delivery | undefined, + step: number, + overflowGuard: OverflowRecoveryGuard, + ) { + return yield* runTurnAttempt(sessionID, promotion, step, undefined, overflowGuard).pipe( Effect.catchDefect( Effect.fnUntraced(function* (defect) { if (!(defect instanceof TurnTransitionError)) return yield* Effect.die(defect) if (defect.transition._tag === "ContinueAfterOverflowCompaction") return yield* Effect.die("Post-compaction provider attempt cannot recover another overflow") yield* Effect.yieldNow - return yield* runAfterOverflowCompaction(sessionID, undefined, defect.transition.step) + return yield* runAfterOverflowCompaction(sessionID, undefined, defect.transition.step, overflowGuard) }), ), ) @@ -368,7 +417,12 @@ export const layer = Layer.effect( if (!(defect instanceof TurnTransitionError)) return yield* Effect.die(defect) yield* Effect.yieldNow if (defect.transition._tag === "ContinueAfterOverflowCompaction") - return yield* runAfterOverflowCompaction(sessionID, undefined, defect.transition.step) + return yield* runAfterOverflowCompaction( + sessionID, + undefined, + defect.transition.step, + defect.transition.guard, + ) return yield* runTurn(sessionID, undefined, defect.transition.step) }), ), diff --git a/packages/core/test/session-runner.test.ts b/packages/core/test/session-runner.test.ts index 2cb286804c06..153f6fee7cf9 100644 --- a/packages/core/test/session-runner.test.ts +++ b/packages/core/test/session-runner.test.ts @@ -53,6 +53,7 @@ import { ReferenceGuidance } from "@opencode-ai/core/reference/guidance" import { ModelV2 } from "@opencode-ai/core/model" import { Location } from "@opencode-ai/core/location" import { ProviderV2 } from "@opencode-ai/core/provider" +import { Token } from "@opencode-ai/core/util/token" import { Cause, DateTime, Deferred, Effect, Exit, Fiber, Layer, Schema, Stream } from "effect" import { asc, eq } from "drizzle-orm" import { testEffect } from "./lib/effect" @@ -107,6 +108,16 @@ const recoveryModel = Model.make({ provider: "fake", route: OpenAIChat.route.with({ limits: { context: 20_000, output: 1_000 } }), }) +const wideRecoveryModel = Model.make({ + id: "wide-recovery", + provider: "fake", + route: OpenAIChat.route.with({ limits: { context: 100_000, output: 1_000 } }), +}) +const narrowBudgetRecoveryModel = Model.make({ + id: "narrow-budget-recovery", + provider: "fake", + route: OpenAIChat.route.with({ limits: { context: 50_000, output: 47_000 } }), +}) const authorizations: Tool.Context[] = [] const executions: string[] = [] const permission = Layer.succeed( @@ -214,61 +225,74 @@ const skillGuidance = Layer.mock(SkillGuidance.Service, { ), }) const referenceGuidance = Layer.mock(ReferenceGuidance.Service, { load: () => Effect.succeed(SystemContext.empty) }) -const config = Layer.succeed( - Config.Service, - Config.Service.of({ - entries: () => - Effect.succeed([ - new Config.Document({ - type: "document", - info: new Config.Info({ - compaction: new ConfigCompaction.Info({ - buffer: 3_000, - keep: new ConfigCompaction.Keep({ tokens: 1_000 }), +const configLayer = (auto?: boolean) => + Layer.succeed( + Config.Service, + Config.Service.of({ + entries: () => + Effect.succeed([ + new Config.Document({ + type: "document", + info: new Config.Info({ + compaction: new ConfigCompaction.Info({ + ...(auto === undefined ? {} : { auto }), + buffer: 3_000, + keep: new ConfigCompaction.Keep({ tokens: 1_000 }), + }), }), }), - }), - ]), - }), -) -const runner = SessionRunnerLLM.layer.pipe( - Layer.provide(Snapshot.noopLayer), - Layer.provide(Database.defaultLayer), - Layer.provide(SessionStore.defaultLayer), - Layer.provide(EventV2.defaultLayer), - Layer.provide(client), - Layer.provide(registry), - Layer.provide(models), - Layer.provide(systemContext), - Layer.provide(location), - Layer.provide(agents), - Layer.provide(skillGuidance), - Layer.provide(referenceGuidance), - Layer.provide(config), -) -const execution = Layer.effect( - SessionExecution.Service, - Effect.gen(function* () { - const sessionRunner = yield* SessionRunner.Service - const coordinator = yield* SessionRunCoordinator.make({ - drain: (sessionID, force) => sessionRunner.run({ sessionID, force }), - }) - return SessionExecution.Service.of({ - active: coordinator.active, - resume: coordinator.run, - wake: coordinator.wake, - interrupt: coordinator.interrupt, - }) - }), -).pipe(Layer.provide(runner)) -const sessions = SessionV2.layer.pipe( - Layer.provide(locationServiceMapLayer), - Layer.provide(EventV2.defaultLayer), - Layer.provide(Database.defaultLayer), - Layer.provide(SessionStore.defaultLayer), - Layer.provide(Project.defaultLayer), - Layer.provide(execution), -) + ]), + }), + ) +const config = configLayer() +const configNoAuto = configLayer(false) +const runnerLayer = (config: typeof configNoAuto) => + SessionRunnerLLM.layer.pipe( + Layer.provide(Snapshot.noopLayer), + Layer.provide(Database.defaultLayer), + Layer.provide(SessionStore.defaultLayer), + Layer.provide(EventV2.defaultLayer), + Layer.provide(client), + Layer.provide(registry), + Layer.provide(models), + Layer.provide(systemContext), + Layer.provide(location), + Layer.provide(agents), + Layer.provide(skillGuidance), + Layer.provide(referenceGuidance), + Layer.provide(config), + ) +const runner = runnerLayer(config) +const runnerNoAuto = runnerLayer(configNoAuto) +const executionLayer = (runner: typeof runnerNoAuto) => + Layer.effect( + SessionExecution.Service, + Effect.gen(function* () { + const sessionRunner = yield* SessionRunner.Service + const coordinator = yield* SessionRunCoordinator.make({ + drain: (sessionID, force) => sessionRunner.run({ sessionID, force }), + }) + return SessionExecution.Service.of({ + active: coordinator.active, + resume: coordinator.run, + wake: coordinator.wake, + interrupt: coordinator.interrupt, + }) + }), + ).pipe(Layer.provide(runner)) +const execution = executionLayer(runner) +const executionNoAuto = executionLayer(runnerNoAuto) +const sessionsLayer = (execution: typeof executionNoAuto) => + SessionV2.layer.pipe( + Layer.provide(locationServiceMapLayer), + Layer.provide(EventV2.defaultLayer), + Layer.provide(Database.defaultLayer), + Layer.provide(SessionStore.defaultLayer), + Layer.provide(Project.defaultLayer), + Layer.provide(execution), + ) +const sessions = sessionsLayer(execution) +const sessionsNoAuto = sessionsLayer(executionNoAuto) const it = testEffect( Layer.mergeAll( Database.defaultLayer, @@ -292,6 +316,29 @@ const it = testEffect( sessions, ), ) +const itNoAuto = testEffect( + Layer.mergeAll( + Database.defaultLayer, + EventV2.defaultLayer, + questions, + SessionProjector.defaultLayer, + SessionStore.defaultLayer, + client, + permission, + applications, + agents, + registry, + echo, + models, + systemContext, + location, + skillGuidance, + configNoAuto, + runnerNoAuto, + executionNoAuto, + sessionsNoAuto, + ), +) const sessionID = SessionV2.ID.make("ses_runner_test") const otherSessionID = SessionV2.ID.make("ses_runner_other") @@ -363,6 +410,19 @@ const setupOverflowRecovery = Effect.gen(function* () { requests.length = 0 return session }) +const setupLargeOverflowRecovery = Effect.gen(function* () { + yield* setup + const session = yield* SessionV2.Service + response = fragmentFixture("text", "text-earlier", ["Earlier answer"]).completeEvents + yield* session.prompt({ + sessionID, + prompt: Prompt.make({ text: "Earlier question ".repeat(5_000) }), + resume: false, + }) + yield* session.resume(sessionID) + requests.length = 0 + return session +}) const messageTexts = (request: LLMRequest, role: "user" | "system") => request.messages.flatMap((message) => @@ -370,6 +430,8 @@ const messageTexts = (request: LLMRequest, role: "user" | "system") => ) const userTexts = (request: LLMRequest) => messageTexts(request, "user") const systemTexts = (request: LLMRequest) => messageTexts(request, "system") +const estimateRequestTokens = (request: LLMRequest) => + Token.estimate(JSON.stringify({ system: request.system, messages: request.messages, tools: request.tools })) const replaySessionProjection = (id: SessionV2.ID) => Effect.gen(function* () { @@ -1201,6 +1263,81 @@ describe("SessionRunnerLLM", () => { }), ) + itNoAuto.effect("stops overflow recovery when compacted request remains over budget", () => + Effect.gen(function* () { + const session = yield* setupLargeOverflowRecovery + currentModel = narrowBudgetRecoveryModel + const mediumSummary = `## Goal\n- ${"Recovered context ".repeat(750)}` + const summaryTokens = Token.estimate(JSON.stringify(mediumSummary)) + expect(summaryTokens).toBeGreaterThan(3_000) + responses = [ + [LLMEvent.providerError({ message: "prompt too long", classification: "context-overflow" })], + fragmentFixture("text", "text-summary", [mediumSummary]).completeEvents, + fragmentFixture("text", "text-final", ["Should not retry"]).completeEvents, + ] + yield* session.prompt({ sessionID, prompt: Prompt.make({ text: "Continue" }), resume: false }) + yield* session.resume(sessionID) + + expect(requests).toHaveLength(2) + expect(estimateRequestTokens(requests[0])).toBeGreaterThan(summaryTokens + 10_000) + expect(yield* session.context(sessionID)).toMatchObject([ + { type: "compaction" }, + { type: "assistant", finish: "error", error: { message: "prompt too long" } }, + ]) + }), + ) + + it.effect("stops overflow recovery when compaction does not reduce request tokens", () => + Effect.gen(function* () { + const session = yield* setupOverflowRecovery + currentModel = wideRecoveryModel + const nonReducingSummary = `## Goal\n- ${"Recovered context ".repeat(5_000)}` + const summaryTokens = Token.estimate(JSON.stringify(nonReducingSummary)) + expect(summaryTokens).toBeLessThan(97_000) + responses = [ + [LLMEvent.providerError({ message: "prompt too long", classification: "context-overflow" })], + fragmentFixture("text", "text-summary", [nonReducingSummary]).completeEvents, + fragmentFixture("text", "text-final", ["Should not retry"]).completeEvents, + ] + yield* session.prompt({ sessionID, prompt: Prompt.make({ text: "Continue" }), resume: false }) + yield* session.resume(sessionID) + + expect(requests).toHaveLength(2) + expect(summaryTokens).toBeGreaterThan(estimateRequestTokens(requests[0])) + expect(yield* session.context(sessionID)).toMatchObject([ + { type: "compaction" }, + { type: "assistant", finish: "error", error: { message: "prompt too long" } }, + ]) + }), + ) + + itNoAuto.effect("fails raw overflow when budget guard stops recovery", () => + Effect.gen(function* () { + const session = yield* setupLargeOverflowRecovery + currentModel = narrowBudgetRecoveryModel + const failure = new LLMError({ + module: "test", + method: "stream", + reason: new InvalidRequestReason({ + message: "prompt too long", + classification: "context-overflow", + }), + }) + const mediumSummary = `## Goal\n- ${"Recovered context ".repeat(750)}` + responseStream = Stream.fail(failure) + responses = [fragmentFixture("text", "text-summary", [mediumSummary]).completeEvents] + yield* session.prompt({ sessionID, prompt: Prompt.make({ text: "Continue" }), resume: false }) + + expect(yield* session.resume(sessionID).pipe(Effect.flip)).toBe(failure) + expect(requests).toHaveLength(2) + yield* replaySessionProjection(sessionID) + expect(yield* session.context(sessionID)).toMatchObject([ + { type: "compaction" }, + { type: "assistant", finish: "error", error: { message: "prompt too long" } }, + ]) + }), + ) + it.effect("recovers once from a raw context overflow failure", () => Effect.gen(function* () { const session = yield* setupOverflowRecovery