fix: max payload size in web fetch
This commit is contained in:
parent
2f4f45080b
commit
a9a307584b
2 changed files with 52 additions and 1 deletions
|
|
@ -4,6 +4,7 @@ import (
|
||||||
"bytes"
|
"bytes"
|
||||||
"context"
|
"context"
|
||||||
"encoding/json"
|
"encoding/json"
|
||||||
|
"errors"
|
||||||
"fmt"
|
"fmt"
|
||||||
"io"
|
"io"
|
||||||
"net/http"
|
"net/http"
|
||||||
|
|
@ -15,6 +16,7 @@ import (
|
||||||
|
|
||||||
const (
|
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
|
// 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 {
|
if err != nil {
|
||||||
return ErrorResult(fmt.Sprintf("request failed: %v", err))
|
return ErrorResult(fmt.Sprintf("request failed: %v", err))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
resp.Body = http.MaxBytesReader(nil, resp.Body, MaxFetchLimitBytes)
|
||||||
|
|
||||||
defer resp.Body.Close()
|
defer resp.Body.Close()
|
||||||
|
|
||||||
body, err := io.ReadAll(resp.Body)
|
body, err := io.ReadAll(resp.Body)
|
||||||
if err != nil {
|
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))
|
return ErrorResult(fmt.Sprintf("failed to read response: %v", err))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -1,8 +1,10 @@
|
||||||
package tools
|
package tools
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"bytes"
|
||||||
"context"
|
"context"
|
||||||
"encoding/json"
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
"net/http"
|
"net/http"
|
||||||
"net/http/httptest"
|
"net/http/httptest"
|
||||||
"strings"
|
"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
|
// TestWebTool_WebSearch_NoApiKey verifies that no tool is created when API key is missing
|
||||||
func TestWebTool_WebSearch_NoApiKey(t *testing.T) {
|
func TestWebTool_WebSearch_NoApiKey(t *testing.T) {
|
||||||
tool := NewWebSearchTool(WebSearchToolOptions{BraveEnabled: true, BraveAPIKey: ""})
|
tool := NewWebSearchTool(WebSearchToolOptions{BraveEnabled: true, BraveAPIKey: ""})
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue