Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 7 additions & 6 deletions server/errors.go
Original file line number Diff line number Diff line change
Expand Up @@ -13,12 +13,13 @@ var (
ErrToolNotFound = errors.New("tool not found")

// Session-related errors
ErrSessionNotFound = errors.New("session not found")
ErrSessionExists = errors.New("session already exists")
ErrSessionNotInitialized = errors.New("session not properly initialized")
ErrSessionDoesNotSupportTools = errors.New("session does not support per-session tools")
ErrSessionDoesNotSupportResources = errors.New("session does not support per-session resources")
ErrSessionDoesNotSupportLogging = errors.New("session does not support setting logging level")
ErrSessionNotFound = errors.New("session not found")
ErrSessionExists = errors.New("session already exists")
ErrSessionNotInitialized = errors.New("session not properly initialized")
ErrSessionDoesNotSupportTools = errors.New("session does not support per-session tools")
ErrSessionDoesNotSupportResources = errors.New("session does not support per-session resources")
ErrSessionDoesNotSupportResourceTemplates = errors.New("session does not support resource templates")
ErrSessionDoesNotSupportLogging = errors.New("session does not support setting logging level")

// Notification-related errors
ErrNotificationNotInitialized = errors.New("notification channel not initialized")
Expand Down
80 changes: 66 additions & 14 deletions server/server.go
Original file line number Diff line number Diff line change
Expand Up @@ -880,12 +880,34 @@ func (s *MCPServer) handleListResourceTemplates(
id any,
request mcp.ListResourceTemplatesRequest,
) (*mcp.ListResourceTemplatesResult, *requestError) {
// Get global templates
s.resourcesMu.RLock()
templates := make([]mcp.ResourceTemplate, 0, len(s.resourceTemplates))
for _, entry := range s.resourceTemplates {
templates = append(templates, entry.template)
templateMap := make(map[string]mcp.ResourceTemplate, len(s.resourceTemplates))
for uri, entry := range s.resourceTemplates {
templateMap[uri] = entry.template
}
s.resourcesMu.RUnlock()

// Check if there are session-specific resource templates
session := ClientSessionFromContext(ctx)
if session != nil {
if sessionWithTemplates, ok := session.(SessionWithResourceTemplates); ok {
if sessionTemplates := sessionWithTemplates.GetSessionResourceTemplates(); sessionTemplates != nil {
// Merge session-specific templates with global templates
// Session templates override global ones
for uriTemplate, serverTemplate := range sessionTemplates {
templateMap[uriTemplate] = serverTemplate.Template
}
}
}
}

// Convert map to slice for sorting and pagination
templates := make([]mcp.ResourceTemplate, 0, len(templateMap))
for _, template := range templateMap {
templates = append(templates, template)
}

sort.Slice(templates, func(i, j int) bool {
return templates[i].Name < templates[j].Name
})
Expand Down Expand Up @@ -971,18 +993,48 @@ func (s *MCPServer) handleReadResource(
// If no direct handler found, try matching against templates
var matchedHandler ResourceTemplateHandlerFunc
var matched bool
for _, entry := range s.resourceTemplates {
template := entry.template
if matchesTemplate(request.Params.URI, template.URITemplate) {
matchedHandler = entry.handler
matched = true
matchedVars := template.URITemplate.Match(request.Params.URI)
// Convert matched variables to a map
request.Params.Arguments = make(map[string]any, len(matchedVars))
for name, value := range matchedVars {
request.Params.Arguments[name] = value.V

// First check session templates if available
if session != nil {
if sessionWithTemplates, ok := session.(SessionWithResourceTemplates); ok {
sessionTemplates := sessionWithTemplates.GetSessionResourceTemplates()
for _, serverTemplate := range sessionTemplates {
if serverTemplate.Template.URITemplate == nil {
continue
}
if matchesTemplate(request.Params.URI, serverTemplate.Template.URITemplate) {
matchedHandler = serverTemplate.Handler
matched = true
matchedVars := serverTemplate.Template.URITemplate.Match(request.Params.URI)
// Convert matched variables to a map
request.Params.Arguments = make(map[string]any, len(matchedVars))
for name, value := range matchedVars {
request.Params.Arguments[name] = value.V
}
break
}
}
}
}

// If not found in session templates, check global templates
if !matched {
for _, entry := range s.resourceTemplates {
template := entry.template
if template.URITemplate == nil {
continue
}
if matchesTemplate(request.Params.URI, template.URITemplate) {
matchedHandler = entry.handler
matched = true
matchedVars := template.URITemplate.Match(request.Params.URI)
// Convert matched variables to a map
request.Params.Arguments = make(map[string]any, len(matchedVars))
for name, value := range matchedVars {
request.Params.Arguments[name] = value.V
}
break
}
break
}
}
s.resourcesMu.RUnlock()
Expand Down
145 changes: 145 additions & 0 deletions server/session.go
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,17 @@ type SessionWithResources interface {
SetSessionResources(resources map[string]ServerResource)
}

// SessionWithResourceTemplates is an extension of ClientSession that can store session-specific resource template data
type SessionWithResourceTemplates interface {
ClientSession
// GetSessionResourceTemplates returns the resource templates specific to this session, if any
// This method must be thread-safe for concurrent access
GetSessionResourceTemplates() map[string]ServerResourceTemplate
// SetSessionResourceTemplates sets resource templates specific to this session
// This method must be thread-safe for concurrent access
SetSessionResourceTemplates(templates map[string]ServerResourceTemplate)
}

// SessionWithClientInfo is an extension of ClientSession that can store client info
type SessionWithClientInfo interface {
ClientSession
Expand Down Expand Up @@ -613,3 +624,137 @@ func (s *MCPServer) DeleteSessionResources(sessionID string, uris ...string) err

return nil
}

// AddSessionResourceTemplate adds a resource template for a specific session
func (s *MCPServer) AddSessionResourceTemplate(sessionID string, template mcp.ResourceTemplate, handler ResourceTemplateHandlerFunc) error {
return s.AddSessionResourceTemplates(sessionID, ServerResourceTemplate{
Template: template,
Handler: handler,
})
}

// AddSessionResourceTemplates adds resource templates for a specific session
func (s *MCPServer) AddSessionResourceTemplates(sessionID string, templates ...ServerResourceTemplate) error {
sessionValue, ok := s.sessions.Load(sessionID)
if !ok {
return ErrSessionNotFound
}

session, ok := sessionValue.(SessionWithResourceTemplates)
if !ok {
return ErrSessionDoesNotSupportResourceTemplates
}

// For session resource templates, enable listChanged by default
// This is the same behavior as session resources
s.implicitlyRegisterCapabilities(
func() bool { return s.capabilities.resources != nil },
func() { s.capabilities.resources = &resourceCapabilities{listChanged: true} },
)

// Get existing templates (this returns a thread-safe copy)
sessionTemplates := session.GetSessionResourceTemplates()

// Create a new map to avoid modifying the returned copy
newTemplates := make(map[string]ServerResourceTemplate, len(sessionTemplates)+len(templates))

// Copy existing templates
for k, v := range sessionTemplates {
newTemplates[k] = v
}

// Validate and add new templates
for _, t := range templates {
if t.Template.URITemplate == nil {
return fmt.Errorf("resource template URITemplate cannot be nil")
}
raw := t.Template.URITemplate.Raw()
if raw == "" {
return fmt.Errorf("resource template URITemplate cannot be empty")
}
if t.Template.Name == "" {
return fmt.Errorf("resource template name cannot be empty")
}
newTemplates[raw] = t
}

// Set the new templates (this method must handle thread-safety)
session.SetSessionResourceTemplates(newTemplates)

// Send notification if the session is initialized and listChanged is enabled
if session.Initialized() && s.capabilities.resources != nil && s.capabilities.resources.listChanged {
if err := s.SendNotificationToSpecificClient(sessionID, "notifications/resources/list_changed", nil); err != nil {
// Log the error but don't fail the operation
if s.hooks != nil && len(s.hooks.OnError) > 0 {
hooks := s.hooks
go func(sID string, hooks *Hooks) {
ctx := context.Background()
hooks.onError(ctx, nil, "notification", map[string]any{
"method": "notifications/resources/list_changed",
"sessionID": sID,
}, fmt.Errorf("failed to send notification after adding resource templates: %w", err))
}(sessionID, hooks)
}
}
}

return nil
}

// DeleteSessionResourceTemplates removes resource templates from a specific session
func (s *MCPServer) DeleteSessionResourceTemplates(sessionID string, uriTemplates ...string) error {
sessionValue, ok := s.sessions.Load(sessionID)
if !ok {
return ErrSessionNotFound
}

session, ok := sessionValue.(SessionWithResourceTemplates)
if !ok {
return ErrSessionDoesNotSupportResourceTemplates
}

// Get existing templates (this returns a thread-safe copy)
sessionTemplates := session.GetSessionResourceTemplates()

// Track if any were actually deleted
deletedAny := false

// Create a new map without the deleted templates
newTemplates := make(map[string]ServerResourceTemplate, len(sessionTemplates))
for k, v := range sessionTemplates {
newTemplates[k] = v
}

// Delete specified templates
for _, uriTemplate := range uriTemplates {
if _, exists := newTemplates[uriTemplate]; exists {
delete(newTemplates, uriTemplate)
deletedAny = true
}
}

// Only update if something was actually deleted
if deletedAny {
// Set the new templates (this method must handle thread-safety)
session.SetSessionResourceTemplates(newTemplates)

// Send notification if the session is initialized and listChanged is enabled
if session.Initialized() && s.capabilities.resources != nil && s.capabilities.resources.listChanged {
if err := s.SendNotificationToSpecificClient(sessionID, "notifications/resources/list_changed", nil); err != nil {
// Log the error but don't fail the operation
if s.hooks != nil && len(s.hooks.OnError) > 0 {
hooks := s.hooks
go func(sID string, hooks *Hooks) {
ctx := context.Background()
hooks.onError(ctx, nil, "notification", map[string]any{
"method": "notifications/resources/list_changed",
"sessionID": sID,
}, fmt.Errorf("failed to send notification after deleting resource templates: %w", err))
}(sessionID, hooks)
}
}
}
}

return nil
}
Loading
Loading