From 1b5d4ea0c2829e4b0ed60957a6a2b3fbdb5a60aa Mon Sep 17 00:00:00 2001 From: adranwit Date: Fri, 9 Oct 2026 17:58:18 -0700 Subject: [PATCH] ENG-54517 Preserve binary MCP responses as embedded resources --- gateway/mcp.go | 119 +++++++++++---- gateway/mcp_binary_output_test.go | 237 ++++++++++++++++++++++++++++++ 2 files changed, 325 insertions(+), 31 deletions(-) create mode 100644 gateway/mcp_binary_output_test.go diff --git a/gateway/mcp.go b/gateway/mcp.go index 57dbb9bf9..ecd37d700 100644 --- a/gateway/mcp.go +++ b/gateway/mcp.go @@ -17,6 +17,9 @@ import ( "strconv" "strings" "unicode" + "unicode/utf8" + + "github.com/google/uuid" furl "github.com/viant/afs/url" "github.com/viant/datly/gateway/router/proxy" @@ -1377,9 +1380,12 @@ func (r *Router) newToolHTTPRequest(ctx context.Context, method, URL string, bod func decodeToolResponseBody(responseWriter *proxy.Writer) ([]byte, error) { data := responseWriter.Body.Bytes() encoding := strings.TrimSpace(responseWriter.HeaderMap.Get("Content-Encoding")) - if !strings.EqualFold(encoding, "gzip") { + if encoding == "" || strings.EqualFold(encoding, "identity") { return data, nil } + if !strings.EqualFold(encoding, "gzip") { + return nil, fmt.Errorf("unsupported response content encoding") + } reader, err := gzip.NewReader(bytes.NewReader(data)) if err != nil { return nil, fmt.Errorf("failed to create gzip reader: %w", err) @@ -1392,23 +1398,71 @@ func decodeToolResponseBody(responseWriter *proxy.Writer) ([]byte, error) { return decoded, nil } -// buildToolCallResult composes a CallToolResult with text content and structured error info if status is not OK. -func (r *Router) buildToolCallResult(responseWriter *proxy.Writer, URL, method string) *schema.CallToolResult { - var result = &schema.CallToolResult{} - mimeType := responseWriter.HeaderMap.Get("Content-Type") +// mcpResponseMediaType classifies attachments and non-text responses before +// converting any bytes to strings. Missing MIME types retain the JSON/text +// default unless the body is not valid UTF-8 or HTTP sniffing identifies binary. +func mcpResponseMediaType(writer *proxy.Writer, data []byte, fallback string) (string, bool) { + mimeType := strings.TrimSpace(writer.HeaderMap.Get("Content-Type")) if mimeType == "" { - mimeType = "application/json" + sniffed, _, _ := mime.ParseMediaType(http.DetectContentType(data)) + if !utf8.Valid(data) || !mcpTextMediaType(sniffed) { + mimeType = "application/octet-stream" + } else if fallback != "" { + mimeType = fallback + } else { + mimeType = "application/json" + } + } + mediaType, _, err := mime.ParseMediaType(mimeType) + disposition, _, _ := mime.ParseMediaType(writer.HeaderMap.Get("Content-Disposition")) + return mimeType, strings.EqualFold(disposition, "attachment") || !utf8.Valid(data) || err != nil || !mcpTextMediaType(mediaType) +} + +func mcpTextMediaType(mediaType string) bool { + mediaType = strings.ToLower(mediaType) + return strings.HasPrefix(mediaType, "text/") || + strings.HasSuffix(mediaType, "+json") || strings.HasSuffix(mediaType, "+xml") || + mediaType == "application/json" || mediaType == "application/xml" || + mediaType == "application/javascript" || mediaType == "application/x-javascript" || + mediaType == "application/x-www-form-urlencoded" || mediaType == "application/yaml" || + mediaType == "application/x-yaml" +} + +func mcpBinaryResourceURI() string { + // Never derive binary resource identifiers from the request URL, query, + // Content-Disposition filename, or private storage location. + return "urn:uuid:" + uuid.NewString() +} + +func mcpResponseError(message string, status int) *schema.CallToolResult { + isError := true + return &schema.CallToolResult{ + IsError: &isError, + Content: []schema.CallToolResultContentElem{schema.TextContent{Type: "text", Text: message}}, + StructuredContent: map[string]interface{}{"status": status, "error": true, "message": message}, } +} + +// buildToolCallResult preserves binary responses as standard embedded resources. +func (r *Router) buildToolCallResult(responseWriter *proxy.Writer, URL, method string) *schema.CallToolResult { data, err := decodeToolResponseBody(responseWriter) if err != nil { - data = []byte(err.Error()) - } - result.Content = append(result.Content, schema.CallToolResultContentElem( - schema.TextContent{ - Type: "text", - Text: string(data), - }, - )) + return mcpResponseError("Unable to decode HTTP response body", responseWriter.Code) + } + mimeType, binary := mcpResponseMediaType(responseWriter, data, "application/json") + if binary { + if responseWriter.Code >= http.StatusBadRequest { + return mcpResponseError("HTTP request failed with binary response body", responseWriter.Code) + } + return &schema.CallToolResult{Content: []schema.CallToolResultContentElem{schema.EmbeddedResource{ + Type: "resource", + Resource: schema.EmbeddedResourceResource{ + Uri: mcpBinaryResourceURI(), MimeType: &mimeType, + Blob: base64.StdEncoding.EncodeToString(data), + }, + }}} + } + result := &schema.CallToolResult{Content: []schema.CallToolResultContentElem{schema.TextContent{Type: "text", Text: string(data)}}} _ = json.Unmarshal(data, &result.StructuredContent) if responseWriter.Code >= http.StatusBadRequest || toolPayloadReportsError(result.StructuredContent) { isErr := true @@ -1416,12 +1470,8 @@ func (r *Router) buildToolCallResult(responseWriter *proxy.Writer, URL, method s } if responseWriter.Code >= http.StatusBadRequest { result.StructuredContent = map[string]interface{}{ - "status": responseWriter.Code, - "error": true, - "message": string(data), - "headers": responseWriter.HeaderMap, - "uri": URL, - "method": method, + "status": responseWriter.Code, "error": true, "message": string(data), + "headers": responseWriter.HeaderMap, "uri": URL, "method": method, } } return result @@ -1949,17 +1999,24 @@ func (r *Router) handleMcpRead(ctx context.Context, params *schema.ReadResourceR if responseWriter.Code == http.StatusUnauthorized { return nil, r.mcpUnauthorizedError() } - var result []schema.ReadResourceResultContentsElem - mimeType := "" + data, err := decodeToolResponseBody(responseWriter) + if err != nil { + return nil, jsonrpc.NewInternalError("Unable to decode HTTP response body", nil) + } + fallback := "" if template.MimeType != nil { - mimeType = *template.MimeType + fallback = *template.MimeType } - result = append(result, schema.ReadResourceResultContentsElem{ - Uri: URL, // The URI of the resource - Text: responseWriter.Body.String(), // The actual data returned from the request - MimeType: &mimeType, // The MIME type of the resource - Blob: responseWriter.Body.String(), - }) - return result, nil - + mimeType, binary := mcpResponseMediaType(responseWriter, data, fallback) + result := schema.ReadResourceResultContentsElem{Uri: URL, MimeType: &mimeType} + if binary { + if responseWriter.Code >= http.StatusBadRequest { + return nil, jsonrpc.NewInternalError("HTTP request failed with binary response body", nil) + } + result.Uri = mcpBinaryResourceURI() + result.Blob = base64.StdEncoding.EncodeToString(data) + } else { + result.Text = string(data) + } + return []schema.ReadResourceResultContentsElem{result}, nil } diff --git a/gateway/mcp_binary_output_test.go b/gateway/mcp_binary_output_test.go new file mode 100644 index 000000000..566c3add7 --- /dev/null +++ b/gateway/mcp_binary_output_test.go @@ -0,0 +1,237 @@ +package gateway + +import ( + "bytes" + "compress/gzip" + "context" + "encoding/base64" + "encoding/json" + "net/http" + "net/url" + "testing" + + "github.com/stretchr/testify/require" + "github.com/viant/datly/gateway/router/proxy" + "github.com/viant/datly/repository" + "github.com/viant/datly/repository/contract" + "github.com/viant/datly/repository/version" + "github.com/viant/datly/view" + "github.com/viant/mcp-protocol/schema" + "github.com/xuri/excelize/v2" +) + +const binarySpreadsheetMIME = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet" + +func binaryOutputFixture(t *testing.T) []byte { + t.Helper() + book := excelize.NewFile() + defer book.Close() + require.NoError(t, book.SetCellValue("Sheet1", "A1", "Campaign")) + _, err := book.NewSheet("Ads") + require.NoError(t, err) + require.NoError(t, book.SetCellValue("Ads", "A1", "Creative")) + buffer, err := book.WriteToBuffer() + require.NoError(t, err) + return buffer.Bytes() +} + +func binaryResponse(body []byte, mediaType, disposition, encoding string, status int) *proxy.Writer { + writer := proxy.NewWriter() + writer.Code = status + writer.HeaderMap.Set("Content-Type", mediaType) + writer.HeaderMap.Set("Content-Disposition", disposition) + writer.HeaderMap.Set("Content-Encoding", encoding) + writer.Body.Write(body) + return writer +} + +func gzipBinaryOutput(t *testing.T, data []byte) []byte { + t.Helper() + var output bytes.Buffer + writer := gzip.NewWriter(&output) + _, err := writer.Write(data) + require.NoError(t, err) + require.NoError(t, writer.Close()) + return output.Bytes() +} + +func assertBinaryOutput(t *testing.T, result *schema.CallToolResult, expected []byte, mediaType string) { + t.Helper() + require.Nil(t, result.IsError) + require.Nil(t, result.StructuredContent) + require.Len(t, result.Content, 1) + wire, err := json.Marshal(result.Content[0]) + require.NoError(t, err) + var envelope struct { + Type string `json:"type"` + Resource struct { + Uri string `json:"uri"` + Blob string `json:"blob"` + MimeType *string `json:"mimeType"` + Text string `json:"text"` + } `json:"resource"` + } + require.NoError(t, json.Unmarshal(wire, &envelope)) + require.Equal(t, "resource", envelope.Type) + var standard schema.EmbeddedResource + require.NoError(t, json.Unmarshal(wire, &standard)) + require.Equal(t, envelope.Resource.Blob, standard.Resource.Blob) + require.Empty(t, standard.Meta) + require.Empty(t, standard.Resource.Meta) + require.Empty(t, envelope.Resource.Text) + require.NotNil(t, envelope.Resource.MimeType) + require.Equal(t, mediaType, *envelope.Resource.MimeType) + decoded, err := base64.StdEncoding.DecodeString(envelope.Resource.Blob) + require.NoError(t, err) + require.True(t, bytes.Equal(expected, decoded), "binary output changed") + uri, err := url.Parse(envelope.Resource.Uri) + require.NoError(t, err) + require.Equal(t, "urn", uri.Scheme) + require.Empty(t, uri.RawQuery) + require.Empty(t, uri.User) + require.NotContains(t, uri.String(), "private") + require.NotContains(t, uri.String(), "secret") +} + +func TestMCPBinaryToolOutput(t *testing.T) { + spreadsheet := binaryOutputFixture(t) + for _, test := range []struct { + name string + data []byte + mediaType, disposition, encoding string + }{ + {"invalidUTF8", []byte{0, 255, 254, 128}, "application/octet-stream", "", ""}, + {"xlsx", spreadsheet, binarySpreadsheetMIME, `attachment; filename="campaign.xlsx"`, ""}, + {"gzipXlsx", spreadsheet, binarySpreadsheetMIME, "", "gzip"}, + {"textAttachment", []byte("name,value\nfoo,42"), "text/csv", `attachment; filename="report.csv"`, ""}, + {"missingMIME", []byte{0, 255, 254}, "", "", ""}, + } { + t.Run(test.name, func(t *testing.T) { + body := test.data + if test.encoding == "gzip" { + body = gzipBinaryOutput(t, body) + } + result := (&Router{}).buildToolCallResult(binaryResponse(body, test.mediaType, test.disposition, test.encoding, 200), "http://user:secret@localhost/private/file?token=secret", "GET") + mediaType := test.mediaType + if mediaType == "" { + mediaType = "application/octet-stream" + } + assertBinaryOutput(t, result, test.data, mediaType) + if test.name == "xlsx" || test.name == "gzipXlsx" { + book, err := excelize.OpenReader(bytes.NewReader(test.data)) + require.NoError(t, err) + defer book.Close() + require.Equal(t, []string{"Sheet1", "Ads"}, book.GetSheetList()) + } + }) + } +} + +func TestMCPBinaryToolErrorsDoNotEchoBytes(t *testing.T) { + for _, test := range []struct { + name, mediaType, encoding string + status int + data []byte + }{ + {"binaryHTTPError", "application/octet-stream", "", 500, []byte("private binary payload\x00\xff")}, + {"badGzip", binarySpreadsheetMIME, "gzip", 200, []byte("private binary payload\x00\xff")}, + {"truncatedGzip", binarySpreadsheetMIME, "gzip", 200, gzipBinaryOutput(t, []byte("private binary payload"))[:15]}, + {"unsupportedEncoding", binarySpreadsheetMIME, "br", 200, []byte("private binary payload")}, + } { + t.Run(test.name, func(t *testing.T) { + result := (&Router{}).buildToolCallResult(binaryResponse(test.data, test.mediaType, "", test.encoding, test.status), "http://localhost/test", "GET") + require.NotNil(t, result.IsError) + require.True(t, *result.IsError) + wire, err := json.Marshal(result) + require.NoError(t, err) + require.NotContains(t, string(wire), "private binary payload") + require.NotContains(t, string(wire), base64.StdEncoding.EncodeToString(test.data)) + }) + } +} + +func TestMCPOrdinaryToolOutputCompatible(t *testing.T) { + for _, test := range []struct { + name, mediaType, encoding string + status int + body string + }{ + {"json", "application/json", "", 200, `{"data":[1]}`}, + {"text", "text/plain; charset=utf-8", "", 200, "ordinary text"}, + {"jsonError", "application/json", "", 400, `{"message":"invalid request"}`}, + {"gzipJSON", "application/json", "gzip", 200, `{"data":[1]}`}, + {"defaultJSON", "", "", 200, `{"data":[1]}`}, + {"jsonSuffix", "application/problem+json", "", 400, `{"message":"invalid request"}`}, + {"textUnicode", "text/plain", "", 200, "広告 – Café"}, + {"gzipJSONError", "application/json", "gzip", 400, `{"message":"invalid request"}`}, + } { + t.Run(test.name, func(t *testing.T) { + body := []byte(test.body) + if test.encoding == "gzip" { + body = gzipBinaryOutput(t, body) + } + result := (&Router{}).buildToolCallResult(binaryResponse(body, test.mediaType, "", test.encoding, test.status), "http://localhost/test", "GET") + require.Equal(t, schema.TextContent{Type: "text", Text: test.body}, result.Content[0]) + if test.status >= 400 { + require.True(t, *result.IsError) + require.Equal(t, test.body, result.StructuredContent.(map[string]interface{})["message"]) + } + }) + } +} + +func TestMCPBinaryResourceRead(t *testing.T) { + for _, test := range []struct { + name, mediaType, encoding string + status int + data []byte + }{ + {"xlsx", binarySpreadsheetMIME, "", 200, binaryOutputFixture(t)}, + {"gzip", "application/octet-stream", "gzip", 200, []byte{0, 255, 128}}, + {"json", "application/json", "", 200, []byte(`{"ok":true}`)}, + {"binaryError", binarySpreadsheetMIME, "", 500, []byte{0, 255, 128}}, + {"badGzip", binarySpreadsheetMIME, "gzip", 200, []byte{0, 255, 128}}, + } { + t.Run(test.name, func(t *testing.T) { + component := &repository.Component{View: &view.View{}} + component.View.SetResource(&view.Resource{}) + provider := repository.NewProvider(contract.Path{}, &version.Control{}, func(context.Context, ...repository.Option) (*repository.Component, error) { return component, nil }) + body := test.data + if test.name == "gzip" { + body = gzipBinaryOutput(t, body) + } + route := &Route{Handler: func(_ context.Context, writer http.ResponseWriter, _ *http.Request) { + writer.Header().Set("Content-Type", test.mediaType) + writer.Header().Set("Content-Encoding", test.encoding) + writer.WriteHeader(test.status) + writer.Write(body) + }} + fallback := "application/json" + results, rpcErr := (&Router{}).handleMcpRead(context.Background(), &schema.ReadResourceRequestParams{Uri: "datly://localhost/private/file?token=secret"}, &schema.ResourceTemplate{MimeType: &fallback}, route, provider) + if test.name == "binaryError" || test.name == "badGzip" { + require.NotNil(t, rpcErr) + require.Empty(t, results) + return + } + require.Nil(t, rpcErr) + require.Len(t, results, 1) + require.Equal(t, test.mediaType, *results[0].MimeType) + if test.name == "json" { + require.Equal(t, string(test.data), results[0].Text) + require.Empty(t, results[0].Blob) + return + } + require.Empty(t, results[0].Text) + wire, err := json.Marshal(results[0]) + require.NoError(t, err) + var standard schema.BlobResourceContents + require.NoError(t, json.Unmarshal(wire, &standard)) + require.Equal(t, results[0].Blob, standard.Blob) + decoded, err := base64.StdEncoding.DecodeString(results[0].Blob) + require.NoError(t, err) + require.True(t, bytes.Equal(test.data, decoded), "binary resource changed") + require.NotContains(t, results[0].Uri, "private") + require.NotContains(t, results[0].Uri, "secret") + }) + } +}