cloudwego--eino/adk/retry_chatmodel.go

/*
 * Copyright 2025 CloudWeGo Authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package adk

import (
	"context"
	"errors"
	"fmt"
	"io"
	"math/rand"
	"time"

	"github.com/cloudwego/eino/components/model"
	"github.com/cloudwego/eino/compose"
	"github.com/cloudwego/eino/schema"
)

var (
	// ErrExceedMaxRetries is returned when the maximum number of retries has been exceeded.
	// Use errors.Is to check if an error is due to max retries being exceeded:
	//
	//   if errors.Is(err, adk.ErrExceedMaxRetries) {
	//       // handle max retries exceeded
	//   }
	//
	// Use errors.As to extract the underlying RetryExhaustedError for the last error details:
	//
	//   var retryErr *adk.RetryExhaustedError
	//   if errors.As(err, &retryErr) {
	//       fmt.Printf("last error was: %v\n", retryErr.LastErr)
	//   }
	ErrExceedMaxRetries = errors.New("exceeds max retries")
)

// RetryExhaustedError is returned when all retry attempts have been exhausted.
// It wraps the last error that occurred during retry attempts.
type RetryExhaustedError struct {
	LastErr      error
	TotalRetries int
}

func (e *RetryExhaustedError) Error() string {
	if e.LastErr != nil {
		return fmt.Sprintf("exceeds max retries: last error: %v", e.LastErr)
	}
	return "exceeds max retries"
}

func (e *RetryExhaustedError) Unwrap() error {
	return ErrExceedMaxRetries
}

// WillRetryError is emitted when a retryable error occurs and a retry will be attempted.
// It allows end-users to observe retry events in real-time via AgentEvent.
//
// Field design rationale:
//   - ErrStr (exported): Stores the error message string for Gob serialization during checkpointing.
//     This ensures the error message is preserved after checkpoint restore.
//   - err (unexported): Stores the original error for Unwrap() support at runtime.
//     This field is intentionally unexported because Gob serialization would fail for unregistered
//     concrete error types. Since end-users only need the original error when the AgentEvent first
//     occurs (not after restoring from checkpoint), skipping serialization is acceptable.
//     After checkpoint restore, err will be nil and Unwrap() returns nil.
//   - rejectReason (unexported): Stores a user-defined value set by the ShouldRetry callback
//     via RetryDecision.RejectReason. This is runtime-only observability data — after checkpoint
//     restore it will be nil. Unexported to avoid Gob serialization of arbitrary types.
type WillRetryError struct {
	ErrStr       string
	RetryAttempt int
	rejectReason any
	err          error
}

func (e *WillRetryError) Error() string {
	return e.ErrStr
}

func (e *WillRetryError) Unwrap() error {
	return e.err
}

// RejectReason returns the user-defined rejection reason set by the ShouldRetry callback
// via RetryDecision.RejectReason. Returns nil if not set or after checkpoint restore.
func (e *WillRetryError) RejectReason() any {
	return e.rejectReason
}

func init() {
	schema.RegisterName[*WillRetryError]("eino_adk_chatmodel_will_retry_error")
}

// TypedRetryContext contains context information passed to TypedModelRetryConfig.ShouldRetry
// during a retry decision.
//
// State combinations for OutputMessage and Err:
//
//	OutputMessage != nil, Err == nil  → successful call; inspect message quality
//	OutputMessage == nil, Err != nil  → failed call (Generate error or Stream() error)
//	OutputMessage != nil, Err != nil  → partial stream (chunks received before mid-stream error)
//	OutputMessage == nil, Err == nil  → empty stream (zero chunks before EOF)
type TypedRetryContext[M MessageType] struct {
	// RetryAttempt is the current retry attempt number (1-based).
	// For the first retry decision (after the initial call), this is 1.
	RetryAttempt int

	// InputMessages is the input messages that were sent to the model for the current attempt.
	InputMessages []M

	// Options is the model options that were used for the current attempt.
	Options []model.Option

	// OutputMessage is the output message from the model, if any.
	// This is non-nil when the model returned a message successfully.
	// For streaming, this is the fully concatenated message (the entire stream is consumed
	// before ShouldRetry is called).
	// For streaming with mid-stream errors, this is the partial concatenation of chunks
	// received before the error occurred.
	// May be nil if the model returned an error without producing a message, or if the
	// stream was empty (zero chunks before EOF).
	OutputMessage M

	// Err is the error from the model call, if any.
	// May be nil if the model produced a message without error.
	// Note: both OutputMessage and Err can be nil simultaneously for empty streams.
	Err error
}

// RetryContext is the default retry context type using *schema.Message.
type RetryContext = TypedRetryContext[*schema.Message]

// TypedRetryDecision represents the decision made by TypedModelRetryConfig.ShouldRetry.
type TypedRetryDecision[M MessageType] struct {
	// Retry indicates whether the model call should be retried.
	// If false, the model output (or error) is accepted as-is, unless RewriteError is set.
	Retry bool

	// RewriteError, when non-nil, overrides the return value of the model call with this error.
	// The agent run will fail with this error.
	//
	// This is useful for two scenarios:
	//   - When the model returns a "seemingly correct" message (no error) that actually
	//     contains unrecoverable issues. RewriteError converts the successful output
	//     into a fatal error.
	//   - When the model returns an error, but you want to replace it with a different,
	//     more descriptive error (e.g., adding context or wrapping).
	//
	// When Retry is true, RewriteError is ignored.
	// When Retry is false and RewriteError is non-nil, the model call returns
	// RewriteError regardless of whether the original call had an error or a message.
	RewriteError error

	// ModifiedInputMessages, when non-nil, replaces the input messages for the next retry.
	//
	// This enables advanced recovery strategies like context compression or message trimming.
	// Only used when Retry is true. Ignored when Retry is false.
	ModifiedInputMessages []M

	// PersistModifiedInputMessages controls whether ModifiedInputMessages are written
	// back to the agent's conversation history, affecting subsequent model calls in
	// the agent loop (not just the next retry attempt).
	//
	// When true, the modified messages replace the current conversation history.
	// When false (default), the modified messages are only used for the next retry attempt
	// within this retry cycle.
	//
	// Only used when Retry is true and ModifiedInputMessages is non-nil.
	PersistModifiedInputMessages bool

	// AdditionalOptions, when non-nil, provides additional model options for the next retry.
	// These options are appended to the existing options, taking precedence via last-wins semantics.
	//
	// This enables adjustments like increasing MaxTokens for the retry attempt.
	// Note: options accumulate across retries within a single retry cycle. If ShouldRetry
	// returns AdditionalOptions on every attempt, each set is appended to the previous ones.
	// Only the last value for each option key takes effect, but earlier values remain in the slice.
	// AdditionalOptions are scoped to the current retry cycle and do not persist to subsequent
	// agent iterations — each new model call in the agent loop starts with the original options.
	// Only used when Retry is true. Ignored when Retry is false.
	AdditionalOptions []model.Option

	// Backoff specifies the duration to wait before the next retry attempt.
	// If zero, the default backoff function (from ModelRetryConfig.BackoffFunc or the
	// built-in exponential backoff) is used.
	//
	// This allows the ShouldRetry callback to dynamically control retry timing based on
	// the specific error or problematic message encountered.
	// Only used when Retry is true. Ignored when Retry is false.
	Backoff time.Duration

	// RejectReason is an optional user-defined value describing why the output was rejected.
	// When Retry is true and the rejected stream/message is observed downstream via
	// AgentEvent, this value is attached to the WillRetryError emitted to the event stream.
	// Consumers can retrieve it via WillRetryError.RejectReason().
	//
	// The ShouldRetry callback has full access to the model output (via retryCtx.OutputMessage)
	// and error (via retryCtx.Err), so it can distill whatever information it wants into
	// RejectReason — a string, a struct, the output message itself, or nil.
	//
	// Only used when Retry is true. Ignored when Retry is false.
	RejectReason any
}

// RetryDecision is the default retry decision type using *schema.Message.
type RetryDecision = TypedRetryDecision[*schema.Message]

// TypedModelRetryConfig configures retry behavior for the ChatModel node.
// It defines how the agent should handle transient failures when calling the ChatModel.
type TypedModelRetryConfig[M MessageType] struct {
	// MaxRetries specifies the maximum number of retry attempts.
	// A value of 0 means no retries will be attempted.
	// A value of 3 means up to 3 retry attempts (4 total calls including the initial attempt).
	MaxRetries int

	// ShouldRetry determines how to handle a model call result.
	// It receives context information about the current attempt including the output message
	// and/or error, and returns a decision on whether to retry, what to modify, etc.
	// Returning nil is treated as &RetryDecision{Retry: false} (accept as-is).
	//
	// If nil, defaults to retrying on any non-nil error (backward compatible with IsRetryAble).
	//
	// Note: When ShouldRetry is set, IsRetryAble is ignored.
	// Note: In streaming mode, the entire stream is consumed before ShouldRetry is called.
	// The event stream is sent to the client in real time regardless; only the retry
	// decision is deferred until the full response is available.
	ShouldRetry func(ctx context.Context, retryCtx *TypedRetryContext[M]) *TypedRetryDecision[M]

	// Deprecated: Use ShouldRetry instead for richer retry control including message
	// inspection, input modification, and option adjustment. When ShouldRetry is set,
	// IsRetryAble is ignored.
	IsRetryAble func(ctx context.Context, err error) bool

	// BackoffFunc calculates the delay before the next retry attempt.
	// The attempt parameter starts at 1 for the first retry.
	// Used as the default when RetryDecision.Backoff is zero.
	// If nil, a default exponential backoff with jitter is used:
	// base delay 100ms, exponentially increasing up to 10s max,
	// with random jitter (0-50% of delay) to prevent thundering herd.
	BackoffFunc func(ctx context.Context, attempt int) time.Duration
}

// ModelRetryConfig is the default retry config type using *schema.Message.
type ModelRetryConfig = TypedModelRetryConfig[*schema.Message]

func defaultIsRetryAble(_ context.Context, err error) bool {
	return err != nil
}

func defaultBackoff(_ context.Context, attempt int) time.Duration {
	baseDelay := 100 * time.Millisecond
	maxDelay := 10 * time.Second

	if attempt <= 0 {
		return baseDelay
	}

	if attempt > 7 {
		return maxDelay + time.Duration(rand.Int63n(int64(maxDelay/2)))
	}

	delay := baseDelay * time.Duration(1<<uint(attempt-1))
	if delay > maxDelay {
		delay = maxDelay
	}

	jitter := time.Duration(rand.Int63n(int64(delay / 2)))
	return delay + jitter
}

func genErrWrapper(ctx context.Context, maxRetries, attempt int, isRetryAbleFunc func(ctx context.Context, err error) bool) func(error) error {
	return func(err error) error {
		isRetryAble := isRetryAbleFunc == nil || isRetryAbleFunc(ctx, err)
		hasRetriesLeft := attempt < maxRetries

		if isRetryAble && hasRetriesLeft {
			return &WillRetryError{ErrStr: err.Error(), RetryAttempt: attempt, err: err}
		}
		return err
	}
}

func consumeStreamForError[M any](stream *schema.StreamReader[M]) error {
	defer stream.Close()
	for {
		_, err := stream.Recv()
		if err == io.EOF {
			return nil
		}
		if err != nil {
			return err
		}
	}
}

type retryVerdictSignal struct {
	ch chan retryVerdict
}

type retryVerdict struct {
	WillRetry    bool
	RetryAttempt int
	Err          error
	RejectReason any
}

// retryModelWrapper wraps a BaseChatModel with retry logic.
// This is used inside the model wrapper chain, positioned between eventSenderModelWrapper
// and stateModelWrapper, so that retry only affects the inner chain (event sending, user wrappers,
// callback injection) without re-running state management (BeforeModelRewriteState/AfterModelRewriteState).
type typedRetryModelWrapper[M MessageType] struct {
	inner  model.BaseModel[M]
	config *TypedModelRetryConfig[M]
}

func newTypedRetryModelWrapper[M MessageType](inner model.BaseModel[M], config *TypedModelRetryConfig[M]) *typedRetryModelWrapper[M] {
	return &typedRetryModelWrapper[M]{inner: inner, config: config}
}

func (r *typedRetryModelWrapper[M]) Generate(ctx context.Context, input []M, opts ...model.Option) (M, error) {
	if r.config.ShouldRetry != nil {
		return generateWithShouldRetry(r, ctx, input, opts...)
	}
	return r.generateLegacy(ctx, input, opts...)
}

func (r *typedRetryModelWrapper[M]) generateLegacy(ctx context.Context, input []M, opts ...model.Option) (zero M, _ error) {
	isRetryAble := r.config.IsRetryAble
	if isRetryAble == nil {
		isRetryAble = defaultIsRetryAble
	}
	backoffFunc := r.config.BackoffFunc
	if backoffFunc == nil {
		backoffFunc = defaultBackoff
	}

	var lastErr error
	for attempt := 0; attempt <= r.config.MaxRetries; attempt++ {
		out, err := r.inner.Generate(ctx, input, opts...)
		if err == nil {
			return out, nil
		}

		if _, ok := compose.ExtractInterruptInfo(err); ok {
			return zero, err
		}

		if errors.Is(err, ErrStreamCanceled) {
			return zero, err
		}

		if !isRetryAble(ctx, err) {
			return zero, err
		}

		lastErr = err
		if attempt < r.config.MaxRetries {
			if err := r.contextAwareSleep(ctx, backoffFunc(ctx, attempt+1)); err != nil {
				return zero, err
			}
		}
	}

	return zero, &RetryExhaustedError{LastErr: lastErr, TotalRetries: r.config.MaxRetries}
}

func generateWithShouldRetry[M MessageType](r *typedRetryModelWrapper[M], ctx context.Context, input []M, opts ...model.Option) (M, error) {
	backoffFunc := r.config.BackoffFunc
	if backoffFunc == nil {
		backoffFunc = defaultBackoff
	}

	execCtx := getTypedChatModelAgentExecCtx[M](ctx)

	currentInput := input
	currentOpts := opts
	var lastErr error
	var zero M

	defer func() {
		_ = compose.ProcessState(ctx, func(_ context.Context, st *typedState[M]) error {
			st.setRetryAttempt(0)
			return nil
		})
	}()

	for attempt := 0; attempt <= r.config.MaxRetries; attempt++ {
		_ = compose.ProcessState(ctx, func(_ context.Context, st *typedState[M]) error {
			st.setRetryAttempt(attempt)
			return nil
		})

		// Suppress event sending during Generate: the ShouldRetry callback must decide whether
		// to accept or reject the result before any event is emitted. If accepted, the event
		// is sent explicitly below (lines after decision check). If rejected, no event leaks.
		if execCtx != nil {
			execCtx.suppressEventSend = true
		}
		out, err := r.inner.Generate(ctx, currentInput, currentOpts...)
		if execCtx != nil {
			execCtx.suppressEventSend = false
		}

		if err != nil {
			if _, ok := compose.ExtractInterruptInfo(err); ok {
				return zero, err
			}

			if errors.Is(err, ErrStreamCanceled) {
				return zero, err
			}
		}

		retryCtx := &TypedRetryContext[M]{
			RetryAttempt:  attempt + 1,
			InputMessages: currentInput,
			Options:       currentOpts,
			OutputMessage: out,
			Err:           err,
		}
		decision := r.config.ShouldRetry(ctx, retryCtx)
		if decision == nil {
			decision = &TypedRetryDecision[M]{}
		}

		if !decision.Retry {
			if decision.RewriteError != nil {
				return zero, decision.RewriteError
			}
			if err != nil {
				return zero, err
			}
			if execCtx != nil && execCtx.generator != nil && out != nil {
				event := typedModelOutputEvent(out, nil)
				execCtx.send(event)
			}
			return out, nil
		}

		lastErr = err
		if lastErr == nil {
			lastErr = fmt.Errorf("model output rejected by ShouldRetry at attempt %d", attempt+1)
		}

		if attempt >= r.config.MaxRetries {
			break
		}

		applyDecisionForRetry(&currentInput, &currentOpts, ctx, decision)

		delay := decision.Backoff
		if delay == 0 {
			delay = backoffFunc(ctx, attempt+1)
		}

		if err := r.contextAwareSleep(ctx, delay); err != nil {
			return zero, err
		}
	}

	return zero, &RetryExhaustedError{LastErr: lastErr, TotalRetries: r.config.MaxRetries}
}

func (r *typedRetryModelWrapper[M]) contextAwareSleep(ctx context.Context, delay time.Duration) error {
	if delay <= 0 {
		return nil
	}
	select {
	case <-ctx.Done():
		return ctx.Err()
	case <-time.After(delay):
		return nil
	}
}

func streamWithShouldRetry[M MessageType](r *typedRetryModelWrapper[M], ctx context.Context, input []M, opts ...model.Option) (
	*schema.StreamReader[M], error) {

	backoffFunc := r.config.BackoffFunc
	if backoffFunc == nil {
		backoffFunc = defaultBackoff
	}

	defer func() {
		_ = compose.ProcessState(ctx, func(_ context.Context, st *typedState[M]) error {
			st.setRetryAttempt(0)
			return nil
		})
	}()

	execCtx := getTypedChatModelAgentExecCtx[M](ctx)

	currentInput := input
	currentOpts := opts
	var lastErr error
	var curSignal *retryVerdictSignal

	// Panic recovery for verdict signal: if ShouldRetry panics, the onEOF/errWrapper closures in
	// buildStreamConvertOptions will block forever on signal.ch, causing a goroutine leak. This
	// defer ensures a verdict is always sent, even on panic, before re-panicking.
	defer func() {
		if p := recover(); p != nil {
			if curSignal != nil {
				select {
				case curSignal.ch <- retryVerdict{WillRetry: false, Err: fmt.Errorf("panic: %v", p)}:
				default:
				}
			}
			panic(p)
		}
	}()

	for attempt := 0; attempt <= r.config.MaxRetries; attempt++ {
		_ = compose.ProcessState(ctx, func(_ context.Context, st *typedState[M]) error {
			st.setRetryAttempt(attempt)
			return nil
		})

		signal := &retryVerdictSignal{ch: make(chan retryVerdict, 1)}
		curSignal = signal
		if execCtx != nil {
			execCtx.retryVerdictSignal = signal
		}

		stream, err := r.inner.Stream(ctx, currentInput, currentOpts...)
		if err != nil {
			// Defensive no-op: when Stream() returns an error, no stream exists, so
			// eventSenderModel never creates the StreamReaderWithConvert hooks that would
			// read from signal.ch. This send has no consumer — it merely fills the
			// buffered(1) slot so the panic-recovery defer (select/default) won't block
			// if a later panic tries to send a second verdict. The signal is discarded
			// when the next iteration creates a new one.
			signal.ch <- retryVerdict{WillRetry: false}

			if _, ok := compose.ExtractInterruptInfo(err); ok {
				return nil, err
			}

			if errors.Is(err, ErrStreamCanceled) {
				return nil, err
			}

			retryCtx := &TypedRetryContext[M]{
				RetryAttempt:  attempt + 1,
				InputMessages: currentInput,
				Options:       currentOpts,
				Err:           err,
			}
			decision := r.config.ShouldRetry(ctx, retryCtx)
			if decision == nil {
				decision = &TypedRetryDecision[M]{}
			}

			if !decision.Retry {
				if decision.RewriteError != nil {
					return nil, decision.RewriteError
				}
				return nil, err
			}

			lastErr = err
			if attempt < r.config.MaxRetries {
				applyDecisionForRetry(&currentInput, &currentOpts, ctx, decision)
				delay := decision.Backoff
				if delay == 0 {
					delay = backoffFunc(ctx, attempt+1)
				}
				if err := r.contextAwareSleep(ctx, delay); err != nil {
					return nil, err
				}
			}
			continue
		}

		// Split the stream: checkCopy is consumed synchronously here to build the complete
		// message for ShouldRetry inspection; returnCopy is returned to the caller and may
		// already be consumed downstream in parallel. The verdict signal bridges the two:
		// once ShouldRetry decides, the signal tells returnCopy's errWrapper/onEOF whether
		// to pass through normally or inject a WillRetryError.
		copies := stream.Copy(2)
		checkCopy := copies[0]
		returnCopy := copies[1]

		msg, streamErr := typedConsumeStream(checkCopy)

		if errors.Is(streamErr, ErrStreamCanceled) {
			signal.ch <- retryVerdict{WillRetry: false}
			returnCopy.Close()
			return nil, streamErr
		}

		retryCtx := &TypedRetryContext[M]{
			RetryAttempt:  attempt + 1,
			InputMessages: currentInput,
			Options:       currentOpts,
			OutputMessage: msg,
			Err:           streamErr,
		}
		decision := r.config.ShouldRetry(ctx, retryCtx)
		if decision == nil {
			decision = &TypedRetryDecision[M]{}
		}

		if !decision.Retry {
			signal.ch <- retryVerdict{WillRetry: false}

			if decision.RewriteError != nil {
				returnCopy.Close()
				return nil, decision.RewriteError
			}
			if streamErr != nil {
				returnCopy.Close()
				return nil, streamErr
			}
			return returnCopy, nil
		}

		verdictErr := streamErr
		if verdictErr == nil {
			verdictErr = fmt.Errorf("model output rejected by ShouldRetry at attempt %d", attempt+1)
		}
		signal.ch <- retryVerdict{
			WillRetry:    true,
			RetryAttempt: attempt,
			Err:          verdictErr,
			RejectReason: decision.RejectReason,
		}
		returnCopy.Close()

		lastErr = verdictErr

		if attempt < r.config.MaxRetries {
			applyDecisionForRetry(&currentInput, &currentOpts, ctx, decision)
			delay := decision.Backoff
			if delay == 0 {
				delay = backoffFunc(ctx, attempt+1)
			}
			if err := r.contextAwareSleep(ctx, delay); err != nil {
				return nil, err
			}
		}
	}

	return nil, &RetryExhaustedError{LastErr: lastErr, TotalRetries: r.config.MaxRetries}
}

func applyDecisionForRetry[M MessageType](currentInput *[]M, currentOpts *[]model.Option, ctx context.Context, decision *TypedRetryDecision[M]) {
	if decision.ModifiedInputMessages != nil {
		*currentInput = decision.ModifiedInputMessages
		if decision.PersistModifiedInputMessages {
			modifiedInput := *currentInput
			_ = compose.ProcessState(ctx, func(_ context.Context, st *typedState[M]) error {
				st.Messages = modifiedInput
				return nil
			})
		}
	}

	if decision.AdditionalOptions != nil {
		cloned := make([]model.Option, len(*currentOpts), len(*currentOpts)+len(decision.AdditionalOptions))
		copy(cloned, *currentOpts)
		*currentOpts = append(cloned, decision.AdditionalOptions...)
	}
}

func (r *typedRetryModelWrapper[M]) Stream(ctx context.Context, input []M, opts ...model.Option) (
	*schema.StreamReader[M], error) {

	if r.config.ShouldRetry != nil {
		return streamWithShouldRetry(r, ctx, input, opts...)
	}
	return r.streamLegacy(ctx, input, opts...)
}

func (r *typedRetryModelWrapper[M]) streamLegacy(ctx context.Context, input []M, opts ...model.Option) (
	*schema.StreamReader[M], error) {

	isRetryAble := r.config.IsRetryAble
	if isRetryAble == nil {
		isRetryAble = defaultIsRetryAble
	}
	backoffFunc := r.config.BackoffFunc
	if backoffFunc == nil {
		backoffFunc = defaultBackoff
	}

	defer func() {
		_ = compose.ProcessState(ctx, func(_ context.Context, st *typedState[M]) error {
			st.setRetryAttempt(0)
			return nil
		})
	}()

	var lastErr error
	for attempt := 0; attempt <= r.config.MaxRetries; attempt++ {
		_ = compose.ProcessState(ctx, func(_ context.Context, st *typedState[M]) error {
			st.setRetryAttempt(attempt)
			return nil
		})

		stream, err := r.inner.Stream(ctx, input, opts...)
		if err != nil {
			if _, ok := compose.ExtractInterruptInfo(err); ok {
				return nil, err
			}
			if errors.Is(err, ErrStreamCanceled) {
				return nil, err
			}
			if !isRetryAble(ctx, err) {
				return nil, err
			}
			lastErr = err
			if attempt < r.config.MaxRetries {
				if err := r.contextAwareSleep(ctx, backoffFunc(ctx, attempt+1)); err != nil {
					return nil, err
				}
			}
			continue
		}

		copies := stream.Copy(2)
		checkCopy := copies[0]
		returnCopy := copies[1]

		streamErr := consumeStreamForError(checkCopy)
		if streamErr == nil {
			return returnCopy, nil
		}

		returnCopy.Close()
		if errors.Is(streamErr, ErrStreamCanceled) {
			return nil, streamErr
		}
		if !isRetryAble(ctx, streamErr) {
			return nil, streamErr
		}

		lastErr = streamErr
		if attempt < r.config.MaxRetries {
			if err := r.contextAwareSleep(ctx, backoffFunc(ctx, attempt+1)); err != nil {
				return nil, err
			}
		}
	}

	return nil, &RetryExhaustedError{LastErr: lastErr, TotalRetries: r.config.MaxRetries}
}