fix(errors): improve check for when to retry requests (#164)

Andrey Nering and yuguorui created

* Follow OpenAI's Go SDK.
* Check status code >= 500.
* Check `x-should-retry` header.

Co-authored-by: yuguorui <yuguorui@pku.edu.cn>

Change summary

errors.go | 33 +++++++++++++++++++++++++++++++--
1 file changed, 31 insertions(+), 2 deletions(-)

Detailed changes

errors.go 🔗

@@ -3,7 +3,9 @@ package fantasy
 import (
 	"errors"
 	"fmt"
+	"io"
 	"net/http"
+	"strconv"
 	"strings"
 
 	"github.com/charmbracelet/x/exp/slice"
@@ -51,9 +53,36 @@ func (m *ProviderError) Error() string {
 	return fmt.Sprintf("%s: %s", m.Title, m.Message)
 }
 
-// IsRetryable checks if the error is retryable based on the status code.
+// IsRetryable reports whether the error should be retried.
+// It returns true if the underlying cause is io.ErrUnexpectedEOF, if the
+// "x-should-retry" response header evaluates to true, or if the HTTP status
+// code indicates a retryable condition (408, 409, 429, or any 5xx).
 func (m *ProviderError) IsRetryable() bool {
-	return m.StatusCode == http.StatusRequestTimeout || m.StatusCode == http.StatusConflict || m.StatusCode == http.StatusTooManyRequests
+	// We're mostly mimicking OpenAI's Go SDK here:
+	// https://github.com/openai/openai-go/blob/b9d280a37149430982e9dfeed16c41d27d45cfc5/internal/requestconfig/requestconfig.go#L244
+	if errors.Is(m.Cause, io.ErrUnexpectedEOF) {
+		return true
+	}
+	if m.shouldRetryHeader() {
+		return true
+	}
+	return m.StatusCode == http.StatusRequestTimeout ||
+		m.StatusCode == http.StatusConflict ||
+		m.StatusCode == http.StatusTooManyRequests ||
+		m.StatusCode >= http.StatusInternalServerError
+}
+
+func (m *ProviderError) shouldRetryHeader() bool {
+	if m.ResponseHeaders == nil {
+		return false
+	}
+	for k, v := range m.ResponseHeaders {
+		if strings.EqualFold(k, "x-should-retry") {
+			b, _ := strconv.ParseBool(v)
+			return b
+		}
+	}
+	return false
 }
 
 // IsContextTooLarge checks if the error is due to the context exceeding the model's limit.