66 "os"
77 "slices"
88 "sort"
9+ "sync"
910
1011 "github.com/modelcontextprotocol/go-sdk/mcp"
1112)
@@ -25,12 +26,21 @@ import (
2526// - Lazy dependency injection during registration via RegisterAll()
2627// - Runtime toolset enabling for dynamic toolsets mode
2728type Registry struct {
28- // tools holds all tools in this group
29+ // tools holds all tools in this group (ordered for iteration)
2930 tools []ServerTool
30- // resourceTemplates holds all resource templates in this group
31+ // toolsByName provides O(1) lookup by tool name (lazy-initialized)
32+ toolsByName map [string ]* ServerTool
33+ toolsByNameOnce sync.Once
34+ // resourceTemplates holds all resource templates in this group (ordered for iteration)
3135 resourceTemplates []ServerResourceTemplate
32- // prompts holds all prompts in this group
36+ // resourcesByURI provides O(1) lookup by URI template (lazy-initialized)
37+ resourcesByURI map [string ]* ServerResourceTemplate
38+ resourcesByURIOnce sync.Once
39+ // prompts holds all prompts in this group (ordered for iteration)
3340 prompts []ServerPrompt
41+ // promptsByName provides O(1) lookup by prompt name (lazy-initialized)
42+ promptsByName map [string ]* ServerPrompt
43+ promptsByNameOnce sync.Once
3444 // deprecatedAliases maps old tool names to new canonical names
3545 deprecatedAliases map [string ]string
3646
@@ -57,6 +67,39 @@ func (r *Registry) UnrecognizedToolsets() []string {
5767 return r .unrecognizedToolsets
5868}
5969
70+ // getToolsByName returns the toolsByName map, initializing it lazily on first call.
71+ func (r * Registry ) getToolsByName () map [string ]* ServerTool {
72+ r .toolsByNameOnce .Do (func () {
73+ r .toolsByName = make (map [string ]* ServerTool , len (r .tools ))
74+ for i := range r .tools {
75+ r .toolsByName [r .tools [i ].Tool .Name ] = & r .tools [i ]
76+ }
77+ })
78+ return r .toolsByName
79+ }
80+
81+ // getResourcesByURI returns the resourcesByURI map, initializing it lazily on first call.
82+ func (r * Registry ) getResourcesByURI () map [string ]* ServerResourceTemplate {
83+ r .resourcesByURIOnce .Do (func () {
84+ r .resourcesByURI = make (map [string ]* ServerResourceTemplate , len (r .resourceTemplates ))
85+ for i := range r .resourceTemplates {
86+ r .resourcesByURI [r .resourceTemplates [i ].Template .URITemplate ] = & r .resourceTemplates [i ]
87+ }
88+ })
89+ return r .resourcesByURI
90+ }
91+
92+ // getPromptsByName returns the promptsByName map, initializing it lazily on first call.
93+ func (r * Registry ) getPromptsByName () map [string ]* ServerPrompt {
94+ r .promptsByNameOnce .Do (func () {
95+ r .promptsByName = make (map [string ]* ServerPrompt , len (r .prompts ))
96+ for i := range r .prompts {
97+ r .promptsByName [r .prompts [i ].Prompt .Name ] = & r .prompts [i ]
98+ }
99+ })
100+ return r .promptsByName
101+ }
102+
60103// MCP method constants for use with ForMCPRequest.
61104const (
62105 MCPMethodInitialize = "initialize"
@@ -90,6 +133,8 @@ const (
90133// All existing filters (read-only, toolsets, etc.) still apply to the returned items.
91134func (r * Registry ) ForMCPRequest (method string , itemName string ) * Registry {
92135 // Create a shallow copy with shared filter settings
136+ // Note: lazy-init maps (toolsByName, etc.) are NOT copied - the new Registry
137+ // will initialize its own maps on first use if needed
93138 result := & Registry {
94139 tools : r .tools ,
95140 resourceTemplates : r .resourceTemplates ,
@@ -269,11 +314,8 @@ func (r *Registry) ResolveToolAliases(toolNames []string) (resolved []string, al
269314// Returns the tool, its toolset ID, and an error if not found.
270315// This searches ALL tools regardless of filters.
271316func (r * Registry ) FindToolByName (toolName string ) (* ServerTool , ToolsetID , error ) {
272- for i := range r .tools {
273- tool := & r .tools [i ]
274- if tool .Tool .Name == toolName {
275- return tool , tool .Toolset .ID , nil
276- }
317+ if tool , ok := r .getToolsByName ()[toolName ]; ok {
318+ return tool , tool .Toolset .ID , nil
277319 }
278320 return nil , "" , NewToolDoesNotExistError (toolName )
279321}
0 commit comments