1package message
2
3import (
4 "encoding/base64"
5 "fmt"
6 "strings"
7 "testing"
8
9 "charm.land/fantasy"
10 "github.com/stretchr/testify/require"
11)
12
13func makeTestAttachments(n int, contentSize int) []Attachment {
14 attachments := make([]Attachment, n)
15 content := []byte(strings.Repeat("x", contentSize))
16 for i := range n {
17 attachments[i] = Attachment{
18 FilePath: fmt.Sprintf("/path/to/file%d.txt", i),
19 MimeType: "text/plain",
20 Content: content,
21 }
22 }
23 return attachments
24}
25
26func TestToAIMessage_CorruptedMediaData(t *testing.T) {
27 t.Parallel()
28
29 msg := &Message{
30 Role: Tool,
31 Parts: []ContentPart{
32 ToolResult{
33 ToolCallID: "call_123",
34 Name: "screenshot",
35 Content: "Loaded image/png content",
36 Data: "abc\x80def",
37 MIMEType: "image/png",
38 },
39 },
40 }
41
42 messages := msg.ToAIMessage()
43 require.Len(t, messages, 1)
44 require.Len(t, messages[0].Content, 1)
45
46 part, ok := messages[0].Content[0].(fantasy.ToolResultPart)
47 require.True(t, ok)
48
49 require.Equal(t, "call_123", part.ToolCallID)
50
51 textContent, ok := part.Output.(fantasy.ToolResultOutputContentText)
52 require.True(t, ok, "corrupted media should be downgraded to text")
53 require.Equal(t, mediaLoadFailedPlaceholder, textContent.Text)
54}
55
56func TestToAIMessage_ValidMediaData(t *testing.T) {
57 t.Parallel()
58
59 validBase64 := base64.StdEncoding.EncodeToString([]byte{0x89, 0x50, 0x4E, 0x47})
60
61 msg := &Message{
62 Role: Tool,
63 Parts: []ContentPart{
64 ToolResult{
65 ToolCallID: "call_456",
66 Name: "screenshot",
67 Content: "Loaded image/png content",
68 Data: validBase64,
69 MIMEType: "image/png",
70 },
71 },
72 }
73
74 messages := msg.ToAIMessage()
75 require.Len(t, messages, 1)
76 require.Len(t, messages[0].Content, 1)
77
78 part, ok := messages[0].Content[0].(fantasy.ToolResultPart)
79 require.True(t, ok)
80
81 require.Equal(t, "call_456", part.ToolCallID)
82
83 mediaContent, ok := part.Output.(fantasy.ToolResultOutputContentMedia)
84 require.True(t, ok, "valid media should remain as media")
85 require.Equal(t, validBase64, mediaContent.Data)
86 require.Equal(t, "image/png", mediaContent.MediaType)
87}
88
89func TestToAIMessage_ASCIIButInvalidBase64(t *testing.T) {
90 t.Parallel()
91
92 msg := &Message{
93 Role: Tool,
94 Parts: []ContentPart{
95 ToolResult{
96 ToolCallID: "call_789",
97 Name: "screenshot",
98 Content: "Loaded image/png content",
99 Data: "not-valid-base64!!!",
100 MIMEType: "image/png",
101 },
102 },
103 }
104
105 messages := msg.ToAIMessage()
106 require.Len(t, messages, 1)
107 require.Len(t, messages[0].Content, 1)
108
109 part, ok := messages[0].Content[0].(fantasy.ToolResultPart)
110 require.True(t, ok)
111
112 require.Equal(t, "call_789", part.ToolCallID)
113
114 textContent, ok := part.Output.(fantasy.ToolResultOutputContentText)
115 require.True(t, ok, "ASCII but invalid base64 should be downgraded to text")
116 require.Equal(t, mediaLoadFailedPlaceholder, textContent.Text)
117}
118
119func BenchmarkPromptWithTextAttachments(b *testing.B) {
120 cases := []struct {
121 name string
122 numFiles int
123 contentSize int
124 }{
125 {"1file_100bytes", 1, 100},
126 {"5files_1KB", 5, 1024},
127 {"10files_10KB", 10, 10 * 1024},
128 {"20files_50KB", 20, 50 * 1024},
129 }
130
131 for _, tc := range cases {
132 attachments := makeTestAttachments(tc.numFiles, tc.contentSize)
133 prompt := "Process these files"
134
135 b.Run(tc.name, func(b *testing.B) {
136 b.ReportAllocs()
137 for range b.N {
138 _ = PromptWithTextAttachments(prompt, attachments)
139 }
140 })
141 }
142}