| 1 | package provider |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "errors" |
| 6 | "fmt" |
| 7 | "io" |
| 8 | "math/rand" |
| 9 | "net" |
| 10 | "net/http" |
| 11 | "strconv" |
| 12 | "strings" |
| 13 | "sync/atomic" |
| 14 | "syscall" |
| 15 | "time" |
| 16 | ) |
| 17 | |
| 18 | // MaxRetries is the number of times SendWithRetry re-attempts the connection + |
| 19 | // header phase after the initial try (so up to MaxRetries+1 total attempts). |
| 20 | const MaxRetries = 10 |
| 21 | |
| 22 | const maxBackoff = 15 * time.Second |
| 23 | |
| 24 | // maxRetryAfter bounds a server-supplied Retry-After. Rate-limit windows are |
| 25 | // routinely longer than our own backoff cap, and clamping to it just spends |
| 26 | // attempts re-hitting the same closed window; the sleep is cancellable, so a |
| 27 | // longer honest wait costs nothing the user can't interrupt. |
| 28 | const maxRetryAfter = 60 * time.Second |
| 29 | |
| 30 | // errorBodyReadTimeout bounds how long draining a non-OK response body may |
| 31 | // block. Proxies and gateways under load (502/524 storms) can send headers and |
| 32 | // then stall the body on a half-open connection; http.Client has no Timeout |
| 33 | // and ResponseHeaderTimeout no longer applies once headers arrive, so without |
| 34 | // this deadline the retry loop blocks in io.ReadAll indefinitely with no |
| 35 | // user-visible progress — the turn looks frozen until the process is killed |
| 36 | // (#6607). A var, not a const, so tests can shrink it. |
| 37 | var errorBodyReadTimeout = 10 * time.Second |
| 38 | |
| 39 | // maxAuthRetries bounds how many times a 401/403 is retried for a key that has |
| 40 | // authenticated before: a transient server-side rejection (quota/gateway/rate) |
| 41 | // usually clears in a couple of attempts, whereas a key that never worked is a |
| 42 | // real config error and fails fast. |
| 43 | const maxAuthRetries = 2 |
| 44 | |
| 45 | // SendOptions carries the per-request context SendWithRetry needs to label |
| 46 | // errors and decide whether a 401 is worth retrying. |
| 47 | type SendOptions struct { |
| 48 | Provider string // stable provider instance id |
| 49 | ProviderDisplayName string // user-editable display label |
| 50 | Protocol string // configured wire adapter id |
| 51 | KeyEnv string // api_key_env the key is read from, when known |
| 52 | KeySource string // human-readable source of KeyEnv, when known |
| 53 | KeyPresent bool // a non-empty key is being sent — separates "rejected" from "missing" |
| 54 | RetryAuth bool // the key has authenticated before — retry transient 401s instead of failing fast |
| 55 | } |
| 56 | |
| 57 | // RetryInfo describes a backoff about to happen: Attempt is the 1-based retry |
| 58 | // number (of Max) and Delay is how long SendWithRetry will wait before it. |
| 59 | type RetryInfo struct { |
| 60 | Attempt int |
| 61 | Max int |
| 62 | Delay time.Duration |
| 63 | Err error |
| 64 | } |
| 65 | |
| 66 | type RetryNotify func(RetryInfo) |
| 67 | |
| 68 | type retryNotifyKey struct{} |
| 69 | |
| 70 | type requestAttemptCounterKey struct{} |
| 71 | |
| 72 | type requestAttemptCounter struct { |
| 73 | count atomic.Int64 |
| 74 | } |
| 75 | |
| 76 | // WithRetryNotify attaches a callback that SendWithRetry invokes before each |
| 77 | // backoff sleep, so the agent can surface a transient "retrying (n/m)" status. |
| 78 | func WithRetryNotify(ctx context.Context, fn RetryNotify) context.Context { |
| 79 | if fn == nil { |
| 80 | return ctx |
| 81 | } |
| 82 | return context.WithValue(ctx, retryNotifyKey{}, fn) |
| 83 | } |
| 84 | |
| 85 | func retryNotifyFromContext(ctx context.Context) RetryNotify { |
| 86 | fn, _ := ctx.Value(retryNotifyKey{}).(RetryNotify) |
| 87 | return fn |
| 88 | } |
| 89 | |
| 90 | // WithRequestAttemptCounter returns a context that counts every HTTP request |
| 91 | // SendWithRetry starts. An existing counter is reused so a caller can observe |
| 92 | // attempts even when the provider returns before producing a Usage chunk. |
| 93 | // Provider implementations use one counter for a logical stream (including |
| 94 | // header retries and safe reconnects), then attach the final count to the |
| 95 | // stream's Usage record. |
| 96 | func WithRequestAttemptCounter(ctx context.Context) context.Context { |
| 97 | if ctx == nil { |
| 98 | ctx = context.Background() |
| 99 | } |
| 100 | if counter, _ := ctx.Value(requestAttemptCounterKey{}).(*requestAttemptCounter); counter != nil { |
| 101 | return ctx |
| 102 | } |
| 103 | return context.WithValue(ctx, requestAttemptCounterKey{}, &requestAttemptCounter{}) |
| 104 | } |
| 105 | |
| 106 | // WithIndependentRequestAttemptCounter gives an auxiliary call its own usage |
| 107 | // count while preserving cancellation and other context values from its parent. |
| 108 | func WithIndependentRequestAttemptCounter(ctx context.Context) context.Context { |
| 109 | return context.WithValue(ctx, requestAttemptCounterKey{}, &requestAttemptCounter{}) |
| 110 | } |
| 111 | |
| 112 | // RequestAttemptCount returns the number of HTTP requests started through |
| 113 | // SendWithRetry for the counter attached to ctx. |
| 114 | func RequestAttemptCount(ctx context.Context) int { |
| 115 | if ctx == nil { |
| 116 | return 0 |
| 117 | } |
| 118 | counter, _ := ctx.Value(requestAttemptCounterKey{}).(*requestAttemptCounter) |
| 119 | if counter == nil { |
| 120 | return 0 |
| 121 | } |
| 122 | return int(counter.count.Load()) |
| 123 | } |
| 124 | |
| 125 | // ApplyRequestAttemptCount copies the stream's exact HTTP request count into a |
| 126 | // Usage record. Contexts without a counter leave the record unchanged so custom |
| 127 | // providers keep the zero-means-one compatibility contract. |
| 128 | func ApplyRequestAttemptCount(ctx context.Context, usage *Usage) { |
| 129 | if usage == nil { |
| 130 | return |
| 131 | } |
| 132 | if count := RequestAttemptCount(ctx); count > 0 { |
| 133 | usage.RequestCount = count |
| 134 | } |
| 135 | } |
| 136 | |
| 137 | // UsageWithRequestAttemptCount returns a copy of usage carrying the exact |
| 138 | // number of HTTP requests observed through ctx. When a provider request fails |
| 139 | // before producing token usage, it returns a request-only Usage record so |
| 140 | // callers can still account for the API calls. If neither usage nor attempts |
| 141 | // exist, it returns nil. |
| 142 | func UsageWithRequestAttemptCount(ctx context.Context, usage *Usage) *Usage { |
| 143 | count := RequestAttemptCount(ctx) |
| 144 | if usage == nil { |
| 145 | if count <= 0 { |
| 146 | return nil |
| 147 | } |
| 148 | return &Usage{RequestCount: count, Unknown: true} |
| 149 | } |
| 150 | result := *usage |
| 151 | if count > 0 { |
| 152 | result.RequestCount = count |
| 153 | } |
| 154 | return &result |
| 155 | } |
| 156 | |
| 157 | func recordRequestAttempt(ctx context.Context) { |
| 158 | if ctx == nil { |
| 159 | return |
| 160 | } |
| 161 | counter, _ := ctx.Value(requestAttemptCounterKey{}).(*requestAttemptCounter) |
| 162 | if counter != nil { |
| 163 | counter.count.Add(1) |
| 164 | } |
| 165 | } |
| 166 | |
| 167 | // APIError reports a non-OK HTTP status that isn't an auth failure. Status |
| 168 | // carries the code so the display layer can map it to an actionable, localized |
| 169 | // message; Body is a trimmed snippet of the response. |
| 170 | type APIError struct { |
| 171 | RetryAfter time.Duration // uncapped server delay for managed recovery |
| 172 | ShouldRetry string // explicit provider retry hint |
| 173 | Provider string // stable provider instance id |
| 174 | ProviderDisplayName string |
| 175 | Protocol string |
| 176 | Status int |
| 177 | Body string |
| 178 | TraceID string // provider trace identifier from the response headers, when present |
| 179 | RequestPath string // path only; query and URL userinfo are never retained |
| 180 | ToolContext string // resolved Reasonix/MCP identity for provider-indexed tool schema errors |
| 181 | } |
| 182 | |
| 183 | func (e *APIError) Error() string { |
| 184 | label := ProviderDisplayLabel(e.Provider, e.ProviderDisplayName, e.Protocol) |
| 185 | var base string |
| 186 | if e.Body == "" { |
| 187 | base = fmt.Sprintf("%s: status %d", label, e.Status) |
| 188 | } else { |
| 189 | base = fmt.Sprintf("%s: status %d: %s", label, e.Status, e.Body) |
| 190 | } |
| 191 | if e.ToolContext != "" { |
| 192 | return base + "\n" + e.ToolContext |
| 193 | } |
| 194 | return base |
| 195 | } |
| 196 | |
| 197 | // RetryableStatus reports whether a backoff can plausibly recover from status s: |
| 198 | // 408 (request timeout), 429 (rate limit) and 5xx (incl. Anthropic's 529). Other |
| 199 | // 4xx (400/401/402/422, …) are caller/config problems retrying can't fix. |
| 200 | func RetryableStatus(s int) bool { |
| 201 | return s == http.StatusRequestTimeout || s == http.StatusTooManyRequests || (s >= 500 && s <= 599) |
| 202 | } |
| 203 | |
| 204 | func transientErr(err error) bool { |
| 205 | if err == nil { |
| 206 | return false |
| 207 | } |
| 208 | if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { |
| 209 | return false |
| 210 | } |
| 211 | return true |
| 212 | } |
| 213 | |
| 214 | // IsConnReset reports whether err is a connection-level drop (peer reset, |
| 215 | // truncated body, closed socket) as opposed to a protocol or caller error. A |
| 216 | // stream cut this way mid-body can be replayed from scratch, unlike a decode or |
| 217 | // 4xx error. The common trigger is a local proxy (v2rayN/sing-box) idle-closing |
| 218 | // the long-lived SSE connection during a reasoner's first-token gap. |
| 219 | func IsConnReset(err error) bool { |
| 220 | if err == nil { |
| 221 | return false |
| 222 | } |
| 223 | if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { |
| 224 | return false |
| 225 | } |
| 226 | if errors.Is(err, io.ErrUnexpectedEOF) || errors.Is(err, io.EOF) || |
| 227 | errors.Is(err, net.ErrClosed) || |
| 228 | errors.Is(err, syscall.ECONNRESET) || errors.Is(err, syscall.ECONNABORTED) { |
| 229 | return true |
| 230 | } |
| 231 | var netErr net.Error |
| 232 | return errors.As(err, &netErr) |
| 233 | } |
| 234 | |
| 235 | func backoffDelay(attempt int, retryAfter time.Duration) time.Duration { |
| 236 | if retryAfter > 0 { |
| 237 | if retryAfter > maxRetryAfter { |
| 238 | return maxRetryAfter |
| 239 | } |
| 240 | return retryAfter |
| 241 | } |
| 242 | d := min(time.Duration(1<<(attempt-1))*500*time.Millisecond, maxBackoff) |
| 243 | return d + time.Duration(rand.Intn(250))*time.Millisecond |
| 244 | } |
| 245 | |
| 246 | func parseRetryAfter(resp *http.Response) time.Duration { |
| 247 | v := strings.TrimSpace(resp.Header.Get("Retry-After")) |
| 248 | if v == "" { |
| 249 | return 0 |
| 250 | } |
| 251 | if secs, err := strconv.Atoi(v); err == nil && secs >= 0 { |
| 252 | return time.Duration(secs) * time.Second |
| 253 | } |
| 254 | // RFC 9110 also allows an HTTP-date; gateways in front of rate-limited |
| 255 | // backends use it more often than the delta-seconds form. |
| 256 | if when, err := http.ParseTime(v); err == nil { |
| 257 | if d := time.Until(when); d > 0 { |
| 258 | return d |
| 259 | } |
| 260 | } |
| 261 | return 0 |
| 262 | } |
| 263 | |
| 264 | // readErrorBody drains a non-OK response body under a hard deadline and |
| 265 | // returns up to the first 4 KiB for the error message. Context cancellation |
| 266 | // already unblocks the read (the transport aborts body reads when the request |
| 267 | // context is canceled); the timer covers the case nobody cancels — a half-open |
| 268 | // upstream that sent headers and then went silent. Closing the body from the |
| 269 | // timer goroutine is the documented way to unblock an in-flight Read; it |
| 270 | // tears down the connection, which is the right call for a stalled peer. |
| 271 | func readErrorBody(resp *http.Response) []byte { |
| 272 | timer := time.AfterFunc(errorBodyReadTimeout, func() { resp.Body.Close() }) |
| 273 | msg, _ := io.ReadAll(io.LimitReader(resp.Body, 4096)) |
| 274 | // Drain the rest so a healthy connection can be reused; the timer still |
| 275 | // arms this read, so a body that stalls after the first 4 KiB cannot |
| 276 | // wedge the retry loop either. |
| 277 | _, _ = io.Copy(io.Discard, resp.Body) |
| 278 | timer.Stop() |
| 279 | resp.Body.Close() |
| 280 | return msg |
| 281 | } |
| 282 | |
| 283 | // SendWithRetry POSTs a streaming request built by newReq and returns the OK |
| 284 | // response. It retries the connection+header phase up to MaxRetries times on |
| 285 | // transient network errors and retryable statuses with capped exponential |
| 286 | // backoff + jitter, honoring Retry-After. A 401/403 becomes *AuthError: it |
| 287 | // fails fast for a key that has never authenticated (opts.RetryAuth false), but |
| 288 | // for a previously-good key it backs off and retries up to maxAuthRetries — |
| 289 | // MiMo and similar gateways return a transient 401 under load. Other non-OK |
| 290 | // statuses become *APIError. A RetryNotify in ctx fires before each sleep. |
| 291 | // Retries cover only the header phase — once the body streams, mid-stream |
| 292 | // failures are not retried (the model has already emitted tokens). |
| 293 | func SendWithRetry(ctx context.Context, httpClient *http.Client, opts SendOptions, newReq func(context.Context) (*http.Request, error)) (*http.Response, error) { |
| 294 | notify := retryNotifyFromContext(ctx) |
| 295 | identity := RequestIdentity{Provider: opts.Provider, DisplayName: opts.ProviderDisplayName, Protocol: opts.Protocol} |
| 296 | var lastErr error |
| 297 | var retryAfter time.Duration |
| 298 | authRetries := 0 |
| 299 | |
| 300 | limit := MaxRetries |
| 301 | if ManagedRecovery(ctx) { |
| 302 | limit = 0 |
| 303 | } |
| 304 | for attempt := 0; attempt <= limit; attempt++ { |
| 305 | if attempt > 0 { |
| 306 | delay := backoffDelay(attempt, retryAfter) |
| 307 | if notify != nil { |
| 308 | notify(RetryInfo{Attempt: attempt, Max: MaxRetries, Delay: delay, Err: lastErr}) |
| 309 | } |
| 310 | select { |
| 311 | case <-ctx.Done(): |
| 312 | return nil, ctx.Err() |
| 313 | case <-time.After(delay): |
| 314 | } |
| 315 | } |
| 316 | retryAfter = 0 |
| 317 | |
| 318 | req, err := newReq(ctx) |
| 319 | if err != nil { |
| 320 | return nil, &RequestFailure{Identity: identity, Operation: "build request", Err: err} |
| 321 | } |
| 322 | recordRequestAttempt(ctx) |
| 323 | resp, err := httpClient.Do(req) |
| 324 | if err != nil { |
| 325 | if !transientErr(err) { |
| 326 | return nil, &RequestFailure{Identity: identity, Operation: "request failed", Err: err} |
| 327 | } |
| 328 | lastErr = &RequestFailure{Identity: identity, Operation: "request failed", Err: err} |
| 329 | continue |
| 330 | } |
| 331 | if resp.StatusCode == http.StatusOK { |
| 332 | return resp, nil |
| 333 | } |
| 334 | |
| 335 | msg := readErrorBody(resp) |
| 336 | retryAfter = parseRetryAfter(resp) |
| 337 | if quota := QuotaErrorFromResponseWithIdentity(opts.Provider, opts.ProviderDisplayName, opts.Protocol, resp.StatusCode, string(msg)); quota != nil { |
| 338 | return nil, quota |
| 339 | } |
| 340 | |
| 341 | if resp.StatusCode == http.StatusUnauthorized || resp.StatusCode == http.StatusForbidden { |
| 342 | authErr := &AuthError{Provider: opts.Provider, ProviderDisplayName: opts.ProviderDisplayName, Protocol: opts.Protocol, KeyEnv: opts.KeyEnv, KeySource: opts.KeySource, Status: resp.StatusCode, HasKey: opts.KeyPresent, Body: strings.TrimSpace(string(msg))} |
| 343 | if !ManagedRecovery(ctx) && opts.RetryAuth && authRetries < maxAuthRetries { |
| 344 | authRetries++ |
| 345 | lastErr = authErr |
| 346 | continue |
| 347 | } |
| 348 | return nil, authErr |
| 349 | } |
| 350 | apiErr := &APIError{ |
| 351 | RetryAfter: retryAfter, |
| 352 | ShouldRetry: resp.Header.Get("x-should-retry"), |
| 353 | Provider: opts.Provider, |
| 354 | ProviderDisplayName: opts.ProviderDisplayName, |
| 355 | Protocol: opts.Protocol, |
| 356 | Status: resp.StatusCode, |
| 357 | Body: strings.TrimSpace(string(msg)), |
| 358 | TraceID: responseTraceID(resp.Header), |
| 359 | RequestPath: responseRequestPath(resp), |
| 360 | } |
| 361 | if !RetryableStatus(resp.StatusCode) { |
| 362 | if limitErr := ParseOutputLimitError(apiErr); limitErr != nil { |
| 363 | return nil, limitErr |
| 364 | } |
| 365 | if limitErr := ParseContextLimitError(apiErr); limitErr != nil { |
| 366 | return nil, limitErr |
| 367 | } |
| 368 | if replayErr := ParseReasoningReplayError(apiErr); replayErr != nil { |
| 369 | return nil, replayErr |
| 370 | } |
| 371 | return nil, apiErr |
| 372 | } |
| 373 | lastErr = apiErr |
| 374 | } |
| 375 | return nil, lastErr |
| 376 | } |
| 377 | |
| 378 | func responseRequestPath(resp *http.Response) string { |
| 379 | if resp == nil || resp.Request == nil || resp.Request.URL == nil { |
| 380 | return "" |
| 381 | } |
| 382 | path := resp.Request.URL.EscapedPath() |
| 383 | if len(path) > 512 { |
| 384 | return path[:512] |
| 385 | } |
| 386 | return path |
| 387 | } |
| 388 | |
| 389 | func responseTraceID(header http.Header) string { |
| 390 | for _, name := range []string{"trace_id", "trace-id", "x-trace-id"} { |
| 391 | if value := strings.TrimSpace(header.Get(name)); value != "" { |
| 392 | return value |
| 393 | } |
| 394 | } |
| 395 | return "" |
| 396 | } |
| 397 |