TUI - Optimize rendering of long streaming responses

pkg/tui/components/markdown/incremental_test.go

@@ -2,13 +2,57 @@ package markdown
 
 import (
 	_ "embed"
+	"strconv"
 	"strings"
 	"testing"
+	"unicode/utf8"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
+func TestIncrementalRenderedPartsPreserveMarkdownSemantics(t *testing.T) {
+	t.Parallel()
+	cases := map[string]string{"long heading": "# A long heading that wraps across multiple terminal rows while retaining heading styling and continuation indentation", "unordered list": "- first list item with enough words to wrap onto a continuation line that stays indented\n- second item", "ordered list": "1. first ordered item with enough words to wrap onto a continuation line that stays indented\n2. second item", "mixed blocks": "# Heading\n\nOpening **paragraph** with a [link](https://example.com).\n\n- list item\n  continued content\n\n1. ordered item\n2. next item\n\nFinal `code` paragraph."}
+	for name, input := range cases {
+		t.Run(name, func(t *testing.T) {
+			for _, width := range []int{24, 47, 80} {
+				r := NewIncrementalRenderer(width)
+				for end := 1; end <= len(input); end++ {
+					parts, err := r.RenderParts(input[:end])
+					require.NoError(t, err)
+					got := r.joinPrefixAndTail(parts.StablePrefix, parts.MutableTail)
+					want, err := NewFastRenderer(width).Render(input[:end])
+					require.NoError(t, err)
+					require.Equal(t, want, got, "width %d byte %d", width, end)
+				}
+			}
+		})
+	}
+}
+
+func TestIncrementalRenderedPartsMatchOneShotAtEveryMarkdownBoundary(t *testing.T) {
+	const input = "Thinking… λ界\n\n# Heading\n\nParagraph with **bold**, `more`, and [link](https://example.com).\n\n- one\n- two\n\n```console\nroot\nmore\n```\n\n## Result\n\nDone."
+	for _, width := range []int{24, 47, 80} {
+		t.Run(strconv.Itoa(width), func(t *testing.T) {
+			r := NewIncrementalRenderer(width)
+			for end := range len(input) + 1 {
+				if end < len(input) && !utf8.RuneStart(input[end]) {
+					continue
+				}
+				prefix := input[:end]
+				parts, err := r.RenderParts(prefix)
+				require.NoError(t, err)
+				got := r.joinPrefixAndTail(parts.StablePrefix, parts.MutableTail)
+				want, wantBlocks, err := NewFastRenderer(width).RenderWithCodeBlocks(prefix)
+				require.NoError(t, err)
+				require.Equal(t, want, got, "byte boundary %d", end)
+				require.Equal(t, wantBlocks, parts.CodeBlocks, "code blocks at byte boundary %d", end)
+			}
+		})
+	}
+}
+
 func TestIncrementalRenderedPartsMatchJoinedOutputAtEveryStep(t *testing.T) {
 	t.Parallel()
 	chunks := []string{"unfinished *em", "phasis* and [li", "nk](https://example.com)\n\n", "```go\nfmt.Print(\"λ界\")", "\n```\n\n- one", "\n- two\n\nfinal"}

pkg/tui/components/message/append_content_test.go

@@ -0,0 +1,153 @@
+package message
+
+import (
+	"bytes"
+	"encoding/base64"
+	stdimage "image"
+	"image/color"
+	"image/png"
+	"strconv"
+	"strings"
+	"testing"
+	"unicode/utf8"
+
+	"github.com/charmbracelet/x/ansi"
+	"github.com/stretchr/testify/require"
+
+	"github.com/docker/docker-agent/pkg/tui/animation"
+	tuiimage "github.com/docker/docker-agent/pkg/tui/image"
+	"github.com/docker/docker-agent/pkg/tui/types"
+)
+
+func testImageURI(t *testing.T, colorValue color.RGBA) string {
+	t.Helper()
+	img := stdimage.NewRGBA(stdimage.Rect(0, 0, 2, 1))
+	img.Set(0, 0, colorValue)
+	var data bytes.Buffer
+	require.NoError(t, png.Encode(&data, img))
+	return "data:image/png;base64," + base64.StdEncoding.EncodeToString(data.Bytes())
+}
+
+func TestAppendContentPreservesSplitImageOpenerAndExactContent(t *testing.T) {
+	msg := types.Agent(types.MessageTypeAssistant, "root", "prefix ")
+	m := New(animation.NewRuntime(), msg, nil)
+	for _, chunk := range []string{"!", "[alt]", "(https://example.com/image.png)", " suffix"} {
+		_ = m.AppendContent(chunk)
+	}
+	require.Equal(t, "prefix ![alt](https://example.com/image.png) suffix", m.message.Content)
+	require.Equal(t, "prefix ![alt](https://example.com/image.png) suffix", m.contentBuf.String())
+}
+
+func TestAppendContentSplitImageSchedulesAndRendersInline(t *testing.T) {
+	tuiimage.SetRenderingEnabled(true)
+	uri := testImageURI(t, color.RGBA{G: 255, A: 255})
+	m := New(animation.NewRuntime(), types.Agent(types.MessageTypeAssistant, "root", ""), nil)
+	m.SetSize(80, 0)
+	chunks := []string{"!", "[alt]", "(" + uri + ")"}
+	for _, chunk := range chunks[:2] {
+		require.Nil(t, m.AppendContent(chunk))
+	}
+	cmd := m.AppendContent(chunks[2])
+	require.NotNil(t, cmd, "completing a pending image reference must schedule loading")
+	_, _ = m.Update(cmd())
+	view := m.View()
+	require.Contains(t, view, "cagent-image")
+	require.Contains(t, ansi.Strip(view), "alt")
+	require.NotContains(t, ansi.Strip(view), strings.Join(chunks, ""))
+}
+
+func TestAppendContentCompletesInitialImageAcrossPartitions(t *testing.T) {
+	uri := testImageURI(t, color.RGBA{R: 128, G: 64, A: 255})
+	full := "prefix ![界](" + uri + ") suffix"
+	for split := range len(full) + 1 {
+		if split == 0 || split == len(full) || !utf8.RuneStart(full[split]) {
+			continue
+		}
+		initial, appended := full[:split], full[split:]
+		if len(tuiimage.MarkdownReferences(initial)) != 0 || !strings.Contains(initial, "![") {
+			continue
+		}
+		t.Run(strconv.Itoa(split), func(t *testing.T) {
+			m := New(animation.NewRuntime(), types.Agent(types.MessageTypeAssistant, "root", initial), nil)
+			_ = m.Init()
+			require.GreaterOrEqual(t, m.imageScanOffset, 0)
+			cmd := m.AppendContent(appended)
+			require.NotNil(t, cmd, "partition %d must schedule the completed initial image", split)
+			_, _ = m.Update(cmd())
+			require.Contains(t, m.markdownImages, uri)
+			require.Equal(t, full, m.message.Content)
+		})
+	}
+}
+
+func TestAppendContentInitialUnresolvedImageForms(t *testing.T) {
+	uri := testImageURI(t, color.RGBA{B: 128, A: 255})
+	for _, tc := range []struct {
+		initial, appended string
+	}{
+		{"![alt", "](" + uri + ")"},
+		{"prefix ![界", "](" + uri + ")"},
+		{"![alt](", uri + ")"},
+	} {
+		m := New(animation.NewRuntime(), types.Agent(types.MessageTypeAssistant, "root", tc.initial), nil)
+		_ = m.Init()
+		require.GreaterOrEqual(t, m.imageScanOffset, 0)
+		cmd := m.AppendContent(tc.appended)
+		require.NotNil(t, cmd)
+		_, _ = m.Update(cmd())
+		require.Contains(t, m.markdownImages, uri)
+	}
+}
+
+func TestAppendContentSchedulesSecondImageSplitAfterFirstCompletion(t *testing.T) {
+	first := testImageURI(t, color.RGBA{R: 255, A: 255})
+	second := testImageURI(t, color.RGBA{B: 255, A: 255})
+	m := New(animation.NewRuntime(), types.Agent(types.MessageTypeAssistant, "root", ""), nil)
+
+	cmd := m.AppendContent("![one](" + first + ") ![two")
+	require.NotNil(t, cmd)
+	_, _ = m.Update(cmd())
+	require.GreaterOrEqual(t, m.imageScanOffset, 0, "incomplete second opener remains tracked")
+
+	cmd = m.AppendContent("](" + second + ")")
+	require.NotNil(t, cmd, "completing the second image must schedule loading")
+	_, _ = m.Update(cmd())
+	require.Contains(t, m.markdownImages, first)
+	require.Contains(t, m.markdownImages, second)
+	require.Equal(t, -1, m.imageScanOffset)
+}
+
+func TestAppendContentSplitImageSyntaxInCodeDoesNotFetch(t *testing.T) {
+	for _, chunks := range [][]string{
+		{"`!", "[inline]", "(https://example.com/inline.png)`"},
+		{"```md\n!", "[fenced]", "(https://example.com/fenced.png)\n```"},
+	} {
+		m := New(animation.NewRuntime(), types.Agent(types.MessageTypeAssistant, "root", ""), nil)
+		for _, chunk := range chunks {
+			require.Nil(t, m.AppendContent(chunk))
+		}
+		require.Empty(t, m.loadingImages)
+	}
+}
+
+func TestAppendContentImageScanStateIsBoundedAndCleared(t *testing.T) {
+	m := New(animation.NewRuntime(), types.Agent(types.MessageTypeAssistant, "root", ""), nil)
+	long := "![abandoned" + strings.Repeat("界", 100_000)
+	require.Nil(t, m.AppendContent(long))
+	require.Equal(t, 0, m.imageScanOffset)
+	require.Equal(t, len(long), m.contentBuf.Len(), "scan state must not duplicate the canonical suffix")
+
+	m.Finalize()
+	require.Equal(t, -1, m.imageScanOffset)
+	_ = m.SetMessage(types.Agent(types.MessageTypeAssistant, "root", "reset"))
+	require.Equal(t, -1, m.imageScanOffset)
+}
+
+func TestAppendContentImageScanOffsetIsUTF8Safe(t *testing.T) {
+	m := New(animation.NewRuntime(), types.Agent(types.MessageTypeAssistant, "root", ""), nil)
+	require.Nil(t, m.AppendContent("界!"))
+	require.Equal(t, -1, m.imageScanOffset)
+	require.Nil(t, m.AppendContent("[界"))
+	require.Equal(t, len("界"), m.imageScanOffset)
+	require.Equal(t, "![界", m.message.Content[m.imageScanOffset:])
+}