Refactor session management and enhance MCP server functionality

- Updated `newMcpServer` function to accept a context parameter, improving the encapsulation of server initialization.
- Refactored Kafka and Pulsar session management to utilize session-based context handling, enhancing code readability and maintainability.
- Modified `PulsarFunctionManagedMcpTools` to use `sessionID` instead of `sessionId`, ensuring consistency across session management.
- Improved error handling and added comments for better understanding of the code structure and functionality.

Author: freeznet

pkg/cmd/mcp/server.go

@@ -31,7 +31,7 @@ import (
 	"github.com/streamnative/streamnative-mcp-server/pkg/pulsar"
 )
 
-func newMcpServer(ctx context.Context, configOpts *ServerOptions, logrusLogger *logrus.Logger) (*mcp.Server, error) {
+func newMcpServer(_ context.Context, configOpts *ServerOptions, logrusLogger *logrus.Logger) (*mcp.Server, error) {
 	snConfig := configOpts.Options.LoadConfigOrDie()
 	var s *server.MCPServer
 	var mcpServer *mcp.Server
@@ -51,8 +51,6 @@ func newMcpServer(ctx context.Context, configOpts *ServerOptions, logrusLogger *
 			}
 			mcpServer = mcp.NewServer("streamnative-mcp-server", "0.0.1", logrusLogger, server.WithInstructions(mcp.GetStreamNativeCloudServerInstructions(userName, snConfig)))
 			mcpServer.SNCloudSession = session
-			mcpServer.KafkaSession = &kafka.Session{}
-			mcpServer.PulsarSession = &pulsar.Session{}
 
 			s = mcpServer.MCPServer
 			mcp.RegisterPrompts(s)
@@ -97,7 +95,7 @@ func newMcpServer(ctx context.Context, configOpts *ServerOptions, logrusLogger *
 				TLSTrustCertsFilePath:         snConfig.ExternalPulsar.TLSTrustCertsFilePath,
 				TLSCertFile:                   snConfig.ExternalPulsar.TLSCertFile,
 				TLSKeyFile:                    snConfig.ExternalPulsar.TLSKeyFile,
-			}, nil, nil)
+			})
 			if err != nil {
 				return nil, errors.Wrap(err, "failed to set external Pulsar context")
 			}

pkg/kafka/connection.go

@@ -76,18 +76,13 @@ func NewSession(ctx KafkaContext) (*Session, error) {
 		Ctx: ctx,
 	}
 
-	if err := session.SetKafkaContext(); err != nil {
+	if err := session.SetKafkaContext(ctx); err != nil {
 		return nil, fmt.Errorf("failed to set kafka context: %w", err)
 	}
 
 	return session, nil
 }
 
