Merge branch 'main' into mcp-tools-support

This commit is contained in:
daming大铭
2026-03-03 00:17:39 +08:00
committed by GitHub
136 changed files with 12924 additions and 2845 deletions
+12 -6
View File
@@ -3,6 +3,7 @@ package tools
import (
"context"
"fmt"
"strings"
"sync"
"time"
@@ -33,15 +34,19 @@ type CronTool struct {
func NewCronTool(
cronService *cron.CronService, executor JobExecutor, msgBus *bus.MessageBus, workspace string, restrict bool,
execTimeout time.Duration, config *config.Config,
) *CronTool {
execTool := NewExecToolWithConfig(workspace, restrict, config)
) (*CronTool, error) {
execTool, err := NewExecToolWithConfig(workspace, restrict, config)
if err != nil {
return nil, fmt.Errorf("unable to configure exec tool: %w", err)
}
execTool.SetTimeout(execTimeout)
return &CronTool{
cronService: cronService,
executor: executor,
msgBus: msgBus,
execTool: execTool,
}
}, nil
}
// Name returns the tool name
@@ -218,7 +223,8 @@ func (t *CronTool) listJobs() *ToolResult {
return SilentResult("No scheduled jobs")
}
result := "Scheduled jobs:\n"
var result strings.Builder
result.WriteString("Scheduled jobs:\n")
for _, j := range jobs {
var scheduleInfo string
if j.Schedule.Kind == "every" && j.Schedule.EveryMS != nil {
@@ -230,10 +236,10 @@ func (t *CronTool) listJobs() *ToolResult {
} else {
scheduleInfo = "unknown"
}
result += fmt.Sprintf("- %s (id: %s, %s)\n", j.Name, j.ID, scheduleInfo)
result.WriteString(fmt.Sprintf("- %s (id: %s, %s)\n", j.Name, j.ID, scheduleInfo))
}
return SilentResult(result)
return SilentResult(result.String())
}
func (t *CronTool) removeJob(args map[string]any) *ToolResult {
+11 -14
View File
@@ -5,6 +5,7 @@ import (
"errors"
"fmt"
"io/fs"
"regexp"
"strings"
)
@@ -15,14 +16,12 @@ type EditFileTool struct {
}
// NewEditFileTool creates a new EditFileTool with optional directory restriction.
func NewEditFileTool(workspace string, restrict bool) *EditFileTool {
var fs fileSystem
if restrict {
fs = &sandboxFs{workspace: workspace}
} else {
fs = &hostFs{}
func NewEditFileTool(workspace string, restrict bool, allowPaths ...[]*regexp.Regexp) *EditFileTool {
var patterns []*regexp.Regexp
if len(allowPaths) > 0 {
patterns = allowPaths[0]
}
return &EditFileTool{fs: fs}
return &EditFileTool{fs: buildFs(workspace, restrict, patterns)}
}
func (t *EditFileTool) Name() string {
@@ -80,14 +79,12 @@ type AppendFileTool struct {
fs fileSystem
}
func NewAppendFileTool(workspace string, restrict bool) *AppendFileTool {
var fs fileSystem
if restrict {
fs = &sandboxFs{workspace: workspace}
} else {
fs = &hostFs{}
func NewAppendFileTool(workspace string, restrict bool, allowPaths ...[]*regexp.Regexp) *AppendFileTool {
var patterns []*regexp.Regexp
if len(allowPaths) > 0 {
patterns = allowPaths[0]
}
return &AppendFileTool{fs: fs}
return &AppendFileTool{fs: buildFs(workspace, restrict, patterns)}
}
func (t *AppendFileTool) Name() string {
+67 -21
View File
@@ -6,6 +6,7 @@ import (
"io/fs"
"os"
"path/filepath"
"regexp"
"strings"
"time"
@@ -87,14 +88,12 @@ type ReadFileTool struct {
fs fileSystem
}
func NewReadFileTool(workspace string, restrict bool) *ReadFileTool {
var fs fileSystem
if restrict {
fs = &sandboxFs{workspace: workspace}
} else {
fs = &hostFs{}
func NewReadFileTool(workspace string, restrict bool, allowPaths ...[]*regexp.Regexp) *ReadFileTool {
var patterns []*regexp.Regexp
if len(allowPaths) > 0 {
patterns = allowPaths[0]
}
return &ReadFileTool{fs: fs}
return &ReadFileTool{fs: buildFs(workspace, restrict, patterns)}
}
func (t *ReadFileTool) Name() string {
@@ -135,14 +134,12 @@ type WriteFileTool struct {
fs fileSystem
}
func NewWriteFileTool(workspace string, restrict bool) *WriteFileTool {
var fs fileSystem
if restrict {
fs = &sandboxFs{workspace: workspace}
} else {
fs = &hostFs{}
func NewWriteFileTool(workspace string, restrict bool, allowPaths ...[]*regexp.Regexp) *WriteFileTool {
var patterns []*regexp.Regexp
if len(allowPaths) > 0 {
patterns = allowPaths[0]
}
return &WriteFileTool{fs: fs}
return &WriteFileTool{fs: buildFs(workspace, restrict, patterns)}
}
func (t *WriteFileTool) Name() string {
@@ -192,14 +189,12 @@ type ListDirTool struct {
fs fileSystem
}
func NewListDirTool(workspace string, restrict bool) *ListDirTool {
var fs fileSystem
if restrict {
fs = &sandboxFs{workspace: workspace}
} else {
fs = &hostFs{}
func NewListDirTool(workspace string, restrict bool, allowPaths ...[]*regexp.Regexp) *ListDirTool {
var patterns []*regexp.Regexp
if len(allowPaths) > 0 {
patterns = allowPaths[0]
}
return &ListDirTool{fs: fs}
return &ListDirTool{fs: buildFs(workspace, restrict, patterns)}
}
func (t *ListDirTool) Name() string {
@@ -394,6 +389,57 @@ func (r *sandboxFs) ReadDir(path string) ([]os.DirEntry, error) {
return entries, err
}
// whitelistFs wraps a sandboxFs and allows access to specific paths outside
// the workspace when they match any of the provided patterns.
type whitelistFs struct {
sandbox *sandboxFs
host hostFs
patterns []*regexp.Regexp
}
func (w *whitelistFs) matches(path string) bool {
for _, p := range w.patterns {
if p.MatchString(path) {
return true
}
}
return false
}
func (w *whitelistFs) ReadFile(path string) ([]byte, error) {
if w.matches(path) {
return w.host.ReadFile(path)
}
return w.sandbox.ReadFile(path)
}
func (w *whitelistFs) WriteFile(path string, data []byte) error {
if w.matches(path) {
return w.host.WriteFile(path, data)
}
return w.sandbox.WriteFile(path, data)
}
func (w *whitelistFs) ReadDir(path string) ([]os.DirEntry, error) {
if w.matches(path) {
return w.host.ReadDir(path)
}
return w.sandbox.ReadDir(path)
}
// buildFs returns the appropriate fileSystem implementation based on restriction
// settings and optional path whitelist patterns.
func buildFs(workspace string, restrict bool, patterns []*regexp.Regexp) fileSystem {
if !restrict {
return &hostFs{}
}
sandbox := &sandboxFs{workspace: workspace}
if len(patterns) > 0 {
return &whitelistFs{sandbox: sandbox, patterns: patterns}
}
return sandbox
}
// Helper to get a safe relative path for os.Root usage
func getSafeRelPath(workspace, path string) (string, error) {
if workspace == "" {
+34
View File
@@ -5,6 +5,7 @@ import (
"io"
"os"
"path/filepath"
"regexp"
"strings"
"testing"
@@ -486,3 +487,36 @@ func TestRootRW_Write(t *testing.T) {
assert.NoError(t, err)
assert.Equal(t, newData, content)
}
// TestWhitelistFs_AllowsMatchingPaths verifies that whitelistFs allows access to
// paths matching the whitelist patterns while blocking non-matching paths.
func TestWhitelistFs_AllowsMatchingPaths(t *testing.T) {
workspace := t.TempDir()
outsideDir := t.TempDir()
outsideFile := filepath.Join(outsideDir, "allowed.txt")
os.WriteFile(outsideFile, []byte("outside content"), 0o644)
// Pattern allows access to the outsideDir.
patterns := []*regexp.Regexp{regexp.MustCompile(`^` + regexp.QuoteMeta(outsideDir))}
tool := NewReadFileTool(workspace, true, patterns)
// Read from whitelisted path should succeed.
result := tool.Execute(context.Background(), map[string]any{"path": outsideFile})
if result.IsError {
t.Errorf("expected whitelisted path to be readable, got: %s", result.ForLLM)
}
if !strings.Contains(result.ForLLM, "outside content") {
t.Errorf("expected file content, got: %s", result.ForLLM)
}
// Read from non-whitelisted path outside workspace should fail.
otherDir := t.TempDir()
otherFile := filepath.Join(otherDir, "blocked.txt")
os.WriteFile(otherFile, []byte("blocked"), 0o644)
result = tool.Execute(context.Background(), map[string]any{"path": otherFile})
if !result.IsError {
t.Errorf("expected non-whitelisted path to be blocked, got: %s", result.ForLLM)
}
}
+1 -1
View File
@@ -329,7 +329,7 @@ func TestToolRegistry_ConcurrentAccess(t *testing.T) {
r := NewToolRegistry()
var wg sync.WaitGroup
for i := 0; i < 50; i++ {
for i := range 50 {
wg.Add(1)
go func(n int) {
defer wg.Done()
+99 -53
View File
@@ -21,60 +21,85 @@ type ExecTool struct {
timeout time.Duration
denyPatterns []*regexp.Regexp
allowPatterns []*regexp.Regexp
customAllowPatterns []*regexp.Regexp
restrictToWorkspace bool
}
var defaultDenyPatterns = []*regexp.Regexp{
regexp.MustCompile(`\brm\s+-[rf]{1,2}\b`),
regexp.MustCompile(`\bdel\s+/[fq]\b`),
regexp.MustCompile(`\brmdir\s+/s\b`),
regexp.MustCompile(`\b(format|mkfs|diskpart)\b\s`), // Match disk wiping commands (must be followed by space/args)
regexp.MustCompile(`\bdd\s+if=`),
regexp.MustCompile(`>\s*/dev/sd[a-z]\b`), // Block writes to disk devices (but allow /dev/null)
regexp.MustCompile(`\b(shutdown|reboot|poweroff)\b`),
regexp.MustCompile(`:\(\)\s*\{.*\};\s*:`),
regexp.MustCompile(`\$\([^)]+\)`),
regexp.MustCompile(`\$\{[^}]+\}`),
regexp.MustCompile("`[^`]+`"),
regexp.MustCompile(`\|\s*sh\b`),
regexp.MustCompile(`\|\s*bash\b`),
regexp.MustCompile(`;\s*rm\s+-[rf]`),
regexp.MustCompile(`&&\s*rm\s+-[rf]`),
regexp.MustCompile(`\|\|\s*rm\s+-[rf]`),
regexp.MustCompile(`>\s*/dev/null\s*>&?\s*\d?`),
regexp.MustCompile(`<<\s*EOF`),
regexp.MustCompile(`\$\(\s*cat\s+`),
regexp.MustCompile(`\$\(\s*curl\s+`),
regexp.MustCompile(`\$\(\s*wget\s+`),
regexp.MustCompile(`\$\(\s*which\s+`),
regexp.MustCompile(`\bsudo\b`),
regexp.MustCompile(`\bchmod\s+[0-7]{3,4}\b`),
regexp.MustCompile(`\bchown\b`),
regexp.MustCompile(`\bpkill\b`),
regexp.MustCompile(`\bkillall\b`),
regexp.MustCompile(`\bkill\s+-[9]\b`),
regexp.MustCompile(`\bcurl\b.*\|\s*(sh|bash)`),
regexp.MustCompile(`\bwget\b.*\|\s*(sh|bash)`),
regexp.MustCompile(`\bnpm\s+install\s+-g\b`),
regexp.MustCompile(`\bpip\s+install\s+--user\b`),
regexp.MustCompile(`\bapt\s+(install|remove|purge)\b`),
regexp.MustCompile(`\byum\s+(install|remove)\b`),
regexp.MustCompile(`\bdnf\s+(install|remove)\b`),
regexp.MustCompile(`\bdocker\s+run\b`),
regexp.MustCompile(`\bdocker\s+exec\b`),
regexp.MustCompile(`\bgit\s+push\b`),
regexp.MustCompile(`\bgit\s+force\b`),
regexp.MustCompile(`\bssh\b.*@`),
regexp.MustCompile(`\beval\b`),
regexp.MustCompile(`\bsource\s+.*\.sh\b`),
}
var (
defaultDenyPatterns = []*regexp.Regexp{
regexp.MustCompile(`\brm\s+-[rf]{1,2}\b`),
regexp.MustCompile(`\bdel\s+/[fq]\b`),
regexp.MustCompile(`\brmdir\s+/s\b`),
// Match disk wiping commands (must be followed by space/args)
regexp.MustCompile(
`\b(format|mkfs|diskpart)\b\s`,
),
regexp.MustCompile(`\bdd\s+if=`),
// Block writes to block devices (all common naming schemes).
regexp.MustCompile(
`>\s*/dev/(sd[a-z]|hd[a-z]|vd[a-z]|xvd[a-z]|nvme\d|mmcblk\d|loop\d|dm-\d|md\d|sr\d|nbd\d)`,
),
regexp.MustCompile(`\b(shutdown|reboot|poweroff)\b`),
regexp.MustCompile(`:\(\)\s*\{.*\};\s*:`),
regexp.MustCompile(`\$\([^)]+\)`),
regexp.MustCompile(`\$\{[^}]+\}`),
regexp.MustCompile("`[^`]+`"),
regexp.MustCompile(`\|\s*sh\b`),
regexp.MustCompile(`\|\s*bash\b`),
regexp.MustCompile(`;\s*rm\s+-[rf]`),
regexp.MustCompile(`&&\s*rm\s+-[rf]`),
regexp.MustCompile(`\|\|\s*rm\s+-[rf]`),
regexp.MustCompile(`<<\s*EOF`),
regexp.MustCompile(`\$\(\s*cat\s+`),
regexp.MustCompile(`\$\(\s*curl\s+`),
regexp.MustCompile(`\$\(\s*wget\s+`),
regexp.MustCompile(`\$\(\s*which\s+`),
regexp.MustCompile(`\bsudo\b`),
regexp.MustCompile(`\bchmod\s+[0-7]{3,4}\b`),
regexp.MustCompile(`\bchown\b`),
regexp.MustCompile(`\bpkill\b`),
regexp.MustCompile(`\bkillall\b`),
regexp.MustCompile(`\bkill\s+-[9]\b`),
regexp.MustCompile(`\bcurl\b.*\|\s*(sh|bash)`),
regexp.MustCompile(`\bwget\b.*\|\s*(sh|bash)`),
regexp.MustCompile(`\bnpm\s+install\s+-g\b`),
regexp.MustCompile(`\bpip\s+install\s+--user\b`),
regexp.MustCompile(`\bapt\s+(install|remove|purge)\b`),
regexp.MustCompile(`\byum\s+(install|remove)\b`),
regexp.MustCompile(`\bdnf\s+(install|remove)\b`),
regexp.MustCompile(`\bdocker\s+run\b`),
regexp.MustCompile(`\bdocker\s+exec\b`),
regexp.MustCompile(`\bgit\s+push\b`),
regexp.MustCompile(`\bgit\s+force\b`),
regexp.MustCompile(`\bssh\b.*@`),
regexp.MustCompile(`\beval\b`),
regexp.MustCompile(`\bsource\s+.*\.sh\b`),
}
func NewExecTool(workingDir string, restrict bool) *ExecTool {
// absolutePathPattern matches absolute file paths in commands (Unix and Windows).
absolutePathPattern = regexp.MustCompile(`[A-Za-z]:\\[^\\\"']+|/[^\s\"']+`)
// safePaths are kernel pseudo-devices that are always safe to reference in
// commands, regardless of workspace restriction. They contain no user data
// and cannot cause destructive writes.
safePaths = map[string]bool{
"/dev/null": true,
"/dev/zero": true,
"/dev/random": true,
"/dev/urandom": true,
"/dev/stdin": true,
"/dev/stdout": true,
"/dev/stderr": true,
}
)
func NewExecTool(workingDir string, restrict bool) (*ExecTool, error) {
return NewExecToolWithConfig(workingDir, restrict, nil)
}
func NewExecToolWithConfig(workingDir string, restrict bool, config *config.Config) *ExecTool {
func NewExecToolWithConfig(workingDir string, restrict bool, config *config.Config) (*ExecTool, error) {
denyPatterns := make([]*regexp.Regexp, 0)
customAllowPatterns := make([]*regexp.Regexp, 0)
if config != nil {
execConfig := config.Tools.Exec
@@ -86,8 +111,7 @@ func NewExecToolWithConfig(workingDir string, restrict bool, config *config.Conf
for _, pattern := range execConfig.CustomDenyPatterns {
re, err := regexp.Compile(pattern)
if err != nil {
fmt.Printf("Invalid custom deny pattern %q: %v\n", pattern, err)
continue
return nil, fmt.Errorf("invalid custom deny pattern %q: %w", pattern, err)
}
denyPatterns = append(denyPatterns, re)
}
@@ -96,6 +120,13 @@ func NewExecToolWithConfig(workingDir string, restrict bool, config *config.Conf
// If deny patterns are disabled, we won't add any patterns, allowing all commands.
fmt.Println("Warning: deny patterns are disabled. All commands will be allowed.")
}
for _, pattern := range execConfig.CustomAllowPatterns {
re, err := regexp.Compile(pattern)
if err != nil {
return nil, fmt.Errorf("invalid custom allow pattern %q: %w", pattern, err)
}
customAllowPatterns = append(customAllowPatterns, re)
}
} else {
denyPatterns = append(denyPatterns, defaultDenyPatterns...)
}
@@ -105,8 +136,9 @@ func NewExecToolWithConfig(workingDir string, restrict bool, config *config.Conf
timeout: 60 * time.Second,
denyPatterns: denyPatterns,
allowPatterns: nil,
customAllowPatterns: customAllowPatterns,
restrictToWorkspace: restrict,
}
}, nil
}
func (t *ExecTool) Name() string {
@@ -259,9 +291,20 @@ func (t *ExecTool) guardCommand(command, cwd string) string {
cmd := strings.TrimSpace(command)
lower := strings.ToLower(cmd)
for _, pattern := range t.denyPatterns {
// Custom allow patterns exempt a command from deny checks.
explicitlyAllowed := false
for _, pattern := range t.customAllowPatterns {
if pattern.MatchString(lower) {
return "Command blocked by safety guard (dangerous pattern detected)"
explicitlyAllowed = true
break
}
}
if !explicitlyAllowed {
for _, pattern := range t.denyPatterns {
if pattern.MatchString(lower) {
return "Command blocked by safety guard (dangerous pattern detected)"
}
}
}
@@ -288,8 +331,7 @@ func (t *ExecTool) guardCommand(command, cwd string) string {
return ""
}
pathPattern := regexp.MustCompile(`[A-Za-z]:\\[^\\\"']+|/[^\s\"']+`)
matches := pathPattern.FindAllString(cmd, -1)
matches := absolutePathPattern.FindAllString(cmd, -1)
for _, raw := range matches {
p, err := filepath.Abs(raw)
@@ -297,6 +339,10 @@ func (t *ExecTool) guardCommand(command, cwd string) string {
continue
}
if safePaths[p] {
continue
}
rel, err := filepath.Rel(cwdPath, p)
if err != nil {
continue
+162 -11
View File
@@ -7,11 +7,16 @@ import (
"strings"
"testing"
"time"
"github.com/sipeed/picoclaw/pkg/config"
)
// TestShellTool_Success verifies successful command execution
func TestShellTool_Success(t *testing.T) {
tool := NewExecTool("", false)
tool, err := NewExecTool("", false)
if err != nil {
t.Errorf("unable to configure exec tool: %s", err)
}
ctx := context.Background()
args := map[string]any{
@@ -38,7 +43,10 @@ func TestShellTool_Success(t *testing.T) {
// TestShellTool_Failure verifies failed command execution
func TestShellTool_Failure(t *testing.T) {
tool := NewExecTool("", false)
tool, err := NewExecTool("", false)
if err != nil {
t.Errorf("unable to configure exec tool: %s", err)
}
ctx := context.Background()
args := map[string]any{
@@ -65,7 +73,11 @@ func TestShellTool_Failure(t *testing.T) {
// TestShellTool_Timeout verifies command timeout handling
func TestShellTool_Timeout(t *testing.T) {
tool := NewExecTool("", false)
tool, err := NewExecTool("", false)
if err != nil {
t.Errorf("unable to configure exec tool: %s", err)
}
tool.SetTimeout(100 * time.Millisecond)
ctx := context.Background()
@@ -93,7 +105,10 @@ func TestShellTool_WorkingDir(t *testing.T) {
testFile := filepath.Join(tmpDir, "test.txt")
os.WriteFile(testFile, []byte("test content"), 0o644)
tool := NewExecTool("", false)
tool, err := NewExecTool("", false)
if err != nil {
t.Errorf("unable to configure exec tool: %s", err)
}
ctx := context.Background()
args := map[string]any{
@@ -114,7 +129,10 @@ func TestShellTool_WorkingDir(t *testing.T) {
// TestShellTool_DangerousCommand verifies safety guard blocks dangerous commands
func TestShellTool_DangerousCommand(t *testing.T) {
tool := NewExecTool("", false)
tool, err := NewExecTool("", false)
if err != nil {
t.Errorf("unable to configure exec tool: %s", err)
}
ctx := context.Background()
args := map[string]any{
@@ -135,7 +153,10 @@ func TestShellTool_DangerousCommand(t *testing.T) {
// TestShellTool_MissingCommand verifies error handling for missing command
func TestShellTool_MissingCommand(t *testing.T) {
tool := NewExecTool("", false)
tool, err := NewExecTool("", false)
if err != nil {
t.Errorf("unable to configure exec tool: %s", err)
}
ctx := context.Background()
args := map[string]any{}
@@ -150,7 +171,10 @@ func TestShellTool_MissingCommand(t *testing.T) {
// TestShellTool_StderrCapture verifies stderr is captured and included
func TestShellTool_StderrCapture(t *testing.T) {
tool := NewExecTool("", false)
tool, err := NewExecTool("", false)
if err != nil {
t.Errorf("unable to configure exec tool: %s", err)
}
ctx := context.Background()
args := map[string]any{
@@ -170,7 +194,10 @@ func TestShellTool_StderrCapture(t *testing.T) {
// TestShellTool_OutputTruncation verifies long output is truncated
func TestShellTool_OutputTruncation(t *testing.T) {
tool := NewExecTool("", false)
tool, err := NewExecTool("", false)
if err != nil {
t.Errorf("unable to configure exec tool: %s", err)
}
ctx := context.Background()
// Generate long output (>10000 chars)
@@ -198,7 +225,11 @@ func TestShellTool_WorkingDir_OutsideWorkspace(t *testing.T) {
t.Fatalf("failed to create outside dir: %v", err)
}
tool := NewExecTool(workspace, true)
tool, err := NewExecTool(workspace, true)
if err != nil {
t.Errorf("unable to configure exec tool: %s", err)
}
result := tool.Execute(context.Background(), map[string]any{
"command": "pwd",
"working_dir": outsideDir,
@@ -232,7 +263,11 @@ func TestShellTool_WorkingDir_SymlinkEscape(t *testing.T) {
t.Skipf("symlinks not supported in this environment: %v", err)
}
tool := NewExecTool(workspace, true)
tool, err := NewExecTool(workspace, true)
if err != nil {
t.Errorf("unable to configure exec tool: %s", err)
}
result := tool.Execute(context.Background(), map[string]any{
"command": "cat secret.txt",
"working_dir": link,
@@ -249,7 +284,11 @@ func TestShellTool_WorkingDir_SymlinkEscape(t *testing.T) {
// TestShellTool_RestrictToWorkspace verifies workspace restriction
func TestShellTool_RestrictToWorkspace(t *testing.T) {
tmpDir := t.TempDir()
tool := NewExecTool(tmpDir, false)
tool, err := NewExecTool(tmpDir, false)
if err != nil {
t.Errorf("unable to configure exec tool: %s", err)
}
tool.SetRestrictToWorkspace(true)
ctx := context.Background()
@@ -272,3 +311,115 @@ func TestShellTool_RestrictToWorkspace(t *testing.T) {
)
}
}
// TestShellTool_DevNullAllowed verifies that /dev/null redirections are not blocked (issue #964).
func TestShellTool_DevNullAllowed(t *testing.T) {
tmpDir := t.TempDir()
tool, err := NewExecTool(tmpDir, true)
if err != nil {
t.Fatalf("unable to configure exec tool: %s", err)
}
commands := []string{
"echo hello 2>/dev/null",
"echo hello >/dev/null",
"echo hello > /dev/null",
"echo hello 2> /dev/null",
"echo hello >/dev/null 2>&1",
"find " + tmpDir + " -name '*.go' 2>/dev/null",
}
for _, cmd := range commands {
result := tool.Execute(context.Background(), map[string]any{"command": cmd})
if result.IsError && strings.Contains(result.ForLLM, "blocked") {
t.Errorf("command should not be blocked: %s\n error: %s", cmd, result.ForLLM)
}
}
}
// TestShellTool_BlockDevices verifies that writes to block devices are blocked (issue #965).
func TestShellTool_BlockDevices(t *testing.T) {
tool, err := NewExecTool("", false)
if err != nil {
t.Fatalf("unable to configure exec tool: %s", err)
}
blocked := []string{
"echo x > /dev/sda",
"echo x > /dev/hda",
"echo x > /dev/vda",
"echo x > /dev/xvda",
"echo x > /dev/nvme0n1",
"echo x > /dev/mmcblk0",
"echo x > /dev/loop0",
"echo x > /dev/dm-0",
"echo x > /dev/md0",
"echo x > /dev/sr0",
"echo x > /dev/nbd0",
}
for _, cmd := range blocked {
result := tool.Execute(context.Background(), map[string]any{"command": cmd})
if !result.IsError {
t.Errorf("expected block device write to be blocked: %s", cmd)
}
}
}
// TestShellTool_SafePathsInWorkspaceRestriction verifies that safe kernel pseudo-devices
// are allowed even when workspace restriction is active.
func TestShellTool_SafePathsInWorkspaceRestriction(t *testing.T) {
tmpDir := t.TempDir()
tool, err := NewExecTool(tmpDir, true)
if err != nil {
t.Fatalf("unable to configure exec tool: %s", err)
}
// These reference paths outside workspace but should be allowed via safePaths.
commands := []string{
"cat /dev/urandom | head -c 16 | od",
"echo test > /dev/null",
"dd if=/dev/zero bs=1 count=1",
}
for _, cmd := range commands {
result := tool.Execute(context.Background(), map[string]any{"command": cmd})
if result.IsError && strings.Contains(result.ForLLM, "path outside working dir") {
t.Errorf("safe path should not be blocked by workspace check: %s\n error: %s", cmd, result.ForLLM)
}
}
}
// TestShellTool_CustomAllowPatterns verifies that custom allow patterns exempt
// commands from deny pattern checks.
func TestShellTool_CustomAllowPatterns(t *testing.T) {
cfg := &config.Config{
Tools: config.ToolsConfig{
Exec: config.ExecConfig{
EnableDenyPatterns: true,
CustomAllowPatterns: []string{`\bgit\s+push\s+origin\b`},
},
},
}
tool, err := NewExecToolWithConfig("", false, cfg)
if err != nil {
t.Fatalf("unable to configure exec tool: %s", err)
}
// "git push origin main" should be allowed by custom allow pattern.
result := tool.Execute(context.Background(), map[string]any{
"command": "git push origin main",
})
if result.IsError && strings.Contains(result.ForLLM, "blocked") {
t.Errorf("custom allow pattern should exempt 'git push origin main', got: %s", result.ForLLM)
}
// "git push upstream main" should still be blocked (does not match allow pattern).
result = tool.Execute(context.Background(), map[string]any{
"command": "git push upstream main",
})
if !result.IsError {
t.Errorf("'git push upstream main' should still be blocked by deny pattern")
}
}
+5 -1
View File
@@ -22,7 +22,11 @@ func processExists(pid int) bool {
}
func TestShellTool_TimeoutKillsChildProcess(t *testing.T) {
tool := NewExecTool(t.TempDir(), false)
tool, err := NewExecTool(t.TempDir(), false)
if err != nil {
t.Errorf("unable to configure exec tool: %s", err)
}
tool.SetTimeout(500 * time.Millisecond)
args := map[string]any{
+82 -61
View File
@@ -4,6 +4,7 @@ import (
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
@@ -15,6 +16,14 @@ import (
const (
userAgent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36"
// HTTP client timeouts for web tool providers.
searchTimeout = 10 * time.Second // Brave, Tavily, DuckDuckGo
perplexityTimeout = 30 * time.Second // Perplexity (LLM-based, slower)
fetchTimeout = 60 * time.Second // WebFetchTool
defaultMaxChars = 50000
maxRedirects = 5
)
// Pre-compiled regexes for HTML text extraction
@@ -74,6 +83,7 @@ type SearchProvider interface {
type BraveSearchProvider struct {
apiKey string
proxy string
client *http.Client
}
func (p *BraveSearchProvider) Search(ctx context.Context, query string, count int) (string, error) {
@@ -88,11 +98,7 @@ func (p *BraveSearchProvider) Search(ctx context.Context, query string, count in
req.Header.Set("Accept", "application/json")
req.Header.Set("X-Subscription-Token", p.apiKey)
client, err := createHTTPClient(p.proxy, 10*time.Second)
if err != nil {
return "", fmt.Errorf("failed to create HTTP client: %w", err)
}
resp, err := client.Do(req)
resp, err := p.client.Do(req)
if err != nil {
return "", fmt.Errorf("request failed: %w", err)
}
@@ -143,6 +149,7 @@ type TavilySearchProvider struct {
apiKey string
baseURL string
proxy string
client *http.Client
}
func (p *TavilySearchProvider) Search(ctx context.Context, query string, count int) (string, error) {
@@ -174,11 +181,7 @@ func (p *TavilySearchProvider) Search(ctx context.Context, query string, count i
req.Header.Set("Content-Type", "application/json")
req.Header.Set("User-Agent", userAgent)
client, err := createHTTPClient(p.proxy, 10*time.Second)
if err != nil {
return "", fmt.Errorf("failed to create HTTP client: %w", err)
}
resp, err := client.Do(req)
resp, err := p.client.Do(req)
if err != nil {
return "", fmt.Errorf("request failed: %w", err)
}
@@ -226,7 +229,8 @@ func (p *TavilySearchProvider) Search(ctx context.Context, query string, count i
}
type DuckDuckGoSearchProvider struct {
proxy string
proxy string
client *http.Client
}
func (p *DuckDuckGoSearchProvider) Search(ctx context.Context, query string, count int) (string, error) {
@@ -239,11 +243,7 @@ func (p *DuckDuckGoSearchProvider) Search(ctx context.Context, query string, cou
req.Header.Set("User-Agent", userAgent)
client, err := createHTTPClient(p.proxy, 10*time.Second)
if err != nil {
return "", fmt.Errorf("failed to create HTTP client: %w", err)
}
resp, err := client.Do(req)
resp, err := p.client.Do(req)
if err != nil {
return "", fmt.Errorf("request failed: %w", err)
}
@@ -285,7 +285,7 @@ func (p *DuckDuckGoSearchProvider) extractResults(html string, count int, query
maxItems := min(len(matches), count)
for i := 0; i < maxItems; i++ {
for i := range maxItems {
urlStr := matches[i][1]
title := stripTags(matches[i][2])
title = strings.TrimSpace(title)
@@ -293,9 +293,9 @@ func (p *DuckDuckGoSearchProvider) extractResults(html string, count int, query
// URL decoding if needed
if strings.Contains(urlStr, "uddg=") {
if u, err := url.QueryUnescape(urlStr); err == nil {
idx := strings.Index(u, "uddg=")
if idx != -1 {
urlStr = u[idx+5:]
_, after, ok := strings.Cut(u, "uddg=")
if ok {
urlStr = after
}
}
}
@@ -322,6 +322,7 @@ func stripTags(content string) string {
type PerplexitySearchProvider struct {
apiKey string
proxy string
client *http.Client
}
func (p *PerplexitySearchProvider) Search(ctx context.Context, query string, count int) (string, error) {
@@ -356,11 +357,7 @@ func (p *PerplexitySearchProvider) Search(ctx context.Context, query string, cou
req.Header.Set("Authorization", "Bearer "+p.apiKey)
req.Header.Set("User-Agent", userAgent)
client, err := createHTTPClient(p.proxy, 30*time.Second)
if err != nil {
return "", fmt.Errorf("failed to create HTTP client: %w", err)
}
resp, err := client.Do(req)
resp, err := p.client.Do(req)
if err != nil {
return "", fmt.Errorf("request failed: %w", err)
}
@@ -415,43 +412,60 @@ type WebSearchToolOptions struct {
Proxy string
}
func NewWebSearchTool(opts WebSearchToolOptions) *WebSearchTool {
func NewWebSearchTool(opts WebSearchToolOptions) (*WebSearchTool, error) {
var provider SearchProvider
maxResults := 5
// Priority: Perplexity > Brave > Tavily > DuckDuckGo
if opts.PerplexityEnabled && opts.PerplexityAPIKey != "" {
provider = &PerplexitySearchProvider{apiKey: opts.PerplexityAPIKey, proxy: opts.Proxy}
client, err := createHTTPClient(opts.Proxy, perplexityTimeout)
if err != nil {
return nil, fmt.Errorf("failed to create HTTP client for Perplexity: %w", err)
}
provider = &PerplexitySearchProvider{apiKey: opts.PerplexityAPIKey, proxy: opts.Proxy, client: client}
if opts.PerplexityMaxResults > 0 {
maxResults = opts.PerplexityMaxResults
}
} else if opts.BraveEnabled && opts.BraveAPIKey != "" {
provider = &BraveSearchProvider{apiKey: opts.BraveAPIKey, proxy: opts.Proxy}
client, err := createHTTPClient(opts.Proxy, searchTimeout)
if err != nil {
return nil, fmt.Errorf("failed to create HTTP client for Brave: %w", err)
}
provider = &BraveSearchProvider{apiKey: opts.BraveAPIKey, proxy: opts.Proxy, client: client}
if opts.BraveMaxResults > 0 {
maxResults = opts.BraveMaxResults
}
} else if opts.TavilyEnabled && opts.TavilyAPIKey != "" {
client, err := createHTTPClient(opts.Proxy, searchTimeout)
if err != nil {
return nil, fmt.Errorf("failed to create HTTP client for Tavily: %w", err)
}
provider = &TavilySearchProvider{
apiKey: opts.TavilyAPIKey,
baseURL: opts.TavilyBaseURL,
proxy: opts.Proxy,
client: client,
}
if opts.TavilyMaxResults > 0 {
maxResults = opts.TavilyMaxResults
}
} else if opts.DuckDuckGoEnabled {
provider = &DuckDuckGoSearchProvider{proxy: opts.Proxy}
client, err := createHTTPClient(opts.Proxy, searchTimeout)
if err != nil {
return nil, fmt.Errorf("failed to create HTTP client for DuckDuckGo: %w", err)
}
provider = &DuckDuckGoSearchProvider{proxy: opts.Proxy, client: client}
if opts.DuckDuckGoMaxResults > 0 {
maxResults = opts.DuckDuckGoMaxResults
}
} else {
return nil
return nil, nil
}
return &WebSearchTool{
provider: provider,
maxResults: maxResults,
}
}, nil
}
func (t *WebSearchTool) Name() string {
@@ -506,27 +520,40 @@ func (t *WebSearchTool) Execute(ctx context.Context, args map[string]any) *ToolR
}
type WebFetchTool struct {
maxChars int
proxy string
maxChars int
proxy string
client *http.Client
fetchLimitBytes int64
}
func NewWebFetchTool(maxChars int) *WebFetchTool {
if maxChars <= 0 {
maxChars = 50000
}
return &WebFetchTool{
maxChars: maxChars,
}
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) *WebFetchTool {
func NewWebFetchToolWithProxy(maxChars int, proxy string, fetchLimitBytes int64) (*WebFetchTool, error) {
if maxChars <= 0 {
maxChars = 50000
maxChars = defaultMaxChars
}
client, err := createHTTPClient(proxy, fetchTimeout)
if err != nil {
return nil, fmt.Errorf("failed to create HTTP client for web fetch: %w", err)
}
client.CheckRedirect = func(req *http.Request, via []*http.Request) error {
if len(via) >= maxRedirects {
return fmt.Errorf("stopped after %d redirects", maxRedirects)
}
return nil
}
if fetchLimitBytes <= 0 {
fetchLimitBytes = 10 * 1024 * 1024 // Security Fallback
}
return &WebFetchTool{
maxChars: maxChars,
proxy: proxy,
}
maxChars: maxChars,
proxy: proxy,
client: client,
fetchLimitBytes: fetchLimitBytes,
}, nil
}
func (t *WebFetchTool) Name() string {
@@ -588,27 +615,21 @@ func (t *WebFetchTool) Execute(ctx context.Context, args map[string]any) *ToolRe
req.Header.Set("User-Agent", userAgent)
client, err := createHTTPClient(t.proxy, 60*time.Second)
if err != nil {
return ErrorResult(fmt.Sprintf("failed to create HTTP client: %v", err))
}
// Configure redirect handling
client.CheckRedirect = func(req *http.Request, via []*http.Request) error {
if len(via) >= 5 {
return fmt.Errorf("stopped after 5 redirects")
}
return nil
}
resp, err := client.Do(req)
resp, err := t.client.Do(req)
if err != nil {
return ErrorResult(fmt.Sprintf("request failed: %v", err))
}
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))
}
@@ -652,14 +673,14 @@ func (t *WebFetchTool) Execute(ctx context.Context, args map[string]any) *ToolRe
resultJSON, _ := json.MarshalIndent(result, "", " ")
return &ToolResult{
ForLLM: fmt.Sprintf(
ForLLM: string(resultJSON),
ForUser: fmt.Sprintf(
"Fetched %d bytes from %s (extractor: %s, truncated: %v)",
len(text),
urlStr,
extractor,
truncated,
),
ForUser: string(resultJSON),
}
}
+144 -35
View File
@@ -1,15 +1,21 @@
package tools
import (
"bytes"
"context"
"encoding/json"
"fmt"
"net/http"
"net/http/httptest"
"strings"
"testing"
"time"
"github.com/sipeed/picoclaw/pkg/logger"
)
const testFetchLimit = int64(10 * 1024 * 1024)
// TestWebTool_WebFetch_Success verifies successful URL fetching
func TestWebTool_WebFetch_Success(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -19,7 +25,11 @@ func TestWebTool_WebFetch_Success(t *testing.T) {
}))
defer server.Close()
tool := NewWebFetchTool(50000)
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,
@@ -32,14 +42,14 @@ func TestWebTool_WebFetch_Success(t *testing.T) {
t.Errorf("Expected success, got IsError=true: %s", result.ForLLM)
}
// ForUser should contain the fetched content
if !strings.Contains(result.ForUser, "Test Page") {
t.Errorf("Expected ForUser to contain 'Test Page', got: %s", result.ForUser)
// 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)
}
// ForLLM should contain summary
if !strings.Contains(result.ForLLM, "bytes") && !strings.Contains(result.ForLLM, "extractor") {
t.Errorf("Expected ForLLM to contain summary, 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)
}
}
@@ -55,7 +65,11 @@ func TestWebTool_WebFetch_JSON(t *testing.T) {
}))
defer server.Close()
tool := NewWebFetchTool(50000)
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,
@@ -68,15 +82,19 @@ func TestWebTool_WebFetch_JSON(t *testing.T) {
t.Errorf("Expected success, got IsError=true: %s", result.ForLLM)
}
// ForUser should contain formatted JSON
if !strings.Contains(result.ForUser, "key") && !strings.Contains(result.ForUser, "value") {
t.Errorf("Expected ForUser to contain JSON data, got: %s", result.ForUser)
// 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 := NewWebFetchTool(50000)
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",
@@ -97,7 +115,11 @@ func TestWebTool_WebFetch_InvalidURL(t *testing.T) {
// TestWebTool_WebFetch_UnsupportedScheme verifies error handling for non-http URLs
func TestWebTool_WebFetch_UnsupportedScheme(t *testing.T) {
tool := NewWebFetchTool(50000)
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",
@@ -118,7 +140,11 @@ func TestWebTool_WebFetch_UnsupportedScheme(t *testing.T) {
// TestWebTool_WebFetch_MissingURL verifies error handling for missing URL
func TestWebTool_WebFetch_MissingURL(t *testing.T) {
tool := NewWebFetchTool(50000)
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{}
@@ -146,7 +172,11 @@ func TestWebTool_WebFetch_Truncation(t *testing.T) {
}))
defer server.Close()
tool := NewWebFetchTool(1000) // Limit to 1000 chars
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,
@@ -159,9 +189,9 @@ func TestWebTool_WebFetch_Truncation(t *testing.T) {
t.Errorf("Expected success, got IsError=true: %s", result.ForLLM)
}
// ForUser should contain truncated content (not the full 20000 chars)
// ForLLM should contain truncated content (not the full 20000 chars)
resultMap := make(map[string]any)
json.Unmarshal([]byte(result.ForUser), &resultMap)
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))
@@ -174,15 +204,64 @@ 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(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 := NewWebSearchTool(WebSearchToolOptions{BraveEnabled: true, BraveAPIKey: ""})
tool, err := NewWebSearchTool(WebSearchToolOptions{BraveEnabled: true, BraveAPIKey: ""})
if err != nil {
t.Fatalf("Unexpected error: %v", err)
}
if tool != nil {
t.Errorf("Expected nil tool when Brave API key is empty")
}
// Also nil when nothing is enabled
tool = NewWebSearchTool(WebSearchToolOptions{})
tool, err = NewWebSearchTool(WebSearchToolOptions{})
if err != nil {
t.Fatalf("Unexpected error: %v", err)
}
if tool != nil {
t.Errorf("Expected nil tool when no provider is enabled")
}
@@ -190,7 +269,10 @@ func TestWebTool_WebSearch_NoApiKey(t *testing.T) {
// TestWebTool_WebSearch_MissingQuery verifies error handling for missing query
func TestWebTool_WebSearch_MissingQuery(t *testing.T) {
tool := NewWebSearchTool(WebSearchToolOptions{BraveEnabled: true, BraveAPIKey: "test-key", BraveMaxResults: 5})
tool, err := NewWebSearchTool(WebSearchToolOptions{BraveEnabled: true, BraveAPIKey: "test-key", BraveMaxResults: 5})
if err != nil {
t.Fatalf("Unexpected error: %v", err)
}
ctx := context.Background()
args := map[string]any{}
@@ -215,7 +297,11 @@ func TestWebTool_WebFetch_HTMLExtraction(t *testing.T) {
}))
defer server.Close()
tool := NewWebFetchTool(50000)
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,
@@ -228,14 +314,14 @@ func TestWebTool_WebFetch_HTMLExtraction(t *testing.T) {
t.Errorf("Expected success, got IsError=true: %s", result.ForLLM)
}
// ForUser should contain extracted text (without script/style tags)
if !strings.Contains(result.ForUser, "Title") && !strings.Contains(result.ForUser, "Content") {
t.Errorf("Expected ForUser to contain extracted text, got: %s", result.ForUser)
// 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
if strings.Contains(result.ForUser, "<script>") || strings.Contains(result.ForUser, "<style>") {
t.Errorf("Expected script/style tags to be removed, got: %s", result.ForUser)
// Should NOT contain script or style tags in ForLLM
if strings.Contains(result.ForLLM, "<script>") || strings.Contains(result.ForLLM, "<style>") {
t.Errorf("Expected script/style tags to be removed, got: %s", result.ForLLM)
}
}
@@ -316,7 +402,11 @@ func TestWebFetchTool_extractText(t *testing.T) {
// TestWebTool_WebFetch_MissingDomain verifies error handling for URL without domain
func TestWebTool_WebFetch_MissingDomain(t *testing.T) {
tool := NewWebFetchTool(50000)
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://",
@@ -438,15 +528,22 @@ func TestCreateHTTPClient_ProxyFromEnvironmentWhenConfigEmpty(t *testing.T) {
}
func TestNewWebFetchToolWithProxy(t *testing.T) {
tool := NewWebFetchToolWithProxy(1024, "http://127.0.0.1:7890")
if tool.maxChars != 1024 {
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 = NewWebFetchToolWithProxy(0, "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)
}
@@ -454,12 +551,15 @@ func TestNewWebFetchToolWithProxy(t *testing.T) {
func TestNewWebSearchTool_PropagatesProxy(t *testing.T) {
t.Run("perplexity", func(t *testing.T) {
tool := NewWebSearchTool(WebSearchToolOptions{
tool, err := NewWebSearchTool(WebSearchToolOptions{
PerplexityEnabled: true,
PerplexityAPIKey: "k",
PerplexityMaxResults: 3,
Proxy: "http://127.0.0.1:7890",
})
if err != nil {
t.Fatalf("NewWebSearchTool() error: %v", err)
}
p, ok := tool.provider.(*PerplexitySearchProvider)
if !ok {
t.Fatalf("provider type = %T, want *PerplexitySearchProvider", tool.provider)
@@ -470,12 +570,15 @@ func TestNewWebSearchTool_PropagatesProxy(t *testing.T) {
})
t.Run("brave", func(t *testing.T) {
tool := NewWebSearchTool(WebSearchToolOptions{
tool, err := NewWebSearchTool(WebSearchToolOptions{
BraveEnabled: true,
BraveAPIKey: "k",
BraveMaxResults: 3,
Proxy: "http://127.0.0.1:7890",
})
if err != nil {
t.Fatalf("NewWebSearchTool() error: %v", err)
}
p, ok := tool.provider.(*BraveSearchProvider)
if !ok {
t.Fatalf("provider type = %T, want *BraveSearchProvider", tool.provider)
@@ -486,11 +589,14 @@ func TestNewWebSearchTool_PropagatesProxy(t *testing.T) {
})
t.Run("duckduckgo", func(t *testing.T) {
tool := NewWebSearchTool(WebSearchToolOptions{
tool, err := NewWebSearchTool(WebSearchToolOptions{
DuckDuckGoEnabled: true,
DuckDuckGoMaxResults: 3,
Proxy: "http://127.0.0.1:7890",
})
if err != nil {
t.Fatalf("NewWebSearchTool() error: %v", err)
}
p, ok := tool.provider.(*DuckDuckGoSearchProvider)
if !ok {
t.Fatalf("provider type = %T, want *DuckDuckGoSearchProvider", tool.provider)
@@ -542,12 +648,15 @@ func TestWebTool_TavilySearch_Success(t *testing.T) {
}))
defer server.Close()
tool := NewWebSearchTool(WebSearchToolOptions{
tool, err := NewWebSearchTool(WebSearchToolOptions{
TavilyEnabled: true,
TavilyAPIKey: "test-key",
TavilyBaseURL: server.URL,
TavilyMaxResults: 5,
})
if err != nil {
t.Fatalf("NewWebSearchTool() error: %v", err)
}
ctx := context.Background()
args := map[string]any{