efe42e8 · crush · git.iko.soy

internal/agent/agent.go+15−6

compare with main

unifiedsplit

@@ -1091,13 +1091,22 @@ func (a *sessionAgent) convertToToolResult(result fantasy.ToolResultContent) mes

1091 }

1092 case fantasy.ToolResultContentTypeMedia:

1093 if r, ok := fantasy.AsToolResultOutputType[fantasy.ToolResultOutputContentMedia](result.Result); ok {

1094+ if !stringext.IsValidBase64(r.Data) {

1095+ slog.Warn("Tool returned media with invalid base64 data, discarding image",

1096+ "tool", result.ToolName,

1097+ "tool_call_id", result.ToolCallID,

1098+ )

1099+ baseResult.Content = "Tool returned image data with invalid encoding"

1100+ baseResult.IsError = true

1101+ } else {

1094− content := r.Text

1102+ content := r.Text

1095− if content == "" {

1103+ if content == "" {

1096− content = fmt.Sprintf("Loaded %s content", r.MediaType)

1104+ content = fmt.Sprintf("Loaded %s content", r.MediaType)

1105+ }

1106+ baseResult.Content = content

1107+ baseResult.Data = r.Data

1108+ baseResult.MIMEType = r.MediaType

1097 }

1109 }

1098− baseResult.Content = content

1099− baseResult.Data = r.Data

1100− baseResult.MIMEType = r.MediaType

1101 }

1110 }

1102 }

1111 }

1103

1112

internal/agent/convert_test.gonew file+94

@@ -0,0 +1,94 @@

1+package agent

2+

3+import (

4+ "encoding/base64"

5+ "testing"

6+

7+ "charm.land/fantasy"

8+ "github.com/stretchr/testify/require"

9+)

10+

11+func TestConvertToToolResult_InvalidBase64(t *testing.T) {

12+ t.Parallel()

13+

14+ a := &sessionAgent{}

15+ result := fantasy.ToolResultContent{

16+ ToolCallID: "call_123",

17+ ToolName: "test_tool",

18+ Result: fantasy.ToolResultOutputContentMedia{

19+ Data: "abc\x80def",

20+ MediaType: "image/png",

21+ },

22+ }

23+

24+ tr := a.convertToToolResult(result)

25+ require.True(t, tr.IsError)

26+ require.Empty(t, tr.Data)

27+ require.Contains(t, tr.Content, "invalid encoding")

28+ require.Equal(t, "call_123", tr.ToolCallID)

29+ require.Equal(t, "test_tool", tr.Name)

30+}

31+

32+func TestConvertToToolResult_ValidMedia(t *testing.T) {

33+ t.Parallel()

34+

35+ a := &sessionAgent{}

36+ validData := base64.StdEncoding.EncodeToString([]byte("test image data"))

37+

38+ result := fantasy.ToolResultContent{

39+ ToolCallID: "call_456",

40+ ToolName: "screenshot",

41+ Result: fantasy.ToolResultOutputContentMedia{

42+ Data: validData,

43+ MediaType: "image/png",

44+ Text: "Screenshot captured",

45+ },

46+ }

47+

48+ tr := a.convertToToolResult(result)

49+ require.False(t, tr.IsError)

50+ require.Equal(t, validData, tr.Data)

51+ require.Equal(t, "image/png", tr.MIMEType)

52+ require.Equal(t, "Screenshot captured", tr.Content)

53+}

54+

55+func TestConvertToToolResult_ValidMediaNoText(t *testing.T) {

56+ t.Parallel()

57+

58+ a := &sessionAgent{}

59+ validData := base64.StdEncoding.EncodeToString([]byte("test image data"))

60+

61+ result := fantasy.ToolResultContent{

62+ ToolCallID: "call_789",

63+ ToolName: "view",

64+ Result: fantasy.ToolResultOutputContentMedia{

65+ Data: validData,

66+ MediaType: "image/jpeg",

67+ },

68+ }

69+

70+ tr := a.convertToToolResult(result)

71+ require.False(t, tr.IsError)

72+ require.Equal(t, validData, tr.Data)

73+ require.Equal(t, "image/jpeg", tr.MIMEType)

74+ require.Equal(t, "Loaded image/jpeg content", tr.Content)

75+}

