Skip to content

Commit 6bce2d5

Browse files
Add MCP_DISABLE_EMBEDDED_RESOURCES feature flag for file contents
- Added FeatureFlagDisableEmbeddedResources constant to feature_flags.go - Implemented NewToolResultResourceWithFlag in utils/result.go - When flag is enabled, returns TextContent for text files - Returns ImageContent with base64 data for binary files - Maintains backward compatibility when flag is disabled - Updated get_file_contents tool to check feature flag and use new function - Added comprehensive tests for both utils function and integration with get_file_contents Co-authored-by: SamMorrowDrums <4811358+SamMorrowDrums@users.noreply.github.com>
1 parent 183705f commit 6bce2d5

File tree

5 files changed

+367
-5
lines changed

5 files changed

+367
-5
lines changed

pkg/github/feature_flags.go

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,3 +5,9 @@ type FeatureFlags struct {
55
LockdownMode bool
66
InsidersMode bool
77
}
8+
9+
// FeatureFlagDisableEmbeddedResources is a feature flag that when enabled,
10+
// returns file contents as regular MCP content (TextContent/ImageContent)
11+
// instead of EmbeddedResource. This provides better compatibility with
12+
// clients that don't support embedded resources.
13+
const FeatureFlagDisableEmbeddedResources = "MCP_DISABLE_EMBEDDED_RESOURCES"

pkg/github/repositories.go

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -774,23 +774,27 @@ func GetFileContents(t translations.TranslationHelperFunc) inventory.ServerTool
774774
Text: string(body),
775775
MIMEType: contentType,
776776
}
777+
// Check if embedded resources should be disabled
778+
disableEmbedded := deps.IsFeatureEnabled(ctx, FeatureFlagDisableEmbeddedResources)
777779
// Include SHA in the result metadata
778780
if fileSHA != "" {
779-
return utils.NewToolResultResource(fmt.Sprintf("successfully downloaded text file (SHA: %s)", fileSHA)+successNote, result), nil, nil
781+
return utils.NewToolResultResourceWithFlag(fmt.Sprintf("successfully downloaded text file (SHA: %s)", fileSHA)+successNote, result, disableEmbedded), nil, nil
780782
}
781-
return utils.NewToolResultResource("successfully downloaded text file"+successNote, result), nil, nil
783+
return utils.NewToolResultResourceWithFlag("successfully downloaded text file"+successNote, result, disableEmbedded), nil, nil
782784
}
783785

784786
result := &mcp.ResourceContents{
785787
URI: resourceURI,
786788
Blob: body,
787789
MIMEType: contentType,
788790
}
791+
// Check if embedded resources should be disabled
792+
disableEmbedded := deps.IsFeatureEnabled(ctx, FeatureFlagDisableEmbeddedResources)
789793
// Include SHA in the result metadata
790794
if fileSHA != "" {
791-
return utils.NewToolResultResource(fmt.Sprintf("successfully downloaded binary file (SHA: %s)", fileSHA)+successNote, result), nil, nil
795+
return utils.NewToolResultResourceWithFlag(fmt.Sprintf("successfully downloaded binary file (SHA: %s)", fileSHA)+successNote, result, disableEmbedded), nil, nil
792796
}
793-
return utils.NewToolResultResource("successfully downloaded binary file"+successNote, result), nil, nil
797+
return utils.NewToolResultResourceWithFlag("successfully downloaded binary file"+successNote, result, disableEmbedded), nil, nil
794798
}
795799

796800
// Raw API call failed

pkg/github/repositories_test.go

Lines changed: 177 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -403,6 +403,183 @@ func Test_GetFileContents(t *testing.T) {
403403
}
404404
}
405405

