Enhance session management and improve MCP server functionality

- Updated `PulsarFunctionManagedMcpTools` to use a fixed `sessionId` instead of an empty string, ensuring proper session handling.
- Refactored API and log client retrieval methods in the `Session` struct to utilize `sync.Once`, improving thread safety and performance.
- Added validation for `sessionID` in `PulsarFunctionManagedMcpTools` to prevent processing with empty session IDs.
- Enhanced context retrieval functions to safely return empty strings when context values are not present, improving robustness.
- Updated function names to follow consistent naming conventions, enhancing code readability and maintainability.

Author: freeznet

pkg/cmd/mcp/sse.go

@@ -89,8 +89,8 @@ func runSseServer(configOpts *ServerOptions) error {
 	}
 
 	// add Pulsar Functions as MCP tools
-	// SSE is not support session-based tools, so we pass an empty sessionId
-	mcpServer.PulsarFunctionManagedMcpTools(configOpts.ReadOnly, configOpts.Features, "")
+	// SSE is not support session-based tools, so we pass an fixed sessionId
+	mcpServer.PulsarFunctionManagedMcpTools(configOpts.ReadOnly, configOpts.Features, "FIXED_SESSION_ID")
 
 	sseServer := server.NewSSEServer(
 		mcpServer.MCPServer,

pkg/config/apiclient.go

@@ -55,6 +55,8 @@ type Session struct {
 	TokenRefresher *OAuth2TokenRefresher
 	Configuration  *sncloud.Configuration
 	mutex          sync.RWMutex
+	apiClientOnce  sync.Once
+	logClientOnce  sync.Once
 }
 
 // OAuth2TokenRefresher implements oauth2.TokenSource interface for refreshing OAuth2 tokens
@@ -182,27 +184,15 @@ func (s *Session) initializeTokenRefresher() error {
 
 // GetAPIClient returns the API client for the session, initializing it if necessary
 func (s *Session) GetAPIClient() (*sncloud.APIClient, error) {
-	s.mutex.RLock()
-	if s.APIClient != nil {
-		defer s.mutex.RUnlock()
-		return s.APIClient, nil
-	}
-	s.mutex.RUnlock()
-
-	// Need to initialize the client
-	s.mutex.Lock()
-	defer s.mutex.Unlock()
-
-	// Double check after acquiring write lock
-	if s.APIClient != nil {
-		return s.APIClient, nil
-	}
-
-	// Initialize the API client
-	if err := s.initializeAPIClient(); err != nil {
+	var err error
+	s.apiClientOnce.Do(func() {
+		err = s.initializeAPIClient()
+	})
+	
+	if err != nil {
 		return nil, errors.Wrap(err, "failed to initialize API client")
 	}
-
+	
 	return s.APIClient, nil
 }
 
@@ -245,27 +235,15 @@ func (s *Session) initializeAPIClient() error {
 
 // GetLogClient returns the log client for the session, initializing it if necessary
 func (s *Session) GetLogClient() (*http.Client, error) {
-	s.mutex.RLock()
-	if s.LogClient != nil {
-		defer s.mutex.RUnlock()
-		return s.LogClient, nil
-	}
-	s.mutex.RUnlock()
-
-	// Need to initialize the client
-	s.mutex.Lock()
-	defer s.mutex.Unlock()
-
-	// Double check after acquiring write lock
-	if s.LogClient != nil {
-		return s.LogClient, nil
-	}
-
-	// Initialize the log client
-	if err := s.initializeLogClient(); err != nil {
+	var err error
+	s.logClientOnce.Do(func() {
+		err = s.initializeLogClient()
+	})
+	
+	if err != nil {
 		return nil, errors.Wrap(err, "failed to initialize log client")
 	}
-
+	
 	return s.LogClient, nil
 }
 

pkg/mcp/ctx.go

@@ -68,17 +68,32 @@ func WithKafkaSession(ctx context.Context, session *kafka.Session) context.Conte
 
 // GetSNCloudOrganization gets the SNCloud organization from the context
 func GetSNCloudOrganization(ctx context.Context) string {
-	return ctx.Value(SNCloudOrganizationContextKey).(string)
+	if val := ctx.Value(SNCloudOrganizationContextKey); val != nil {
+		if str, ok := val.(string); ok {
+			return str
+		}
+	}
+	return ""
 }
 
 // GetSNCloudInstance gets the SNCloud instance from the context
 func GetSNCloudInstance(ctx context.Context) string {
-	return ctx.Value(SNCloudInstanceContextKey).(string)
+	if val := ctx.Value(SNCloudInstanceContextKey); val != nil {
+		if str, ok := val.(string); ok {
+			return str
+		}
+	}
+	return ""
 }
 
 // GetSNCloudCluster gets the SNCloud cluster from the context
 func GetSNCloudCluster(ctx context.Context) string {
-	return ctx.Value(SNCloudClusterContextKey).(string)
+	if val := ctx.Value(SNCloudClusterContextKey); val != nil {
+		if str, ok := val.(string); ok {
+			return str
+		}
+	}
+	return ""
 }
 
 // GetSNCloudSession gets the SNCloud session from the context

pkg/mcp/pulsar_functions_as_tools.go

@@ -58,6 +58,12 @@ func (s *Server) PulsarFunctionManagedMcpTools(readOnly bool, features []string,
 		return
 	}
 
+	// Validate sessionID
+	if sessionID == "" {
+		log.Printf("Skipping Pulsar Functions as MCP Tools because sessionID is empty")
+		return
+	}
+
 	options := pftools.DefaultManagerOptions()
 
 	if s.SNCloudSession.Ctx.Organization == "" || s.SNCloudSession.Ctx.PulsarInstance == "" || s.SNCloudSession.Ctx.PulsarCluster == "" {

pkg/mcp/server.go

@@ -35,13 +35,13 @@ type Server struct {
 
 func NewServer(name, version string, logger *logrus.Logger, opts ...server.ServerOption) *Server {
 	// Create a new MCP server
-	opts = addOpts(opts...)
+	opts = AddOpts(opts...)
 	s := server.NewMCPServer(name, version, opts...)
-	mcpserver := createSNCloudMCPServer(s, logger)
+	mcpserver := CreateSNCloudMCPServer(s, logger)
 	return mcpserver
 }
 
-func addOpts(opts ...server.ServerOption) []server.ServerOption {
+func AddOpts(opts ...server.ServerOption) []server.ServerOption {
 	defaultOpts := []server.ServerOption{
 		server.WithResourceCapabilities(true, true),
 		server.WithRecovery(),
@@ -51,7 +51,7 @@ func addOpts(opts ...server.ServerOption) []server.ServerOption {
 	return opts
 }
 
-func createSNCloudMCPServer(s *server.MCPServer, logger *logrus.Logger) *Server {
+func CreateSNCloudMCPServer(s *server.MCPServer, logger *logrus.Logger) *Server {
 	mcpserver := &Server{
 		MCPServer:      s,
 		logger:         logger,