Merge pull request #3721 from docker/worktree-board-570607f2c61e5e8e

fix(tui): address PR #3720 review feedback on markdown image rendering

Author: Sayt-0

pkg/tui/components/message/message.go

@@ -87,8 +87,9 @@ type messageModel struct {
 }
 
 type markdownImagesLoadedMsg struct {
-	target *messageModel
-	images map[string]tuiimage.Inline
+	target    *messageModel
+	requested []tuiimage.MarkdownReference
+	images    map[string]tuiimage.Inline
 }
 
 type markdownImageRenderedMsg struct{}
@@ -187,7 +188,7 @@ func (mv *messageModel) loadMarkdownImages(msg *types.Message) tea.Cmd {
 				loaded[ref.Source] = image
 			}
 		}
-		return markdownImagesLoadedMsg{target: mv, images: loaded}
+		return markdownImagesLoadedMsg{target: mv, requested: pending, images: loaded}
 	}
 }
 
@@ -208,6 +209,12 @@ func (mv *messageModel) SetHovered(hovered bool) {
 // Update handles messages and updates the message view state
 func (mv *messageModel) Update(msg tea.Msg) (layout.Model, tea.Cmd) {
 	if loaded, ok := msg.(markdownImagesLoadedMsg); ok && loaded.target == mv {
+		// Unmark failed URLs so a later SetMessage can retry them.
+		for _, ref := range loaded.requested {
+			if _, success := loaded.images[ref.Source]; !success {
+				delete(mv.loadingImages, ref.Source)
+			}
+		}
 		if len(loaded.images) > 0 {
 			if mv.markdownImages == nil {
 				mv.markdownImages = make(map[string]tuiimage.Inline)
@@ -532,8 +539,11 @@ func replaceMarkdownImagePlaceholders(rendered string, codeBlocks []markdown.Cod
 		shifts = append(shifts, lineShift{line: lineIndex, delta: len(replacement) - 1})
 	}
 	for i := range codeBlocks {
+		// Compare against the pre-replacement line so later shifts don't
+		// re-match against already-adjusted positions.
+		orig := codeBlocks[i].Line
 		for _, shift := range shifts {
-			if shift.line < codeBlocks[i].Line {
+			if shift.line < orig {
 				codeBlocks[i].Line += shift.delta
 			}
 		}

pkg/tui/components/message/message_test.go

@@ -15,6 +15,7 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
+	"github.com/docker/docker-agent/pkg/tui/components/markdown"
 	"github.com/docker/docker-agent/pkg/tui/components/spinner"
 	tuiimage "github.com/docker/docker-agent/pkg/tui/image"
 	"github.com/docker/docker-agent/pkg/tui/types"
@@ -50,6 +51,42 @@ func TestAssistantMarkdownImageRendersInline(t *testing.T) {
 	assert.Less(t, strings.Index(view, "chart"), strings.Index(view, "cagent-image"), "alt label must immediately precede the image")
 }
 
+func TestFailedMarkdownImageLoadCanRetry(t *testing.T) {
+	t.Parallel()
+
+	source := "https://example.com/missing.png"
+	msg := types.Agent(types.MessageTypeAssistant, "assistant", "![img]("+source+")")
+	mv := New(msg, nil)
+	mv.loadingImages = map[string]bool{source: true}
+
+	_, _ = mv.Update(markdownImagesLoadedMsg{
+		target:    mv,
+		requested: []tuiimage.MarkdownReference{{Source: source}},
+		images:    map[string]tuiimage.Inline{},
+	})
+
+	assert.Empty(t, mv.loadingImages, "failed sources must be cleared so SetMessage can retry")
+	assert.NotNil(t, mv.loadMarkdownImages(msg), "a retry fetch must be scheduled")
+}
+
+func TestReplaceMarkdownImagePlaceholdersShiftsCodeBlocksByOriginalLine(t *testing.T) {
+	t.Parallel()
+
+	lines := make([]string, 13)
+	lines[3] = "TOKENA"
+	lines[12] = "TOKENB"
+	placeholders := []markdownImagePlaceholder{
+		{token: "TOKENA", lines: []string{"a1", "a2", "a3", "a4", "a5", "a6"}}, // delta +5
+		{token: "TOKENB", lines: []string{"b1", "b2", "b3"}},                   // delta +2, after the code block
+	}
+
+	_, adjusted := replaceMarkdownImagePlaceholders(strings.Join(lines, "\n"), []markdown.CodeBlock{{Line: 10}}, placeholders)
+
+	// Only the placeholder before the code block's original line shifts it.
+	require.Len(t, adjusted, 1)
+	assert.Equal(t, 15, adjusted[0].Line)
+}
+
 func TestErrorMessageWrapping(t *testing.T) {
 	t.Parallel()
 

pkg/tui/image/image.go

@@ -224,15 +224,13 @@ func LoadMarkdownReference(ctx context.Context, ref MarkdownReference) (Inline,
 			return Inline{}, false
 		}
 	} else {
-		path := ref.Source
-		switch {
-		case parsed == nil || parsed.Scheme == "":
-		case parsed.Scheme == "file", parsed.Scheme == "sandbox":
-			path = parsed.Path
-		default:
+		// Sources come from LLM-generated markdown, so local schemes like
+		// file:// and sandbox:// would let a prompt-injected response read
+		// arbitrary local files (confused deputy). Only bare paths pass.
+		if parsed != nil && parsed.Scheme != "" {
 			return Inline{}, false
 		}
-		file, err := os.Open(path)
+		file, err := os.Open(ref.Source)
 		if err != nil {
 			return Inline{}, false
 		}

pkg/tui/image/image_test.go

@@ -7,6 +7,8 @@ import (
 	stdimage "image"
 	"image/color"
 	"image/png"
+	"os"
+	"path/filepath"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -50,6 +52,26 @@ func TestLoadMarkdownReferenceDataURI(t *testing.T) {
 	assert.NotEmpty(t, inline.PNGData)
 }
 
+func TestLoadMarkdownReferenceRejectsLocalSchemes(t *testing.T) {
+	t.Parallel()
+
+	img := stdimage.NewRGBA(stdimage.Rect(0, 0, 2, 1))
+	var encoded bytes.Buffer
+	require.NoError(t, png.Encode(&encoded, img))
+	path := filepath.Join(t.TempDir(), "chart.png")
+	require.NoError(t, os.WriteFile(path, encoded.Bytes(), 0o600))
+
+	for _, source := range []string{"file://" + path, "sandbox://" + path} {
+		_, ok := LoadMarkdownReference(t.Context(), MarkdownReference{Alt: "chart", Source: source})
+		assert.False(t, ok, source)
+	}
+
+	// Bare paths remain supported for agent-generated local images.
+	inline, ok := LoadMarkdownReference(t.Context(), MarkdownReference{Alt: "chart", Source: path})
+	require.True(t, ok)
+	assert.Equal(t, "chart", inline.Name)
+}
+
 func TestInlineRegistryIsBounded(t *testing.T) {
 	resetInlineRegistry()
 	defer resetInlineRegistry()