76+

77+func TestConvertToToolResult_ASCIIButInvalidBase64(t *testing.T) {

78+ t.Parallel()

79+

80+ a := &sessionAgent{}

81+ result := fantasy.ToolResultContent{

82+ ToolCallID: "call_abc",

83+ ToolName: "mcp_tool",

84+ Result: fantasy.ToolResultOutputContentMedia{

85+ Data: "not-valid-base64!!!",

86+ MediaType: "image/png",

87+ },

88+ }

89+

90+ tr := a.convertToToolResult(result)

91+ require.True(t, tr.IsError)

92+ require.Empty(t, tr.Data)

93+ require.Contains(t, tr.Content, "invalid encoding")

94+}

internal/agent/tools/mcp/tools.go+17−32

@@ -88,7 +88,7 @@ func RunTool(ctx context.Context, cfg *config.ConfigStore, name, toolName string

88 return ToolResult{

89 Type: "image",

90 Content: textContent,

91− Data: ensureBase64(imageData),

91+ Data: ensureRawBytes(imageData),

92 MediaType: imageMimeType,

93 }, nil

94 }

@@ -97,7 +97,7 @@ func RunTool(ctx context.Context, cfg *config.ConfigStore, name, toolName string

97 return ToolResult{

98 Type: "media",

99 Content: textContent,

100− Data: ensureBase64(audioData),

100+ Data: ensureRawBytes(audioData),

101 MediaType: audioMimeType,

102 }, nil

103 }

@@ -167,23 +167,29 @@ func filterDisabledTools(cfg *config.ConfigStore, mcpName string, tools []*Tool)

167 return filtered

168}

169

170−// ensureBase64 normalizes valid base64 input and guarantees padded

170+// ensureRawBytes normalizes MCP media data into raw binary bytes.

171−// base64.StdEncoding output; otherwise it encodes raw binary data.

171+//

172+// The MCP Go SDK's json.Unmarshal normally base64-decodes

173+// ImageContent.Data into raw bytes automatically. However, some MCP

174+// transports (notably Docker over stdio) can deliver data in

175+// unexpected formats. This function handles both cases:

176+//

177+// - If data looks like a valid base64 string (ASCII-only, decodable)

178+// it is decoded and the raw bytes are returned.

179+// - If data is already raw binary (contains bytes > 127) it is

180+// returned as-is.

172−func ensureBase64(data []byte) []byte {

181+func ensureRawBytes(data []byte) []byte {

173 if len(data) == 0 {

182 if len(data) == 0 {

174 return data

183 return data

175 }

184 }

176

185

177 normalized := normalizeBase64Input(data)

186 normalized := normalizeBase64Input(data)

178 if decoded, ok := decodeBase64(normalized); ok {

187 if decoded, ok := decodeBase64(normalized); ok {

179− encoded := make([]byte, base64.StdEncoding.EncodedLen(len(decoded)))

180− base64.StdEncoding.Encode(encoded, decoded)

181− return encoded

188+ return decoded

182 }

189 }

183

190

184− encoded := make([]byte, base64.StdEncoding.EncodedLen(len(data)))

191+ // Already raw binary — return unchanged.

185− base64.StdEncoding.Encode(encoded, data)

186− return encoded

192+ return data

187}

193}

188

194

189func normalizeBase64Input(data []byte) []byte {

195func normalizeBase64Input(data []byte) []byte {

@@ -213,24 +219,3 @@ func decodeBase64(data []byte) ([]byte, bool) {

213 }

219 }

214 return nil, false

220 return nil, false

215}

221}

216−

217−// isValidBase64 checks if the data appears to be valid base64-encoded content.

