Abstract session for SNCloud, Pulsar & Kafka Clients (#49)

Author: freeznet

.gitignore

@@ -9,3 +9,4 @@ tmp/
 # Go
 vendor
 .cursor/
+agents/

README.md

@@ -89,6 +89,10 @@ If you want to access to your StreamNative Cloud, you will need to have followin
 # Start MCP server with StreamNative Cloud authentication
 bin/snmcp stdio --organization my-org --key-file /path/to/key-file.json
 
+# Start MCP server with StreamNative Cloud authentication and pre-configured context
+# When --pulsar-instance and --pulsar-cluster are provided, context management tools are disabled
+bin/snmcp stdio --organization my-org --key-file /path/to/key-file.json --pulsar-instance my-instance --pulsar-cluster my-cluster
+
 # Start MCP server with external Kafka
 bin/snmcp stdio --use-external-kafka --kafka-bootstrap-servers localhost:9092 --kafka-auth-type SASL_SSL --kafka-auth-mechanism PLAIN --kafka-auth-user user --kafka-auth-pass pass --kafka-use-tls --kafka-schema-registry-url https://sr.local --kafka-schema-registry-auth-user user --kafka-schema-registry-auth-pass pass
 
@@ -106,6 +110,10 @@ docker run -i --rm -e SNMCP_ORGANIZATION=my-org -e SNMCP_KEY_FILE=/key.json -v /
 # Start MCP server with SSE and StreamNative Cloud authentication
 snmcp sse --http-addr :9090 --http-path /mcp --organization my-org --key-file /path/to/key-file.json
 
+# Start MCP server with SSE and pre-configured StreamNative Cloud context
+# When --pulsar-instance and --pulsar-cluster are provided, context management tools are disabled
+snmcp sse --http-addr :9090 --http-path /mcp --organization my-org --key-file /path/to/key-file.json --pulsar-instance my-instance --pulsar-cluster my-cluster
+
 # Start MCP server with SSE and external Kafka
 snmcp sse --http-addr :9090 --http-path /mcp --use-external-kafka --kafka-bootstrap-servers localhost:9092
 
@@ -235,6 +243,8 @@ The StreamNative MCP Server allows you to enable or disable specific groups of f
 | `streamnative-cloud`| Manage StreamNative Cloud context and check resource logs         | [streamnative_cloud.md](docs/tools/streamnative_cloud.md) |
 | `functions-as-tools`     | Dynamically exposes deployed Pulsar Functions as invokable MCP tools, with automatic input/output schema handling. | [functions_as_tools.md](docs/tools/functions_as_tools.md)  |
 
+> **Note:** When using `--pulsar-instance` and `--pulsar-cluster` flags together, context management tools (`sncloud_context_use_cluster`) are automatically disabled since the context is pre-configured.
+
 You can combine these features as needed using the `--features` flag. For example, to enable only Pulsar client features:
 ```bash
 # Enable only Pulsar client features

go.mod

@@ -1,14 +1,14 @@
 module github.com/streamnative/streamnative-mcp-server
 
-go 1.24.3
+go 1.24.4
 
 require (
 	github.com/99designs/keyring v1.2.2
 	github.com/apache/pulsar-client-go v0.13.1
 	github.com/dgrijalva/jwt-go v3.2.0+incompatible
 	github.com/google/go-cmp v0.7.0
 	github.com/hamba/avro/v2 v2.28.0
-	github.com/mark3labs/mcp-go v0.31.0
+	github.com/mark3labs/mcp-go v0.34.0
 	github.com/mitchellh/go-homedir v1.1.0
 	github.com/pkg/errors v0.9.1
 	github.com/sirupsen/logrus v1.9.3
@@ -80,6 +80,7 @@ require (
 	github.com/spf13/afero v1.12.0 // indirect
 	github.com/spf13/cast v1.7.1 // indirect
 	github.com/spf13/pflag v1.0.6 // indirect
+	github.com/stretchr/objx v0.5.2 // indirect
 	github.com/subosito/gotenv v1.6.0 // indirect
 	github.com/twmb/franz-go/pkg/kmsg v1.9.0 // indirect
 	github.com/yosida95/uritemplate/v3 v3.0.2 // indirect

go.sum

@@ -129,8 +129,8 @@ github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0 h1:6E+4a0GO5zZEnZ
 github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0/go.mod h1:zJYVVT2jmtg6P3p1VtQj7WsuWi/y4VnjVBn7F8KPB3I=
 github.com/magiconair/properties v1.8.7 h1:IeQXZAiQcpL9mgcAe1Nu6cX9LLw6ExEHKjN0VQdvPDY=
 github.com/magiconair/properties v1.8.7/go.mod h1:Dhd985XPs7jluiymwWYZ0G4Z61jb3vdS329zhj2hYo0=
-github.com/mark3labs/mcp-go v0.31.0 h1:4UxSV8aM770OPmTvaVe/b1rA2oZAjBMhGBfUgOGut+4=
-github.com/mark3labs/mcp-go v0.31.0/go.mod h1:rXqOudj/djTORU/ThxYx8fqEVj/5pvTuuebQ2RC7uk4=
+github.com/mark3labs/mcp-go v0.34.0 h1:eWy7WBGvhk6EyAAyVzivTCprE52iXJwNtvHV6Cv3bR0=
+github.com/mark3labs/mcp-go v0.34.0/go.mod h1:rXqOudj/djTORU/ThxYx8fqEVj/5pvTuuebQ2RC7uk4=
 github.com/mattn/go-colorable v0.1.2 h1:/bC9yWikZXAL9uJdulbSfyVNIR3n3trXl+v8+1sx8mU=
 github.com/mattn/go-colorable v0.1.2/go.mod h1:U0ppj6V5qS13XJ6of8GYAs25YV2eR4EVcfRqFIhoBtE=
 github.com/mattn/go-isatty v0.0.8 h1:HLtExJ+uU2HOZ+wI0Tt5DtUDrx8yhUqDcp7fYERX4CE=

go.work

@@ -1,4 +1,4 @@
-go 1.24.3
+go 1.24.4
 
 use .
 

go.work.sum

@@ -138,8 +138,8 @@ github.com/kylelemons/godebug v1.1.0/go.mod h1:9/0rRGxNHcop5bhtWyNeEfOS8JIWk580+
 github.com/lithammer/dedent v1.1.0/go.mod h1:jrXYCQtgg0nJiN+StA2KgR7w6CiQNv9Fd/Z9BP0jIOc=
 github.com/mark3labs/mcp-go v0.23.1 h1:RzTzZ5kJ+HxwnutKA4rll8N/pKV6Wh5dhCmiJUu5S9I=
 github.com/mark3labs/mcp-go v0.23.1/go.mod h1:rXqOudj/djTORU/ThxYx8fqEVj/5pvTuuebQ2RC7uk4=
-github.com/mark3labs/mcp-go v0.31.0 h1:4UxSV8aM770OPmTvaVe/b1rA2oZAjBMhGBfUgOGut+4=
-github.com/mark3labs/mcp-go v0.31.0/go.mod h1:rXqOudj/djTORU/ThxYx8fqEVj/5pvTuuebQ2RC7uk4=
+github.com/mark3labs/mcp-go v0.34.0 h1:eWy7WBGvhk6EyAAyVzivTCprE52iXJwNtvHV6Cv3bR0=
+github.com/mark3labs/mcp-go v0.34.0/go.mod h1:rXqOudj/djTORU/ThxYx8fqEVj/5pvTuuebQ2RC7uk4=
 github.com/mattn/go-runewidth v0.0.4 h1:2BvfKmzob6Bmd4YsL0zygOqfdFnK7GR4QL06Do4/p7Y=
 github.com/mattn/go-runewidth v0.0.4/go.mod h1:LwmH8dsx7+W8Uxz3IHJYH5QSwggIsqBzpuz5H//U1FU=
 github.com/matttproud/golang_protobuf_extensions v1.0.4 h1:mmDVorXM7PCGKw94cs5zkfA9PSy5pEvNWRP0ET0TIVo=

pkg/cmd/mcp/mcp.go

@@ -19,15 +19,12 @@ package mcp
 
 import (
 	"slices"
-	"time"
 
 	"github.com/pkg/errors"
 	"github.com/spf13/cobra"
 	"github.com/streamnative/streamnative-mcp-server/pkg/auth"
 	"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"
 )
 
 // ServerOptions is the options for the MCP server commands
@@ -75,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,
@@ -108,46 +91,13 @@ func (o *ServerOptions) Complete() error {
 				return errors.New("kafka-only mode does not support additional features")
 			}
 			o.Features = []string{string(mcp.FeatureKafkaClient), string(mcp.FeatureKafkaAdmin), string(mcp.FeatureKafkaAdminSchemaRegistry)}
-			err := kafka.NewCurrentKafkaContext(kafka.KafkaContext{
-				BootstrapServers:          snConfig.ExternalKafka.BootstrapServers,
-				AuthType:                  snConfig.ExternalKafka.AuthType,
-				AuthMechanism:             snConfig.ExternalKafka.AuthMechanism,
-				AuthUser:                  snConfig.ExternalKafka.AuthUser,
-				AuthPass:                  snConfig.ExternalKafka.AuthPass,
-				UseTLS:                    snConfig.ExternalKafka.UseTLS,
-				ClientKeyFile:             snConfig.ExternalKafka.ClientKeyFile,
-				ClientCertFile:            snConfig.ExternalKafka.ClientCertFile,
-				CaFile:                    snConfig.ExternalKafka.CaFile,
-				SchemaRegistryURL:         snConfig.ExternalKafka.SchemaRegistryURL,
-				SchemaRegistryAuthUser:    snConfig.ExternalKafka.SchemaRegistryAuthUser,
-				SchemaRegistryAuthPass:    snConfig.ExternalKafka.SchemaRegistryAuthPass,
-				SchemaRegistryBearerToken: snConfig.ExternalKafka.SchemaRegistryBearerToken,
-			})
-			if err != nil {
-				return errors.Wrap(err, "failed to set external Kafka context")
-			}
 		}
 	case snConfig.ExternalPulsar != nil:
 		{
 			if len(o.Features) != 0 {
 				return errors.New("pulsar-only mode does not support additional features")
 			}
 			o.Features = []string{string(mcp.FeatureAllPulsar)}
-			err := pulsar.NewCurrentPulsarContext(pulsar.PulsarContext{
-				ServiceURL:                    snConfig.ExternalPulsar.ServiceURL,
-				WebServiceURL:                 snConfig.ExternalPulsar.WebServiceURL,
-				AuthPlugin:                    snConfig.ExternalPulsar.AuthPlugin,
-				AuthParams:                    snConfig.ExternalPulsar.AuthParams,
-				Token:                         snConfig.ExternalPulsar.Token,
-				TLSAllowInsecureConnection:    snConfig.ExternalPulsar.TLSAllowInsecureConnection,
-				TLSEnableHostnameVerification: snConfig.ExternalPulsar.TLSEnableHostnameVerification,
-				TLSTrustCertsFilePath:         snConfig.ExternalPulsar.TLSTrustCertsFilePath,
-				TLSCertFile:                   snConfig.ExternalPulsar.TLSCertFile,
-				TLSKeyFile:                    snConfig.ExternalPulsar.TLSKeyFile,
-			}, nil, nil)
-			if err != nil {
-				return errors.Wrap(err, "failed to set external Pulsar context")
-			}
 		}
 	default:
 		{

pkg/cmd/mcp/server.go

@@ -18,17 +18,23 @@
 package mcp
 
 import (
+	"context"
 	stdlog "log"
 	"os"
 
 	"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"
 )
 
-func newMcpServer(configOpts *ServerOptions, logrusLogger *logrus.Logger) *server.MCPServer {
+func newMcpServer(_ context.Context, configOpts *ServerOptions, logrusLogger *logrus.Logger) (*mcp.Server, error) {
 	snConfig := configOpts.Options.LoadConfigOrDie()
 	var s *server.MCPServer
+	var mcpServer *mcp.Server
 	switch {
 	case snConfig.KeyFile != "":
 		{
@@ -38,36 +44,66 @@ func newMcpServer(configOpts *ServerOptions, logrusLogger *logrus.Logger) *serve
 				stdlog.Fatalf("failed to get user name: %v", err)
 				os.Exit(1)
 			}
-			// Create a new MCP server
-			s = server.NewMCPServer(
-				"streamnative-mcp-server",
-				"0.0.1",
-				server.WithResourceCapabilities(true, true),
-				server.WithInstructions(mcp.GetStreamNativeCloudServerInstructions(userName, snConfig)),
-				server.WithLogging())
+			// 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
 
+			s = mcpServer.MCPServer
 			mcp.RegisterPrompts(s)
-			mcp.RegisterContextTools(s, configOpts.Features)
+			// Skip context tools if pulsar instance and cluster are provided via CLI
+			skipContextTools := snConfig.Context.PulsarInstance != "" && snConfig.Context.PulsarCluster != ""
+			mcp.RegisterContextTools(s, configOpts.Features, skipContextTools)
 			mcp.StreamNativeAddLogTools(s, configOpts.ReadOnly, configOpts.Features)
 			mcp.StreamNativeAddResourceTools(s, configOpts.ReadOnly, configOpts.Features)
 		}
 	case snConfig.ExternalKafka != nil:
 		{
-			s = server.NewMCPServer(
-				"streamnative-mcp-server/kafka",
-				"0.0.1",
-				server.WithResourceCapabilities(true, true),
-				server.WithInstructions(mcp.GetExternalKafkaServerInstructions(snConfig.ExternalKafka.BootstrapServers)),
-				server.WithLogging())
+			ksession, err := kafka.NewSession(kafka.KafkaContext{
+				BootstrapServers:          snConfig.ExternalKafka.BootstrapServers,
+				AuthType:                  snConfig.ExternalKafka.AuthType,
+				AuthMechanism:             snConfig.ExternalKafka.AuthMechanism,
+				AuthUser:                  snConfig.ExternalKafka.AuthUser,
+				AuthPass:                  snConfig.ExternalKafka.AuthPass,
+				UseTLS:                    snConfig.ExternalKafka.UseTLS,
+				ClientKeyFile:             snConfig.ExternalKafka.ClientKeyFile,
+				ClientCertFile:            snConfig.ExternalKafka.ClientCertFile,
+				CaFile:                    snConfig.ExternalKafka.CaFile,
+				SchemaRegistryURL:         snConfig.ExternalKafka.SchemaRegistryURL,
+				SchemaRegistryAuthUser:    snConfig.ExternalKafka.SchemaRegistryAuthUser,
+				SchemaRegistryAuthPass:    snConfig.ExternalKafka.SchemaRegistryAuthPass,
+				SchemaRegistryBearerToken: snConfig.ExternalKafka.SchemaRegistryBearerToken,
+			})
+			if err != nil {
+				return nil, errors.Wrap(err, "failed to set external Kafka context")
+			}
+			mcpServer = mcp.NewServer("streamnative-mcp-server", "0.0.1", logrusLogger, server.WithInstructions(mcp.GetExternalKafkaServerInstructions(snConfig.ExternalKafka.BootstrapServers)))
+			mcpServer.KafkaSession = ksession
+			s = mcpServer.MCPServer
 		}
 	case snConfig.ExternalPulsar != nil:
 		{
-			s = server.NewMCPServer(
-				"streamnative-mcp-server/pulsar",
-				"0.0.1",
-				server.WithResourceCapabilities(true, true),
-				server.WithInstructions(mcp.GetExternalPulsarServerInstructions(snConfig.ExternalPulsar.WebServiceURL)),
-				server.WithLogging())
+			psession, err := pulsar.NewSession(pulsar.PulsarContext{
+				ServiceURL:                    snConfig.ExternalPulsar.ServiceURL,
+				WebServiceURL:                 snConfig.ExternalPulsar.WebServiceURL,
+				AuthPlugin:                    snConfig.ExternalPulsar.AuthPlugin,
+				AuthParams:                    snConfig.ExternalPulsar.AuthParams,
+				Token:                         snConfig.ExternalPulsar.Token,
+				TLSAllowInsecureConnection:    snConfig.ExternalPulsar.TLSAllowInsecureConnection,
+				TLSEnableHostnameVerification: snConfig.ExternalPulsar.TLSEnableHostnameVerification,
+				TLSTrustCertsFilePath:         snConfig.ExternalPulsar.TLSTrustCertsFilePath,
+				TLSCertFile:                   snConfig.ExternalPulsar.TLSCertFile,
+				TLSKeyFile:                    snConfig.ExternalPulsar.TLSKeyFile,
+			})
+			if err != nil {
+				return nil, errors.Wrap(err, "failed to set external Pulsar context")
+			}
+			mcpServer = mcp.NewServer("streamnative-mcp-server", "0.0.1", logrusLogger, server.WithInstructions(mcp.GetExternalPulsarServerInstructions(snConfig.ExternalPulsar.WebServiceURL)))
+			mcpServer.PulsarSession = psession
+			s = mcpServer.MCPServer
 		}
 	default:
 		{
@@ -103,5 +139,5 @@ func newMcpServer(configOpts *ServerOptions, logrusLogger *logrus.Logger) *serve
 	mcp.KafkaAdminAddKafkaConnectTools(s, configOpts.ReadOnly, configOpts.Features)
 	mcp.KafkaClientAddConsumeTools(s, configOpts.ReadOnly, logrusLogger, configOpts.Features)
 	mcp.KafkaClientAddProduceTools(s, configOpts.ReadOnly, configOpts.Features)
-	return s
+	return mcpServer, nil
 }

pkg/cmd/mcp/sse.go

@@ -33,6 +33,7 @@ import (
 	"github.com/spf13/cobra"
 	"github.com/streamnative/streamnative-mcp-server/pkg/common"
 	"github.com/streamnative/streamnative-mcp-server/pkg/mcp"
+	context2 "github.com/streamnative/streamnative-mcp-server/pkg/mcp"
 )
 
 func NewCmdMcpSseServer(configOpts *ServerOptions) *cobra.Command {
@@ -70,16 +71,37 @@ func runSseServer(configOpts *ServerOptions) error {
 
 	// 3. Create a new MCP server
 	ctx = context.WithValue(ctx, common.OptionsKey, configOpts.Options)
-	mcpServer := newMcpServer(configOpts, logger)
+	mcpServer, err := newMcpServer(ctx, configOpts, logger)
+	if err != nil {
+		return fmt.Errorf("failed to create MCP server: %w", err)
+	}
+
+	// 4. Set the context
+	ctx = context2.WithSNCloudSession(ctx, mcpServer.SNCloudSession)
+	ctx = context2.WithPulsarSession(ctx, mcpServer.PulsarSession)
+	ctx = context2.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
-	mcp.PulsarFunctionManagedMcpTools(mcpServer, false, 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.MCPServer,
 		server.WithStaticBasePath(configOpts.HTTPPath),
 		server.WithSSEContextFunc(func(ctx context.Context, _ *http.Request) context.Context {
-			return context.WithValue(ctx, common.OptionsKey, configOpts.Options)
+			c := context.WithValue(ctx, common.OptionsKey, configOpts.Options)
+			c = context2.WithKafkaSession(c, mcpServer.KafkaSession)
+			c = context2.WithPulsarSession(c, mcpServer.PulsarSession)
+			c = context2.WithSNCloudSession(c, mcpServer.SNCloudSession)
+			return c
 		}),
 	)
 

pkg/cmd/mcp/stdio.go

@@ -65,8 +65,12 @@ func runStdioServer(configOpts *ServerOptions) error {
 	// Create a new MCP server
 	ctx = context.WithValue(ctx, common.OptionsKey, configOpts.Options)
 	stdLogger := stdlog.New(logger.Writer(), "snmcp-server", 0)
-	stdioServer := server.NewStdioServer(newMcpServer(configOpts, logger))
+	mcpServer, err := newMcpServer(ctx, configOpts, logger)
+	if err != nil {
+		return fmt.Errorf("failed to create MCP server: %w", err)
+	}
 
+	stdioServer := server.NewStdioServer(mcpServer.MCPServer)
 	stdioServer.SetErrorLogger(stdLogger)
 
 	// Start listening for messages