From e89913a28bdef19853ddcd4761d9bb1b36f27fe4 Mon Sep 17 00:00:00 2001 From: xiongfei <278837230@qq.com> Date: Tue, 31 Mar 2026 14:43:52 +0800 Subject: [PATCH] feat:add openrouter appinfo --- pkg/config/useragent_transport.go | 49 +++++ pkg/config/useragent_transport_test.go | 105 ++++++++++ pkg/providers/anthropic_messages/provider.go | 4 +- pkg/providers/bedrock/provider_bedrock.go | 5 +- .../bedrock/provider_bedrock_test.go | 12 +- pkg/providers/common/common.go | 13 +- pkg/providers/common/common_test.go | 11 +- pkg/providers/github_copilot_provider.go | 4 +- pkg/providers/openai_compat/provider.go | 192 ++++++++++++++---- pkg/providers/openai_compat/provider_test.go | 154 +++++++++++++- pkg/skills/clawhub_registry.go | 5 +- pkg/skills/installer_test.go | 7 +- pkg/utils/http_client.go | 24 ++- pkg/utils/http_client_test.go | 17 +- 14 files changed, 523 insertions(+), 79 deletions(-) create mode 100644 pkg/config/useragent_transport.go create mode 100644 pkg/config/useragent_transport_test.go diff --git a/pkg/config/useragent_transport.go b/pkg/config/useragent_transport.go new file mode 100644 index 000000000..e02f8d4d0 --- /dev/null +++ b/pkg/config/useragent_transport.go @@ -0,0 +1,49 @@ +package config + +import ( + "net/http" + "strings" +) + +// HTTPUserAgent is the default User-Agent for outbound HTTP requests (e.g. "PicoClaw/0.2.4"). +func HTTPUserAgent() string { + v := strings.TrimSpace(Version) + if v == "" { + v = "dev" + } + return "PicoClaw/" + v +} + +// userAgentTransport wraps an http.RoundTripper and sets User-Agent to HTTPUserAgent when +// the request does not already specify one. +type userAgentTransport struct { + base http.RoundTripper + ua string +} + +// WrapTransportUserAgent wraps base so requests without User-Agent receive HTTPUserAgent(). +// If base is nil, http.DefaultTransport is used. +func WrapTransportUserAgent(base http.RoundTripper) http.RoundTripper { + if base == nil { + base = http.DefaultTransport + } + return &userAgentTransport{base: base, ua: HTTPUserAgent()} +} + +// UnwrapUserAgent returns the inner RoundTripper if rt was produced by WrapTransportUserAgent; otherwise rt. +func UnwrapUserAgent(rt http.RoundTripper) http.RoundTripper { + if t, ok := rt.(*userAgentTransport); ok { + return t.base + } + return rt +} + +// RoundTrip implements http.RoundTripper. +func (t *userAgentTransport) RoundTrip(req *http.Request) (*http.Response, error) { + if req.Header.Get("User-Agent") != "" { + return t.base.RoundTrip(req) + } + r2 := req.Clone(req.Context()) + r2.Header.Set("User-Agent", t.ua) + return t.base.RoundTrip(r2) +} diff --git a/pkg/config/useragent_transport_test.go b/pkg/config/useragent_transport_test.go new file mode 100644 index 000000000..51676acc9 --- /dev/null +++ b/pkg/config/useragent_transport_test.go @@ -0,0 +1,105 @@ +package config + +import ( + "net/http" + "net/http/httptest" + "strings" + "testing" +) + +func TestHTTPUserAgent_format(t *testing.T) { + ua := HTTPUserAgent() + if !strings.HasPrefix(ua, "PicoClaw/") { + t.Fatalf("want PicoClaw/ prefix, got %q", ua) + } + suffix := strings.TrimPrefix(ua, "PicoClaw/") + if strings.TrimSpace(suffix) == "" { + t.Fatal("empty version suffix") + } +} + +func TestWrapTransportUserAgent_SetsDefaultUA(t *testing.T) { + var gotUA string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotUA = r.Header.Get("User-Agent") + w.WriteHeader(http.StatusOK) + })) + t.Cleanup(srv.Close) + + client := &http.Client{Transport: WrapTransportUserAgent(http.DefaultTransport)} + req, err := http.NewRequest(http.MethodGet, srv.URL, nil) + if err != nil { + t.Fatal(err) + } + resp, err := client.Do(req) + if err != nil { + t.Fatal(err) + } + resp.Body.Close() + + want := HTTPUserAgent() + if gotUA != want { + t.Fatalf("User-Agent = %q, want %q", gotUA, want) + } +} + +func TestWrapTransportUserAgent_PreservesExplicitUA(t *testing.T) { + var gotUA string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotUA = r.Header.Get("User-Agent") + w.WriteHeader(http.StatusOK) + })) + t.Cleanup(srv.Close) + + client := &http.Client{Transport: WrapTransportUserAgent(http.DefaultTransport)} + req, err := http.NewRequest(http.MethodGet, srv.URL, nil) + if err != nil { + t.Fatal(err) + } + req.Header.Set("User-Agent", "custom-agent/1") + resp, err := client.Do(req) + if err != nil { + t.Fatal(err) + } + resp.Body.Close() + + if gotUA != "custom-agent/1" { + t.Fatalf("User-Agent = %q, want custom-agent/1", gotUA) + } +} + +func TestWrapTransportUserAgent_DoesNotMutateOriginalRequest(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + })) + t.Cleanup(srv.Close) + + rt := WrapTransportUserAgent(http.DefaultTransport) + req, err := http.NewRequest(http.MethodGet, srv.URL, nil) + if err != nil { + t.Fatal(err) + } + if req.Header.Get("User-Agent") != "" { + t.Fatal("expected empty User-Agent on original request") + } + resp, err := rt.RoundTrip(req) + if err != nil { + t.Fatal(err) + } + resp.Body.Close() + + if req.Header.Get("User-Agent") != "" { + t.Fatal("RoundTrip must not mutate the original request's headers") + } +} + +func TestUnwrapUserAgent(t *testing.T) { + inner := http.DefaultTransport + wrapped := WrapTransportUserAgent(inner) + if UnwrapUserAgent(wrapped) != inner { + t.Fatal("UnwrapUserAgent should return inner transport") + } + if UnwrapUserAgent(inner) != inner { + t.Fatal("UnwrapUserAgent on non-wrapped should return same") + } +} diff --git a/pkg/providers/anthropic_messages/provider.go b/pkg/providers/anthropic_messages/provider.go index 6a1c473dd..83303fa43 100644 --- a/pkg/providers/anthropic_messages/provider.go +++ b/pkg/providers/anthropic_messages/provider.go @@ -16,6 +16,7 @@ import ( "strings" "time" + "github.com/sipeed/picoclaw/pkg/config" "github.com/sipeed/picoclaw/pkg/providers/protocoltypes" ) @@ -60,7 +61,8 @@ func NewProviderWithTimeout(apiKey, apiBase string, timeoutSeconds int) *Provide apiKey: apiKey, apiBase: baseURL, httpClient: &http.Client{ - Timeout: timeout, + Timeout: timeout, + Transport: config.WrapTransportUserAgent(nil), }, } } diff --git a/pkg/providers/bedrock/provider_bedrock.go b/pkg/providers/bedrock/provider_bedrock.go index 9ca29455f..3798c5fd8 100644 --- a/pkg/providers/bedrock/provider_bedrock.go +++ b/pkg/providers/bedrock/provider_bedrock.go @@ -208,7 +208,10 @@ func (p *Provider) Chat( if err != nil { // Check for SSO token expiration errors and provide actionable guidance if isSSOTokenError(err) { - return nil, fmt.Errorf("bedrock converse: AWS credentials may have expired. If using AWS SSO, run 'aws sso login' to refresh: %w", err) + return nil, fmt.Errorf( + "bedrock converse: AWS credentials may have expired. If using AWS SSO, run 'aws sso login' to refresh: %w", + err, + ) } return nil, fmt.Errorf("bedrock converse: %w", err) } diff --git a/pkg/providers/bedrock/provider_bedrock_test.go b/pkg/providers/bedrock/provider_bedrock_test.go index 882c2971c..38a5e26da 100644 --- a/pkg/providers/bedrock/provider_bedrock_test.go +++ b/pkg/providers/bedrock/provider_bedrock_test.go @@ -583,13 +583,17 @@ func TestIsSSOTokenError(t *testing.T) { expected: true, }, { - name: "full SSO error message", - err: fmt.Errorf("get identity: get credentials: failed to refresh cached credentials, refresh cached SSO token failed, unable to refresh SSO token"), + name: "full SSO error message", + err: fmt.Errorf( + "get identity: get credentials: failed to refresh cached credentials, refresh cached SSO token failed, unable to refresh SSO token", + ), expected: true, }, { - name: "SSO token file missing", - err: fmt.Errorf("get identity: get credentials: failed to refresh cached credentials, failed to read cached SSO token file, open ~/.aws/sso/cache/abc123.json: no such file or directory"), + name: "SSO token file missing", + err: fmt.Errorf( + "get identity: get credentials: failed to refresh cached credentials, failed to read cached SSO token file, open ~/.aws/sso/cache/abc123.json: no such file or directory", + ), expected: true, }, } diff --git a/pkg/providers/common/common.go b/pkg/providers/common/common.go index 90142fb8b..4595eacc3 100644 --- a/pkg/providers/common/common.go +++ b/pkg/providers/common/common.go @@ -19,6 +19,7 @@ import ( "strings" "time" + "github.com/sipeed/picoclaw/pkg/config" "github.com/sipeed/picoclaw/pkg/providers/protocoltypes" ) @@ -43,17 +44,18 @@ func NewHTTPClient(proxy string) *http.Client { client := &http.Client{ Timeout: DefaultRequestTimeout, } + base := http.DefaultTransport if proxy != "" { parsed, err := url.Parse(proxy) if err == nil { // Preserve http.DefaultTransport settings (TLS, HTTP/2, timeouts, etc.) - if base, ok := http.DefaultTransport.(*http.Transport); ok { - tr := base.Clone() - tr.Proxy = http.ProxyURL(parsed) - client.Transport = tr + if tr, ok := http.DefaultTransport.(*http.Transport); ok { + clone := tr.Clone() + clone.Proxy = http.ProxyURL(parsed) + base = clone } else { // Fallback: minimal transport if DefaultTransport is not *http.Transport. - client.Transport = &http.Transport{ + base = &http.Transport{ Proxy: http.ProxyURL(parsed), } } @@ -61,6 +63,7 @@ func NewHTTPClient(proxy string) *http.Client { log.Printf("common: invalid proxy URL %q: %v", proxy, err) } } + client.Transport = config.WrapTransportUserAgent(base) return client } diff --git a/pkg/providers/common/common_test.go b/pkg/providers/common/common_test.go index 79a637d48..e9bf59037 100644 --- a/pkg/providers/common/common_test.go +++ b/pkg/providers/common/common_test.go @@ -8,6 +8,7 @@ import ( "strings" "testing" + "github.com/sipeed/picoclaw/pkg/config" "github.com/sipeed/picoclaw/pkg/providers/protocoltypes" ) @@ -22,9 +23,10 @@ func TestNewHTTPClient_DefaultTimeout(t *testing.T) { func TestNewHTTPClient_WithProxy(t *testing.T) { client := NewHTTPClient("http://127.0.0.1:8080") - transport, ok := client.Transport.(*http.Transport) + base := config.UnwrapUserAgent(client.Transport) + transport, ok := base.(*http.Transport) if !ok || transport == nil { - t.Fatalf("expected http.Transport with proxy, got %T", client.Transport) + t.Fatalf("expected http.Transport with proxy, got %T", base) } req := &http.Request{URL: &url.URL{Scheme: "https", Host: "api.example.com"}} gotProxy, err := transport.Proxy(req) @@ -38,8 +40,9 @@ func TestNewHTTPClient_WithProxy(t *testing.T) { func TestNewHTTPClient_NoProxy(t *testing.T) { client := NewHTTPClient("") - if client.Transport != nil { - t.Errorf("expected nil transport without proxy, got %T", client.Transport) + base := config.UnwrapUserAgent(client.Transport) + if base == nil { + t.Fatal("expected non-nil base transport") } } diff --git a/pkg/providers/github_copilot_provider.go b/pkg/providers/github_copilot_provider.go index e2d1d7d98..472c14257 100644 --- a/pkg/providers/github_copilot_provider.go +++ b/pkg/providers/github_copilot_provider.go @@ -41,9 +41,9 @@ func NewGitHubCopilotProvider(uri string, connectMode string, model string) (*Gi } session, err := client.CreateSession(context.Background(), &copilot.SessionConfig{ - Model: model, + Model: model, OnPermissionRequest: copilot.PermissionHandler.ApproveAll, - Hooks: &copilot.SessionHooks{}, + Hooks: &copilot.SessionHooks{}, }) if err != nil { client.Stop() diff --git a/pkg/providers/openai_compat/provider.go b/pkg/providers/openai_compat/provider.go index 90bc683b8..a3a724dbe 100644 --- a/pkg/providers/openai_compat/provider.go +++ b/pkg/providers/openai_compat/provider.go @@ -13,6 +13,7 @@ import ( "strings" "time" + "github.com/sipeed/picoclaw/pkg/logger" "github.com/sipeed/picoclaw/pkg/providers/common" "github.com/sipeed/picoclaw/pkg/providers/protocoltypes" ) @@ -42,6 +43,63 @@ type Option func(*Provider) const defaultRequestTimeout = common.DefaultRequestTimeout +// OpenRouter free-tier models often return transient HTTP 429; retry a few times before surfacing. +const ( + openRouter429MaxAttempts = 10 + openRouter429Backoff = time.Second +) + +// OpenRouter app attribution (optional headers for rankings/analytics). +// See https://openrouter.ai/docs/app-attribution +const ( + openRouterAttributionReferer = "https://picoclaw.io/" + openRouterAttributionTitle = "PicoClaw" + openRouterAttributionCategories = "personal-agent,general-chat" +) + +func isOpenRouterHost(apiBase string) bool { + return strings.Contains(strings.ToLower(strings.TrimSpace(apiBase)), "openrouter.ai") +} + +func openRouterLogKind(stream bool) string { + if stream { + return "ChatStream" + } + return "Chat" +} + +func logOpenRouter429Retry(stream bool, resolvedModel string, attempt, maxAttempts int) { + logger.WarnC("agent", + fmt.Sprintf("openai_compat OpenRouter %s: HTTP 429 (attempt %d/%d), backing off %v then retry (model=%q)", + openRouterLogKind(stream), attempt, maxAttempts, openRouter429Backoff, resolvedModel, + )) +} + +func logOpenRouter429Exhausted(stream bool, resolvedModel string, maxAttempts int) { + logger.WarnC("agent", + fmt.Sprintf("openai_compat OpenRouter %s: HTTP 429 after %d attempts, giving up (model=%q)", + openRouterLogKind(stream), maxAttempts, resolvedModel, + )) +} + +func logOpenRouter429BackoffCancelled(stream bool, resolvedModel string, attempt, maxAttempts int, cancelErr error) { + log.Printf( + "openai_compat OpenRouter %s: HTTP 429 retry cancel during backoff (attempt %d/%d, model=%q): %v", + openRouterLogKind(stream), attempt, maxAttempts, resolvedModel, cancelErr, + ) +} + +// applyOpenRouterAttributionHeaders sets Referer, X-OpenRouter-Title, and X-OpenRouter-Categories +// when the API base is OpenRouter. Uses the standard Referer header name (HTTP-Referer in OpenRouter docs). +func applyOpenRouterAttributionHeaders(h http.Header, apiBase string) { + if !isOpenRouterHost(apiBase) { + return + } + h.Set("Referer", openRouterAttributionReferer) + h.Set("X-Openrouter-Title", openRouterAttributionTitle) + h.Set("X-Openrouter-Categories", openRouterAttributionCategories) +} + func WithMaxTokensField(maxTokensField string) Option { return func(p *Provider) { p.maxTokensField = maxTokensField @@ -174,27 +232,56 @@ func (p *Provider) Chat( return nil, fmt.Errorf("failed to marshal request: %w", err) } - req, err := http.NewRequestWithContext(ctx, "POST", p.apiBase+"/chat/completions", bytes.NewReader(jsonData)) - if err != nil { - return nil, fmt.Errorf("failed to create request: %w", err) + maxAttempts := 1 + if isOpenRouterHost(p.apiBase) { + maxAttempts = openRouter429MaxAttempts } + resolvedModel := normalizeModel(model, p.apiBase) - req.Header.Set("Content-Type", "application/json") - if p.apiKey != "" { - req.Header.Set("Authorization", "Bearer "+p.apiKey) + for attempt := 1; attempt <= maxAttempts; attempt++ { + req, err := http.NewRequestWithContext(ctx, "POST", p.apiBase+"/chat/completions", bytes.NewReader(jsonData)) + if err != nil { + return nil, fmt.Errorf("failed to create request: %w", err) + } + + req.Header.Set("Content-Type", "application/json") + if p.apiKey != "" { + req.Header.Set("Authorization", "Bearer "+p.apiKey) + } + applyOpenRouterAttributionHeaders(req.Header, p.apiBase) + + resp, err := p.httpClient.Do(req) + if err != nil { + return nil, fmt.Errorf("failed to send request: %w", err) + } + + if resp.StatusCode == http.StatusOK { + out, readErr := common.ReadAndParseResponse(resp, p.apiBase) + resp.Body.Close() + return out, readErr + } + + if resp.StatusCode == http.StatusTooManyRequests && isOpenRouterHost(p.apiBase) && attempt < maxAttempts { + logOpenRouter429Retry(false, resolvedModel, attempt, maxAttempts) + _, _ = io.Copy(io.Discard, resp.Body) + resp.Body.Close() + select { + case <-ctx.Done(): + logOpenRouter429BackoffCancelled(false, resolvedModel, attempt, maxAttempts, ctx.Err()) + return nil, ctx.Err() + case <-time.After(openRouter429Backoff): + } + continue + } + + if resp.StatusCode == http.StatusTooManyRequests && isOpenRouterHost(p.apiBase) { + logOpenRouter429Exhausted(false, resolvedModel, maxAttempts) + } + err = common.HandleErrorResponse(resp, p.apiBase) + resp.Body.Close() + return nil, err } - - resp, err := p.httpClient.Do(req) - if err != nil { - return nil, fmt.Errorf("failed to send request: %w", err) - } - defer resp.Body.Close() - - if resp.StatusCode != http.StatusOK { - return nil, common.HandleErrorResponse(resp, p.apiBase) - } - - return common.ReadAndParseResponse(resp, p.apiBase) + return nil, fmt.Errorf("internal error: chat request loop exited without return") } // ChatStream implements streaming via OpenAI-compatible SSE (stream: true). @@ -219,32 +306,59 @@ func (p *Provider) ChatStream( return nil, fmt.Errorf("failed to marshal request: %w", err) } - req, err := http.NewRequestWithContext(ctx, "POST", p.apiBase+"/chat/completions", bytes.NewReader(jsonData)) - if err != nil { - return nil, fmt.Errorf("failed to create request: %w", err) + maxAttempts := 1 + if isOpenRouterHost(p.apiBase) { + maxAttempts = openRouter429MaxAttempts } + resolvedModel := normalizeModel(model, p.apiBase) - req.Header.Set("Content-Type", "application/json") - req.Header.Set("Accept", "text/event-stream") - if p.apiKey != "" { - req.Header.Set("Authorization", "Bearer "+p.apiKey) - } - - // Use a client without Timeout for streaming — the http.Client.Timeout covers - // the entire request lifecycle including body reads, which would kill long streams. - // Context cancellation still provides the safety net. streamClient := &http.Client{Transport: p.httpClient.Transport} - resp, err := streamClient.Do(req) - if err != nil { - return nil, fmt.Errorf("failed to send request: %w", err) - } - defer resp.Body.Close() - if resp.StatusCode != http.StatusOK { - return nil, common.HandleErrorResponse(resp, p.apiBase) - } + for attempt := 1; attempt <= maxAttempts; attempt++ { + req, err := http.NewRequestWithContext(ctx, "POST", p.apiBase+"/chat/completions", bytes.NewReader(jsonData)) + if err != nil { + return nil, fmt.Errorf("failed to create request: %w", err) + } - return parseStreamResponse(ctx, resp.Body, onChunk) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "text/event-stream") + if p.apiKey != "" { + req.Header.Set("Authorization", "Bearer "+p.apiKey) + } + applyOpenRouterAttributionHeaders(req.Header, p.apiBase) + + resp, err := streamClient.Do(req) + if err != nil { + return nil, fmt.Errorf("failed to send request: %w", err) + } + + if resp.StatusCode == http.StatusOK { + out, streamErr := parseStreamResponse(ctx, resp.Body, onChunk) + resp.Body.Close() + return out, streamErr + } + + if resp.StatusCode == http.StatusTooManyRequests && isOpenRouterHost(p.apiBase) && attempt < maxAttempts { + logOpenRouter429Retry(true, resolvedModel, attempt, maxAttempts) + _, _ = io.Copy(io.Discard, resp.Body) + resp.Body.Close() + select { + case <-ctx.Done(): + logOpenRouter429BackoffCancelled(true, resolvedModel, attempt, maxAttempts, ctx.Err()) + return nil, ctx.Err() + case <-time.After(openRouter429Backoff): + } + continue + } + + if resp.StatusCode == http.StatusTooManyRequests && isOpenRouterHost(p.apiBase) { + logOpenRouter429Exhausted(true, resolvedModel, maxAttempts) + } + err = common.HandleErrorResponse(resp, p.apiBase) + resp.Body.Close() + return nil, err + } + return nil, fmt.Errorf("internal error: chat stream loop exited without return") } // parseStreamResponse parses an OpenAI-compatible SSE stream. diff --git a/pkg/providers/openai_compat/provider_test.go b/pkg/providers/openai_compat/provider_test.go index ab632ccf3..f22b63fe5 100644 --- a/pkg/providers/openai_compat/provider_test.go +++ b/pkg/providers/openai_compat/provider_test.go @@ -12,6 +12,7 @@ import ( "testing" "time" + "github.com/sipeed/picoclaw/pkg/config" "github.com/sipeed/picoclaw/pkg/providers/common" "github.com/sipeed/picoclaw/pkg/providers/protocoltypes" ) @@ -519,9 +520,9 @@ 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) + transport, ok := config.UnwrapUserAgent(p.httpClient.Transport).(*http.Transport) if !ok || transport == nil { - t.Fatalf("expected http transport with proxy, got %T", p.httpClient.Transport) + t.Fatalf("expected http transport with proxy, got %T", config.UnwrapUserAgent(p.httpClient.Transport)) } req := &http.Request{URL: &url.URL{Scheme: "https", Host: "api.example.com"}} @@ -1173,3 +1174,152 @@ func TestSerializeMessages_StripsSystemParts(t *testing.T) { t.Fatal("system_parts should not appear in serialized output") } } + +func TestProviderChat_OpenRouter429RetriesThenSucceeds(t *testing.T) { + var n int + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if !strings.HasSuffix(r.URL.Path, "/chat/completions") { + http.NotFound(w, r) + return + } + n++ + if n < 3 { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusTooManyRequests) + _, _ = w.Write([]byte(`{"error":{"message":"rate limit"}}`)) + return + } + resp := map[string]any{ + "choices": []map[string]any{ + {"message": map[string]any{"content": "ok"}, "finish_reason": "stop"}, + }, + } + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(resp) + })) + defer server.Close() + + apiBase := server.URL + "/via-openrouter.ai" + p := NewProvider("key", apiBase, "") + start := time.Now() + out, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "m", nil) + if err != nil { + t.Fatalf("Chat() error = %v", err) + } + if out == nil || out.Content != "ok" { + t.Fatalf("unexpected response: %+v", out) + } + if n != 3 { + t.Fatalf("request count = %d, want 3", n) + } + if d := time.Since(start); d < 2*time.Second-100*time.Millisecond { + t.Fatalf("expected ~2s backoff between retries, got %v", d) + } +} + +func TestProviderChat_OpenRouter429Exhausted(t *testing.T) { + var n int + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if !strings.HasSuffix(r.URL.Path, "/chat/completions") { + http.NotFound(w, r) + return + } + n++ + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusTooManyRequests) + _, _ = w.Write([]byte(`{"error":{"message":"rate limit"}}`)) + })) + defer server.Close() + + apiBase := server.URL + "/via-openrouter.ai" + p := NewProvider("key", apiBase, "") + _, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "m", nil) + if err == nil { + t.Fatal("expected error after 429 exhaustion") + } + if !strings.Contains(err.Error(), "429") { + t.Fatalf("error should mention 429: %v", err) + } + if n != 10 { + t.Fatalf("request count = %d, want 3", n) + } +} + +func TestProviderChat_OpenRouterSendsAttributionHeaders(t *testing.T) { + var gotReferer, gotTitle, gotCat string + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if !strings.HasSuffix(r.URL.Path, "/chat/completions") { + http.NotFound(w, r) + return + } + gotReferer = r.Header.Get("Referer") + gotTitle = r.Header.Get("X-OpenRouter-Title") + gotCat = r.Header.Get("X-OpenRouter-Categories") + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]any{ + "choices": []map[string]any{ + {"message": map[string]any{"content": "ok"}, "finish_reason": "stop"}, + }, + }) + })) + defer server.Close() + + apiBase := server.URL + "/via-openrouter.ai" + p := NewProvider("key", apiBase, "") + _, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "m", nil) + if err != nil { + t.Fatalf("Chat: %v", err) + } + if gotReferer != openRouterAttributionReferer { + t.Errorf("Referer = %q, want %q", gotReferer, openRouterAttributionReferer) + } + if gotTitle != openRouterAttributionTitle { + t.Errorf("X-OpenRouter-Title = %q, want %q", gotTitle, openRouterAttributionTitle) + } + if gotCat != openRouterAttributionCategories { + t.Errorf("X-OpenRouter-Categories = %q, want %q", gotCat, openRouterAttributionCategories) + } +} + +func TestProviderChat_NonOpenRouterOmitsAttributionHeaders(t *testing.T) { + var gotReferer string + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotReferer = r.Header.Get("Referer") + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]any{ + "choices": []map[string]any{ + {"message": map[string]any{"content": "ok"}, "finish_reason": "stop"}, + }, + }) + })) + defer server.Close() + + p := NewProvider("key", server.URL, "") + _, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "m", nil) + if err != nil { + t.Fatalf("Chat: %v", err) + } + if gotReferer != "" { + t.Errorf("non-OpenRouter base should not set Referer, got %q", gotReferer) + } +} + +func TestProviderChat_NonOpenRouterSingle429NoRetry(t *testing.T) { + var n int + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + n++ + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusTooManyRequests) + _, _ = w.Write([]byte(`{"error":{"message":"slow down"}}`)) + })) + defer server.Close() + + p := NewProvider("key", server.URL, "") + _, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "m", nil) + if err == nil { + t.Fatal("expected error") + } + if n != 1 { + t.Fatalf("request count = %d, want 1 (no OpenRouter 429 retry)", n) + } +} diff --git a/pkg/skills/clawhub_registry.go b/pkg/skills/clawhub_registry.go index bd4bed8fb..35b3ed0ae 100644 --- a/pkg/skills/clawhub_registry.go +++ b/pkg/skills/clawhub_registry.go @@ -10,6 +10,7 @@ import ( "os" "time" + "github.com/sipeed/picoclaw/pkg/config" "github.com/sipeed/picoclaw/pkg/utils" ) @@ -75,11 +76,11 @@ func NewClawHubRegistry(cfg ClawHubConfig) *ClawHubRegistry { maxResponseSize: maxResp, client: &http.Client{ Timeout: timeout, - Transport: &http.Transport{ + Transport: config.WrapTransportUserAgent(&http.Transport{ MaxIdleConns: 5, IdleConnTimeout: 30 * time.Second, TLSHandshakeTimeout: 10 * time.Second, - }, + }), }, } } diff --git a/pkg/skills/installer_test.go b/pkg/skills/installer_test.go index 759cfc489..bc26735c0 100644 --- a/pkg/skills/installer_test.go +++ b/pkg/skills/installer_test.go @@ -10,6 +10,8 @@ import ( "strings" "testing" "time" + + "github.com/sipeed/picoclaw/pkg/config" ) func TestParseGitHubRef(t *testing.T) { @@ -224,9 +226,10 @@ func TestNewSkillInstaller_WithProxy(t *testing.T) { } // Verify the transport has proxy configured - transport, ok := installer.client.Transport.(*http.Transport) + base := config.UnwrapUserAgent(installer.client.Transport) + transport, ok := base.(*http.Transport) if !ok { - t.Fatal("client.Transport is not *http.Transport") + t.Fatalf("client base transport is not *http.Transport, got %T", base) } if transport.Proxy == nil { diff --git a/pkg/utils/http_client.go b/pkg/utils/http_client.go index bda7c5c83..1c5ea6c26 100644 --- a/pkg/utils/http_client.go +++ b/pkg/utils/http_client.go @@ -6,20 +6,19 @@ import ( "net/url" "strings" "time" + + "github.com/sipeed/picoclaw/pkg/config" ) // CreateHTTPClient creates an HTTP client with optional proxy support. // If proxyURL is empty, it uses the system environment proxy settings. // Supported proxy schemes: http, https, socks5, socks5h. func CreateHTTPClient(proxyURL string, timeout time.Duration) (*http.Client, error) { - client := &http.Client{ - Timeout: timeout, - Transport: &http.Transport{ - MaxIdleConns: 10, - IdleConnTimeout: 30 * time.Second, - DisableCompression: false, - TLSHandshakeTimeout: 15 * time.Second, - }, + tr := &http.Transport{ + MaxIdleConns: 10, + IdleConnTimeout: 30 * time.Second, + DisableCompression: false, + TLSHandshakeTimeout: 15 * time.Second, } if proxyURL != "" { @@ -39,10 +38,13 @@ func CreateHTTPClient(proxyURL string, timeout time.Duration) (*http.Client, err if proxy.Host == "" { return nil, fmt.Errorf("invalid proxy URL: missing host") } - client.Transport.(*http.Transport).Proxy = http.ProxyURL(proxy) + tr.Proxy = http.ProxyURL(proxy) } else { - client.Transport.(*http.Transport).Proxy = http.ProxyFromEnvironment + tr.Proxy = http.ProxyFromEnvironment } - return client, nil + return &http.Client{ + Timeout: timeout, + Transport: config.WrapTransportUserAgent(tr), + }, nil } diff --git a/pkg/utils/http_client_test.go b/pkg/utils/http_client_test.go index ff3d0429b..e31899d99 100644 --- a/pkg/utils/http_client_test.go +++ b/pkg/utils/http_client_test.go @@ -5,6 +5,8 @@ import ( "strings" "testing" "time" + + "github.com/sipeed/picoclaw/pkg/config" ) func TestCreateHTTPClient_ProxyConfigured(t *testing.T) { @@ -16,9 +18,10 @@ func TestCreateHTTPClient_ProxyConfigured(t *testing.T) { t.Fatalf("client.Timeout = %v, want %v", client.Timeout, 12*time.Second) } - tr, ok := client.Transport.(*http.Transport) + base := config.UnwrapUserAgent(client.Transport) + tr, ok := base.(*http.Transport) if !ok { - t.Fatalf("client.Transport type = %T, want *http.Transport", client.Transport) + t.Fatalf("base transport type = %T, want *http.Transport", base) } if tr.Proxy == nil { t.Fatal("transport.Proxy is nil, want non-nil") @@ -50,9 +53,10 @@ func TestCreateHTTPClient_Socks5ProxyConfigured(t *testing.T) { t.Fatalf("createHTTPClient() error: %v", err) } - tr, ok := client.Transport.(*http.Transport) + base := config.UnwrapUserAgent(client.Transport) + tr, ok := base.(*http.Transport) if !ok { - t.Fatalf("client.Transport type = %T, want *http.Transport", client.Transport) + t.Fatalf("base transport type = %T, want *http.Transport", base) } req, err := http.NewRequest("GET", "https://example.com", nil) if err != nil { @@ -92,9 +96,10 @@ func TestCreateHTTPClient_ProxyFromEnvironmentWhenConfigEmpty(t *testing.T) { t.Fatalf("createHTTPClient() error: %v", err) } - tr, ok := client.Transport.(*http.Transport) + base := config.UnwrapUserAgent(client.Transport) + tr, ok := base.(*http.Transport) if !ok { - t.Fatalf("client.Transport type = %T, want *http.Transport", client.Transport) + t.Fatalf("base transport type = %T, want *http.Transport", base) } if tr.Proxy == nil { t.Fatal("transport.Proxy is nil, want proxy function from environment")