Skip to content

Commit 1350e4c

Browse files
ezynda3ezynda3
committed
refactor: address CodeRabbit review comments
- Use atomic.Bool for initialized field in test mock to prevent data races - Add nil checks for URITemplate in handleReadResource for both session and global templates - Add validation in AddSessionResourceTemplates to prevent nil URITemplate, empty URI, or empty Name - Simplify Get/SetSessionResourceTemplates using maps.Clone - Fix test initialization states to match expected behavior All tests pass with race detector.
1 parent 952ecbd commit 1350e4c

File tree

3 files changed

+31
-35
lines changed

3 files changed

+31
-35
lines changed

‎server/server.go‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -999,6 +999,9 @@ func (s *MCPServer) handleReadResource(
999999
if sessionWithTemplates, ok := session.(SessionWithResourceTemplates); ok {
10001000
sessionTemplates := sessionWithTemplates.GetSessionResourceTemplates()
10011001
for _, serverTemplate := range sessionTemplates {
1002+
if serverTemplate.Template.URITemplate == nil {
1003+
continue
1004+
}
10021005
if matchesTemplate(request.Params.URI, serverTemplate.Template.URITemplate) {
10031006
matchedHandler = serverTemplate.Handler
10041007
matched = true
@@ -1018,6 +1021,9 @@ func (s *MCPServer) handleReadResource(
10181021
if !matched {
10191022
for _, entry := range s.resourceTemplates {
10201023
template := entry.template
1024+
if template.URITemplate == nil {
1025+
continue
1026+
}
10211027
if matchesTemplate(request.Params.URI, template.URITemplate) {
10221028
matchedHandler = entry.handler
10231029
matched = true

‎server/session.go‎

Lines changed: 13 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -663,10 +663,19 @@ func (s *MCPServer) AddSessionResourceTemplates(sessionID string, templates ...S
663663
newTemplates[k] = v
664664
}
665665

666-
// Add new templates
667-
for _, template := range templates {
668-
key := template.Template.URITemplate.Raw()
669-
newTemplates[key] = template
666+
// Validate and add new templates
667+
for _, t := range templates {
668+
if t.Template.URITemplate == nil {
669+
return fmt.Errorf("resource template URITemplate cannot be nil")
670+
}
671+
raw := t.Template.URITemplate.Raw()
672+
if raw == "" {
673+
return fmt.Errorf("resource template URITemplate cannot be empty")
674+
}
675+
if t.Template.Name == "" {
676+
return fmt.Errorf("resource template name cannot be empty")
677+
}
678+
newTemplates[raw] = t
670679
}
671680

672681
// Set the new templates (this method must handle thread-safety)

‎server/session_resource_templates_test.go‎

Lines changed: 12 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import (
55
"encoding/json"
66
"maps"
77
"sync"
8+
"sync/atomic"
89
"testing"
910
"time"
1011

@@ -17,7 +18,7 @@ import (
1718
type sessionTestClientWithResourceTemplates struct {
1819
sessionID string
1920
notificationChannel chan mcp.JSONRPCNotification
20-
initialized bool
21+
initialized atomic.Bool
2122
sessionResourceTemplates map[string]ServerResourceTemplate
2223
mu sync.RWMutex
2324
}
@@ -31,38 +32,23 @@ func (f *sessionTestClientWithResourceTemplates) NotificationChannel() chan<- mc
3132
}
3233

3334
func (f *sessionTestClientWithResourceTemplates) Initialize() {
34-
f.initialized = true
35+
f.initialized.Store(true)
3536
}
3637

3738
func (f *sessionTestClientWithResourceTemplates) Initialized() bool {
38-
return f.initialized
39+
return f.initialized.Load()
3940
}
4041

4142
func (f *sessionTestClientWithResourceTemplates) GetSessionResourceTemplates() map[string]ServerResourceTemplate {
4243
f.mu.RLock()
4344
defer f.mu.RUnlock()
44-
45-
if f.sessionResourceTemplates == nil {
46-
return nil
47-
}
48-
49-
templatesCopy := make(map[string]ServerResourceTemplate, len(f.sessionResourceTemplates))
50-
maps.Copy(templatesCopy, f.sessionResourceTemplates)
51-
return templatesCopy
45+
return maps.Clone(f.sessionResourceTemplates)
5246
}
5347

5448
func (f *sessionTestClientWithResourceTemplates) SetSessionResourceTemplates(templates map[string]ServerResourceTemplate) {
5549
f.mu.Lock()
5650
defer f.mu.Unlock()
57-
58-
if templates == nil {
59-
f.sessionResourceTemplates = nil
60-
return
61-
}
62-
63-
templatesCopy := make(map[string]ServerResourceTemplate, len(templates))
64-
maps.Copy(templatesCopy, templates)
65-
f.sessionResourceTemplates = templatesCopy
51+
f.sessionResourceTemplates = maps.Clone(templates)
6652
}
6753

6854
var _ SessionWithResourceTemplates = (*sessionTestClientWithResourceTemplates)(nil)
@@ -83,11 +69,11 @@ func TestSessionWithResourceTemplates_Integration(t *testing.T) {
8369
session := &sessionTestClientWithResourceTemplates{
8470
sessionID: "session-1",
8571
notificationChannel: make(chan mcp.JSONRPCNotification, 10),
86-
initialized: true,
8772
sessionResourceTemplates: map[string]ServerResourceTemplate{
8873
"test://session/{id}": sessionTemplate,
8974
},
9075
}
76+
session.initialized.Store(true)
9177

9278
err := server.RegisterSession(context.Background(), session)
9379
require.NoError(t, err)
@@ -151,7 +137,6 @@ func TestMCPServer_ResourceTemplatesWithSessionResourceTemplates(t *testing.T) {
151137
session := &sessionTestClientWithResourceTemplates{
152138
sessionID: "session-1",
153139
notificationChannel: make(chan mcp.JSONRPCNotification, 10),
154-
initialized: true,
155140
sessionResourceTemplates: map[string]ServerResourceTemplate{
156141
"test://global/{id}": {
157142
Template: mcp.NewResourceTemplate("test://global/{id}", "global-template-1-overridden"),
@@ -164,6 +149,7 @@ func TestMCPServer_ResourceTemplatesWithSessionResourceTemplates(t *testing.T) {
164149
},
165150
},
166151
}
152+
session.initialized.Store(true)
167153

168154
err := server.RegisterSession(context.Background(), session)
169155
require.NoError(t, err)
@@ -225,8 +211,8 @@ func TestMCPServer_AddSessionResourceTemplates(t *testing.T) {
225211
session := &sessionTestClientWithResourceTemplates{
226212
sessionID: "session-1",
227213
notificationChannel: sessionChan,
228-
initialized: true,
229214
}
215+
session.initialized.Store(true)
230216

231217
err := server.RegisterSession(ctx, session)
232218
require.NoError(t, err)
@@ -255,8 +241,8 @@ func TestMCPServer_AddSessionResourceTemplate(t *testing.T) {
255241
session := &sessionTestClientWithResourceTemplates{
256242
sessionID: "session-1",
257243
notificationChannel: sessionChan,
258-
initialized: true,
259244
}
245+
session.initialized.Store(true)
260246

261247
err := server.RegisterSession(ctx, session)
262248
require.NoError(t, err)
@@ -301,9 +287,8 @@ func TestMCPServer_AddSessionResourceTemplatesUninitialized(t *testing.T) {
301287

302288
sessionChan := make(chan mcp.JSONRPCNotification, 1)
303289
session := &sessionTestClientWithResourceTemplates{
304-
sessionID: "uninitialized-session",
290+
sessionID: "session-1",
305291
notificationChannel: sessionChan,
306-
initialized: false,
307292
}
308293

309294
err := server.RegisterSession(ctx, session)
@@ -373,7 +358,6 @@ func TestMCPServer_DeleteSessionResourceTemplatesUninitialized(t *testing.T) {
373358
session := &sessionTestClientWithResourceTemplates{
374359
sessionID: "uninitialized-session",
375360
notificationChannel: sessionChan,
376-
initialized: false,
377361
sessionResourceTemplates: map[string]ServerResourceTemplate{
378362
"test://delete/{id}": {Template: mcp.NewResourceTemplate("test://delete/{id}", "template-to-delete")},
379363
"test://keep/{id}": {Template: mcp.NewResourceTemplate("test://keep/{id}", "template-to-keep")},
@@ -440,7 +424,6 @@ func TestMCPServer_CallSessionResourceTemplate(t *testing.T) {
440424
session := &sessionTestClientWithResourceTemplates{
441425
sessionID: "session-1",
442426
notificationChannel: sessionChan,
443-
initialized: true,
444427
}
445428

446429
err := server.RegisterSession(context.Background(), session)
@@ -492,7 +475,6 @@ func TestMCPServer_DeleteSessionResourceTemplates(t *testing.T) {
492475
session := &sessionTestClientWithResourceTemplates{
493476
sessionID: "session-1",
494477
notificationChannel: sessionChan,
495-
initialized: true,
496478
sessionResourceTemplates: map[string]ServerResourceTemplate{
497479
"test://template1/{id}": {
498480
Template: mcp.NewResourceTemplate("test://template1/{id}", "session-template-1"),
@@ -502,6 +484,7 @@ func TestMCPServer_DeleteSessionResourceTemplates(t *testing.T) {
502484
},
503485
},
504486
}
487+
session.initialized.Store(true)
505488

506489
err := server.RegisterSession(ctx, session)
507490
require.NoError(t, err)
@@ -528,7 +511,6 @@ func TestMCPServer_SessionResourceTemplateError(t *testing.T) {
528511
session := &sessionTestClient{
529512
sessionID: "session-1",
530513
notificationChannel: make(chan mcp.JSONRPCNotification, 10),
531-
initialized: true,
532514
}
533515

534516
err := server.RegisterSession(ctx, session)
@@ -549,7 +531,6 @@ func TestMCPServer_ResourceTemplatesNotificationsDisabled(t *testing.T) {
549531
session := &sessionTestClientWithResourceTemplates{
550532
sessionID: "session-1",
551533
notificationChannel: sessionChan,
552-
initialized: true,
553534
}
554535

555536
err := server.RegisterSession(ctx, session)

0 commit comments

Comments
 (0)