Refactor SNCloud session management and enhance MCP server functionality

- Introduced a new `Session` struct to encapsulate SNCloud context and client management, improving thread safety with mutex locks.
- Updated various functions to utilize the new session methods for obtaining SNCloud clients, enhancing code readability and maintainability.
- Refactored MCP server initialization to support SNCloud sessions, including error handling improvements.
- Removed deprecated context management code and replaced it with session-based approaches for better encapsulation.
- Added comments for better understanding of the code structure and functionality.

Author: freeznet

pkg/cmd/mcp/mcp.go

@@ -19,7 +19,6 @@ package mcp
 
 import (
 	"slices"
-	"time"
 
 	"github.com/pkg/errors"
 	"github.com/spf13/cobra"
@@ -73,20 +72,6 @@ func (o *ServerOptions) Complete() error {
 				return errors.Wrap(err, "Unable to store the authorization data")
 			}
 
-			err = config.InitSNCloudClient(
-				issuer.IssuerEndpoint, issuer.Audience, o.KeyFile, o.Options.Server, 30*time.Second, o.Options.Store)
-			if err != nil {
-				return errors.Wrap(err, "failed to initialize StreamNative Cloud client")
-			}
-
-			if o.Options.PulsarInstance != "" && o.Options.PulsarCluster != "" {
-				err = mcp.SetContext(o.Options, o.Options.PulsarInstance, o.Options.PulsarCluster)
-				if err != nil {
-					mcp.ResetMcpContext()
-					return errors.Wrap(err, "failed to set StreamNative Cloud context")
-				}
-			}
-
 			if len(o.Features) != 0 {
 				requiredFeatures := []mcp.Feature{
 					mcp.FeatureStreamNativeCloud,

pkg/cmd/mcp/server.go

@@ -25,6 +25,7 @@ import (
 	"github.com/mark3labs/mcp-go/server"
 	"github.com/pkg/errors"
 	"github.com/sirupsen/logrus"
+	"github.com/streamnative/streamnative-mcp-server/pkg/config"
 	"github.com/streamnative/streamnative-mcp-server/pkg/kafka"
 	"github.com/streamnative/streamnative-mcp-server/pkg/mcp"
 	"github.com/streamnative/streamnative-mcp-server/pkg/pulsar"
@@ -43,17 +44,16 @@ func newMcpServer(ctx context.Context, configOpts *ServerOptions, logrusLogger *
 				stdlog.Fatalf("failed to get user name: %v", err)
 				os.Exit(1)
 			}
-			// sncloudClient, err := config.GetAPIClient()
-			// if err != nil {
-			// 	stdlog.Fatalf("failed to get SNCloud client: %v", err)
-			// 	os.Exit(1)
-			// }
-			// sncloudLogClient, err := config.GetSNCloudLogClient()
-			// if err != nil {
-			// 	stdlog.Fatalf("failed to get SNCloud log client: %v", err)
-			// 	os.Exit(1)
-			// }
+			// Create StreamNative Cloud session and set as default
+			session, err := config.NewSNCloudSessionFromOptions(configOpts.Options)
+			if err != nil {
+				return nil, errors.Wrap(err, "failed to create StreamNative Cloud session")
+			}
 			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)
 			mcp.RegisterContextTools(s, configOpts.Features)

pkg/cmd/mcp/sse.go

@@ -75,6 +75,19 @@ func runSseServer(configOpts *ServerOptions) error {
 		return fmt.Errorf("failed to create MCP server: %w", err)
 	}
 
+	// 4. Set the context
+	ctx = mcp.WithSNCloudSession(ctx, mcpServer.SNCloudSession)
+	ctx = mcp.WithPulsarSession(ctx, mcpServer.PulsarSession)
+	ctx = mcp.WithKafkaSession(ctx, mcpServer.KafkaSession)
+	if configOpts.Options.KeyFile != "" {
+		if configOpts.Options.PulsarInstance != "" && configOpts.Options.PulsarCluster != "" {
+			err = mcp.SetContext(ctx, configOpts.Options, configOpts.Options.PulsarInstance, configOpts.Options.PulsarCluster)
+			if err != nil {
+				return errors.Wrap(err, "failed to set StreamNative Cloud context")
+			}
+		}
+	}
+
 	// 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, "")
@@ -86,6 +99,7 @@ func runSseServer(configOpts *ServerOptions) error {
 			c := context.WithValue(ctx, common.OptionsKey, configOpts.Options)
 			c = mcp.WithKafkaSession(c, mcpServer.KafkaSession)
 			c = mcp.WithPulsarSession(c, mcpServer.PulsarSession)
+			c = mcp.WithSNCloudSession(c, mcpServer.SNCloudSession)
 			return c
 		}),
 	)