Skip to content

fix(core): guard non-reducing compaction #34261

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
cwj2001 wants to merge 1 commit into
base: dev
Choose a base branch
from
+295 −92
Open

fix(core): guard non-reducing compaction #34261

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
40 changes: 26 additions & 14 deletions packages/core/src/session/compaction.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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]`
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand Down
104 changes: 79 additions & 25 deletions packages/core/src/session/runner/llm.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -144,11 +144,18 @@ export const layer = Layer.effect(
const isQuestionRejected = (cause: Cause.Cause<unknown>) =>
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) {
Expand All @@ -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()], {
Expand All @@ -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)
Expand Down Expand Up @@ -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<string> = []) =>
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<string> = []) =>
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) =>
Expand Down Expand Up @@ -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()) {
Expand Down Expand Up @@ -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)
}),
),
)
Expand All @@ -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)
}),
),
Expand Down
Loading
Loading