Skip to content

Commit e00d4be

Browse files
nikolubbenikolubbe
authored
Add ResourceTemplates (#73)
1 parent 6f6550e commit e00d4be

File tree

3 files changed

+267
-5
lines changed

3 files changed

+267
-5
lines changed

‎resource_response_types.go

Lines changed: 26 additions & 0 deletions
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+
}

‎server.go

Lines changed: 107 additions & 5 deletions
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 {
@@ -829,6 +859,78 @@ func (s *Server) handleListResources(ctx context.Context, request *transport.Bas
829859
}, nil
830860
}
831861

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

‎server_test.go

Lines changed: 134 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -572,3 +572,137 @@ func TestHandleListResourcesPagination(t *testing.T) {
572572
t.Error("Expected no next cursor when pagination is disabled")
573573
}
574574
}
575+
576+
func TestHandleListResourceTemplatesPagination(t *testing.T) {
577+
mockTransport := testingutils.NewMockTransport()
578+
server := NewServer(mockTransport)
579+
err := server.Serve()
580+
if err != nil {
581+
t.Fatal(err)
582+
}
583+
584+
// Register templates in a non alphabetical order
585+
templateURIs := []string{
586+
"b://{param}/resource",
587+
"a://{param}/resource",
588+
"c://{param}/resource",
589+
"e://{param}/resource",
590+
"d://{param}/resource",
591+
}
592+
for _, uri := range templateURIs {
593+
err = server.RegisterResourceTemplate(
594+
uri,
595+
"template-"+uri,
596+
"Test template "+uri,
597+
"text/plain",
598+
)
599+
if err != nil {
600+
t.Fatal(err)
601+
}
602+
}
603+
604+
// Set pagination limit to 2 items per page
605+
limit := 2
606+
server.paginationLimit = &limit
607+
608+
// Test first page (no cursor)
609+
resp, err := server.handleListResourceTemplates(context.Background(), &transport.BaseJSONRPCRequest{
610+
Params: []byte(`{}`),
611+
}, protocol.RequestHandlerExtra{})
612+
if err != nil {
613+
t.Fatal(err)
614+
}
615+
616+
templatesResp, ok := resp.(ListResourceTemplatesResponse)
617+
if !ok {
618+
t.Fatal("Expected ListResourceTemplatesResponse")
619+
}
620+
621+
// Verify first page
622+
if len(templatesResp.Templates) != 2 {
623+
t.Errorf("Expected 2 templates, got %d", len(templatesResp.Templates))
624+
}
625+
if templatesResp.Templates[0].UriTemplate != "a://{param}/resource" || templatesResp.Templates[1].UriTemplate != "b://{param}/resource" {
626+
t.Errorf("Unexpected templates in first page: %v", templatesResp.Templates)
627+
}
628+
if templatesResp.NextCursor == nil {
629+
t.Fatal("Expected next cursor for first page")
630+
}
631+
632+
// Test second page
633+
resp, err = server.handleListResourceTemplates(context.Background(), &transport.BaseJSONRPCRequest{
634+
Params: []byte(`{"cursor":"` + *templatesResp.NextCursor + `"}`),
635+
}, protocol.RequestHandlerExtra{})
636+
if err != nil {
637+
t.Fatal(err)
638+
}
639+
640+
templatesResp, ok = resp.(ListResourceTemplatesResponse)
641+
if !ok {
642+
t.Fatal("Expected ListResourceTemplatesResponse")
643+
}
644+
645+
// Verify second page
646+
if len(templatesResp.Templates) != 2 {
647+
t.Errorf("Expected 2 templates, got %d", len(templatesResp.Templates))
648+
}
649+
if templatesResp.Templates[0].UriTemplate != "c://{param}/resource" || templatesResp.Templates[1].UriTemplate != "d://{param}/resource" {
650+
t.Errorf("Unexpected templates in second page: %v", templatesResp.Templates)
651+
}
652+
if templatesResp.NextCursor == nil {
653+
t.Fatal("Expected next cursor for second page")
654+
}
655+
656+
// Test last page
657+
resp, err = server.handleListResourceTemplates(context.Background(), &transport.BaseJSONRPCRequest{
658+
Params: []byte(`{"cursor":"` + *templatesResp.NextCursor + `"}`),
659+
}, protocol.RequestHandlerExtra{})
660+
if err != nil {
661+
t.Fatal(err)
662+
}
663+
664+
templatesResp, ok = resp.(ListResourceTemplatesResponse)
665+
if !ok {
666+
t.Fatal("Expected ListResourceTemplatesResponse")
667+
}
668+
669+
// Verify last page
670+
if len(templatesResp.Templates) != 1 {
671+
t.Errorf("Expected 1 template, got %d", len(templatesResp.Templates))
672+
}
673+
if templatesResp.Templates[0].UriTemplate != "e://{param}/resource" {
674+
t.Errorf("Unexpected template in last page: %v", templatesResp.Templates)
675+
}
676+
if templatesResp.NextCursor != nil {
677+
t.Error("Expected no next cursor for last page")
678+
}
679+
680+
// Test invalid cursor
681+
_, err = server.handleListResourceTemplates(context.Background(), &transport.BaseJSONRPCRequest{
682+
Params: []byte(`{"cursor":"invalid-cursor"}`),
683+
}, protocol.RequestHandlerExtra{})
684+
if err == nil {
685+
t.Error("Expected error for invalid cursor")
686+
}
687+
688+
// Test without pagination (should return all templates)
689+
server.paginationLimit = nil
690+
resp, err = server.handleListResourceTemplates(context.Background(), &transport.BaseJSONRPCRequest{
691+
Params: []byte(`{}`),
692+
}, protocol.RequestHandlerExtra{})
693+
if err != nil {
694+
t.Fatal(err)
695+
}
696+
697+
templatesResp, ok = resp.(ListResourceTemplatesResponse)
698+
if !ok {
699+
t.Fatal("Expected ListResourceTemplatesResponse")
700+
}
701+
702+
if len(templatesResp.Templates) != 5 {
703+
t.Errorf("Expected 5 templates, got %d", len(templatesResp.Templates))
704+
}
705+
if templatesResp.NextCursor != nil {
706+
t.Error("Expected no next cursor when pagination is disabled")
707+
}
708+
}

0 commit comments

Comments
 (0)