diff --git a/packages/opencode/src/session/processor.ts b/packages/opencode/src/session/processor.ts index 1bdc820b5..15e74557c 100644 --- a/packages/opencode/src/session/processor.ts +++ b/packages/opencode/src/session/processor.ts @@ -96,6 +96,7 @@ interface ProcessorContext extends Input { needsCompaction: boolean currentText: SessionV1.TextPart | undefined reasoningMap: Record + outputLimitUsage: Pick | undefined } type StreamEvent = LLMEvent @@ -135,6 +136,7 @@ const layer = Layer.effect( needsCompaction: false, currentText: undefined, reasoningMap: {}, + outputLimitUsage: undefined, } let aborted = false @@ -165,6 +167,26 @@ const layer = Layer.effect( return { call, part } }) + const resetOutputLimit = Effect.fn("SessionProcessor.resetOutputLimit")(function* () { + const parts = yield* MessageV2.parts(ctx.assistantMessage.id).pipe( + Effect.provideService(Database.Service, database), + ) + // Replace the streamed attempt before resampling the unchanged request. + // Its usage is carried into the next step-finish part. + yield* Effect.forEach( + parts, + (part) => + session.removePart({ + sessionID: part.sessionID, + messageID: part.messageID, + partID: part.id, + }), + { concurrency: "unbounded" }, + ) + ctx.assistantMessage.finish = undefined + yield* session.updateMessage(ctx.assistantMessage) + }) + const updateToolCall = Effect.fn("SessionProcessor.updateToolCall")(function* ( toolCallID: string, update: (part: SessionV1.ToolPart) => SessionV1.ToolPart, @@ -466,9 +488,28 @@ const layer = Layer.effect( usage: value.usage ?? new Usage({}), metadata: value.providerMetadata, }) + const previous = ctx.outputLimitUsage + const total = + previous?.tokens.total === undefined && usage.tokens.total === undefined + ? undefined + : (previous?.tokens.total ?? 0) + (usage.tokens.total ?? 0) + const accounted = { + cost: (previous?.cost ?? 0) + usage.cost, + tokens: { + ...(total === undefined ? {} : { total }), + input: (previous?.tokens.input ?? 0) + usage.tokens.input, + output: (previous?.tokens.output ?? 0) + usage.tokens.output, + reasoning: (previous?.tokens.reasoning ?? 0) + usage.tokens.reasoning, + cache: { + read: (previous?.tokens.cache.read ?? 0) + usage.tokens.cache.read, + write: (previous?.tokens.cache.write ?? 0) + usage.tokens.cache.write, + }, + }, + } + ctx.outputLimitUsage = value.reason === "length" ? accounted : undefined ctx.assistantMessage.finish = value.reason ctx.assistantMessage.cost += usage.cost - ctx.assistantMessage.tokens = usage.tokens + ctx.assistantMessage.tokens = accounted.tokens yield* session.updatePart({ id: PartID.ascending(), reason: value.reason, @@ -476,10 +517,11 @@ const layer = Layer.effect( messageID: ctx.assistantMessage.id, sessionID: ctx.assistantMessage.sessionID, type: "step-finish", - tokens: usage.tokens, - cost: usage.cost, + tokens: accounted.tokens, + cost: accounted.cost, }) yield* session.updateMessage(ctx.assistantMessage) + if (value.reason === "length") throw new SessionV1.OutputLengthError({}) if (ctx.snapshot) { const patch = yield* snapshot.patch(ctx.snapshot) if (patch.files.length) { @@ -692,6 +734,10 @@ const layer = Layer.effect( SessionRetry.policy({ provider: input.model.providerID, parse, + // Only replace attempts that will be retried. Cloud intentionally + // returns the terminal partial next to the truncation error. + onRetry: (error) => + SessionV1.OutputLengthError.isInstance(error) ? resetOutputLimit() : Effect.void, set: (info) => { return status.set(ctx.sessionID, { type: "retry", diff --git a/packages/opencode/src/session/retry.ts b/packages/opencode/src/session/retry.ts index 33c373d82..8d47b7330 100644 --- a/packages/opencode/src/session/retry.ts +++ b/packages/opencode/src/session/retry.ts @@ -23,6 +23,7 @@ export type RetryReason = "free_tier_limit" | "account_rate_limit" | (string & { export type Retryable = { message: string + maxAttempts?: number action?: { reason: RetryReason provider: string @@ -76,6 +77,9 @@ export function delay(attempt: number, error?: SessionV1.APIError) { } export function retryable(error: Err, provider: string) { + if (SessionV1.OutputLengthError.isInstance(error)) { + return { message: "Model hit its output limit", maxAttempts: 3 } + } // context overflow errors should not be retried if (SessionV1.ContextOverflowError.isInstance(error)) return undefined if (SessionV1.APIError.isInstance(error)) { @@ -186,6 +190,7 @@ function parseJSON(value: unknown) { export function policy(opts: { provider: string parse: (error: unknown) => Err + onRetry?: (error: Err) => Effect.Effect set: (input: { attempt: number; message: string; action?: Retryable["action"]; next: number }) => Effect.Effect }) { return Schedule.fromStepWithMetadata( @@ -193,7 +198,9 @@ export function policy(opts: { const error = opts.parse(meta.input) const retry = retryable(error, opts.provider) if (!retry) return Cause.done(meta.attempt) + if (retry.maxAttempts !== undefined && meta.attempt >= retry.maxAttempts) return Cause.done(meta.attempt) return Effect.gen(function* () { + if (opts.onRetry) yield* opts.onRetry(error) const wait = delay(meta.attempt, SessionV1.APIError.isInstance(error) ? error : undefined) const now = yield* Clock.currentTimeMillis yield* opts.set({ diff --git a/packages/opencode/test/session/processor-effect.test.ts b/packages/opencode/test/session/processor-effect.test.ts index 528760543..09d108bf1 100644 --- a/packages/opencode/test/session/processor-effect.test.ts +++ b/packages/opencode/test/session/processor-effect.test.ts @@ -41,6 +41,28 @@ const ref = { modelID: ModelV2.ID.make("test-model"), } +const outputRetryModel: Provider.Model = { + id: ref.modelID, + providerID: ref.providerID, + api: { id: "test-model", url: "https://example.com", npm: "@ai-sdk/openai" }, + name: "Test Model", + capabilities: { + temperature: true, + reasoning: false, + attachment: false, + toolcall: true, + input: { text: true, audio: false, image: false, video: false, pdf: false }, + output: { text: true, audio: false, image: false, video: false, pdf: false }, + interleaved: false, + }, + cost: { input: 1, output: 1, cache: { read: 0, write: 0 } }, + limit: { context: 100_000, input: 100_000, output: 10_000 }, + status: "active", + options: {}, + headers: {}, + release_date: "2026-01-01", +} + const cfg = { provider: { test: { @@ -226,6 +248,38 @@ const fragmentFailureLLM = Layer.succeed( const fragmentFailureEnv = LayerNode.compile(root, [...replacements, [LLM.node, fragmentFailureLLM]]) const itFragmentFailure = testEffect(fragmentFailureEnv) +const outputRetryInputs: LLM.StreamInput[] = [] +const outputRetryUsage = { + truncated: { input: 3, output: 5 }, + complete: { input: 7, output: 11 }, +} as const +const outputRetryLLM = Layer.succeed( + LLM.Service, + LLM.Service.of({ + stream: (input) => { + outputRetryInputs.push(input) + const first = outputRetryInputs.length === 1 + return Stream.make( + LLMEvent.stepStart({ index: 0 }), + LLMEvent.textStart({ id: "text-1" }), + LLMEvent.textDelta({ id: "text-1", text: first ? "truncated" : "complete" }), + LLMEvent.textEnd({ id: "text-1" }), + LLMEvent.stepFinish({ + index: 0, + reason: first ? "length" : "stop", + usage: { + inputTokens: first ? outputRetryUsage.truncated.input : outputRetryUsage.complete.input, + outputTokens: first ? outputRetryUsage.truncated.output : outputRetryUsage.complete.output, + }, + }), + LLMEvent.finish({ reason: first ? "length" : "stop" }), + ) + }, + }), +) +const outputRetryEnv = LayerNode.compile(root, [...replacements, [LLM.node, outputRetryLLM]]) +const itOutputRetry = testEffect(outputRetryEnv) + const boot = Effect.fn("test.boot")(function* () { const processors = yield* SessionProcessor.Service const session = yield* Session.Service @@ -514,6 +568,58 @@ it.live("session.processor effect tests reset reasoning state across retries", ( ), ) +itOutputRetry.live("session.processor effect tests resample the exact request after an output limit", () => + provideTmpdirInstance((dir) => + Effect.gen(function* () { + const { processors, session } = yield* boot() + outputRetryInputs.length = 0 + + const chat = yield* session.create({}) + const parent = yield* user(chat.id, "resample") + const msg = yield* assistant(chat.id, parent.id, path.resolve(dir)) + const handle = yield* processors.create({ + assistantMessage: msg, + sessionID: chat.id, + model: outputRetryModel, + }) + + const value = yield* handle.process({ + user: { + id: parent.id, + sessionID: chat.id, + role: "user", + time: parent.time, + agent: parent.agent, + model: { providerID: ref.providerID, modelID: ref.modelID }, + } satisfies SessionV1.User, + sessionID: chat.id, + model: outputRetryModel, + agent: agent(), + system: [], + messages: [{ role: "user", content: "resample" }], + tools: {}, + }) + + const parts = yield* MessageV2.parts(msg.id) + + expect(value).toBe("continue") + expect(outputRetryInputs).toHaveLength(2) + expect(outputRetryInputs[1]).toBe(outputRetryInputs[0]) + expect(parts.filter((part) => part.type === "text").map((part) => part.text)).toStrictEqual(["complete"]) + const finishes = parts.filter((part) => part.type === "step-finish") + const input = outputRetryUsage.truncated.input + outputRetryUsage.complete.input + const output = outputRetryUsage.truncated.output + outputRetryUsage.complete.output + expect(finishes).toHaveLength(1) + expect(finishes[0]).toMatchObject({ + reason: "stop", + tokens: { input, output }, + }) + expect(finishes[0]?.cost).toBeCloseTo((input + output) / 1_000_000) + expect(handle.message.finish).toBe("stop") + }), + ), +) + it.live("session.processor effect tests do not retry unknown json errors", () => provideTmpdirServer( ({ dir, llm }) => diff --git a/packages/opencode/test/session/retry.test.ts b/packages/opencode/test/session/retry.test.ts index ea6d596a1..cbf43c459 100644 --- a/packages/opencode/test/session/retry.test.ts +++ b/packages/opencode/test/session/retry.test.ts @@ -4,7 +4,7 @@ import { SessionV1 } from "@opencode-ai/core/v1/session" import type { NamedError } from "@opencode-ai/core/util/error" import { APICallError } from "ai" import { setTimeout as sleep } from "node:timers/promises" -import { Effect, Schedule, Schema } from "effect" +import { Effect, Exit, Schedule, Schema } from "effect" import { CrossSpawnSpawner } from "@opencode-ai/core/cross-spawn-spawner" import { SessionRetry } from "../../src/session/retry" import { MessageV2 } from "../../src/session/message-v2" @@ -115,6 +115,30 @@ describe("session.retry.delay", () => { }) }), ) + + it.effect("policy caps output-length errors at three total calls", () => + Effect.gen(function* () { + const error = new SessionV1.OutputLengthError({}).toObject() + const attempts: number[] = [] + let retries = 0 + const step = yield* Schedule.toStep( + SessionRetry.policy({ + provider: "test", + parse: () => error, + onRetry: () => Effect.sync(() => retries++), + set: (info) => Effect.sync(() => attempts.push(info.attempt)), + }), + ) + + yield* step(0, error) + yield* step(0, error) + const third = yield* step(0, error).pipe(Effect.exit) + + expect(attempts).toStrictEqual([1, 2]) + expect(retries).toBe(2) + expect(Exit.isFailure(third)).toBe(true) + }), + ) }) describe("session.retry.retryable", () => {