opencode/picoclaw
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

fix(tools): use magic-bytes MIME detection and add file size limit to send_file

Browse source 
- Replace hardcoded extension-to-MIME map with h2non/filetype (magic
  bytes) + mime.TypeByExtension fallback, consistent with the vision
  pipeline in resolveMediaRefs
- Add configurable max file size check (defaults to config.DefaultMaxMediaSize,
  20 MB) to prevent oversized uploads
- Add tests for magic-bytes detection, extension fallback, size limit,
  and default max size

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
shikihane 2026-03-04 10:25:31 +08:00
parent 8803e283e0
commit 46b87248b1
3 changed files with 116 additions and 46 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
pkg/agent/loop.go
View file

@ -171,6 +171,7 @@ func registerSharedTools(
sendFileTool := tools.NewSendFileTool(
agent.Workspace,
cfg.Agents.Defaults.RestrictToWorkspace,
cfg.Agents.Defaults.GetMaxMediaSize(),
nil,
)
agent.Tools.Register(sendFileTool)

56
pkg/tools/send_file.go
View file

@ -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"
}

105
pkg/tools/send_file_test.go
View file

@ -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)
}
}