fix(agent): extract MEDIA: paths from LLM response before sanitize

When a tool returns its artifact via ForLLM "MEDIA:<path>" 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
This commit is contained in:
viettranx committed 2026-04-18 15:34:46 +07:00
1 parent 226aa8ed11
commit 3ed3fd9832
3 files changed
+180

No files matched your search

+9
View File
@@ -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:<path> 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)
+70
View File
@@ -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:<path> 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 {
+101
View File
@@ -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)
}
}
})
}
}