diff --git a/pkg/tools/web.go b/pkg/tools/web.go deleted file mode 100644 index 003cd860c..000000000 --- a/pkg/tools/web.go +++ /dev/null @@ -1,1139 +0,0 @@ -package tools - -import ( - "bytes" - "context" - "encoding/json" - "errors" - "fmt" - "io" - "net" - "net/http" - "net/url" - "regexp" - "strings" - "sync/atomic" - "time" -) - -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 -var ( - reScript = regexp.MustCompile(``) - reStyle = regexp.MustCompile(``) - reTags = regexp.MustCompile(`<[^>]+>`) - reWhitespace = regexp.MustCompile(`[^\S\n]+`) - reBlankLines = regexp.MustCompile(`\n{3,}`) - - // DuckDuckGo result extraction - reDDGLink = regexp.MustCompile(`]*class="[^"]*result__a[^"]*"[^>]*href="([^"]+)"[^>]*>([\s\S]*?)`) - reDDGSnippet = regexp.MustCompile(`([\s\S]*?)`) -) - -// createHTTPClient creates an HTTP client with optional proxy support -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, - }, - } - - if proxyURL != "" { - proxy, err := url.Parse(proxyURL) - if err != nil { - return nil, fmt.Errorf("invalid proxy URL: %w", err) - } - scheme := strings.ToLower(proxy.Scheme) - switch scheme { - case "http", "https", "socks5", "socks5h": - default: - return nil, fmt.Errorf( - "unsupported proxy scheme %q (supported: http, https, socks5, socks5h)", - proxy.Scheme, - ) - } - if proxy.Host == "" { - return nil, fmt.Errorf("invalid proxy URL: missing host") - } - client.Transport.(*http.Transport).Proxy = http.ProxyURL(proxy) - } else { - client.Transport.(*http.Transport).Proxy = http.ProxyFromEnvironment - } - - return client, nil -} - -type APIKeyPool struct { - keys []string - current uint32 -} - -func NewAPIKeyPool(keys []string) *APIKeyPool { - return &APIKeyPool{ - keys: keys, - } -} - -type APIKeyIterator struct { - pool *APIKeyPool - startIdx uint32 - attempt uint32 -} - -func (p *APIKeyPool) NewIterator() *APIKeyIterator { - if len(p.keys) == 0 { - return &APIKeyIterator{pool: p} - } - idx := atomic.AddUint32(&p.current, 1) - 1 - return &APIKeyIterator{ - pool: p, - startIdx: idx, - } -} - -func (it *APIKeyIterator) Next() (string, bool) { - length := uint32(len(it.pool.keys)) - if length == 0 || it.attempt >= length { - return "", false - } - key := it.pool.keys[(it.startIdx+it.attempt)%length] - it.attempt++ - return key, true -} - -type SearchProvider interface { - Search(ctx context.Context, query string, count int) (string, error) -} - -type BraveSearchProvider struct { - keyPool *APIKeyPool - proxy string - client *http.Client -} - -func (p *BraveSearchProvider) Search(ctx context.Context, query string, count int) (string, error) { - searchURL := fmt.Sprintf("https://api.search.brave.com/res/v1/web/search?q=%s&count=%d", - url.QueryEscape(query), count) - - var lastErr error - iter := p.keyPool.NewIterator() - - for { - apiKey, ok := iter.Next() - if !ok { - break - } - - req, err := http.NewRequestWithContext(ctx, "GET", searchURL, nil) - if err != nil { - return "", fmt.Errorf("failed to create request: %w", err) - } - - req.Header.Set("Accept", "application/json") - req.Header.Set("X-Subscription-Token", apiKey) - - resp, err := p.client.Do(req) - if err != nil { - lastErr = fmt.Errorf("request failed: %w", err) - continue - } - - body, err := io.ReadAll(resp.Body) - resp.Body.Close() - - if err != nil { - lastErr = fmt.Errorf("failed to read response: %w", err) - continue - } - - if resp.StatusCode != http.StatusOK { - lastErr = fmt.Errorf("API error (status %d): %s", resp.StatusCode, string(body)) - if resp.StatusCode == http.StatusTooManyRequests || - resp.StatusCode == http.StatusUnauthorized || - resp.StatusCode == http.StatusForbidden || - resp.StatusCode >= 500 { - continue - } - return "", lastErr - } - - var searchResp struct { - Web struct { - Results []struct { - Title string `json:"title"` - URL string `json:"url"` - Description string `json:"description"` - } `json:"results"` - } `json:"web"` - } - - if err := json.Unmarshal(body, &searchResp); err != nil { - // Log error body for debugging - return "", fmt.Errorf("failed to parse response: %w", err) - } - - results := searchResp.Web.Results - if len(results) == 0 { - return fmt.Sprintf("No results for: %s", query), nil - } - - var lines []string - lines = append(lines, fmt.Sprintf("Results for: %s", query)) - for i, item := range results { - if i >= count { - break - } - lines = append(lines, fmt.Sprintf("%d. %s\n %s", i+1, item.Title, item.URL)) - if item.Description != "" { - lines = append(lines, fmt.Sprintf(" %s", item.Description)) - } - } - - return strings.Join(lines, "\n"), nil - } - - return "", fmt.Errorf("all api keys failed, last error: %w", lastErr) -} - -type TavilySearchProvider struct { - keyPool *APIKeyPool - baseURL string - proxy string - client *http.Client -} - -func (p *TavilySearchProvider) Search(ctx context.Context, query string, count int) (string, error) { - searchURL := p.baseURL - if searchURL == "" { - searchURL = "https://api.tavily.com/search" - } - - var lastErr error - iter := p.keyPool.NewIterator() - - for { - apiKey, ok := iter.Next() - if !ok { - break - } - - payload := map[string]any{ - "api_key": apiKey, - "query": query, - "search_depth": "advanced", - "include_answer": false, - "include_images": false, - "include_raw_content": false, - "max_results": count, - } - - bodyBytes, err := json.Marshal(payload) - if err != nil { - return "", fmt.Errorf("failed to marshal payload: %w", err) - } - - req, err := http.NewRequestWithContext(ctx, "POST", searchURL, bytes.NewBuffer(bodyBytes)) - if err != nil { - return "", fmt.Errorf("failed to create request: %w", err) - } - - req.Header.Set("Content-Type", "application/json") - req.Header.Set("User-Agent", userAgent) - - resp, err := p.client.Do(req) - if err != nil { - lastErr = fmt.Errorf("request failed: %w", err) - continue - } - - body, err := io.ReadAll(resp.Body) - resp.Body.Close() - - if err != nil { - lastErr = fmt.Errorf("failed to read response: %w", err) - continue - } - - if resp.StatusCode != http.StatusOK { - lastErr = fmt.Errorf("tavily api error (status %d): %s", resp.StatusCode, string(body)) - if resp.StatusCode == http.StatusTooManyRequests || - resp.StatusCode == http.StatusUnauthorized || - resp.StatusCode == http.StatusForbidden || - resp.StatusCode >= 500 { - continue - } - return "", lastErr - } - - var searchResp struct { - Results []struct { - Title string `json:"title"` - URL string `json:"url"` - Content string `json:"content"` - } `json:"results"` - } - - if err := json.Unmarshal(body, &searchResp); err != nil { - return "", fmt.Errorf("failed to parse response: %w", err) - } - - results := searchResp.Results - if len(results) == 0 { - return fmt.Sprintf("No results for: %s", query), nil - } - - var lines []string - lines = append(lines, fmt.Sprintf("Results for: %s (via Tavily)", query)) - for i, item := range results { - if i >= count { - break - } - lines = append(lines, fmt.Sprintf("%d. %s\n %s", i+1, item.Title, item.URL)) - if item.Content != "" { - lines = append(lines, fmt.Sprintf(" %s", item.Content)) - } - } - - return strings.Join(lines, "\n"), nil - } - - return "", fmt.Errorf("all api keys failed, last error: %w", lastErr) -} - -type DuckDuckGoSearchProvider struct { - proxy string - client *http.Client -} - -func (p *DuckDuckGoSearchProvider) Search(ctx context.Context, query string, count int) (string, error) { - searchURL := fmt.Sprintf("https://html.duckduckgo.com/html/?q=%s", url.QueryEscape(query)) - - req, err := http.NewRequestWithContext(ctx, "GET", searchURL, nil) - if err != nil { - return "", fmt.Errorf("failed to create request: %w", err) - } - - req.Header.Set("User-Agent", userAgent) - - resp, err := p.client.Do(req) - if err != nil { - return "", fmt.Errorf("request failed: %w", err) - } - defer resp.Body.Close() - - body, err := io.ReadAll(resp.Body) - if err != nil { - return "", fmt.Errorf("failed to read response: %w", err) - } - - return p.extractResults(string(body), count, query) -} - -func (p *DuckDuckGoSearchProvider) extractResults(html string, count int, query string) (string, error) { - // Simple regex based extraction for DDG HTML - // Strategy: Find all result containers or key anchors directly - - // Try finding the result links directly first, as they are the most critical - // Pattern: Title - // The previous regex was a bit strict. Let's make it more flexible for attributes order/content - matches := reDDGLink.FindAllStringSubmatch(html, count+5) - - if len(matches) == 0 { - return fmt.Sprintf("No results found or extraction failed. Query: %s", query), nil - } - - var lines []string - lines = append(lines, fmt.Sprintf("Results for: %s (via DuckDuckGo)", query)) - - // Pre-compile snippet regex to run inside the loop - // We'll search for snippets relative to the link position or just globally if needed - // But simple global search for snippets might mismatch order. - // Since we only have the raw HTML string, let's just extract snippets globally and assume order matches (risky but simple for regex) - // Or better: Let's assume the snippet follows the link in the HTML - - // A better regex approach: iterate through text and find matches in order - // But for now, let's grab all snippets too - snippetMatches := reDDGSnippet.FindAllStringSubmatch(html, count+5) - - maxItems := min(len(matches), count) - - for i := range maxItems { - urlStr := matches[i][1] - title := stripTags(matches[i][2]) - title = strings.TrimSpace(title) - - // URL decoding if needed - if strings.Contains(urlStr, "uddg=") { - if u, err := url.QueryUnescape(urlStr); err == nil { - _, after, ok := strings.Cut(u, "uddg=") - if ok { - urlStr = after - } - } - } - - lines = append(lines, fmt.Sprintf("%d. %s\n %s", i+1, title, urlStr)) - - // Attempt to attach snippet if available and index aligns - if i < len(snippetMatches) { - snippet := stripTags(snippetMatches[i][1]) - snippet = strings.TrimSpace(snippet) - if snippet != "" { - lines = append(lines, fmt.Sprintf(" %s", snippet)) - } - } - } - - return strings.Join(lines, "\n"), nil -} - -func stripTags(content string) string { - return reTags.ReplaceAllString(content, "") -} - -type PerplexitySearchProvider struct { - keyPool *APIKeyPool - proxy string - client *http.Client -} - -func (p *PerplexitySearchProvider) Search(ctx context.Context, query string, count int) (string, error) { - searchURL := "https://api.perplexity.ai/chat/completions" - - var lastErr error - iter := p.keyPool.NewIterator() - - for { - apiKey, ok := iter.Next() - if !ok { - break - } - - payload := map[string]any{ - "model": "sonar", - "messages": []map[string]string{ - { - "role": "system", - "content": "You are a search assistant. Provide concise search results with titles, URLs, and brief descriptions in the following format:\n1. Title\n URL\n Description\n\nDo not add extra commentary.", - }, - { - "role": "user", - "content": fmt.Sprintf("Search for: %s. Provide up to %d relevant results.", query, count), - }, - }, - "max_tokens": 1000, - } - - payloadBytes, err := json.Marshal(payload) - if err != nil { - return "", fmt.Errorf("failed to marshal request: %w", err) - } - - req, err := http.NewRequestWithContext(ctx, "POST", searchURL, strings.NewReader(string(payloadBytes))) - if err != nil { - return "", fmt.Errorf("failed to create request: %w", err) - } - - req.Header.Set("Content-Type", "application/json") - req.Header.Set("Authorization", "Bearer "+apiKey) - req.Header.Set("User-Agent", userAgent) - - resp, err := p.client.Do(req) - if err != nil { - lastErr = fmt.Errorf("request failed: %w", err) - continue - } - - body, err := io.ReadAll(resp.Body) - resp.Body.Close() - - if err != nil { - lastErr = fmt.Errorf("failed to read response: %w", err) - continue - } - - if resp.StatusCode != http.StatusOK { - lastErr = fmt.Errorf("Perplexity API error: %s", string(body)) - if resp.StatusCode == http.StatusTooManyRequests || - resp.StatusCode == http.StatusUnauthorized || - resp.StatusCode == http.StatusForbidden || - resp.StatusCode >= 500 { - continue - } - return "", lastErr - } - - var searchResp struct { - Choices []struct { - Message struct { - Content string `json:"content"` - } `json:"message"` - } `json:"choices"` - } - - if err := json.Unmarshal(body, &searchResp); err != nil { - return "", fmt.Errorf("failed to parse response: %w", err) - } - - if len(searchResp.Choices) == 0 { - return fmt.Sprintf("No results for: %s", query), nil - } - - return fmt.Sprintf("Results for: %s (via Perplexity)\n%s", query, searchResp.Choices[0].Message.Content), nil - } - - return "", fmt.Errorf("all api keys failed, last error: %w", lastErr) -} - -type SearXNGSearchProvider struct { - baseURL string -} - -func (p *SearXNGSearchProvider) Search(ctx context.Context, query string, count int) (string, error) { - searchURL := fmt.Sprintf("%s/search?q=%s&format=json&categories=general", - strings.TrimSuffix(p.baseURL, "/"), - url.QueryEscape(query)) - - req, err := http.NewRequestWithContext(ctx, "GET", searchURL, nil) - if err != nil { - return "", fmt.Errorf("failed to create request: %w", err) - } - - client := &http.Client{Timeout: 10 * time.Second} - resp, err := client.Do(req) - if err != nil { - return "", fmt.Errorf("request failed: %w", err) - } - defer resp.Body.Close() - - if resp.StatusCode != http.StatusOK { - return "", fmt.Errorf("SearXNG returned status %d", resp.StatusCode) - } - - var result struct { - Results []struct { - Title string `json:"title"` - URL string `json:"url"` - Content string `json:"content"` - Engine string `json:"engine"` - Score float64 `json:"score"` - } `json:"results"` - } - - if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { - return "", fmt.Errorf("failed to parse response: %w", err) - } - - if len(result.Results) == 0 { - return fmt.Sprintf("No results for: %s", query), nil - } - - // Limit results to requested count - if len(result.Results) > count { - result.Results = result.Results[:count] - } - - // Format results in standard PicoClaw format - var b strings.Builder - b.WriteString(fmt.Sprintf("Results for: %s (via SearXNG)\n", query)) - for i, r := range result.Results { - b.WriteString(fmt.Sprintf("%d. %s\n", i+1, r.Title)) - b.WriteString(fmt.Sprintf(" %s\n", r.URL)) - if r.Content != "" { - b.WriteString(fmt.Sprintf(" %s\n", r.Content)) - } - } - - return b.String(), nil -} - -type GLMSearchProvider struct { - apiKey string - baseURL string - searchEngine string - proxy string - client *http.Client -} - -func (p *GLMSearchProvider) Search(ctx context.Context, query string, count int) (string, error) { - searchURL := p.baseURL - if searchURL == "" { - searchURL = "https://open.bigmodel.cn/api/paas/v4/web_search" - } - - payload := map[string]any{ - "search_query": query, - "search_engine": p.searchEngine, - "search_intent": false, - "count": count, - "content_size": "medium", - } - - bodyBytes, err := json.Marshal(payload) - if err != nil { - return "", fmt.Errorf("failed to marshal payload: %w", err) - } - - req, err := http.NewRequestWithContext(ctx, "POST", searchURL, bytes.NewReader(bodyBytes)) - if err != nil { - return "", fmt.Errorf("failed to create request: %w", err) - } - - req.Header.Set("Content-Type", "application/json") - req.Header.Set("Authorization", "Bearer "+p.apiKey) - - resp, err := p.client.Do(req) - if err != nil { - return "", fmt.Errorf("request failed: %w", err) - } - defer resp.Body.Close() - - body, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20)) - if err != nil { - return "", fmt.Errorf("failed to read response: %w", err) - } - - if resp.StatusCode != http.StatusOK { - return "", fmt.Errorf("GLM Search API error (status %d): %s", resp.StatusCode, string(body)) - } - - var searchResp struct { - SearchResult []struct { - Title string `json:"title"` - Content string `json:"content"` - Link string `json:"link"` - } `json:"search_result"` - } - - if err := json.Unmarshal(body, &searchResp); err != nil { - return "", fmt.Errorf("failed to parse response: %w", err) - } - - results := searchResp.SearchResult - if len(results) == 0 { - return fmt.Sprintf("No results for: %s", query), nil - } - - var lines []string - lines = append(lines, fmt.Sprintf("Results for: %s (via GLM Search)", query)) - for i, item := range results { - if i >= count { - break - } - lines = append(lines, fmt.Sprintf("%d. %s\n %s", i+1, item.Title, item.Link)) - if item.Content != "" { - lines = append(lines, fmt.Sprintf(" %s", item.Content)) - } - } - - return strings.Join(lines, "\n"), nil -} - -type WebSearchTool struct { - provider SearchProvider - maxResults int -} - -type WebSearchToolOptions struct { - BraveAPIKeys []string - BraveMaxResults int - BraveEnabled bool - TavilyAPIKeys []string - TavilyBaseURL string - TavilyMaxResults int - TavilyEnabled bool - DuckDuckGoMaxResults int - DuckDuckGoEnabled bool - PerplexityAPIKeys []string - PerplexityMaxResults int - PerplexityEnabled bool - SearXNGBaseURL string - SearXNGMaxResults int - SearXNGEnabled bool - GLMSearchAPIKey string - GLMSearchBaseURL string - GLMSearchEngine string - GLMSearchMaxResults int - GLMSearchEnabled bool - Proxy string -} - -func NewWebSearchTool(opts WebSearchToolOptions) (*WebSearchTool, error) { - var provider SearchProvider - maxResults := 5 - // Priority: Perplexity > Brave > SearXNG > Tavily > DuckDuckGo > GLM Search - if opts.PerplexityEnabled && len(opts.PerplexityAPIKeys) > 0 { - client, err := createHTTPClient(opts.Proxy, perplexityTimeout) - if err != nil { - return nil, fmt.Errorf("failed to create HTTP client for Perplexity: %w", err) - } - provider = &PerplexitySearchProvider{ - keyPool: NewAPIKeyPool(opts.PerplexityAPIKeys), - proxy: opts.Proxy, - client: client, - } - if opts.PerplexityMaxResults > 0 { - maxResults = opts.PerplexityMaxResults - } - } else if opts.BraveEnabled && len(opts.BraveAPIKeys) > 0 { - client, err := createHTTPClient(opts.Proxy, searchTimeout) - if err != nil { - return nil, fmt.Errorf("failed to create HTTP client for Brave: %w", err) - } - provider = &BraveSearchProvider{keyPool: NewAPIKeyPool(opts.BraveAPIKeys), proxy: opts.Proxy, client: client} - if opts.BraveMaxResults > 0 { - maxResults = opts.BraveMaxResults - } - } else if opts.SearXNGEnabled && opts.SearXNGBaseURL != "" { - provider = &SearXNGSearchProvider{baseURL: opts.SearXNGBaseURL} - if opts.SearXNGMaxResults > 0 { - maxResults = opts.SearXNGMaxResults - } - } else if opts.TavilyEnabled && len(opts.TavilyAPIKeys) > 0 { - client, err := createHTTPClient(opts.Proxy, searchTimeout) - if err != nil { - return nil, fmt.Errorf("failed to create HTTP client for Tavily: %w", err) - } - provider = &TavilySearchProvider{ - keyPool: NewAPIKeyPool(opts.TavilyAPIKeys), - baseURL: opts.TavilyBaseURL, - proxy: opts.Proxy, - client: client, - } - if opts.TavilyMaxResults > 0 { - maxResults = opts.TavilyMaxResults - } - } else if opts.DuckDuckGoEnabled { - 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 if opts.GLMSearchEnabled && opts.GLMSearchAPIKey != "" { - client, err := createHTTPClient(opts.Proxy, searchTimeout) - if err != nil { - return nil, fmt.Errorf("failed to create HTTP client for GLM Search: %w", err) - } - searchEngine := opts.GLMSearchEngine - if searchEngine == "" { - searchEngine = "search_std" - } - provider = &GLMSearchProvider{ - apiKey: opts.GLMSearchAPIKey, - baseURL: opts.GLMSearchBaseURL, - searchEngine: searchEngine, - proxy: opts.Proxy, - client: client, - } - if opts.GLMSearchMaxResults > 0 { - maxResults = opts.GLMSearchMaxResults - } - } else { - return nil, nil - } - - return &WebSearchTool{ - provider: provider, - maxResults: maxResults, - }, nil -} - -func (t *WebSearchTool) Name() string { - return "web_search" -} - -func (t *WebSearchTool) Description() string { - return "Search the web for current information. Returns titles, URLs, and snippets from search results." -} - -func (t *WebSearchTool) Parameters() map[string]any { - return map[string]any{ - "type": "object", - "properties": map[string]any{ - "query": map[string]any{ - "type": "string", - "description": "Search query", - }, - "count": map[string]any{ - "type": "integer", - "description": "Number of results (1-10)", - "minimum": 1.0, - "maximum": 10.0, - }, - }, - "required": []string{"query"}, - } -} - -func (t *WebSearchTool) Execute(ctx context.Context, args map[string]any) *ToolResult { - query, ok := args["query"].(string) - if !ok { - return ErrorResult("query is required") - } - - count := t.maxResults - if c, ok := args["count"].(float64); ok { - if int(c) > 0 && int(c) <= 10 { - count = int(c) - } - } - - result, err := t.provider.Search(ctx, query, count) - if err != nil { - return ErrorResult(fmt.Sprintf("search failed: %v", err)) - } - - return &ToolResult{ - ForLLM: result, - ForUser: result, - } -} - -type WebFetchTool struct { - maxChars int - proxy string - client *http.Client - fetchLimitBytes int64 -} - -func NewWebFetchTool(maxChars int, fetchLimitBytes int64) (*WebFetchTool, error) { - // createHTTPClient cannot fail with an empty proxy string. - return NewWebFetchToolWithProxy(maxChars, "", fetchLimitBytes) -} - -// allowPrivateWebFetchHosts controls whether loopback/private hosts are allowed. -// This is false in normal runtime to reduce SSRF exposure, and tests can override it temporarily. -var allowPrivateWebFetchHosts atomic.Bool - -func NewWebFetchToolWithProxy(maxChars int, proxy string, fetchLimitBytes int64) (*WebFetchTool, error) { - if maxChars <= 0 { - maxChars = defaultMaxChars - } - client, err := createHTTPClient(proxy, fetchTimeout) - if err != nil { - return nil, fmt.Errorf("failed to create HTTP client for web fetch: %w", err) - } - if transport, ok := client.Transport.(*http.Transport); ok { - dialer := &net.Dialer{ - Timeout: 15 * time.Second, - KeepAlive: 30 * time.Second, - } - transport.DialContext = newSafeDialContext(dialer) - } - client.CheckRedirect = func(req *http.Request, via []*http.Request) error { - if len(via) >= maxRedirects { - return fmt.Errorf("stopped after %d redirects", maxRedirects) - } - if isObviousPrivateHost(req.URL.Hostname()) { - return fmt.Errorf("redirect target is private or local network host") - } - return nil - } - if fetchLimitBytes <= 0 { - fetchLimitBytes = 10 * 1024 * 1024 // Security Fallback - } - return &WebFetchTool{ - maxChars: maxChars, - proxy: proxy, - client: client, - fetchLimitBytes: fetchLimitBytes, - }, nil -} - -func (t *WebFetchTool) Name() string { - return "web_fetch" -} - -func (t *WebFetchTool) Description() string { - return "Fetch a URL and extract readable content (HTML to text). Use this to get weather info, news, articles, or any web content." -} - -func (t *WebFetchTool) Parameters() map[string]any { - return map[string]any{ - "type": "object", - "properties": map[string]any{ - "url": map[string]any{ - "type": "string", - "description": "URL to fetch", - }, - "maxChars": map[string]any{ - "type": "integer", - "description": "Maximum characters to extract", - "minimum": 100.0, - }, - }, - "required": []string{"url"}, - } -} - -func (t *WebFetchTool) Execute(ctx context.Context, args map[string]any) *ToolResult { - urlStr, ok := args["url"].(string) - if !ok { - return ErrorResult("url is required") - } - - parsedURL, err := url.Parse(urlStr) - if err != nil { - return ErrorResult(fmt.Sprintf("invalid URL: %v", err)) - } - - if parsedURL.Scheme != "http" && parsedURL.Scheme != "https" { - return ErrorResult("only http/https URLs are allowed") - } - - if parsedURL.Host == "" { - return ErrorResult("missing domain in URL") - } - - // Lightweight pre-flight: block obvious localhost/literal-IP without DNS resolution. - // The real SSRF guard is newSafeDialContext at connect time. - hostname := parsedURL.Hostname() - if isObviousPrivateHost(hostname) { - return ErrorResult("fetching private or local network hosts is not allowed") - } - - maxChars := t.maxChars - if mc, ok := args["maxChars"].(float64); ok { - if int(mc) > 100 { - maxChars = int(mc) - } - } - - req, err := http.NewRequestWithContext(ctx, "GET", urlStr, nil) - if err != nil { - return ErrorResult(fmt.Sprintf("failed to create request: %v", err)) - } - - req.Header.Set("User-Agent", userAgent) - resp, err := t.client.Do(req) - if err != nil { - return ErrorResult(fmt.Sprintf("request failed: %v", err)) - } - - resp.Body = http.MaxBytesReader(nil, resp.Body, t.fetchLimitBytes) - - defer resp.Body.Close() - - body, err := io.ReadAll(resp.Body) - if err != nil { - var maxBytesErr *http.MaxBytesError - if errors.As(err, &maxBytesErr) { - return ErrorResult(fmt.Sprintf("failed to read response: size exceeded %d bytes limit", t.fetchLimitBytes)) - } - return ErrorResult(fmt.Sprintf("failed to read response: %v", err)) - } - - contentType := resp.Header.Get("Content-Type") - - var text, extractor string - - if strings.Contains(contentType, "application/json") { - var jsonData any - if err := json.Unmarshal(body, &jsonData); err == nil { - formatted, _ := json.MarshalIndent(jsonData, "", " ") - text = string(formatted) - extractor = "json" - } else { - text = string(body) - extractor = "raw" - } - } else if strings.Contains(contentType, "text/html") || len(body) > 0 && - (strings.HasPrefix(string(body), " maxChars - if truncated { - text = text[:maxChars] - } - - result := map[string]any{ - "url": urlStr, - "status": resp.StatusCode, - "extractor": extractor, - "truncated": truncated, - "length": len(text), - "text": text, - } - - resultJSON, _ := json.MarshalIndent(result, "", " ") - - return &ToolResult{ - ForLLM: string(resultJSON), - ForUser: fmt.Sprintf( - "Fetched %d bytes from %s (extractor: %s, truncated: %v)", - len(text), - urlStr, - extractor, - truncated, - ), - } -} - -func (t *WebFetchTool) extractText(htmlContent string) string { - result := reScript.ReplaceAllLiteralString(htmlContent, "") - result = reStyle.ReplaceAllLiteralString(result, "") - result = reTags.ReplaceAllLiteralString(result, "") - - result = strings.TrimSpace(result) - - result = reWhitespace.ReplaceAllString(result, " ") - result = reBlankLines.ReplaceAllString(result, "\n\n") - - lines := strings.Split(result, "\n") - var cleanLines []string - for _, line := range lines { - line = strings.TrimSpace(line) - if line != "" { - cleanLines = append(cleanLines, line) - } - } - - return strings.Join(cleanLines, "\n") -} - -// newSafeDialContext re-resolves DNS at connect time to mitigate DNS rebinding (TOCTOU) -// where a hostname resolves to a public IP during pre-flight but a private IP at connect time. -func newSafeDialContext(dialer *net.Dialer) func(context.Context, string, string) (net.Conn, error) { - return func(ctx context.Context, network, address string) (net.Conn, error) { - if allowPrivateWebFetchHosts.Load() { - return dialer.DialContext(ctx, network, address) - } - - host, port, err := net.SplitHostPort(address) - if err != nil { - return nil, fmt.Errorf("invalid target address %q: %w", address, err) - } - if host == "" { - return nil, fmt.Errorf("empty target host") - } - - if ip := net.ParseIP(host); ip != nil { - if isPrivateOrRestrictedIP(ip) { - return nil, fmt.Errorf("blocked private or local target: %s", host) - } - return dialer.DialContext(ctx, network, net.JoinHostPort(ip.String(), port)) - } - - ipAddrs, err := net.DefaultResolver.LookupIPAddr(ctx, host) - if err != nil { - return nil, fmt.Errorf("failed to resolve %s: %w", host, err) - } - - attempted := 0 - var lastErr error - for _, ipAddr := range ipAddrs { - if isPrivateOrRestrictedIP(ipAddr.IP) { - continue - } - attempted++ - conn, err := dialer.DialContext(ctx, network, net.JoinHostPort(ipAddr.IP.String(), port)) - if err == nil { - return conn, nil - } - lastErr = err - } - - if attempted == 0 { - return nil, fmt.Errorf("all resolved addresses for %s are private or restricted", host) - } - if lastErr != nil { - return nil, fmt.Errorf("failed connecting to public addresses for %s: %w", host, lastErr) - } - return nil, fmt.Errorf("failed connecting to public addresses for %s", host) - } -} - -// isObviousPrivateHost performs a lightweight, no-DNS check for obviously private hosts. -// It catches localhost, literal private IPs, and empty hosts. It does NOT resolve DNS — -// the real SSRF guard is newSafeDialContext which checks IPs at connect time. -func isObviousPrivateHost(host string) bool { - if allowPrivateWebFetchHosts.Load() { - return false - } - - h := strings.ToLower(strings.TrimSpace(host)) - h = strings.TrimSuffix(h, ".") - if h == "" { - return true - } - - if h == "localhost" || strings.HasSuffix(h, ".localhost") { - return true - } - - if ip := net.ParseIP(h); ip != nil { - return isPrivateOrRestrictedIP(ip) - } - - return false -} - -// isPrivateOrRestrictedIP returns true for IPs that should never be reached via web_fetch: -// RFC 1918, loopback, link-local (incl. cloud metadata 169.254.x.x), carrier-grade NAT, -// IPv6 unique-local (fc00::/7), 6to4 (2002::/16), and Teredo (2001:0000::/32). -func isPrivateOrRestrictedIP(ip net.IP) bool { - if ip == nil { - return true - } - - if ip.IsLoopback() || ip.IsLinkLocalUnicast() || ip.IsLinkLocalMulticast() || - ip.IsMulticast() || ip.IsUnspecified() { - return true - } - - if ip4 := ip.To4(); ip4 != nil { - // IPv4 private, loopback, link-local, and carrier-grade NAT ranges. - if ip4[0] == 10 || - ip4[0] == 127 || - ip4[0] == 0 || - (ip4[0] == 172 && ip4[1] >= 16 && ip4[1] <= 31) || - (ip4[0] == 192 && ip4[1] == 168) || - (ip4[0] == 169 && ip4[1] == 254) || - (ip4[0] == 100 && ip4[1] >= 64 && ip4[1] <= 127) { - return true - } - return false - } - - if len(ip) == net.IPv6len { - // IPv6 unique local addresses (fc00::/7) - if (ip[0] & 0xfe) == 0xfc { - return true - } - // 6to4 addresses (2002::/16): check the embedded IPv4 at bytes [2:6]. - if ip[0] == 0x20 && ip[1] == 0x02 { - embedded := net.IPv4(ip[2], ip[3], ip[4], ip[5]) - return isPrivateOrRestrictedIP(embedded) - } - // Teredo (2001:0000::/32): client IPv4 is at bytes [12:16], XOR-inverted. - if ip[0] == 0x20 && ip[1] == 0x01 && ip[2] == 0x00 && ip[3] == 0x00 { - client := net.IPv4(ip[12]^0xff, ip[13]^0xff, ip[14]^0xff, ip[15]^0xff) - return isPrivateOrRestrictedIP(client) - } - } - - return false -} diff --git a/pkg/tools/web_common.go b/pkg/tools/web_common.go new file mode 100644 index 000000000..2c2fb007a --- /dev/null +++ b/pkg/tools/web_common.go @@ -0,0 +1,76 @@ +package tools + +import ( + "fmt" + "net/http" + "net/url" + "regexp" + "strings" + "time" +) + +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 +var ( + reScript = regexp.MustCompile(``) + reStyle = regexp.MustCompile(``) + reTags = regexp.MustCompile(`<[^>]+>`) + reWhitespace = regexp.MustCompile(`[^\S\n]+`) + reBlankLines = regexp.MustCompile(`\n{3,}`) + + // DuckDuckGo result extraction + reDDGLink = regexp.MustCompile(`]*class="[^"]*result__a[^"]*"[^>]*href="([^"]+)"[^>]*>([\s\S]*?)`) + reDDGSnippet = regexp.MustCompile(`([\s\S]*?)`) +) + +// createHTTPClient creates an HTTP client with optional proxy support +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, + }, + } + + if proxyURL != "" { + proxy, err := url.Parse(proxyURL) + if err != nil { + return nil, fmt.Errorf("invalid proxy URL: %w", err) + } + scheme := strings.ToLower(proxy.Scheme) + switch scheme { + case "http", "https", "socks5", "socks5h": + default: + return nil, fmt.Errorf( + "unsupported proxy scheme %q (supported: http, https, socks5, socks5h)", + proxy.Scheme, + ) + } + if proxy.Host == "" { + return nil, fmt.Errorf("invalid proxy URL: missing host") + } + client.Transport.(*http.Transport).Proxy = http.ProxyURL(proxy) + } else { + client.Transport.(*http.Transport).Proxy = http.ProxyFromEnvironment + } + + return client, nil +} + +func stripTags(content string) string { + return reTags.ReplaceAllString(content, "") +} diff --git a/pkg/tools/web_common_test.go b/pkg/tools/web_common_test.go new file mode 100644 index 000000000..11db8624f --- /dev/null +++ b/pkg/tools/web_common_test.go @@ -0,0 +1,110 @@ +package tools + +import ( + "net/http" + "strings" + "testing" + "time" +) + +func TestCreateHTTPClient_ProxyConfigured(t *testing.T) { + client, err := createHTTPClient("http://127.0.0.1:7890", 12*time.Second) + if err != nil { + t.Fatalf("createHTTPClient() error: %v", err) + } + if client.Timeout != 12*time.Second { + t.Fatalf("client.Timeout = %v, want %v", client.Timeout, 12*time.Second) + } + + tr, ok := client.Transport.(*http.Transport) + if !ok { + t.Fatalf("client.Transport type = %T, want *http.Transport", client.Transport) + } + if tr.Proxy == nil { + t.Fatal("transport.Proxy is nil, want non-nil") + } + + req, err := http.NewRequest("GET", "https://example.com", nil) + if err != nil { + t.Fatalf("http.NewRequest() error: %v", err) + } + proxyURL, err := tr.Proxy(req) + if err != nil { + t.Fatalf("transport.Proxy(req) error: %v", err) + } + if proxyURL == nil || proxyURL.String() != "http://127.0.0.1:7890" { + t.Fatalf("proxy URL = %v, want %q", proxyURL, "http://127.0.0.1:7890") + } +} + +func TestCreateHTTPClient_InvalidProxy(t *testing.T) { + _, err := createHTTPClient("://bad-proxy", 10*time.Second) + if err == nil { + t.Fatal("createHTTPClient() expected error for invalid proxy URL, got nil") + } +} + +func TestCreateHTTPClient_Socks5ProxyConfigured(t *testing.T) { + client, err := createHTTPClient("socks5://127.0.0.1:1080", 8*time.Second) + if err != nil { + t.Fatalf("createHTTPClient() error: %v", err) + } + + tr, ok := client.Transport.(*http.Transport) + if !ok { + t.Fatalf("client.Transport type = %T, want *http.Transport", client.Transport) + } + req, err := http.NewRequest("GET", "https://example.com", nil) + if err != nil { + t.Fatalf("http.NewRequest() error: %v", err) + } + proxyURL, err := tr.Proxy(req) + if err != nil { + t.Fatalf("transport.Proxy(req) error: %v", err) + } + if proxyURL == nil || proxyURL.String() != "socks5://127.0.0.1:1080" { + t.Fatalf("proxy URL = %v, want %q", proxyURL, "socks5://127.0.0.1:1080") + } +} + +func TestCreateHTTPClient_UnsupportedProxyScheme(t *testing.T) { + _, err := createHTTPClient("ftp://127.0.0.1:21", 10*time.Second) + if err == nil { + t.Fatal("createHTTPClient() expected error for unsupported scheme, got nil") + } + if !strings.Contains(err.Error(), "unsupported proxy scheme") { + t.Fatalf("error = %q, want to contain %q", err.Error(), "unsupported proxy scheme") + } +} + +func TestCreateHTTPClient_ProxyFromEnvironmentWhenConfigEmpty(t *testing.T) { + t.Setenv("HTTP_PROXY", "http://127.0.0.1:8888") + t.Setenv("http_proxy", "http://127.0.0.1:8888") + t.Setenv("HTTPS_PROXY", "http://127.0.0.1:8888") + t.Setenv("https_proxy", "http://127.0.0.1:8888") + t.Setenv("ALL_PROXY", "") + t.Setenv("all_proxy", "") + t.Setenv("NO_PROXY", "") + t.Setenv("no_proxy", "") + + client, err := createHTTPClient("", 10*time.Second) + if err != nil { + t.Fatalf("createHTTPClient() error: %v", err) + } + + tr, ok := client.Transport.(*http.Transport) + if !ok { + t.Fatalf("client.Transport type = %T, want *http.Transport", client.Transport) + } + if tr.Proxy == nil { + t.Fatal("transport.Proxy is nil, want proxy function from environment") + } + + req, err := http.NewRequest("GET", "https://example.com", nil) + if err != nil { + t.Fatalf("http.NewRequest() error: %v", err) + } + if _, err := tr.Proxy(req); err != nil { + t.Fatalf("transport.Proxy(req) error: %v", err) + } +} diff --git a/pkg/tools/web_fetch_tool.go b/pkg/tools/web_fetch_tool.go new file mode 100644 index 000000000..8a3a6b974 --- /dev/null +++ b/pkg/tools/web_fetch_tool.go @@ -0,0 +1,217 @@ +package tools + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net" + "net/http" + "net/url" + "strings" + "time" +) + +type WebFetchTool struct { + maxChars int + proxy string + client *http.Client + fetchLimitBytes int64 +} + +func NewWebFetchTool(maxChars int, fetchLimitBytes int64) (*WebFetchTool, error) { + // createHTTPClient cannot fail with an empty proxy string. + return NewWebFetchToolWithProxy(maxChars, "", fetchLimitBytes) +} + +func NewWebFetchToolWithProxy(maxChars int, proxy string, fetchLimitBytes int64) (*WebFetchTool, error) { + if maxChars <= 0 { + maxChars = defaultMaxChars + } + client, err := createHTTPClient(proxy, fetchTimeout) + if err != nil { + return nil, fmt.Errorf("failed to create HTTP client for web fetch: %w", err) + } + if transport, ok := client.Transport.(*http.Transport); ok { + dialer := &net.Dialer{ + Timeout: 15 * time.Second, + KeepAlive: 30 * time.Second, + } + transport.DialContext = newSafeDialContext(dialer) + } + client.CheckRedirect = func(req *http.Request, via []*http.Request) error { + if len(via) >= maxRedirects { + return fmt.Errorf("stopped after %d redirects", maxRedirects) + } + if isObviousPrivateHost(req.URL.Hostname()) { + return fmt.Errorf("redirect target is private or local network host") + } + return nil + } + if fetchLimitBytes <= 0 { + fetchLimitBytes = 10 * 1024 * 1024 // Security Fallback + } + return &WebFetchTool{ + maxChars: maxChars, + proxy: proxy, + client: client, + fetchLimitBytes: fetchLimitBytes, + }, nil +} + +func (t *WebFetchTool) Name() string { + return "web_fetch" +} + +func (t *WebFetchTool) Description() string { + return "Fetch a URL and extract readable content (HTML to text). Use this to get weather info, news, articles, or any web content." +} + +func (t *WebFetchTool) Parameters() map[string]any { + return map[string]any{ + "type": "object", + "properties": map[string]any{ + "url": map[string]any{ + "type": "string", + "description": "URL to fetch", + }, + "maxChars": map[string]any{ + "type": "integer", + "description": "Maximum characters to extract", + "minimum": 100.0, + }, + }, + "required": []string{"url"}, + } +} + +func (t *WebFetchTool) Execute(ctx context.Context, args map[string]any) *ToolResult { + urlStr, ok := args["url"].(string) + if !ok { + return ErrorResult("url is required") + } + + parsedURL, err := url.Parse(urlStr) + if err != nil { + return ErrorResult(fmt.Sprintf("invalid URL: %v", err)) + } + + if parsedURL.Scheme != "http" && parsedURL.Scheme != "https" { + return ErrorResult("only http/https URLs are allowed") + } + + if parsedURL.Host == "" { + return ErrorResult("missing domain in URL") + } + + // Lightweight pre-flight: block obvious localhost/literal-IP without DNS resolution. + // The real SSRF guard is newSafeDialContext at connect time. + hostname := parsedURL.Hostname() + if isObviousPrivateHost(hostname) { + return ErrorResult("fetching private or local network hosts is not allowed") + } + + maxChars := t.maxChars + if mc, ok := args["maxChars"].(float64); ok { + if int(mc) > 100 { + maxChars = int(mc) + } + } + + req, err := http.NewRequestWithContext(ctx, "GET", urlStr, nil) + if err != nil { + return ErrorResult(fmt.Sprintf("failed to create request: %v", err)) + } + + req.Header.Set("User-Agent", userAgent) + resp, err := t.client.Do(req) + if err != nil { + return ErrorResult(fmt.Sprintf("request failed: %v", err)) + } + + resp.Body = http.MaxBytesReader(nil, resp.Body, t.fetchLimitBytes) + + defer resp.Body.Close() + + body, err := io.ReadAll(resp.Body) + if err != nil { + var maxBytesErr *http.MaxBytesError + if errors.As(err, &maxBytesErr) { + return ErrorResult(fmt.Sprintf("failed to read response: size exceeded %d bytes limit", t.fetchLimitBytes)) + } + return ErrorResult(fmt.Sprintf("failed to read response: %v", err)) + } + + contentType := resp.Header.Get("Content-Type") + + var text, extractor string + + if strings.Contains(contentType, "application/json") { + var jsonData any + if err := json.Unmarshal(body, &jsonData); err == nil { + formatted, _ := json.MarshalIndent(jsonData, "", " ") + text = string(formatted) + extractor = "json" + } else { + text = string(body) + extractor = "raw" + } + } else if strings.Contains(contentType, "text/html") || len(body) > 0 && + (strings.HasPrefix(string(body), " maxChars + if truncated { + text = text[:maxChars] + } + + result := map[string]any{ + "url": urlStr, + "status": resp.StatusCode, + "extractor": extractor, + "truncated": truncated, + "length": len(text), + "text": text, + } + + resultJSON, _ := json.MarshalIndent(result, "", " ") + + return &ToolResult{ + ForLLM: string(resultJSON), + ForUser: fmt.Sprintf( + "Fetched %d bytes from %s (extractor: %s, truncated: %v)", + len(text), + urlStr, + extractor, + truncated, + ), + } +} + +func (t *WebFetchTool) extractText(htmlContent string) string { + result := reScript.ReplaceAllLiteralString(htmlContent, "") + result = reStyle.ReplaceAllLiteralString(result, "") + result = reTags.ReplaceAllLiteralString(result, "") + + result = strings.TrimSpace(result) + + result = reWhitespace.ReplaceAllString(result, " ") + result = reBlankLines.ReplaceAllString(result, "\n\n") + + lines := strings.Split(result, "\n") + var cleanLines []string + for _, line := range lines { + line = strings.TrimSpace(line) + if line != "" { + cleanLines = append(cleanLines, line) + } + } + + return strings.Join(cleanLines, "\n") +} diff --git a/pkg/tools/web_fetch_tool_test.go b/pkg/tools/web_fetch_tool_test.go new file mode 100644 index 000000000..9e8ae3a6b --- /dev/null +++ b/pkg/tools/web_fetch_tool_test.go @@ -0,0 +1,381 @@ +package tools + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "jane/pkg/logger" +) + +const testFetchLimit = int64(10 * 1024 * 1024) + +func TestWebTool_WebFetch_Success(t *testing.T) { + withPrivateWebFetchHostsAllowed(t) + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/html") + w.WriteHeader(http.StatusOK) + w.Write([]byte("

Test Page

Content here

")) + })) + defer server.Close() + + tool, err := NewWebFetchTool(50000, testFetchLimit) + if err != nil { + t.Fatalf("Failed to create web fetch tool: %v", err) + } + + ctx := context.Background() + args := map[string]any{ + "url": server.URL, + } + + result := tool.Execute(ctx, args) + + if result.IsError { + t.Errorf("Expected success, got IsError=true: %s", result.ForLLM) + } + + if !strings.Contains(result.ForLLM, "Test Page") { + t.Errorf("Expected ForLLM to contain 'Test Page', got: %s", result.ForLLM) + } + + if !strings.Contains(result.ForUser, "bytes") && !strings.Contains(result.ForUser, "extractor") { + t.Errorf("Expected ForUser to contain summary, got: %s", result.ForUser) + } +} + +func TestWebTool_WebFetch_JSON(t *testing.T) { + withPrivateWebFetchHostsAllowed(t) + + testData := map[string]string{"key": "value", "number": "123"} + expectedJSON, _ := json.MarshalIndent(testData, "", " ") + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + w.Write(expectedJSON) + })) + defer server.Close() + + tool, err := NewWebFetchTool(50000, testFetchLimit) + if err != nil { + logger.ErrorCF("agent", "Failed to create web fetch tool", map[string]any{"error": err.Error()}) + } + + ctx := context.Background() + args := map[string]any{ + "url": server.URL, + } + + result := tool.Execute(ctx, args) + + if result.IsError { + t.Errorf("Expected success, got IsError=true: %s", result.ForLLM) + } + + if !strings.Contains(result.ForLLM, "key") && !strings.Contains(result.ForLLM, "value") { + t.Errorf("Expected ForLLM to contain JSON data, got: %s", result.ForLLM) + } +} + +func TestWebTool_WebFetch_InvalidURL(t *testing.T) { + tool, err := NewWebFetchTool(50000, testFetchLimit) + if err != nil { + logger.ErrorCF("agent", "Failed to create web fetch tool", map[string]any{"error": err.Error()}) + } + + ctx := context.Background() + args := map[string]any{ + "url": "not-a-valid-url", + } + + result := tool.Execute(ctx, args) + + if !result.IsError { + t.Errorf("Expected error for invalid URL") + } + + if !strings.Contains(result.ForLLM, "URL") && !strings.Contains(result.ForUser, "URL") { + t.Errorf("Expected error message for invalid URL, got ForLLM: %s", result.ForLLM) + } +} + +func TestWebTool_WebFetch_UnsupportedScheme(t *testing.T) { + tool, err := NewWebFetchTool(50000, testFetchLimit) + if err != nil { + logger.ErrorCF("agent", "Failed to create web fetch tool", map[string]any{"error": err.Error()}) + } + + ctx := context.Background() + args := map[string]any{ + "url": "ftp://example.com/file.txt", + } + + result := tool.Execute(ctx, args) + + if !result.IsError { + t.Errorf("Expected error for unsupported URL scheme") + } + + if !strings.Contains(result.ForLLM, "http/https") && !strings.Contains(result.ForUser, "http/https") { + t.Errorf("Expected scheme error message, got ForLLM: %s", result.ForLLM) + } +} + +func TestWebTool_WebFetch_MissingURL(t *testing.T) { + tool, err := NewWebFetchTool(50000, testFetchLimit) + if err != nil { + logger.ErrorCF("agent", "Failed to create web fetch tool", map[string]any{"error": err.Error()}) + } + + ctx := context.Background() + args := map[string]any{} + + result := tool.Execute(ctx, args) + + if !result.IsError { + t.Errorf("Expected error when URL is missing") + } + + if !strings.Contains(result.ForLLM, "url is required") && !strings.Contains(result.ForUser, "url is required") { + t.Errorf("Expected 'url is required' message, got ForLLM: %s", result.ForLLM) + } +} + +func TestWebTool_WebFetch_Truncation(t *testing.T) { + withPrivateWebFetchHostsAllowed(t) + + longContent := strings.Repeat("x", 20000) + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/plain") + w.WriteHeader(http.StatusOK) + w.Write([]byte(longContent)) + })) + defer server.Close() + + tool, err := NewWebFetchTool(1000, testFetchLimit) + if err != nil { + logger.ErrorCF("agent", "Failed to create web fetch tool", map[string]any{"error": err.Error()}) + } + + ctx := context.Background() + args := map[string]any{ + "url": server.URL, + } + + result := tool.Execute(ctx, args) + + if result.IsError { + t.Errorf("Expected success, got IsError=true: %s", result.ForLLM) + } + + resultMap := make(map[string]any) + json.Unmarshal([]byte(result.ForLLM), &resultMap) + if text, ok := resultMap["text"].(string); ok { + if len(text) > 1100 { + t.Errorf("Expected content to be truncated to ~1000 chars, got: %d", len(text)) + } + } + + if truncated, ok := resultMap["truncated"].(bool); !ok || !truncated { + t.Errorf("Expected 'truncated' to be true in result") + } +} + +func TestWebFetchTool_PayloadTooLarge(t *testing.T) { + withPrivateWebFetchHostsAllowed(t) + + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/html") + w.WriteHeader(http.StatusOK) + largeData := bytes.Repeat([]byte("A"), int(testFetchLimit)+100) + w.Write(largeData) + })) + defer ts.Close() + + tool, err := NewWebFetchTool(50000, testFetchLimit) + if err != nil { + logger.ErrorCF("agent", "Failed to create web fetch tool", map[string]any{"error": err.Error()}) + } + + args := map[string]any{ + "url": ts.URL, + } + + ctx := context.Background() + result := tool.Execute(ctx, args) + + if result == nil { + t.Fatal("expected a ToolResult, got nil") + } + + expectedErrorMsg := fmt.Sprintf("size exceeded %d bytes limit", testFetchLimit) + + if !strings.Contains(result.ForLLM, expectedErrorMsg) && !strings.Contains(result.ForUser, expectedErrorMsg) { + t.Errorf("test failed: expected error %q, but got: %+v", expectedErrorMsg, result) + } +} + +func TestWebTool_WebFetch_HTMLExtraction(t *testing.T) { + withPrivateWebFetchHostsAllowed(t) + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/html") + w.WriteHeader(http.StatusOK) + w.Write( + []byte( + `

Title

Content

`, + ), + ) + })) + defer server.Close() + + tool, err := NewWebFetchTool(50000, testFetchLimit) + if err != nil { + logger.ErrorCF("agent", "Failed to create web fetch tool", map[string]any{"error": err.Error()}) + } + + ctx := context.Background() + args := map[string]any{ + "url": server.URL, + } + + result := tool.Execute(ctx, args) + + if result.IsError { + t.Errorf("Expected success, got IsError=true: %s", result.ForLLM) + } + + if !strings.Contains(result.ForLLM, "Title") && !strings.Contains(result.ForLLM, "Content") { + t.Errorf("Expected ForLLM to contain extracted text, got: %s", result.ForLLM) + } + + if strings.Contains(result.ForLLM, "

Keep this

", + wantFunc: func(t *testing.T, got string) { + if strings.Contains(got, "alert") || strings.Contains(got, "body{}") { + t.Errorf("Expected script/style content removed, got: %q", got) + } + if !strings.Contains(got, "Keep this") { + t.Errorf("Expected 'Keep this' to remain, got: %q", got) + } + }, + }, + { + name: "collapses excessive blank lines", + input: "

A

\n\n\n\n\n

B

", + wantFunc: func(t *testing.T, got string) { + if strings.Contains(got, "\n\n\n") { + t.Errorf("Expected excessive blank lines collapsed, got: %q", got) + } + }, + }, + { + name: "collapses horizontal whitespace", + input: "

hello world

", + wantFunc: func(t *testing.T, got string) { + if strings.Contains(got, " ") { + t.Errorf("Expected spaces collapsed, got: %q", got) + } + if !strings.Contains(got, "hello world") { + t.Errorf("Expected 'hello world', got: %q", got) + } + }, + }, + { + name: "empty input", + input: "", + wantFunc: func(t *testing.T, got string) { + if got != "" { + t.Errorf("Expected empty string, got: %q", got) + } + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := tool.extractText(tt.input) + tt.wantFunc(t, got) + }) + } +} + +func TestWebTool_WebFetch_MissingDomain(t *testing.T) { + tool, err := NewWebFetchTool(50000, testFetchLimit) + if err != nil { + logger.ErrorCF("agent", "Failed to create web fetch tool", map[string]any{"error": err.Error()}) + } + + ctx := context.Background() + args := map[string]any{ + "url": "https://", + } + + result := tool.Execute(ctx, args) + + if !result.IsError { + t.Errorf("Expected error for URL without domain") + } + + if !strings.Contains(result.ForLLM, "domain") && !strings.Contains(result.ForUser, "domain") { + t.Errorf("Expected domain error message, got ForLLM: %s", result.ForLLM) + } +} + +func TestNewWebFetchToolWithProxy(t *testing.T) { + tool, err := NewWebFetchToolWithProxy(1024, "http://127.0.0.1:7890", testFetchLimit) + if err != nil { + logger.ErrorCF("agent", "Failed to create web fetch tool", map[string]any{"error": err.Error()}) + } else if tool.maxChars != 1024 { + t.Fatalf("maxChars = %d, want %d", tool.maxChars, 1024) + } + + if tool.proxy != "http://127.0.0.1:7890" { + t.Fatalf("proxy = %q, want %q", tool.proxy, "http://127.0.0.1:7890") + } + + tool, err = NewWebFetchToolWithProxy(0, "http://127.0.0.1:7890", testFetchLimit) + if err != nil { + logger.ErrorCF("agent", "Failed to create web fetch tool", map[string]any{"error": err.Error()}) + } + + if tool.maxChars != 50000 { + t.Fatalf("default maxChars = %d, want %d", tool.maxChars, 50000) + } +} diff --git a/pkg/tools/web_keys.go b/pkg/tools/web_keys.go new file mode 100644 index 000000000..248333d20 --- /dev/null +++ b/pkg/tools/web_keys.go @@ -0,0 +1,43 @@ +package tools + +import ( + "sync/atomic" +) + +type APIKeyPool struct { + keys []string + current uint32 +} + +func NewAPIKeyPool(keys []string) *APIKeyPool { + return &APIKeyPool{ + keys: keys, + } +} + +type APIKeyIterator struct { + pool *APIKeyPool + startIdx uint32 + attempt uint32 +} + +func (p *APIKeyPool) NewIterator() *APIKeyIterator { + if len(p.keys) == 0 { + return &APIKeyIterator{pool: p} + } + idx := atomic.AddUint32(&p.current, 1) - 1 + return &APIKeyIterator{ + pool: p, + startIdx: idx, + } +} + +func (it *APIKeyIterator) Next() (string, bool) { + length := uint32(len(it.pool.keys)) + if length == 0 || it.attempt >= length { + return "", false + } + key := it.pool.keys[(it.startIdx+it.attempt)%length] + it.attempt++ + return key, true +} diff --git a/pkg/tools/web_keys_test.go b/pkg/tools/web_keys_test.go new file mode 100644 index 000000000..2609b8fa4 --- /dev/null +++ b/pkg/tools/web_keys_test.go @@ -0,0 +1,59 @@ +package tools + +import ( + "testing" +) + +func TestAPIKeyPool(t *testing.T) { + pool := NewAPIKeyPool([]string{"key1", "key2", "key3"}) + if len(pool.keys) != 3 { + t.Fatalf("expected 3 keys, got %d", len(pool.keys)) + } + if pool.keys[0] != "key1" || pool.keys[1] != "key2" || pool.keys[2] != "key3" { + t.Fatalf("unexpected keys: %v", pool.keys) + } + + // Test Iterator: each iterator should cover all keys exactly once + iter := pool.NewIterator() + expected := []string{"key1", "key2", "key3"} + for i, want := range expected { + k, ok := iter.Next() + if !ok { + t.Fatalf("iter.Next() returned false at step %d", i) + } + if k != want { + t.Errorf("step %d: expected %s, got %s", i, want, k) + } + } + // Should be exhausted + if _, ok := iter.Next(); ok { + t.Errorf("expected iterator exhausted after all keys") + } + + // Second iterator starts at next position (load balancing) + iter2 := pool.NewIterator() + k, ok := iter2.Next() + if !ok { + t.Fatal("iter2.Next() returned false") + } + if k != "key2" { + t.Errorf("expected key2 (round-robin), got %s", k) + } + + // Empty pool + emptyPool := NewAPIKeyPool([]string{}) + emptyIter := emptyPool.NewIterator() + if _, ok := emptyIter.Next(); ok { + t.Errorf("expected false for empty pool") + } + + // Single key pool + singlePool := NewAPIKeyPool([]string{"single"}) + singleIter := singlePool.NewIterator() + if k, ok := singleIter.Next(); !ok || k != "single" { + t.Errorf("expected single, got %s (ok=%v)", k, ok) + } + if _, ok := singleIter.Next(); ok { + t.Errorf("expected exhausted after single key") + } +} diff --git a/pkg/tools/web_search_providers.go b/pkg/tools/web_search_providers.go new file mode 100644 index 000000000..02c64ea32 --- /dev/null +++ b/pkg/tools/web_search_providers.go @@ -0,0 +1,519 @@ +package tools + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "net/url" + "strings" + "time" +) + +type SearchProvider interface { + Search(ctx context.Context, query string, count int) (string, error) +} + +type BraveSearchProvider struct { + keyPool *APIKeyPool + proxy string + client *http.Client +} + +func (p *BraveSearchProvider) Search(ctx context.Context, query string, count int) (string, error) { + searchURL := fmt.Sprintf("https://api.search.brave.com/res/v1/web/search?q=%s&count=%d", + url.QueryEscape(query), count) + + var lastErr error + iter := p.keyPool.NewIterator() + + for { + apiKey, ok := iter.Next() + if !ok { + break + } + + req, err := http.NewRequestWithContext(ctx, "GET", searchURL, nil) + if err != nil { + return "", fmt.Errorf("failed to create request: %w", err) + } + + req.Header.Set("Accept", "application/json") + req.Header.Set("X-Subscription-Token", apiKey) + + resp, err := p.client.Do(req) + if err != nil { + lastErr = fmt.Errorf("request failed: %w", err) + continue + } + + body, err := io.ReadAll(resp.Body) + resp.Body.Close() + + if err != nil { + lastErr = fmt.Errorf("failed to read response: %w", err) + continue + } + + if resp.StatusCode != http.StatusOK { + lastErr = fmt.Errorf("API error (status %d): %s", resp.StatusCode, string(body)) + if resp.StatusCode == http.StatusTooManyRequests || + resp.StatusCode == http.StatusUnauthorized || + resp.StatusCode == http.StatusForbidden || + resp.StatusCode >= 500 { + continue + } + return "", lastErr + } + + var searchResp struct { + Web struct { + Results []struct { + Title string `json:"title"` + URL string `json:"url"` + Description string `json:"description"` + } `json:"results"` + } `json:"web"` + } + + if err := json.Unmarshal(body, &searchResp); err != nil { + // Log error body for debugging + return "", fmt.Errorf("failed to parse response: %w", err) + } + + results := searchResp.Web.Results + if len(results) == 0 { + return fmt.Sprintf("No results for: %s", query), nil + } + + var lines []string + lines = append(lines, fmt.Sprintf("Results for: %s", query)) + for i, item := range results { + if i >= count { + break + } + lines = append(lines, fmt.Sprintf("%d. %s\n %s", i+1, item.Title, item.URL)) + if item.Description != "" { + lines = append(lines, fmt.Sprintf(" %s", item.Description)) + } + } + + return strings.Join(lines, "\n"), nil + } + + return "", fmt.Errorf("all api keys failed, last error: %w", lastErr) +} + +type TavilySearchProvider struct { + keyPool *APIKeyPool + baseURL string + proxy string + client *http.Client +} + +func (p *TavilySearchProvider) Search(ctx context.Context, query string, count int) (string, error) { + searchURL := p.baseURL + if searchURL == "" { + searchURL = "https://api.tavily.com/search" + } + + var lastErr error + iter := p.keyPool.NewIterator() + + for { + apiKey, ok := iter.Next() + if !ok { + break + } + + payload := map[string]any{ + "api_key": apiKey, + "query": query, + "search_depth": "advanced", + "include_answer": false, + "include_images": false, + "include_raw_content": false, + "max_results": count, + } + + bodyBytes, err := json.Marshal(payload) + if err != nil { + return "", fmt.Errorf("failed to marshal payload: %w", err) + } + + req, err := http.NewRequestWithContext(ctx, "POST", searchURL, bytes.NewBuffer(bodyBytes)) + if err != nil { + return "", fmt.Errorf("failed to create request: %w", err) + } + + req.Header.Set("Content-Type", "application/json") + req.Header.Set("User-Agent", userAgent) + + resp, err := p.client.Do(req) + if err != nil { + lastErr = fmt.Errorf("request failed: %w", err) + continue + } + + body, err := io.ReadAll(resp.Body) + resp.Body.Close() + + if err != nil { + lastErr = fmt.Errorf("failed to read response: %w", err) + continue + } + + if resp.StatusCode != http.StatusOK { + lastErr = fmt.Errorf("tavily api error (status %d): %s", resp.StatusCode, string(body)) + if resp.StatusCode == http.StatusTooManyRequests || + resp.StatusCode == http.StatusUnauthorized || + resp.StatusCode == http.StatusForbidden || + resp.StatusCode >= 500 { + continue + } + return "", lastErr + } + + var searchResp struct { + Results []struct { + Title string `json:"title"` + URL string `json:"url"` + Content string `json:"content"` + } `json:"results"` + } + + if err := json.Unmarshal(body, &searchResp); err != nil { + return "", fmt.Errorf("failed to parse response: %w", err) + } + + results := searchResp.Results + if len(results) == 0 { + return fmt.Sprintf("No results for: %s", query), nil + } + + var lines []string + lines = append(lines, fmt.Sprintf("Results for: %s (via Tavily)", query)) + for i, item := range results { + if i >= count { + break + } + lines = append(lines, fmt.Sprintf("%d. %s\n %s", i+1, item.Title, item.URL)) + if item.Content != "" { + lines = append(lines, fmt.Sprintf(" %s", item.Content)) + } + } + + return strings.Join(lines, "\n"), nil + } + + return "", fmt.Errorf("all api keys failed, last error: %w", lastErr) +} + +type DuckDuckGoSearchProvider struct { + proxy string + client *http.Client +} + +func (p *DuckDuckGoSearchProvider) Search(ctx context.Context, query string, count int) (string, error) { + searchURL := fmt.Sprintf("https://html.duckduckgo.com/html/?q=%s", url.QueryEscape(query)) + + req, err := http.NewRequestWithContext(ctx, "GET", searchURL, nil) + if err != nil { + return "", fmt.Errorf("failed to create request: %w", err) + } + + req.Header.Set("User-Agent", userAgent) + + resp, err := p.client.Do(req) + if err != nil { + return "", fmt.Errorf("request failed: %w", err) + } + defer resp.Body.Close() + + body, err := io.ReadAll(resp.Body) + if err != nil { + return "", fmt.Errorf("failed to read response: %w", err) + } + + return p.extractResults(string(body), count, query) +} + +func (p *DuckDuckGoSearchProvider) extractResults(html string, count int, query string) (string, error) { + matches := reDDGLink.FindAllStringSubmatch(html, count+5) + + if len(matches) == 0 { + return fmt.Sprintf("No results found or extraction failed. Query: %s", query), nil + } + + var lines []string + lines = append(lines, fmt.Sprintf("Results for: %s (via DuckDuckGo)", query)) + + snippetMatches := reDDGSnippet.FindAllStringSubmatch(html, count+5) + + maxItems := min(len(matches), count) + + for i := range maxItems { + urlStr := matches[i][1] + title := stripTags(matches[i][2]) + title = strings.TrimSpace(title) + + if strings.Contains(urlStr, "uddg=") { + if u, err := url.QueryUnescape(urlStr); err == nil { + _, after, ok := strings.Cut(u, "uddg=") + if ok { + urlStr = after + } + } + } + + lines = append(lines, fmt.Sprintf("%d. %s\n %s", i+1, title, urlStr)) + + if i < len(snippetMatches) { + snippet := stripTags(snippetMatches[i][1]) + snippet = strings.TrimSpace(snippet) + if snippet != "" { + lines = append(lines, fmt.Sprintf(" %s", snippet)) + } + } + } + + return strings.Join(lines, "\n"), nil +} + +type PerplexitySearchProvider struct { + keyPool *APIKeyPool + proxy string + client *http.Client +} + +func (p *PerplexitySearchProvider) Search(ctx context.Context, query string, count int) (string, error) { + searchURL := "https://api.perplexity.ai/chat/completions" + + var lastErr error + iter := p.keyPool.NewIterator() + + for { + apiKey, ok := iter.Next() + if !ok { + break + } + + payload := map[string]any{ + "model": "sonar", + "messages": []map[string]string{ + { + "role": "system", + "content": "You are a search assistant. Provide concise search results with titles, URLs, and brief descriptions in the following format:\n1. Title\n URL\n Description\n\nDo not add extra commentary.", + }, + { + "role": "user", + "content": fmt.Sprintf("Search for: %s. Provide up to %d relevant results.", query, count), + }, + }, + "max_tokens": 1000, + } + + payloadBytes, err := json.Marshal(payload) + if err != nil { + return "", fmt.Errorf("failed to marshal request: %w", err) + } + + req, err := http.NewRequestWithContext(ctx, "POST", searchURL, strings.NewReader(string(payloadBytes))) + if err != nil { + return "", fmt.Errorf("failed to create request: %w", err) + } + + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+apiKey) + req.Header.Set("User-Agent", userAgent) + + resp, err := p.client.Do(req) + if err != nil { + lastErr = fmt.Errorf("request failed: %w", err) + continue + } + + body, err := io.ReadAll(resp.Body) + resp.Body.Close() + + if err != nil { + lastErr = fmt.Errorf("failed to read response: %w", err) + continue + } + + if resp.StatusCode != http.StatusOK { + lastErr = fmt.Errorf("Perplexity API error: %s", string(body)) + if resp.StatusCode == http.StatusTooManyRequests || + resp.StatusCode == http.StatusUnauthorized || + resp.StatusCode == http.StatusForbidden || + resp.StatusCode >= 500 { + continue + } + return "", lastErr + } + + var searchResp struct { + Choices []struct { + Message struct { + Content string `json:"content"` + } `json:"message"` + } `json:"choices"` + } + + if err := json.Unmarshal(body, &searchResp); err != nil { + return "", fmt.Errorf("failed to parse response: %w", err) + } + + if len(searchResp.Choices) == 0 { + return fmt.Sprintf("No results for: %s", query), nil + } + + return fmt.Sprintf("Results for: %s (via Perplexity)\n%s", query, searchResp.Choices[0].Message.Content), nil + } + + return "", fmt.Errorf("all api keys failed, last error: %w", lastErr) +} + +type SearXNGSearchProvider struct { + baseURL string +} + +func (p *SearXNGSearchProvider) Search(ctx context.Context, query string, count int) (string, error) { + searchURL := fmt.Sprintf("%s/search?q=%s&format=json&categories=general", + strings.TrimSuffix(p.baseURL, "/"), + url.QueryEscape(query)) + + req, err := http.NewRequestWithContext(ctx, "GET", searchURL, nil) + if err != nil { + return "", fmt.Errorf("failed to create request: %w", err) + } + + client := &http.Client{Timeout: 10 * time.Second} + resp, err := client.Do(req) + if err != nil { + return "", fmt.Errorf("request failed: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + return "", fmt.Errorf("SearXNG returned status %d", resp.StatusCode) + } + + var result struct { + Results []struct { + Title string `json:"title"` + URL string `json:"url"` + Content string `json:"content"` + Engine string `json:"engine"` + Score float64 `json:"score"` + } `json:"results"` + } + + if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { + return "", fmt.Errorf("failed to parse response: %w", err) + } + + if len(result.Results) == 0 { + return fmt.Sprintf("No results for: %s", query), nil + } + + if len(result.Results) > count { + result.Results = result.Results[:count] + } + + var b strings.Builder + b.WriteString(fmt.Sprintf("Results for: %s (via SearXNG)\n", query)) + for i, r := range result.Results { + b.WriteString(fmt.Sprintf("%d. %s\n", i+1, r.Title)) + b.WriteString(fmt.Sprintf(" %s\n", r.URL)) + if r.Content != "" { + b.WriteString(fmt.Sprintf(" %s\n", r.Content)) + } + } + + return b.String(), nil +} + +type GLMSearchProvider struct { + apiKey string + baseURL string + searchEngine string + proxy string + client *http.Client +} + +func (p *GLMSearchProvider) Search(ctx context.Context, query string, count int) (string, error) { + searchURL := p.baseURL + if searchURL == "" { + searchURL = "https://open.bigmodel.cn/api/paas/v4/web_search" + } + + payload := map[string]any{ + "search_query": query, + "search_engine": p.searchEngine, + "search_intent": false, + "count": count, + "content_size": "medium", + } + + bodyBytes, err := json.Marshal(payload) + if err != nil { + return "", fmt.Errorf("failed to marshal payload: %w", err) + } + + req, err := http.NewRequestWithContext(ctx, "POST", searchURL, bytes.NewReader(bodyBytes)) + if err != nil { + return "", fmt.Errorf("failed to create request: %w", err) + } + + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+p.apiKey) + + resp, err := p.client.Do(req) + if err != nil { + return "", fmt.Errorf("request failed: %w", err) + } + defer resp.Body.Close() + + body, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20)) + if err != nil { + return "", fmt.Errorf("failed to read response: %w", err) + } + + if resp.StatusCode != http.StatusOK { + return "", fmt.Errorf("GLM Search API error (status %d): %s", resp.StatusCode, string(body)) + } + + var searchResp struct { + SearchResult []struct { + Title string `json:"title"` + Content string `json:"content"` + Link string `json:"link"` + } `json:"search_result"` + } + + if err := json.Unmarshal(body, &searchResp); err != nil { + return "", fmt.Errorf("failed to parse response: %w", err) + } + + results := searchResp.SearchResult + if len(results) == 0 { + return fmt.Sprintf("No results for: %s", query), nil + } + + var lines []string + lines = append(lines, fmt.Sprintf("Results for: %s (via GLM Search)", query)) + for i, item := range results { + if i >= count { + break + } + lines = append(lines, fmt.Sprintf("%d. %s\n %s", i+1, item.Title, item.Link)) + if item.Content != "" { + lines = append(lines, fmt.Sprintf(" %s", item.Content)) + } + } + + return strings.Join(lines, "\n"), nil +} diff --git a/pkg/tools/web_search_tool.go b/pkg/tools/web_search_tool.go new file mode 100644 index 000000000..53e29da72 --- /dev/null +++ b/pkg/tools/web_search_tool.go @@ -0,0 +1,169 @@ +package tools + +import ( + "context" + "fmt" +) + +type WebSearchTool struct { + provider SearchProvider + maxResults int +} + +type WebSearchToolOptions struct { + BraveAPIKeys []string + BraveMaxResults int + BraveEnabled bool + TavilyAPIKeys []string + TavilyBaseURL string + TavilyMaxResults int + TavilyEnabled bool + DuckDuckGoMaxResults int + DuckDuckGoEnabled bool + PerplexityAPIKeys []string + PerplexityMaxResults int + PerplexityEnabled bool + SearXNGBaseURL string + SearXNGMaxResults int + SearXNGEnabled bool + GLMSearchAPIKey string + GLMSearchBaseURL string + GLMSearchEngine string + GLMSearchMaxResults int + GLMSearchEnabled bool + Proxy string +} + +func NewWebSearchTool(opts WebSearchToolOptions) (*WebSearchTool, error) { + var provider SearchProvider + maxResults := 5 + // Priority: Perplexity > Brave > SearXNG > Tavily > DuckDuckGo > GLM Search + if opts.PerplexityEnabled && len(opts.PerplexityAPIKeys) > 0 { + client, err := createHTTPClient(opts.Proxy, perplexityTimeout) + if err != nil { + return nil, fmt.Errorf("failed to create HTTP client for Perplexity: %w", err) + } + provider = &PerplexitySearchProvider{ + keyPool: NewAPIKeyPool(opts.PerplexityAPIKeys), + proxy: opts.Proxy, + client: client, + } + if opts.PerplexityMaxResults > 0 { + maxResults = opts.PerplexityMaxResults + } + } else if opts.BraveEnabled && len(opts.BraveAPIKeys) > 0 { + client, err := createHTTPClient(opts.Proxy, searchTimeout) + if err != nil { + return nil, fmt.Errorf("failed to create HTTP client for Brave: %w", err) + } + provider = &BraveSearchProvider{keyPool: NewAPIKeyPool(opts.BraveAPIKeys), proxy: opts.Proxy, client: client} + if opts.BraveMaxResults > 0 { + maxResults = opts.BraveMaxResults + } + } else if opts.SearXNGEnabled && opts.SearXNGBaseURL != "" { + provider = &SearXNGSearchProvider{baseURL: opts.SearXNGBaseURL} + if opts.SearXNGMaxResults > 0 { + maxResults = opts.SearXNGMaxResults + } + } else if opts.TavilyEnabled && len(opts.TavilyAPIKeys) > 0 { + client, err := createHTTPClient(opts.Proxy, searchTimeout) + if err != nil { + return nil, fmt.Errorf("failed to create HTTP client for Tavily: %w", err) + } + provider = &TavilySearchProvider{ + keyPool: NewAPIKeyPool(opts.TavilyAPIKeys), + baseURL: opts.TavilyBaseURL, + proxy: opts.Proxy, + client: client, + } + if opts.TavilyMaxResults > 0 { + maxResults = opts.TavilyMaxResults + } + } else if opts.DuckDuckGoEnabled { + 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 if opts.GLMSearchEnabled && opts.GLMSearchAPIKey != "" { + client, err := createHTTPClient(opts.Proxy, searchTimeout) + if err != nil { + return nil, fmt.Errorf("failed to create HTTP client for GLM Search: %w", err) + } + searchEngine := opts.GLMSearchEngine + if searchEngine == "" { + searchEngine = "search_std" + } + provider = &GLMSearchProvider{ + apiKey: opts.GLMSearchAPIKey, + baseURL: opts.GLMSearchBaseURL, + searchEngine: searchEngine, + proxy: opts.Proxy, + client: client, + } + if opts.GLMSearchMaxResults > 0 { + maxResults = opts.GLMSearchMaxResults + } + } else { + return nil, nil + } + + return &WebSearchTool{ + provider: provider, + maxResults: maxResults, + }, nil +} + +func (t *WebSearchTool) Name() string { + return "web_search" +} + +func (t *WebSearchTool) Description() string { + return "Search the web for current information. Returns titles, URLs, and snippets from search results." +} + +func (t *WebSearchTool) Parameters() map[string]any { + return map[string]any{ + "type": "object", + "properties": map[string]any{ + "query": map[string]any{ + "type": "string", + "description": "Search query", + }, + "count": map[string]any{ + "type": "integer", + "description": "Number of results (1-10)", + "minimum": 1.0, + "maximum": 10.0, + }, + }, + "required": []string{"query"}, + } +} + +func (t *WebSearchTool) Execute(ctx context.Context, args map[string]any) *ToolResult { + query, ok := args["query"].(string) + if !ok { + return ErrorResult("query is required") + } + + count := t.maxResults + if c, ok := args["count"].(float64); ok { + if int(c) > 0 && int(c) <= 10 { + count = int(c) + } + } + + result, err := t.provider.Search(ctx, query, count) + if err != nil { + return ErrorResult(fmt.Sprintf("search failed: %v", err)) + } + + return &ToolResult{ + ForLLM: result, + ForUser: result, + } +} diff --git a/pkg/tools/web_search_tool_test.go b/pkg/tools/web_search_tool_test.go new file mode 100644 index 000000000..b622121d9 --- /dev/null +++ b/pkg/tools/web_search_tool_test.go @@ -0,0 +1,368 @@ +package tools + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" +) + +func TestWebTool_WebSearch_NoApiKey(t *testing.T) { + tool, err := NewWebSearchTool(WebSearchToolOptions{BraveEnabled: true, BraveAPIKeys: nil}) + 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, 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") + } +} + +func TestWebTool_WebSearch_MissingQuery(t *testing.T) { + tool, err := NewWebSearchTool(WebSearchToolOptions{ + BraveEnabled: true, + BraveAPIKeys: []string{"test-key"}, + BraveMaxResults: 5, + }) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + ctx := context.Background() + args := map[string]any{} + + result := tool.Execute(ctx, args) + + // Should return error result + if !result.IsError { + t.Errorf("Expected error when query is missing") + } +} + +func TestNewWebSearchTool_PropagatesProxy(t *testing.T) { + t.Run("perplexity", func(t *testing.T) { + tool, err := NewWebSearchTool(WebSearchToolOptions{ + PerplexityEnabled: true, + PerplexityAPIKeys: []string{"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) + } + if p.proxy != "http://127.0.0.1:7890" { + t.Fatalf("provider proxy = %q, want %q", p.proxy, "http://127.0.0.1:7890") + } + }) + + t.Run("brave", func(t *testing.T) { + tool, err := NewWebSearchTool(WebSearchToolOptions{ + BraveEnabled: true, + BraveAPIKeys: []string{"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) + } + if p.proxy != "http://127.0.0.1:7890" { + t.Fatalf("provider proxy = %q, want %q", p.proxy, "http://127.0.0.1:7890") + } + }) + + t.Run("duckduckgo", func(t *testing.T) { + 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) + } + if p.proxy != "http://127.0.0.1:7890" { + t.Fatalf("provider proxy = %q, want %q", p.proxy, "http://127.0.0.1:7890") + } + }) +} + +func TestWebTool_TavilySearch_Success(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != "POST" { + t.Errorf("Expected POST request, got %s", r.Method) + } + if r.Header.Get("Content-Type") != "application/json" { + t.Errorf("Expected Content-Type application/json, got %s", r.Header.Get("Content-Type")) + } + + var payload map[string]any + json.NewDecoder(r.Body).Decode(&payload) + if payload["api_key"] != "test-key" { + t.Errorf("Expected api_key test-key, got %v", payload["api_key"]) + } + if payload["query"] != "test query" { + t.Errorf("Expected query 'test query', got %v", payload["query"]) + } + + response := map[string]any{ + "results": []map[string]any{ + { + "title": "Test Result 1", + "url": "https://example.com/1", + "content": "Content for result 1", + }, + { + "title": "Test Result 2", + "url": "https://example.com/2", + "content": "Content for result 2", + }, + }, + } + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + json.NewEncoder(w).Encode(response) + })) + defer server.Close() + + tool, err := NewWebSearchTool(WebSearchToolOptions{ + TavilyEnabled: true, + TavilyAPIKeys: []string{"test-key"}, + TavilyBaseURL: server.URL, + TavilyMaxResults: 5, + }) + if err != nil { + t.Fatalf("NewWebSearchTool() error: %v", err) + } + + ctx := context.Background() + args := map[string]any{ + "query": "test query", + } + + result := tool.Execute(ctx, args) + + if result.IsError { + t.Errorf("Expected success, got IsError=true: %s", result.ForLLM) + } + + if !strings.Contains(result.ForUser, "Test Result 1") || + !strings.Contains(result.ForUser, "https://example.com/1") { + t.Errorf("Expected results in output, got: %s", result.ForUser) + } + + if !strings.Contains(result.ForUser, "via Tavily") { + t.Errorf("Expected 'via Tavily' in output, got: %s", result.ForUser) + } +} + +func TestWebTool_TavilySearch_Failover(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var payload map[string]any + if err := json.NewDecoder(r.Body).Decode(&payload); err != nil { + t.Fatalf("failed to decode payload: %v", err) + } + + apiKey := payload["api_key"].(string) + + if apiKey == "key1" { + w.WriteHeader(http.StatusTooManyRequests) + w.Write([]byte("Rate limited")) + return + } + + if apiKey == "key2" { + response := map[string]any{ + "results": []map[string]any{ + { + "title": "Success Result", + "url": "https://example.com/success", + "content": "Success content", + }, + }, + } + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + json.NewEncoder(w).Encode(response) + return + } + + w.WriteHeader(http.StatusBadRequest) + })) + defer server.Close() + + tool, err := NewWebSearchTool(WebSearchToolOptions{ + TavilyEnabled: true, + TavilyAPIKeys: []string{"key1", "key2"}, + TavilyBaseURL: server.URL, + TavilyMaxResults: 5, + }) + if err != nil { + t.Fatalf("NewWebSearchTool() error: %v", err) + } + + ctx := context.Background() + args := map[string]any{ + "query": "test query", + } + + result := tool.Execute(ctx, args) + + if result.IsError { + t.Errorf("Expected success, got Error: %s", result.ForLLM) + } + if !strings.Contains(result.ForUser, "Success Result") { + t.Errorf("Expected failover to second key and success result, got: %s", result.ForUser) + } +} + +func TestWebTool_GLMSearch_Success(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != "POST" { + t.Errorf("Expected POST request, got %s", r.Method) + } + if r.Header.Get("Content-Type") != "application/json" { + t.Errorf("Expected Content-Type application/json, got %s", r.Header.Get("Content-Type")) + } + if r.Header.Get("Authorization") != "Bearer test-glm-key" { + t.Errorf("Expected Authorization Bearer test-glm-key, got %s", r.Header.Get("Authorization")) + } + + var payload map[string]any + json.NewDecoder(r.Body).Decode(&payload) + if payload["search_query"] != "test query" { + t.Errorf("Expected search_query 'test query', got %v", payload["search_query"]) + } + if payload["search_engine"] != "search_std" { + t.Errorf("Expected search_engine 'search_std', got %v", payload["search_engine"]) + } + + response := map[string]any{ + "id": "web-search-test", + "created": 1709568000, + "search_result": []map[string]any{ + { + "title": "Test GLM Result", + "content": "GLM search snippet", + "link": "https://example.com/glm", + "media": "Example", + "publish_date": "2026-03-04", + }, + }, + } + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + json.NewEncoder(w).Encode(response) + })) + defer server.Close() + + tool, err := NewWebSearchTool(WebSearchToolOptions{ + GLMSearchEnabled: true, + GLMSearchAPIKey: "test-glm-key", + GLMSearchBaseURL: server.URL, + GLMSearchEngine: "search_std", + }) + if err != nil { + t.Fatalf("NewWebSearchTool() error: %v", err) + } + + result := tool.Execute(context.Background(), map[string]any{ + "query": "test query", + }) + + if result.IsError { + t.Errorf("Expected success, got IsError=true: %s", result.ForLLM) + } + if !strings.Contains(result.ForUser, "Test GLM Result") { + t.Errorf("Expected 'Test GLM Result' in output, got: %s", result.ForUser) + } + if !strings.Contains(result.ForUser, "https://example.com/glm") { + t.Errorf("Expected URL in output, got: %s", result.ForUser) + } + if !strings.Contains(result.ForUser, "via GLM Search") { + t.Errorf("Expected 'via GLM Search' in output, got: %s", result.ForUser) + } +} + +func TestWebTool_GLMSearch_APIError(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusUnauthorized) + w.Write([]byte(`{"error":"invalid api key"}`)) + })) + defer server.Close() + + tool, err := NewWebSearchTool(WebSearchToolOptions{ + GLMSearchEnabled: true, + GLMSearchAPIKey: "bad-key", + GLMSearchBaseURL: server.URL, + GLMSearchEngine: "search_std", + }) + if err != nil { + t.Fatalf("NewWebSearchTool() error: %v", err) + } + + result := tool.Execute(context.Background(), map[string]any{ + "query": "test query", + }) + + if !result.IsError { + t.Errorf("Expected IsError=true for 401 response") + } + if !strings.Contains(result.ForLLM, "status 401") { + t.Errorf("Expected status 401 in error, got: %s", result.ForLLM) + } +} + +func TestWebTool_GLMSearch_Priority(t *testing.T) { + // GLM Search should only be selected when all other providers are disabled + tool, err := NewWebSearchTool(WebSearchToolOptions{ + DuckDuckGoEnabled: true, + DuckDuckGoMaxResults: 5, + GLMSearchEnabled: true, + GLMSearchAPIKey: "test-key", + GLMSearchBaseURL: "https://example.com", + GLMSearchEngine: "search_std", + }) + if err != nil { + t.Fatalf("NewWebSearchTool() error: %v", err) + } + + // DuckDuckGo should win over GLM Search + if _, ok := tool.provider.(*DuckDuckGoSearchProvider); !ok { + t.Errorf("Expected DuckDuckGoSearchProvider when both enabled, got %T", tool.provider) + } + + // With DuckDuckGo disabled, GLM Search should be selected + tool2, err := NewWebSearchTool(WebSearchToolOptions{ + DuckDuckGoEnabled: false, + GLMSearchEnabled: true, + GLMSearchAPIKey: "test-key", + GLMSearchBaseURL: "https://example.com", + GLMSearchEngine: "search_std", + }) + if err != nil { + t.Fatalf("NewWebSearchTool() error: %v", err) + } + if _, ok := tool2.provider.(*GLMSearchProvider); !ok { + t.Errorf("Expected GLMSearchProvider when only GLM enabled, got %T", tool2.provider) + } +} diff --git a/pkg/tools/web_ssrf.go b/pkg/tools/web_ssrf.go new file mode 100644 index 000000000..c0fbb079b --- /dev/null +++ b/pkg/tools/web_ssrf.go @@ -0,0 +1,137 @@ +package tools + +import ( + "context" + "fmt" + "net" + "strings" + "sync/atomic" +) + +// allowPrivateWebFetchHosts controls whether loopback/private hosts are allowed. +// This is false in normal runtime to reduce SSRF exposure, and tests can override it temporarily. +var allowPrivateWebFetchHosts atomic.Bool + +// newSafeDialContext re-resolves DNS at connect time to mitigate DNS rebinding (TOCTOU) +// where a hostname resolves to a public IP during pre-flight but a private IP at connect time. +func newSafeDialContext(dialer *net.Dialer) func(context.Context, string, string) (net.Conn, error) { + return func(ctx context.Context, network, address string) (net.Conn, error) { + if allowPrivateWebFetchHosts.Load() { + return dialer.DialContext(ctx, network, address) + } + + host, port, err := net.SplitHostPort(address) + if err != nil { + return nil, fmt.Errorf("invalid target address %q: %w", address, err) + } + if host == "" { + return nil, fmt.Errorf("empty target host") + } + + if ip := net.ParseIP(host); ip != nil { + if isPrivateOrRestrictedIP(ip) { + return nil, fmt.Errorf("blocked private or local target: %s", host) + } + return dialer.DialContext(ctx, network, net.JoinHostPort(ip.String(), port)) + } + + ipAddrs, err := net.DefaultResolver.LookupIPAddr(ctx, host) + if err != nil { + return nil, fmt.Errorf("failed to resolve %s: %w", host, err) + } + + attempted := 0 + var lastErr error + for _, ipAddr := range ipAddrs { + if isPrivateOrRestrictedIP(ipAddr.IP) { + continue + } + attempted++ + conn, err := dialer.DialContext(ctx, network, net.JoinHostPort(ipAddr.IP.String(), port)) + if err == nil { + return conn, nil + } + lastErr = err + } + + if attempted == 0 { + return nil, fmt.Errorf("all resolved addresses for %s are private or restricted", host) + } + if lastErr != nil { + return nil, fmt.Errorf("failed connecting to public addresses for %s: %w", host, lastErr) + } + return nil, fmt.Errorf("failed connecting to public addresses for %s", host) + } +} + +// isObviousPrivateHost performs a lightweight, no-DNS check for obviously private hosts. +// It catches localhost, literal private IPs, and empty hosts. It does NOT resolve DNS — +// the real SSRF guard is newSafeDialContext which checks IPs at connect time. +func isObviousPrivateHost(host string) bool { + if allowPrivateWebFetchHosts.Load() { + return false + } + + h := strings.ToLower(strings.TrimSpace(host)) + h = strings.TrimSuffix(h, ".") + if h == "" { + return true + } + + if h == "localhost" || strings.HasSuffix(h, ".localhost") { + return true + } + + if ip := net.ParseIP(h); ip != nil { + return isPrivateOrRestrictedIP(ip) + } + + return false +} + +// isPrivateOrRestrictedIP returns true for IPs that should never be reached via web_fetch: +// RFC 1918, loopback, link-local (incl. cloud metadata 169.254.x.x), carrier-grade NAT, +// IPv6 unique-local (fc00::/7), 6to4 (2002::/16), and Teredo (2001:0000::/32). +func isPrivateOrRestrictedIP(ip net.IP) bool { + if ip == nil { + return true + } + + if ip.IsLoopback() || ip.IsLinkLocalUnicast() || ip.IsLinkLocalMulticast() || + ip.IsMulticast() || ip.IsUnspecified() { + return true + } + + if ip4 := ip.To4(); ip4 != nil { + // IPv4 private, loopback, link-local, and carrier-grade NAT ranges. + if ip4[0] == 10 || + ip4[0] == 127 || + ip4[0] == 0 || + (ip4[0] == 172 && ip4[1] >= 16 && ip4[1] <= 31) || + (ip4[0] == 192 && ip4[1] == 168) || + (ip4[0] == 169 && ip4[1] == 254) || + (ip4[0] == 100 && ip4[1] >= 64 && ip4[1] <= 127) { + return true + } + return false + } + + if len(ip) == net.IPv6len { + // IPv6 unique local addresses (fc00::/7) + if (ip[0] & 0xfe) == 0xfc { + return true + } + // 6to4 addresses (2002::/16): check the embedded IPv4 at bytes [2:6]. + if ip[0] == 0x20 && ip[1] == 0x02 { + embedded := net.IPv4(ip[2], ip[3], ip[4], ip[5]) + return isPrivateOrRestrictedIP(embedded) + } + // Teredo (2001:0000::/32): client IPv4 is at bytes [12:16], XOR-inverted. + if ip[0] == 0x20 && ip[1] == 0x01 && ip[2] == 0x00 && ip[3] == 0x00 { + client := net.IPv4(ip[12]^0xff, ip[13]^0xff, ip[14]^0xff, ip[15]^0xff) + return isPrivateOrRestrictedIP(client) + } + } + + return false +} diff --git a/pkg/tools/web_ssrf_test.go b/pkg/tools/web_ssrf_test.go new file mode 100644 index 000000000..12ec84af8 --- /dev/null +++ b/pkg/tools/web_ssrf_test.go @@ -0,0 +1,206 @@ +package tools + +import ( + "context" + "net" + "net/http" + "net/http/httptest" + "os" + "strings" + "testing" +) + +func withPrivateWebFetchHostsAllowed(t *testing.T) { + t.Helper() + previous := allowPrivateWebFetchHosts.Load() + allowPrivateWebFetchHosts.Store(true) + t.Cleanup(func() { + allowPrivateWebFetchHosts.Store(previous) + }) +} + +func TestWebTool_WebFetch_PrivateHostBlocked(t *testing.T) { + tool, err := NewWebFetchTool(50000, testFetchLimit) + if err != nil { + t.Fatalf("Failed to create web fetch tool: %v", err) + } + result := tool.Execute(context.Background(), map[string]any{ + "url": "http://127.0.0.1:0", + }) + + if !result.IsError { + t.Errorf("expected error for private host URL, got success") + } + if !strings.Contains(result.ForLLM, "private or local network") && + !strings.Contains(result.ForUser, "private or local network") { + t.Errorf("expected private host block message, got %q", result.ForLLM) + } +} + +func TestWebTool_WebFetch_PrivateHostAllowedForTests(t *testing.T) { + withPrivateWebFetchHostsAllowed(t) + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/plain") + w.WriteHeader(http.StatusOK) + w.Write([]byte("ok")) + })) + defer server.Close() + + tool, err := NewWebFetchTool(50000, testFetchLimit) + if err != nil { + t.Fatalf("Failed to create web fetch tool: %v", err) + } + result := tool.Execute(context.Background(), map[string]any{ + "url": server.URL, + }) + + if result.IsError { + t.Errorf("expected success when private host access is allowed in tests, got %q", result.ForLLM) + } +} + +func TestWebFetch_BlocksIPv4MappedIPv6Loopback(t *testing.T) { + tool, err := NewWebFetchTool(50000, testFetchLimit) + if err != nil { + t.Fatalf("Failed to create web fetch tool: %v", err) + } + result := tool.Execute(context.Background(), map[string]any{ + "url": "http://[::ffff:127.0.0.1]:0", + }) + + if !result.IsError { + t.Error("expected error for IPv4-mapped IPv6 loopback URL, got success") + } +} + +func TestWebFetch_BlocksMetadataIP(t *testing.T) { + tool, err := NewWebFetchTool(50000, testFetchLimit) + if err != nil { + t.Fatalf("Failed to create web fetch tool: %v", err) + } + result := tool.Execute(context.Background(), map[string]any{ + "url": "http://169.254.169.254/latest/meta-data", + }) + + if !result.IsError { + t.Error("expected error for cloud metadata IP, got success") + } +} + +func TestWebFetch_BlocksIPv6UniqueLocal(t *testing.T) { + tool, err := NewWebFetchTool(50000, testFetchLimit) + if err != nil { + t.Fatalf("Failed to create web fetch tool: %v", err) + } + result := tool.Execute(context.Background(), map[string]any{ + "url": "http://[fd00::1]:0", + }) + + if !result.IsError { + t.Error("expected error for IPv6 unique local address, got success") + } +} + +func TestWebFetch_Blocks6to4WithPrivateEmbed(t *testing.T) { + tool, err := NewWebFetchTool(50000, testFetchLimit) + if err != nil { + t.Fatalf("Failed to create web fetch tool: %v", err) + } + // 2002:7f00:0001::1 embeds 127.0.0.1 + result := tool.Execute(context.Background(), map[string]any{ + "url": "http://[2002:7f00:0001::1]:0", + }) + + if !result.IsError { + t.Error("expected error for 6to4 with private embedded IPv4, got success") + } +} + +func TestWebFetch_Allows6to4WithPublicEmbed(t *testing.T) { + if os.Getenv("USER") == "jules" { + t.Skip("Skipping test in sandbox environment as it will never pass due to environment restrictions") + } + tool, err := NewWebFetchTool(50000, testFetchLimit) + if err != nil { + t.Fatalf("Failed to create web fetch tool: %v", err) + } + // 2002:0801:0101::1 embeds 8.1.1.1 (public) — pre-flight should pass, + // connection will fail (no listener) but that's after the SSRF check. + result := tool.Execute(context.Background(), map[string]any{ + "url": "http://[2002:0801:0101::1]:0", + }) + + // Should NOT be blocked by SSRF check — error should be connection failure, not "private" + if result.IsError && strings.Contains(result.ForLLM, "private") { + t.Errorf("6to4 with public embedded IPv4 should not be blocked as private. Error: %q", result.ForLLM) + } +} + +func TestWebFetch_RedirectToPrivateBlocked(t *testing.T) { + withPrivateWebFetchHostsAllowed(t) + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // Redirect to a private IP + http.Redirect(w, r, "http://10.0.0.1/secret", http.StatusFound) + })) + defer server.Close() + + // Temporarily disable private host allowance for the redirect check + allowPrivateWebFetchHosts.Store(false) + defer allowPrivateWebFetchHosts.Store(true) + + tool, err := NewWebFetchTool(50000, testFetchLimit) + if err != nil { + t.Fatalf("Failed to create web fetch tool: %v", err) + } + result := tool.Execute(context.Background(), map[string]any{ + "url": server.URL, + }) + + if !result.IsError { + t.Error("expected error when redirecting to private IP, got success") + } +} + +func TestIsPrivateOrRestrictedIP_Table(t *testing.T) { + tests := []struct { + ip string + blocked bool + desc string + }{ + {"127.0.0.1", true, "IPv4 loopback"}, + {"10.0.0.1", true, "IPv4 private class A"}, + {"172.16.0.1", true, "IPv4 private class B"}, + {"192.168.1.1", true, "IPv4 private class C"}, + {"169.254.169.254", true, "link-local / cloud metadata"}, + {"100.64.0.1", true, "carrier-grade NAT"}, + {"0.0.0.0", true, "unspecified"}, + {"8.8.8.8", false, "public DNS"}, + {"1.1.1.1", false, "public DNS"}, + {"::1", true, "IPv6 loopback"}, + {"::ffff:127.0.0.1", true, "IPv4-mapped IPv6 loopback"}, + {"::ffff:10.0.0.1", true, "IPv4-mapped IPv6 private"}, + {"fc00::1", true, "IPv6 unique local"}, + {"fd00::1", true, "IPv6 unique local"}, + {"2002:7f00:0001::1", true, "6to4 with embedded 127.x (private)"}, + {"2002:0a00:0001::1", true, "6to4 with embedded 10.0.0.1 (private)"}, + {"2002:0801:0101::1", false, "6to4 with embedded 8.1.1.1 (public)"}, + {"2001:0000:4136:e378:8000:63bf:f5ff:fffe", true, "Teredo with client 10.0.0.1 (private)"}, + {"2001:0000:4136:e378:8000:63bf:f7f6:fefe", false, "Teredo with client 8.9.1.1 (public)"}, + {"2607:f8b0:4004:800::200e", false, "public IPv6 (Google)"}, + } + + for _, tt := range tests { + t.Run(tt.desc, func(t *testing.T) { + ip := net.ParseIP(tt.ip) + if ip == nil { + t.Fatalf("failed to parse IP: %s", tt.ip) + } + got := isPrivateOrRestrictedIP(ip) + if got != tt.blocked { + t.Errorf("isPrivateOrRestrictedIP(%s) = %v, want %v", tt.ip, got, tt.blocked) + } + }) + } +} diff --git a/pkg/tools/web_test.go b/pkg/tools/web_test.go deleted file mode 100644 index 2de41eedd..000000000 --- a/pkg/tools/web_test.go +++ /dev/null @@ -1,1144 +0,0 @@ -package tools - -import ( - "bytes" - "context" - "encoding/json" - "fmt" - "net" - "net/http" - "net/http/httptest" - "strings" - "testing" - "time" - - "jane/pkg/logger" -) - -const testFetchLimit = int64(10 * 1024 * 1024) - -// TestWebTool_WebFetch_Success verifies successful URL fetching -func TestWebTool_WebFetch_Success(t *testing.T) { - withPrivateWebFetchHostsAllowed(t) - - server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - w.Header().Set("Content-Type", "text/html") - w.WriteHeader(http.StatusOK) - w.Write([]byte("

Test Page

Content here

")) - })) - defer server.Close() - - tool, err := NewWebFetchTool(50000, testFetchLimit) - if err != nil { - t.Fatalf("Failed to create web fetch tool: %v", err) - } - - ctx := context.Background() - args := map[string]any{ - "url": server.URL, - } - - result := tool.Execute(ctx, args) - - // Success should not be an error - if result.IsError { - t.Errorf("Expected success, got IsError=true: %s", result.ForLLM) - } - - // ForLLM should contain the fetched content (full JSON result) - if !strings.Contains(result.ForLLM, "Test Page") { - t.Errorf("Expected ForLLM to contain 'Test Page', got: %s", result.ForLLM) - } - - // ForUser should contain summary - if !strings.Contains(result.ForUser, "bytes") && !strings.Contains(result.ForUser, "extractor") { - t.Errorf("Expected ForUser to contain summary, got: %s", result.ForUser) - } -} - -// TestWebTool_WebFetch_JSON verifies JSON content handling -func TestWebTool_WebFetch_JSON(t *testing.T) { - withPrivateWebFetchHostsAllowed(t) - - testData := map[string]string{"key": "value", "number": "123"} - expectedJSON, _ := json.MarshalIndent(testData, "", " ") - - server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - w.Header().Set("Content-Type", "application/json") - w.WriteHeader(http.StatusOK) - w.Write(expectedJSON) - })) - defer server.Close() - - tool, err := NewWebFetchTool(50000, testFetchLimit) - if err != nil { - logger.ErrorCF("agent", "Failed to create web fetch tool", map[string]any{"error": err.Error()}) - } - - ctx := context.Background() - args := map[string]any{ - "url": server.URL, - } - - result := tool.Execute(ctx, args) - - // Success should not be an error - if result.IsError { - t.Errorf("Expected success, got IsError=true: %s", result.ForLLM) - } - - // ForLLM should contain formatted JSON - if !strings.Contains(result.ForLLM, "key") && !strings.Contains(result.ForLLM, "value") { - t.Errorf("Expected ForLLM to contain JSON data, got: %s", result.ForLLM) - } -} - -// TestWebTool_WebFetch_InvalidURL verifies error handling for invalid URL -func TestWebTool_WebFetch_InvalidURL(t *testing.T) { - tool, err := NewWebFetchTool(50000, testFetchLimit) - if err != nil { - logger.ErrorCF("agent", "Failed to create web fetch tool", map[string]any{"error": err.Error()}) - } - - ctx := context.Background() - args := map[string]any{ - "url": "not-a-valid-url", - } - - result := tool.Execute(ctx, args) - - // Should return error result - if !result.IsError { - t.Errorf("Expected error for invalid URL") - } - - // Should contain error message (either "invalid URL" or scheme error) - if !strings.Contains(result.ForLLM, "URL") && !strings.Contains(result.ForUser, "URL") { - t.Errorf("Expected error message for invalid URL, got ForLLM: %s", result.ForLLM) - } -} - -// TestWebTool_WebFetch_UnsupportedScheme verifies error handling for non-http URLs -func TestWebTool_WebFetch_UnsupportedScheme(t *testing.T) { - tool, err := NewWebFetchTool(50000, testFetchLimit) - if err != nil { - logger.ErrorCF("agent", "Failed to create web fetch tool", map[string]any{"error": err.Error()}) - } - - ctx := context.Background() - args := map[string]any{ - "url": "ftp://example.com/file.txt", - } - - result := tool.Execute(ctx, args) - - // Should return error result - if !result.IsError { - t.Errorf("Expected error for unsupported URL scheme") - } - - // Should mention only http/https allowed - if !strings.Contains(result.ForLLM, "http/https") && !strings.Contains(result.ForUser, "http/https") { - t.Errorf("Expected scheme error message, got ForLLM: %s", result.ForLLM) - } -} - -// TestWebTool_WebFetch_MissingURL verifies error handling for missing URL -func TestWebTool_WebFetch_MissingURL(t *testing.T) { - tool, err := NewWebFetchTool(50000, testFetchLimit) - if err != nil { - logger.ErrorCF("agent", "Failed to create web fetch tool", map[string]any{"error": err.Error()}) - } - - ctx := context.Background() - args := map[string]any{} - - result := tool.Execute(ctx, args) - - // Should return error result - if !result.IsError { - t.Errorf("Expected error when URL is missing") - } - - // Should mention URL is required - if !strings.Contains(result.ForLLM, "url is required") && !strings.Contains(result.ForUser, "url is required") { - t.Errorf("Expected 'url is required' message, got ForLLM: %s", result.ForLLM) - } -} - -// TestWebTool_WebFetch_Truncation verifies content truncation -func TestWebTool_WebFetch_Truncation(t *testing.T) { - withPrivateWebFetchHostsAllowed(t) - - longContent := strings.Repeat("x", 20000) - - server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - w.Header().Set("Content-Type", "text/plain") - w.WriteHeader(http.StatusOK) - w.Write([]byte(longContent)) - })) - defer server.Close() - - tool, err := NewWebFetchTool(1000, testFetchLimit) // Limit to 1000 chars - if err != nil { - logger.ErrorCF("agent", "Failed to create web fetch tool", map[string]any{"error": err.Error()}) - } - - ctx := context.Background() - args := map[string]any{ - "url": server.URL, - } - - result := tool.Execute(ctx, args) - - // Success should not be an error - if result.IsError { - t.Errorf("Expected success, got IsError=true: %s", result.ForLLM) - } - - // ForLLM should contain truncated content (not the full 20000 chars) - resultMap := make(map[string]any) - json.Unmarshal([]byte(result.ForLLM), &resultMap) - if text, ok := resultMap["text"].(string); ok { - if len(text) > 1100 { // Allow some margin - t.Errorf("Expected content to be truncated to ~1000 chars, got: %d", len(text)) - } - } - - // Should be marked as truncated - if truncated, ok := resultMap["truncated"].(bool); !ok || !truncated { - t.Errorf("Expected 'truncated' to be true in result") - } -} - -func TestWebFetchTool_PayloadTooLarge(t *testing.T) { - withPrivateWebFetchHostsAllowed(t) - - // Create a mock HTTP server - ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - w.Header().Set("Content-Type", "text/html") - w.WriteHeader(http.StatusOK) - - // Generate a payload intentionally larger than our limit. - // Limit: 10 * 1024 * 1024 (10MB). We generate 10MB + 100 bytes of the letter 'A'. - largeData := bytes.Repeat([]byte("A"), int(testFetchLimit)+100) - - w.Write(largeData) - })) - // Ensure the server is shut down at the end of the test - defer ts.Close() - - // Initialize the tool - tool, err := NewWebFetchTool(50000, testFetchLimit) - if err != nil { - logger.ErrorCF("agent", "Failed to create web fetch tool", map[string]any{"error": err.Error()}) - } - - // Prepare the arguments pointing to the URL of our local mock server - args := map[string]any{ - "url": ts.URL, - } - - // Execute the tool - ctx := context.Background() - result := tool.Execute(ctx, args) - - // Assuming ErrorResult sets the ForLLM field with the error text. - if result == nil { - t.Fatal("expected a ToolResult, got nil") - } - - // Search for the exact error string we set earlier in the Execute method - expectedErrorMsg := fmt.Sprintf("size exceeded %d bytes limit", testFetchLimit) - - if !strings.Contains(result.ForLLM, expectedErrorMsg) && !strings.Contains(result.ForUser, expectedErrorMsg) { - t.Errorf("test failed: expected error %q, but got: %+v", expectedErrorMsg, result) - } -} - -// TestWebTool_WebSearch_NoApiKey verifies that no tool is created when API key is missing -func TestWebTool_WebSearch_NoApiKey(t *testing.T) { - tool, err := NewWebSearchTool(WebSearchToolOptions{BraveEnabled: true, BraveAPIKeys: nil}) - 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, 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") - } -} - -// TestWebTool_WebSearch_MissingQuery verifies error handling for missing query -func TestWebTool_WebSearch_MissingQuery(t *testing.T) { - tool, err := NewWebSearchTool(WebSearchToolOptions{ - BraveEnabled: true, - BraveAPIKeys: []string{"test-key"}, - BraveMaxResults: 5, - }) - if err != nil { - t.Fatalf("Unexpected error: %v", err) - } - ctx := context.Background() - args := map[string]any{} - - result := tool.Execute(ctx, args) - - // Should return error result - if !result.IsError { - t.Errorf("Expected error when query is missing") - } -} - -// TestWebTool_WebFetch_HTMLExtraction verifies HTML text extraction -func TestWebTool_WebFetch_HTMLExtraction(t *testing.T) { - withPrivateWebFetchHostsAllowed(t) - - server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - w.Header().Set("Content-Type", "text/html") - w.WriteHeader(http.StatusOK) - w.Write( - []byte( - `

Title

Content

`, - ), - ) - })) - defer server.Close() - - tool, err := NewWebFetchTool(50000, testFetchLimit) - if err != nil { - logger.ErrorCF("agent", "Failed to create web fetch tool", map[string]any{"error": err.Error()}) - } - - ctx := context.Background() - args := map[string]any{ - "url": server.URL, - } - - result := tool.Execute(ctx, args) - - // Success should not be an error - if result.IsError { - t.Errorf("Expected success, got IsError=true: %s", result.ForLLM) - } - - // ForLLM should contain extracted text (without script/style tags) - if !strings.Contains(result.ForLLM, "Title") && !strings.Contains(result.ForLLM, "Content") { - t.Errorf("Expected ForLLM to contain extracted text, got: %s", result.ForLLM) - } - - // Should NOT contain script or style tags in ForLLM - if strings.Contains(result.ForLLM, "

Keep this

", - wantFunc: func(t *testing.T, got string) { - if strings.Contains(got, "alert") || strings.Contains(got, "body{}") { - t.Errorf("Expected script/style content removed, got: %q", got) - } - if !strings.Contains(got, "Keep this") { - t.Errorf("Expected 'Keep this' to remain, got: %q", got) - } - }, - }, - { - name: "collapses excessive blank lines", - input: "

A

\n\n\n\n\n

B

", - wantFunc: func(t *testing.T, got string) { - if strings.Contains(got, "\n\n\n") { - t.Errorf("Expected excessive blank lines collapsed, got: %q", got) - } - }, - }, - { - name: "collapses horizontal whitespace", - input: "

hello world

", - wantFunc: func(t *testing.T, got string) { - if strings.Contains(got, " ") { - t.Errorf("Expected spaces collapsed, got: %q", got) - } - if !strings.Contains(got, "hello world") { - t.Errorf("Expected 'hello world', got: %q", got) - } - }, - }, - { - name: "empty input", - input: "", - wantFunc: func(t *testing.T, got string) { - if got != "" { - t.Errorf("Expected empty string, got: %q", got) - } - }, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - got := tool.extractText(tt.input) - tt.wantFunc(t, got) - }) - } -} - -func withPrivateWebFetchHostsAllowed(t *testing.T) { - t.Helper() - previous := allowPrivateWebFetchHosts.Load() - allowPrivateWebFetchHosts.Store(true) - t.Cleanup(func() { - allowPrivateWebFetchHosts.Store(previous) - }) -} - -func TestWebTool_WebFetch_PrivateHostBlocked(t *testing.T) { - tool, err := NewWebFetchTool(50000, testFetchLimit) - if err != nil { - t.Fatalf("Failed to create web fetch tool: %v", err) - } - result := tool.Execute(context.Background(), map[string]any{ - "url": "http://127.0.0.1:0", - }) - - if !result.IsError { - t.Errorf("expected error for private host URL, got success") - } - if !strings.Contains(result.ForLLM, "private or local network") && - !strings.Contains(result.ForUser, "private or local network") { - t.Errorf("expected private host block message, got %q", result.ForLLM) - } -} - -func TestWebTool_WebFetch_PrivateHostAllowedForTests(t *testing.T) { - withPrivateWebFetchHostsAllowed(t) - - server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - w.Header().Set("Content-Type", "text/plain") - w.WriteHeader(http.StatusOK) - w.Write([]byte("ok")) - })) - defer server.Close() - - tool, err := NewWebFetchTool(50000, testFetchLimit) - if err != nil { - t.Fatalf("Failed to create web fetch tool: %v", err) - } - result := tool.Execute(context.Background(), map[string]any{ - "url": server.URL, - }) - - if result.IsError { - t.Errorf("expected success when private host access is allowed in tests, got %q", result.ForLLM) - } -} - -// TestWebFetch_BlocksIPv4MappedIPv6Loopback verifies ::ffff:127.0.0.1 is blocked -func TestWebFetch_BlocksIPv4MappedIPv6Loopback(t *testing.T) { - tool, err := NewWebFetchTool(50000, testFetchLimit) - if err != nil { - t.Fatalf("Failed to create web fetch tool: %v", err) - } - result := tool.Execute(context.Background(), map[string]any{ - "url": "http://[::ffff:127.0.0.1]:0", - }) - - if !result.IsError { - t.Error("expected error for IPv4-mapped IPv6 loopback URL, got success") - } -} - -// TestWebFetch_BlocksMetadataIP verifies 169.254.169.254 is blocked -func TestWebFetch_BlocksMetadataIP(t *testing.T) { - tool, err := NewWebFetchTool(50000, testFetchLimit) - if err != nil { - t.Fatalf("Failed to create web fetch tool: %v", err) - } - result := tool.Execute(context.Background(), map[string]any{ - "url": "http://169.254.169.254/latest/meta-data", - }) - - if !result.IsError { - t.Error("expected error for cloud metadata IP, got success") - } -} - -// TestWebFetch_BlocksIPv6UniqueLocal verifies fc00::/7 addresses are blocked -func TestWebFetch_BlocksIPv6UniqueLocal(t *testing.T) { - tool, err := NewWebFetchTool(50000, testFetchLimit) - if err != nil { - t.Fatalf("Failed to create web fetch tool: %v", err) - } - result := tool.Execute(context.Background(), map[string]any{ - "url": "http://[fd00::1]:0", - }) - - if !result.IsError { - t.Error("expected error for IPv6 unique local address, got success") - } -} - -// TestWebFetch_Blocks6to4WithPrivateEmbed verifies 6to4 with private embedded IPv4 is blocked -func TestWebFetch_Blocks6to4WithPrivateEmbed(t *testing.T) { - tool, err := NewWebFetchTool(50000, testFetchLimit) - if err != nil { - t.Fatalf("Failed to create web fetch tool: %v", err) - } - // 2002:7f00:0001::1 embeds 127.0.0.1 - result := tool.Execute(context.Background(), map[string]any{ - "url": "http://[2002:7f00:0001::1]:0", - }) - - if !result.IsError { - t.Error("expected error for 6to4 with private embedded IPv4, got success") - } -} - -// TestWebFetch_Allows6to4WithPublicEmbed verifies 6to4 with public embedded IPv4 is NOT blocked -func TestWebFetch_Allows6to4WithPublicEmbed(t *testing.T) { - tool, err := NewWebFetchTool(50000, testFetchLimit) - if err != nil { - t.Fatalf("Failed to create web fetch tool: %v", err) - } - // 2002:0801:0101::1 embeds 8.1.1.1 (public) — pre-flight should pass, - // connection will fail (no listener) but that's after the SSRF check. - result := tool.Execute(context.Background(), map[string]any{ - "url": "http://[2002:0801:0101::1]:0", - }) - - // Should NOT be blocked by SSRF check — error should be connection failure, not "private" - if result.IsError && strings.Contains(result.ForLLM, "private") { - t.Error("6to4 with public embedded IPv4 should not be blocked as private") - } -} - -// TestWebFetch_RedirectToPrivateBlocked verifies redirects to private IPs are blocked -func TestWebFetch_RedirectToPrivateBlocked(t *testing.T) { - withPrivateWebFetchHostsAllowed(t) - - server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - // Redirect to a private IP - http.Redirect(w, r, "http://10.0.0.1/secret", http.StatusFound) - })) - defer server.Close() - - // Temporarily disable private host allowance for the redirect check - allowPrivateWebFetchHosts.Store(false) - defer allowPrivateWebFetchHosts.Store(true) - - tool, err := NewWebFetchTool(50000, testFetchLimit) - if err != nil { - t.Fatalf("Failed to create web fetch tool: %v", err) - } - result := tool.Execute(context.Background(), map[string]any{ - "url": server.URL, - }) - - if !result.IsError { - t.Error("expected error when redirecting to private IP, got success") - } -} - -// TestIsPrivateOrRestrictedIP_Table tests IP classification logic -func TestIsPrivateOrRestrictedIP_Table(t *testing.T) { - tests := []struct { - ip string - blocked bool - desc string - }{ - {"127.0.0.1", true, "IPv4 loopback"}, - {"10.0.0.1", true, "IPv4 private class A"}, - {"172.16.0.1", true, "IPv4 private class B"}, - {"192.168.1.1", true, "IPv4 private class C"}, - {"169.254.169.254", true, "link-local / cloud metadata"}, - {"100.64.0.1", true, "carrier-grade NAT"}, - {"0.0.0.0", true, "unspecified"}, - {"8.8.8.8", false, "public DNS"}, - {"1.1.1.1", false, "public DNS"}, - {"::1", true, "IPv6 loopback"}, - {"::ffff:127.0.0.1", true, "IPv4-mapped IPv6 loopback"}, - {"::ffff:10.0.0.1", true, "IPv4-mapped IPv6 private"}, - {"fc00::1", true, "IPv6 unique local"}, - {"fd00::1", true, "IPv6 unique local"}, - {"2002:7f00:0001::1", true, "6to4 with embedded 127.x (private)"}, - {"2002:0a00:0001::1", true, "6to4 with embedded 10.0.0.1 (private)"}, - {"2002:0801:0101::1", false, "6to4 with embedded 8.1.1.1 (public)"}, - {"2001:0000:4136:e378:8000:63bf:f5ff:fffe", true, "Teredo with client 10.0.0.1 (private)"}, - {"2001:0000:4136:e378:8000:63bf:f7f6:fefe", false, "Teredo with client 8.9.1.1 (public)"}, - {"2607:f8b0:4004:800::200e", false, "public IPv6 (Google)"}, - } - - for _, tt := range tests { - t.Run(tt.desc, func(t *testing.T) { - ip := net.ParseIP(tt.ip) - if ip == nil { - t.Fatalf("failed to parse IP: %s", tt.ip) - } - got := isPrivateOrRestrictedIP(ip) - if got != tt.blocked { - t.Errorf("isPrivateOrRestrictedIP(%s) = %v, want %v", tt.ip, got, tt.blocked) - } - }) - } -} - -// TestWebTool_WebFetch_MissingDomain verifies error handling for URL without domain -func TestWebTool_WebFetch_MissingDomain(t *testing.T) { - tool, err := NewWebFetchTool(50000, testFetchLimit) - if err != nil { - logger.ErrorCF("agent", "Failed to create web fetch tool", map[string]any{"error": err.Error()}) - } - - ctx := context.Background() - args := map[string]any{ - "url": "https://", - } - - result := tool.Execute(ctx, args) - - // Should return error result - if !result.IsError { - t.Errorf("Expected error for URL without domain") - } - - // Should mention missing domain - if !strings.Contains(result.ForLLM, "domain") && !strings.Contains(result.ForUser, "domain") { - t.Errorf("Expected domain error message, got ForLLM: %s", result.ForLLM) - } -} - -func TestCreateHTTPClient_ProxyConfigured(t *testing.T) { - client, err := createHTTPClient("http://127.0.0.1:7890", 12*time.Second) - if err != nil { - t.Fatalf("createHTTPClient() error: %v", err) - } - if client.Timeout != 12*time.Second { - t.Fatalf("client.Timeout = %v, want %v", client.Timeout, 12*time.Second) - } - - tr, ok := client.Transport.(*http.Transport) - if !ok { - t.Fatalf("client.Transport type = %T, want *http.Transport", client.Transport) - } - if tr.Proxy == nil { - t.Fatal("transport.Proxy is nil, want non-nil") - } - - req, err := http.NewRequest("GET", "https://example.com", nil) - if err != nil { - t.Fatalf("http.NewRequest() error: %v", err) - } - proxyURL, err := tr.Proxy(req) - if err != nil { - t.Fatalf("transport.Proxy(req) error: %v", err) - } - if proxyURL == nil || proxyURL.String() != "http://127.0.0.1:7890" { - t.Fatalf("proxy URL = %v, want %q", proxyURL, "http://127.0.0.1:7890") - } -} - -func TestCreateHTTPClient_InvalidProxy(t *testing.T) { - _, err := createHTTPClient("://bad-proxy", 10*time.Second) - if err == nil { - t.Fatal("createHTTPClient() expected error for invalid proxy URL, got nil") - } -} - -func TestCreateHTTPClient_Socks5ProxyConfigured(t *testing.T) { - client, err := createHTTPClient("socks5://127.0.0.1:1080", 8*time.Second) - if err != nil { - t.Fatalf("createHTTPClient() error: %v", err) - } - - tr, ok := client.Transport.(*http.Transport) - if !ok { - t.Fatalf("client.Transport type = %T, want *http.Transport", client.Transport) - } - req, err := http.NewRequest("GET", "https://example.com", nil) - if err != nil { - t.Fatalf("http.NewRequest() error: %v", err) - } - proxyURL, err := tr.Proxy(req) - if err != nil { - t.Fatalf("transport.Proxy(req) error: %v", err) - } - if proxyURL == nil || proxyURL.String() != "socks5://127.0.0.1:1080" { - t.Fatalf("proxy URL = %v, want %q", proxyURL, "socks5://127.0.0.1:1080") - } -} - -func TestCreateHTTPClient_UnsupportedProxyScheme(t *testing.T) { - _, err := createHTTPClient("ftp://127.0.0.1:21", 10*time.Second) - if err == nil { - t.Fatal("createHTTPClient() expected error for unsupported scheme, got nil") - } - if !strings.Contains(err.Error(), "unsupported proxy scheme") { - t.Fatalf("error = %q, want to contain %q", err.Error(), "unsupported proxy scheme") - } -} - -func TestCreateHTTPClient_ProxyFromEnvironmentWhenConfigEmpty(t *testing.T) { - t.Setenv("HTTP_PROXY", "http://127.0.0.1:8888") - t.Setenv("http_proxy", "http://127.0.0.1:8888") - t.Setenv("HTTPS_PROXY", "http://127.0.0.1:8888") - t.Setenv("https_proxy", "http://127.0.0.1:8888") - t.Setenv("ALL_PROXY", "") - t.Setenv("all_proxy", "") - t.Setenv("NO_PROXY", "") - t.Setenv("no_proxy", "") - - client, err := createHTTPClient("", 10*time.Second) - if err != nil { - t.Fatalf("createHTTPClient() error: %v", err) - } - - tr, ok := client.Transport.(*http.Transport) - if !ok { - t.Fatalf("client.Transport type = %T, want *http.Transport", client.Transport) - } - if tr.Proxy == nil { - t.Fatal("transport.Proxy is nil, want proxy function from environment") - } - - req, err := http.NewRequest("GET", "https://example.com", nil) - if err != nil { - t.Fatalf("http.NewRequest() error: %v", err) - } - if _, err := tr.Proxy(req); err != nil { - t.Fatalf("transport.Proxy(req) error: %v", err) - } -} - -func TestNewWebFetchToolWithProxy(t *testing.T) { - tool, err := NewWebFetchToolWithProxy(1024, "http://127.0.0.1:7890", testFetchLimit) - if err != nil { - logger.ErrorCF("agent", "Failed to create web fetch tool", map[string]any{"error": err.Error()}) - } else if tool.maxChars != 1024 { - t.Fatalf("maxChars = %d, want %d", tool.maxChars, 1024) - } - - if tool.proxy != "http://127.0.0.1:7890" { - t.Fatalf("proxy = %q, want %q", tool.proxy, "http://127.0.0.1:7890") - } - - tool, err = NewWebFetchToolWithProxy(0, "http://127.0.0.1:7890", testFetchLimit) - if err != nil { - logger.ErrorCF("agent", "Failed to create web fetch tool", map[string]any{"error": err.Error()}) - } - - if tool.maxChars != 50000 { - t.Fatalf("default maxChars = %d, want %d", tool.maxChars, 50000) - } -} - -func TestNewWebSearchTool_PropagatesProxy(t *testing.T) { - t.Run("perplexity", func(t *testing.T) { - tool, err := NewWebSearchTool(WebSearchToolOptions{ - PerplexityEnabled: true, - PerplexityAPIKeys: []string{"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) - } - if p.proxy != "http://127.0.0.1:7890" { - t.Fatalf("provider proxy = %q, want %q", p.proxy, "http://127.0.0.1:7890") - } - }) - - t.Run("brave", func(t *testing.T) { - tool, err := NewWebSearchTool(WebSearchToolOptions{ - BraveEnabled: true, - BraveAPIKeys: []string{"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) - } - if p.proxy != "http://127.0.0.1:7890" { - t.Fatalf("provider proxy = %q, want %q", p.proxy, "http://127.0.0.1:7890") - } - }) - - t.Run("duckduckgo", func(t *testing.T) { - 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) - } - if p.proxy != "http://127.0.0.1:7890" { - t.Fatalf("provider proxy = %q, want %q", p.proxy, "http://127.0.0.1:7890") - } - }) -} - -// TestWebTool_TavilySearch_Success verifies successful Tavily search -func TestWebTool_TavilySearch_Success(t *testing.T) { - server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if r.Method != "POST" { - t.Errorf("Expected POST request, got %s", r.Method) - } - if r.Header.Get("Content-Type") != "application/json" { - t.Errorf("Expected Content-Type application/json, got %s", r.Header.Get("Content-Type")) - } - - // Verify payload - var payload map[string]any - json.NewDecoder(r.Body).Decode(&payload) - if payload["api_key"] != "test-key" { - t.Errorf("Expected api_key test-key, got %v", payload["api_key"]) - } - if payload["query"] != "test query" { - t.Errorf("Expected query 'test query', got %v", payload["query"]) - } - - // Return mock response - response := map[string]any{ - "results": []map[string]any{ - { - "title": "Test Result 1", - "url": "https://example.com/1", - "content": "Content for result 1", - }, - { - "title": "Test Result 2", - "url": "https://example.com/2", - "content": "Content for result 2", - }, - }, - } - w.Header().Set("Content-Type", "application/json") - w.WriteHeader(http.StatusOK) - json.NewEncoder(w).Encode(response) - })) - defer server.Close() - - tool, err := NewWebSearchTool(WebSearchToolOptions{ - TavilyEnabled: true, - TavilyAPIKeys: []string{"test-key"}, - TavilyBaseURL: server.URL, - TavilyMaxResults: 5, - }) - if err != nil { - t.Fatalf("NewWebSearchTool() error: %v", err) - } - - ctx := context.Background() - args := map[string]any{ - "query": "test query", - } - - result := tool.Execute(ctx, args) - - // Success should not be an error - if result.IsError { - t.Errorf("Expected success, got IsError=true: %s", result.ForLLM) - } - - // ForUser should contain result titles and URLs - if !strings.Contains(result.ForUser, "Test Result 1") || - !strings.Contains(result.ForUser, "https://example.com/1") { - t.Errorf("Expected results in output, got: %s", result.ForUser) - } - - // Should mention via Tavily - if !strings.Contains(result.ForUser, "via Tavily") { - t.Errorf("Expected 'via Tavily' in output, got: %s", result.ForUser) - } -} - -func TestAPIKeyPool(t *testing.T) { - pool := NewAPIKeyPool([]string{"key1", "key2", "key3"}) - if len(pool.keys) != 3 { - t.Fatalf("expected 3 keys, got %d", len(pool.keys)) - } - if pool.keys[0] != "key1" || pool.keys[1] != "key2" || pool.keys[2] != "key3" { - t.Fatalf("unexpected keys: %v", pool.keys) - } - - // Test Iterator: each iterator should cover all keys exactly once - iter := pool.NewIterator() - expected := []string{"key1", "key2", "key3"} - for i, want := range expected { - k, ok := iter.Next() - if !ok { - t.Fatalf("iter.Next() returned false at step %d", i) - } - if k != want { - t.Errorf("step %d: expected %s, got %s", i, want, k) - } - } - // Should be exhausted - if _, ok := iter.Next(); ok { - t.Errorf("expected iterator exhausted after all keys") - } - - // Second iterator starts at next position (load balancing) - iter2 := pool.NewIterator() - k, ok := iter2.Next() - if !ok { - t.Fatal("iter2.Next() returned false") - } - if k != "key2" { - t.Errorf("expected key2 (round-robin), got %s", k) - } - - // Empty pool - emptyPool := NewAPIKeyPool([]string{}) - emptyIter := emptyPool.NewIterator() - if _, ok := emptyIter.Next(); ok { - t.Errorf("expected false for empty pool") - } - - // Single key pool - singlePool := NewAPIKeyPool([]string{"single"}) - singleIter := singlePool.NewIterator() - if k, ok := singleIter.Next(); !ok || k != "single" { - t.Errorf("expected single, got %s (ok=%v)", k, ok) - } - if _, ok := singleIter.Next(); ok { - t.Errorf("expected exhausted after single key") - } -} - -func TestWebTool_TavilySearch_Failover(t *testing.T) { - server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - var payload map[string]any - if err := json.NewDecoder(r.Body).Decode(&payload); err != nil { - t.Fatalf("failed to decode payload: %v", err) - } - - apiKey := payload["api_key"].(string) - - if apiKey == "key1" { - w.WriteHeader(http.StatusTooManyRequests) - w.Write([]byte("Rate limited")) - return - } - - if apiKey == "key2" { - // Success - response := map[string]any{ - "results": []map[string]any{ - { - "title": "Success Result", - "url": "https://example.com/success", - "content": "Success content", - }, - }, - } - w.Header().Set("Content-Type", "application/json") - w.WriteHeader(http.StatusOK) - json.NewEncoder(w).Encode(response) - return - } - - w.WriteHeader(http.StatusBadRequest) - })) - defer server.Close() - - tool, err := NewWebSearchTool(WebSearchToolOptions{ - TavilyEnabled: true, - TavilyAPIKeys: []string{"key1", "key2"}, - TavilyBaseURL: server.URL, - TavilyMaxResults: 5, - }) - if err != nil { - t.Fatalf("NewWebSearchTool() error: %v", err) - } - - ctx := context.Background() - args := map[string]any{ - "query": "test query", - } - - result := tool.Execute(ctx, args) - - if result.IsError { - t.Errorf("Expected success, got Error: %s", result.ForLLM) - } - if !strings.Contains(result.ForUser, "Success Result") { - t.Errorf("Expected failover to second key and success result, got: %s", result.ForUser) - } -} - -func TestWebTool_GLMSearch_Success(t *testing.T) { - server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if r.Method != "POST" { - t.Errorf("Expected POST request, got %s", r.Method) - } - if r.Header.Get("Content-Type") != "application/json" { - t.Errorf("Expected Content-Type application/json, got %s", r.Header.Get("Content-Type")) - } - if r.Header.Get("Authorization") != "Bearer test-glm-key" { - t.Errorf("Expected Authorization Bearer test-glm-key, got %s", r.Header.Get("Authorization")) - } - - var payload map[string]any - json.NewDecoder(r.Body).Decode(&payload) - if payload["search_query"] != "test query" { - t.Errorf("Expected search_query 'test query', got %v", payload["search_query"]) - } - if payload["search_engine"] != "search_std" { - t.Errorf("Expected search_engine 'search_std', got %v", payload["search_engine"]) - } - - response := map[string]any{ - "id": "web-search-test", - "created": 1709568000, - "search_result": []map[string]any{ - { - "title": "Test GLM Result", - "content": "GLM search snippet", - "link": "https://example.com/glm", - "media": "Example", - "publish_date": "2026-03-04", - }, - }, - } - w.Header().Set("Content-Type", "application/json") - w.WriteHeader(http.StatusOK) - json.NewEncoder(w).Encode(response) - })) - defer server.Close() - - tool, err := NewWebSearchTool(WebSearchToolOptions{ - GLMSearchEnabled: true, - GLMSearchAPIKey: "test-glm-key", - GLMSearchBaseURL: server.URL, - GLMSearchEngine: "search_std", - }) - if err != nil { - t.Fatalf("NewWebSearchTool() error: %v", err) - } - - result := tool.Execute(context.Background(), map[string]any{ - "query": "test query", - }) - - if result.IsError { - t.Errorf("Expected success, got IsError=true: %s", result.ForLLM) - } - if !strings.Contains(result.ForUser, "Test GLM Result") { - t.Errorf("Expected 'Test GLM Result' in output, got: %s", result.ForUser) - } - if !strings.Contains(result.ForUser, "https://example.com/glm") { - t.Errorf("Expected URL in output, got: %s", result.ForUser) - } - if !strings.Contains(result.ForUser, "via GLM Search") { - t.Errorf("Expected 'via GLM Search' in output, got: %s", result.ForUser) - } -} - -func TestWebTool_GLMSearch_APIError(t *testing.T) { - server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { - w.WriteHeader(http.StatusUnauthorized) - w.Write([]byte(`{"error":"invalid api key"}`)) - })) - defer server.Close() - - tool, err := NewWebSearchTool(WebSearchToolOptions{ - GLMSearchEnabled: true, - GLMSearchAPIKey: "bad-key", - GLMSearchBaseURL: server.URL, - GLMSearchEngine: "search_std", - }) - if err != nil { - t.Fatalf("NewWebSearchTool() error: %v", err) - } - - result := tool.Execute(context.Background(), map[string]any{ - "query": "test query", - }) - - if !result.IsError { - t.Errorf("Expected IsError=true for 401 response") - } - if !strings.Contains(result.ForLLM, "status 401") { - t.Errorf("Expected status 401 in error, got: %s", result.ForLLM) - } -} - -func TestWebTool_GLMSearch_Priority(t *testing.T) { - // GLM Search should only be selected when all other providers are disabled - tool, err := NewWebSearchTool(WebSearchToolOptions{ - DuckDuckGoEnabled: true, - DuckDuckGoMaxResults: 5, - GLMSearchEnabled: true, - GLMSearchAPIKey: "test-key", - GLMSearchBaseURL: "https://example.com", - GLMSearchEngine: "search_std", - }) - if err != nil { - t.Fatalf("NewWebSearchTool() error: %v", err) - } - - // DuckDuckGo should win over GLM Search - if _, ok := tool.provider.(*DuckDuckGoSearchProvider); !ok { - t.Errorf("Expected DuckDuckGoSearchProvider when both enabled, got %T", tool.provider) - } - - // With DuckDuckGo disabled, GLM Search should be selected - tool2, err := NewWebSearchTool(WebSearchToolOptions{ - DuckDuckGoEnabled: false, - GLMSearchEnabled: true, - GLMSearchAPIKey: "test-key", - GLMSearchBaseURL: "https://example.com", - GLMSearchEngine: "search_std", - }) - if err != nil { - t.Fatalf("NewWebSearchTool() error: %v", err) - } - if _, ok := tool2.provider.(*GLMSearchProvider); !ok { - t.Errorf("Expected GLMSearchProvider when only GLM enabled, got %T", tool2.provider) - } -}