406+
func Test_GetFileContents_WithDisabledEmbeddedResourcesFlag(t *testing.T) {
407+
// Verify that when MCP_DISABLE_EMBEDDED_RESOURCES flag is enabled,
408+
// file contents are returned as TextContent/ImageContent instead of EmbeddedResource
409+
serverTool := GetFileContents(translations.NullTranslationHelper)
410+
411+
mockRawContent := []byte("# Test Content\n\nSample text.")
412+
mockBinaryContent := []byte{0x89, 0x50, 0x4E, 0x47} // PNG header
413+
414+
tests := []struct {
415+
name string
416+
mockedClient *http.Client
417+
requestArgs map[string]interface{}
418+
flagEnabled bool
419+
expectTextContent bool // true for TextContent, false for ImageContent
420+
expectedText string
421+
expectedMimeType string
422+
}{
423+
{
424+
name: "text file with flag enabled",
425+
mockedClient: MockHTTPClientWithHandlers(map[string]http.HandlerFunc{
426+
GetReposGitRefByOwnerByRepoByRef: mockResponse(t, http.StatusOK, "{\"ref\": \"refs/heads/main\", \"object\": {\"sha\": \"\"}}"),
427+
GetReposByOwnerByRepo: mockResponse(t, http.StatusOK, "{\"name\": \"repo\", \"default_branch\": \"main\"}"),
428+
GetReposContentsByOwnerByRepoByPath: func(w http.ResponseWriter, _ *http.Request) {
429+
w.WriteHeader(http.StatusOK)
430+
fileContent := &github.RepositoryContent{
431+
Name: github.Ptr("test.md"),
432+
Path: github.Ptr("test.md"),
433+
SHA: github.Ptr("abc123"),
434+
Type: github.Ptr("file"),
435+
}
436+
contentBytes, _ := json.Marshal(fileContent)
437+
_, _ = w.Write(contentBytes)
438+
},
439+
GetRawReposContentsByOwnerByRepoByBranchByPath: func(w http.ResponseWriter, _ *http.Request) {
440+
w.Header().Set("Content-Type", "text/markdown")
441+
_, _ = w.Write(mockRawContent)
442+
},
443+
}),
444+
requestArgs: map[string]interface{}{
445+
"owner": "owner",
446+
"repo": "repo",
447+
"path": "test.md",
448+
"ref": "refs/heads/main",
449+
},
450+
flagEnabled: true,
451+
expectTextContent: true,
452+
expectedText: string(mockRawContent),
453+
expectedMimeType: "text/markdown",
454+
},
455+
{
456+
name: "binary file with flag enabled",
457+
mockedClient: MockHTTPClientWithHandlers(map[string]http.HandlerFunc{
458+
GetReposGitRefByOwnerByRepoByRef: mockResponse(t, http.StatusOK, "{\"ref\": \"refs/heads/main\", \"object\": {\"sha\": \"\"}}"),
459+
GetReposByOwnerByRepo: mockResponse(t, http.StatusOK, "{\"name\": \"repo\", \"default_branch\": \"main\"}"),
460+
GetReposContentsByOwnerByRepoByPath: func(w http.ResponseWriter, _ *http.Request) {
461+
w.WriteHeader(http.StatusOK)
462+
fileContent := &github.RepositoryContent{
463+
Name: github.Ptr("image.png"),
464+
Path: github.Ptr("image.png"),
465+
SHA: github.Ptr("def456"),
466+
Type: github.Ptr("file"),
467+
}
468+
contentBytes, _ := json.Marshal(fileContent)
469+
_, _ = w.Write(contentBytes)
470+
},
471+
GetRawReposContentsByOwnerByRepoByBranchByPath: func(w http.ResponseWriter, _ *http.Request) {
472+
w.Header().Set("Content-Type", "image/png")
473+
_, _ = w.Write(mockBinaryContent)
474+
},
475+
}),
476+
requestArgs: map[string]interface{}{
477+
"owner": "owner",
478+
"repo": "repo",
479+
"path": "image.png",
480+
"ref": "refs/heads/main",
481+
},
482+
flagEnabled: true,
483+
expectTextContent: false,
484+
expectedMimeType: "image/png",
485+
},
486+
{
487+
name: "text file with flag disabled (default behavior)",
488+
mockedClient: MockHTTPClientWithHandlers(map[string]http.HandlerFunc{
489+
GetReposGitRefByOwnerByRepoByRef: mockResponse(t, http.StatusOK, "{\"ref\": \"refs/heads/main\", \"object\": {\"sha\": \"\"}}"),
490+
GetReposByOwnerByRepo: mockResponse(t, http.StatusOK, "{\"name\": \"repo\", \"default_branch\": \"main\"}"),
491+
GetReposContentsByOwnerByRepoByPath: func(w http.ResponseWriter, _ *http.Request) {
492+
w.WriteHeader(http.StatusOK)
493+
fileContent := &github.RepositoryContent{
494+
Name: github.Ptr("test.md"),
495+
Path: github.Ptr("test.md"),
496+
SHA: github.Ptr("abc123"),
497+
Type: github.Ptr("file"),
498+
}
499+
contentBytes, _ := json.Marshal(fileContent)
500+
_, _ = w.Write(contentBytes)
501+
},
502+
GetRawReposContentsByOwnerByRepoByBranchByPath: func(w http.ResponseWriter, _ *http.Request) {
503+
w.Header().Set("Content-Type", "text/markdown")
504+
_, _ = w.Write(mockRawContent)
505+
},
506+
}),
507+
requestArgs: map[string]interface{}{
508+
"owner": "owner",
509+
"repo": "repo",
510+
"path": "test.md",
511+
"ref": "refs/heads/main",
512+
},
513+
flagEnabled: false,
514+
},
515+
}
516+
517+
for _, tc := range tests {
518+
t.Run(tc.name, func(t *testing.T) {
519+
// Setup client with mock
520+
client := github.NewClient(tc.mockedClient)
521+
mockRawClient := raw.NewClient(client, &url.URL{Scheme: "https", Host: "raw.example.com", Path: "/"})
522+
523+
// Create feature flag checker
524+
featureChecker := func(_ context.Context, flagName string) (bool, error) {
525+
if flagName == FeatureFlagDisableEmbeddedResources {
526+
return tc.flagEnabled, nil
527+
}
528+
return false, nil
529+
}
530+
531+
deps := BaseDeps{
532+
Client: client,
533+
RawClient: mockRawClient,
534+
featureChecker: featureChecker,
535+
}
536+
handler := serverTool.Handler(deps)
537+
538+
// Create call request
539+
request := createMCPRequest(tc.requestArgs)
540+
541+
// Call handler
542+
result, err := handler(ContextWithDeps(context.Background(), deps), &request)
543+
require.NoError(t, err)
544+
545+
// Verify results
546+
require.Len(t, result.Content, 2, "Expected 2 content items")
547+
assert.False(t, result.IsError)
548+
549+
// First item should always be the message text
550+
textContent, ok := result.Content[0].(*mcp.TextContent)
551+
require.True(t, ok, "First content should be TextContent")
552+
assert.Contains(t, textContent.Text, "successfully downloaded")
553+
554+
if tc.flagEnabled {
555+
// When flag is enabled, second item should be TextContent or ImageContent
556+
if tc.expectTextContent {
557+
// Expecting TextContent for text files
558+
content, ok := result.Content[1].(*mcp.TextContent)
559+
require.True(t, ok, "Expected TextContent for text file with flag enabled, got %T", result.Content[1])
560+
assert.Equal(t, tc.expectedText, content.Text)
561+
assert.NotNil(t, content.Meta)
562+
assert.Equal(t, tc.expectedMimeType, content.Meta["mimeType"])
563+
assert.NotNil(t, content.Annotations)
564+
} else {
565+
// Expecting ImageContent for binary files
566+
content, ok := result.Content[1].(*mcp.ImageContent)
567+
require.True(t, ok, "Expected ImageContent for binary file with flag enabled, got %T", result.Content[1])
568+
assert.Equal(t, tc.expectedMimeType, content.MIMEType)
569+
assert.NotNil(t, content.Meta)
570+
assert.NotNil(t, content.Annotations)
571+
// Verify data is base64 encoded
572+
assert.NotEmpty(t, content.Data)
573+
}
574+
} else {
575+
// When flag is disabled, should use EmbeddedResource (default)
576+
_, ok := result.Content[1].(*mcp.EmbeddedResource)
577+
assert.True(t, ok, "Expected EmbeddedResource when flag is disabled, got %T", result.Content[1])
578+
}
579+
})
580+
}
581+
}
582+
406583
func Test_ForkRepository(t *testing.T) {
407584
// Verify tool definition once
408585
serverTool := ForkRepository(translations.NullTranslationHelper)

pkg/utils/result.go

Lines changed: 58 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,10 @@
11
package utils //nolint:revive //TODO: figure out a better name for this package
22

3-
import "github.com/modelcontextprotocol/go-sdk/mcp"
3+
import (
4+
"encoding/base64"
5+
6+
"github.com/modelcontextprotocol/go-sdk/mcp"
7+
)
48

59
func NewToolResultText(message string) *mcp.CallToolResult {
610
return &mcp.CallToolResult{
@@ -47,3 +51,56 @@ func NewToolResultResource(message string, contents *mcp.ResourceContents) *mcp.
4751
IsError: false,
4852
}
4953
}
54+
55+
// NewToolResultResourceWithFlag returns a CallToolResult with either an embedded resource
56+
// or regular content based on the disableEmbeddedResources flag.
57+
// When disableEmbeddedResources is true, text content is returned as TextContent and
58+
// binary content is returned as ImageContent, providing better client compatibility.
59+
func NewToolResultResourceWithFlag(message string, contents *mcp.ResourceContents, disableEmbeddedResources bool) *mcp.CallToolResult {
60+
if !disableEmbeddedResources {
61+
// Default behavior - return as embedded resource
62+
return NewToolResultResource(message, contents)
63+
}
64+
65+
// When flag is enabled, return as regular content
66+
var content mcp.Content
67+
switch {
68+
case contents.Text != "":
69+
// Text content - use TextContent with mime type in annotations
70+
content = &mcp.TextContent{
71+
Text: contents.Text,
72+
Annotations: &mcp.Annotations{
73+
Audience: []mcp.Role{"user"},
74+
},
75+
Meta: mcp.Meta{
76+
"mimeType": contents.MIMEType,
77+
"uri": contents.URI,
78+
},
79+
}
80+
case len(contents.Blob) > 0:
81+
// Binary content - use ImageContent with base64 data
82+
content = &mcp.ImageContent{
83+
Data: []byte(base64.StdEncoding.EncodeToString(contents.Blob)),
84+
MIMEType: contents.MIMEType,
85+
Meta: mcp.Meta{
86+
"uri": contents.URI,
87+
},
88+
Annotations: &mcp.Annotations{
89+
Audience: []mcp.Role{"user"},
90+
},
91+
}
92+
default:
93+
// Fallback to embedded resource if neither text nor blob
94+
return NewToolResultResource(message, contents)
95+
}
96+
97+
return &mcp.CallToolResult{
98+
Content: []mcp.Content{
99+
&mcp.TextContent{
100+
Text: message,
101+
},
102+
content,
103+
},
104+
IsError: false,
105+
}
106+
}

0 commit comments

Comments
 (0)