-func (s *Session) ChangeContext(ctx KafkaContext) error {
-	s.Ctx = ctx
-	return s.SetKafkaContext()
-}
-
 type SASLConfig struct {
 	Mechanism string
 	Username  string
@@ -152,7 +147,8 @@ func saslOpt(config *SASLConfig, opts []kgo.Opt) ([]kgo.Opt, error) {
 	return opts, nil
 }
 
-func (s *Session) SetKafkaContext() error {
+func (s *Session) SetKafkaContext(ctx KafkaContext) error {
+	s.Ctx = ctx
 	kc := &s.Ctx
 	var err error
 	s.Options = []kgo.Opt{}

pkg/mcp/internal/pftools/manager.go

@@ -62,7 +62,7 @@ type Server struct {
 }
 
 // NewPulsarFunctionManager creates a new PulsarFunctionManager
-func NewPulsarFunctionManager(snServer *Server, readOnly bool, options *ManagerOptions, sessionId string) (*PulsarFunctionManager, error) {
+func NewPulsarFunctionManager(snServer *Server, readOnly bool, options *ManagerOptions, sessionID string) (*PulsarFunctionManager, error) {
 	// Get Pulsar client and admin client
 	if snServer.PulsarSession == nil {
 		return nil, fmt.Errorf("Pulsar session not found in context")
@@ -104,7 +104,7 @@ func NewPulsarFunctionManager(snServer *Server, readOnly bool, options *ManagerO
 		circuitBreakers:   make(map[string]*CircuitBreaker),
 		tenantNamespaces:  options.TenantNamespaces,
 		strictExport:      options.StrictExport,
-		sessionId:         sessionId,
+		sessionID:         sessionID,
 	}
 
 	return manager, nil
@@ -198,19 +198,19 @@ func (m *PulsarFunctionManager) updateFunctions() {
 		}
 
 		if changed {
-			if m.sessionId != "" {
-				err := m.mcpServer.DeleteSessionTools(m.sessionId, fnTool.Tool.Name)
+			if m.sessionID != "" {
+				err := m.mcpServer.DeleteSessionTools(m.sessionID, fnTool.Tool.Name)
 				if err != nil {
-					log.Printf("Failed to delete tool %s from session %s: %v", fnTool.Tool.Name, m.sessionId, err)
+					log.Printf("Failed to delete tool %s from session %s: %v", fnTool.Tool.Name, m.sessionID, err)
 				}
 			} else {
 				m.mcpServer.DeleteTools(fnTool.Tool.Name)
 			}
 		}
-		if m.sessionId != "" {
-			err := m.mcpServer.AddSessionTool(m.sessionId, fnTool.Tool, m.handleToolCall(fnTool))
+		if m.sessionID != "" {
+			err := m.mcpServer.AddSessionTool(m.sessionID, fnTool.Tool, m.handleToolCall(fnTool))
 			if err != nil {
-				log.Printf("Failed to add tool %s to session %s: %v", fnTool.Tool.Name, m.sessionId, err)
+				log.Printf("Failed to add tool %s to session %s: %v", fnTool.Tool.Name, m.sessionID, err)
 			}
 		} else {
 			m.mcpServer.AddTool(fnTool.Tool, m.handleToolCall(fnTool))
@@ -232,10 +232,10 @@ func (m *PulsarFunctionManager) updateFunctions() {
 	m.mutex.Lock()
 	for fullName, fnTool := range m.fnToToolMap {
 		if !seenFunctions[fullName] {
-			if m.sessionId != "" {
-				err := m.mcpServer.DeleteSessionTools(m.sessionId, fnTool.Tool.Name)
+			if m.sessionID != "" {
+				err := m.mcpServer.DeleteSessionTools(m.sessionID, fnTool.Tool.Name)
 				if err != nil {
-					log.Printf("Failed to delete tool %s from session %s: %v", fnTool.Tool.Name, m.sessionId, err)
+					log.Printf("Failed to delete tool %s from session %s: %v", fnTool.Tool.Name, m.sessionID, err)
 				}
 			} else {
 				m.mcpServer.DeleteTools(fnTool.Tool.Name)

pkg/mcp/internal/pftools/types.go

@@ -47,7 +47,7 @@ type PulsarFunctionManager struct {
 	circuitBreakers   map[string]*CircuitBreaker
 	tenantNamespaces  []string
 	strictExport      bool
-	sessionId         string
+	sessionID         string
 }
 
 type FunctionTool struct {

pkg/mcp/pulsar_admin_topic_policy_tools.go

@@ -27,7 +27,6 @@ import (
 	"github.com/apache/pulsar-client-go/pulsaradmin/pkg/utils"
 	"github.com/mark3labs/mcp-go/mcp"
 	"github.com/mark3labs/mcp-go/server"
-	
 )
 
 // PulsarAdminAddTopicPolicyTools adds topic policy-related tools to the MCP server

pkg/mcp/pulsar_functions_as_tools.go

@@ -51,7 +51,7 @@ func StopAllPulsarFunctionManagers() {
 	log.Println("All Pulsar Function managers stopped")
 }
 
-func (s *Server) PulsarFunctionManagedMcpTools(readOnly bool, features []string, sessionId string) {
+func (s *Server) PulsarFunctionManagedMcpTools(readOnly bool, features []string, sessionID string) {
 	if !slices.Contains(features, string(FeatureAll)) &&
 		!slices.Contains(features, string(FeatureFunctionsAsTools)) &&
 		!slices.Contains(features, string(FeatureStreamNativeCloud)) {
@@ -60,6 +60,11 @@ func (s *Server) PulsarFunctionManagedMcpTools(readOnly bool, features []string,
 
 	options := pftools.DefaultManagerOptions()
 
+	if s.SNCloudSession.Ctx.Organization == "" || s.SNCloudSession.Ctx.PulsarInstance == "" || s.SNCloudSession.Ctx.PulsarCluster == "" {
+		log.Printf("Skipping Pulsar Functions as MCP Tools because both organization, pulsar instance and pulsar cluster are not set")
+		return
+	}
+
 	if pollIntervalStr := os.Getenv("FUNCTIONS_AS_TOOLS_POLL_INTERVAL"); pollIntervalStr != "" {
 		if seconds, err := strconv.Atoi(pollIntervalStr); err == nil && seconds > 0 {
 			options.PollInterval = time.Duration(seconds) * time.Second
@@ -106,7 +111,7 @@ func (s *Server) PulsarFunctionManagedMcpTools(readOnly bool, features []string,
 		Logger:        s.logger,
 	}
 
-	manager, err := pftools.NewPulsarFunctionManager(pftoolsServer, readOnly, options, sessionId)
+	manager, err := pftools.NewPulsarFunctionManager(pftoolsServer, readOnly, options, sessionID)
 	if err != nil {
 		log.Printf("Failed to create Pulsar Function manager: %v", err)
 		return

pkg/mcp/server.go

@@ -53,8 +53,11 @@ func addOpts(opts ...server.ServerOption) []server.ServerOption {
 
 func createSNCloudMCPServer(s *server.MCPServer, logger *logrus.Logger) *Server {
 	mcpserver := &Server{
-		MCPServer: s,
-		logger:    logger,
+		MCPServer:      s,
+		logger:         logger,
+		SNCloudSession: &config.Session{},
+		KafkaSession:   &kafka.Session{},
+		PulsarSession:  &pulsar.Session{},
 	}
 
 	return mcpserver

pkg/mcp/sncontext_utils.go

@@ -159,11 +159,11 @@ func SetContext(ctx context.Context, options *config.Options, instanceName, clus
 	if psession == nil {
 		return fmt.Errorf("failed to get pulsar session")
 	}
-	err = psession.ChangeContext(pulsar.PulsarContext{
+	err = psession.SetPulsarContext(pulsar.PulsarContext{
 		WebServiceURL: getBasePath(snConfig.ProxyLocation, options.Organization, clusterUID),
 		ServiceURL:    getServiceURL(dnsName),
 		Token:         accessToken,
-	}, issuer, &options.AuthOptions.Store)
+	})
 	if err != nil {
 		return fmt.Errorf("failed to change pulsar context: %v", err)
 	}
@@ -187,10 +187,14 @@ func SetContext(ctx context.Context, options *config.Options, instanceName, clus
 	if ksession == nil {
 		return fmt.Errorf("failed to get kafka session")
 	}
-	err = ksession.ChangeContext(kctx)
+	err = ksession.SetKafkaContext(kctx)
 	if err != nil {
 		return fmt.Errorf("failed to change kafka context: %v", err)
 	}
 
+	// TODO: check if need to set log client
+	// if issuer != nil && options.AuthOptions.Store != nil {
+	// }
+
 	return nil
 }

pkg/pulsar/connection.go

@@ -25,8 +25,6 @@ import (
 	"github.com/apache/pulsar-client-go/pulsar"
 	pulsaradminconfig "github.com/apache/pulsar-client-go/pulsaradmin/pkg/admin/config"
 	"github.com/streamnative/pulsarctl/pkg/cmdutils"
-	"github.com/streamnative/streamnative-mcp-server/pkg/auth"
-	"github.com/streamnative/streamnative-mcp-server/pkg/auth/store"
 )
 
 const (
@@ -49,51 +47,39 @@ type PulsarContext struct {
 
 // Session represents a Pulsar session
 type Session struct {
-	Ctx           PulsarContext
-	Client        pulsar.Client
-	AdminClient   cmdutils.Client
-	AdminV3Client cmdutils.Client
-	ClientOptions pulsar.ClientOptions
-	mutex         sync.RWMutex
-}
-
-func init() {
-	cmdutils.PulsarCtlConfig = &cmdutils.ClusterConfig{}
+	Ctx             PulsarContext
+	Client          pulsar.Client
+	AdminClient     cmdutils.Client
+	AdminV3Client   cmdutils.Client
+	ClientOptions   pulsar.ClientOptions
+	PulsarCtlConfig *cmdutils.ClusterConfig
+	mutex           sync.RWMutex
 }
 
 // NewSession creates a new Pulsar session with the given context
 // This function dynamically constructs clients without relying on global state
-func NewSession(ctx PulsarContext, issuer *auth.Issuer, tokenStore *store.Store) (*Session, error) {
+func NewSession(ctx PulsarContext) (*Session, error) {
 	session := &Session{
 		Ctx: ctx,
 	}
 
-	if err := session.SetPulsarContext(); err != nil {
+	if err := session.SetPulsarContext(ctx); err != nil {
 		return nil, fmt.Errorf("failed to set pulsar context: %w", err)
 	}
 
 	return session, nil
 }
 
-func (s *Session) ChangeContext(ctx PulsarContext, issuer *auth.Issuer, tokenStore *store.Store) error {
+func (s *Session) SetPulsarContext(ctx PulsarContext) error {
 	s.mutex.Lock()
 	defer s.mutex.Unlock()
-
 	s.Ctx = ctx
-
-	return s.SetPulsarContext()
-}
-
-func (s *Session) SetPulsarContext() error {
-	s.mutex.Lock()
-	defer s.mutex.Unlock()
-
 	pc := &s.Ctx
 	var err error
 	// Configure pulsarctl with the token
 	switch {
 	case pc.Token != "":
-		cmdutils.PulsarCtlConfig = &cmdutils.ClusterConfig{
+		s.PulsarCtlConfig = &cmdutils.ClusterConfig{
 			WebServiceURL:                 pc.WebServiceURL,
 			AuthPlugin:                    "org.apache.pulsar.client.impl.auth.AuthenticationToken",
 			AuthParams:                    fmt.Sprintf("token:%s", pc.Token),
@@ -117,7 +103,7 @@ func (s *Session) SetPulsarContext() error {
 			TLSKeyFilePath:             pc.TLSKeyFile,
 		}
 	case pc.AuthPlugin != "" && pc.AuthParams != "":
-		cmdutils.PulsarCtlConfig = &cmdutils.ClusterConfig{
+		s.PulsarCtlConfig = &cmdutils.ClusterConfig{
 			WebServiceURL:                 pc.WebServiceURL,
 			AuthPlugin:                    pc.AuthPlugin,
 			AuthParams:                    pc.AuthParams,
@@ -145,7 +131,7 @@ func (s *Session) SetPulsarContext() error {
 		}
 	default:
 		// No authentication provided
-		cmdutils.PulsarCtlConfig = &cmdutils.ClusterConfig{
+		s.PulsarCtlConfig = &cmdutils.ClusterConfig{
 			WebServiceURL:                 pc.WebServiceURL,
 			TLSAllowInsecureConnection:    pc.TLSAllowInsecureConnection,
 			TLSEnableHostnameVerification: pc.TLSEnableHostnameVerification,
@@ -167,8 +153,8 @@ func (s *Session) SetPulsarContext() error {
 		}
 	}
 
-	s.AdminClient = cmdutils.NewPulsarClient()
-	s.AdminV3Client = cmdutils.NewPulsarClientWithAPIVersion(pulsaradminconfig.V3)
+	s.AdminClient = s.PulsarCtlConfig.Client(pulsaradminconfig.V2)
+	s.AdminV3Client = s.PulsarCtlConfig.Client(pulsaradminconfig.V3)
 
 	s.Client, err = pulsar.NewClient(s.ClientOptions)
 	if err != nil {
@@ -182,7 +168,7 @@ func (s *Session) GetAdminClient() (cmdutils.Client, error) {
 	s.mutex.RLock()
 	defer s.mutex.RUnlock()
 
-	if cmdutils.PulsarCtlConfig.WebServiceURL == "" {
+	if s.PulsarCtlConfig.WebServiceURL == "" {
 		return nil, fmt.Errorf("err: ContextNotSetErr: Please set the cluster context first")
 	}
 	return s.AdminClient, nil
@@ -192,7 +178,7 @@ func (s *Session) GetAdminV3Client() (cmdutils.Client, error) {
 	s.mutex.RLock()
 	defer s.mutex.RUnlock()
 
-	if cmdutils.PulsarCtlConfig.WebServiceURL == "" {
+	if s.PulsarCtlConfig.WebServiceURL == "" {
 		return nil, fmt.Errorf("err: ContextNotSetErr: Please set the cluster context first")
 	}
 	return s.AdminV3Client, nil