From 3ed3fd9832158ff38f9df9af90bea7e363264f80 Mon Sep 17 00:00:00 2001 From: viettranx Date: Sat, 18 Apr 2026 15:34:46 +0700 Subject: [PATCH] fix(agent): extract MEDIA: paths from LLM response before sanitize MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When a tool returns its artifact via ForLLM "MEDIA:" and the LLM echoes that path in its final response (PDF from exec/weasyprint, TTS bridged through claude_cli MCP, etc.), sanitize stripped the token without populating rs.mediaResults — attachments never reached the channel. Add extractMediaFromContent() and call it right before SanitizeAssistantContent. Tight security: workspace must be set, relative paths resolve against workspace, absolute paths must stay inside workspace (blocks ../ escapes and sibling-tenant access), file must exist and be a regular file. Dedup prevents duplicates when a tool already populated Result.Media. Closes #882 --- internal/agent/loop_finalize.go | 9 +++ internal/agent/loop_media.go | 70 +++++++++++++++++++++ internal/agent/loop_media_test.go | 101 ++++++++++++++++++++++++++++++ 3 files changed, 180 insertions(+) create mode 100644 internal/agent/loop_media_test.go diff --git a/internal/agent/loop_finalize.go b/internal/agent/loop_finalize.go index d15c1bcd..dae304b3 100644 --- a/internal/agent/loop_finalize.go +++ b/internal/agent/loop_finalize.go @@ -13,6 +13,7 @@ import ( "github.com/nextlevelbuilder/goclaw/internal/i18n" "github.com/nextlevelbuilder/goclaw/internal/providers" "github.com/nextlevelbuilder/goclaw/internal/store" + "github.com/nextlevelbuilder/goclaw/internal/tools" ) // isUserFilePopulated checks if USER.md has been filled with actual user data @@ -42,6 +43,14 @@ func (l *Loop) finalizeRun( hadBootstrap bool, toolTiming ToolTimingMap, ) *RunResult { + // Extract MEDIA: tokens the LLM echoed in its final response + // BEFORE sanitize strips them. Covers cases where a tool returned its + // artifact via the ForLLM MEDIA: prefix but the agent relayed it as plain + // text (e.g. PDF from exec/weasyprint, TTS mp3 paths the LLM quotes back). + if extracted := extractMediaFromContent(rs.finalContent, tools.ToolWorkspaceFromCtx(ctx)); len(extracted) > 0 { + rs.mediaResults = append(rs.mediaResults, extracted...) + } + // 5. Full sanitization pipeline (matching TS extractAssistantText + sanitizeUserFacingText) rs.finalContent = SanitizeAssistantContent(rs.finalContent) diff --git a/internal/agent/loop_media.go b/internal/agent/loop_media.go index ac5152ad..ab6e1c6c 100644 --- a/internal/agent/loop_media.go +++ b/internal/agent/loop_media.go @@ -1,6 +1,7 @@ package agent import ( + "os" "path/filepath" "strings" ) @@ -44,6 +45,75 @@ func parseMediaResult(toolOutput string) *MediaResult { } } +// extractMediaFromContent scans text for MEDIA: tokens the LLM may echo +// in its final response (e.g. when a tool returned the MEDIA: prefix as plain +// text instead of setting Result.Media). Relative paths are resolved against +// workspace. Called before sanitize strips the tokens so the attachments are +// still delivered. +// +// Security: only paths that (a) exist on disk and (b) resolve inside the +// workspace root are accepted. An LLM cannot inject attachments pointing at +// /etc/passwd, a sibling tenant's workspace, or a hallucinated path — the +// extractor silently drops them. When workspace is empty, only legacy absolute +// paths from tool outputs (via parseMediaResult's upstream flow) are trusted; +// LLM-echoed absolute paths without a workspace context are dropped. +func extractMediaFromContent(content, workspace string) []MediaResult { + if !strings.Contains(content, "MEDIA:") || workspace == "" { + return nil + } + matches := mediaPathPattern.FindAllString(content, -1) + if len(matches) == 0 { + return nil + } + wsRoot := "" + if abs, err := filepath.Abs(workspace); err == nil { + wsRoot = filepath.Clean(abs) + } + if wsRoot == "" { + return nil + } + results := make([]MediaResult, 0, len(matches)) + seen := make(map[string]struct{}, len(matches)) + for _, m := range matches { + path := strings.TrimSpace(strings.TrimPrefix(m, "MEDIA:")) + if path == "" { + continue + } + // Drop markdown/JSON trailing punctuation that would otherwise stick: + // ")", "]", "\"", "'", ",", ";", ".". + path = strings.TrimRight(path, `)]"',;.`) + if path == "" { + continue + } + // Resolve relative paths against workspace. + if !filepath.IsAbs(path) { + path = filepath.Join(wsRoot, path) + } + cleaned := filepath.Clean(path) + // Workspace containment check — blocks "../" escapes and absolute paths + // outside the agent's allowed area. + rel, err := filepath.Rel(wsRoot, cleaned) + if err != nil || rel == ".." || strings.HasPrefix(rel, ".."+string(filepath.Separator)) { + continue + } + // Existence + regular-file check prevents hallucinated paths and + // directory/symlink-device shenanigans from reaching the outbound queue. + info, err := os.Lstat(cleaned) + if err != nil || !info.Mode().IsRegular() { + continue + } + if _, dup := seen[cleaned]; dup { + continue + } + seen[cleaned] = struct{}{} + results = append(results, MediaResult{ + Path: cleaned, + ContentType: mimeFromExt(filepath.Ext(cleaned)), + }) + } + return results +} + // deduplicateMedia removes duplicate media results by path, keeping the first occurrence. func deduplicateMedia(media []MediaResult) []MediaResult { if len(media) <= 1 { diff --git a/internal/agent/loop_media_test.go b/internal/agent/loop_media_test.go new file mode 100644 index 00000000..4d4d65eb --- /dev/null +++ b/internal/agent/loop_media_test.go @@ -0,0 +1,101 @@ +package agent + +import ( + "os" + "path/filepath" + "testing" +) + +// writeTempFile drops a zero-byte file at workspace/relPath, creating dirs. +func writeTempFile(t *testing.T, workspace, relPath string) string { + t.Helper() + full := filepath.Join(workspace, relPath) + if err := os.MkdirAll(filepath.Dir(full), 0o755); err != nil { + t.Fatalf("mkdir: %v", err) + } + if err := os.WriteFile(full, nil, 0o644); err != nil { + t.Fatalf("write: %v", err) + } + return full +} + +func TestExtractMediaFromContent(t *testing.T) { + ws := t.TempDir() + reportPath := writeTempFile(t, ws, "deliver/report.pdf") + audioA := writeTempFile(t, ws, "a.mp3") + audioB := writeTempFile(t, ws, "b.mp3") + chartPath := writeTempFile(t, ws, "charts/q4.png") + + // Outside-workspace file: should be rejected by containment check. + outsideDir := t.TempDir() + outsidePath := filepath.Join(outsideDir, "leak.pdf") + if err := os.WriteFile(outsidePath, nil, 0o644); err != nil { + t.Fatal(err) + } + + tests := []struct { + name string + content string + workspace string + wantPaths []string + }{ + { + name: "empty content", + content: "", + }, + { + name: "no media prefix", + content: "Just a regular response with no attachments.", + }, + { + name: "relative path resolved + exists", + content: "MEDIA:deliver/report.pdf", + workspace: ws, + wantPaths: []string{reportPath}, + }, + { + name: "multiple tokens deduped", + content: "First: MEDIA:a.mp3\nSecond: MEDIA:b.mp3\nAgain: MEDIA:a.mp3", + workspace: ws, + wantPaths: []string{audioA, audioB}, + }, + { + name: "markdown wrapped and punctuation stripped", + content: `![chart](MEDIA:charts/q4.png). See "MEDIA:deliver/report.pdf".`, + workspace: ws, + wantPaths: []string{chartPath, reportPath}, + }, + { + name: "hallucinated path dropped (file missing)", + content: "MEDIA:not-real.pdf", + workspace: ws, + }, + { + name: "path traversal escape blocked", + content: "MEDIA:../leak.pdf", + workspace: ws, + }, + { + name: "absolute path outside workspace blocked", + content: "MEDIA:" + outsidePath, + workspace: ws, + }, + { + name: "absolute path with no workspace dropped", + content: "MEDIA:" + reportPath, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := extractMediaFromContent(tt.content, tt.workspace) + if len(got) != len(tt.wantPaths) { + t.Fatalf("count = %d, want %d; got=%+v", len(got), len(tt.wantPaths), got) + } + for i, want := range tt.wantPaths { + if got[i].Path != want { + t.Errorf("path[%d] = %q, want %q", i, got[i].Path, want) + } + } + }) + } +}