From f164e0d3a2b578f8de329b7c1cd4610eeba86438 Mon Sep 17 00:00:00 2001 From: tong3jie <14191774+tong3jie@users.noreply.github.com> Date: Wed, 11 Mar 2026 17:57:13 +0800 Subject: [PATCH] fix(provider): improve error handling and add retry tests for HTML responses --- pkg/providers/openai_compat/provider.go | 11 +-- pkg/providers/openai_compat/provider_test.go | 75 ++++++++++++++++++-- 2 files changed, 78 insertions(+), 8 deletions(-) diff --git a/pkg/providers/openai_compat/provider.go b/pkg/providers/openai_compat/provider.go index 2b9f40e17..4781e35e5 100644 --- a/pkg/providers/openai_compat/provider.go +++ b/pkg/providers/openai_compat/provider.go @@ -14,6 +14,7 @@ import ( "time" "github.com/hashicorp/go-retryablehttp" + "github.com/sipeed/picoclaw/pkg/providers/protocoltypes" ) @@ -87,7 +88,6 @@ func WithRetry(maxRetries int, minWait, maxWait time.Duration) Option { } func NewProvider(apiKey, apiBase, proxy string, opts ...Option) *Provider { - retryClient := retryablehttp.NewClient() retryClient.RetryMax = 3 retryClient.RetryWaitMin = 1 * time.Second @@ -97,7 +97,6 @@ func NewProvider(apiKey, apiBase, proxy string, opts ...Option) *Provider { transport := &http.Transport{} if proxy != "" { - if parsed, err := url.Parse(proxy); err == nil { transport.Proxy = http.ProxyURL(parsed) } else { @@ -114,6 +113,8 @@ func NewProvider(apiKey, apiBase, proxy string, opts ...Option) *Provider { httpClient: retryClient.StandardClient(), } + p.httpClient.Timeout = defaultRequestTimeout + for _, opt := range opts { if opt != nil { opt(p) @@ -236,7 +237,7 @@ 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) + // 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. @@ -268,7 +269,9 @@ func (p *Provider) resolveMaxTokenField(model string) string { } lowerModel := strings.ToLower(model) - if strings.Contains(lowerModel, "o1") || strings.Contains(lowerModel, "glm-4") || strings.Contains(lowerModel, "gpt-5") { + if strings.Contains(lowerModel, "o1") || + strings.Contains(lowerModel, "glm-4") || + strings.Contains(lowerModel, "gpt-5") { return "max_completion_tokens" } return "max_tokens" diff --git a/pkg/providers/openai_compat/provider_test.go b/pkg/providers/openai_compat/provider_test.go index 9a3a7acc5..a75582c96 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" ) @@ -252,7 +254,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", }, { @@ -272,7 +274,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") @@ -290,6 +292,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"}]}` @@ -323,7 +384,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() @@ -454,12 +515,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)