fix(agent): use forward-scan replaceFirstMediaTag for correct multi-ref ordering

Replace replaceLastMediaTag with replaceFirstMediaTag across all 5 media enrichment functions. Forward iteration + first-match produces natural positional pairing, fixing reversed tag alignment when multiple media refs exist in one message. Also fixes same latent bug in enrichDocumentPaths, enrichAudioIDs, and enrichVideoIDs. Supersedes #608.
2026-10-14 20:17:20 +00:00 · 2026-03-31 19:26:40 +07:00
1 parent 49f51da81c
commit 1d39626a88
@@ -151,7 +151,6 @@ func (l *Loop) enrichDocumentPaths(messages []providers.Message, refs []provider
 	if len(messages) == 0 {
 		return
 	}
-	// Find last user message
 	lastIdx := -1
 	for i := len(messages) - 1; i >= 0; i-- {
 		if messages[i].Role == "user" {
@@ -168,7 +167,6 @@ func (l *Loop) enrichDocumentPaths(messages []providers.Message, refs []provider
 		if ref.Kind != "document" {
 			continue
 		}
-		// Use persisted workspace path; fall back to legacy .media/ lookup.
 		p := ref.Path
 		if p == "" && l.mediaStore != nil {
 			var err error
@@ -180,41 +178,20 @@ func (l *Loop) enrichDocumentPaths(messages []providers.Message, refs []provider
 		if p == "" {
 			continue
 		}
-		// Replace <media:document> or <media:document name="X"> with version that includes path.
-		// The hint tells the agent the file is directly accessible (no copy needed).
 		pathAttr := fmt.Sprintf(" path=%q", p)
-		old1 := "<media:document>"
-		new1 := "<media:document" + pathAttr + ">"
-		// Replace the LAST bare tag (current message, not group history).
-		if idx := strings.LastIndex(content, old1); idx >= 0 {
-			content = content[:idx] + new1 + content[idx+len(old1):]
-			continue
-		}
-		// For named variant, inject path attribute (last occurrence)
-		if idx := strings.LastIndex(content, "<media:document name="); idx >= 0 {
-			closeIdx := strings.Index(content[idx:], ">")
-			if closeIdx >= 0 {
-				tag := content[idx : idx+closeIdx]
-				content = content[:idx] + tag + pathAttr + ">" + content[idx+closeIdx+1:]
-			}
-		}
-		// For Slack variant with file= attribute (last occurrence)
-		if idx := strings.LastIndex(content, "<media:document file="); idx >= 0 {
-			closeIdx := strings.Index(content[idx:], ">")
-			if closeIdx >= 0 {
-				tag := content[idx : idx+closeIdx]
-				content = content[:idx] + tag + pathAttr + ">" + content[idx+closeIdx+1:]
-			}
-		}
+
+		// Match first <media:document> without a path — covers bare, named, and file= variants.
+		content, _ = replaceFirstMediaTag(content, "<media:document", func(tag string) bool {
+			return !tagHasAttr(tag, "path")
+		}, func(tag string) string {
+			return appendTagAttrs(tag, pathAttr)
+		})
 	}
 	messages[lastIdx].Content = content
 }

 // enrichAudioIDs updates the last user message to embed persisted media IDs
 // in <media:audio> and <media:voice> tags so the LLM can reference them.
-// Without this, the LLM sees plain <media:audio> and cannot pass a valid media_id.
-// Replaces the LAST bare tag (current message) rather than the first (which may be
-// in group history context), so the current turn's media gets the correct ID.
 func (l *Loop) enrichAudioIDs(messages []providers.Message, refs []providers.MediaRef) {
 	if len(messages) == 0 {
 		return
@@ -237,27 +214,27 @@ func (l *Loop) enrichAudioIDs(messages []providers.Message, refs []providers.Med
 		}
 		idAttr := fmt.Sprintf(" id=%q", ref.ID)

-		// Replace the LAST bare <media:audio> with <media:audio id="uuid">
-		bare := "<media:audio>"
-		if idx := strings.LastIndex(content, bare); idx >= 0 {
-			content = content[:idx] + "<media:audio" + idAttr + ">" + content[idx+len(bare):]
-			continue
-		}
-		// Replace the LAST bare <media:voice> with <media:voice id="uuid">
-		bareVoice := "<media:voice>"
-		if idx := strings.LastIndex(content, bareVoice); idx >= 0 {
-			content = content[:idx] + "<media:voice" + idAttr + ">" + content[idx+len(bareVoice):]
+		var replaced bool
+		content, replaced = replaceFirstMediaTag(content, "<media:audio", func(tag string) bool {
+			return !tagHasAttr(tag, "id")
+		}, func(tag string) string {
+			return appendTagAttrs(tag, idAttr)
+		})
+		if replaced {
 			continue
 		}
+
+		content, _ = replaceFirstMediaTag(content, "<media:voice", func(tag string) bool {
+			return !tagHasAttr(tag, "id")
+		}, func(tag string) string {
+			return appendTagAttrs(tag, idAttr)
+		})
 	}
 	messages[lastIdx].Content = content
 }

 // enrichVideoIDs updates the last user message to embed persisted media IDs
 // in <media:video> tags so the LLM can reference them via read_video tool.
-// Without this, the LLM sees plain <media:video> and hallucinates a media_id.
-// Replaces the LAST bare tag (current message) rather than the first (which may be
-// in group history context), so the current turn's media gets the correct ID.
 func (l *Loop) enrichVideoIDs(messages []providers.Message, refs []providers.MediaRef) {
 	if len(messages) == 0 {
 		return
@@ -280,12 +257,11 @@ func (l *Loop) enrichVideoIDs(messages []providers.Message, refs []providers.Med
 		}
 		idAttr := fmt.Sprintf(" id=%q", ref.ID)

-		// Replace the LAST bare <media:video> with <media:video id="uuid">
-		bare := "<media:video>"
-		if idx := strings.LastIndex(content, bare); idx >= 0 {
-			content = content[:idx] + "<media:video" + idAttr + ">" + content[idx+len(bare):]
-			continue
-		}
+		content, _ = replaceFirstMediaTag(content, "<media:video", func(tag string) bool {
+			return !tagHasAttr(tag, "id")
+		}, func(tag string) string {
+			return appendTagAttrs(tag, idAttr)
+		})
 	}
 	messages[lastIdx].Content = content
 }
@@ -295,8 +271,6 @@ func (l *Loop) enrichVideoIDs(messages []providers.Message, refs []providers.Med
 // and stored. The path attribute allows tools called via MCP bridge (e.g.
 // claude-cli) to access images via read_image(path=...) even though the
 // bridge context does not carry WithMediaImages.
-// Iterates refs in reverse order so that when multiple images are present,
-// each ref maps to the correct positional tag (last ref → last tag, etc.).
 func (l *Loop) enrichImageIDs(messages []providers.Message, refs []providers.MediaRef) {
 	if len(messages) == 0 {
 		return
@@ -313,9 +287,7 @@ func (l *Loop) enrichImageIDs(messages []providers.Message, refs []providers.Med
 	}

 	content := messages[lastIdx].Content
-	// Iterate in reverse so first ref matches first tag when using LastIndex replacements.
-	for i := len(refs) - 1; i >= 0; i-- {
-		ref := refs[i]
+	for _, ref := range refs {
 		if ref.Kind != "image" {
 			continue
 		}
@@ -325,8 +297,7 @@ func (l *Loop) enrichImageIDs(messages []providers.Message, refs []providers.Med
 			pathAttr = fmt.Sprintf(" path=%q", ref.Path)
 		}

-		var replaced bool
-		content, replaced = replaceLastMediaTag(content, "<media:image", func(tag string) bool {
+		content, _ = replaceFirstMediaTag(content, "<media:image", func(tag string) bool {
 			return !tagHasAttr(tag, "id")
 		}, func(tag string) string {
 			attrs := []string{idAttr}
@@ -335,9 +306,6 @@ func (l *Loop) enrichImageIDs(messages []providers.Message, refs []providers.Med
 			}
 			return appendTagAttrs(tag, attrs...)
 		})
-		if replaced {
-			continue
-		}
 	}
 	messages[lastIdx].Content = content
 }
@@ -374,9 +342,9 @@ func (l *Loop) enrichImagePaths(messages []providers.Message) {
 			}
 			pathAttr := fmt.Sprintf(" path=%q", p)

-			// Prefer tags that already carry the matching media ID, regardless of attribute order.
+			// Prefer tags that already carry the matching media ID.
 			var replaced bool
-			content, replaced = replaceLastMediaTag(content, "<media:image", func(tag string) bool {
+			content, replaced = replaceFirstMediaTag(content, "<media:image", func(tag string) bool {
 				return tagHasAttrValue(tag, "id", ref.ID) && !tagHasAttr(tag, "path")
 			}, func(tag string) string {
 				return appendTagAttrs(tag, pathAttr)
@@ -386,8 +354,8 @@ func (l *Loop) enrichImagePaths(messages []providers.Message) {
 				continue
 			}

-			// Fallback to the last image tag without an ID and attach both id and path.
-			content, replaced = replaceLastMediaTag(content, "<media:image", func(tag string) bool {
+			// Fallback: first image tag without an ID — attach both id and path.
+			content, replaced = replaceFirstMediaTag(content, "<media:image", func(tag string) bool {
 				return !tagHasAttr(tag, "id")
 			}, func(tag string) string {
 				return appendTagAttrs(tag, fmt.Sprintf(` id=%q`, ref.ID), pathAttr)
@@ -417,13 +385,17 @@ func mediaKindFromMime(mime string) string {
 	}
 }

-func replaceLastMediaTag(content, prefix string, match func(tag string) bool, replace func(tag string) string) (string, bool) {
-	searchFrom := len(content)
-	for searchFrom > 0 {
-		idx := strings.LastIndex(content[:searchFrom], prefix)
+// replaceFirstMediaTag finds the first tag in content starting with prefix
+// whose full text satisfies match, and replaces it using replace.
+// Forward scanning ensures natural positional pairing when iterating refs in order.
+func replaceFirstMediaTag(content, prefix string, match func(tag string) bool, replace func(tag string) string) (string, bool) {
+	pos := 0
+	for pos < len(content) {
+		idx := strings.Index(content[pos:], prefix)
 		if idx < 0 {
 			return content, false
 		}
+		idx += pos
 		endRel := strings.IndexByte(content[idx:], '>')
 		if endRel < 0 {
 			return content, false
@@ -433,7 +405,7 @@ func replaceLastMediaTag(content, prefix string, match func(tag string) bool, re
 		if match(tag) {
 			return content[:idx] + replace(tag) + content[end:], true
 		}
-		searchFrom = idx
+		pos = end
 	}
 	return content, false
 }
@@ -1,6 +1,8 @@
 package agent

 import (
+	"os"
+	"path/filepath"
 	"strings"
 	"testing"

@@ -131,3 +133,190 @@ func TestEnrichImagePaths_AttributeOrderIndependence(t *testing.T) {
 		t.Fatalf("expected id to be preserved, got %q", got)
 	}
 }
+
+func TestEnrichImageIDs_MultipleRefs(t *testing.T) {
+	messages := []providers.Message{{
+		Role:    "user",
+		Content: "first <media:image>\nsecond <media:image>",
+	}}
+	refs := []providers.MediaRef{
+		{ID: "img-a", Kind: "image", Path: "/tmp/a.jpg"},
+		{ID: "img-b", Kind: "image", Path: "/tmp/b.jpg"},
+	}
+
+	var loop Loop
+	loop.enrichImageIDs(messages, refs)
+
+	want := `first <media:image id="img-a" path="/tmp/a.jpg">` + "\n" + `second <media:image id="img-b" path="/tmp/b.jpg">`
+	if messages[0].Content != want {
+		t.Fatalf("multi-ref alignment:\n got %q\nwant %q", messages[0].Content, want)
+	}
+}
+
+func TestEnrichImagePaths_MultipleRefsKeepTagAlignment(t *testing.T) {
+	storeDir := t.TempDir()
+	mediaStore, err := media.NewStore(storeDir)
+	if err != nil {
+		t.Fatalf("NewStore() error = %v", err)
+	}
+
+	const (
+		sessionKey = "session-1"
+		pathA      = "/persisted/a.jpg"
+		pathB      = "/persisted/b.jpg"
+	)
+
+	srcA := filepath.Join(storeDir, "a.jpg")
+	if err := os.WriteFile(srcA, []byte("a"), 0644); err != nil {
+		t.Fatalf("WriteFile(a) error = %v", err)
+	}
+	idA, _, err := mediaStore.SaveFile(sessionKey, srcA, "image/jpeg")
+	if err != nil {
+		t.Fatalf("SaveFile(a) error = %v", err)
+	}
+
+	srcB := filepath.Join(storeDir, "b.jpg")
+	if err := os.WriteFile(srcB, []byte("b"), 0644); err != nil {
+		t.Fatalf("WriteFile(b) error = %v", err)
+	}
+	idB, _, err := mediaStore.SaveFile(sessionKey, srcB, "image/jpeg")
+	if err != nil {
+		t.Fatalf("SaveFile(b) error = %v", err)
+	}
+
+	messages := []providers.Message{{
+		Role:    "user",
+		Content: "first <media:image>\nsecond <media:image>",
+		MediaRefs: []providers.MediaRef{
+			{ID: idA, Kind: "image", Path: pathA},
+			{ID: idB, Kind: "image", Path: pathB},
+		},
+	}}
+
+	var loop Loop
+	loop.mediaStore = mediaStore
+	loop.enrichImagePaths(messages)
+
+	want := `first <media:image id="` + idA + `" path="` + pathA + `">` +
+		"\n" + `second <media:image id="` + idB + `" path="` + pathB + `">`
+	if messages[0].Content != want {
+		t.Fatalf("enrichImagePaths() content = %q, want %q", messages[0].Content, want)
+	}
+}
+
+func TestEnrichDocumentPaths_MultipleRefs(t *testing.T) {
+	messages := []providers.Message{{
+		Role:    "user",
+		Content: "first <media:document>\nsecond <media:document>",
+	}}
+	refs := []providers.MediaRef{
+		{ID: "doc-a", Kind: "document", Path: "/tmp/a.pdf"},
+		{ID: "doc-b", Kind: "document", Path: "/tmp/b.pdf"},
+	}
+
+	var loop Loop
+	loop.enrichDocumentPaths(messages, refs)
+
+	want := `first <media:document path="/tmp/a.pdf">` + "\n" + `second <media:document path="/tmp/b.pdf">`
+	if messages[0].Content != want {
+		t.Fatalf("multi-ref alignment:\n got %q\nwant %q", messages[0].Content, want)
+	}
+}
+
+func TestEnrichAudioIDs_MultipleRefs(t *testing.T) {
+	messages := []providers.Message{{
+		Role:    "user",
+		Content: "first <media:audio>\nsecond <media:audio>",
+	}}
+	refs := []providers.MediaRef{
+		{ID: "aud-a", Kind: "audio"},
+		{ID: "aud-b", Kind: "audio"},
+	}
+
+	var loop Loop
+	loop.enrichAudioIDs(messages, refs)
+
+	want := `first <media:audio id="aud-a">` + "\n" + `second <media:audio id="aud-b">`
+	if messages[0].Content != want {
+		t.Fatalf("multi-ref alignment:\n got %q\nwant %q", messages[0].Content, want)
+	}
+}
+
+func TestEnrichVideoIDs_MultipleRefs(t *testing.T) {
+	messages := []providers.Message{{
+		Role:    "user",
+		Content: "first <media:video>\nsecond <media:video>",
+	}}
+	refs := []providers.MediaRef{
+		{ID: "vid-a", Kind: "video"},
+		{ID: "vid-b", Kind: "video"},
+	}
+
+	var loop Loop
+	loop.enrichVideoIDs(messages, refs)
+
+	want := `first <media:video id="vid-a">` + "\n" + `second <media:video id="vid-b">`
+	if messages[0].Content != want {
+		t.Fatalf("multi-ref alignment:\n got %q\nwant %q", messages[0].Content, want)
+	}
+}
+
+// TestEnrichImageIDs_MoreRefsThanTags verifies that extra refs are silently skipped.
+func TestEnrichImageIDs_MoreRefsThanTags(t *testing.T) {
+	messages := []providers.Message{{
+		Role:    "user",
+		Content: "only one <media:image>",
+	}}
+	refs := []providers.MediaRef{
+		{ID: "img-a", Kind: "image", Path: "/tmp/a.jpg"},
+		{ID: "img-b", Kind: "image", Path: "/tmp/b.jpg"},
+	}
+
+	var loop Loop
+	loop.enrichImageIDs(messages, refs)
+
+	want := `only one <media:image id="img-a" path="/tmp/a.jpg">`
+	if messages[0].Content != want {
+		t.Fatalf("more refs than tags:\n got %q\nwant %q", messages[0].Content, want)
+	}
+}
+
+// TestEnrichImageIDs_MoreTagsThanRefs verifies that extra tags are left bare.
+func TestEnrichImageIDs_MoreTagsThanRefs(t *testing.T) {
+	messages := []providers.Message{{
+		Role:    "user",
+		Content: "first <media:image>\nsecond <media:image>",
+	}}
+	refs := []providers.MediaRef{
+		{ID: "img-a", Kind: "image", Path: "/tmp/a.jpg"},
+	}
+
+	var loop Loop
+	loop.enrichImageIDs(messages, refs)
+
+	want := `first <media:image id="img-a" path="/tmp/a.jpg">` + "\n" + `second <media:image>`
+	if messages[0].Content != want {
+		t.Fatalf("more tags than refs:\n got %q\nwant %q", messages[0].Content, want)
+	}
+}
+
+// TestEnrichAudioIDs_MixedAudioAndVoice verifies correct pairing when both
+// <media:audio> and <media:voice> tags appear in the same message.
+func TestEnrichAudioIDs_MixedAudioAndVoice(t *testing.T) {
+	messages := []providers.Message{{
+		Role:    "user",
+		Content: "hear this <media:audio>\nand this <media:voice>",
+	}}
+	refs := []providers.MediaRef{
+		{ID: "aud-1", Kind: "audio"},
+		{ID: "aud-2", Kind: "audio"},
+	}
+
+	var loop Loop
+	loop.enrichAudioIDs(messages, refs)
+
+	want := `hear this <media:audio id="aud-1">` + "\n" + `and this <media:voice id="aud-2">`
+	if messages[0].Content != want {
+		t.Fatalf("mixed audio/voice:\n got %q\nwant %q", messages[0].Content, want)
+	}
+}