add logging tools

freeznet · freeznet · commit 3c629d121d6c · 2025-05-02T16:27:58.000+08:00
diff --git a/go.mod b/go.mod
@@ -6,7 +6,7 @@ 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/mark3labs/mcp-go v0.23.1
+	github.com/mark3labs/mcp-go v0.25.0
 	github.com/mitchellh/go-homedir v1.1.0
 	github.com/pkg/errors v0.9.1
 	github.com/sirupsen/logrus v1.9.3
@@ -37,6 +37,7 @@ require (
 	github.com/dvsekhvalnov/jose2go v1.6.0 // indirect
 	github.com/fatih/color v1.7.0 // indirect
 	github.com/ghodss/yaml v1.0.0 // indirect
+	github.com/gmlewis/go-httpdebug v0.0.9 // indirect
 	github.com/godbus/dbus v0.0.0-20190726142602-4481cbc300e2 // indirect
 	github.com/golang-jwt/jwt/v4 v4.5.1 // indirect
 	github.com/golang-jwt/jwt/v5 v5.2.1 // indirect
diff --git a/go.sum b/go.sum
@@ -69,6 +69,8 @@ github.com/fsnotify/fsnotify v1.6.0 h1:n+5WquG0fcWoWp6xPWfHdbskMCQaFnG6PfBrh1Ky4
 github.com/fsnotify/fsnotify v1.6.0/go.mod h1:sl3t1tCWJFWoRz9R8WJCbQihKKwmorjAbSClcnxKAGw=
 github.com/ghodss/yaml v1.0.0 h1:wQHKEahhL6wmXdzwWG11gIVCkOv05bNOh+Rxn0yngAk=
 github.com/ghodss/yaml v1.0.0/go.mod h1:4dBDuWmgqj2HViK6kFavaiC9ZROes6MMH2rRYeMEF04=
+github.com/gmlewis/go-httpdebug v0.0.9 h1:zNweu7aEyDOL2bDt5tpwAMMNifai9YtAh479bZEML6g=
+github.com/gmlewis/go-httpdebug v0.0.9/go.mod h1:t6xDKFIjenEc4rouOt79Oo6N09EiHFgMjbCOiQnOnOc=
 github.com/go-logr/logr v1.4.2 h1:6pFjapn8bFcIbiKo3XT4j/BhANplGihG6tvd+8rYgrY=
 github.com/go-logr/logr v1.4.2/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY=
 github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag=
@@ -129,8 +131,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.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.25.0 h1:UUpcMT3L5hIhuDy7aifj4Bphw4Pfx1Rf8mzMXDe8RQw=
+github.com/mark3labs/mcp-go v0.25.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=
diff --git a/go.work.sum b/go.work.sum
@@ -135,6 +135,7 @@ golang.org/x/image v0.0.0-20190802002840-cff245a6509b h1:+qEpEAPhDZ1o0x3tHzZTQDA
 golang.org/x/lint v0.0.0-20200302205851-738671d3881b h1:Wh+f8QHJXR411sJR8/vRBTZ7YapZaRvUcLFFJhusH0k=
 golang.org/x/mobile v0.0.0-20190719004257-d2bd2a29d028 h1:4+4C/Iv2U4fMZBiMCc98MG1In4gJY5YRhtpDNeDeHWs=
 golang.org/x/mod v0.3.0 h1:RM4zey1++hCTbCVQfnWeKs9/IEsaBLA8vTkd0WVtmH4=
+golang.org/x/net v0.38.0/go.mod h1:ivrbrMbzFq5J41QOQh0siUuly180yBYtLp+CKbEaFx8=
 golang.org/x/sync v0.0.0-20200625203802-6e8e738ad208 h1:qwRHBd0NqMbJxfbotnDhm2ByMI1Shq4Y6oRJo21SGJA=
 golang.org/x/sync v0.8.0 h1:3NFvSEYkUoMifnESzZl15y791HH1qU2xm6eCJU5ZPXQ=
 golang.org/x/sync v0.8.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
diff --git a/pkg/cmd/mcp/mcp.go b/pkg/cmd/mcp/mcp.go
@@ -82,6 +82,7 @@ func (o *ServerOptions) Complete() error {
 			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")
 				}
 			}
@@ -141,7 +142,7 @@ func (o *ServerOptions) Complete() error {
 				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")
 			}
diff --git a/pkg/cmd/mcp/stdio.go b/pkg/cmd/mcp/stdio.go
@@ -143,6 +143,7 @@ func newStdioServer(configOpts *ServerOptions, logrusLogger *logrus.Logger) *ser
 
 			mcp.RegisterPrompts(s)
 			mcp.RegisterContextTools(s, configOpts.Features)
