diff --git a/pkg/tools/web.go b/pkg/tools/web.go index 8ba2a723a..bac05a862 100644 --- a/pkg/tools/web.go +++ b/pkg/tools/web.go @@ -4,6 +4,7 @@ import ( "bytes" "context" "encoding/json" + "errors" "fmt" "io" "net/http" @@ -14,7 +15,8 @@ import ( ) const ( - userAgent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36" + userAgent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36" + MaxFetchLimitBytes = int64(10 * 1024 * 1024) // 10MB limit ) // Pre-compiled regexes for HTML text extraction @@ -605,10 +607,17 @@ func (t *WebFetchTool) Execute(ctx context.Context, args map[string]any) *ToolRe if err != nil { return ErrorResult(fmt.Sprintf("request failed: %v", err)) } + + resp.Body = http.MaxBytesReader(nil, resp.Body, MaxFetchLimitBytes) + 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", MaxFetchLimitBytes)) + } return ErrorResult(fmt.Sprintf("failed to read response: %v", err)) } diff --git a/pkg/tools/web_test.go b/pkg/tools/web_test.go index 2cd79eb24..6735d8b05 100644 --- a/pkg/tools/web_test.go +++ b/pkg/tools/web_test.go @@ -1,8 +1,10 @@ package tools import ( + "bytes" "context" "encoding/json" + "fmt" "net/http" "net/http/httptest" "strings" @@ -174,6 +176,46 @@ func TestWebTool_WebFetch_Truncation(t *testing.T) { } } +func TestWebFetchTool_PayloadTooLarge(t *testing.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(MaxFetchLimitBytes)+100) + + w.Write(largeData) + })) + // Ensure the server is shut down at the end of the test + defer ts.Close() + + // Initialize the tool + tool := NewWebFetchTool(50000) + + // 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", MaxFetchLimitBytes) + + 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 := NewWebSearchTool(WebSearchToolOptions{BraveEnabled: true, BraveAPIKey: ""})