Skip to content

Commit 33e8dec

Browse files
nikolubbenikolubbe
committed
Add ResourceTemplates
1 parent 1fa18e0 commit 33e8dec

File tree

3 files changed

+267
-5
lines changed

3 files changed

+267
-5
lines changed

Diff for: ‎resource_response_types.go

+26
Original file line numberDiff line numberDiff line change
@@ -36,3 +36,29 @@ type ResourceSchema struct {
3636
// The URI of this resource.
3737
Uri string `json:"uri" yaml:"uri" mapstructure:"uri"`
3838
}
39+
40+
// A resource template that defines a pattern for dynamic resources.
41+
type ResourceTemplateSchema struct {
42+
// Annotations corresponds to the JSON schema field "annotations".
43+
Annotations *Annotations `json:"annotations,omitempty" yaml:"annotations,omitempty" mapstructure:"annotations,omitempty"`
44+
45+
// A description of what resources matching this template represent.
46+
Description *string `json:"description,omitempty" yaml:"description,omitempty" mapstructure:"description,omitempty"`
47+
48+
// The MIME type of resources matching this template, if known.
49+
MimeType *string `json:"mimeType,omitempty" yaml:"mimeType,omitempty" mapstructure:"mimeType,omitempty"`
50+
51+
// A human-readable name for this template.
52+
Name string `json:"name" yaml:"name" mapstructure:"name"`
53+
54+
// The URI template following RFC 6570.
55+
UriTemplate string `json:"uriTemplate" yaml:"uriTemplate" mapstructure:"uriTemplate"`
56+
}
57+
58+
// The server's response to a resources/templates/list request from the client.
59+
type ListResourceTemplatesResponse struct {
60+
// Templates corresponds to the JSON schema field "templates".
61+
Templates []*ResourceTemplateSchema `json:"resourceTemplates" yaml:"resourceTemplates" mapstructure:"resourceTemplates"`
62+
// NextCursor is a cursor for pagination. If not nil, there are more templates available.
63+
NextCursor *string `json:"nextCursor,omitempty" yaml:"nextCursor,omitempty" mapstructure:"nextCursor,omitempty"`
64+
}

Diff for: ‎server.go

