Method RoundTrip

aibridge/keypool/failover.go:60–107

RoundTrip is invoked by the proxy once per outer client request, after Rewrite has applied proxy headers. For centralized requests it walks the key pool, retrying on key-specific failures until one key succeeds or the pool is exhausted. BYOK requests skip the failover loop.

(req *http.Request)

// key-specific failures until one key succeeds or the pool is
// exhausted. BYOK requests skip the failover loop.
func (t *keyFailoverTransport) RoundTrip(req *http.Request) (*http.Response, error) {
	if t.config.IsBYOK(req) {
		return t.inner.RoundTrip(req)
	}

	// Buffer once so retries can replay the body.
	body, err := bufferBody(req)
	if err != nil {
		return nil, err
	}

	// Fresh walker per request, independent of other inflight requests.
	walker := t.config.Pool.Walker()
	for {
		key, keyPoolErr := walker.Next()
		if keyPoolErr != nil {
			resp := t.config.BuildKeyPoolResponse(keyPoolErr)
			if resp == nil {
				// Fallback if BuildKeyPoolResponse returns nil.
				body := []byte(`{"error":"key pool unavailable"}`)
				resp = utils.NewJSONErrorResponse(http.StatusBadGateway, 0, body)
			}
			return resp, nil
		}

		// Clone per attempt so the original request isn't mutated.
		outReq := req.Clone(req.Context())
		if body != nil {
			outReq.Body = io.NopCloser(bytes.NewReader(body))
		}
		t.config.InjectAuthKey(&outReq.Header, key.Value())

		resp, rtErr := t.inner.RoundTrip(outReq)
		if rtErr != nil {
			// Transport-level error, not a key issue.
			return resp, rtErr
		}
		// MarkKeyOnStatus returns true on key-specific failures (e.g. 401/403/429).
		if MarkKeyOnStatus(req.Context(), key, resp, t.config.Logger, t.config.ProviderName) {
			// Drain and retry with the next key.
			_, _ = io.Copy(io.Discard, resp.Body)
			_ = resp.Body.Close()
			continue
		}
		// Success or non-key error, forward as-is.
		return resp, nil
	}
}

// bufferBody reads the request body fully so it can be replayed
// across key-failover retries. Returns nil for a nil body.