218−func isValidBase64(data []byte) bool {

219− if len(data) == 0 {

220− return true

221− }

222−

223− // Base64 strings should only contain ASCII characters.

224− for _, b := range data {

225− if b > 127 {

226− return false

227− }

228− }

229−

230− s := string(data)

231− if _, err := base64.StdEncoding.DecodeString(s); err == nil {

232− return true

233− }

234− _, err := base64.RawStdEncoding.DecodeString(s)

235− return err == nil

236−}

internal/agent/tools/mcp/tools_test.go+14−70

@@ -1,34 +1,35 @@

1package mcp

2

3import (

4+ "bytes"

4 "encoding/base64"

5 "encoding/base64"

5 "testing"

6 "testing"

6

7

7 "github.com/stretchr/testify/require"

8 "github.com/stretchr/testify/require"

8)

9)

9

10

10−func TestEnsureBase64(t *testing.T) {

11+func TestEnsureRawBytes(t *testing.T) {

11 t.Parallel()

12 t.Parallel()

12

13

13 tests := []struct {

14 tests := []struct {

14 name string

15 name string

15 input []byte

16 input []byte

16− wantData []byte // expected output

17+ wantData []byte

17 }{

18 }{

18 {

19 {

19 name: "already base64 encoded",

20 name: "already base64 encoded",

20 input: []byte("SGVsbG8gV29ybGQh"), // "Hello World!" in base64

21 input: []byte("SGVsbG8gV29ybGQh"), // "Hello World!" in base64

21− wantData: []byte("SGVsbG8gV29ybGQh"),

22+ wantData: []byte("Hello World!"),

22 },

23 },

23 {

24 {

24 name: "raw binary data (PNG header)",

25 name: "raw binary data (PNG header)",

25 input: []byte{0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A},

26 input: []byte{0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A},

26− wantData: []byte(base64.StdEncoding.EncodeToString([]byte{0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A})),

27+ wantData: []byte{0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A},

27 },

28 },

28 {

29 {

29 name: "raw binary with high bytes",

30 name: "raw binary with high bytes",

30 input: []byte{0xFF, 0xD8, 0xFF, 0xE0}, // JPEG header

31 input: []byte{0xFF, 0xD8, 0xFF, 0xE0}, // JPEG header

31− wantData: []byte(base64.StdEncoding.EncodeToString([]byte{0xFF, 0xD8, 0xFF, 0xE0})),

32+ wantData: []byte{0xFF, 0xD8, 0xFF, 0xE0},

32 },

33 },

33 {

34 {

34 name: "empty data",

35 name: "empty data",

@@ -38,88 +39,31 @@ func TestEnsureBase64(t *testing.T) {

38 {

39 {

39 name: "base64 with padding",

40 name: "base64 with padding",

40 input: []byte("YQ=="), // "a" in base64

41 input: []byte("YQ=="), // "a" in base64

41− wantData: []byte("YQ=="),

42+ wantData: []byte("a"),

42 },

43 },

43 {

44 {

44 name: "base64 without padding",

45 name: "base64 without padding",

45 input: []byte("YQ"),

46 input: []byte("YQ"),

46− wantData: []byte("YQ=="),

47+ wantData: []byte("a"),

47 },

48 },

48 {

49 {

49 name: "base64 with whitespace",

50 name: "base64 with whitespace",

50 input: []byte("U0dWc2JHOGdWMjl5YkdRaA==\n"),

51 input: []byte("U0dWc2JHOGdWMjl5YkdRaA==\n"),

51− wantData: []byte("U0dWc2JHOGdWMjl5YkdRaA=="),

52+ wantData: []byte("SGVsbG8gV29ybGQh"),

52 },

53 },

53 }

54 }

54

55

55 for _, tt := range tests {

56 for _, tt := range tests {

56 t.Run(tt.name, func(t *testing.T) {

57 t.Run(tt.name, func(t *testing.T) {

57 t.Parallel()

58 t.Parallel()

58− result := ensureBase64(tt.input)

59+ result := ensureRawBytes(tt.input)

59 require.Equal(t, tt.wantData, result)

60 require.Equal(t, tt.wantData, result)

60

61

61− // Verify the result is valid base64 that can be decoded.

62− if len(result) > 0 {

62+ if len(result) > 0 && !bytes.Equal(result, tt.input) {

63− _, err := base64.StdEncoding.DecodeString(string(result))

63+ reEncoded := base64.StdEncoding.EncodeToString(result)

64− if err != nil {

65− _, err = base64.RawStdEncoding.DecodeString(string(result))

64+ _, err := base64.StdEncoding.DecodeString(reEncoded)

66− }

67− require.NoError(t, err, "result should be valid base64")

65+ require.NoError(t, err, "re-encoded result should be valid base64")

68 }

66 }

69 })

67 })

70 }

68 }

71}

69}

72−

73−func TestIsValidBase64(t *testing.T) {

74− t.Parallel()

75−

76− tests := []struct {

77− name string

78− input []byte

79− want bool

80− }{

81− {

82− name: "valid base64",

83− input: []byte("SGVsbG8gV29ybGQh"),

84− want: true,

85− },

86− {

87− name: "valid base64 with padding",

88− input: []byte("YQ=="),

89− want: true,

90− },

91− {

92− name: "raw binary with high bytes",

93− input: []byte{0xFF, 0xD8, 0xFF},

94− want: false,

95− },

96− {

97− name: "empty",

98− input: []byte{},

99− want: true,

100− },

101− {

102− name: "valid raw base64 without padding",

103− input: []byte("YQ"),

104− want: true,

105− },

106− {

107− name: "valid base64 with whitespace",

108− input: normalizeBase64Input([]byte("U0dWc2JHOGdWMjl5YkdRaA==\n")),

109− want: true,

110− },

111− {

112− name: "invalid base64 characters",

113− input: []byte("SGVsbG8!@#$"),

114− want: false,

115− },

116− }

117−

118− for _, tt := range tests {

119− t.Run(tt.name, func(t *testing.T) {

120− t.Parallel()

121− got := isValidBase64(tt.input)

122− require.Equal(t, tt.want, got)

123− })

124− }

125−}

internal/agent/tools/view.go+1−3

@@ -4,7 +4,6 @@ import (

4 "bufio"

5 "context"

6 _ "embed"

7− "encoding/base64"

8 "fmt"

7 "fmt"

9 "io"

8 "io"

10 "io/fs"

9 "io/fs"

@@ -189,8 +188,7 @@ func NewViewTool(

189 return fantasy.ToolResponse{}, fmt.Errorf("error reading image file: %w", readErr)

188 return fantasy.ToolResponse{}, fmt.Errorf("error reading image file: %w", readErr)

190 }

189 }

191

190

192− encoded := base64.StdEncoding.EncodeToString(imageData)

193− return fantasy.NewImageResponse([]byte(encoded), mimeType), nil

191+ return fantasy.NewImageResponse(imageData, mimeType), nil

194 }

192 }

195

193

196 // Read the file content

194 // Read the file content

internal/message/content.go+14−3

@@ -13,6 +13,7 @@ import (

13 "charm.land/fantasy/providers/anthropic"

14 "charm.land/fantasy/providers/google"

15 "charm.land/fantasy/providers/openai"

16+ "github.com/charmbracelet/crush/internal/stringext"

16)

17)