+			mcp.StreamNativeAddLogTools(s, configOpts.ReadOnly, configOpts.Features)
 		}
 	case snConfig.ExternalKafka != nil:
 		{
diff --git a/pkg/config/apiclient.go b/pkg/config/apiclient.go
@@ -27,6 +27,7 @@ import (
 	"golang.org/x/oauth2"
 	"k8s.io/utils/clock"
 
+	dbg "github.com/gmlewis/go-httpdebug/httpdebug"
 	"github.com/streamnative/streamnative-mcp-server/pkg/auth"
 	"github.com/streamnative/streamnative-mcp-server/pkg/auth/cache"
 	"github.com/streamnative/streamnative-mcp-server/pkg/auth/store"
@@ -35,6 +36,7 @@ import (
 
 var SNCloudClientConfiguration *sncloud.Configuration
 var SNCloudClient *sncloud.APIClient
+var SNCloudLogClient *http.Client
 
 // OAuth2TokenRefresher implements oauth2.TokenSource interface for refreshing OAuth2 tokens
 // This is now a wrapper around the cache.CachingTokenSource to leverage the existing token caching
@@ -183,3 +185,40 @@ func TokenRefreshed(audience string, token *oauth2.Token, tokenStore store.Store
 	// Save back to store
 	return tokenStore.SaveGrant(audience, *grant)
 }
+
+func InitSNCloudLogClient(issuerData auth.Issuer, tokenStore store.Store) error {
+	refresher, err := auth.NewDefaultClientCredentialsGrantRefresher(issuerData, clock.RealClock{})
+	if err != nil {
+		return errors.Wrap(err, "failed to create token refresher")
+	}
+
+	tokenRefresher, err := NewOAuth2TokenRefresher(tokenStore, issuerData.Audience, refresher)
+	if err != nil {
+		return errors.Wrap(err, "failed to create token refresher")
+	}
+
+	tokenSource := oauth2.ReuseTokenSource(nil, tokenRefresher)
+	b := dbg.New()
+	b.RedactEntireJWT = false
+	SNCloudLogClient = &http.Client{
+		Timeout: 10 * time.Second,
+		Transport: &oauth2.Transport{
+			Source: tokenSource,
+			// Base:   http.DefaultTransport,
+			Base: b,
+		},
+	}
+
+	return nil
+}
+
+func ResetSNCloudLogClient() {
+	SNCloudLogClient = nil
+}
+
+func GetSNCloudLogClient() (*http.Client, error) {
+	if SNCloudLogClient == nil {
+		return nil, errors.New("log tools are for StreamNative Cloud only, please check your context")
+	}
+	return SNCloudLogClient, nil
+}
diff --git a/pkg/config/config.go b/pkg/config/config.go
@@ -36,6 +36,7 @@ type SnConfig struct {
 	// settable context
 	Context       Context `yaml:"context"`
 	ProxyLocation string  `yaml:"proxy-location"`
+	LogLocation   string  `yaml:"log-location"`
 	KeyFile       string  `yaml:"key-file"`
 
 	ExternalKafka  *ExternalKafka  `yaml:"external-kafka"`
diff --git a/pkg/config/options.go b/pkg/config/options.go
@@ -34,6 +34,7 @@ const (
 	GlobalDefaultAudience      = "https://api.streamnative.cloud"
 	GlobalDefaultAPIServer     = "https://api.streamnative.cloud"
 	GlobalDefaultProxyLocation = "https://proxy.streamnative.cloud"
+	GlobalDefaultLogLocation   = "https://log.streamnative.cloud"
 )
 
 // Options represents the common options used throughout the program.
@@ -52,6 +53,7 @@ type Options struct {
 	PulsarInstance string
 	PulsarCluster  string
 	ProxyLocation  string
+	LogLocation    string
 	KeyFile        string
 
 	UseExternalKafka  bool
@@ -85,6 +87,8 @@ func (o *Options) AddFlags(cmd *cobra.Command) {
 		"The organization to use for the API server")
 	cmd.PersistentFlags().StringVar(&o.ProxyLocation, "proxy-location", GlobalDefaultProxyLocation,
 		"The proxy location to use for the API server")
+	cmd.PersistentFlags().StringVar(&o.LogLocation, "log-location", GlobalDefaultLogLocation,
+		"The log location to use for the API server")
 	_ = cmd.MarkFlagRequired("organization")
 	cmd.PersistentFlags().StringVar(&o.PulsarInstance, "pulsar-instance", "",
 		"The default instance to use for the API server")
@@ -192,6 +196,7 @@ func (o *Options) LoadConfig() (*SnConfig, error) {
 			PulsarCluster:  o.PulsarCluster,
 		},
 		ProxyLocation: o.ProxyLocation,
+		LogLocation:   o.LogLocation,
 		KeyFile:       o.KeyFile,
 	}
 	if o.UseExternalKafka {
diff --git a/pkg/kafka/connection.go b/pkg/kafka/connection.go
@@ -65,6 +65,14 @@ func NewCurrentKafkaContext(kc KafkaContext) error {
 	return kc.SetKafkaContext()
 }
 
+func ResetCurrentKafkaContext() {
+	CurrentKafkaContext = KafkaContext{}
+	KafkaAdminClient = nil
+	KafkaClient = nil
+	KafkaSchemaRegistryClient = nil
+	KafkaConnectClient = nil
+}
+
 type SASLConfig struct {
 	Mechanism string
 	Username  string
diff --git a/pkg/mcp/context.go b/pkg/mcp/context.go
@@ -0,0 +1,38 @@
+// Licensed to Elasticsearch B.V. under one or more contributor
+// license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright
+// ownership. Elasticsearch B.V. licenses this file to you under
+// the Apache License, Version 2.0 (the "License"); you may
+// not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package mcp
+
+var McpContextPulsarInstance *string
+var McpContextPulsarCluster *string
+var McpContextOrganization *string
+
+func SetMcpContext(instance, cluster, organization string) {
+	McpContextPulsarInstance = &instance
+	McpContextPulsarCluster = &cluster
+	McpContextOrganization = &organization
+}
+
+func GetMcpContext() (string, string, string) {
+	return *McpContextPulsarInstance, *McpContextPulsarCluster, *McpContextOrganization
+}
+
+func ResetMcpContext() {
+	McpContextPulsarInstance = nil
+	McpContextPulsarCluster = nil
+	McpContextOrganization = nil
+}
diff --git a/pkg/mcp/context_tools.go b/pkg/mcp/context_tools.go
@@ -104,6 +104,7 @@ func handleSetContext(ctx context.Context, request mcp.CallToolRequest) (*mcp.Ca
 
 	err = SetContext(options, instanceName, clusterName)
 	if err != nil {
+		ResetMcpContext()
 		return mcp.NewToolResultError(fmt.Sprintf("Failed to set context: %v", err)), nil
 	}
 
diff --git a/pkg/mcp/context_utils.go b/pkg/mcp/context_utils.go
@@ -102,7 +102,8 @@ func SetContext(options *config.Options, instanceName, clusterName string) error
 		return fmt.Errorf("failed to get issuer: %v", err)
 	}
 
-	tokenKey := buildTokenKey(options.Organization, clusterUID, myselfGrant.ClientCredentials.ClientEmail)
+	// tokenKey := buildTokenKey(options.Organization, clusterUID, myselfGrant.ClientCredentials.ClientEmail)
+	tokenKey := issuer.Audience
 
 	accessToken := ""
 	refreshToken := true
@@ -152,7 +153,7 @@ func SetContext(options *config.Options, instanceName, clusterName string) error
 		WebServiceURL: getBasePath(snConfig.ProxyLocation, options.Organization, clusterUID),
 		ServiceURL:    getServiceURL(dnsName),
 		Token:         accessToken,
-	})
+	}, issuer, &options.AuthOptions.Store)
 	if err != nil {
 		return fmt.Errorf("failed to set pulsar context: %v", err)
 	}
@@ -175,5 +176,13 @@ func SetContext(options *config.Options, instanceName, clusterName string) error
 		return fmt.Errorf("failed to set kafka context: %v", err)
 	}
 
+	SetMcpContext(instanceName, clusterName, options.Organization)
+
 	return nil
 }
