perf: pre-compute allowedToolSets at init time for O(1) lookup in isToolAllowed and registerToolsFromBackend

Agent-Logs-Url: https://github.com/github/gh-aw-mcpg/sessions/00582754-0d76-4bb6-a950-0cd6ff534f1b

Co-authored-by: lpcox <15877973+lpcox@users.noreply.github.com>

Author: Copilot

internal/server/call_backend_tool_test.go

@@ -401,13 +401,14 @@ func TestIsToolAllowed(t *testing.T) {
 
 	for _, tc := range tests {
 		t.Run(tc.name, func(t *testing.T) {
-			us := &UnifiedServer{
-				cfg: &config.Config{
-					Servers: map[string]*config.ServerConfig{
-						"s": {Tools: tc.tools},
-					},
+			cfg := &config.Config{
+				Servers: map[string]*config.ServerConfig{
+					"s": {Tools: tc.tools},
 				},
 			}
+			us := &UnifiedServer{
+				allowedToolSets: buildAllowedToolSets(cfg),
+			}
 			got := us.isToolAllowed("s", tc.toolName)
 			assert.Equal(t, tc.wantAllow, got)
 		})
@@ -416,17 +417,14 @@ func TestIsToolAllowed(t *testing.T) {
 
 // TestIsToolAllowed_NilConfig verifies that a nil config allows all tools.
 func TestIsToolAllowed_NilConfig(t *testing.T) {
-	us := &UnifiedServer{cfg: nil}
+	us := &UnifiedServer{allowedToolSets: buildAllowedToolSets(nil)}
 	assert.True(t, us.isToolAllowed("s", "anything"), "nil cfg should allow all tools")
 }
 
 // TestIsToolAllowed_UnknownServer verifies that an unknown server ID allows all tools.
 func TestIsToolAllowed_UnknownServer(t *testing.T) {
-	us := &UnifiedServer{
-		cfg: &config.Config{
-			Servers: map[string]*config.ServerConfig{},
-		},
-	}
+	cfg := &config.Config{Servers: map[string]*config.ServerConfig{}}
+	us := &UnifiedServer{allowedToolSets: buildAllowedToolSets(cfg)}
 	assert.True(t, us.isToolAllowed("unknown", "tool"), "unknown server should allow all tools")
 }
 

internal/server/tool_registry.go

@@ -185,25 +185,19 @@ func (us *UnifiedServer) registerToolsFromBackend(serverID string) error {
 	// Filter tools by the server's allowed-tools list (if configured).
 	// This prevents non-allowed tools from appearing in tools/list responses
 	// and is defense-in-depth alongside the callBackendTool enforcement.
-	if us.cfg != nil {
-		if serverCfg, ok := us.cfg.Servers[serverID]; ok && len(serverCfg.Tools) > 0 {
-			allowedSet := make(map[string]bool, len(serverCfg.Tools))
-			for _, t := range serverCfg.Tools {
-				allowedSet[t] = true
+	if allowedSet, ok := us.allowedToolSets[serverID]; ok && len(allowedSet) > 0 {
+		n := 0
+		for _, tool := range listResult.Tools {
+			if allowedSet[tool.Name] {
+				listResult.Tools[n] = tool
+				n++
 			}
-			n := 0
-			for _, tool := range listResult.Tools {
-				if allowedSet[tool.Name] {
-					listResult.Tools[n] = tool
-					n++
-				}
-			}
-			if n < len(listResult.Tools) {
-				logger.LogInfo("backend", "[allowed-tools] Filtered %d tools from %s: keeping %d of %d",
-					len(listResult.Tools)-n, serverID, n, len(listResult.Tools))
-			}
-			listResult.Tools = listResult.Tools[:n]
 		}
+		if n < len(listResult.Tools) {
+			logger.LogInfo("backend", "[allowed-tools] Filtered %d tools from %s: keeping %d of %d",
+				len(listResult.Tools)-n, serverID, n, len(listResult.Tools))
+		}
+		listResult.Tools = listResult.Tools[:n]
 	}
 
 	// Collect tools for logging

internal/server/unified.go

@@ -90,6 +90,11 @@ type UnifiedServer struct {
 	payloadPathPrefix    string // Path prefix to use when returning payloadPath to clients (allows remapping host paths to client/agent container paths)
 	payloadSizeThreshold int    // Size threshold (in bytes) for storing payloads to disk. Payloads larger than this are stored to disk, smaller ones are returned inline.
 
+	// allowedToolSets holds a pre-computed set of allowed tool names per server ID.
+	// Built once during NewUnified from the config Tools lists. A missing or nil entry
+	// means all tools are permitted for that server.
+	allowedToolSets map[string]map[string]bool
+
 	// DIFC components
 	guardRegistry *guard.Registry
 	agentRegistry *difc.AgentRegistry
@@ -144,6 +149,7 @@ func NewUnified(ctx context.Context, cfg *config.Config) (*UnifiedServer, error)
 		payloadDir:           payloadDir,
 		payloadPathPrefix:    payloadPathPrefix,
 		payloadSizeThreshold: payloadSizeThreshold,
+		allowedToolSets:      buildAllowedToolSets(cfg),
 
 		// Initialize DIFC components
 		guardRegistry: guard.NewRegistry(),
@@ -371,25 +377,35 @@ func newErrorCallToolResult(err error) (*sdk.CallToolResult, interface{}, error)
 	}, nil, err
 }
 
+// buildAllowedToolSets converts the per-server Tools lists from the config into pre-computed
+// map[string]bool sets for O(1) lookup. Servers with no Tools list are not added to the map,
+// which signals that all tools are permitted.
+func buildAllowedToolSets(cfg *config.Config) map[string]map[string]bool {
+	sets := make(map[string]map[string]bool)
+	if cfg == nil {
+		return sets
+	}
+	for serverID, serverCfg := range cfg.Servers {
+		if len(serverCfg.Tools) > 0 {
+			set := make(map[string]bool, len(serverCfg.Tools))
+			for _, t := range serverCfg.Tools {
+				set[t] = true
+			}
+			sets[serverID] = set
+		}
+	}
+	return sets
+}
+
 // isToolAllowed reports whether toolName is permitted by the server's configured
 // allowed-tools list. When no list is configured (empty), all tools are allowed.
+// Uses the pre-computed allowedToolSets map for O(1) lookup.
 func (us *UnifiedServer) isToolAllowed(serverID, toolName string) bool {
-	if us.cfg == nil {
+	set, ok := us.allowedToolSets[serverID]
+	if !ok || set == nil {
 		return true
 	}
-	serverCfg, ok := us.cfg.Servers[serverID]
-	if !ok {
-		return true
-	}
-	if len(serverCfg.Tools) == 0 {
-		return true
-	}
-	for _, allowed := range serverCfg.Tools {
-		if allowed == toolName {
-			return true
-		}
-	}
-	return false
+	return set[toolName]
 }
 
 // callBackendTool calls a tool on a backend server with DIFC enforcement