mirror of
https://github.com/coder/coder.git
synced 2026-10-29 17:44:33 +08:00
fix(coderd/x/chatd): prevent invalid tool results from poisoning chat history (#24663)
- **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
> 🤖
This commit is contained in:
@@ -6177,7 +6177,7 @@ func TestComputerUseSubagentToolsAndModel(t *testing.T) {
|
|||||||
Return(workspacesdk.DesktopActionResponse{
|
Return(workspacesdk.DesktopActionResponse{
|
||||||
ScreenshotWidth: 1920,
|
ScreenshotWidth: 1920,
|
||||||
ScreenshotHeight: 1080,
|
ScreenshotHeight: 1080,
|
||||||
ScreenshotData: "iVBOR",
|
ScreenshotData: "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGP4n539HwAHFwLVF8kc1wAAAABJRU5ErkJggg==",
|
||||||
}, nil).
|
}, nil).
|
||||||
AnyTimes()
|
AnyTimes()
|
||||||
mockConn.EXPECT().
|
mockConn.EXPECT().
|
||||||
|
|||||||
@@ -3,6 +3,7 @@ package chatloop
|
|||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
"database/sql"
|
"database/sql"
|
||||||
|
"encoding/base64"
|
||||||
"encoding/json"
|
"encoding/json"
|
||||||
"errors"
|
"errors"
|
||||||
"maps"
|
"maps"
|
||||||
@@ -1213,13 +1214,13 @@ func executeSingleTool(
|
|||||||
)
|
)
|
||||||
case resp.Type == "image" || resp.Type == "media":
|
case resp.Type == "image" || resp.Type == "media":
|
||||||
result.Result = fantasy.ToolResultOutputContentMedia{
|
result.Result = fantasy.ToolResultOutputContentMedia{
|
||||||
Data: string(resp.Data),
|
Data: base64.StdEncoding.EncodeToString(resp.Data),
|
||||||
MediaType: resp.MediaType,
|
MediaType: resp.MediaType,
|
||||||
Text: resp.Content,
|
Text: strings.ToValidUTF8(resp.Content, "\uFFFD"),
|
||||||
}
|
}
|
||||||
default:
|
default:
|
||||||
result.Result = fantasy.ToolResultOutputContentText{
|
result.Result = fantasy.ToolResultOutputContentText{
|
||||||
Text: resp.Content,
|
Text: strings.ToValidUTF8(resp.Content, "\uFFFD"),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
return result
|
return result
|
||||||
|
|||||||
@@ -2,6 +2,7 @@ package chatloop //nolint:testpackage // Uses internal symbols.
|
|||||||
|
|
||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
|
"encoding/base64"
|
||||||
"errors"
|
"errors"
|
||||||
"iter"
|
"iter"
|
||||||
"strings"
|
"strings"
|
||||||
@@ -9,13 +10,16 @@ import (
|
|||||||
"sync/atomic"
|
"sync/atomic"
|
||||||
"testing"
|
"testing"
|
||||||
"time"
|
"time"
|
||||||
|
"unicode/utf8"
|
||||||
|
|
||||||
"charm.land/fantasy"
|
"charm.land/fantasy"
|
||||||
fantasyanthropic "charm.land/fantasy/providers/anthropic"
|
fantasyanthropic "charm.land/fantasy/providers/anthropic"
|
||||||
|
"github.com/prometheus/client_golang/prometheus"
|
||||||
"github.com/stretchr/testify/assert"
|
"github.com/stretchr/testify/assert"
|
||||||
"github.com/stretchr/testify/require"
|
"github.com/stretchr/testify/require"
|
||||||
"golang.org/x/xerrors"
|
"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/chaterror"
|
||||||
"github.com/coder/coder/v2/coderd/x/chatd/chatretry"
|
"github.com/coder/coder/v2/coderd/x/chatd/chatretry"
|
||||||
"github.com/coder/coder/v2/coderd/x/chatd/chattest"
|
"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.
|
// PrepareMessages is called before each of the 3 steps.
|
||||||
require.Equal(t, 3, int(prepareCalls.Load()))
|
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")
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|||||||
@@ -3,6 +3,7 @@ package chatprompt
|
|||||||
import (
|
import (
|
||||||
"bytes"
|
"bytes"
|
||||||
"context"
|
"context"
|
||||||
|
"encoding/base64"
|
||||||
"encoding/json"
|
"encoding/json"
|
||||||
"fmt"
|
"fmt"
|
||||||
"mime"
|
"mime"
|
||||||
@@ -703,7 +704,7 @@ func MarshalToolResult(toolCallID, toolName string, result json.RawMessage, isEr
|
|||||||
// PartFromContent converts fantasy content into a SDK chat message
|
// PartFromContent converts fantasy content into a SDK chat message
|
||||||
// part, preserving ProviderMetadata and ProviderExecuted fields.
|
// part, preserving ProviderMetadata and ProviderExecuted fields.
|
||||||
func PartFromContent(block fantasy.Content) codersdk.ChatMessagePart {
|
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
|
// PartFromContentWithLogger is for call sites that can surface malformed
|
||||||
@@ -713,7 +714,7 @@ func PartFromContentWithLogger(
|
|||||||
logger slog.Logger,
|
logger slog.Logger,
|
||||||
block fantasy.Content,
|
block fantasy.Content,
|
||||||
) codersdk.ChatMessagePart {
|
) 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",
|
logger.Warn(ctx, "skipping malformed tool attachment metadata",
|
||||||
slog.F("tool_name", content.ToolName),
|
slog.F("tool_name", content.ToolName),
|
||||||
slog.F("tool_call_id", content.ToolCallID),
|
slog.F("tool_call_id", content.ToolCallID),
|
||||||
@@ -723,6 +724,7 @@ func PartFromContentWithLogger(
|
|||||||
}
|
}
|
||||||
|
|
||||||
func sdkPartFromContent(
|
func sdkPartFromContent(
|
||||||
|
logger slog.Logger,
|
||||||
block fantasy.Content,
|
block fantasy.Content,
|
||||||
logMalformedAttachmentMetadata func(fantasy.ToolResultContent, error),
|
logMalformedAttachmentMetadata func(fantasy.ToolResultContent, error),
|
||||||
) codersdk.ChatMessagePart {
|
) codersdk.ChatMessagePart {
|
||||||
@@ -800,9 +802,9 @@ func sdkPartFromContent(
|
|||||||
ProviderMetadata: marshalProviderMetadata(value.ProviderMetadata),
|
ProviderMetadata: marshalProviderMetadata(value.ProviderMetadata),
|
||||||
}
|
}
|
||||||
case fantasy.ToolResultContent:
|
case fantasy.ToolResultContent:
|
||||||
return toolResultContentToPart(value, logMalformedAttachmentMetadata)
|
return toolResultContentToPart(logger, value, logMalformedAttachmentMetadata)
|
||||||
case *fantasy.ToolResultContent:
|
case *fantasy.ToolResultContent:
|
||||||
return toolResultContentToPart(*value, logMalformedAttachmentMetadata)
|
return toolResultContentToPart(logger, *value, logMalformedAttachmentMetadata)
|
||||||
default:
|
default:
|
||||||
return codersdk.ChatMessagePart{}
|
return codersdk.ChatMessagePart{}
|
||||||
}
|
}
|
||||||
@@ -819,6 +821,7 @@ func ToolResultToPart(toolCallID, toolName string, result json.RawMessage, isErr
|
|||||||
// toolResultContentToPart converts a fantasy ToolResultContent into a
|
// toolResultContentToPart converts a fantasy ToolResultContent into a
|
||||||
// ChatMessagePart.
|
// ChatMessagePart.
|
||||||
func toolResultContentToPart(
|
func toolResultContentToPart(
|
||||||
|
logger slog.Logger,
|
||||||
content fantasy.ToolResultContent,
|
content fantasy.ToolResultContent,
|
||||||
logMalformedAttachmentMetadata func(fantasy.ToolResultContent, error),
|
logMalformedAttachmentMetadata func(fantasy.ToolResultContent, error),
|
||||||
) codersdk.ChatMessagePart {
|
) codersdk.ChatMessagePart {
|
||||||
@@ -834,23 +837,42 @@ func toolResultContentToPart(
|
|||||||
if isSubagentLifecycleToolName(content.ToolName) && hasErrorField(raw) {
|
if isSubagentLifecycleToolName(content.ToolName) && hasErrorField(raw) {
|
||||||
result = raw
|
result = raw
|
||||||
} else {
|
} 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 {
|
} else {
|
||||||
result = []byte(`{"error":""}`)
|
result = []byte(`{"error":""}`)
|
||||||
}
|
}
|
||||||
case fantasy.ToolResultOutputContentText:
|
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.
|
// Ensure valid JSON; wrap in an object if not.
|
||||||
if !json.Valid(result) {
|
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:
|
case fantasy.ToolResultOutputContentMedia:
|
||||||
isMedia = true
|
isMedia = true
|
||||||
persisted := persistedMediaResult{
|
persisted := persistedMediaResult{
|
||||||
Data: output.Data,
|
Data: output.Data,
|
||||||
MimeType: output.MediaType,
|
MimeType: output.MediaType,
|
||||||
Text: output.Text,
|
Text: strings.ToValidUTF8(output.Text, "\uFFFD"),
|
||||||
}
|
}
|
||||||
// Tool renderers only receive the persisted result JSON, while
|
// Tool renderers only receive the persisted result JSON, while
|
||||||
// ClientMetadata is consumed later to append sibling file parts.
|
// 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 != "" {
|
if extracted := extractErrorString(part.Result); extracted != "" {
|
||||||
message = 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{
|
return fantasy.ToolResultPart{
|
||||||
ToolCallID: toolCallID,
|
ToolCallID: toolCallID,
|
||||||
ProviderExecuted: part.ProviderExecuted,
|
ProviderExecuted: part.ProviderExecuted,
|
||||||
@@ -1336,38 +1362,59 @@ func toolResultPartToMessagePart(logger slog.Logger, part codersdk.ChatMessagePa
|
|||||||
var media persistedMediaResult
|
var media persistedMediaResult
|
||||||
unmarshalErr := json.Unmarshal(part.Result, &media)
|
unmarshalErr := json.Unmarshal(part.Result, &media)
|
||||||
if unmarshalErr == nil && media.Data != "" && media.MimeType != "" {
|
if unmarshalErr == nil && media.Data != "" && media.MimeType != "" {
|
||||||
return fantasy.ToolResultPart{
|
_, decErr := base64.StdEncoding.DecodeString(media.Data)
|
||||||
ToolCallID: toolCallID,
|
if decErr == nil {
|
||||||
ProviderExecuted: part.ProviderExecuted,
|
return fantasy.ToolResultPart{
|
||||||
Output: fantasy.ToolResultOutputContentMedia{
|
ToolCallID: toolCallID,
|
||||||
Data: media.Data,
|
ProviderExecuted: part.ProviderExecuted,
|
||||||
MediaType: media.MimeType,
|
Output: fantasy.ToolResultOutputContentMedia{
|
||||||
Text: media.Text,
|
Data: media.Data,
|
||||||
},
|
MediaType: media.MimeType,
|
||||||
ProviderOptions: opts,
|
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{
|
return fantasy.ToolResultPart{
|
||||||
ToolCallID: toolCallID,
|
ToolCallID: toolCallID,
|
||||||
ProviderExecuted: part.ProviderExecuted,
|
ProviderExecuted: part.ProviderExecuted,
|
||||||
Output: fantasy.ToolResultOutputContentText{
|
Output: fantasy.ToolResultOutputContentText{
|
||||||
Text: resultText,
|
Text: sanitizedResult,
|
||||||
},
|
},
|
||||||
ProviderOptions: opts,
|
ProviderOptions: opts,
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -6,6 +6,7 @@ import (
|
|||||||
"encoding/json"
|
"encoding/json"
|
||||||
"strings"
|
"strings"
|
||||||
"testing"
|
"testing"
|
||||||
|
"unicode/utf8"
|
||||||
|
|
||||||
"charm.land/fantasy"
|
"charm.land/fantasy"
|
||||||
fantasyanthropic "charm.land/fantasy/providers/anthropic"
|
fantasyanthropic "charm.land/fantasy/providers/anthropic"
|
||||||
@@ -2781,3 +2782,215 @@ func TestPartFromContent_CreatedAtNotStamped(t *testing.T) {
|
|||||||
assert.Nil(t, part.CreatedAt)
|
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")
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|||||||
@@ -1,4 +1,21 @@
|
|||||||
package chatprompt
|
package chatprompt
|
||||||
|
|
||||||
|
import (
|
||||||
|
"charm.land/fantasy"
|
||||||
|
|
||||||
|
"cdr.dev/slog/v3"
|
||||||
|
"github.com/coder/coder/v2/codersdk"
|
||||||
|
)
|
||||||
|
|
||||||
// IsSyntheticPasteForTest exposes isSyntheticPaste for external tests.
|
// IsSyntheticPasteForTest exposes isSyntheticPaste for external tests.
|
||||||
var IsSyntheticPasteForTest = isSyntheticPaste
|
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)
|
||||||
|
}
|
||||||
|
|||||||
@@ -2,6 +2,7 @@ package chattool
|
|||||||
|
|
||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
|
"encoding/base64"
|
||||||
"fmt"
|
"fmt"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
@@ -168,7 +169,7 @@ func (t *computerUseTool) Run(ctx context.Context, call fantasy.ToolCall) (fanta
|
|||||||
return t.captureScreenshot(ctx, conn, declaredWidth, declaredHeight)
|
return t.captureScreenshot(ctx, conn, declaredWidth, declaredHeight)
|
||||||
}
|
}
|
||||||
|
|
||||||
func (*computerUseTool) captureScreenshot(
|
func (t *computerUseTool) captureScreenshot(
|
||||||
ctx context.Context,
|
ctx context.Context,
|
||||||
conn workspacesdk.AgentConn,
|
conn workspacesdk.AgentConn,
|
||||||
declaredWidth, declaredHeight int,
|
declaredWidth, declaredHeight int,
|
||||||
@@ -179,7 +180,16 @@ func (*computerUseTool) captureScreenshot(
|
|||||||
fmt.Sprintf("screenshot failed: %v", err),
|
fmt.Sprintf("screenshot failed: %v", err),
|
||||||
), nil
|
), 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(
|
func (t *computerUseTool) captureSharedScreenshot(
|
||||||
@@ -194,22 +204,33 @@ func (t *computerUseTool) captureSharedScreenshot(
|
|||||||
), nil
|
), 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(
|
attachmentName := fmt.Sprintf(
|
||||||
"screenshot-%s.png",
|
"screenshot-%s.png",
|
||||||
t.clock.Now().UTC().Format("2006-01-02T15-04-05Z"),
|
t.clock.Now().UTC().Format("2006-01-02T15-04-05Z"),
|
||||||
)
|
)
|
||||||
if t.storeFile == nil {
|
if t.storeFile == nil {
|
||||||
t.logger.Warn(ctx, "screenshot attachment storage is not configured")
|
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(
|
attachment, err := storeScreenshotAttachment(
|
||||||
ctx,
|
ctx,
|
||||||
t.storeFile,
|
t.storeFile,
|
||||||
attachmentName,
|
attachmentName,
|
||||||
screenResp.ScreenshotData,
|
screenResp.ScreenshotData,
|
||||||
)
|
)
|
||||||
response := fantasy.NewImageResponse([]byte(screenResp.ScreenshotData), "image/png")
|
|
||||||
if err != nil {
|
if err != nil {
|
||||||
t.logger.Warn(ctx, "failed to persist screenshot attachment",
|
t.logger.Warn(ctx, "failed to persist screenshot attachment",
|
||||||
slog.F("attachment_name", attachmentName),
|
slog.F("attachment_name", attachmentName),
|
||||||
|
|||||||
@@ -70,7 +70,9 @@ func TestComputerUseTool_Run_Screenshot(t *testing.T) {
|
|||||||
require.NoError(t, err)
|
require.NoError(t, err)
|
||||||
assert.Equal(t, "image", resp.Type)
|
assert.Equal(t, "image", resp.Type)
|
||||||
assert.Equal(t, "image/png", resp.MediaType)
|
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)
|
assert.False(t, resp.IsError)
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -118,7 +120,9 @@ func TestComputerUseTool_Run_Screenshot_PersistsAttachment(t *testing.T) {
|
|||||||
require.NoError(t, err)
|
require.NoError(t, err)
|
||||||
assert.Equal(t, "image", resp.Type)
|
assert.Equal(t, "image", resp.Type)
|
||||||
assert.Equal(t, "image/png", resp.MediaType)
|
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.Contains(t, storedName, "screenshot-")
|
||||||
assert.Equal(t, "image/png", storedType)
|
assert.Equal(t, "image/png", storedType)
|
||||||
expectedPNG, decodeErr := base64.StdEncoding.DecodeString(screenshotPNG)
|
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", resp.Type)
|
||||||
assert.Equal(t, "image/png", resp.MediaType)
|
assert.Equal(t, "image/png", resp.MediaType)
|
||||||
assert.False(t, resp.IsError)
|
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)
|
attachments, err := chattool.AttachmentsFromMetadata(resp.Metadata)
|
||||||
require.NoError(t, err)
|
require.NoError(t, err)
|
||||||
assert.Empty(t, attachments)
|
assert.Empty(t, attachments)
|
||||||
@@ -260,7 +266,9 @@ func TestComputerUseTool_Run_LeftClick(t *testing.T) {
|
|||||||
resp, err := tool.Run(context.Background(), call)
|
resp, err := tool.Run(context.Background(), call)
|
||||||
require.NoError(t, err)
|
require.NoError(t, err)
|
||||||
assert.Equal(t, "image", resp.Type)
|
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)
|
attachments, err := chattool.AttachmentsFromMetadata(resp.Metadata)
|
||||||
require.NoError(t, err)
|
require.NoError(t, err)
|
||||||
assert.Empty(t, attachments)
|
assert.Empty(t, attachments)
|
||||||
@@ -307,13 +315,73 @@ func TestComputerUseTool_Run_Wait(t *testing.T) {
|
|||||||
require.NoError(t, err)
|
require.NoError(t, err)
|
||||||
assert.Equal(t, "image", resp.Type)
|
assert.Equal(t, "image", resp.Type)
|
||||||
assert.Equal(t, "image/png", resp.MediaType)
|
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)
|
assert.False(t, resp.IsError)
|
||||||
attachments, err := chattool.AttachmentsFromMetadata(resp.Metadata)
|
attachments, err := chattool.AttachmentsFromMetadata(resp.Metadata)
|
||||||
require.NoError(t, err)
|
require.NoError(t, err)
|
||||||
assert.Empty(t, attachments)
|
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) {
|
func TestComputerUseTool_Run_ConnError(t *testing.T) {
|
||||||
t.Parallel()
|
t.Parallel()
|
||||||
|
|
||||||
|
|||||||
@@ -107,7 +107,7 @@ func convertMCPToolResponse(
|
|||||||
for _, c := range resp.Content {
|
for _, c := range resp.Content {
|
||||||
switch c.Type {
|
switch c.Type {
|
||||||
case "text":
|
case "text":
|
||||||
textParts = append(textParts, c.Text)
|
textParts = append(textParts, strings.ToValidUTF8(c.Text, "\uFFFD"))
|
||||||
case "image", "audio":
|
case "image", "audio":
|
||||||
if c.Data == "" {
|
if c.Data == "" {
|
||||||
continue
|
continue
|
||||||
@@ -129,7 +129,7 @@ func convertMCPToolResponse(
|
|||||||
binaryResult = &r
|
binaryResult = &r
|
||||||
}
|
}
|
||||||
default:
|
default:
|
||||||
textParts = append(textParts, c.Text)
|
textParts = append(textParts, strings.ToValidUTF8(c.Text, "\uFFFD"))
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -0,0 +1,5 @@
|
|||||||
|
package mcpclient
|
||||||
|
|
||||||
|
// ConvertCallResultForTest exposes convertCallResult for external
|
||||||
|
// tests.
|
||||||
|
var ConvertCallResultForTest = convertCallResult
|
||||||
@@ -607,7 +607,7 @@ func convertCallResult(
|
|||||||
for _, item := range result.Content {
|
for _, item := range result.Content {
|
||||||
switch c := item.(type) {
|
switch c := item.(type) {
|
||||||
case mcp.TextContent:
|
case mcp.TextContent:
|
||||||
textParts = append(textParts, c.Text)
|
textParts = append(textParts, strings.ToValidUTF8(c.Text, "\uFFFD"))
|
||||||
case mcp.ImageContent:
|
case mcp.ImageContent:
|
||||||
data, err := base64.StdEncoding.DecodeString(
|
data, err := base64.StdEncoding.DecodeString(
|
||||||
c.Data,
|
c.Data,
|
||||||
@@ -653,7 +653,7 @@ func convertCallResult(
|
|||||||
// regardless of form.
|
// regardless of form.
|
||||||
switch r := c.Resource.(type) {
|
switch r := c.Resource.(type) {
|
||||||
case mcp.TextResourceContents:
|
case mcp.TextResourceContents:
|
||||||
textParts = append(textParts, r.Text)
|
textParts = append(textParts, strings.ToValidUTF8(r.Text, "\uFFFD"))
|
||||||
case mcp.BlobResourceContents:
|
case mcp.BlobResourceContents:
|
||||||
data, err := base64.StdEncoding.DecodeString(
|
data, err := base64.StdEncoding.DecodeString(
|
||||||
r.Blob,
|
r.Blob,
|
||||||
|
|||||||
@@ -9,6 +9,7 @@ import (
|
|||||||
"sync"
|
"sync"
|
||||||
"testing"
|
"testing"
|
||||||
"time"
|
"time"
|
||||||
|
"unicode/utf8"
|
||||||
|
|
||||||
"charm.land/fantasy"
|
"charm.land/fantasy"
|
||||||
"github.com/google/uuid"
|
"github.com/google/uuid"
|
||||||
@@ -1271,3 +1272,77 @@ func TestModelIntent_Run_FallbackOnBadJSON(t *testing.T) {
|
|||||||
require.NoError(t, err)
|
require.NoError(t, err)
|
||||||
assert.True(t, resp.IsError, "malformed input should produce an error response")
|
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)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user