fix(mcp): preserve raw MCP payload in text artifacts

This commit is contained in:
afjcjsbx 2026-04-03 17:22:40 +02:00
parent 969ff6cfd5
commit b6aa585dc0
2 changed files with 65 additions and 11 deletions

View file

@ -272,14 +272,19 @@ func extractContentText(content []mcp.Content) string {
func (t *MCPTool) normalizeResultContent(ctx context.Context, content []mcp.Content) *ToolResult { func (t *MCPTool) normalizeResultContent(ctx context.Context, content []mcp.Content) *ToolResult {
llmParts := make([]string, 0, len(content)) llmParts := make([]string, 0, len(content))
rawTextParts := make([]string, 0, len(content))
mediaRefs := make([]string, 0, len(content)) mediaRefs := make([]string, 0, len(content))
for _, c := range content { for _, c := range content {
switch v := c.(type) { switch v := c.(type) {
case *mcp.TextContent: case *mcp.TextContent:
text := strings.TrimSpace(sanitizeToolLLMContent(v.Text)) rawText := strings.TrimSpace(v.Text)
if text != "" { if rawText != "" {
llmParts = append(llmParts, text) rawTextParts = append(rawTextParts, rawText)
}
safeText := strings.TrimSpace(sanitizeToolLLMContent(v.Text))
if safeText != "" {
llmParts = append(llmParts, safeText)
} }
case *mcp.ImageContent: case *mcp.ImageContent:
ref, note := t.storeBinaryContent( ref, note := t.storeBinaryContent(
@ -312,10 +317,13 @@ func (t *MCPTool) normalizeResultContent(ctx context.Context, content []mcp.Cont
case *mcp.ResourceLink: case *mcp.ResourceLink:
llmParts = append(llmParts, summarizeResourceLink(v)) llmParts = append(llmParts, summarizeResourceLink(v))
case *mcp.EmbeddedResource: case *mcp.EmbeddedResource:
ref, note := t.storeEmbeddedResource(ctx, v) ref, note, rawText := t.storeEmbeddedResource(ctx, v)
if ref != "" { if ref != "" {
mediaRefs = append(mediaRefs, ref) mediaRefs = append(mediaRefs, ref)
} }
if rawText != "" {
rawTextParts = append(rawTextParts, rawText)
}
if note != "" { if note != "" {
llmParts = append(llmParts, note) llmParts = append(llmParts, note)
} }
@ -325,7 +333,8 @@ func (t *MCPTool) normalizeResultContent(ctx context.Context, content []mcp.Cont
} }
forLLM := strings.Join(compactStrings(llmParts), "\n") forLLM := strings.Join(compactStrings(llmParts), "\n")
if artifactResult := t.persistLargeTextArtifact(forLLM); artifactResult != nil { rawText := strings.Join(compactStrings(rawTextParts), "\n")
if artifactResult := t.persistLargeTextArtifact(rawText); artifactResult != nil {
artifactResult.Media = mediaRefs artifactResult.Media = mediaRefs
return artifactResult return artifactResult
} }
@ -381,27 +390,29 @@ func (t *MCPTool) persistLargeTextArtifact(text string) *ToolResult {
} }
} }
func (t *MCPTool) storeEmbeddedResource(ctx context.Context, content *mcp.EmbeddedResource) (string, string) { func (t *MCPTool) storeEmbeddedResource(ctx context.Context, content *mcp.EmbeddedResource) (string, string, string) {
if content == nil || content.Resource == nil { if content == nil || content.Resource == nil {
return "", "[MCP returned an embedded resource without data.]" return "", "[MCP returned an embedded resource without data.]", ""
} }
resource := content.Resource resource := content.Resource
if len(resource.Blob) > 0 { if len(resource.Blob) > 0 {
return t.storeBinaryContent( ref, note := t.storeBinaryContent(
ctx, ctx,
"resource", "resource",
normalizedMIMEType(resource.MIMEType), normalizedMIMEType(resource.MIMEType),
resource.Blob, resource.Blob,
content.Annotations, content.Annotations,
) )
return ref, note, ""
} }
if strings.TrimSpace(resource.Text) != "" { rawText := strings.TrimSpace(resource.Text)
return "", sanitizeToolLLMContent(resource.Text) if rawText != "" {
return "", sanitizeToolLLMContent(resource.Text), rawText
} }
return "", summarizeEmbeddedResource(content) return "", summarizeEmbeddedResource(content), ""
} }
func (t *MCPTool) storeBinaryContent( func (t *MCPTool) storeBinaryContent(

View file

@ -635,6 +635,49 @@ func TestMCPTool_Execute_LargeBase64TextIsOmittedFromContext(t *testing.T) {
} }
} }
func TestMCPTool_Execute_LargeBase64TextArtifactPreservesRawPayload(t *testing.T) {
workspace := t.TempDir()
largeBase64 := strings.Repeat("QUJD", 400)
manager := &MockMCPManager{
callToolFunc: func(ctx context.Context, serverName, toolName string, arguments map[string]any) (*mcp.CallToolResult, error) {
return &mcp.CallToolResult{
Content: []mcp.Content{
&mcp.TextContent{Text: largeBase64},
},
}, nil
},
}
mcpTool := NewMCPTool(manager, "test_server", &mcp.Tool{Name: "dump_payload"})
mcpTool.SetWorkspace(workspace)
mcpTool.SetMaxInlineTextRunes(32)
result := mcpTool.Execute(context.Background(), nil)
if !strings.Contains(result.ForLLM, "saved as a local artifact") {
t.Fatalf("expected artifact note, got %q", result.ForLLM)
}
if result.ForLLM == largeBase64OmittedMessage {
t.Fatalf("expected artifact note instead of sanitized base64 placeholder")
}
if len(result.ArtifactTags) != 1 {
t.Fatalf("expected 1 artifact tag, got %d", len(result.ArtifactTags))
}
tag := result.ArtifactTags[0]
const prefix = "[file:"
if !strings.HasPrefix(tag, prefix) || !strings.HasSuffix(tag, "]") {
t.Fatalf("expected file artifact tag, got %q", tag)
}
path := strings.TrimSuffix(strings.TrimPrefix(tag, prefix), "]")
data, err := os.ReadFile(path)
if err != nil {
t.Fatalf("expected artifact file to be readable: %v", err)
}
if string(data) != largeBase64 {
t.Fatalf("expected artifact file contents to preserve raw MCP payload")
}
}
func TestMCPTool_Execute_LargeTextStoredAsArtifact(t *testing.T) { func TestMCPTool_Execute_LargeTextStoredAsArtifact(t *testing.T) {
workspace := t.TempDir() workspace := t.TempDir()
largeText := strings.Repeat("This is a large MCP text payload.\n", 800) largeText := strings.Repeat("This is a large MCP text payload.\n", 800)