fix(core): recover v2 context overflow
This commit is contained in:
parent
760d523847
commit
f8f648e5ce
17 changed files with 360 additions and 52 deletions
|
|
@ -165,13 +165,15 @@ export const make = (dependencies: Dependencies) => {
|
|||
readonly entries: readonly Entry[]
|
||||
readonly model: Model
|
||||
readonly request: LLMRequest
|
||||
readonly trigger?: "threshold" | "overflow"
|
||||
}) {
|
||||
const context = input.model.route.defaults.limits?.context
|
||||
if (!config.auto || context === undefined || context <= 0) return false
|
||||
if ((!config.auto && input.trigger !== "overflow") || context === undefined || context <= 0) return false
|
||||
const output = input.request.generation?.maxTokens ?? input.model.route.defaults.limits?.output ?? 0
|
||||
if (
|
||||
input.trigger !== "overflow" &&
|
||||
estimate({ system: input.request.system, messages: input.request.messages, tools: input.request.tools }) <=
|
||||
context - Math.max(output, config.buffer)
|
||||
context - Math.max(output, config.buffer)
|
||||
)
|
||||
return false
|
||||
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
import { LLM, LLMClient, LLMError, LLMEvent, SystemPart } from "@opencode-ai/llm"
|
||||
import { LLM, LLMClient, LLMError, LLMEvent, SystemPart, type ProviderErrorEvent } from "@opencode-ai/llm"
|
||||
import { Cause, DateTime, Effect, FiberSet, Layer, Schema, Semaphore, Stream } from "effect"
|
||||
import { AgentV2 } from "../../agent"
|
||||
import { Config } from "../../config"
|
||||
|
|
@ -131,7 +131,10 @@ export const layer = Layer.effect(
|
|||
cause.reasons.some((reason) => Cause.isDieReason(reason) && reason.defect instanceof QuestionV2.RejectedError)
|
||||
|
||||
class RetryTurn extends Error {
|
||||
constructor(readonly promotion: SessionInput.Delivery | undefined) {
|
||||
constructor(
|
||||
readonly promotion: SessionInput.Delivery | undefined,
|
||||
readonly consumeOverflowRetry = false,
|
||||
) {
|
||||
super()
|
||||
}
|
||||
}
|
||||
|
|
@ -149,6 +152,7 @@ export const layer = Layer.effect(
|
|||
const runTurnAttempt = Effect.fn("SessionRunner.runTurn")(function* (
|
||||
sessionID: SessionSchema.ID,
|
||||
promotion: SessionInput.Delivery | undefined,
|
||||
overflowRetryAvailable: boolean,
|
||||
) {
|
||||
const session = yield* getSession(sessionID)
|
||||
if (session.location.directory !== location.directory || session.location.workspaceID !== location.workspaceID)
|
||||
|
|
@ -208,11 +212,19 @@ export const layer = Layer.effect(
|
|||
})
|
||||
const withPublication = Semaphore.makeUnsafe(1).withPermit
|
||||
const publish = (event: LLMEvent) => withPublication(publisher.publish(event))
|
||||
let overflowFailure: ProviderErrorEvent | undefined
|
||||
if (!(yield* SessionContextEpoch.current(db, session.id, agent.id, system.revision)))
|
||||
return yield* Effect.die(new RetryTurn(undefined))
|
||||
const providerStream = llm.stream(request).pipe(
|
||||
Stream.runForEach((event) =>
|
||||
Effect.gen(function* () {
|
||||
if (overflowFailure || publisher.hasProviderError()) return
|
||||
if (event.type === "provider-error") {
|
||||
if (event.classification === "context-overflow" && !publisher.hasAssistantStarted()) {
|
||||
overflowFailure = event
|
||||
return
|
||||
}
|
||||
}
|
||||
yield* publish(event)
|
||||
if (event.type !== "tool-call" || event.providerExecuted) return
|
||||
needsContinuation = true
|
||||
|
|
@ -251,6 +263,14 @@ export const layer = Layer.effect(
|
|||
if (reason.error instanceof LLMError) llmFailure = reason.error
|
||||
}
|
||||
}
|
||||
const overflow =
|
||||
overflowFailure !== undefined ||
|
||||
(llmFailure?.reason._tag === "InvalidRequest" && llmFailure.reason.classification === "context-overflow")
|
||||
if (overflowRetryAvailable && overflow && !publisher.hasAssistantStarted()) {
|
||||
const compacted = yield* compact({ sessionID: session.id, entries, model, request, trigger: "overflow" })
|
||||
if (compacted) return yield* Effect.die(new RetryTurn(undefined, true))
|
||||
}
|
||||
if (overflowFailure) yield* withPublication(publisher.publish(overflowFailure))
|
||||
if (llmFailure && !publisher.hasProviderError()) {
|
||||
yield* withPublication(publisher.failUnsettledTools("Provider did not return a tool result", true))
|
||||
yield* withPublication(
|
||||
|
|
@ -289,11 +309,16 @@ export const layer = Layer.effect(
|
|||
const runTurn: (
|
||||
sessionID: SessionSchema.ID,
|
||||
promotion: SessionInput.Delivery | undefined,
|
||||
) => Effect.Effect<boolean, RunError> = (sessionID, promotion) =>
|
||||
runTurnAttempt(sessionID, promotion).pipe(
|
||||
overflowRetryAvailable?: boolean,
|
||||
) => Effect.Effect<boolean, RunError> = (sessionID, promotion, overflowRetryAvailable = true) =>
|
||||
runTurnAttempt(sessionID, promotion, overflowRetryAvailable).pipe(
|
||||
Effect.catchDefect((defect) =>
|
||||
defect instanceof RetryTurn
|
||||
? Effect.yieldNow.pipe(Effect.andThen(runTurn(sessionID, defect.promotion)))
|
||||
? Effect.yieldNow.pipe(
|
||||
Effect.andThen(
|
||||
runTurn(sessionID, defect.promotion, defect.consumeOverflowRetry ? false : overflowRetryAvailable),
|
||||
),
|
||||
)
|
||||
: Effect.die(defect),
|
||||
),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -165,7 +165,7 @@ export const createLLMEventPublisher = (events: EventV2.Interface, input: Input)
|
|||
|
||||
const startToolInput = Effect.fnUntraced(function* (event: { readonly id: string; readonly name: string }) {
|
||||
if (tools.has(event.id)) return yield* Effect.die(`Duplicate tool input start: ${event.id}`)
|
||||
const assistantMessageID = yield* currentAssistantMessageID()
|
||||
const assistantMessageID = yield* startAssistant()
|
||||
tools.set(event.id, {
|
||||
assistantMessageID,
|
||||
name: event.name,
|
||||
|
|
@ -221,7 +221,6 @@ export const createLLMEventPublisher = (events: EventV2.Interface, input: Input)
|
|||
const publish = Effect.fn("SessionRunner.publishLLMEvent")(function* (event: LLMEvent) {
|
||||
switch (event.type) {
|
||||
case "step-start":
|
||||
yield* startAssistant()
|
||||
return
|
||||
case "text-start":
|
||||
yield* text.start(event.id)
|
||||
|
|
@ -377,7 +376,7 @@ export const createLLMEventPublisher = (events: EventV2.Interface, input: Input)
|
|||
yield* events.publish(SessionEvent.Step.Ended, {
|
||||
sessionID: input.sessionID,
|
||||
timestamp: yield* timestamp,
|
||||
assistantMessageID: yield* currentAssistantMessageID(),
|
||||
assistantMessageID: yield* startAssistant(),
|
||||
finish: event.reason,
|
||||
cost: 0,
|
||||
tokens: tokens(event.usage),
|
||||
|
|
@ -398,5 +397,12 @@ export const createLLMEventPublisher = (events: EventV2.Interface, input: Input)
|
|||
}
|
||||
})
|
||||
|
||||
return { publish, flush, failUnsettledTools, hasProviderError: () => providerFailed, startAssistant }
|
||||
return {
|
||||
publish,
|
||||
flush,
|
||||
failUnsettledTools,
|
||||
hasAssistantStarted: () => assistantMessageID !== undefined,
|
||||
hasProviderError: () => providerFailed,
|
||||
startAssistant,
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ import {
|
|||
Model,
|
||||
Tool,
|
||||
TransportReason,
|
||||
InvalidRequestReason,
|
||||
type LLMClientShape,
|
||||
type LLMRequest,
|
||||
} from "@opencode-ai/llm"
|
||||
|
|
@ -103,6 +104,11 @@ const compactModel = Model.make({
|
|||
provider: "fake",
|
||||
route: OpenAIChat.route.with({ limits: { context: 4_000, output: 50 } }),
|
||||
})
|
||||
const recoveryModel = Model.make({
|
||||
id: "recovery",
|
||||
provider: "fake",
|
||||
route: OpenAIChat.route.with({ limits: { context: 20_000, output: 1_000 } }),
|
||||
})
|
||||
const authorizations: ToolRegistry.AuthorizeInput[] = []
|
||||
const executions: string[] = []
|
||||
const permission = Layer.succeed(
|
||||
|
|
@ -1456,6 +1462,119 @@ describe("SessionRunnerLLM", () => {
|
|||
}),
|
||||
)
|
||||
|
||||
it.effect("forces one compaction and retries after provider context overflow", () =>
|
||||
Effect.gen(function* () {
|
||||
yield* setup
|
||||
const session = yield* SessionV2.Service
|
||||
response = fragmentFixture("text", "text-earlier", ["Earlier answer"]).completeEvents
|
||||
yield* session.prompt({
|
||||
sessionID,
|
||||
prompt: new Prompt({ text: "Earlier question ".repeat(700) }),
|
||||
resume: false,
|
||||
})
|
||||
yield* session.resume(sessionID)
|
||||
|
||||
currentModel = recoveryModel
|
||||
requests.length = 0
|
||||
responses = [
|
||||
[
|
||||
LLMEvent.stepStart({ index: 0 }),
|
||||
LLMEvent.providerError({ message: "prompt too long", classification: "context-overflow" }),
|
||||
],
|
||||
fragmentFixture("text", "text-summary", ["## Goal\n- Recover overflow"]).completeEvents,
|
||||
fragmentFixture("text", "text-final", ["Recovered"]).completeEvents,
|
||||
]
|
||||
yield* session.prompt({ sessionID, prompt: new Prompt({ text: "Continue" }), resume: false })
|
||||
yield* session.resume(sessionID)
|
||||
|
||||
expect(requests).toHaveLength(3)
|
||||
expect(userTexts(requests[1])[0]).toContain("## Goal")
|
||||
expect(userTexts(requests[2])[0]).toContain("<summary>\n## Goal\n- Recover overflow\n</summary>")
|
||||
expect(yield* session.context(sessionID)).toMatchObject([
|
||||
{ type: "compaction", summary: "## Goal\n- Recover overflow" },
|
||||
{ type: "assistant", finish: "stop" },
|
||||
])
|
||||
yield* replaySessionProjection(sessionID)
|
||||
expect(yield* session.context(sessionID)).toMatchObject([
|
||||
{ type: "compaction" },
|
||||
{ type: "assistant", finish: "stop" },
|
||||
])
|
||||
}),
|
||||
)
|
||||
|
||||
it.effect("persists a second context overflow after one recovery", () =>
|
||||
Effect.gen(function* () {
|
||||
yield* setup
|
||||
const session = yield* SessionV2.Service
|
||||
response = fragmentFixture("text", "text-earlier", ["Earlier answer"]).completeEvents
|
||||
yield* session.prompt({
|
||||
sessionID,
|
||||
prompt: new Prompt({ text: "Earlier question ".repeat(700) }),
|
||||
resume: false,
|
||||
})
|
||||
yield* session.resume(sessionID)
|
||||
|
||||
currentModel = recoveryModel
|
||||
requests.length = 0
|
||||
const overflow = () => [
|
||||
LLMEvent.stepStart({ index: 0 }),
|
||||
LLMEvent.providerError({ message: "prompt too long", classification: "context-overflow" }),
|
||||
]
|
||||
responses = [
|
||||
overflow(),
|
||||
fragmentFixture("text", "text-summary", ["## Goal\n- Recover once"]).completeEvents,
|
||||
overflow(),
|
||||
]
|
||||
yield* session.prompt({ sessionID, prompt: new Prompt({ text: "Continue" }), resume: false })
|
||||
yield* session.resume(sessionID)
|
||||
|
||||
expect(requests).toHaveLength(3)
|
||||
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* () {
|
||||
yield* setup
|
||||
const session = yield* SessionV2.Service
|
||||
response = fragmentFixture("text", "text-earlier", ["Earlier answer"]).completeEvents
|
||||
yield* session.prompt({
|
||||
sessionID,
|
||||
prompt: new Prompt({ text: "Earlier question ".repeat(700) }),
|
||||
resume: false,
|
||||
})
|
||||
yield* session.resume(sessionID)
|
||||
|
||||
currentModel = recoveryModel
|
||||
requests.length = 0
|
||||
responseStream = Stream.fail(
|
||||
new LLMError({
|
||||
module: "test",
|
||||
method: "stream",
|
||||
reason: new InvalidRequestReason({
|
||||
message: "prompt too long",
|
||||
classification: "context-overflow",
|
||||
}),
|
||||
}),
|
||||
)
|
||||
responses = [
|
||||
fragmentFixture("text", "text-summary", ["## Goal\n- Recover raw overflow"]).completeEvents,
|
||||
fragmentFixture("text", "text-final", ["Recovered"]).completeEvents,
|
||||
]
|
||||
yield* session.prompt({ sessionID, prompt: new Prompt({ text: "Continue" }), resume: false })
|
||||
yield* session.resume(sessionID)
|
||||
|
||||
expect(requests).toHaveLength(3)
|
||||
expect(yield* session.context(sessionID)).toMatchObject([
|
||||
{ type: "compaction", summary: "## Goal\n- Recover raw overflow" },
|
||||
{ type: "assistant", finish: "stop" },
|
||||
])
|
||||
}),
|
||||
)
|
||||
|
||||
it.effect("preserves effective System updates while compaction replacement is blocked", () =>
|
||||
Effect.gen(function* () {
|
||||
yield* setup
|
||||
|
|
@ -3108,6 +3227,35 @@ describe("SessionRunnerLLM", () => {
|
|||
}),
|
||||
)
|
||||
|
||||
it.effect("does not recover context overflow after durable assistant output", () =>
|
||||
Effect.gen(function* () {
|
||||
yield* setup
|
||||
const session = yield* SessionV2.Service
|
||||
yield* session.prompt({ sessionID, prompt: new Prompt({ text: "Fail after output" }), resume: false })
|
||||
|
||||
requests.length = 0
|
||||
response = [
|
||||
LLMEvent.stepStart({ index: 0 }),
|
||||
LLMEvent.textStart({ id: "text-partial" }),
|
||||
LLMEvent.textDelta({ id: "text-partial", text: "Partial" }),
|
||||
LLMEvent.textEnd({ id: "text-partial" }),
|
||||
LLMEvent.providerError({ message: "prompt too long", classification: "context-overflow" }),
|
||||
]
|
||||
yield* session.resume(sessionID)
|
||||
|
||||
expect(requests).toHaveLength(1)
|
||||
expect(yield* session.context(sessionID)).toMatchObject([
|
||||
{ type: "user", text: "Fail after output" },
|
||||
{
|
||||
type: "assistant",
|
||||
finish: "error",
|
||||
error: { message: "prompt too long" },
|
||||
content: [{ type: "text", text: "Partial" }],
|
||||
},
|
||||
])
|
||||
}),
|
||||
)
|
||||
|
||||
it.effect("projects raw provider stream failures as terminal assistant step failures", () =>
|
||||
Effect.gen(function* () {
|
||||
yield* setup
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue