From 297f64103b9ed687edcbefc63380f733adfbfe7d Mon Sep 17 00:00:00 2001
From: tong3jie <14191774+tong3jie@users.noreply.github.com>
Date: Wed, 11 Mar 2026 14:09:07 +0800
Subject: [PATCH] feat(provider): add retry logic and vendor prefix handling in
 OpenAI provider

---
 go.mod                                  |  2 +
 go.sum                                  |  4 ++
 pkg/providers/common/common.go          | 15 +++-
 pkg/providers/openai_compat/provider.go | 93 ++++++++++++++++++-------
 4 files changed, 88 insertions(+), 26 deletions(-)

diff --git a/go.mod b/go.mod
index d283f7f5e..4af65a823 100644
--- a/go.mod
+++ b/go.mod
@@ -47,6 +47,7 @@ require (
 	github.com/elliotchance/orderedmap/v3 v3.1.0 // indirect
 	github.com/gdamore/encoding v1.0.1 // indirect
 	github.com/godbus/dbus/v5 v5.1.0 // indirect
+	github.com/hashicorp/go-cleanhttp v0.5.2 // indirect
 	github.com/inconshreveable/mousetrap v1.1.0 // indirect
 	github.com/lucasb-eyer/go-colorful v1.3.0 // indirect
 	github.com/mattn/go-colorable v0.1.14 // indirect
@@ -81,6 +82,7 @@ require (
 	github.com/gogo/protobuf v1.3.2 // indirect
 	github.com/google/jsonschema-go v0.4.2 // indirect
 	github.com/grbit/go-json v0.11.0 // indirect
+	github.com/hashicorp/go-retryablehttp v0.7.8
 	github.com/klauspost/compress v1.18.4 // indirect
 	github.com/klauspost/cpuid/v2 v2.3.0 // indirect
 	github.com/tidwall/gjson v1.18.0 // indirect
diff --git a/go.sum b/go.sum
index f24b997d4..d1ed7eac6 100644
--- a/go.sum
+++ b/go.sum
@@ -108,6 +108,10 @@ github.com/grbit/go-json v0.11.0 h1:bAbyMdYrYl/OjYsSqLH99N2DyQ291mHy726Mx+sYrnc=
 github.com/grbit/go-json v0.11.0/go.mod h1:IYpHsdybQ386+6g3VE6AXQ3uTGa5mquBme5/ZWmtzek=
 github.com/h2non/filetype v1.1.3 h1:FKkx9QbD7HR/zjK1Ia5XiBsq9zdLi5Kf3zGyFTAFkGg=
 github.com/h2non/filetype v1.1.3/go.mod h1:319b3zT68BvV+WRj7cwy856M2ehB3HqNOt6sy1HndBY=
+github.com/hashicorp/go-cleanhttp v0.5.2 h1:035FKYIWjmULyFRBKPs8TBQoi0x6d9G4xc9neXJWAZQ=
+github.com/hashicorp/go-cleanhttp v0.5.2/go.mod h1:kO/YDlP8L1346E6Sodw+PrpBSV4/SoxCXGY6BqNFT48=
+github.com/hashicorp/go-retryablehttp v0.7.8 h1:ylXZWnqa7Lhqpk0L1P1LzDtGcCR0rPVUrx/c8Unxc48=
+github.com/hashicorp/go-retryablehttp v0.7.8/go.mod h1:rjiScheydd+CxvumBsIrFKlx3iS0jrZ7LvzFGFmuKbw=
 github.com/hashicorp/golang-lru/v2 v2.0.7 h1:a+bsQ5rvGLjzHuww6tVxozPZFVghXaHOwFs4luLUK2k=
 github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM=
 github.com/hpcloud/tail v1.0.0/go.mod h1:ab1qPbhIpdTxEkNHXyeSf5vhxWSCs/tWer42PpOxQnU=
diff --git a/pkg/providers/common/common.go b/pkg/providers/common/common.go
index 90142fb8b..67fcae8aa 100644
--- a/pkg/providers/common/common.go
+++ b/pkg/providers/common/common.go
@@ -36,7 +36,12 @@ type (
 	ReasoningDetail        = protocoltypes.ReasoningDetail
 )
 
-const DefaultRequestTimeout = 120 * time.Second
+const (
+	DefaultRequestTimeout = 120 * time.Second
+
+	// set response size limit to prevent OOM if server returns a huge response (e.g., an HTML error page instead of JSON)
+	maxResponseSize = 10 * 1024 * 1024
+)
 
 // NewHTTPClient creates an *http.Client with an optional proxy and the default timeout.
 func NewHTTPClient(proxy string) *http.Client {
@@ -316,7 +321,9 @@ func HandleErrorResponse(resp *http.Response, apiBase string) error {
 // then parses the JSON response into an LLMResponse.
 func ReadAndParseResponse(resp *http.Response, apiBase string) (*LLMResponse, error) {
 	contentType := resp.Header.Get("Content-Type")
-	reader := bufio.NewReader(resp.Body)
+
+	safeReader := io.LimitReader(resp.Body, maxResponseSize)
+	reader := bufio.NewReader(safeReader)
 	prefix, err := reader.Peek(256)
 	if err != nil && err != io.EOF && err != bufio.ErrBufferFull {
 		return nil, fmt.Errorf("failed to inspect response: %w", err)
@@ -326,6 +333,10 @@ func ReadAndParseResponse(resp *http.Response, apiBase string) (*LLMResponse, er
 	}
 	out, err := ParseResponse(reader)
 	if err != nil {
+		// some APIs return 200 with an HTML error page, so check for that before giving up on JSON parsing
+		if LooksLikeHTML(prefix, contentType) {
+			return nil, WrapHTMLResponseError(resp.StatusCode, prefix, contentType, apiBase)
+		}
 		return nil, fmt.Errorf("failed to parse JSON response: %w", err)
 	}
 	return out, nil
diff --git a/pkg/providers/openai_compat/provider.go b/pkg/providers/openai_compat/provider.go
index 90bc683b8..90cff136c 100644
--- a/pkg/providers/openai_compat/provider.go
+++ b/pkg/providers/openai_compat/provider.go
@@ -8,11 +8,14 @@ import (
 	"fmt"
 	"io"
 	"log"
+	"maps"
 	"net/http"
 	"net/url"
 	"strings"
 	"time"
 
+	"github.com/hashicorp/go-retryablehttp"
+
 	"github.com/sipeed/picoclaw/pkg/providers/common"
 	"github.com/sipeed/picoclaw/pkg/providers/protocoltypes"
 )
@@ -30,6 +33,21 @@ type (
 	ReasoningDetail        = protocoltypes.ReasoningDetail
 )
 
+var vendorPrefixes = map[string]struct{}{
+	"litellm":    {},
+	"moonshot":   {},
+	"nvidia":     {},
+	"groq":       {},
+	"ollama":     {},
+	"deepseek":   {},
+	"google":     {},
+	"openrouter": {},
+	"zhipu":      {},
+	"mistral":    {},
+	"vivgrid":    {},
+	"minimax":    {},
+}
+
 type Provider struct {
 	apiKey         string
 	apiBase        string
@@ -62,11 +80,46 @@ func WithExtraBody(extraBody map[string]any) Option {
 	}
 }
 
+func WithRetry(maxRetries int, minWait, maxWait time.Duration) Option {
+	return func(p *Provider) {
+		if rc, ok := p.httpClient.Transport.(*retryablehttp.RoundTripper); ok {
+			if maxRetries >= 0 {
+				rc.Client.RetryMax = maxRetries
+			}
+			if minWait > 0 {
+				rc.Client.RetryWaitMin = minWait
+			}
+			if maxWait > 0 {
+				rc.Client.RetryWaitMax = maxWait
+			}
+		}
+	}
+}
+
 func NewProvider(apiKey, apiBase, proxy string, opts ...Option) *Provider {
+	retryClient := retryablehttp.NewClient()
+	retryClient.RetryMax = 3
+	retryClient.RetryWaitMin = 1 * time.Second
+	retryClient.RetryWaitMax = 30 * time.Second
+	retryClient.Backoff = retryablehttp.LinearJitterBackoff
+	retryClient.Logger = nil
+
+	transport := &http.Transport{}
+	if proxy != "" {
+		if parsed, err := url.Parse(proxy); err == nil {
+			transport.Proxy = http.ProxyURL(parsed)
+		} else {
+			log.Printf("openai_compat: invalid proxy URL %q: %v", proxy, err)
+		}
+	}
+
+	retryClient.HTTPClient.Transport = transport
+	retryClient.HTTPClient.Timeout = defaultRequestTimeout
+
 	p := &Provider{
 		apiKey:     apiKey,
 		apiBase:    strings.TrimRight(apiBase, "/"),
-		httpClient: common.NewHTTPClient(proxy),
+		httpClient: retryClient.StandardClient(),
 	}
 
 	for _, opt := range opts {
@@ -75,6 +128,10 @@ func NewProvider(apiKey, apiBase, proxy string, opts ...Option) *Provider {
 		}
 	}
 
+	if p.maxTokensField == "" {
+		p.maxTokensField = "max_tokens"
+	}
+
 	return p
 }
 
@@ -115,17 +172,7 @@ func (p *Provider) buildRequestBody(
 	}
 
 	if maxTokens, ok := common.AsInt(options["max_tokens"]); ok {
-		fieldName := p.maxTokensField
-		if fieldName == "" {
-			lowerModel := strings.ToLower(model)
-			if strings.Contains(lowerModel, "glm") || strings.Contains(lowerModel, "o1") ||
-				strings.Contains(lowerModel, "gpt-5") {
-				fieldName = "max_completion_tokens"
-			} else {
-				fieldName = "max_tokens"
-			}
-		}
-		requestBody[fieldName] = maxTokens
+		requestBody[p.maxTokensField] = maxTokens
 	}
 
 	if temperature, ok := common.AsFloat(options["temperature"]); ok {
@@ -149,8 +196,8 @@ func (p *Provider) buildRequestBody(
 
 	// Merge extra body fields configured per-provider/model.
 	// These are injected last so they take precedence over defaults.
-	for k, v := range p.extraBody {
-		requestBody[k] = v
+	if p.extraBody != nil {
+		maps.Copy(requestBody, p.extraBody)
 	}
 
 	return requestBody
@@ -190,6 +237,7 @@ func (p *Provider) Chat(
 	}
 	defer resp.Body.Close()
 
+	// Non-200: read a prefix to tell HTML error page apart from JSON error body.
 	if resp.StatusCode != http.StatusOK {
 		return nil, common.HandleErrorResponse(resp, p.apiBase)
 	}
@@ -387,23 +435,20 @@ func parseStreamResponse(
 }
 
 func normalizeModel(model, apiBase string) string {
+	if strings.Contains(strings.ToLower(apiBase), "openrouter.ai") {
+		return model
+	}
+
 	before, after, ok := strings.Cut(model, "/")
 	if !ok {
 		return model
 	}
 
-	if strings.Contains(strings.ToLower(apiBase), "openrouter.ai") {
-		return model
+	if _, exists := vendorPrefixes[strings.ToLower(before)]; exists {
+		return after
 	}
 
-	prefix := strings.ToLower(before)
-	switch prefix {
-	case "litellm", "moonshot", "nvidia", "groq", "ollama", "deepseek", "google",
-		"openrouter", "zhipu", "mistral", "vivgrid", "minimax", "novita":
-		return after
-	default:
-		return model
-	}
+	return model
 }
 
 func buildToolsList(tools []ToolDefinition, nativeSearch bool) []any {