feat: persist media images to disk and add copy_file tool
Save user/tool media (camera, screenshots) as files in data/media/ and preserve base64 in session via AddFullMessage so images survive across turns. Add copy_file tool to let the LLM copy media files into the workspace. Clean up media files when messages are dropped by forceCompression or summarization. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
parent
a3a37c0c67
commit
daac4033d5
3 changed files with 283 additions and 5 deletions
|
|
@ -50,6 +50,7 @@ type AgentLoop struct {
|
|||
mcpManager *mcp.Manager
|
||||
activeProcs map[string]*activeProcess
|
||||
procsMu sync.Mutex
|
||||
mediaDir string
|
||||
}
|
||||
|
||||
type activeProcess struct {
|
||||
|
|
@ -74,7 +75,7 @@ type processOptions struct {
|
|||
|
||||
// createToolRegistry creates a tool registry with common tools.
|
||||
// This is shared between main agent and subagents.
|
||||
func createToolRegistry(workspace string, restrict bool, cfg *config.Config, msgBus *bus.MessageBus) *tools.ToolRegistry {
|
||||
func createToolRegistry(workspace string, restrict bool, cfg *config.Config, msgBus *bus.MessageBus, dataDir string) *tools.ToolRegistry {
|
||||
registry := tools.NewToolRegistry()
|
||||
|
||||
// File system tools
|
||||
|
|
@ -84,6 +85,10 @@ func createToolRegistry(workspace string, restrict bool, cfg *config.Config, msg
|
|||
registry.Register(tools.NewEditFileTool(workspace, restrict))
|
||||
registry.Register(tools.NewAppendFileTool(workspace, restrict))
|
||||
|
||||
// Copy file tool (allows copying from media dir to workspace)
|
||||
mediaDir := filepath.Join(dataDir, "media")
|
||||
registry.Register(tools.NewCopyFileTool(workspace, mediaDir, restrict))
|
||||
|
||||
// Shell execution (disabled by default for security)
|
||||
if cfg.Tools.Exec.Enabled {
|
||||
registry.Register(tools.NewExecTool(workspace, restrict))
|
||||
|
|
@ -148,12 +153,16 @@ func NewAgentLoop(cfg *config.Config, msgBus *bus.MessageBus, provider providers
|
|||
|
||||
restrict := cfg.Agents.Defaults.RestrictToWorkspace
|
||||
|
||||
// Create media directory for persisting images
|
||||
mediaDir := filepath.Join(dataDir, "media")
|
||||
os.MkdirAll(mediaDir, 0755)
|
||||
|
||||
// Create tool registry for main agent
|
||||
toolsRegistry := createToolRegistry(workspace, restrict, cfg, msgBus)
|
||||
toolsRegistry := createToolRegistry(workspace, restrict, cfg, msgBus, dataDir)
|
||||
|
||||
// Create subagent manager with its own tool registry
|
||||
subagentManager := tools.NewSubagentManager(provider, cfg.Agents.Defaults.Model, workspace, msgBus)
|
||||
subagentTools := createToolRegistry(workspace, restrict, cfg, msgBus)
|
||||
subagentTools := createToolRegistry(workspace, restrict, cfg, msgBus, dataDir)
|
||||
// Subagent doesn't need spawn/subagent tools to avoid recursion
|
||||
subagentManager.SetTools(subagentTools)
|
||||
|
||||
|
|
@ -225,6 +234,7 @@ func NewAgentLoop(cfg *config.Config, msgBus *bus.MessageBus, provider providers
|
|||
rateLimiter: newRateLimiter(cfg.RateLimits.MaxToolCallsPerMinute, cfg.RateLimits.MaxRequestsPerMinute),
|
||||
mcpManager: mcpManager,
|
||||
activeProcs: make(map[string]*activeProcess),
|
||||
mediaDir: mediaDir,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -541,8 +551,19 @@ func (al *AgentLoop) runAgentLoop(ctx context.Context, opts processOptions) (str
|
|||
opts.InputMode,
|
||||
)
|
||||
|
||||
// 3. Save user message to session
|
||||
al.sessions.AddMessage(opts.SessionKey, "user", opts.UserMessage)
|
||||
// 3. Save user message to session (with media if present)
|
||||
userContent := opts.UserMessage
|
||||
if len(opts.Media) > 0 {
|
||||
paths := PersistMedia(opts.Media, al.mediaDir)
|
||||
for _, p := range paths {
|
||||
userContent += fmt.Sprintf("\n[Image: %s]", p)
|
||||
}
|
||||
}
|
||||
al.sessions.AddFullMessage(opts.SessionKey, providers.Message{
|
||||
Role: "user",
|
||||
Content: userContent,
|
||||
Media: opts.Media,
|
||||
})
|
||||
|
||||
// 4. Emit thinking status
|
||||
if !constants.IsInternalChannel(opts.Channel) {
|
||||
|
|
@ -955,6 +976,14 @@ func (al *AgentLoop) runLLMIteration(ctx context.Context, messages []providers.M
|
|||
contentForLLM = toolResult.Err.Error()
|
||||
}
|
||||
|
||||
// Persist media files from tool results (e.g. screenshots)
|
||||
if len(toolResult.Media) > 0 {
|
||||
paths := PersistMedia(toolResult.Media, al.mediaDir)
|
||||
for _, p := range paths {
|
||||
contentForLLM += fmt.Sprintf("\n[Image: %s]", p)
|
||||
}
|
||||
}
|
||||
|
||||
toolResultMsg := providers.Message{
|
||||
Role: "tool",
|
||||
Content: contentForLLM,
|
||||
|
|
@ -1076,6 +1105,8 @@ func (al *AgentLoop) forceCompression(sessionKey string) {
|
|||
}
|
||||
|
||||
droppedCount := mid
|
||||
// Clean up media files from dropped messages
|
||||
CleanupMediaFiles(conversation[:mid])
|
||||
keptConversation := conversation[mid:]
|
||||
|
||||
newHistory := make([]providers.Message, 0)
|
||||
|
|
@ -1245,6 +1276,8 @@ func (al *AgentLoop) summarizeSession(sessionKey string) {
|
|||
}
|
||||
|
||||
if finalSummary != "" {
|
||||
// Clean up media files from messages being summarized
|
||||
CleanupMediaFiles(toSummarize)
|
||||
al.sessions.SetSummary(sessionKey, finalSummary)
|
||||
al.sessions.TruncateHistory(sessionKey, 4)
|
||||
al.sessions.Save(sessionKey)
|
||||
|
|
|
|||
120
pkg/agent/media.go
Normal file
120
pkg/agent/media.go
Normal file
|
|
@ -0,0 +1,120 @@
|
|||
package agent
|
||||
|
||||
import (
|
||||
"encoding/base64"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/logger"
|
||||
"github.com/sipeed/picoclaw/pkg/providers"
|
||||
)
|
||||
|
||||
// PersistMedia saves base64 data URL images to the mediaDir as files.
|
||||
// It returns the list of saved file paths. Items that are not data URLs
|
||||
// (e.g. already file paths) are skipped.
|
||||
func PersistMedia(media []string, mediaDir string) []string {
|
||||
if len(media) == 0 || mediaDir == "" {
|
||||
return nil
|
||||
}
|
||||
|
||||
var paths []string
|
||||
ts := time.Now().Format("20060102_150405")
|
||||
|
||||
for i, item := range media {
|
||||
if !strings.HasPrefix(item, "data:") {
|
||||
continue
|
||||
}
|
||||
|
||||
// Parse data URL: data:<mime>;base64,<data>
|
||||
ext, data, err := parseDataURL(item)
|
||||
if err != nil {
|
||||
logger.WarnCF("media", "Failed to parse data URL",
|
||||
map[string]interface{}{"index": i, "error": err.Error()})
|
||||
continue
|
||||
}
|
||||
|
||||
filename := fmt.Sprintf("%s_%d%s", ts, i, ext)
|
||||
filePath := filepath.Join(mediaDir, filename)
|
||||
|
||||
if err := os.WriteFile(filePath, data, 0644); err != nil {
|
||||
logger.WarnCF("media", "Failed to write media file",
|
||||
map[string]interface{}{"path": filePath, "error": err.Error()})
|
||||
continue
|
||||
}
|
||||
|
||||
paths = append(paths, filePath)
|
||||
}
|
||||
|
||||
return paths
|
||||
}
|
||||
|
||||
// parseDataURL extracts extension and decoded bytes from a data URL.
|
||||
func parseDataURL(dataURL string) (ext string, data []byte, err error) {
|
||||
// Expected format: data:<mime>;base64,<base64data>
|
||||
if !strings.HasPrefix(dataURL, "data:") {
|
||||
return "", nil, fmt.Errorf("not a data URL")
|
||||
}
|
||||
|
||||
commaIdx := strings.Index(dataURL, ",")
|
||||
if commaIdx < 0 {
|
||||
return "", nil, fmt.Errorf("invalid data URL: no comma separator")
|
||||
}
|
||||
|
||||
header := dataURL[5:commaIdx] // after "data:"
|
||||
encoded := dataURL[commaIdx+1:]
|
||||
|
||||
// Extract MIME type (before ;base64)
|
||||
mime := header
|
||||
if idx := strings.Index(header, ";"); idx >= 0 {
|
||||
mime = header[:idx]
|
||||
}
|
||||
|
||||
ext = mimeToExt(mime)
|
||||
|
||||
data, err = base64.StdEncoding.DecodeString(encoded)
|
||||
if err != nil {
|
||||
return "", nil, fmt.Errorf("base64 decode failed: %w", err)
|
||||
}
|
||||
|
||||
return ext, data, nil
|
||||
}
|
||||
|
||||
// mimeToExt maps common MIME types to file extensions.
|
||||
func mimeToExt(mime string) string {
|
||||
switch strings.ToLower(mime) {
|
||||
case "image/jpeg":
|
||||
return ".jpg"
|
||||
case "image/png":
|
||||
return ".png"
|
||||
case "image/gif":
|
||||
return ".gif"
|
||||
case "image/webp":
|
||||
return ".webp"
|
||||
case "image/bmp":
|
||||
return ".bmp"
|
||||
default:
|
||||
return ".bin"
|
||||
}
|
||||
}
|
||||
|
||||
// imagePathRe matches [Image: <path>] tags embedded in message content.
|
||||
var imagePathRe = regexp.MustCompile(`\[Image: ([^\]]+)\]`)
|
||||
|
||||
// CleanupMediaFiles extracts [Image: <path>] references from messages
|
||||
// and deletes the corresponding files.
|
||||
func CleanupMediaFiles(messages []providers.Message) {
|
||||
for _, msg := range messages {
|
||||
matches := imagePathRe.FindAllStringSubmatch(msg.Content, -1)
|
||||
for _, m := range matches {
|
||||
path := m[1]
|
||||
if err := os.Remove(path); err != nil && !os.IsNotExist(err) {
|
||||
logger.WarnCF("media", "Failed to remove media file",
|
||||
map[string]interface{}{"path": path, "error": err.Error()})
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
125
pkg/tools/copy_file.go
Normal file
125
pkg/tools/copy_file.go
Normal file
|
|
@ -0,0 +1,125 @@
|
|||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
)
|
||||
|
||||
// CopyFileTool copies a file from source to destination.
|
||||
// Source is allowed from mediaDir or workspace (when restrict=true).
|
||||
// Destination must be within workspace.
|
||||
type CopyFileTool struct {
|
||||
workspace string
|
||||
mediaDir string
|
||||
restrict bool
|
||||
}
|
||||
|
||||
func NewCopyFileTool(workspace, mediaDir string, restrict bool) *CopyFileTool {
|
||||
return &CopyFileTool{
|
||||
workspace: workspace,
|
||||
mediaDir: mediaDir,
|
||||
restrict: restrict,
|
||||
}
|
||||
}
|
||||
|
||||
func (t *CopyFileTool) Name() string {
|
||||
return "copy_file"
|
||||
}
|
||||
|
||||
func (t *CopyFileTool) Description() string {
|
||||
return "Copy a file from source to destination. Source can be a media file (from camera/screenshot) or a workspace file. Destination must be within the workspace."
|
||||
}
|
||||
|
||||
func (t *CopyFileTool) Parameters() map[string]interface{} {
|
||||
return map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"source": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "Path to the source file to copy",
|
||||
},
|
||||
"destination": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "Path to the destination file (within workspace)",
|
||||
},
|
||||
},
|
||||
"required": []string{"source", "destination"},
|
||||
}
|
||||
}
|
||||
|
||||
func (t *CopyFileTool) Execute(ctx context.Context, args map[string]interface{}) *ToolResult {
|
||||
source, ok := args["source"].(string)
|
||||
if !ok {
|
||||
return ErrorResult("source is required")
|
||||
}
|
||||
|
||||
destination, ok := args["destination"].(string)
|
||||
if !ok {
|
||||
return ErrorResult("destination is required")
|
||||
}
|
||||
|
||||
// Validate source path
|
||||
srcPath, err := t.validateSource(source)
|
||||
if err != nil {
|
||||
return ErrorResult(err.Error())
|
||||
}
|
||||
|
||||
// Validate destination path (must be within workspace)
|
||||
dstPath, err := validatePath(destination, t.workspace, true)
|
||||
if err != nil {
|
||||
return ErrorResult(fmt.Sprintf("destination: %s", err.Error()))
|
||||
}
|
||||
|
||||
// Read source file
|
||||
data, err := os.ReadFile(srcPath)
|
||||
if err != nil {
|
||||
return ErrorResult(fmt.Sprintf("failed to read source file: %v", err))
|
||||
}
|
||||
|
||||
// Create destination directory if needed
|
||||
if err := os.MkdirAll(filepath.Dir(dstPath), 0755); err != nil {
|
||||
return ErrorResult(fmt.Sprintf("failed to create directory: %v", err))
|
||||
}
|
||||
|
||||
// Write destination file
|
||||
if err := os.WriteFile(dstPath, data, 0644); err != nil {
|
||||
return ErrorResult(fmt.Sprintf("failed to write destination file: %v", err))
|
||||
}
|
||||
|
||||
return SilentResult(fmt.Sprintf("File copied: %s → %s", source, destination))
|
||||
}
|
||||
|
||||
// validateSource checks that the source path is allowed.
|
||||
// When restrict=true, source must be within workspace or mediaDir.
|
||||
func (t *CopyFileTool) validateSource(source string) (string, error) {
|
||||
// First try workspace validation
|
||||
srcPath, err := validatePath(source, t.workspace, t.restrict)
|
||||
if err == nil {
|
||||
return srcPath, nil
|
||||
}
|
||||
|
||||
// If restrict mode and workspace validation failed, check mediaDir
|
||||
if t.restrict && t.mediaDir != "" {
|
||||
absMediaDir, merr := filepath.Abs(t.mediaDir)
|
||||
if merr != nil {
|
||||
return "", fmt.Errorf("access denied: path is outside the workspace")
|
||||
}
|
||||
|
||||
var absSource string
|
||||
if filepath.IsAbs(source) {
|
||||
absSource = filepath.Clean(source)
|
||||
} else {
|
||||
absSource = filepath.Clean(filepath.Join(absMediaDir, source))
|
||||
}
|
||||
|
||||
if isWithinWorkspace(absSource, absMediaDir) {
|
||||
return absSource, nil
|
||||
}
|
||||
|
||||
return "", fmt.Errorf("access denied: source path is outside workspace and media directory")
|
||||
}
|
||||
|
||||
return "", err
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue