From daac4033d518bb61781c6fb27d7ffccafef89558 Mon Sep 17 00:00:00 2001 From: Kohei Date: Sun, 22 Feb 2026 22:11:53 +0900 Subject: [PATCH] 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 --- pkg/agent/loop.go | 43 ++++++++++++-- pkg/agent/media.go | 120 +++++++++++++++++++++++++++++++++++++++ pkg/tools/copy_file.go | 125 +++++++++++++++++++++++++++++++++++++++++ 3 files changed, 283 insertions(+), 5 deletions(-) create mode 100644 pkg/agent/media.go create mode 100644 pkg/tools/copy_file.go diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index 4a0459dcd..e7790446f 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -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) diff --git a/pkg/agent/media.go b/pkg/agent/media.go new file mode 100644 index 000000000..763a75c33 --- /dev/null +++ b/pkg/agent/media.go @@ -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:;base64, + 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:;base64, + 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: ] tags embedded in message content. +var imagePathRe = regexp.MustCompile(`\[Image: ([^\]]+)\]`) + +// CleanupMediaFiles extracts [Image: ] 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()}) + } + } + } +} diff --git a/pkg/tools/copy_file.go b/pkg/tools/copy_file.go new file mode 100644 index 000000000..625fe48b0 --- /dev/null +++ b/pkg/tools/copy_file.go @@ -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 +}