+107-5
Original file line numberDiff line numberDiff line change
@@ -107,6 +107,7 @@ type Server struct {
107107
tools *datastructures.SyncMap[string, *tool]
108108
prompts *datastructures.SyncMap[string, *prompt]
109109
resources *datastructures.SyncMap[string, *resource]
110+
resourceTemplates *datastructures.SyncMap[string, *resourceTemplate]
110111
serverInstructions *string
111112
serverName string
112113
serverVersion string
@@ -134,6 +135,13 @@ type resource struct {
134135
Handler func(context.Context) *resourceResponseSent
135136
}
136137

138+
type resourceTemplate struct {
139+
Name string
140+
Description string
141+
UriTemplate string
142+
MimeType string
143+
}
144+
137145
type ServerOptions func(*Server)
138146

139147
func WithProtocol(protocol *protocol.Protocol) ServerOptions {
@@ -163,11 +171,12 @@ func WithVersion(version string) ServerOptions {
163171

164172
func NewServer(transport transport.Transport, options ...ServerOptions) *Server {
165173
server := &Server{
166-
protocol: protocol.NewProtocol(nil),
167-
transport: transport,
168-
tools: new(datastructures.SyncMap[string, *tool]),
169-
prompts: new(datastructures.SyncMap[string, *prompt]),
170-
resources: new(datastructures.SyncMap[string, *resource]),
174+
protocol: protocol.NewProtocol(nil),
175+
transport: transport,
176+
tools: new(datastructures.SyncMap[string, *tool]),
177+
prompts: new(datastructures.SyncMap[string, *prompt]),
178+
resources: new(datastructures.SyncMap[string, *resource]),
179+
resourceTemplates: new(datastructures.SyncMap[string, *resourceTemplate]),
171180
}
172181
for _, option := range options {
173182
option(server)
@@ -299,6 +308,26 @@ func validateResourceHandler(handler any) error {
299308
return nil
300309
}
301310

311+
func (s *Server) RegisterResourceTemplate(uriTemplate string, name string, description string, mimeType string) error {
312+
s.resourceTemplates.Store(uriTemplate, &resourceTemplate{
313+
Name: name,
314+
Description: description,
315+
UriTemplate: uriTemplate,
316+
MimeType: mimeType,
317+
})
318+
return s.sendResourceListChangedNotification()
319+
}
320+
321+
func (s *Server) CheckResourceTemplateRegistered(uriTemplate string) bool {
322+
_, ok := s.resourceTemplates.Load(uriTemplate)
323+
return ok
324+
}
325+
326+
func (s *Server) DeregisterResourceTemplate(uriTemplate string) error {
327+
s.resourceTemplates.Delete(uriTemplate)
328+
return s.sendResourceListChangedNotification()
329+
}
330+
302331
func (s *Server) RegisterPrompt(name string, description string, handler any) error {
303332
err := validatePromptHandler(handler)
304333
if err != nil {
@@ -553,6 +582,7 @@ func (s *Server) Serve() error {
553582
pr.SetRequestHandler("prompts/list", s.handleListPrompts)
554583
pr.SetRequestHandler("prompts/get", s.handlePromptCalls)
555584
pr.SetRequestHandler("resources/list", s.handleListResources)
585+
pr.SetRequestHandler("resources/templates/list", s.handleListResourceTemplates)
556586
pr.SetRequestHandler("resources/read", s.handleResourceCalls)
557587
err := pr.Connect(s.transport)
558588
if err != nil {
@@ -825,6 +855,78 @@ func (s *Server) handleListResources(ctx context.Context, request *transport.Bas
825855
}, nil
826856
}
827857

858+
func (s *Server) handleListResourceTemplates(ctx context.Context, request *transport.BaseJSONRPCRequest, extra protocol.RequestHandlerExtra) (transport.JsonRpcBody, error) {
859+
type resourceTemplateRequestParams struct {
860+
Cursor *string `json:"cursor"`
861+
}
862+
var params resourceTemplateRequestParams
863+
if request.Params == nil {
864+
params = resourceTemplateRequestParams{}
865+
} else {
866+
err := json.Unmarshal(request.Params, &params)
867+
if err != nil {
868+
return nil, errors.Wrap(err, "failed to unmarshal arguments")
869+
}
870+
}
871+
872+
// Order by URI template for pagination
873+
var orderedTemplates []*resourceTemplate
874+
s.resourceTemplates.Range(func(k string, t *resourceTemplate) bool {
875+
orderedTemplates = append(orderedTemplates, t)
876+
return true
877+
})
878+
sort.Slice(orderedTemplates, func(i, j int) bool {
879+
return orderedTemplates[i].UriTemplate < orderedTemplates[j].UriTemplate
880+
})
881+
882+
startPosition := 0
883+
if params.Cursor != nil {
884+
// Base64 decode the cursor
885+
c, err := base64.StdEncoding.DecodeString(*params.Cursor)
886+
if err != nil {
887+
return nil, errors.Wrap(err, "failed to decode cursor")
888+
}
889+
cString := string(c)
890+
// Iterate through the templates until we find an entry > the cursor
891+
for i := 0; i < len(orderedTemplates); i++ {
892+
if orderedTemplates[i].UriTemplate > cString {
893+
startPosition = i
894+
break
895+
}
896+
}
897+
}
898+
endPosition := len(orderedTemplates)
899+
if s.paginationLimit != nil {
900+
// Make sure we don't go out of bounds
901+
if len(orderedTemplates) > startPosition+*s.paginationLimit {
902+
endPosition = startPosition + *s.paginationLimit
903+
}
904+
}
905+
906+
templatesToReturn := make([]*ResourceTemplateSchema, 0)
907+
for i := startPosition; i < endPosition; i++ {
908+
t := orderedTemplates[i]
909+
templatesToReturn = append(templatesToReturn, &ResourceTemplateSchema{
910+
Annotations: nil,
911+
Description: &t.Description,
912+
MimeType: &t.MimeType,
913+
Name: t.Name,
914+
UriTemplate: t.UriTemplate,
915+
})
916+
}
917+
918+
return ListResourceTemplatesResponse{
919+
Templates: templatesToReturn,
920+
NextCursor: func() *string {
921+
if s.paginationLimit != nil && len(templatesToReturn) >= *s.paginationLimit {
922+
toString := base64.StdEncoding.EncodeToString([]byte(templatesToReturn[len(templatesToReturn)-1].UriTemplate))
923+
return &toString
924+
}
925+
return nil
926+
}(),
927+
}, nil
928+
}
929+
828930
func (s *Server) handlePromptCalls(ctx context.Context, req *transport.BaseJSONRPCRequest, extra protocol.RequestHandlerExtra) (transport.JsonRpcBody, error) {
829931
params := baseGetPromptRequestParamsArguments{}
830932
// Instantiate a struct of the type of the arguments

Diff for: ‎server_test.go

+134
Original file line numberDiff line numberDiff line change
@@ -536,3 +536,137 @@ func TestHandleListResourcesPagination(t *testing.T) {
536536
t.Error("Expected no next cursor when pagination is disabled")
537537
}
538538
}
539+
540+
func TestHandleListResourceTemplatesPagination(t *testing.T) {
541+
mockTransport := testingutils.NewMockTransport()
542+
server := NewServer(mockTransport)
543+
err := server.Serve()
544+
if err != nil {
545+
t.Fatal(err)
546+
}
547+
548+
// Register templates in a non alphabetical order
549+
templateURIs := []string{
550+
"b://{param}/resource",
551+
"a://{param}/resource",
552+
"c://{param}/resource",
553+
"e://{param}/resource",
554+
"d://{param}/resource",
555+
}
556+
for _, uri := range templateURIs {
557+
err = server.RegisterResourceTemplate(
558+
uri,
559+
"template-"+uri,
560+
"Test template "+uri,
561+
"text/plain",
562+
)
563+
if err != nil {
564+
t.Fatal(err)
565+
}
566+
}
567+
568+
// Set pagination limit to 2 items per page
569+
limit := 2
570+
server.paginationLimit = &limit
571+
572+
// Test first page (no cursor)
573+
resp, err := server.handleListResourceTemplates(context.Background(), &transport.BaseJSONRPCRequest{
574+
Params: []byte(`{}`),
575+
}, protocol.RequestHandlerExtra{})
576+
if err != nil {
577+
t.Fatal(err)
578+
}
579+
580+
templatesResp, ok := resp.(ListResourceTemplatesResponse)
581+
if !ok {
582+
t.Fatal("Expected ListResourceTemplatesResponse")
583+
}
584+
585+
// Verify first page
586+
if len(templatesResp.Templates) != 2 {
587+
t.Errorf("Expected 2 templates, got %d", len(templatesResp.Templates))
588+
}
589+
if templatesResp.Templates[0].UriTemplate != "a://{param}/resource" || templatesResp.Templates[1].UriTemplate != "b://{param}/resource" {
590+
t.Errorf("Unexpected templates in first page: %v", templatesResp.Templates)
591+
}
592+
if templatesResp.NextCursor == nil {
593+
t.Fatal("Expected next cursor for first page")
594+
}
595+
596+
// Test second page
597+
resp, err = server.handleListResourceTemplates(context.Background(), &transport.BaseJSONRPCRequest{
598+
Params: []byte(`{"cursor":"` + *templatesResp.NextCursor + `"}`),
599+
}, protocol.RequestHandlerExtra{})
600+
if err != nil {
601+
t.Fatal(err)
602+
}
603+
604+
templatesResp, ok = resp.(ListResourceTemplatesResponse)
605+
if !ok {
606+
t.Fatal("Expected ListResourceTemplatesResponse")
607+
}
608+
609+
// Verify second page
610+
if len(templatesResp.Templates) != 2 {
611+
t.Errorf("Expected 2 templates, got %d", len(templatesResp.Templates))
612+
}
613+
if templatesResp.Templates[0].UriTemplate != "c://{param}/resource" || templatesResp.Templates[1].UriTemplate != "d://{param}/resource" {
614+
t.Errorf("Unexpected templates in second page: %v", templatesResp.Templates)
615+
}
616+
if templatesResp.NextCursor == nil {
617+
t.Fatal("Expected next cursor for second page")
618+
}
619+
620+
// Test last page
621+
resp, err = server.handleListResourceTemplates(context.Background(), &transport.BaseJSONRPCRequest{
622+
Params: []byte(`{"cursor":"` + *templatesResp.NextCursor + `"}`),
623+
}, protocol.RequestHandlerExtra{})
624+
if err != nil {
625+
t.Fatal(err)
626+
}
627+
628+
templatesResp, ok = resp.(ListResourceTemplatesResponse)
629+
if !ok {
630+
t.Fatal("Expected ListResourceTemplatesResponse")
631+
}
632+
633+
// Verify last page
634+
if len(templatesResp.Templates) != 1 {
635+
t.Errorf("Expected 1 template, got %d", len(templatesResp.Templates))
636+
}
637+
if templatesResp.Templates[0].UriTemplate != "e://{param}/resource" {
638+
t.Errorf("Unexpected template in last page: %v", templatesResp.Templates)
639+
}
640+
if templatesResp.NextCursor != nil {
641+
t.Error("Expected no next cursor for last page")
642+
}
643+
644+
// Test invalid cursor
645+
_, err = server.handleListResourceTemplates(context.Background(), &transport.BaseJSONRPCRequest{
646+
Params: []byte(`{"cursor":"invalid-cursor"}`),
647+
}, protocol.RequestHandlerExtra{})
648+
if err == nil {
649+
t.Error("Expected error for invalid cursor")
650+
}
651+
652+
// Test without pagination (should return all templates)
653+
server.paginationLimit = nil
654+
resp, err = server.handleListResourceTemplates(context.Background(), &transport.BaseJSONRPCRequest{
655+
Params: []byte(`{}`),
656+
}, protocol.RequestHandlerExtra{})
657+
if err != nil {
658+
t.Fatal(err)
659+
}
660+
661+
templatesResp, ok = resp.(ListResourceTemplatesResponse)
662+
if !ok {
663+
t.Fatal("Expected ListResourceTemplatesResponse")
664+
}
665+
666+
if len(templatesResp.Templates) != 5 {
667+
t.Errorf("Expected 5 templates, got %d", len(templatesResp.Templates))
668+
}
669+
if templatesResp.NextCursor != nil {
670+
t.Error("Expected no next cursor when pagination is disabled")
671+
}
672+
}

0 commit comments

Comments
 (0)