From a02339c66a3f44fdbe2b430fa723ddcd6343f950 Mon Sep 17 00:00:00 2001 From: Cian Johnston Date: Thu, 23 Apr 2026 19:58:38 +0100 Subject: [PATCH] fix(coderd/x/chatd): prevent invalid tool results from poisoning chat history (#24663) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - **computeruse.go**: Decode base64 screenshot data before storing in `ToolResponse.Data` (was casting base64 string to bytes without decoding) - **chatloop.go**: Re-encode `ToolResponse.Data` to base64 via `base64.StdEncoding.EncodeToString` instead of `string()` cast - **mcpclient.go**: UTF-8 validate all text from MCP responses in `convertCallResult()` using `strings.ToValidUTF8` - **chatprompt.go (persist)**: Defense-in-depth UTF-8 sanitization of text and media Text fields before database storage - **chatprompt.go (replay)**: Antivenom layer that validates base64 and UTF-8 at read time, auto-healing already-poisoned chats without requiring a migration - `TestToolResultAntivenom`: 4 subtests covering poisoned text, poisoned media, valid media round-trip, and media with invalid UTF-8 text - Adds `TestConvertCallResult_UTF8Sanitization`: 4 subtests covering invalid UTF-8 in TextContent, EmbeddedResource, valid passthrough, and multi-part - Adds `TestComputerUseTool_Run_ScreenshotDataIsDecodedBinary`: Verifies no double-encode in the computer-use path - Updated existing computer-use tests for the new decoded-binary contract > πŸ€– --- coderd/x/chatd/chatd_test.go | 2 +- coderd/x/chatd/chatloop/chatloop.go | 7 +- coderd/x/chatd/chatloop/chatloop_test.go | 142 +++++++++++++ coderd/x/chatd/chatprompt/chatprompt.go | 111 +++++++--- coderd/x/chatd/chatprompt/chatprompt_test.go | 213 +++++++++++++++++++ coderd/x/chatd/chatprompt/export_test.go | 17 ++ coderd/x/chatd/chattool/computeruse.go | 29 ++- coderd/x/chatd/chattool/computeruse_test.go | 78 ++++++- coderd/x/chatd/chattool/mcpworkspace.go | 4 +- coderd/x/chatd/mcpclient/export_test.go | 5 + coderd/x/chatd/mcpclient/mcpclient.go | 4 +- coderd/x/chatd/mcpclient/mcpclient_test.go | 75 +++++++ 12 files changed, 638 insertions(+), 49 deletions(-) create mode 100644 coderd/x/chatd/mcpclient/export_test.go diff --git a/coderd/x/chatd/chatd_test.go b/coderd/x/chatd/chatd_test.go index 605dce907e..7f545974c3 100644 --- a/coderd/x/chatd/chatd_test.go +++ b/coderd/x/chatd/chatd_test.go @@ -6177,7 +6177,7 @@ func TestComputerUseSubagentToolsAndModel(t *testing.T) { Return(workspacesdk.DesktopActionResponse{ ScreenshotWidth: 1920, ScreenshotHeight: 1080, - ScreenshotData: "iVBOR", + ScreenshotData: "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGP4n539HwAHFwLVF8kc1wAAAABJRU5ErkJggg==", }, nil). AnyTimes() mockConn.EXPECT(). diff --git a/coderd/x/chatd/chatloop/chatloop.go b/coderd/x/chatd/chatloop/chatloop.go index ad647de15d..9113297f16 100644 --- a/coderd/x/chatd/chatloop/chatloop.go +++ b/coderd/x/chatd/chatloop/chatloop.go @@ -3,6 +3,7 @@ package chatloop import ( "context" "database/sql" + "encoding/base64" "encoding/json" "errors" "maps" @@ -1213,13 +1214,13 @@ func executeSingleTool( ) case resp.Type == "image" || resp.Type == "media": result.Result = fantasy.ToolResultOutputContentMedia{ - Data: string(resp.Data), + Data: base64.StdEncoding.EncodeToString(resp.Data), MediaType: resp.MediaType, - Text: resp.Content, + Text: strings.ToValidUTF8(resp.Content, "\uFFFD"), } default: result.Result = fantasy.ToolResultOutputContentText{ - Text: resp.Content, + Text: strings.ToValidUTF8(resp.Content, "\uFFFD"), } } return result diff --git a/coderd/x/chatd/chatloop/chatloop_test.go b/coderd/x/chatd/chatloop/chatloop_test.go index 2e7539bdb8..57ca1174d9 100644 --- a/coderd/x/chatd/chatloop/chatloop_test.go +++ b/coderd/x/chatd/chatloop/chatloop_test.go @@ -2,6 +2,7 @@ package chatloop //nolint:testpackage // Uses internal symbols. import ( "context" + "encoding/base64" "errors" "iter" "strings" @@ -9,13 +10,16 @@ import ( "sync/atomic" "testing" "time" + "unicode/utf8" "charm.land/fantasy" fantasyanthropic "charm.land/fantasy/providers/anthropic" + "github.com/prometheus/client_golang/prometheus" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "golang.org/x/xerrors" + "cdr.dev/slog/v3" "github.com/coder/coder/v2/coderd/x/chatd/chaterror" "github.com/coder/coder/v2/coderd/x/chatd/chatretry" "github.com/coder/coder/v2/coderd/x/chatd/chattest" @@ -2115,3 +2119,141 @@ func TestRun_PrepareMessagesOnlyFiresOnce(t *testing.T) { // PrepareMessages is called before each of the 3 steps. require.Equal(t, 3, int(prepareCalls.Load())) } + +func TestExecuteSingleTool_MediaBase64Encoding(t *testing.T) { + t.Parallel() + + originalBytes := []byte{0xFF, 0xD8, 0xFF, 0xE0, 0x00, 0x10} + metrics := NewMetrics(prometheus.NewRegistry()) + logger := slog.Make() + + t.Run("EncodesRawBytesToBase64", func(t *testing.T) { + t.Parallel() + + tool := fantasy.NewAgentTool( + "screenshot", + "takes a screenshot", + func(_ context.Context, _ struct{}, _ fantasy.ToolCall) (fantasy.ToolResponse, error) { + return fantasy.ToolResponse{ + Type: "image", + Data: originalBytes, + MediaType: "image/jpeg", + }, nil + }, + ) + + toolMap := map[string]fantasy.AgentTool{ + "screenshot": tool, + } + tc := fantasy.ToolCallContent{ + ToolCallID: "call-1", + ToolName: "screenshot", + Input: "{}", + } + + result := executeSingleTool( + context.Background(), + toolMap, + tc, + metrics, + logger, + "fake", "fake-model", + map[string]bool{}, + []string{"screenshot"}, + map[string]struct{}{}, + ) + + media, ok := result.Result.(fantasy.ToolResultOutputContentMedia) + require.True(t, ok, "expected ToolResultOutputContentMedia") + require.Equal(t, "image/jpeg", media.MediaType) + + decoded, err := base64.StdEncoding.DecodeString(media.Data) + require.NoError(t, err, "Data should be valid base64") + require.Equal(t, originalBytes, decoded) + }) + + t.Run("SanitizesInvalidUTF8InContent", func(t *testing.T) { + t.Parallel() + + tool := fantasy.NewAgentTool( + "screenshot", + "takes a screenshot", + func(_ context.Context, _ struct{}, _ fantasy.ToolCall) (fantasy.ToolResponse, error) { + return fantasy.ToolResponse{ + Type: "image", + Data: originalBytes, + MediaType: "image/png", + Content: "hello\xffworld", + }, nil + }, + ) + + toolMap := map[string]fantasy.AgentTool{ + "screenshot": tool, + } + tc := fantasy.ToolCallContent{ + ToolCallID: "call-2", + ToolName: "screenshot", + Input: "{}", + } + + result := executeSingleTool( + context.Background(), + toolMap, + tc, + metrics, + logger, + "fake", "fake-model", + map[string]bool{}, + []string{"screenshot"}, + map[string]struct{}{}, + ) + + media, ok := result.Result.(fantasy.ToolResultOutputContentMedia) + require.True(t, ok, "expected ToolResultOutputContentMedia") + require.True(t, utf8.ValidString(media.Text), "Text should be valid UTF-8") + require.Contains(t, media.Text, "hello") + require.Contains(t, media.Text, "world") + }) + + t.Run("SanitizesInvalidUTF8InTextResult", func(t *testing.T) { + t.Parallel() + + tool := fantasy.NewAgentTool( + "echo", + "echoes input", + func(_ context.Context, _ struct{}, _ fantasy.ToolCall) (fantasy.ToolResponse, error) { + return fantasy.ToolResponse{ + Content: "hello\xffworld", + }, nil + }, + ) + + toolMap := map[string]fantasy.AgentTool{ + "echo": tool, + } + tc := fantasy.ToolCallContent{ + ToolCallID: "call-3", + ToolName: "echo", + Input: "{}", + } + + result := executeSingleTool( + context.Background(), + toolMap, + tc, + metrics, + logger, + "fake", "fake-model", + map[string]bool{}, + []string{"echo"}, + map[string]struct{}{}, + ) + + textOutput, ok := result.Result.(fantasy.ToolResultOutputContentText) + require.True(t, ok, "expected ToolResultOutputContentText, got %T", result.Result) + require.True(t, utf8.ValidString(textOutput.Text), "Text should be valid UTF-8") + require.Contains(t, textOutput.Text, "hello") + require.Contains(t, textOutput.Text, "world") + }) +} diff --git a/coderd/x/chatd/chatprompt/chatprompt.go b/coderd/x/chatd/chatprompt/chatprompt.go index 452fd77144..e6e855e8f8 100644 --- a/coderd/x/chatd/chatprompt/chatprompt.go +++ b/coderd/x/chatd/chatprompt/chatprompt.go @@ -3,6 +3,7 @@ package chatprompt import ( "bytes" "context" + "encoding/base64" "encoding/json" "fmt" "mime" @@ -703,7 +704,7 @@ func MarshalToolResult(toolCallID, toolName string, result json.RawMessage, isEr // PartFromContent converts fantasy content into a SDK chat message // part, preserving ProviderMetadata and ProviderExecuted fields. func PartFromContent(block fantasy.Content) codersdk.ChatMessagePart { - return sdkPartFromContent(block, nil) + return sdkPartFromContent(slog.Logger{}, block, nil) } // PartFromContentWithLogger is for call sites that can surface malformed @@ -713,7 +714,7 @@ func PartFromContentWithLogger( logger slog.Logger, block fantasy.Content, ) codersdk.ChatMessagePart { - return sdkPartFromContent(block, func(content fantasy.ToolResultContent, err error) { + return sdkPartFromContent(logger, block, func(content fantasy.ToolResultContent, err error) { logger.Warn(ctx, "skipping malformed tool attachment metadata", slog.F("tool_name", content.ToolName), slog.F("tool_call_id", content.ToolCallID), @@ -723,6 +724,7 @@ func PartFromContentWithLogger( } func sdkPartFromContent( + logger slog.Logger, block fantasy.Content, logMalformedAttachmentMetadata func(fantasy.ToolResultContent, error), ) codersdk.ChatMessagePart { @@ -800,9 +802,9 @@ func sdkPartFromContent( ProviderMetadata: marshalProviderMetadata(value.ProviderMetadata), } case fantasy.ToolResultContent: - return toolResultContentToPart(value, logMalformedAttachmentMetadata) + return toolResultContentToPart(logger, value, logMalformedAttachmentMetadata) case *fantasy.ToolResultContent: - return toolResultContentToPart(*value, logMalformedAttachmentMetadata) + return toolResultContentToPart(logger, *value, logMalformedAttachmentMetadata) default: return codersdk.ChatMessagePart{} } @@ -819,6 +821,7 @@ func ToolResultToPart(toolCallID, toolName string, result json.RawMessage, isErr // toolResultContentToPart converts a fantasy ToolResultContent into a // ChatMessagePart. func toolResultContentToPart( + logger slog.Logger, content fantasy.ToolResultContent, logMalformedAttachmentMetadata func(fantasy.ToolResultContent, error), ) codersdk.ChatMessagePart { @@ -834,23 +837,42 @@ func toolResultContentToPart( if isSubagentLifecycleToolName(content.ToolName) && hasErrorField(raw) { result = raw } else { - result, _ = json.Marshal(map[string]any{"error": output.Error.Error()}) + var marshalErr error + result, marshalErr = json.Marshal(map[string]any{"error": output.Error.Error()}) + if marshalErr != nil { + logger.Error(context.Background(), "failed to marshal error tool result", + slog.F("tool_name", content.ToolName), + slog.F("tool_call_id", content.ToolCallID), + slog.Error(marshalErr), + ) + result = []byte(`{"error":"marshal failure"}`) + } } } else { result = []byte(`{"error":""}`) } case fantasy.ToolResultOutputContentText: - result = json.RawMessage(output.Text) + sanitized := strings.ToValidUTF8(output.Text, "\uFFFD") + result = json.RawMessage(sanitized) // Ensure valid JSON; wrap in an object if not. if !json.Valid(result) { - result, _ = json.Marshal(map[string]any{"output": output.Text}) + var marshalErr error + result, marshalErr = json.Marshal(map[string]any{"output": sanitized}) + if marshalErr != nil { + logger.Error(context.Background(), "failed to marshal text tool result", + slog.F("tool_name", content.ToolName), + slog.F("tool_call_id", content.ToolCallID), + slog.Error(marshalErr), + ) + result = []byte(`{}`) + } } case fantasy.ToolResultOutputContentMedia: isMedia = true persisted := persistedMediaResult{ Data: output.Data, MimeType: output.MediaType, - Text: output.Text, + Text: strings.ToValidUTF8(output.Text, "\uFFFD"), } // Tool renderers only receive the persisted result JSON, while // ClientMetadata is consumed later to append sibling file parts. @@ -1314,6 +1336,10 @@ func toolResultPartToMessagePart(logger slog.Logger, part codersdk.ChatMessagePa if extracted := extractErrorString(part.Result); extracted != "" { message = extracted } + // Sanitize before wrapping in an error so that invalid + // byte sequences from tool output do not propagate into + // the LLM message stream. + message = strings.ToValidUTF8(message, "\uFFFD") return fantasy.ToolResultPart{ ToolCallID: toolCallID, ProviderExecuted: part.ProviderExecuted, @@ -1336,38 +1362,59 @@ func toolResultPartToMessagePart(logger slog.Logger, part codersdk.ChatMessagePa var media persistedMediaResult unmarshalErr := json.Unmarshal(part.Result, &media) if unmarshalErr == nil && media.Data != "" && media.MimeType != "" { - return fantasy.ToolResultPart{ - ToolCallID: toolCallID, - ProviderExecuted: part.ProviderExecuted, - Output: fantasy.ToolResultOutputContentMedia{ - Data: media.Data, - MediaType: media.MimeType, - Text: media.Text, - }, - ProviderOptions: opts, + _, decErr := base64.StdEncoding.DecodeString(media.Data) + if decErr == nil { + return fantasy.ToolResultPart{ + ToolCallID: toolCallID, + ProviderExecuted: part.ProviderExecuted, + Output: fantasy.ToolResultOutputContentMedia{ + Data: media.Data, + MediaType: media.MimeType, + Text: strings.ToValidUTF8(media.Text, "\uFFFD"), + }, + ProviderOptions: opts, + } } + // Base64 invalid. Use the human-readable annotation + // instead of the full JSON blob to preserve context. + logger.Warn(context.Background(), + "tool result not valid base64, falling through to text", + slog.F("tool_call_id", toolCallID), + slog.F("mime_type", media.MimeType), + slog.Error(decErr), + ) + if media.Text != "" { + resultText = strings.ToValidUTF8(media.Text, "\uFFFD") + } else { + resultText = "[media content unavailable: corrupted data]" + } + } else { + // Generic warning: unmarshal failure or missing fields. + fields := []slog.Field{ + slog.F("tool_call_id", toolCallID), + slog.F("tool_name", part.ToolName), + slog.F("has_data", media.Data != ""), + slog.F("has_mime_type", media.MimeType != ""), + } + if unmarshalErr != nil { + fields = append(fields, slog.Error(unmarshalErr)) + } + logger.Warn(context.Background(), + "media tool result failed reconstruction, falling through to text", + fields..., + ) } - - fields := []slog.Field{ - slog.F("tool_call_id", toolCallID), - slog.F("tool_name", part.ToolName), - slog.F("has_data", media.Data != ""), - slog.F("has_mime_type", media.MimeType != ""), - } - if unmarshalErr != nil { - fields = append(fields, slog.Error(unmarshalErr)) - } - logger.Warn(context.Background(), - "media tool result failed reconstruction, falling through to text", - fields..., - ) } + // Sanitize invalid UTF-8 in text results before sending + // to the LLM. This repairs stored messages that were + // poisoned by raw binary in tool results. + sanitizedResult := strings.ToValidUTF8(resultText, "\uFFFD") return fantasy.ToolResultPart{ ToolCallID: toolCallID, ProviderExecuted: part.ProviderExecuted, Output: fantasy.ToolResultOutputContentText{ - Text: resultText, + Text: sanitizedResult, }, ProviderOptions: opts, } diff --git a/coderd/x/chatd/chatprompt/chatprompt_test.go b/coderd/x/chatd/chatprompt/chatprompt_test.go index 06c707d7b8..0db9c5d5b2 100644 --- a/coderd/x/chatd/chatprompt/chatprompt_test.go +++ b/coderd/x/chatd/chatprompt/chatprompt_test.go @@ -6,6 +6,7 @@ import ( "encoding/json" "strings" "testing" + "unicode/utf8" "charm.land/fantasy" fantasyanthropic "charm.land/fantasy/providers/anthropic" @@ -2781,3 +2782,215 @@ func TestPartFromContent_CreatedAtNotStamped(t *testing.T) { assert.Nil(t, part.CreatedAt) }) } + +func TestToolResultAntivenom(t *testing.T) { + t.Parallel() + logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) + + t.Run("PoisonedTextResultSanitized", func(t *testing.T) { + t.Parallel() + + // Simulate raw binary bytes stored as json.RawMessage. + // This reproduces the crash where tool output containing + // invalid UTF-8 was passed verbatim to the LLM provider. + poisonedBytes := json.RawMessage(string([]byte{0xFF, 0xD8, 0xFF})) + part := codersdk.ChatMessagePart{ + Type: codersdk.ChatMessagePartTypeToolResult, + ToolCallID: "call-1", + ToolName: "test_tool", + Result: poisonedBytes, + IsError: false, + IsMedia: false, + } + + result := chatprompt.ToolResultPartToMessagePartForTest(logger, part) + + textOutput, ok := fantasy.AsToolResultOutputType[fantasy.ToolResultOutputContentText](result.Output) + require.True(t, ok, "expected text output, got %T", result.Output) + require.True(t, utf8.ValidString(textOutput.Text), "output text must be valid UTF-8") + require.NotEmpty(t, textOutput.Text) + }) + + t.Run("PoisonedMediaResultDegradesToText", func(t *testing.T) { + t.Parallel() + + // Simulate raw JPEG bytes stored where base64 is expected. + // The base64 validation guard should reject this and fall + // through to the text path. + corruptedData := string([]byte{0xFF, 0xD8, 0xFF, 0xE0}) + media := struct { + Data string `json:"data"` + MimeType string `json:"mime_type"` + Text string `json:"text,omitempty"` + }{ + Data: corruptedData, + MimeType: "image/jpeg", + } + mediaJSON, err := json.Marshal(media) + require.NoError(t, err) + + part := codersdk.ChatMessagePart{ + Type: codersdk.ChatMessagePartTypeToolResult, + ToolCallID: "call-2", + ToolName: "computer", + Result: json.RawMessage(mediaJSON), + IsError: false, + IsMedia: true, + } + + result := chatprompt.ToolResultPartToMessagePartForTest(logger, part) + + // Should degrade to text since the data is not valid base64. + _, isMedia := fantasy.AsToolResultOutputType[fantasy.ToolResultOutputContentMedia](result.Output) + require.False(t, isMedia, "corrupted media should not be returned as media") + + textOutput, isText := fantasy.AsToolResultOutputType[fantasy.ToolResultOutputContentText](result.Output) + require.True(t, isText, "should fall through to text, got %T", result.Output) + require.True(t, utf8.ValidString(textOutput.Text), "fallback text must be valid UTF-8") + }) + + t.Run("ValidMediaResultRoundTrips", func(t *testing.T) { + t.Parallel() + + // Valid base64 media should pass through the guard and + // be returned as ToolResultOutputContentMedia. + validBase64 := "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAAC0lEQVQI12NgAAIABQAB" + media := struct { + Data string `json:"data"` + MimeType string `json:"mime_type"` + Text string `json:"text,omitempty"` + }{ + Data: validBase64, + MimeType: "image/png", + Text: "screenshot", + } + mediaJSON, err := json.Marshal(media) + require.NoError(t, err) + + part := codersdk.ChatMessagePart{ + Type: codersdk.ChatMessagePartTypeToolResult, + ToolCallID: "call-3", + ToolName: "computer", + Result: json.RawMessage(mediaJSON), + IsError: false, + IsMedia: true, + } + + result := chatprompt.ToolResultPartToMessagePartForTest(logger, part) + + mediaOutput, ok := fantasy.AsToolResultOutputType[fantasy.ToolResultOutputContentMedia](result.Output) + require.True(t, ok, "valid media should round-trip as media, got %T", result.Output) + require.Equal(t, validBase64, mediaOutput.Data) + require.Equal(t, "image/png", mediaOutput.MediaType) + require.Equal(t, "screenshot", mediaOutput.Text) + }) + + t.Run("MediaWithInvalidUTF8TextSanitized", func(t *testing.T) { + t.Parallel() + + // Valid base64 data with an invalid UTF-8 text annotation. + // The media should survive but the text field must be + // sanitized to valid UTF-8. + validBase64 := "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAAC0lEQVQI12NgAAIABQAB" + invalidText := "hello" + string([]byte{0xFF, 0xFE}) + "world" + media := struct { + Data string `json:"data"` + MimeType string `json:"mime_type"` + Text string `json:"text,omitempty"` + }{ + Data: validBase64, + MimeType: "image/png", + Text: invalidText, + } + mediaJSON, err := json.Marshal(media) + require.NoError(t, err) + + part := codersdk.ChatMessagePart{ + Type: codersdk.ChatMessagePartTypeToolResult, + ToolCallID: "call-4", + ToolName: "computer", + Result: json.RawMessage(mediaJSON), + IsError: false, + IsMedia: true, + } + + result := chatprompt.ToolResultPartToMessagePartForTest(logger, part) + + mediaOutput, ok := fantasy.AsToolResultOutputType[fantasy.ToolResultOutputContentMedia](result.Output) + require.True(t, ok, "media with valid base64 should stay as media, got %T", result.Output) + require.Equal(t, validBase64, mediaOutput.Data) + require.True(t, utf8.ValidString(mediaOutput.Text), "text must be sanitized to valid UTF-8") + require.Contains(t, mediaOutput.Text, "hello") + require.Contains(t, mediaOutput.Text, "world") + }) + + t.Run("PoisonedErrorResultSanitized", func(t *testing.T) { + t.Parallel() + // Simulate invalid UTF-8 in an error tool result. + poisonedError := json.RawMessage(`{"error":"fail` + string([]byte{0xFF, 0xFE}) + `ed"}`) + part := codersdk.ChatMessagePart{ + Type: codersdk.ChatMessagePartTypeToolResult, + ToolCallID: "call-5", + ToolName: "broken_tool", + Result: poisonedError, + IsError: true, + IsMedia: false, + } + + result := chatprompt.ToolResultPartToMessagePartForTest(logger, part) + + errOutput, ok := fantasy.AsToolResultOutputType[fantasy.ToolResultOutputContentError](result.Output) + require.True(t, ok, "expected error output, got %T", result.Output) + require.True(t, utf8.ValidString(errOutput.Error.Error()), + "error message must be valid UTF-8") + require.Contains(t, errOutput.Error.Error(), "fail") + require.Contains(t, errOutput.Error.Error(), "ed") + }) +} + +func TestToolResultContentToPart_UTF8Sanitization(t *testing.T) { + t.Parallel() + logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) + + t.Run("TextWithInvalidUTF8", func(t *testing.T) { + t.Parallel() + part := chatprompt.ToolResultContentToPartForTest(logger, fantasy.ToolResultContent{ + ToolCallID: "call-1", + ToolName: "test", + Result: fantasy.ToolResultOutputContentText{ + Text: "hello\xffworld", + }, + }) + + require.True(t, utf8.Valid(part.Result), + "persisted result must be valid UTF-8, got: %q", string(part.Result)) + }) + + t.Run("MediaTextWithInvalidUTF8", func(t *testing.T) { + t.Parallel() + validBase64 := "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAAC0lEQVQI12NgAAIABQAB" + part := chatprompt.ToolResultContentToPartForTest(logger, fantasy.ToolResultContent{ + ToolCallID: "call-2", + ToolName: "computer", + Result: fantasy.ToolResultOutputContentMedia{ + Data: validBase64, + MediaType: "image/png", + Text: "screenshot\xfe\xffdone", + }, + }) + + require.True(t, part.IsMedia) + // Unmarshal the persisted media and check Text field. + var media struct { + Data string `json:"data"` + MimeType string `json:"mime_type"` + Text string `json:"text"` + } + err := json.Unmarshal(part.Result, &media) + require.NoError(t, err) + require.True(t, utf8.ValidString(media.Text), + "persisted media text must be valid UTF-8") + require.Contains(t, media.Text, "screenshot") + require.Contains(t, media.Text, "done") + }) +} diff --git a/coderd/x/chatd/chatprompt/export_test.go b/coderd/x/chatd/chatprompt/export_test.go index 63c3e1f9ca..588664a0a7 100644 --- a/coderd/x/chatd/chatprompt/export_test.go +++ b/coderd/x/chatd/chatprompt/export_test.go @@ -1,4 +1,21 @@ package chatprompt +import ( + "charm.land/fantasy" + + "cdr.dev/slog/v3" + "github.com/coder/coder/v2/codersdk" +) + // IsSyntheticPasteForTest exposes isSyntheticPaste for external tests. var IsSyntheticPasteForTest = isSyntheticPaste + +// ToolResultPartToMessagePartForTest exposes toolResultPartToMessagePart +// for external tests. +var ToolResultPartToMessagePartForTest = toolResultPartToMessagePart + +// ToolResultContentToPartForTest exposes toolResultContentToPart +// for external tests. +var ToolResultContentToPartForTest = func(logger slog.Logger, content fantasy.ToolResultContent) codersdk.ChatMessagePart { + return toolResultContentToPart(logger, content, nil) +} diff --git a/coderd/x/chatd/chattool/computeruse.go b/coderd/x/chatd/chattool/computeruse.go index 7f65d44fb9..1b7a639148 100644 --- a/coderd/x/chatd/chattool/computeruse.go +++ b/coderd/x/chatd/chattool/computeruse.go @@ -2,6 +2,7 @@ package chattool import ( "context" + "encoding/base64" "fmt" "time" @@ -168,7 +169,7 @@ func (t *computerUseTool) Run(ctx context.Context, call fantasy.ToolCall) (fanta return t.captureScreenshot(ctx, conn, declaredWidth, declaredHeight) } -func (*computerUseTool) captureScreenshot( +func (t *computerUseTool) captureScreenshot( ctx context.Context, conn workspacesdk.AgentConn, declaredWidth, declaredHeight int, @@ -179,7 +180,16 @@ func (*computerUseTool) captureScreenshot( fmt.Sprintf("screenshot failed: %v", err), ), nil } - return fantasy.NewImageResponse([]byte(screenResp.ScreenshotData), "image/png"), nil + screenData, err := base64.StdEncoding.DecodeString(screenResp.ScreenshotData) + if err != nil { + t.logger.Error(ctx, "failed to decode screenshot base64 in captureScreenshot", + slog.Error(err), + ) + return fantasy.NewTextErrorResponse( + fmt.Sprintf("failed to decode screenshot data: %v", err), + ), nil + } + return fantasy.NewImageResponse(screenData, "image/png"), nil } func (t *computerUseTool) captureSharedScreenshot( @@ -194,22 +204,33 @@ func (t *computerUseTool) captureSharedScreenshot( ), nil } + screenData, err := base64.StdEncoding.DecodeString(screenResp.ScreenshotData) + if err != nil { + t.logger.Error(ctx, "failed to decode screenshot base64 in captureSharedScreenshot", + slog.Error(err), + ) + return fantasy.NewTextErrorResponse( + fmt.Sprintf("failed to decode screenshot data: %v", err), + ), nil + } + attachmentName := fmt.Sprintf( "screenshot-%s.png", t.clock.Now().UTC().Format("2006-01-02T15-04-05Z"), ) if t.storeFile == nil { t.logger.Warn(ctx, "screenshot attachment storage is not configured") - return fantasy.NewImageResponse([]byte(screenResp.ScreenshotData), "image/png"), nil + return fantasy.NewImageResponse(screenData, "image/png"), nil } + response := fantasy.NewImageResponse(screenData, "image/png") + attachment, err := storeScreenshotAttachment( ctx, t.storeFile, attachmentName, screenResp.ScreenshotData, ) - response := fantasy.NewImageResponse([]byte(screenResp.ScreenshotData), "image/png") if err != nil { t.logger.Warn(ctx, "failed to persist screenshot attachment", slog.F("attachment_name", attachmentName), diff --git a/coderd/x/chatd/chattool/computeruse_test.go b/coderd/x/chatd/chattool/computeruse_test.go index 4ba5308d2d..fa92a55449 100644 --- a/coderd/x/chatd/chattool/computeruse_test.go +++ b/coderd/x/chatd/chattool/computeruse_test.go @@ -70,7 +70,9 @@ func TestComputerUseTool_Run_Screenshot(t *testing.T) { require.NoError(t, err) assert.Equal(t, "image", resp.Type) assert.Equal(t, "image/png", resp.MediaType) - assert.Equal(t, []byte("iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGP4n539HwAHFwLVF8kc1wAAAABJRU5ErkJggg=="), resp.Data) + expectedBinary, decErr := base64.StdEncoding.DecodeString("iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGP4n539HwAHFwLVF8kc1wAAAABJRU5ErkJggg==") + require.NoError(t, decErr) + assert.Equal(t, expectedBinary, resp.Data) assert.False(t, resp.IsError) } @@ -118,7 +120,9 @@ func TestComputerUseTool_Run_Screenshot_PersistsAttachment(t *testing.T) { require.NoError(t, err) assert.Equal(t, "image", resp.Type) assert.Equal(t, "image/png", resp.MediaType) - assert.Equal(t, []byte(screenshotPNG), resp.Data) + expectedBinary, decErr := base64.StdEncoding.DecodeString(screenshotPNG) + require.NoError(t, decErr) + assert.Equal(t, expectedBinary, resp.Data) assert.Contains(t, storedName, "screenshot-") assert.Equal(t, "image/png", storedType) expectedPNG, decodeErr := base64.StdEncoding.DecodeString(screenshotPNG) @@ -200,7 +204,9 @@ func TestComputerUseTool_Run_Screenshot_OversizedAttachmentFallsBackToImage(t *t assert.Equal(t, "image", resp.Type) assert.Equal(t, "image/png", resp.MediaType) assert.False(t, resp.IsError) - require.Len(t, resp.Data, len(oversizedScreenshot)) + expectedOversized, decErr := base64.StdEncoding.DecodeString(oversizedScreenshot) + require.NoError(t, decErr) + require.Len(t, resp.Data, len(expectedOversized)) attachments, err := chattool.AttachmentsFromMetadata(resp.Metadata) require.NoError(t, err) assert.Empty(t, attachments) @@ -260,7 +266,9 @@ func TestComputerUseTool_Run_LeftClick(t *testing.T) { resp, err := tool.Run(context.Background(), call) require.NoError(t, err) assert.Equal(t, "image", resp.Type) - assert.Equal(t, []byte(followUpScreenshot), resp.Data) + expectedBinary, decErr := base64.StdEncoding.DecodeString(followUpScreenshot) + require.NoError(t, decErr) + assert.Equal(t, expectedBinary, resp.Data) attachments, err := chattool.AttachmentsFromMetadata(resp.Metadata) require.NoError(t, err) assert.Empty(t, attachments) @@ -307,13 +315,73 @@ func TestComputerUseTool_Run_Wait(t *testing.T) { require.NoError(t, err) assert.Equal(t, "image", resp.Type) assert.Equal(t, "image/png", resp.MediaType) - assert.Equal(t, []byte(followUpScreenshot), resp.Data) + expectedBinary, decErr := base64.StdEncoding.DecodeString(followUpScreenshot) + require.NoError(t, decErr) + assert.Equal(t, expectedBinary, resp.Data) assert.False(t, resp.IsError) attachments, err := chattool.AttachmentsFromMetadata(resp.Metadata) require.NoError(t, err) assert.Empty(t, attachments) } +func TestComputerUseTool_Run_ScreenshotDataIsDecodedBinary(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + mockConn := agentconnmock.NewMockAgentConn(ctrl) + geometry := workspacesdk.DefaultDesktopGeometry() + + // A known base64 string (1x1 red PNG). + const screenshotBase64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGP4z8BQDwAEgAF/pooBPQAAAABJRU5ErkJggg==" + + mockConn.EXPECT().ExecuteDesktopAction( + gomock.Any(), + gomock.AssignableToTypeOf(workspacesdk.DesktopAction{}), + ).Return(workspacesdk.DesktopActionResponse{ + Output: "screenshot", + ScreenshotData: screenshotBase64, + ScreenshotWidth: geometry.DeclaredWidth, + ScreenshotHeight: geometry.DeclaredHeight, + }, nil) + + tool := chattool.NewComputerUseTool( + geometry.DeclaredWidth, + geometry.DeclaredHeight, + func(_ context.Context) (workspacesdk.AgentConn, error) { + return mockConn, nil + }, + nil, + quartz.NewReal(), + slogtest.Make(t, nil), + ) + + call := fantasy.ToolCall{ + ID: "test-decode-1", + Name: "computer", + Input: `{"action":"screenshot"}`, + } + + resp, err := tool.Run(context.Background(), call) + require.NoError(t, err) + + assert.Equal(t, "image", resp.Type) + assert.Equal(t, "image/png", resp.MediaType) + + // Data must contain decoded binary, not the base64 string + // reinterpreted as bytes. + expectedBinary, err := base64.StdEncoding.DecodeString(screenshotBase64) + require.NoError(t, err) + assert.Equal(t, expectedBinary, resp.Data, + "ToolResponse.Data should contain decoded binary, not base64-as-bytes") + + // Verify that re-encoding produces the original base64 string. + // This is the round-trip that the chat loop performs when + // building the API response. + reEncoded := base64.StdEncoding.EncodeToString(resp.Data) + assert.Equal(t, screenshotBase64, reEncoded, + "re-encoding Data should produce the original base64 string (no double-encode)") +} + func TestComputerUseTool_Run_ConnError(t *testing.T) { t.Parallel() diff --git a/coderd/x/chatd/chattool/mcpworkspace.go b/coderd/x/chatd/chattool/mcpworkspace.go index dc175ab43e..5748c5b03f 100644 --- a/coderd/x/chatd/chattool/mcpworkspace.go +++ b/coderd/x/chatd/chattool/mcpworkspace.go @@ -107,7 +107,7 @@ func convertMCPToolResponse( for _, c := range resp.Content { switch c.Type { case "text": - textParts = append(textParts, c.Text) + textParts = append(textParts, strings.ToValidUTF8(c.Text, "\uFFFD")) case "image", "audio": if c.Data == "" { continue @@ -129,7 +129,7 @@ func convertMCPToolResponse( binaryResult = &r } default: - textParts = append(textParts, c.Text) + textParts = append(textParts, strings.ToValidUTF8(c.Text, "\uFFFD")) } } diff --git a/coderd/x/chatd/mcpclient/export_test.go b/coderd/x/chatd/mcpclient/export_test.go new file mode 100644 index 0000000000..dbdca8c638 --- /dev/null +++ b/coderd/x/chatd/mcpclient/export_test.go @@ -0,0 +1,5 @@ +package mcpclient + +// ConvertCallResultForTest exposes convertCallResult for external +// tests. +var ConvertCallResultForTest = convertCallResult diff --git a/coderd/x/chatd/mcpclient/mcpclient.go b/coderd/x/chatd/mcpclient/mcpclient.go index e91d76b9e0..ce6755cfa6 100644 --- a/coderd/x/chatd/mcpclient/mcpclient.go +++ b/coderd/x/chatd/mcpclient/mcpclient.go @@ -607,7 +607,7 @@ func convertCallResult( for _, item := range result.Content { switch c := item.(type) { case mcp.TextContent: - textParts = append(textParts, c.Text) + textParts = append(textParts, strings.ToValidUTF8(c.Text, "\uFFFD")) case mcp.ImageContent: data, err := base64.StdEncoding.DecodeString( c.Data, @@ -653,7 +653,7 @@ func convertCallResult( // regardless of form. switch r := c.Resource.(type) { case mcp.TextResourceContents: - textParts = append(textParts, r.Text) + textParts = append(textParts, strings.ToValidUTF8(r.Text, "\uFFFD")) case mcp.BlobResourceContents: data, err := base64.StdEncoding.DecodeString( r.Blob, diff --git a/coderd/x/chatd/mcpclient/mcpclient_test.go b/coderd/x/chatd/mcpclient/mcpclient_test.go index c0550d5ae4..7676296360 100644 --- a/coderd/x/chatd/mcpclient/mcpclient_test.go +++ b/coderd/x/chatd/mcpclient/mcpclient_test.go @@ -9,6 +9,7 @@ import ( "sync" "testing" "time" + "unicode/utf8" "charm.land/fantasy" "github.com/google/uuid" @@ -1271,3 +1272,77 @@ func TestModelIntent_Run_FallbackOnBadJSON(t *testing.T) { require.NoError(t, err) assert.True(t, resp.IsError, "malformed input should produce an error response") } + +func TestConvertCallResult_UTF8Sanitization(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + result *mcp.CallToolResult + wantContains []string + }{ + { + name: "InvalidUTF8InTextContent", + result: &mcp.CallToolResult{ + Content: []mcp.Content{ + mcp.TextContent{ + Text: "Hello" + string([]byte{0xFF, 0xFE, 0x80}) + "World", + }, + }, + }, + wantContains: []string{"Hello", "World", "\uFFFD"}, + }, + { + name: "InvalidUTF8InEmbeddedResourceText", + result: &mcp.CallToolResult{ + Content: []mcp.Content{ + mcp.EmbeddedResource{ + Resource: mcp.TextResourceContents{ + Text: "Content" + string([]byte{0x80, 0x81, 0x82}), + }, + }, + }, + }, + wantContains: []string{"Content"}, + }, + { + name: "ValidUTF8PassesThrough", + result: &mcp.CallToolResult{ + Content: []mcp.Content{ + mcp.TextContent{ + Text: "Hello, δΈ–η•Œ! 🌍", + }, + }, + }, + wantContains: []string{"Hello, δΈ–η•Œ! 🌍"}, + }, + { + name: "MultipleTextPartsAllSanitized", + result: &mcp.CallToolResult{ + Content: []mcp.Content{ + mcp.TextContent{ + Text: "Part1" + string([]byte{0xFF}), + }, + mcp.TextContent{ + Text: "Part2" + string([]byte{0xFE}), + }, + }, + }, + wantContains: []string{"Part1", "Part2"}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + resp := mcpclient.ConvertCallResultForTest(tt.result) + + require.True(t, utf8.ValidString(resp.Content), + "response content must be valid UTF-8") + for _, want := range tt.wantContains { + require.Contains(t, resp.Content, want) + } + }) + } +}