| 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 // provider instance name, surfaced in errors |
| 49 | KeyEnv string // api_key_env the key is read from, when known |
| 50 | KeySource string // human-readable source of KeyEnv, when known |
| 51 | KeyPresent bool // a non-empty key is being sent — separates "rejected" from "missing" |
| 52 | RetryAuth bool // the key has authenticated before — retry transient 401s instead of failing fast |
| 53 | } |
| 54 | |
| 55 | // RetryInfo describes a backoff about to happen: Attempt is the 1-based retry |
| 56 | // number (of Max) and Delay is how long SendWithRetry will wait before it. |
| 57 | type RetryInfo struct { |
| 58 | Attempt int |
| 59 | Max int |
| 60 | Delay time.Duration |
| 61 | Err error |
| 62 | } |
| 63 | |
| 64 | type RetryNotify func(RetryInfo) |
| 65 | |
| 66 | type retryNotifyKey struct{} |
| 67 | |
| 68 | type requestAttemptCounterKey struct{} |
| 69 | |
| 70 | type requestAttemptCounter struct { |
| 71 | count atomic.Int64 |
| 72 | } |
| 73 | |
| 74 | // WithRetryNotify attaches a callback that SendWithRetry invokes before each |
| 75 | // backoff sleep, so the agent can surface a transient "retrying (n/m)" status. |
| 76 | func WithRetryNotify(ctx context.Context, fn RetryNotify) context.Context { |
| 77 | if fn == nil { |
| 78 | return ctx |
| 79 | } |
| 80 | return context.WithValue(ctx, retryNotifyKey{}, fn) |
| 81 | } |
| 82 | |
| 83 | func retryNotifyFromContext(ctx context.Context) RetryNotify { |
| 84 | fn, _ := ctx.Value(retryNotifyKey{}).(RetryNotify) |
| 85 | return fn |
| 86 | } |
| 87 | |
| 88 | // WithRequestAttemptCounter returns a context that counts every HTTP request |
| 89 | // SendWithRetry starts. An existing counter is reused so a caller can observe |
| 90 | // attempts even when the provider returns before producing a Usage chunk. |
| 91 | // Provider implementations use one counter for a logical stream (including |
| 92 | // header retries and safe reconnects), then attach the final count to the |
| 93 | // stream's Usage record. |
| 94 | func WithRequestAttemptCounter(ctx context.Context) context.Context { |
| 95 | if ctx == nil { |
| 96 | ctx = context.Background() |
| 97 | } |
| 98 | if counter, _ := ctx.Value(requestAttemptCounterKey{}).(*requestAttemptCounter); counter != nil { |
| 99 | return ctx |
| 100 | } |
| 101 | return context.WithValue(ctx, requestAttemptCounterKey{}, &requestAttemptCounter{}) |
| 102 | } |
| 103 | |
| 104 | // RequestAttemptCount returns the number of HTTP requests started through |
| 105 | // SendWithRetry for the counter attached to ctx. |
| 106 | func RequestAttemptCount(ctx context.Context) int { |
| 107 | if ctx == nil { |
| 108 | return 0 |
| 109 | } |
| 110 | counter, _ := ctx.Value(requestAttemptCounterKey{}).(*requestAttemptCounter) |
| 111 | if counter == nil { |
| 112 | return 0 |
| 113 | } |
| 114 | return int(counter.count.Load()) |
| 115 | } |
| 116 | |
| 117 | // ApplyRequestAttemptCount copies the stream's exact HTTP request count into a |
| 118 | // Usage record. Contexts without a counter leave the record unchanged so custom |
| 119 | // providers keep the zero-means-one compatibility contract. |
| 120 | func ApplyRequestAttemptCount(ctx context.Context, usage *Usage) { |
| 121 | if usage == nil { |
| 122 | return |
| 123 | } |
| 124 | if count := RequestAttemptCount(ctx); count > 0 { |
| 125 | usage.RequestCount = count |
| 126 | } |
| 127 | } |
| 128 | |
| 129 | // UsageWithRequestAttemptCount returns a copy of usage carrying the exact |
| 130 | // number of HTTP requests observed through ctx. When a provider request fails |
| 131 | // before producing token usage, it returns a request-only Usage record so |
| 132 | // callers can still account for the API calls. If neither usage nor attempts |
| 133 | // exist, it returns nil. |
| 134 | func UsageWithRequestAttemptCount(ctx context.Context, usage *Usage) *Usage { |
| 135 | count := RequestAttemptCount(ctx) |
| 136 | if usage == nil { |
| 137 | if count <= 0 { |
| 138 | return nil |
| 139 | } |
| 140 | return &Usage{RequestCount: count} |
| 141 | } |
| 142 | result := *usage |
| 143 | if count > 0 { |
| 144 | result.RequestCount = count |
| 145 | } |
| 146 | return &result |
| 147 | } |
| 148 | |
| 149 | func recordRequestAttempt(ctx context.Context) { |
| 150 | if ctx == nil { |
| 151 | return |
| 152 | } |
| 153 | counter, _ := ctx.Value(requestAttemptCounterKey{}).(*requestAttemptCounter) |
| 154 | if counter != nil { |
| 155 | counter.count.Add(1) |
| 156 | } |
| 157 | } |
| 158 | |
| 159 | // APIError reports a non-OK HTTP status that isn't an auth failure. Status |
| 160 | // carries the code so the display layer can map it to an actionable, localized |
| 161 | // message; Body is a trimmed snippet of the response. |
| 162 | type APIError struct { |
| 163 | Provider string |
| 164 | Status int |
| 165 | Body string |
| 166 | TraceID string // provider trace identifier from the response headers, when present |
| 167 | ToolContext string // resolved Reasonix/MCP identity for provider-indexed tool schema errors |
| 168 | } |
| 169 | |
| 170 | func (e *APIError) Error() string { |
| 171 | var base string |
| 172 | if e.Body == "" { |
| 173 | base = fmt.Sprintf("%s: status %d", e.Provider, e.Status) |
| 174 | } else { |
| 175 | base = fmt.Sprintf("%s: status %d: %s", e.Provider, e.Status, e.Body) |
| 176 | } |
| 177 | if e.ToolContext != "" { |
| 178 | return base + "\n" + e.ToolContext |
| 179 | } |
| 180 | return base |
| 181 | } |
| 182 | |
| 183 | // RetryableStatus reports whether a backoff can plausibly recover from status s: |
| 184 | // 408 (request timeout), 429 (rate limit) and 5xx (incl. Anthropic's 529). Other |
| 185 | // 4xx (400/401/402/422, …) are caller/config problems retrying can't fix. |
| 186 | func RetryableStatus(s int) bool { |
| 187 | return s == http.StatusRequestTimeout || s == http.StatusTooManyRequests || (s >= 500 && s <= 599) |
| 188 | } |
| 189 | |
| 190 | func transientErr(err error) bool { |
| 191 | if err == nil { |
| 192 | return false |
| 193 | } |
| 194 | if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { |
| 195 | return false |
| 196 | } |
| 197 | return true |
| 198 | } |
| 199 | |
| 200 | // IsConnReset reports whether err is a connection-level drop (peer reset, |
| 201 | // truncated body, closed socket) as opposed to a protocol or caller error. A |
| 202 | // stream cut this way mid-body can be replayed from scratch, unlike a decode or |
| 203 | // 4xx error. The common trigger is a local proxy (v2rayN/sing-box) idle-closing |
| 204 | // the long-lived SSE connection during a reasoner's first-token gap. |
| 205 | func IsConnReset(err error) bool { |
| 206 | if err == nil { |
| 207 | return false |
| 208 | } |
| 209 | if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { |
| 210 | return false |
| 211 | } |
| 212 | if errors.Is(err, io.ErrUnexpectedEOF) || errors.Is(err, io.EOF) || |
| 213 | errors.Is(err, net.ErrClosed) || |
| 214 | errors.Is(err, syscall.ECONNRESET) || errors.Is(err, syscall.ECONNABORTED) { |
| 215 | return true |
| 216 | } |
| 217 | var netErr net.Error |
| 218 | return errors.As(err, &netErr) |
| 219 | } |
| 220 | |
| 221 | func backoffDelay(attempt int, retryAfter time.Duration) time.Duration { |
| 222 | if retryAfter > 0 { |
| 223 | if retryAfter > maxRetryAfter { |
| 224 | return maxRetryAfter |
| 225 | } |
| 226 | return retryAfter |
| 227 | } |
| 228 | d := time.Duration(1<<(attempt-1)) * 500 * time.Millisecond |
| 229 | if d > maxBackoff { |
| 230 | d = maxBackoff |
| 231 | } |
| 232 | return d + time.Duration(rand.Intn(250))*time.Millisecond |
| 233 | } |
| 234 | |
| 235 | func parseRetryAfter(resp *http.Response) time.Duration { |
| 236 | v := strings.TrimSpace(resp.Header.Get("Retry-After")) |
| 237 | if v == "" { |
| 238 | return 0 |
| 239 | } |
| 240 | if secs, err := strconv.Atoi(v); err == nil && secs >= 0 { |
| 241 | return time.Duration(secs) * time.Second |
| 242 | } |
| 243 | // RFC 9110 also allows an HTTP-date; gateways in front of rate-limited |
| 244 | // backends use it more often than the delta-seconds form. |
| 245 | if when, err := http.ParseTime(v); err == nil { |
| 246 | if d := time.Until(when); d > 0 { |
| 247 | return d |
| 248 | } |
| 249 | } |
| 250 | return 0 |
| 251 | } |
| 252 | |
| 253 | // readErrorBody drains a non-OK response body under a hard deadline and |
| 254 | // returns up to the first 4 KiB for the error message. Context cancellation |
| 255 | // already unblocks the read (the transport aborts body reads when the request |
| 256 | // context is canceled); the timer covers the case nobody cancels — a half-open |
| 257 | // upstream that sent headers and then went silent. Closing the body from the |
| 258 | // timer goroutine is the documented way to unblock an in-flight Read; it |
| 259 | // tears down the connection, which is the right call for a stalled peer. |
| 260 | func readErrorBody(resp *http.Response) []byte { |
| 261 | timer := time.AfterFunc(errorBodyReadTimeout, func() { resp.Body.Close() }) |
| 262 | msg, _ := io.ReadAll(io.LimitReader(resp.Body, 4096)) |
| 263 | // Drain the rest so a healthy connection can be reused; the timer still |
| 264 | // arms this read, so a body that stalls after the first 4 KiB cannot |
| 265 | // wedge the retry loop either. |
| 266 | _, _ = io.Copy(io.Discard, resp.Body) |
| 267 | timer.Stop() |
| 268 | resp.Body.Close() |
| 269 | return msg |
| 270 | } |
| 271 | |
| 272 | // SendWithRetry POSTs a streaming request built by newReq and returns the OK |
| 273 | // response. It retries the connection+header phase up to MaxRetries times on |
| 274 | // transient network errors and retryable statuses with capped exponential |
| 275 | // backoff + jitter, honoring Retry-After. A 401/403 becomes *AuthError: it |
| 276 | // fails fast for a key that has never authenticated (opts.RetryAuth false), but |
| 277 | // for a previously-good key it backs off and retries up to maxAuthRetries — |
| 278 | // MiMo and similar gateways return a transient 401 under load. Other non-OK |
| 279 | // statuses become *APIError. A RetryNotify in ctx fires before each sleep. |
| 280 | // Retries cover only the header phase — once the body streams, mid-stream |
| 281 | // failures are not retried (the model has already emitted tokens). |
| 282 | func SendWithRetry(ctx context.Context, httpClient *http.Client, opts SendOptions, newReq func(context.Context) (*http.Request, error)) (*http.Response, error) { |
| 283 | notify := retryNotifyFromContext(ctx) |
| 284 | var lastErr error |
| 285 | var retryAfter time.Duration |
| 286 | authRetries := 0 |
| 287 | |
| 288 | for attempt := 0; attempt <= MaxRetries; attempt++ { |
| 289 | if attempt > 0 { |
| 290 | delay := backoffDelay(attempt, retryAfter) |
| 291 | if notify != nil { |
| 292 | notify(RetryInfo{Attempt: attempt, Max: MaxRetries, Delay: delay, Err: lastErr}) |
| 293 | } |
| 294 | select { |
| 295 | case <-ctx.Done(): |
| 296 | return nil, ctx.Err() |
| 297 | case <-time.After(delay): |
| 298 | } |
| 299 | } |
| 300 | retryAfter = 0 |
| 301 | |
| 302 | req, err := newReq(ctx) |
| 303 | if err != nil { |
| 304 | return nil, fmt.Errorf("%s: build request: %w", opts.Provider, err) |
| 305 | } |
| 306 | recordRequestAttempt(ctx) |
| 307 | resp, err := httpClient.Do(req) |
| 308 | if err != nil { |
| 309 | if !transientErr(err) { |
| 310 | return nil, fmt.Errorf("%s: request failed: %w", opts.Provider, err) |
| 311 | } |
| 312 | lastErr = fmt.Errorf("%s: request failed: %w", opts.Provider, err) |
| 313 | continue |
| 314 | } |
| 315 | if resp.StatusCode == http.StatusOK { |
| 316 | return resp, nil |
| 317 | } |
| 318 | |
| 319 | msg := readErrorBody(resp) |
| 320 | retryAfter = parseRetryAfter(resp) |
| 321 | |
| 322 | if resp.StatusCode == http.StatusUnauthorized || resp.StatusCode == http.StatusForbidden { |
| 323 | authErr := &AuthError{Provider: opts.Provider, KeyEnv: opts.KeyEnv, KeySource: opts.KeySource, Status: resp.StatusCode, HasKey: opts.KeyPresent, Body: strings.TrimSpace(string(msg))} |
| 324 | if opts.RetryAuth && authRetries < maxAuthRetries { |
| 325 | authRetries++ |
| 326 | lastErr = authErr |
| 327 | continue |
| 328 | } |
| 329 | return nil, authErr |
| 330 | } |
| 331 | apiErr := &APIError{ |
| 332 | Provider: opts.Provider, |
| 333 | Status: resp.StatusCode, |
| 334 | Body: strings.TrimSpace(string(msg)), |
| 335 | TraceID: responseTraceID(resp.Header), |
| 336 | } |
| 337 | if !RetryableStatus(resp.StatusCode) { |
| 338 | return nil, apiErr |
| 339 | } |
| 340 | lastErr = apiErr |
| 341 | } |
| 342 | return nil, lastErr |
| 343 | } |
| 344 | |
| 345 | func responseTraceID(header http.Header) string { |
| 346 | for _, name := range []string{"trace_id", "trace-id", "x-trace-id"} { |
| 347 | if value := strings.TrimSpace(header.Get(name)); value != "" { |
| 348 | return value |
| 349 | } |
| 350 | } |
| 351 | return "" |
| 352 | } |
| 353 |