+
+func ResetContext() {
+	pulsar.ResetCurrentPulsarContext()
+	kafka.ResetCurrentKafkaContext()
+	ResetMcpContext()
+}
diff --git a/pkg/mcp/streamnative_resources_log_tools.go b/pkg/mcp/streamnative_resources_log_tools.go
diff --git a/pkg/pulsar/connection.go b/pkg/pulsar/connection.go

Original file line number	Diff line number	Diff line change
`@@ -82,6 +82,7 @@ func (o *ServerOptions) Complete() error {`
`82`	`82`	`if o.Options.PulsarInstance != "" && o.Options.PulsarCluster != "" {`
`83`	`83`	`err = mcp.SetContext(o.Options, o.Options.PulsarInstance, o.Options.PulsarCluster)`
`84`	`84`	`if err != nil {`
	`85`	`+ mcp.ResetMcpContext()`
`85`	`86`	`return errors.Wrap(err, "failed to set StreamNative Cloud context")`
`86`	`87`	`}`
`87`	`88`	`}`
`@@ -141,7 +142,7 @@ func (o *ServerOptions) Complete() error {`
`141`	`142`	`TLSTrustCertsFilePath: snConfig.ExternalPulsar.TLSTrustCertsFilePath,`
`142`	`143`	`TLSCertFile: snConfig.ExternalPulsar.TLSCertFile,`
`143`	`144`	`TLSKeyFile: snConfig.ExternalPulsar.TLSKeyFile,`
`144`		`- })`
	`145`	`+ }, nil, nil)`
`145`	`146`	`if err != nil {`
`146`	`147`	`return errors.Wrap(err, "failed to set external Pulsar context")`
`147`	`148`	`}`
Original file line number	Diff line number	Diff line change
`@@ -143,6 +143,7 @@ func newStdioServer(configOpts ServerOptions, logrusLogger logrus.Logger) *ser`
`143`	`143`
`144`	`144`	`mcp.RegisterPrompts(s)`
`145`	`145`	`mcp.RegisterContextTools(s, configOpts.Features)`
	`146`	`+ mcp.StreamNativeAddLogTools(s, configOpts.ReadOnly, configOpts.Features)`
`146`	`147`	`}`
`147`	`148`	`case snConfig.ExternalKafka != nil:`
`148`	`149`	`{`