Pārlūkot izejas kodu

fix: adjust tui retry dialog logic to be more provider specific and error case specific (#26366)

Aiden Cline 3 mēneši atpakaļ
vecāks
revīzija
799996db76

+ 29 - 6
packages/opencode/src/cli/cmd/tui/routes/session/index.tsx

@@ -93,9 +93,29 @@ import { useBindings, useCommandShortcut } from "../../keymap"
 
 addDefaultParsers(parsers.parsers)
 
-const GO_UPSELL_LAST_SEEN_AT = "go_upsell_last_seen_at"
-const GO_UPSELL_DONT_SHOW = "go_upsell_dont_show"
+const GO_UPSELL_FREE_TIER_LAST_SEEN_AT = "go_upsell_last_seen_at"
+const GO_UPSELL_FREE_TIER_DONT_SHOW = "go_upsell_dont_show"
+const GO_UPSELL_ACCOUNT_RATE_LIMIT_LAST_SEEN_AT = "go_upsell_account_rate_limit_last_seen_at"
+const GO_UPSELL_ACCOUNT_RATE_LIMIT_DONT_SHOW = "go_upsell_account_rate_limit_dont_show"
 const GO_UPSELL_WINDOW = 86_400_000 // 24 hrs
+const GO_UPSELL_PROVIDERS = new Set(["opencode", "opencode-go"])
+
+function goUpsellKeys(action: SessionRetry.Retryable["action"]) {
+  if (!action) return
+  if (!GO_UPSELL_PROVIDERS.has(action.provider)) return
+  if (action.reason === "free_tier_limit") {
+    return {
+      lastSeenAt: GO_UPSELL_FREE_TIER_LAST_SEEN_AT,
+      dontShow: GO_UPSELL_FREE_TIER_DONT_SHOW,
+    }
+  }
+  if (action.reason === "account_rate_limit") {
+    return {
+      lastSeenAt: GO_UPSELL_ACCOUNT_RATE_LIMIT_LAST_SEEN_AT,
+      dontShow: GO_UPSELL_ACCOUNT_RATE_LIMIT_DONT_SHOW,
+    }
+  }
+}
 
 const context = createContext<{
   width: number
@@ -263,14 +283,17 @@ export function Session() {
     if (!evt.properties.status.action) return
     if (dialog.stack.length > 0) return
 
-    const seen = kv.get(GO_UPSELL_LAST_SEEN_AT)
+    const keys = goUpsellKeys(evt.properties.status.action)
+    if (!keys) return
+
+    const seen = kv.get(keys.lastSeenAt)
     if (typeof seen === "number" && Date.now() - seen < GO_UPSELL_WINDOW) return
 
-    if (kv.get(GO_UPSELL_DONT_SHOW)) return
+    if (kv.get(keys.dontShow)) return
 
     void DialogRetryAction.show(dialog, evt.properties.status.action).then((dontShowAgain) => {
-      if (dontShowAgain) kv.set(GO_UPSELL_DONT_SHOW, true)
-      kv.set(GO_UPSELL_LAST_SEEN_AT, Date.now())
+      if (dontShowAgain) kv.set(keys.dontShow, true)
+      kv.set(keys.lastSeenAt, Date.now())
     })
   })
 

+ 1 - 0
packages/opencode/src/session/processor.ts

@@ -701,6 +701,7 @@ export const layer: Layer.Layer<
             ),
             Effect.retry(
               SessionRetry.policy({
+                provider: input.model.providerID,
                 parse,
                 set: (info) => {
                   // TODO(v2): Temporary dual-write while migrating session messages to v2 events.

+ 11 - 3
packages/opencode/src/session/retry.ts

@@ -7,10 +7,13 @@ export type Err = ReturnType<NamedError["toObject"]>
 
 export const GO_UPSELL_MESSAGE = "Free usage exceeded, subscribe to Go"
 export const GO_UPSELL_URL = "https://opencode.ai/go"
+export type RetryReason = "free_tier_limit" | "account_rate_limit" | (string & {})
 
 export type Retryable = {
   message: string
   action?: {
+    reason: RetryReason
+    provider: string
     title: string
     message: string
     label: string
@@ -60,7 +63,7 @@ export function delay(attempt: number, error?: MessageV2.APIError) {
   return cap(Math.min(RETRY_INITIAL_DELAY * Math.pow(RETRY_BACKOFF_FACTOR, attempt - 1), RETRY_MAX_DELAY_NO_HEADERS))
 }
 
-export function retryable(error: Err) {
+export function retryable(error: Err, provider: string) {
   // context overflow errors should not be retried
   if (MessageV2.ContextOverflowError.isInstance(error)) return undefined
   if (MessageV2.APIError.isInstance(error)) {
@@ -72,6 +75,8 @@ export function retryable(error: Err) {
       return {
         message: GO_UPSELL_MESSAGE,
         action: {
+          reason: "free_tier_limit",
+          provider,
           title: "Free limit reached",
           message: "Subscribe to OpenCode Go for reliable access to the best open-source models, starting at $5/month.",
           label: "subscribe",
@@ -97,12 +102,14 @@ export function retryable(error: Err) {
         return minutes > 0 ? unit(minutes, "minute") : "less than a minute"
       })
 
-      const message = `${limitName} usage limit reached. It will reset in ${resetIn}. To continue using this model now, enable usage from your available balance`
+      const message = `${limitName ? `${limitName} usage limit` : "Usage limit"} reached. It will reset in ${resetIn}. To continue using this model now, enable usage from your available balance`
 
       const link = `https://opencode.ai/workspace/${workspace}/go`
       return {
         message: `${message} - ${link}`,
         action: {
+          reason: "account_rate_limit",
+          provider,
           title: "Go limit reached",
           message,
           label: "open settings",
@@ -165,13 +172,14 @@ function parseJSON(value: unknown) {
 }
 
 export function policy(opts: {
+  provider: string
   parse: (error: unknown) => Err
   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)
+      const retry = retryable(error, opts.provider)
       if (!retry) return Cause.done(meta.attempt)
       return Effect.gen(function* () {
         const wait = delay(meta.attempt, MessageV2.APIError.isInstance(error) ? error : undefined)

+ 2 - 0
packages/opencode/src/session/status.ts

@@ -17,6 +17,8 @@ export const Info = Schema.Union([
     message: Schema.String,
     action: Schema.optional(
       Schema.Struct({
+        reason: Schema.String,
+        provider: Schema.String,
         title: Schema.String,
         message: Schema.String,
         label: Schema.String,

+ 51 - 18
packages/opencode/test/session/retry.test.ts

@@ -13,6 +13,7 @@ import { provideTmpdirInstance } from "../fixture/fixture"
 import { testEffect } from "../lib/effect"
 
 const providerID = ProviderID.make("test")
+const retryProvider = "test"
 const it = testEffect(Layer.mergeAll(SessionStatus.defaultLayer, CrossSpawnSpawner.defaultLayer))
 
 function apiError(headers?: Record<string, string>): MessageV2.APIError {
@@ -92,6 +93,7 @@ describe("session.retry.delay", () => {
 
         const step = yield* Schedule.toStepWithMetadata(
           SessionRetry.policy({
+            provider: "test",
             parse: (err) => MessageV2.APIError.Schema.parse(err),
             set: (info) =>
               status.set(sessionID, {
@@ -118,47 +120,47 @@ describe("session.retry.delay", () => {
 describe("session.retry.retryable", () => {
   test("maps too_many_requests json messages", () => {
     const error = wrap(JSON.stringify({ type: "error", error: { type: "too_many_requests" } }))
-    expect(SessionRetry.retryable(error)).toEqual({ message: "Too Many Requests" })
+    expect(SessionRetry.retryable(error, retryProvider)).toEqual({ message: "Too Many Requests" })
   })
 
   test("maps overloaded provider codes", () => {
     const error = wrap(JSON.stringify({ code: "resource_exhausted" }))
-    expect(SessionRetry.retryable(error)).toEqual({ message: "Provider is overloaded" })
+    expect(SessionRetry.retryable(error, retryProvider)).toEqual({ message: "Provider is overloaded" })
   })
 
   test("does not retry unknown json messages", () => {
     const error = wrap(JSON.stringify({ error: { message: "no_kv_space" } }))
-    expect(SessionRetry.retryable(error)).toBeUndefined()
+    expect(SessionRetry.retryable(error, retryProvider)).toBeUndefined()
   })
 
   test("does not throw on numeric error codes", () => {
     const error = wrap(JSON.stringify({ type: "error", error: { code: 123 } }))
-    const result = SessionRetry.retryable(error)
+    const result = SessionRetry.retryable(error, retryProvider)
     expect(result).toBeUndefined()
   })
 
   test("returns undefined for non-json message", () => {
     const error = wrap("not-json")
-    expect(SessionRetry.retryable(error)).toBeUndefined()
+    expect(SessionRetry.retryable(error, retryProvider)).toBeUndefined()
   })
 
   test("retries plain text rate limit errors from Alibaba", () => {
     const msg =
       "Upstream error from Alibaba: Request rate increased too quickly. To ensure system stability, please adjust your client logic to scale requests more smoothly over time."
     const error = wrap(msg)
-    expect(SessionRetry.retryable(error)).toEqual({ message: msg })
+    expect(SessionRetry.retryable(error, retryProvider)).toEqual({ message: msg })
   })
 
   test("retries plain text rate limit errors", () => {
     const msg = "Rate limit exceeded, please try again later"
     const error = wrap(msg)
-    expect(SessionRetry.retryable(error)).toEqual({ message: msg })
+    expect(SessionRetry.retryable(error, retryProvider)).toEqual({ message: msg })
   })
 
   test("retries too many requests in plain text", () => {
     const msg = "Too many requests, please slow down"
     const error = wrap(msg)
-    expect(SessionRetry.retryable(error)).toEqual({ message: msg })
+    expect(SessionRetry.retryable(error, retryProvider)).toEqual({ message: msg })
   })
 
   test("does not retry context overflow errors", () => {
@@ -167,7 +169,7 @@ describe("session.retry.retryable", () => {
       responseBody: '{"error":{"code":"context_length_exceeded"}}',
     }).toObject()
 
-    expect(SessionRetry.retryable(error)).toBeUndefined()
+    expect(SessionRetry.retryable(error, retryProvider)).toBeUndefined()
   })
 
   test("retries 500 errors even when isRetryable is false", () => {
@@ -180,7 +182,7 @@ describe("session.retry.retryable", () => {
       }).toObject(),
     )
 
-    expect(SessionRetry.retryable(error)).toEqual({ message: "Internal server error" })
+    expect(SessionRetry.retryable(error, retryProvider)).toEqual({ message: "Internal server error" })
   })
 
   test("retries 502 bad gateway errors", () => {
@@ -192,7 +194,7 @@ describe("session.retry.retryable", () => {
       }).toObject(),
     )
 
-    expect(SessionRetry.retryable(error)).toEqual({ message: "Bad gateway" })
+    expect(SessionRetry.retryable(error, retryProvider)).toEqual({ message: "Bad gateway" })
   })
 
   test("retries 503 service unavailable errors", () => {
@@ -204,7 +206,7 @@ describe("session.retry.retryable", () => {
       }).toObject(),
     )
 
-    expect(SessionRetry.retryable(error)).toEqual({ message: "Service unavailable" })
+    expect(SessionRetry.retryable(error, retryProvider)).toEqual({ message: "Service unavailable" })
   })
 
   test("does not retry 4xx errors when isRetryable is false", () => {
@@ -216,7 +218,7 @@ describe("session.retry.retryable", () => {
       }).toObject(),
     )
 
-    expect(SessionRetry.retryable(error)).toBeUndefined()
+    expect(SessionRetry.retryable(error, retryProvider)).toBeUndefined()
   })
 
   test("retries ZlibError decompression failures", () => {
@@ -228,7 +230,7 @@ describe("session.retry.retryable", () => {
       }).toObject(),
     )
 
-    const retryable = SessionRetry.retryable(error)
+    const retryable = SessionRetry.retryable(error, retryProvider)
     expect(retryable).toBeDefined()
     expect(retryable).toEqual({ message: "Response decompression failed" })
   })
@@ -246,9 +248,11 @@ describe("session.retry.retryable", () => {
       }).toObject(),
     )
 
-    expect(SessionRetry.retryable(error)).toEqual({
+    expect(SessionRetry.retryable(error, "opencode")).toEqual({
       message: SessionRetry.GO_UPSELL_MESSAGE,
       action: {
+        reason: "free_tier_limit",
+        provider: "opencode",
         title: "Free limit reached",
         message: "Subscribe to OpenCode Go for reliable access to the best open-source models, starting at $5/month.",
         label: "subscribe",
@@ -280,10 +284,12 @@ describe("session.retry.retryable", () => {
       }).toObject(),
     )
 
-    expect(SessionRetry.retryable(error)).toEqual({
+    expect(SessionRetry.retryable(error, "opencode-go")).toEqual({
       message:
         "5 hour usage limit reached. It will reset in 5 hours 23 minutes. To continue using this model now, enable usage from your available balance - https://opencode.ai/workspace/wrk_01K6XGM22R6FM8JVABE9XDQXGH/go",
       action: {
+        reason: "account_rate_limit",
+        provider: "opencode-go",
         title: "Go limit reached",
         message:
           "5 hour usage limit reached. It will reset in 5 hours 23 minutes. To continue using this model now, enable usage from your available balance",
@@ -292,6 +298,33 @@ describe("session.retry.retryable", () => {
       },
     })
   })
+
+  test("maps Go subscription limits without limit metadata", () => {
+    const error = MessageV2.APIError.Schema.parse(
+      new MessageV2.APIError({
+        message: "Subscription quota exceeded. You can continue using free models.",
+        isRetryable: true,
+        statusCode: 429,
+        responseHeaders: {
+          "retry-after": "900",
+        },
+        responseBody: JSON.stringify({
+          type: "error",
+          error: {
+            type: "GoUsageLimitError",
+            message: "Subscription quota exceeded. You can continue using free models.",
+          },
+          metadata: {
+            workspace: "wrk_01K6XGM22R6FM8JVABE9XDQXGH",
+          },
+        }),
+      }).toObject(),
+    )
+
+    expect(SessionRetry.retryable(error, "opencode-go")?.action?.message).toBe(
+      "Usage limit reached. It will reset in 15 minutes. To continue using this model now, enable usage from your available balance",
+    )
+  })
 })
 
 describe("session.message-v2.fromError", () => {
@@ -341,7 +374,7 @@ describe("session.message-v2.fromError", () => {
       }).toObject(),
     )
 
-    const retryable = SessionRetry.retryable(error)
+    const retryable = SessionRetry.retryable(error, retryProvider)
     expect(retryable).toBeDefined()
     expect(retryable).toEqual({ message: "Connection reset by server" })
   })
@@ -381,6 +414,6 @@ describe("session.message-v2.fromError", () => {
     expect(MessageV2.APIError.isInstance(result)).toBe(true)
     if (!MessageV2.APIError.isInstance(result)) throw new Error("expected APIError")
     expect(result.data.isRetryable).toBe(true)
-    expect(SessionRetry.retryable(result)).toEqual({ message: "An error occurred while processing your request." })
+    expect(SessionRetry.retryable(result, retryProvider)).toEqual({ message: "An error occurred while processing your request." })
   })
 })

+ 2 - 0
packages/opencode/test/session/schema-decoding.test.ts

@@ -236,6 +236,8 @@ describe("SessionStatus.Info", () => {
       attempt: 1,
       message: "transient",
       action: {
+        reason: "free_tier_limit",
+        provider: "opencode",
         title: "Free limit reached",
         message: "Subscribe to OpenCode Go.",
         label: "subscribe",

+ 2 - 0
packages/sdk/js/src/v2/gen/types.gen.ts

@@ -267,6 +267,8 @@ export type SessionStatus =
       attempt: number
       message: string
       action?: {
+        reason: string
+        provider: string
         title: string
         message: string
         label: string

+ 0 - 37
script/zen-limit-server.ts

@@ -1,37 +0,0 @@
-const retryAfterSeconds = 15 * 60
-
-// const response = {
-//   type: "error",
-//   error: {
-//     type: "FreeUsageLimitError",
-//     message: "Free usage exceeded, subscribe to Go https://opencode.ai/go",
-//   },
-//   metadata: {},
-// }
-
-const response = {
-  type: "error",
-  error: {
-    type: "GoUsageLimitError",
-    message: "Subscription quota exceeded. You can continue using free models.",
-  },
-  metadata: {
-    workspace: "wrk_01K6XGM22R6FM8JVABE9XDQXGH",
-    limit: "5 hour",
-    resetAt: retryAfterSeconds,
-  },
-}
-
-Bun.serve({
-  port: 4141,
-  fetch() {
-    return Response.json(response, {
-      status: 429,
-      headers: {
-        "retry-after": String(retryAfterSeconds),
-      },
-    })
-  },
-})
-
-console.log("Zen limit repro server listening on http://localhost:4141")