feat: switch to lazy loading of dynamic tools

Slach · aider-chat-bot · Slach · commit b3c73ea2c5a1 · 2025-11-22T10:08:10.000+04:00
Co-authored-by: aider (gemini/gemini-3-pro-preview) &lt;aider@aider.chat&gt;
diff --git a/cmd/altinity-mcp/main.go b/cmd/altinity-mcp/main.go
@@ -327,6 +327,17 @@ func (a *application) createTokenInjector() func(http.Handler) http.Handler {
 	}
 }
 
+// dynamicToolsInjector creates a middleware that ensures dynamic tools are loaded
+func (a *application) dynamicToolsInjector(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if err := a.mcpServer.EnsureDynamicTools(r.Context()); err != nil {
+			// Log error but continue, static tools should still work
+			log.Warn().Err(err).Msg("Failed to ensure dynamic tools")
+		}
+		next.ServeHTTP(w, r)
+	})
+}
+
 // stripTrailingSlash normalizes paths to remove a single trailing slash (except root)
 func stripTrailingSlash(next http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -535,11 +546,12 @@ func (a *application) startHTTPServer(cfg config.Config, mcpServer *server.MCPSe
 		log.Info().Msg("Using dynamic base path for JWE authentication")
 
 		tokenInjector := a.createTokenInjector()
+		dtInjector := a.dynamicToolsInjector
 		httpServer := server.NewStreamableHTTPServer(mcpServer)
 
 		// Register custom handlers to ensure token is in the path and inject it into context
 		mux := http.NewServeMux()
-		mux.Handle("/{token}/http", serverInjector(tokenInjector(httpServer)))
+		mux.Handle("/{token}/http", serverInjector(tokenInjector(dtInjector(httpServer))))
         if cfg.Server.OpenAPI.Enabled {
             mux.HandleFunc("/openapi", a.mcpServer.ServeOpenAPISchema)
             mux.HandleFunc("/{token}/openapi", serverInjectorOpenAPI)
@@ -555,8 +567,9 @@ func (a *application) startHTTPServer(cfg config.Config, mcpServer *server.MCPSe
 	} else {
 		// Use standard HTTP server without dynamic paths
 		httpServer := server.NewStreamableHTTPServer(mcpServer)
+		dtInjector := a.dynamicToolsInjector
 		mux := http.NewServeMux()
-		mux.Handle("/http", serverInjector(httpServer))
+		mux.Handle("/http", serverInjector(dtInjector(httpServer)))
         if cfg.Server.OpenAPI.Enabled {
             mux.HandleFunc("/openapi", serverInjectorOpenAPI)
             mux.HandleFunc("/openapi/", serverInjectorOpenAPI)
@@ -629,6 +642,7 @@ func (a *application) startSSEServer(cfg config.Config, mcpServer *server.MCPSer
 		log.Info().Msg("Using dynamic base path for JWE authentication")
 
 		tokenInjector := a.createTokenInjector()
+		dtInjector := a.dynamicToolsInjector
 
 		sseServer := server.NewSSEServer(
 			mcpServer,
@@ -645,8 +659,8 @@ func (a *application) startSSEServer(cfg config.Config, mcpServer *server.MCPSer
 		)
 
 		mux := http.NewServeMux()
-		mux.Handle("/{token}/sse", serverInjector(tokenInjector(sseServer.SSEHandler())))
-		mux.Handle("/{token}/message", serverInjector(tokenInjector(sseServer.MessageHandler())))
+		mux.Handle("/{token}/sse", serverInjector(tokenInjector(dtInjector(sseServer.SSEHandler()))))
+		mux.Handle("/{token}/message", serverInjector(tokenInjector(dtInjector(sseServer.MessageHandler()))))
         if cfg.Server.OpenAPI.Enabled {
             mux.HandleFunc("/openapi", a.mcpServer.ServeOpenAPISchema)
             mux.HandleFunc("/{token}/openapi", serverInjectorOpenAPI)
@@ -662,9 +676,10 @@ func (a *application) startSSEServer(cfg config.Config, mcpServer *server.MCPSer
 	} else {
 		// Use standard SSE server without dynamic paths
 		sseServer := server.NewSSEServer(mcpServer)
+		dtInjector := a.dynamicToolsInjector
 		mux := http.NewServeMux()
-		mux.Handle("/sse", serverInjector(sseServer))
-		mux.Handle("/message", serverInjector(sseServer.MessageHandler()))
+		mux.Handle("/sse", serverInjector(dtInjector(sseServer)))
+		mux.Handle("/message", serverInjector(dtInjector(sseServer.MessageHandler())))
         if cfg.Server.OpenAPI.Enabled {
             mux.HandleFunc("/openapi", serverInjectorOpenAPI)
             mux.HandleFunc("/openapi/", serverInjectorOpenAPI)
diff --git a/pkg/server/server.go b/pkg/server/server.go
@@ -10,6 +10,7 @@ import (
 	"regexp"
 	"strconv"
 	"strings"
+	"sync"
 
 	"github.com/altinity/altinity-mcp/pkg/clickhouse"
 	"github.com/altinity/altinity-mcp/pkg/config"
@@ -25,7 +26,9 @@ type ClickHouseJWEServer struct {
 	Config  config.Config
 	Version string
 	// dynamic tools metadata for OpenAPI routing and schema
-	dynamicTools map[string]dynamicToolMeta
+	dynamicTools     map[string]dynamicToolMeta
+	dynamicToolsMu   sync.RWMutex
+	dynamicToolsInit bool
 }
 
 type dynamicToolParam struct {
@@ -75,8 +78,7 @@ func NewClickHouseMCPServer(cfg config.Config, version string) *ClickHouseJWESer
 
 	// Register tools, resources, and prompts
 	RegisterTools(chJweServer)
-	// dynamic tools registered after static ones
-	registerDynamicTools(chJweServer)
+	// dynamic tools registered lazily via EnsureDynamicTools
 	RegisterResources(chJweServer)
 	RegisterPrompts(chJweServer)
 
@@ -414,19 +416,28 @@ func RegisterPrompts(srv AltinityMCPServer) {
 	log.Info().Int("prompt_count", 0).Msg("ClickHouse prompts registered")
 }
 
-// registerDynamicTools discovers ClickHouse views and registers MCP/OpenAPI tools
-func registerDynamicTools(s *ClickHouseJWEServer) {
+// EnsureDynamicTools discovers ClickHouse views and registers MCP/OpenAPI tools
+func (s *ClickHouseJWEServer) EnsureDynamicTools(ctx context.Context) error {
+	s.dynamicToolsMu.Lock()
+	defer s.dynamicToolsMu.Unlock()
+
+	if s.dynamicToolsInit {
+		return nil
+	}
+
 	if len(s.Config.Server.DynamicTools) == 0 {
-		return
+		s.dynamicToolsInit = true
+		return nil
 	}
 
-	ctx := context.WithValue(context.Background(), "clickhouse_jwe_server", s)
-	token := ""
+	// Check if we have a valid client/token to proceed
+	token := s.ExtractTokenFromCtx(ctx)
 	// Get ClickHouse client
 	chClient, err := s.GetClickHouseClient(ctx, token)
 	if err != nil {
-		log.Error().Err(err).Msg("dynamic_tools: failed to get ClickHouse client")
-		return
+		// If we can't get a client (e.g. missing token when JWE enabled), we can't register dynamic tools yet
+		// Return error so we retry later
+		return fmt.Errorf("dynamic_tools: failed to get ClickHouse client: %w", err)
 	}
 	defer func() {
 		if closeErr := chClient.Close(); closeErr != nil {
@@ -438,8 +449,7 @@ func registerDynamicTools(s *ClickHouseJWEServer) {
 	q := "SELECT database, name, create_table_query, comment FROM system.tables WHERE engine='View'"
 	result, err := chClient.ExecuteQuery(ctx, q)
 	if err != nil {
-		log.Error().Err(err).Str("query", q).Msg("dynamic_tools: failed to list views")
-		return
+		return fmt.Errorf("dynamic_tools: failed to list views: %w", err)
 	}
 
 	// compile regex rules
@@ -556,6 +566,9 @@ func registerDynamicTools(s *ClickHouseJWEServer) {
 		log.Error().Msg("dynamic_tools: overlaps detected; conflicting views were skipped as per policy 'error on overlap'")
 	}
 	log.Info().Int("tool_count", dynamicCount).Msg("Dynamic ClickHouse view tools registered")
+
+	s.dynamicToolsInit = true
+	return nil
 }
 
 func makeDynamicToolHandler(meta dynamicToolMeta) server.ToolHandlerFunc {
@@ -830,11 +843,21 @@ func (s *ClickHouseJWEServer) OpenAPIHandler(w http.ResponseWriter, r *http.Requ
 	case strings.HasSuffix(r.URL.Path, "/openapi/execute_query"):
 		s.handleExecuteQueryOpenAPI(w, r, token)
 	case strings.Contains(r.URL.Path, "/openapi/") && r.Method == http.MethodPost:
+		// Ensure dynamic tools are loaded
+		if err := s.EnsureDynamicTools(r.Context()); err != nil {
+			log.Warn().Err(err).Msg("Failed to ensure dynamic tools in OpenAPI handler")
+		}
+
 		// dynamic tool endpoint: /openapi/{tool}
 		parts := strings.Split(r.URL.Path, "/openapi/")
 		if len(parts) == 2 {
 			tool := strings.Trim(parts[1], "/")
-			if meta, ok := s.dynamicTools[tool]; ok {
+
+			s.dynamicToolsMu.RLock()
+			meta, ok := s.dynamicTools[tool]
+			s.dynamicToolsMu.RUnlock()
+
+			if ok {
 				s.handleDynamicToolOpenAPI(w, r, token, meta)
 				return
 			}
@@ -847,6 +870,11 @@ func (s *ClickHouseJWEServer) OpenAPIHandler(w http.ResponseWriter, r *http.Requ
 }
 
 func (s *ClickHouseJWEServer) ServeOpenAPISchema(w http.ResponseWriter, r *http.Request) {
+	// Ensure dynamic tools are loaded
+	if err := s.EnsureDynamicTools(r.Context()); err != nil {
+		log.Warn().Err(err).Msg("Failed to ensure dynamic tools in ServeOpenAPISchema")
+	}
+
 	// Get host URL based on OpenAPI TLS configuration
 	protocol := "http"
 	if s.Config.Server.OpenAPI.TLS {
@@ -936,6 +964,10 @@ func (s *ClickHouseJWEServer) ServeOpenAPISchema(w http.ResponseWriter, r *http.
 
 	// add dynamic tool paths (POST)
 	paths := schema["paths"].(map[string]interface{})
+
+	s.dynamicToolsMu.RLock()
+	defer s.dynamicToolsMu.RUnlock()
+
 	for toolName, meta := range s.dynamicTools {
 		path := "/{jwe_token}/openapi/" + toolName
 		// request body schema
diff --git a/pkg/server/server_test.go b/pkg/server/server_test.go
@@ -1263,7 +1263,8 @@ func TestRegisterDynamicTools_SuccessAndOverlap(t *testing.T) {
         {Regexp: "default\\.v_a", Prefix: "other_"},
     }}}, dynamicTools: make(map[string]dynamicToolMeta)}
 
-    registerDynamicTools(s)
+    err = s.EnsureDynamicTools(ctx)
+    require.NoError(t, err)
 
     // v_a matches two rules -> should be skipped
     _, existsA1 := s.dynamicTools["custom_default_v_a"]
