diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index 40c1d411e..a5202f1b9 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -171,6 +171,7 @@ func registerSharedTools( sendFileTool := tools.NewSendFileTool( agent.Workspace, cfg.Agents.Defaults.RestrictToWorkspace, + cfg.Agents.Defaults.GetMaxMediaSize(), nil, ) agent.Tools.Register(sendFileTool) diff --git a/pkg/tools/send_file.go b/pkg/tools/send_file.go index 00dc5b21d..e54f86acc 100644 --- a/pkg/tools/send_file.go +++ b/pkg/tools/send_file.go @@ -3,29 +3,38 @@ package tools import ( "context" "fmt" + "mime" "os" "path/filepath" "strings" + "github.com/h2non/filetype" + + "github.com/sipeed/picoclaw/pkg/config" "github.com/sipeed/picoclaw/pkg/media" ) // SendFileTool allows the LLM to send a local file (image, document, etc.) // to the user on the current chat channel via the MediaStore pipeline. type SendFileTool struct { - workspace string - restrict bool - mediaStore media.MediaStore + workspace string + restrict bool + maxFileSize int + mediaStore media.MediaStore defaultChannel string defaultChatID string } -func NewSendFileTool(workspace string, restrict bool, store media.MediaStore) *SendFileTool { +func NewSendFileTool(workspace string, restrict bool, maxFileSize int, store media.MediaStore) *SendFileTool { + if maxFileSize <= 0 { + maxFileSize = config.DefaultMaxMediaSize + } return &SendFileTool{ - workspace: workspace, - restrict: restrict, - mediaStore: store, + workspace: workspace, + restrict: restrict, + maxFileSize: maxFileSize, + mediaStore: store, } } @@ -86,6 +95,12 @@ func (t *SendFileTool) Execute(_ context.Context, args map[string]any) *ToolResu if info.IsDir() { return ErrorResult("path is a directory, expected a file") } + if info.Size() > int64(t.maxFileSize) { + return ErrorResult(fmt.Sprintf( + "file too large: %d bytes (max %d bytes)", + info.Size(), t.maxFileSize, + )) + } filename, _ := args["filename"].(string) if filename == "" { @@ -107,19 +122,20 @@ func (t *SendFileTool) Execute(_ context.Context, args map[string]any) *ToolResu return MediaResult(fmt.Sprintf("File %q sent to user", filename), []string{ref}) } +// detectMediaType determines the MIME type of a file. +// Uses magic-bytes detection (h2non/filetype) first, then falls back to +// extension-based lookup via mime.TypeByExtension. func detectMediaType(path string) string { - switch strings.ToLower(filepath.Ext(path)) { - case ".jpg", ".jpeg": - return "image/jpeg" - case ".png": - return "image/png" - case ".gif": - return "image/gif" - case ".webp": - return "image/webp" - case ".pdf": - return "application/pdf" - default: - return "application/octet-stream" + kind, err := filetype.MatchFile(path) + if err == nil && kind != filetype.Unknown { + return kind.MIME.Value } + + if ext := filepath.Ext(path); ext != "" { + if t := mime.TypeByExtension(ext); t != "" { + return t + } + } + + return "application/octet-stream" } diff --git a/pkg/tools/send_file_test.go b/pkg/tools/send_file_test.go index 83825c1d9..08d129674 100644 --- a/pkg/tools/send_file_test.go +++ b/pkg/tools/send_file_test.go @@ -4,14 +4,16 @@ import ( "context" "os" "path/filepath" + "strings" "testing" + "github.com/sipeed/picoclaw/pkg/config" "github.com/sipeed/picoclaw/pkg/media" ) func TestSendFileTool_MissingPath(t *testing.T) { store := media.NewFileMediaStore() - tool := NewSendFileTool("/tmp", false, store) + tool := NewSendFileTool("/tmp", false, 0, store) tool.SetContext("feishu", "chat123") result := tool.Execute(context.Background(), map[string]any{}) @@ -22,7 +24,7 @@ func TestSendFileTool_MissingPath(t *testing.T) { func TestSendFileTool_NoContext(t *testing.T) { store := media.NewFileMediaStore() - tool := NewSendFileTool("/tmp", false, store) + tool := NewSendFileTool("/tmp", false, 0, store) // no SetContext call result := tool.Execute(context.Background(), map[string]any{"path": "/tmp/test.txt"}) @@ -32,7 +34,7 @@ func TestSendFileTool_NoContext(t *testing.T) { } func TestSendFileTool_NoMediaStore(t *testing.T) { - tool := NewSendFileTool("/tmp", false, nil) + tool := NewSendFileTool("/tmp", false, 0, nil) tool.SetContext("feishu", "chat123") result := tool.Execute(context.Background(), map[string]any{"path": "/tmp/test.txt"}) @@ -43,7 +45,7 @@ func TestSendFileTool_NoMediaStore(t *testing.T) { func TestSendFileTool_Directory(t *testing.T) { store := media.NewFileMediaStore() - tool := NewSendFileTool("/tmp", false, store) + tool := NewSendFileTool("/tmp", false, 0, store) tool.SetContext("feishu", "chat123") result := tool.Execute(context.Background(), map[string]any{"path": "/tmp"}) @@ -52,6 +54,34 @@ func TestSendFileTool_Directory(t *testing.T) { } } +func TestSendFileTool_FileTooLarge(t *testing.T) { + dir := t.TempDir() + testFile := filepath.Join(dir, "big.bin") + // Create a file larger than the limit + if err := os.WriteFile(testFile, make([]byte, 1024), 0o644); err != nil { + t.Fatal(err) + } + + store := media.NewFileMediaStore() + tool := NewSendFileTool(dir, false, 512, store) // 512 byte limit + tool.SetContext("feishu", "chat123") + + result := tool.Execute(context.Background(), map[string]any{"path": testFile}) + if !result.IsError { + t.Fatal("expected error for oversized file") + } + if !strings.Contains(result.ForLLM, "too large") { + t.Errorf("expected 'too large' in error, got %q", result.ForLLM) + } +} + +func TestSendFileTool_DefaultMaxSize(t *testing.T) { + tool := NewSendFileTool("/tmp", false, 0, nil) + if tool.maxFileSize != config.DefaultMaxMediaSize { + t.Errorf("expected default max size %d, got %d", config.DefaultMaxMediaSize, tool.maxFileSize) + } +} + func TestSendFileTool_Success(t *testing.T) { dir := t.TempDir() testFile := filepath.Join(dir, "photo.png") @@ -60,7 +90,7 @@ func TestSendFileTool_Success(t *testing.T) { } store := media.NewFileMediaStore() - tool := NewSendFileTool(dir, false, store) + tool := NewSendFileTool(dir, false, 0, store) tool.SetContext("feishu", "chat123") result := tool.Execute(context.Background(), map[string]any{"path": testFile}) @@ -83,7 +113,7 @@ func TestSendFileTool_CustomFilename(t *testing.T) { } store := media.NewFileMediaStore() - tool := NewSendFileTool(dir, false, store) + tool := NewSendFileTool(dir, false, 0, store) tool.SetContext("telegram", "chat456") result := tool.Execute(context.Background(), map[string]any{ @@ -98,26 +128,49 @@ func TestSendFileTool_CustomFilename(t *testing.T) { } } -func TestDetectMediaType(t *testing.T) { - tests := []struct { - path string - want string - }{ - {"photo.jpg", "image/jpeg"}, - {"photo.jpeg", "image/jpeg"}, - {"photo.png", "image/png"}, - {"anim.gif", "image/gif"}, - {"photo.webp", "image/webp"}, - {"doc.pdf", "application/pdf"}, - {"data.bin", "application/octet-stream"}, - {"noext", "application/octet-stream"}, +func TestDetectMediaType_MagicBytes(t *testing.T) { + dir := t.TempDir() + + // Minimal valid PNG header + pngHeader := []byte{0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A} + pngFile := filepath.Join(dir, "image.dat") // wrong extension, but valid PNG bytes + if err := os.WriteFile(pngFile, pngHeader, 0o644); err != nil { + t.Fatal(err) } - for _, tt := range tests { - t.Run(tt.path, func(t *testing.T) { - got := detectMediaType(tt.path) - if got != tt.want { - t.Errorf("detectMediaType(%q) = %q, want %q", tt.path, got, tt.want) - } - }) + + got := detectMediaType(pngFile) + if got != "image/png" { + t.Errorf("expected image/png from magic bytes, got %q", got) + } +} + +func TestDetectMediaType_FallbackToExtension(t *testing.T) { + dir := t.TempDir() + + // File with unrecognizable content but known extension + txtFile := filepath.Join(dir, "readme.txt") + if err := os.WriteFile(txtFile, []byte("hello world"), 0o644); err != nil { + t.Fatal(err) + } + + got := detectMediaType(txtFile) + // text/plain or similar — just verify it's not application/octet-stream + if got == "application/octet-stream" { + t.Errorf("expected extension-based MIME for .txt, got %q", got) + } +} + +func TestDetectMediaType_UnknownFallsToOctetStream(t *testing.T) { + dir := t.TempDir() + + // File with no extension and random bytes + unknownFile := filepath.Join(dir, "mystery") + if err := os.WriteFile(unknownFile, []byte{0x00, 0x01, 0x02}, 0o644); err != nil { + t.Fatal(err) + } + + got := detectMediaType(unknownFile) + if got != "application/octet-stream" { + t.Errorf("expected application/octet-stream, got %q", got) } }