diff --git a/go.mod b/go.mod index f29ef7207..0116e492d 100644 --- a/go.mod +++ b/go.mod @@ -43,6 +43,7 @@ require ( github.com/dustin/go-humanize v1.0.1 // indirect github.com/elliotchance/orderedmap/v3 v3.1.0 // indirect github.com/gdamore/encoding v1.0.1 // 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 @@ -78,6 +79,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 addbab56c..ab2b6b1f6 100644 --- a/go.sum +++ b/go.sum @@ -100,6 +100,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/openai_compat/provider.go b/pkg/providers/openai_compat/provider.go index f97bf3acd..f045f782d 100644 --- a/pkg/providers/openai_compat/provider.go +++ b/pkg/providers/openai_compat/provider.go @@ -13,6 +13,8 @@ import ( "strings" "time" + "github.com/hashicorp/go-retryablehttp" + "github.com/sipeed/picoclaw/pkg/providers/protocoltypes" ) @@ -29,6 +31,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 @@ -54,28 +71,50 @@ func WithRequestTimeout(timeout time.Duration) Option { } } -func NewProvider(apiKey, apiBase, proxy string, opts ...Option) *Provider { - client := &http.Client{ - Timeout: defaultRequestTimeout, - } - - if proxy != "" { - parsed, err := url.Parse(proxy) - if err == nil { - client.Transport = &http.Transport{ - Proxy: http.ProxyURL(parsed), +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: client, + httpClient: retryClient.StandardClient(), } + p.httpClient.Timeout = defaultRequestTimeout + for _, opt := range opts { if opt != nil { opt(p) @@ -126,18 +165,7 @@ func (p *Provider) Chat( } if maxTokens, ok := asInt(options["max_tokens"]); ok { - // Use configured maxTokensField if specified, otherwise fallback to model-based detection - fieldName := p.maxTokensField - if fieldName == "" { - // Fallback: detect from model name for backward compatibility - 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" - } - } + fieldName := p.resolveMaxTokenField(model) requestBody[fieldName] = maxTokens } @@ -185,6 +213,13 @@ func (p *Provider) Chat( } defer resp.Body.Close() + select { + case <-ctx.Done(): + return nil, ctx.Err() + default: + // continue processing response + } + contentType := resp.Header.Get("Content-Type") // Non-200: read a prefix to tell HTML error page apart from JSON error body. @@ -203,8 +238,12 @@ func (p *Provider) Chat( ) } + // set response size limit to prevent OOM if server returns a huge response (e.g., an HTML error page instead of JSON) + const maxResponseSize = 10 * 1024 * 1024 + // Peek without consuming so the full stream reaches the JSON decoder. - reader := bufio.NewReader(resp.Body) + safeReader := io.LimitReader(resp.Body, maxResponseSize) + reader := bufio.NewReader(safeReader) prefix, err := reader.Peek(256) // io.EOF/ErrBufferFull are normal; only real errors abort if err != nil && err != io.EOF && err != bufio.ErrBufferFull { return nil, fmt.Errorf("failed to inspect response: %w", err) @@ -215,12 +254,30 @@ func (p *Provider) Chat( 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, p.apiBase) + } return nil, fmt.Errorf("failed to parse JSON response: %w", err) } return out, nil } +func (p *Provider) resolveMaxTokenField(model string) string { + if p.maxTokensField != "" { + return p.maxTokensField + } + + lowerModel := strings.ToLower(model) + if strings.Contains(lowerModel, "o1") || + strings.Contains(lowerModel, "glm-4") || + strings.Contains(lowerModel, "gpt-5") { + return "max_completion_tokens" + } + return "max_tokens" +} + func wrapHTMLResponseError(statusCode int, body []byte, contentType, apiBase string) error { respPreview := responsePreview(body, 128) return fmt.Errorf( @@ -457,23 +514,20 @@ func serializeMessages(messages []Message) []any { } 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": - return after - default: - return model - } + return model } func asInt(v any) (int, bool) { diff --git a/pkg/providers/openai_compat/provider_test.go b/pkg/providers/openai_compat/provider_test.go index 41f278a1b..70e1c9c3e 100644 --- a/pkg/providers/openai_compat/provider_test.go +++ b/pkg/providers/openai_compat/provider_test.go @@ -12,6 +12,8 @@ import ( "testing" "time" + "github.com/hashicorp/go-retryablehttp" + "github.com/sipeed/picoclaw/pkg/providers/protocoltypes" ) @@ -301,7 +303,7 @@ func TestProviderChat_HTMLResponsesReturnHelpfulError(t *testing.T) { { name: "html error response", contentType: "text/html; charset=utf-8", - statusCode: http.StatusBadGateway, + statusCode: http.StatusOK, body: "bad gateway", }, { @@ -321,7 +323,7 @@ func TestProviderChat_HTMLResponsesReturnHelpfulError(t *testing.T) { })) defer server.Close() - p := NewProvider("key", server.URL, "") + p := NewProvider("key", server.URL, "", WithRetry(0, time.Second*1, time.Second*3)) _, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "gpt-4o", nil) if err == nil { t.Fatal("expected error, got nil") @@ -339,6 +341,65 @@ func TestProviderChat_HTMLResponsesReturnHelpfulError(t *testing.T) { } } +func TestProviderChatErrorPost(t *testing.T) { + tests := []struct { + name string + statusCode int + body string + reTry [3]int + }{ + { + name: "retry on 502 with html response", + statusCode: http.StatusBadGateway, + body: "gateway login", + reTry: [3]int{0, 2, 3}, + }, + { + name: "retry times is 0", + statusCode: http.StatusBadGateway, + body: "bad gateway", + reTry: [3]int{1, 0, 0}, + }, + { + name: "default retry 3 times with 502 and html response", + statusCode: http.StatusBadGateway, + body: " \r\n\tgateway login", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(tt.statusCode) + _, _ = w.Write([]byte(tt.body)) + })) + defer server.Close() + + p := NewProvider( + "key", + server.URL, + "", + WithRetry(tt.reTry[0], time.Second*time.Duration(tt.reTry[1]), time.Second*time.Duration(tt.reTry[2])), + ) + _, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "gpt-4o", nil) + if err == nil { + t.Fatal("expected error, got nil") + } + if !strings.Contains(err.Error(), "failed to send request") { + t.Fatalf("expected status code in error, got %v", err) + } + times := 1 + if tt.reTry[0] != 0 { + times = tt.reTry[0] + 1 + } + + if !strings.Contains(err.Error(), fmt.Sprintf("giving up after %d attempt(s)", times)) { + t.Fatalf("expected retry count in error, got %v", err) + } + }) + } +} + func TestProviderChat_SuccessResponseUsesStreamingDecoder(t *testing.T) { content := strings.Repeat("a", 1024) body := `{"choices":[{"message":{"content":"` + content + `"},"finish_reason":"stop"}]}` @@ -372,7 +433,7 @@ func TestProviderChat_LargeHTMLResponsePreviewIsTruncated(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "text/html; charset=utf-8") - w.WriteHeader(http.StatusBadGateway) + w.WriteHeader(http.StatusOK) _, _ = w.Write(body) })) defer server.Close() @@ -503,12 +564,18 @@ func TestProvider_ProxyConfigured(t *testing.T) { proxyURL := "http://127.0.0.1:8080" p := NewProvider("key", "https://example.com", proxyURL) - transport, ok := p.httpClient.Transport.(*http.Transport) + transports, ok := p.httpClient.Transport.(*retryablehttp.RoundTripper) + if !ok || transports == nil { + t.Fatalf("expected retryablehttp transport, got %T", p.httpClient.Transport) + } + + transport, ok := transports.Client.HTTPClient.Transport.(*http.Transport) if !ok || transport == nil { t.Fatalf("expected http transport with proxy, got %T", p.httpClient.Transport) } req := &http.Request{URL: &url.URL{Scheme: "https", Host: "api.example.com"}} + gotProxy, err := transport.Proxy(req) if err != nil { t.Fatalf("proxy function returned error: %v", err)