diff --git a/pkg/providers/openai_compat/provider.go b/pkg/providers/openai_compat/provider.go index 90cff136c..e8ff7b330 100644 --- a/pkg/providers/openai_compat/provider.go +++ b/pkg/providers/openai_compat/provider.go @@ -122,6 +122,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) diff --git a/pkg/providers/openai_compat/provider_test.go b/pkg/providers/openai_compat/provider_test.go index ab632ccf3..6c0f25d34 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/common" "github.com/sipeed/picoclaw/pkg/providers/protocoltypes" ) @@ -302,7 +304,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", }, { @@ -322,7 +324,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") @@ -340,6 +342,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"}]}` @@ -373,7 +434,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() @@ -519,12 +580,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)