17

18

18type MessageRole string

19type MessageRole string

@@ -24,6 +25,10 @@ const (

24 Tool MessageRole = "tool"

25 Tool MessageRole = "tool"

25)

26)

26

27

28+// mediaLoadFailedPlaceholder is the text substituted for image data that

29+// cannot be decoded during session replay.

30+const mediaLoadFailedPlaceholder = "[Image data could not be loaded]"

31+

27type FinishReason string

32type FinishReason string

28

33

29const (

34const (

@@ -542,9 +547,15 @@ func (m *Message) ToAIMessage() []fantasy.Message {

542 Error: errors.New(result.Content),

547 Error: errors.New(result.Content),

543 }

548 }

544 } else if result.Data != "" {

549 } else if result.Data != "" {

550+ if stringext.IsValidBase64(result.Data) {

545− content = fantasy.ToolResultOutputContentMedia{

551+ content = fantasy.ToolResultOutputContentMedia{

546− Data: result.Data,

552+ Data: result.Data,

547− MediaType: result.MIMEType,

553+ MediaType: result.MIMEType,

554+ }

555+ } else {

556+ content = fantasy.ToolResultOutputContentText{

557+ Text: mediaLoadFailedPlaceholder,

558+ }

548 }

559 }

549 } else {

560 } else {

550 content = fantasy.ToolResultOutputContentText{

561 content = fantasy.ToolResultOutputContentText{

internal/message/content_test.go+97

@@ -1,9 +1,13 @@

1package message

2

3import (

4+ "encoding/base64"

4 "fmt"

5 "fmt"

5 "strings"

6 "strings"

6 "testing"

7 "testing"

8+

9+ "charm.land/fantasy"

10+ "github.com/stretchr/testify/require"

7)

11)

8

12

9func makeTestAttachments(n int, contentSize int) []Attachment {

13func makeTestAttachments(n int, contentSize int) []Attachment {

@@ -19,6 +23,99 @@ func makeTestAttachments(n int, contentSize int) []Attachment {

19 return attachments

23 return attachments

20}

24}

21

25

26+func 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+

56+func 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+

89+func 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+

22func BenchmarkPromptWithTextAttachments(b *testing.B) {

119func BenchmarkPromptWithTextAttachments(b *testing.B) {

23 cases := []struct {

120 cases := []struct {

24 name string

121 name string

internal/stringext/string.go+18

@@ -1,6 +1,7 @@

1package stringext

2

3import (

4+ "encoding/base64"

4 "strings"

5 "strings"

5

6

6 "golang.org/x/text/cases"

7 "golang.org/x/text/cases"

@@ -20,3 +21,20 @@ func NormalizeSpace(content string) string {

20 content = strings.TrimSpace(content)

21 content = strings.TrimSpace(content)

21 return content

22 return content

22}

23}

24+

25+// IsValidBase64 reports whether s is canonical base64 under standard

26+// encoding (RFC 4648). It requires that s round-trips through

27+// decode/encode unchanged — rejecting whitespace, missing padding,

28+// and other leniencies that DecodeString alone would accept.

29+func IsValidBase64(s string) bool {

30+ if s == "" {

31+ return false

32+ }

33+ decoded, err := base64.StdEncoding.DecodeString(s)

34+ if err != nil {

35+ return false

36+ }

37+ // Round-trip check rejects whitespace, missing padding, and other

38+ // leniencies that DecodeString silently accepts.

39+ return base64.StdEncoding.EncodeToString(decoded) == s

40+}

internal/stringext/string_test.gonew file+38

@@ -0,0 +1,38 @@

1+package stringext

2+

3+import (

4+ "encoding/base64"

5+ "testing"

6+

7+ "github.com/stretchr/testify/require"

8+)

9+

10+func TestIsValidBase64(t *testing.T) {

11+ t.Parallel()

12+

13+ // Real PNG header encoded in standard base64.

14+ pngHeader := []byte{0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A}

15+ pngBase64 := base64.StdEncoding.EncodeToString(pngHeader)

16+

17+ tests := []struct {

18+ name string

19+ input string

20+ expected bool

21+ }{

22+ {name: "empty string", input: "", expected: false},

23+ {name: "valid no padding", input: "SGVsbG8gV29ybGQh", expected: true},

24+ {name: "valid with padding", input: "YQ==", expected: true},

25+ {name: "non-ASCII bytes", input: "abc\x80def", expected: false},

26+ {name: "ASCII but not base64", input: "hello world!!!", expected: false},

27+ {name: "raw encoding no padding", input: "YQ", expected: false},

28+ {name: "trailing whitespace", input: "YQ==\n", expected: false},

29+ {name: "valid PNG header base64", input: pngBase64, expected: true},

30+ }

31+

32+ for _, tt := range tests {

33+ t.Run(tt.name, func(t *testing.T) {

34+ t.Parallel()

35+ require.Equal(t, tt.expected, IsValidBase64(tt.input))

36+ })

37+ }

38+}