From 222d1a30864ab451221a54309747d73794fbb90c Mon Sep 17 00:00:00 2001 From: Petrichor Date: Fri, 27 Feb 2026 16:35:07 +0800 Subject: [PATCH 1/4] refactor(modernize): apply safe modernize fixes --- pkg/agent/context.go | 7 +++---- pkg/agent/context_cache_test.go | 4 ++-- pkg/agent/loop_test.go | 17 +++-------------- pkg/agent/memory.go | 2 +- pkg/channels/onebot.go | 9 +++------ pkg/channels/wecom.go | 2 +- pkg/channels/wecom_app_test.go | 2 +- pkg/channels/wecom_test.go | 2 +- pkg/config/model_config_test.go | 12 +++++------- pkg/health/server.go | 5 ++--- pkg/providers/anthropic/provider.go | 10 +++++----- pkg/providers/codex_provider.go | 4 ++-- pkg/providers/cooldown_test.go | 4 ++-- pkg/providers/openai_compat/provider.go | 8 ++++---- pkg/routing/agent_id_test.go | 10 ++++++---- pkg/skills/loader.go | 2 +- pkg/skills/search_cache.go | 4 ++-- pkg/skills/search_cache_test.go | 4 ++-- pkg/state/state_test.go | 4 ++-- pkg/tools/cron.go | 8 +++++--- pkg/tools/registry_test.go | 2 +- pkg/tools/web.go | 8 ++++---- pkg/utils/message.go | 20 ++++---------------- 23 files changed, 62 insertions(+), 88 deletions(-) diff --git a/pkg/agent/context.go b/pkg/agent/context.go index b7c6e1108..6fccbaf53 100644 --- a/pkg/agent/context.go +++ b/pkg/agent/context.go @@ -7,6 +7,7 @@ import ( "os" "path/filepath" "runtime" + "slices" "strings" "sync" "time" @@ -249,10 +250,8 @@ func (cb *ContextBuilder) sourceFilesChangedLocked() bool { } // Check tracked source files (bootstrap + memory). - for _, p := range cb.sourcePaths() { - if cb.fileChangedSince(p) { - return true - } + if slices.ContainsFunc(cb.sourcePaths(), cb.fileChangedSince) { + return true } // --- Skills directory (handled separately from sourcePaths) --- diff --git a/pkg/agent/context_cache_test.go b/pkg/agent/context_cache_test.go index ba70d4c0d..0905e8a46 100644 --- a/pkg/agent/context_cache_test.go +++ b/pkg/agent/context_cache_test.go @@ -404,11 +404,11 @@ func TestConcurrentBuildSystemPromptWithCache(t *testing.T) { var wg sync.WaitGroup errs := make(chan string, goroutines*iterations) - for g := 0; g < goroutines; g++ { + for g := range goroutines { wg.Add(1) go func(id int) { defer wg.Done() - for i := 0; i < iterations; i++ { + for i := range iterations { result := cb.BuildSystemPromptWithCache() if result == "" { errs <- "empty prompt returned" diff --git a/pkg/agent/loop_test.go b/pkg/agent/loop_test.go index 4414398b1..aa9f823b7 100644 --- a/pkg/agent/loop_test.go +++ b/pkg/agent/loop_test.go @@ -5,6 +5,7 @@ import ( "fmt" "os" "path/filepath" + "slices" "testing" "time" @@ -175,13 +176,7 @@ func TestToolRegistry_ToolRegistration(t *testing.T) { toolsList := toolsInfo["names"].([]string) // Check that our custom tool name is in the list - found := false - for _, name := range toolsList { - if name == "mock_custom" { - found = true - break - } - } + found := slices.Contains(toolsList, "mock_custom") if !found { t.Error("Expected custom tool to be registered") } @@ -250,13 +245,7 @@ func TestToolRegistry_GetDefinitions(t *testing.T) { toolsList := toolsInfo["names"].([]string) // Check that our custom tool name is in the list - found := false - for _, name := range toolsList { - if name == "mock_custom" { - found = true - break - } - } + found := slices.Contains(toolsList, "mock_custom") if !found { t.Error("Expected custom tool to be registered") } diff --git a/pkg/agent/memory.go b/pkg/agent/memory.go index 87a687479..01e682f3b 100644 --- a/pkg/agent/memory.go +++ b/pkg/agent/memory.go @@ -111,7 +111,7 @@ func (ms *MemoryStore) GetRecentDailyNotes(days int) string { var sb strings.Builder first := true - for i := 0; i < days; i++ { + for i := range days { date := time.Now().AddDate(0, 0, -i) dateStr := date.Format("20060102") // YYYYMMDD monthDir := dateStr[:6] // YYYYMM diff --git a/pkg/channels/onebot.go b/pkg/channels/onebot.go index 4576a11ce..1b0cbc4ab 100644 --- a/pkg/channels/onebot.go +++ b/pkg/channels/onebot.go @@ -318,10 +318,7 @@ func (c *OneBotChannel) sendAPIRequest(action string, params any, timeout time.D } func (c *OneBotChannel) reconnectLoop() { - interval := time.Duration(c.config.ReconnectInterval) * time.Second - if interval < 5*time.Second { - interval = 5 * time.Second - } + interval := max(time.Duration(c.config.ReconnectInterval)*time.Second, 5*time.Second) for { select { @@ -975,8 +972,8 @@ func (c *OneBotChannel) checkGroupTrigger( if prefix == "" { continue } - if strings.HasPrefix(content, prefix) { - return true, strings.TrimSpace(strings.TrimPrefix(content, prefix)) + if after, ok := strings.CutPrefix(content, prefix); ok { + return true, strings.TrimSpace(after) } } diff --git a/pkg/channels/wecom.go b/pkg/channels/wecom.go index f8daf89de..e24157f5c 100644 --- a/pkg/channels/wecom.go +++ b/pkg/channels/wecom.go @@ -596,7 +596,7 @@ func pkcs7UnpadWeCom(data []byte) ([]byte, error) { return nil, fmt.Errorf("padding size larger than data") } // Verify all padding bytes - for i := 0; i < padding; i++ { + for i := range padding { if data[len(data)-1-i] != byte(padding) { return nil, fmt.Errorf("invalid padding byte at position %d", i) } diff --git a/pkg/channels/wecom_app_test.go b/pkg/channels/wecom_app_test.go index abf15c52b..ba911d49f 100644 --- a/pkg/channels/wecom_app_test.go +++ b/pkg/channels/wecom_app_test.go @@ -46,7 +46,7 @@ func encryptTestMessageApp(message, aesKey string) (string, error) { // Prepare message: random(16) + msg_len(4) + msg + corp_id random := make([]byte, 0, 16) - for i := 0; i < 16; i++ { + for i := range 16 { random = append(random, byte(i+1)) } diff --git a/pkg/channels/wecom_test.go b/pkg/channels/wecom_test.go index 8afa7e8c3..88aed8d2b 100644 --- a/pkg/channels/wecom_test.go +++ b/pkg/channels/wecom_test.go @@ -45,7 +45,7 @@ func encryptTestMessage(message, aesKey string) (string, error) { // Prepare message: random(16) + msg_len(4) + msg + receiveid random := make([]byte, 0, 16) - for i := 0; i < 16; i++ { + for i := range 16 { random = append(random, byte(i)) } diff --git a/pkg/config/model_config_test.go b/pkg/config/model_config_test.go index 084f50a82..da6e506f8 100644 --- a/pkg/config/model_config_test.go +++ b/pkg/config/model_config_test.go @@ -64,7 +64,7 @@ func TestGetModelConfig_RoundRobin(t *testing.T) { // Test round-robin distribution results := make(map[string]int) - for i := 0; i < 30; i++ { + for range 30 { result, err := cfg.GetModelConfig("lb-model") if err != nil { t.Fatalf("GetModelConfig() error = %v", err) @@ -94,17 +94,15 @@ func TestGetModelConfig_Concurrent(t *testing.T) { var wg sync.WaitGroup errors := make(chan error, goroutines*iterations) - for i := 0; i < goroutines; i++ { - wg.Add(1) - go func() { - defer wg.Done() - for j := 0; j < iterations; j++ { + for range goroutines { + wg.Go(func() { + for range iterations { _, err := cfg.GetModelConfig("concurrent-model") if err != nil { errors <- err } } - }() + }) } wg.Wait() diff --git a/pkg/health/server.go b/pkg/health/server.go index 77b36034d..d1acfb662 100644 --- a/pkg/health/server.go +++ b/pkg/health/server.go @@ -4,6 +4,7 @@ import ( "context" "encoding/json" "fmt" + "maps" "net/http" "sync" "time" @@ -122,9 +123,7 @@ func (s *Server) readyHandler(w http.ResponseWriter, r *http.Request) { s.mu.RLock() ready := s.ready checks := make(map[string]Check) - for k, v := range s.checks { - checks[k] = v - } + maps.Copy(checks, s.checks) s.mu.RUnlock() if !ready { diff --git a/pkg/providers/anthropic/provider.go b/pkg/providers/anthropic/provider.go index 9162174c9..1bb15f771 100644 --- a/pkg/providers/anthropic/provider.go +++ b/pkg/providers/anthropic/provider.go @@ -212,14 +212,14 @@ func translateTools(tools []ToolDefinition) []anthropic.ToolUnionParam { } func parseResponse(resp *anthropic.Message) *LLMResponse { - var content string + var content strings.Builder var toolCalls []ToolCall for _, block := range resp.Content { switch block.Type { case "text": tb := block.AsText() - content += tb.Text + content.WriteString(tb.Text) case "tool_use": tu := block.AsToolUse() var args map[string]any @@ -246,7 +246,7 @@ func parseResponse(resp *anthropic.Message) *LLMResponse { } return &LLMResponse{ - Content: content, + Content: content.String(), ToolCalls: toolCalls, FinishReason: finishReason, Usage: &UsageInfo{ @@ -264,8 +264,8 @@ func normalizeBaseURL(apiBase string) string { } base = strings.TrimRight(base, "/") - if strings.HasSuffix(base, "/v1") { - base = strings.TrimSuffix(base, "/v1") + if before, ok := strings.CutSuffix(base, "/v1"); ok { + base = before } if base == "" { return defaultBaseURL diff --git a/pkg/providers/codex_provider.go b/pkg/providers/codex_provider.go index dcc740ba4..47618300a 100644 --- a/pkg/providers/codex_provider.go +++ b/pkg/providers/codex_provider.go @@ -163,8 +163,8 @@ func resolveCodexModel(model string) (string, string) { return codexDefaultModel, "empty model" } - if strings.HasPrefix(m, "openai/") { - m = strings.TrimPrefix(m, "openai/") + if after, ok := strings.CutPrefix(m, "openai/"); ok { + m = after } else if strings.Contains(m, "/") { return codexDefaultModel, "non-openai model namespace" } diff --git a/pkg/providers/cooldown_test.go b/pkg/providers/cooldown_test.go index 47f43ad5c..b517e7feb 100644 --- a/pkg/providers/cooldown_test.go +++ b/pkg/providers/cooldown_test.go @@ -138,7 +138,7 @@ func TestCooldown_FailureWindowReset(t *testing.T) { ct, current := newTestTracker(now) // 4 errors → 1h cooldown - for i := 0; i < 4; i++ { + for range 4 { ct.MarkFailure("openai", FailoverRateLimit) *current = current.Add(2 * time.Second) // small advance between errors } @@ -230,7 +230,7 @@ func TestCooldown_ConcurrentAccess(t *testing.T) { ct := NewCooldownTracker() var wg sync.WaitGroup - for i := 0; i < 100; i++ { + for range 100 { wg.Add(3) go func() { defer wg.Done() diff --git a/pkg/providers/openai_compat/provider.go b/pkg/providers/openai_compat/provider.go index 7dace71f2..cd606d533 100644 --- a/pkg/providers/openai_compat/provider.go +++ b/pkg/providers/openai_compat/provider.go @@ -307,8 +307,8 @@ func stripSystemParts(messages []Message) []openaiMessage { } func normalizeModel(model, apiBase string) string { - idx := strings.Index(model, "/") - if idx == -1 { + before, after, ok := strings.Cut(model, "/") + if !ok { return model } @@ -316,10 +316,10 @@ func normalizeModel(model, apiBase string) string { return model } - prefix := strings.ToLower(model[:idx]) + prefix := strings.ToLower(before) switch prefix { case "moonshot", "nvidia", "groq", "ollama", "deepseek", "google", "openrouter", "zhipu", "mistral": - return model[idx+1:] + return after default: return model } diff --git a/pkg/routing/agent_id_test.go b/pkg/routing/agent_id_test.go index 050fe0645..ea3b8b6aa 100644 --- a/pkg/routing/agent_id_test.go +++ b/pkg/routing/agent_id_test.go @@ -1,5 +1,7 @@ package routing +import "strings" + import "testing" func TestNormalizeAgentID_Empty(t *testing.T) { @@ -57,11 +59,11 @@ func TestNormalizeAgentID_AllInvalid(t *testing.T) { } func TestNormalizeAgentID_TruncatesAt64(t *testing.T) { - long := "" - for i := 0; i < 100; i++ { - long += "a" + var long strings.Builder + for range 100 { + long.WriteString("a") } - got := NormalizeAgentID(long) + got := NormalizeAgentID(long.String()) if len(got) > MaxAgentIDLength { t.Errorf("length = %d, want <= %d", len(got), MaxAgentIDLength) } diff --git a/pkg/skills/loader.go b/pkg/skills/loader.go index 67d3e70e0..fcbcf934b 100644 --- a/pkg/skills/loader.go +++ b/pkg/skills/loader.go @@ -240,7 +240,7 @@ func (sl *SkillsLoader) parseSimpleYAML(content string) map[string]string { normalized := strings.ReplaceAll(content, "\r\n", "\n") normalized = strings.ReplaceAll(normalized, "\r", "\n") - for _, line := range strings.Split(normalized, "\n") { + for line := range strings.SplitSeq(normalized, "\n") { line = strings.TrimSpace(line) if line == "" || strings.HasPrefix(line, "#") { continue diff --git a/pkg/skills/search_cache.go b/pkg/skills/search_cache.go index 5d7d2797e..1686e3f98 100644 --- a/pkg/skills/search_cache.go +++ b/pkg/skills/search_cache.go @@ -1,7 +1,7 @@ package skills import ( - "sort" + "slices" "strings" "sync" "time" @@ -183,7 +183,7 @@ func buildTrigrams(s string) []uint32 { } // Sort and Deduplication - sort.Slice(trigrams, func(i, j int) bool { return trigrams[i] < trigrams[j] }) + slices.Sort(trigrams) n := 1 for i := 1; i < len(trigrams); i++ { if trigrams[i] != trigrams[i-1] { diff --git a/pkg/skills/search_cache_test.go b/pkg/skills/search_cache_test.go index 816bdfb93..6bbb0e6eb 100644 --- a/pkg/skills/search_cache_test.go +++ b/pkg/skills/search_cache_test.go @@ -153,7 +153,7 @@ func TestSearchCacheConcurrency(t *testing.T) { // Concurrent writes go func() { - for i := 0; i < 100; i++ { + for i := range 100 { cache.Put("query-write-"+string(rune('a'+i%26)), []SearchResult{{Slug: "x"}}) } done <- struct{}{} @@ -161,7 +161,7 @@ func TestSearchCacheConcurrency(t *testing.T) { // Concurrent reads go func() { - for i := 0; i < 100; i++ { + for range 100 { cache.Get("query-write-a") } done <- struct{}{} diff --git a/pkg/state/state_test.go b/pkg/state/state_test.go index f717a5bb4..70117ad61 100644 --- a/pkg/state/state_test.go +++ b/pkg/state/state_test.go @@ -135,7 +135,7 @@ func TestConcurrentAccess(t *testing.T) { // Test concurrent writes done := make(chan bool, 10) - for i := 0; i < 10; i++ { + for i := range 10 { go func(idx int) { channel := fmt.Sprintf("channel-%d", idx) sm.SetLastChannel(channel) @@ -144,7 +144,7 @@ func TestConcurrentAccess(t *testing.T) { } // Wait for all goroutines to complete - for i := 0; i < 10; i++ { + for range 10 { <-done } diff --git a/pkg/tools/cron.go b/pkg/tools/cron.go index 562fffc84..3140e5e25 100644 --- a/pkg/tools/cron.go +++ b/pkg/tools/cron.go @@ -3,6 +3,7 @@ package tools import ( "context" "fmt" + "strings" "sync" "time" @@ -218,7 +219,8 @@ func (t *CronTool) listJobs() *ToolResult { return SilentResult("No scheduled jobs") } - result := "Scheduled jobs:\n" + var result strings.Builder + result.WriteString("Scheduled jobs:\n") for _, j := range jobs { var scheduleInfo string if j.Schedule.Kind == "every" && j.Schedule.EveryMS != nil { @@ -230,10 +232,10 @@ func (t *CronTool) listJobs() *ToolResult { } else { scheduleInfo = "unknown" } - result += fmt.Sprintf("- %s (id: %s, %s)\n", j.Name, j.ID, scheduleInfo) + result.WriteString(fmt.Sprintf("- %s (id: %s, %s)\n", j.Name, j.ID, scheduleInfo)) } - return SilentResult(result) + return SilentResult(result.String()) } func (t *CronTool) removeJob(args map[string]any) *ToolResult { diff --git a/pkg/tools/registry_test.go b/pkg/tools/registry_test.go index 8ae13b20c..8fe88ca78 100644 --- a/pkg/tools/registry_test.go +++ b/pkg/tools/registry_test.go @@ -329,7 +329,7 @@ func TestToolRegistry_ConcurrentAccess(t *testing.T) { r := NewToolRegistry() var wg sync.WaitGroup - for i := 0; i < 50; i++ { + for i := range 50 { wg.Add(1) go func(n int) { defer wg.Done() diff --git a/pkg/tools/web.go b/pkg/tools/web.go index 8ba2a723a..de0c991ff 100644 --- a/pkg/tools/web.go +++ b/pkg/tools/web.go @@ -285,7 +285,7 @@ func (p *DuckDuckGoSearchProvider) extractResults(html string, count int, query maxItems := min(len(matches), count) - for i := 0; i < maxItems; i++ { + for i := range maxItems { urlStr := matches[i][1] title := stripTags(matches[i][2]) title = strings.TrimSpace(title) @@ -293,9 +293,9 @@ func (p *DuckDuckGoSearchProvider) extractResults(html string, count int, query // URL decoding if needed if strings.Contains(urlStr, "uddg=") { if u, err := url.QueryUnescape(urlStr); err == nil { - idx := strings.Index(u, "uddg=") - if idx != -1 { - urlStr = u[idx+5:] + _, after, ok := strings.Cut(u, "uddg=") + if ok { + urlStr = after } } } diff --git a/pkg/utils/message.go b/pkg/utils/message.go index 1d05950d9..a65506edc 100644 --- a/pkg/utils/message.go +++ b/pkg/utils/message.go @@ -13,10 +13,7 @@ func SplitMessage(content string, maxLen int) []string { var messages []string // Dynamic buffer: 10% of maxLen, but at least 50 chars if possible - codeBlockBuffer := maxLen / 10 - if codeBlockBuffer < 50 { - codeBlockBuffer = 50 - } + codeBlockBuffer := max(maxLen/10, 50) if codeBlockBuffer > maxLen/2 { codeBlockBuffer = maxLen / 2 } @@ -28,10 +25,7 @@ func SplitMessage(content string, maxLen int) []string { } // Effective split point: maxLen minus buffer, to leave room for code blocks - effectiveLimit := maxLen - codeBlockBuffer - if effectiveLimit < maxLen/2 { - effectiveLimit = maxLen / 2 - } + effectiveLimit := max(maxLen-codeBlockBuffer, maxLen/2) // Find natural split point within the effective limit msgEnd := findLastNewline(content[:effectiveLimit], 200) @@ -151,10 +145,7 @@ func findNextClosingCodeBlock(text string, startIdx int) int { // findLastNewline finds the last newline character within the last N characters // Returns the position of the newline or -1 if not found func findLastNewline(s string, searchWindow int) int { - searchStart := len(s) - searchWindow - if searchStart < 0 { - searchStart = 0 - } + searchStart := max(len(s)-searchWindow, 0) for i := len(s) - 1; i >= searchStart; i-- { if s[i] == '\n' { return i @@ -166,10 +157,7 @@ func findLastNewline(s string, searchWindow int) int { // findLastSpace finds the last space character within the last N characters // Returns the position of the space or -1 if not found func findLastSpace(s string, searchWindow int) int { - searchStart := len(s) - searchWindow - if searchStart < 0 { - searchStart = 0 - } + searchStart := max(len(s)-searchWindow, 0) for i := len(s) - 1; i >= searchStart; i-- { if s[i] == ' ' || s[i] == '\t' { return i From f2a71ca8247c66d789fda19a29a52f3da91ba292 Mon Sep 17 00:00:00 2001 From: Petrichor Date: Fri, 27 Feb 2026 21:20:56 +0800 Subject: [PATCH 2/4] fix(lint): format imports in agent_id_test --- pkg/routing/agent_id_test.go | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/pkg/routing/agent_id_test.go b/pkg/routing/agent_id_test.go index ea3b8b6aa..f9a65c969 100644 --- a/pkg/routing/agent_id_test.go +++ b/pkg/routing/agent_id_test.go @@ -1,8 +1,9 @@ package routing -import "strings" - -import "testing" +import ( + "strings" + "testing" +) func TestNormalizeAgentID_Empty(t *testing.T) { if got := NormalizeAgentID(""); got != DefaultAgentID { From b3c3b02666bd188a7f0f62a7a086b23d1b52d634 Mon Sep 17 00:00:00 2001 From: Owen Wu <1449069+yubing744@users.noreply.github.com> Date: Sat, 28 Feb 2026 21:25:31 -0800 Subject: [PATCH 3/4] fix(onboard): use AGENTS.md template instead of AGENT.md (#931) * fix(onboard): use AGENTS.md workspace template * test(onboard): move AGENTS template regression test to new package --- cmd/picoclaw/internal/onboard/helpers_test.go | 25 +++++++++++++++++++ workspace/{AGENT.md => AGENTS.md} | 0 2 files changed, 25 insertions(+) create mode 100644 cmd/picoclaw/internal/onboard/helpers_test.go rename workspace/{AGENT.md => AGENTS.md} (100%) diff --git a/cmd/picoclaw/internal/onboard/helpers_test.go b/cmd/picoclaw/internal/onboard/helpers_test.go new file mode 100644 index 000000000..f3e0c92e0 --- /dev/null +++ b/cmd/picoclaw/internal/onboard/helpers_test.go @@ -0,0 +1,25 @@ +package onboard + +import ( + "os" + "path/filepath" + "testing" +) + +func TestCopyEmbeddedToTargetUsesAgentsMarkdown(t *testing.T) { + targetDir := t.TempDir() + + if err := copyEmbeddedToTarget(targetDir); err != nil { + t.Fatalf("copyEmbeddedToTarget() error = %v", err) + } + + agentsPath := filepath.Join(targetDir, "AGENTS.md") + if _, err := os.Stat(agentsPath); err != nil { + t.Fatalf("expected %s to exist: %v", agentsPath, err) + } + + legacyPath := filepath.Join(targetDir, "AGENT.md") + if _, err := os.Stat(legacyPath); !os.IsNotExist(err) { + t.Fatalf("expected legacy file %s to be absent, got err=%v", legacyPath, err) + } +} diff --git a/workspace/AGENT.md b/workspace/AGENTS.md similarity index 100% rename from workspace/AGENT.md rename to workspace/AGENTS.md From 44a52c0cf646beee1522b85c67777df5fdba0386 Mon Sep 17 00:00:00 2001 From: Tong Niu Date: Sun, 1 Mar 2026 16:55:46 +1100 Subject: [PATCH 4/4] fix(tools): close resp.Body on retry cancel and cache http.Client instances (#940) * fix(tools): close resp.Body on retry cancel and cache http.Client instances Fix resp.Body leak in DoRequestWithRetry where req.Body (request) was incorrectly closed instead of resp.Body (response) on context cancel. Cache http.Client on web search/fetch provider structs and channel adapters (WeCom, LINE) to avoid per-call allocation overhead. * fix(channels): preserve original http client timeouts for LINE and WeCom Split LINE single 60s client into infoClient (10s) for bot info lookups and apiClient (30s) for messaging API calls. Lower WeCom cached client base timeout from 60s to 30s (matching uploadMedia), and ensure it is always >= the configured ReplyTimeout so the per-request context deadline remains the effective limit. * refactor(tools): extract timeout consts and deduplicate WebFetchTool constructors Address PR review feedback from xiaket: - Define searchTimeout, perplexityTimeout, fetchTimeout, defaultMaxChars, and maxRedirects as package-level consts instead of magic numbers. - Remove misleading "No proxy" comment in NewWebFetchTool. - Deduplicate NewWebFetchTool by delegating to NewWebFetchToolWithProxy. * test(utils): add context cancellation test for DoRequestWithRetry Verify that resp.Body is properly closed when the context is canceled during retry sleep, covering the C8 resp.Body leak fix. * fix(utils): close resp in test to satisfy bodyclose linter * fix(utils): eliminate flakiness in context cancellation retry test Synchronize cancellation using an onRoundTrip callback from the transport wrapper instead of a timing-based context timeout. This ensures the first client.Do completes before cancel fires, so cancellation always hits during sleepWithCtx. --- pkg/agent/loop.go | 14 ++++- pkg/channels/line/line.go | 20 ++++--- pkg/channels/wecom/app.go | 18 ++++-- pkg/channels/wecom/bot.go | 12 +++- pkg/tools/web.go | 109 +++++++++++++++++++---------------- pkg/tools/web_test.go | 45 ++++++++++++--- pkg/utils/http_retry.go | 3 + pkg/utils/http_retry_test.go | 88 ++++++++++++++++++++++++++++ 8 files changed, 230 insertions(+), 79 deletions(-) diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index 8fd7328d1..a72f95bb1 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -99,7 +99,7 @@ func registerSharedTools( } // Web tools - if searchTool := tools.NewWebSearchTool(tools.WebSearchToolOptions{ + searchTool, err := tools.NewWebSearchTool(tools.WebSearchToolOptions{ BraveAPIKey: cfg.Tools.Web.Brave.APIKey, BraveMaxResults: cfg.Tools.Web.Brave.MaxResults, BraveEnabled: cfg.Tools.Web.Brave.Enabled, @@ -113,10 +113,18 @@ func registerSharedTools( PerplexityMaxResults: cfg.Tools.Web.Perplexity.MaxResults, PerplexityEnabled: cfg.Tools.Web.Perplexity.Enabled, Proxy: cfg.Tools.Web.Proxy, - }); searchTool != nil { + }) + if err != nil { + logger.ErrorCF("agent", "Failed to create web search tool", map[string]any{"error": err.Error()}) + } else if searchTool != nil { agent.Tools.Register(searchTool) } - agent.Tools.Register(tools.NewWebFetchToolWithProxy(50000, cfg.Tools.Web.Proxy)) + fetchTool, err := tools.NewWebFetchToolWithProxy(50000, cfg.Tools.Web.Proxy) + if err != nil { + logger.ErrorCF("agent", "Failed to create web fetch tool", map[string]any{"error": err.Error()}) + } else { + agent.Tools.Register(fetchTool) + } // Hardware tools (I2C, SPI) - Linux only, returns error on other platforms agent.Tools.Register(tools.NewI2CTool()) diff --git a/pkg/channels/line/line.go b/pkg/channels/line/line.go index 9fac2831c..398f12e6b 100644 --- a/pkg/channels/line/line.go +++ b/pkg/channels/line/line.go @@ -45,11 +45,13 @@ type replyTokenEntry struct { type LINEChannel struct { *channels.BaseChannel config config.LINEConfig - botUserID string // Bot's user ID - botBasicID string // Bot's basic ID (e.g. @216ru...) - botDisplayName string // Bot's display name for text-based mention detection - replyTokens sync.Map // chatID -> replyTokenEntry - quoteTokens sync.Map // chatID -> quoteToken (string) + infoClient *http.Client // for bot info lookups (short timeout) + apiClient *http.Client // for messaging API calls + botUserID string // Bot's user ID + botBasicID string // Bot's basic ID (e.g. @216ru...) + botDisplayName string // Bot's display name for text-based mention detection + replyTokens sync.Map // chatID -> replyTokenEntry + quoteTokens sync.Map // chatID -> quoteToken (string) ctx context.Context cancel context.CancelFunc } @@ -69,6 +71,8 @@ func NewLINEChannel(cfg config.LINEConfig, messageBus *bus.MessageBus) (*LINECha return &LINEChannel{ BaseChannel: base, config: cfg, + infoClient: &http.Client{Timeout: 10 * time.Second}, + apiClient: &http.Client{Timeout: 30 * time.Second}, }, nil } @@ -104,8 +108,7 @@ func (c *LINEChannel) fetchBotInfo() error { } req.Header.Set("Authorization", "Bearer "+c.config.ChannelAccessToken) - client := &http.Client{Timeout: 10 * time.Second} - resp, err := client.Do(req) + resp, err := c.infoClient.Do(req) if err != nil { return err } @@ -644,8 +647,7 @@ func (c *LINEChannel) callAPI(ctx context.Context, endpoint string, payload any) req.Header.Set("Content-Type", "application/json") req.Header.Set("Authorization", "Bearer "+c.config.ChannelAccessToken) - client := &http.Client{Timeout: 30 * time.Second} - resp, err := client.Do(req) + resp, err := c.apiClient.Do(req) if err != nil { return channels.ClassifyNetError(err) } diff --git a/pkg/channels/wecom/app.go b/pkg/channels/wecom/app.go index 7a23f9617..292a71fd2 100644 --- a/pkg/channels/wecom/app.go +++ b/pkg/channels/wecom/app.go @@ -32,6 +32,7 @@ const ( type WeComAppChannel struct { *channels.BaseChannel config config.WeComAppConfig + client *http.Client accessToken string tokenExpiry time.Time tokenMu sync.RWMutex @@ -129,10 +130,18 @@ func NewWeComAppChannel(cfg config.WeComAppConfig, messageBus *bus.MessageBus) ( channels.WithReasoningChannelID(cfg.ReasoningChannelID), ) + // Client timeout must be >= the configured ReplyTimeout so the + // per-request context deadline is always the effective limit. + clientTimeout := 30 * time.Second + if d := time.Duration(cfg.ReplyTimeout) * time.Second; d > clientTimeout { + clientTimeout = d + } + ctx, cancel := context.WithCancel(context.Background()) return &WeComAppChannel{ BaseChannel: base, config: cfg, + client: &http.Client{Timeout: clientTimeout}, ctx: ctx, cancel: cancel, processedMsgs: make(map[string]bool), @@ -306,8 +315,7 @@ func (c *WeComAppChannel) uploadMedia(ctx context.Context, accessToken, mediaTyp } req.Header.Set("Content-Type", writer.FormDataContentType()) - client := &http.Client{Timeout: 30 * time.Second} - resp, err := client.Do(req) + resp, err := c.client.Do(req) if err != nil { return "", channels.ClassifyNetError(err) } @@ -364,8 +372,7 @@ func (c *WeComAppChannel) sendImageMessage(ctx context.Context, accessToken, use } req.Header.Set("Content-Type", "application/json") - client := &http.Client{Timeout: time.Duration(timeout) * time.Second} - resp, err := client.Do(req) + resp, err := c.client.Do(req) if err != nil { return channels.ClassifyNetError(err) } @@ -746,8 +753,7 @@ func (c *WeComAppChannel) sendTextMessage(ctx context.Context, accessToken, user } req.Header.Set("Content-Type", "application/json") - client := &http.Client{Timeout: time.Duration(timeout) * time.Second} - resp, err := client.Do(req) + resp, err := c.client.Do(req) if err != nil { return channels.ClassifyNetError(err) } diff --git a/pkg/channels/wecom/bot.go b/pkg/channels/wecom/bot.go index 39f84d55c..0d0426c0d 100644 --- a/pkg/channels/wecom/bot.go +++ b/pkg/channels/wecom/bot.go @@ -25,6 +25,7 @@ import ( type WeComBotChannel struct { *channels.BaseChannel config config.WeComConfig + client *http.Client ctx context.Context cancel context.CancelFunc processedMsgs map[string]bool // Message deduplication: msg_id -> processed @@ -93,10 +94,18 @@ func NewWeComBotChannel(cfg config.WeComConfig, messageBus *bus.MessageBus) (*We channels.WithReasoningChannelID(cfg.ReasoningChannelID), ) + // Client timeout must be >= the configured ReplyTimeout so the + // per-request context deadline is always the effective limit. + clientTimeout := 30 * time.Second + if d := time.Duration(cfg.ReplyTimeout) * time.Second; d > clientTimeout { + clientTimeout = d + } + ctx, cancel := context.WithCancel(context.Background()) return &WeComBotChannel{ BaseChannel: base, config: cfg, + client: &http.Client{Timeout: clientTimeout}, ctx: ctx, cancel: cancel, processedMsgs: make(map[string]bool), @@ -450,8 +459,7 @@ func (c *WeComBotChannel) sendWebhookReply(ctx context.Context, userID, content } req.Header.Set("Content-Type", "application/json") - client := &http.Client{Timeout: time.Duration(timeout) * time.Second} - resp, err := client.Do(req) + resp, err := c.client.Do(req) if err != nil { return channels.ClassifyNetError(err) } diff --git a/pkg/tools/web.go b/pkg/tools/web.go index 8ba2a723a..834e7bfc7 100644 --- a/pkg/tools/web.go +++ b/pkg/tools/web.go @@ -15,6 +15,14 @@ import ( const ( userAgent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36" + + // HTTP client timeouts for web tool providers. + searchTimeout = 10 * time.Second // Brave, Tavily, DuckDuckGo + perplexityTimeout = 30 * time.Second // Perplexity (LLM-based, slower) + fetchTimeout = 60 * time.Second // WebFetchTool + + defaultMaxChars = 50000 + maxRedirects = 5 ) // Pre-compiled regexes for HTML text extraction @@ -74,6 +82,7 @@ type SearchProvider interface { type BraveSearchProvider struct { apiKey string proxy string + client *http.Client } func (p *BraveSearchProvider) Search(ctx context.Context, query string, count int) (string, error) { @@ -88,11 +97,7 @@ func (p *BraveSearchProvider) Search(ctx context.Context, query string, count in req.Header.Set("Accept", "application/json") req.Header.Set("X-Subscription-Token", p.apiKey) - client, err := createHTTPClient(p.proxy, 10*time.Second) - if err != nil { - return "", fmt.Errorf("failed to create HTTP client: %w", err) - } - resp, err := client.Do(req) + resp, err := p.client.Do(req) if err != nil { return "", fmt.Errorf("request failed: %w", err) } @@ -143,6 +148,7 @@ type TavilySearchProvider struct { apiKey string baseURL string proxy string + client *http.Client } func (p *TavilySearchProvider) Search(ctx context.Context, query string, count int) (string, error) { @@ -174,11 +180,7 @@ func (p *TavilySearchProvider) Search(ctx context.Context, query string, count i req.Header.Set("Content-Type", "application/json") req.Header.Set("User-Agent", userAgent) - client, err := createHTTPClient(p.proxy, 10*time.Second) - if err != nil { - return "", fmt.Errorf("failed to create HTTP client: %w", err) - } - resp, err := client.Do(req) + resp, err := p.client.Do(req) if err != nil { return "", fmt.Errorf("request failed: %w", err) } @@ -226,7 +228,8 @@ func (p *TavilySearchProvider) Search(ctx context.Context, query string, count i } type DuckDuckGoSearchProvider struct { - proxy string + proxy string + client *http.Client } func (p *DuckDuckGoSearchProvider) Search(ctx context.Context, query string, count int) (string, error) { @@ -239,11 +242,7 @@ func (p *DuckDuckGoSearchProvider) Search(ctx context.Context, query string, cou req.Header.Set("User-Agent", userAgent) - client, err := createHTTPClient(p.proxy, 10*time.Second) - if err != nil { - return "", fmt.Errorf("failed to create HTTP client: %w", err) - } - resp, err := client.Do(req) + resp, err := p.client.Do(req) if err != nil { return "", fmt.Errorf("request failed: %w", err) } @@ -322,6 +321,7 @@ func stripTags(content string) string { type PerplexitySearchProvider struct { apiKey string proxy string + client *http.Client } func (p *PerplexitySearchProvider) Search(ctx context.Context, query string, count int) (string, error) { @@ -356,11 +356,7 @@ func (p *PerplexitySearchProvider) Search(ctx context.Context, query string, cou req.Header.Set("Authorization", "Bearer "+p.apiKey) req.Header.Set("User-Agent", userAgent) - client, err := createHTTPClient(p.proxy, 30*time.Second) - if err != nil { - return "", fmt.Errorf("failed to create HTTP client: %w", err) - } - resp, err := client.Do(req) + resp, err := p.client.Do(req) if err != nil { return "", fmt.Errorf("request failed: %w", err) } @@ -415,43 +411,60 @@ type WebSearchToolOptions struct { Proxy string } -func NewWebSearchTool(opts WebSearchToolOptions) *WebSearchTool { +func NewWebSearchTool(opts WebSearchToolOptions) (*WebSearchTool, error) { var provider SearchProvider maxResults := 5 // Priority: Perplexity > Brave > Tavily > DuckDuckGo if opts.PerplexityEnabled && opts.PerplexityAPIKey != "" { - provider = &PerplexitySearchProvider{apiKey: opts.PerplexityAPIKey, proxy: opts.Proxy} + client, err := createHTTPClient(opts.Proxy, perplexityTimeout) + if err != nil { + return nil, fmt.Errorf("failed to create HTTP client for Perplexity: %w", err) + } + provider = &PerplexitySearchProvider{apiKey: opts.PerplexityAPIKey, proxy: opts.Proxy, client: client} if opts.PerplexityMaxResults > 0 { maxResults = opts.PerplexityMaxResults } } else if opts.BraveEnabled && opts.BraveAPIKey != "" { - provider = &BraveSearchProvider{apiKey: opts.BraveAPIKey, proxy: opts.Proxy} + client, err := createHTTPClient(opts.Proxy, searchTimeout) + if err != nil { + return nil, fmt.Errorf("failed to create HTTP client for Brave: %w", err) + } + provider = &BraveSearchProvider{apiKey: opts.BraveAPIKey, proxy: opts.Proxy, client: client} if opts.BraveMaxResults > 0 { maxResults = opts.BraveMaxResults } } else if opts.TavilyEnabled && opts.TavilyAPIKey != "" { + client, err := createHTTPClient(opts.Proxy, searchTimeout) + if err != nil { + return nil, fmt.Errorf("failed to create HTTP client for Tavily: %w", err) + } provider = &TavilySearchProvider{ apiKey: opts.TavilyAPIKey, baseURL: opts.TavilyBaseURL, proxy: opts.Proxy, + client: client, } if opts.TavilyMaxResults > 0 { maxResults = opts.TavilyMaxResults } } else if opts.DuckDuckGoEnabled { - provider = &DuckDuckGoSearchProvider{proxy: opts.Proxy} + client, err := createHTTPClient(opts.Proxy, searchTimeout) + if err != nil { + return nil, fmt.Errorf("failed to create HTTP client for DuckDuckGo: %w", err) + } + provider = &DuckDuckGoSearchProvider{proxy: opts.Proxy, client: client} if opts.DuckDuckGoMaxResults > 0 { maxResults = opts.DuckDuckGoMaxResults } } else { - return nil + return nil, nil } return &WebSearchTool{ provider: provider, maxResults: maxResults, - } + }, nil } func (t *WebSearchTool) Name() string { @@ -508,25 +521,34 @@ func (t *WebSearchTool) Execute(ctx context.Context, args map[string]any) *ToolR type WebFetchTool struct { maxChars int proxy string + client *http.Client } func NewWebFetchTool(maxChars int) *WebFetchTool { - if maxChars <= 0 { - maxChars = 50000 - } - return &WebFetchTool{ - maxChars: maxChars, - } + // createHTTPClient cannot fail with an empty proxy string. + tool, _ := NewWebFetchToolWithProxy(maxChars, "") + return tool } -func NewWebFetchToolWithProxy(maxChars int, proxy string) *WebFetchTool { +func NewWebFetchToolWithProxy(maxChars int, proxy string) (*WebFetchTool, error) { if maxChars <= 0 { - maxChars = 50000 + maxChars = defaultMaxChars + } + client, err := createHTTPClient(proxy, fetchTimeout) + if err != nil { + return nil, fmt.Errorf("failed to create HTTP client for web fetch: %w", err) + } + client.CheckRedirect = func(req *http.Request, via []*http.Request) error { + if len(via) >= maxRedirects { + return fmt.Errorf("stopped after %d redirects", maxRedirects) + } + return nil } return &WebFetchTool{ maxChars: maxChars, proxy: proxy, - } + client: client, + }, nil } func (t *WebFetchTool) Name() string { @@ -588,20 +610,7 @@ func (t *WebFetchTool) Execute(ctx context.Context, args map[string]any) *ToolRe req.Header.Set("User-Agent", userAgent) - client, err := createHTTPClient(t.proxy, 60*time.Second) - if err != nil { - return ErrorResult(fmt.Sprintf("failed to create HTTP client: %v", err)) - } - - // Configure redirect handling - client.CheckRedirect = func(req *http.Request, via []*http.Request) error { - if len(via) >= 5 { - return fmt.Errorf("stopped after 5 redirects") - } - return nil - } - - resp, err := client.Do(req) + resp, err := t.client.Do(req) if err != nil { return ErrorResult(fmt.Sprintf("request failed: %v", err)) } diff --git a/pkg/tools/web_test.go b/pkg/tools/web_test.go index 2cd79eb24..db3c08ba6 100644 --- a/pkg/tools/web_test.go +++ b/pkg/tools/web_test.go @@ -176,13 +176,19 @@ func TestWebTool_WebFetch_Truncation(t *testing.T) { // TestWebTool_WebSearch_NoApiKey verifies that no tool is created when API key is missing func TestWebTool_WebSearch_NoApiKey(t *testing.T) { - tool := NewWebSearchTool(WebSearchToolOptions{BraveEnabled: true, BraveAPIKey: ""}) + tool, err := NewWebSearchTool(WebSearchToolOptions{BraveEnabled: true, BraveAPIKey: ""}) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } if tool != nil { t.Errorf("Expected nil tool when Brave API key is empty") } // Also nil when nothing is enabled - tool = NewWebSearchTool(WebSearchToolOptions{}) + tool, err = NewWebSearchTool(WebSearchToolOptions{}) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } if tool != nil { t.Errorf("Expected nil tool when no provider is enabled") } @@ -190,7 +196,10 @@ func TestWebTool_WebSearch_NoApiKey(t *testing.T) { // TestWebTool_WebSearch_MissingQuery verifies error handling for missing query func TestWebTool_WebSearch_MissingQuery(t *testing.T) { - tool := NewWebSearchTool(WebSearchToolOptions{BraveEnabled: true, BraveAPIKey: "test-key", BraveMaxResults: 5}) + tool, err := NewWebSearchTool(WebSearchToolOptions{BraveEnabled: true, BraveAPIKey: "test-key", BraveMaxResults: 5}) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } ctx := context.Background() args := map[string]any{} @@ -438,7 +447,10 @@ func TestCreateHTTPClient_ProxyFromEnvironmentWhenConfigEmpty(t *testing.T) { } func TestNewWebFetchToolWithProxy(t *testing.T) { - tool := NewWebFetchToolWithProxy(1024, "http://127.0.0.1:7890") + tool, err := NewWebFetchToolWithProxy(1024, "http://127.0.0.1:7890") + if err != nil { + t.Fatalf("NewWebFetchToolWithProxy() error: %v", err) + } if tool.maxChars != 1024 { t.Fatalf("maxChars = %d, want %d", tool.maxChars, 1024) } @@ -446,7 +458,10 @@ func TestNewWebFetchToolWithProxy(t *testing.T) { t.Fatalf("proxy = %q, want %q", tool.proxy, "http://127.0.0.1:7890") } - tool = NewWebFetchToolWithProxy(0, "http://127.0.0.1:7890") + tool, err = NewWebFetchToolWithProxy(0, "http://127.0.0.1:7890") + if err != nil { + t.Fatalf("NewWebFetchToolWithProxy() error: %v", err) + } if tool.maxChars != 50000 { t.Fatalf("default maxChars = %d, want %d", tool.maxChars, 50000) } @@ -454,12 +469,15 @@ func TestNewWebFetchToolWithProxy(t *testing.T) { func TestNewWebSearchTool_PropagatesProxy(t *testing.T) { t.Run("perplexity", func(t *testing.T) { - tool := NewWebSearchTool(WebSearchToolOptions{ + tool, err := NewWebSearchTool(WebSearchToolOptions{ PerplexityEnabled: true, PerplexityAPIKey: "k", PerplexityMaxResults: 3, Proxy: "http://127.0.0.1:7890", }) + if err != nil { + t.Fatalf("NewWebSearchTool() error: %v", err) + } p, ok := tool.provider.(*PerplexitySearchProvider) if !ok { t.Fatalf("provider type = %T, want *PerplexitySearchProvider", tool.provider) @@ -470,12 +488,15 @@ func TestNewWebSearchTool_PropagatesProxy(t *testing.T) { }) t.Run("brave", func(t *testing.T) { - tool := NewWebSearchTool(WebSearchToolOptions{ + tool, err := NewWebSearchTool(WebSearchToolOptions{ BraveEnabled: true, BraveAPIKey: "k", BraveMaxResults: 3, Proxy: "http://127.0.0.1:7890", }) + if err != nil { + t.Fatalf("NewWebSearchTool() error: %v", err) + } p, ok := tool.provider.(*BraveSearchProvider) if !ok { t.Fatalf("provider type = %T, want *BraveSearchProvider", tool.provider) @@ -486,11 +507,14 @@ func TestNewWebSearchTool_PropagatesProxy(t *testing.T) { }) t.Run("duckduckgo", func(t *testing.T) { - tool := NewWebSearchTool(WebSearchToolOptions{ + tool, err := NewWebSearchTool(WebSearchToolOptions{ DuckDuckGoEnabled: true, DuckDuckGoMaxResults: 3, Proxy: "http://127.0.0.1:7890", }) + if err != nil { + t.Fatalf("NewWebSearchTool() error: %v", err) + } p, ok := tool.provider.(*DuckDuckGoSearchProvider) if !ok { t.Fatalf("provider type = %T, want *DuckDuckGoSearchProvider", tool.provider) @@ -542,12 +566,15 @@ func TestWebTool_TavilySearch_Success(t *testing.T) { })) defer server.Close() - tool := NewWebSearchTool(WebSearchToolOptions{ + tool, err := NewWebSearchTool(WebSearchToolOptions{ TavilyEnabled: true, TavilyAPIKey: "test-key", TavilyBaseURL: server.URL, TavilyMaxResults: 5, }) + if err != nil { + t.Fatalf("NewWebSearchTool() error: %v", err) + } ctx := context.Background() args := map[string]any{ diff --git a/pkg/utils/http_retry.go b/pkg/utils/http_retry.go index e90fa2129..135ea0ef5 100644 --- a/pkg/utils/http_retry.go +++ b/pkg/utils/http_retry.go @@ -37,6 +37,9 @@ func DoRequestWithRetry(client *http.Client, req *http.Request) (*http.Response, if i < maxRetries-1 { if err = sleepWithCtx(req.Context(), retryDelayUnit*time.Duration(i+1)); err != nil { + if resp != nil { + resp.Body.Close() + } return nil, fmt.Errorf("failed to sleep: %w", err) } } diff --git a/pkg/utils/http_retry_test.go b/pkg/utils/http_retry_test.go index 1c2dbe115..d64cd5eda 100644 --- a/pkg/utils/http_retry_test.go +++ b/pkg/utils/http_retry_test.go @@ -1,8 +1,11 @@ package utils import ( + "context" + "io" "net/http" "net/http/httptest" + "strings" "testing" "time" @@ -77,6 +80,91 @@ func TestDoRequestWithRetry(t *testing.T) { } } +func TestDoRequestWithRetry_ContextCancel(t *testing.T) { + // Use a long retry delay so cancellation always hits during sleepWithCtx. + retryDelayUnit = 10 * time.Second + t.Cleanup(func() { retryDelayUnit = time.Second }) + + bodyClosed := false + firstRoundTripDone := make(chan struct{}, 1) + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + w.Write([]byte("error")) + })) + defer server.Close() + + client := server.Client() + client.Timeout = 30 * time.Second + client.Transport = &bodyCloseTracker{ + rt: client.Transport, + onClose: func() { bodyClosed = true }, + // Signal after the first round-trip response is fully constructed on the client side. + onRoundTrip: func() { + select { + case firstRoundTripDone <- struct{}{}: + default: + } + }, + trackURL: server.URL, + } + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + // Cancel the context after the first round-trip completes on the client side. + // This ensures client.Do has returned a valid resp (with body) and the retry + // loop is about to enter sleepWithCtx, where the cancel will be detected. + go func() { + <-firstRoundTripDone + cancel() + }() + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, server.URL, nil) + require.NoError(t, err) + + resp, err := DoRequestWithRetry(client, req) + if resp != nil { + resp.Body.Close() + } + require.Error(t, err, "expected error from context cancellation") + assert.Nil(t, resp, "expected nil response when context is canceled") + assert.True(t, bodyClosed, "expected resp.Body to be closed on context cancellation") +} + +// bodyCloseTracker wraps an http.RoundTripper and records when response bodies are closed. +type bodyCloseTracker struct { + rt http.RoundTripper + onClose func() + onRoundTrip func() // called after each successful round-trip + trackURL string +} + +func (t *bodyCloseTracker) RoundTrip(req *http.Request) (*http.Response, error) { + resp, err := t.rt.RoundTrip(req) + if err != nil { + return resp, err + } + if strings.HasPrefix(req.URL.String(), t.trackURL) { + resp.Body = &closeNotifier{ReadCloser: resp.Body, onClose: t.onClose} + if t.onRoundTrip != nil { + t.onRoundTrip() + } + } + return resp, nil +} + +// closeNotifier wraps an io.ReadCloser to detect Close calls. +type closeNotifier struct { + io.ReadCloser + onClose func() +} + +func (c *closeNotifier) Close() error { + c.onClose() + return c.ReadCloser.Close() +} + func TestDoRequestWithRetry_Delay(t *testing.T) { retryDelayUnit = time.Millisecond t.Cleanup(func() { retryDelayUnit = time.Second })