Skip to content

Commit 9c0be09

Browse files
fix: ensure resource/template/list is marshalled safely (#56)
1 parent 8aa25c8 commit 9c0be09

File tree

4 files changed

+56
-7
lines changed

4 files changed

+56
-7
lines changed

mcp/resources.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,7 @@ type ResourceTemplateOption func(*ResourceTemplate)
6262
// Options are applied in order, allowing for flexible template configuration.
6363
func NewResourceTemplate(uriTemplate string, name string, opts ...ResourceTemplateOption) ResourceTemplate {
6464
template := ResourceTemplate{
65-
URITemplate: uritemplate.MustNew(uriTemplate),
65+
URITemplate: &URITemplate{Template: uritemplate.MustNew(uriTemplate)},
6666
Name: name,
6767
}
6868

mcp/types.go

Lines changed: 22 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,27 @@ import (
88
"github.com/yosida95/uritemplate/v3"
99
)
1010

11+
type URITemplate struct {
12+
*uritemplate.Template
13+
}
14+
15+
func (t *URITemplate) MarshalJSON() ([]byte, error) {
16+
return json.Marshal(t.Template.Raw())
17+
}
18+
19+
func (t *URITemplate) UnmarshalJSON(data []byte) error {
20+
var raw string
21+
if err := json.Unmarshal(data, &raw); err != nil {
22+
return err
23+
}
24+
template, err := uritemplate.New(raw)
25+
if err != nil {
26+
return err
27+
}
28+
t.Template = template
29+
return nil
30+
}
31+
1132
/* JSON-RPC types */
1233

1334
// JSONRPCMessage represents either a JSONRPCRequest, JSONRPCNotification, JSONRPCResponse, or JSONRPCError
@@ -446,7 +467,7 @@ type ResourceTemplate struct {
446467
Annotated
447468
// A URI template (according to RFC 6570) that can be used to construct
448469
// resource URIs.
449-
URITemplate *uritemplate.Template `json:"uriTemplate"`
470+
URITemplate *URITemplate `json:"uriTemplate"`
450471
// A human-readable name for the type of resource this template refers to.
451472
//
452473
// This can be used by clients to populate UI elements.

server/server.go

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,6 @@ import (
1010
"sync/atomic"
1111

1212
"github.com/mark3labs/mcp-go/mcp"
13-
"github.com/yosida95/uritemplate/v3"
1413
)
1514

1615
// resourceEntry holds both a resource and its handler
@@ -731,7 +730,7 @@ func (s *MCPServer) handleReadResource(
731730
}
732731

733732
// matchesTemplate checks if a URI matches a URI template pattern
734-
func matchesTemplate(uri string, template *uritemplate.Template) bool {
733+
func matchesTemplate(uri string, template *mcp.URITemplate) bool {
735734
return template.Regexp().MatchString(uri)
736735
}
737736

server/server_test.go

Lines changed: 32 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -923,7 +923,6 @@ func TestMCPServer_Instructions(t *testing.T) {
923923
func TestMCPServer_ResourceTemplates(t *testing.T) {
924924
server := NewMCPServer("test-server", "1.0.0",
925925
WithResourceCapabilities(true, true),
926-
WithPromptCapabilities(true),
927926
)
928927

929928
server.AddResourceTemplate(
@@ -947,9 +946,15 @@ func TestMCPServer_ResourceTemplates(t *testing.T) {
947946
},
948947
)
949948

950-
message := `{
949+
listMessage := `{
951950
"jsonrpc": "2.0",
952951
"id": 1,
952+
"method": "resources/templates/list"
953+
}`
954+
955+
message := `{
956+
"jsonrpc": "2.0",
957+
"id": 2,
953958
"method": "resources/read",
954959
"params": {
955960
"uri": "test://something/test-resource/a/b/c"
@@ -959,12 +964,35 @@ func TestMCPServer_ResourceTemplates(t *testing.T) {
959964
t.Run("Get resource template", func(t *testing.T) {
960965
response := server.HandleMessage(
961966
context.Background(),
962-
[]byte(message),
967+
[]byte(listMessage),
963968
)
964969
assert.NotNil(t, response)
965970

966971
resp, ok := response.(mcp.JSONRPCResponse)
967972
assert.True(t, ok)
973+
listResult, ok := resp.Result.(mcp.ListResourceTemplatesResult)
974+
assert.True(t, ok)
975+
assert.Len(t, listResult.ResourceTemplates, 1)
976+
assert.Equal(t, "My Resource", listResult.ResourceTemplates[0].Name)
977+
template, err := json.Marshal(listResult.ResourceTemplates[0])
978+
assert.NoError(t, err)
979+
980+
// Need to serialize the json to map[string]string to validate the URITemplate is correctly marshalled
981+
var resourceTemplate map[string]string
982+
err = json.Unmarshal(template, &resourceTemplate)
983+
assert.NoError(t, err)
984+
985+
assert.Equal(t, "test://{a}/test-resource{/b*}", resourceTemplate["uriTemplate"])
986+
987+
response = server.HandleMessage(
988+
context.Background(),
989+
[]byte(message),
990+
)
991+
992+
assert.NotNil(t, response)
993+
994+
resp, ok = response.(mcp.JSONRPCResponse)
995+
assert.True(t, ok)
968996
// Validate that the resource values are returned correctly
969997
result, ok := resp.Result.(mcp.ReadResourceResult)
970998
assert.True(t, ok)
@@ -974,6 +1002,7 @@ func TestMCPServer_ResourceTemplates(t *testing.T) {
9741002
assert.Equal(t, "test://something/test-resource/a/b/c", resultContent.URI)
9751003
assert.Equal(t, "text/plain", resultContent.MIMEType)
9761004
assert.Equal(t, "test content: something", resultContent.Text)
1005+
9771006
})
9781007
}
9791008

0 commit comments

Comments
 (0)