Merge pull request #137 from browser-use/retry-output-limit

fix(opencode): retry output-limit responses

Author: MagMueller

packages/opencode/src/session/processor.ts

@@ -96,6 +96,7 @@ interface ProcessorContext extends Input {
   needsCompaction: boolean
   currentText: SessionV1.TextPart | undefined
   reasoningMap: Record<string, SessionV1.ReasoningPart>
+  outputLimitUsage: Pick<SessionV1.StepFinishPart, "cost" | "tokens"> | 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,20 +488,40 @@ 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,
               snapshot: completedSnapshot,
               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",

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,14 +190,17 @@ function parseJSON(value: unknown) {
 export function policy(opts: {
   provider: string
   parse: (error: unknown) => Err
+  onRetry?: (error: Err) => Effect.Effect<void>
   set: (input: { attempt: number; message: string; action?: Retryable["action"]; next: number }) => Effect.Effect<void>
 }) {
   return Schedule.fromStepWithMetadata(
     Effect.succeed((meta: Schedule.InputMetadata<unknown>) => {
       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({

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 }) =>

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", () => {