Skip to content

Commit 101ad93

Browse files
weibaohuiweibaohui
committed
refactor(mcp): 重构工具注册和元数据处理逻辑
将工具注册函数统一增加`ServerConfig`参数，以便在注册时传递配置信息。同时将`metadata`包中的`ResourceInfo`和`ResourceMetadata`类型整合，并移除重复的`resourceMap`定义。此外，将`ParseFromRequest`函数移至`metadata`包，并优化其逻辑以支持上下文传递。
1 parent 2654ee0 commit 101ad93

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

53 files changed

+413
-256
lines changed

‎main.go

+26-1
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,14 @@
11
package main
22

33
import (
4+
"context"
45
"flag"
6+
"net/http"
57

8+
"github.com/mark3labs/mcp-go/server"
69
"github.com/weibaohui/kom/example"
710
"github.com/weibaohui/kom/mcp"
11+
"github.com/weibaohui/kom/mcp/metadata"
812
"k8s.io/klog/v2"
913
)
1014

@@ -13,6 +17,27 @@ func main() {
1317
flag.Set("v", "8")
1418
example.Connect()
1519
// example.Example()
16-
mcp.RunMCPServer("kom mcp server", "0.0.1", 9096)
20+
// mcp.RunMCPServer("kom mcp server", "0.0.1", 9096)
1721

22+
authKey := "Authorization"
23+
var ctxFn = func(ctx context.Context, r *http.Request) context.Context {
24+
authVal := r.Header.Get(authKey)
25+
klog.Infof("%s: %s", authKey, authVal)
26+
return context.WithValue(ctx, authKey, authVal)
27+
}
28+
cfg := metadata.ServerConfig{
29+
Name: "kom mcp server",
30+
Version: "0.0.1",
31+
Port: 9096,
32+
ServerOptions: []server.ServerOption{
33+
server.WithResourceCapabilities(false, false),
34+
server.WithPromptCapabilities(false),
35+
server.WithLogging(),
36+
},
37+
SSEOption: []server.SSEOption{
38+
server.WithSSEContextFunc(ctxFn),
39+
},
40+
AuthKey: authKey,
41+
}
42+
mcp.RunMCPServerWithOption(&cfg)
1843
}

‎mcp/metadata/type.go

+22-49
Original file line numberDiff line numberDiff line change
@@ -1,58 +1,31 @@
11
package metadata
22

3-
import "strings"
3+
import (
4+
"github.com/mark3labs/mcp-go/server"
5+
)
6+
7+
// ResourceMetadata 封装资源的元数据信息
8+
type ResourceMetadata struct {
9+
Cluster string
10+
Namespace string
11+
Name string
12+
Group string
13+
Version string
14+
Kind string
15+
}
16+
type ServerConfig struct {
17+
Name string
18+
Version string
19+
Port int
20+
ServerOptions []server.ServerOption
21+
SSEOption []server.SSEOption
22+
Metadata map[string]string // 元数据
23+
AuthKey string // 认证key
24+
}
425

526
type ResourceInfo struct {
627
Group string
728
Version string
829
Kind string
930
Namespaced bool
1031
}
11-
12-
var resourceMap = map[string]ResourceInfo{
13-
// 命名空间级别资源
14-
"pod": {Group: "", Version: "v1", Kind: "Pod", Namespaced: true},
15-
"deployment": {Group: "apps", Version: "v1", Kind: "Deployment", Namespaced: true},
16-
"statefulset": {Group: "apps", Version: "v1", Kind: "StatefulSet", Namespaced: true},
17-
"daemonset": {Group: "apps", Version: "v1", Kind: "DaemonSet", Namespaced: true},
18-
"replicaset": {Group: "apps", Version: "v1", Kind: "ReplicaSet", Namespaced: true},
19-
"service": {Group: "", Version: "v1", Kind: "Service", Namespaced: true},
20-
"configmap": {Group: "", Version: "v1", Kind: "ConfigMap", Namespaced: true},
21-
"secret": {Group: "", Version: "v1", Kind: "Secret", Namespaced: true},
22-
"ingress": {Group: "networking.k8s.io", Version: "v1", Kind: "Ingress", Namespaced: true},
23-
"networkpolicy": {Group: "networking.k8s.io", Version: "v1", Kind: "NetworkPolicy", Namespaced: true},
24-
"role": {Group: "rbac.authorization.k8s.io", Version: "v1", Kind: "Role", Namespaced: true},
25-
"rolebinding": {Group: "rbac.authorization.k8s.io", Version: "v1", Kind: "RoleBinding", Namespaced: true},
26-
"serviceaccount": {Group: "", Version: "v1", Kind: "ServiceAccount", Namespaced: true},
27-
"persistentvolumeclaim": {Group: "", Version: "v1", Kind: "PersistentVolumeClaim", Namespaced: true},
28-
"horizontalpodautoscaler": {Group: "autoscaling", Version: "v2", Kind: "HorizontalPodAutoscaler", Namespaced: true},
29-
"cronjob": {Group: "batch", Version: "v1", Kind: "CronJob", Namespaced: true},
30-
"job": {Group: "batch", Version: "v1", Kind: "Job", Namespaced: true},
31-
"node": {Group: "", Version: "v1", Kind: "Node", Namespaced: false},
32-
"namespace": {Group: "", Version: "v1", Kind: "Namespace", Namespaced: false},
33-
"persistentvolume": {Group: "", Version: "v1", Kind: "PersistentVolume", Namespaced: false},
34-
"clusterrole": {Group: "rbac.authorization.k8s.io", Version: "v1", Kind: "ClusterRole", Namespaced: false},
35-
"clusterrolebinding": {Group: "rbac.authorization.k8s.io", Version: "v1", Kind: "ClusterRoleBinding", Namespaced: false},
36-
"storageclass": {Group: "storage.k8s.io", Version: "v1", Kind: "StorageClass", Namespaced: false},
37-
"customresourcedefinition": {Group: "apiextensions.k8s.io", Version: "v1", Kind: "CustomResourceDefinition", Namespaced: false},
38-
"mutatingwebhookconfiguration": {Group: "admissionregistration.k8s.io", Version: "v1", Kind: "MutatingWebhookConfiguration", Namespaced: false},
39-
"validatingwebhookconfiguration": {Group: "admissionregistration.k8s.io", Version: "v1", Kind: "ValidatingWebhookConfiguration", Namespaced: false},
40-
}
41-
42-
// GetResourceInfo 根据资源类型字符串返回资源信息
43-
func GetResourceInfo(resourceType string) (ResourceInfo, bool) {
44-
resourceType = strings.ToLower(resourceType)
45-
if info, exists := resourceMap[resourceType]; exists {
46-
return info, true
47-
}
48-
return ResourceInfo{}, false
49-
}
50-
51-
// IsNamespaced 判断资源是否为命名空间级别
52-
func IsNamespaced(resourceType string) bool {
53-
resourceType = strings.ToLower(resourceType)
54-
if info, exists := resourceMap[resourceType]; exists {
55-
return info.Namespaced
56-
}
57-
return false
58-
}

‎mcp/metadata/utils.go

+6-7
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,13 @@
1-
package utils
1+
package metadata
22

33
import (
44
"context"
55
"strings"
66

77
"github.com/mark3labs/mcp-go/mcp"
8-
"github.com/weibaohui/kom/mcp/metadata"
98
)
109

11-
var resourceMap = map[string]metadata.ResourceInfo{
10+
var resourceMap = map[string]ResourceInfo{
1211
// 命名空间级别资源
1312
"pod": {Group: "", Version: "v1", Kind: "Pod", Namespaced: true},
1413
"deployment": {Group: "apps", Version: "v1", Kind: "Deployment", Namespaced: true},
@@ -39,12 +38,12 @@ var resourceMap = map[string]metadata.ResourceInfo{
3938
}
4039

4140
// GetResourceInfo 根据资源类型字符串返回资源信息
42-
func GetResourceInfo(resourceType string) (metadata.ResourceInfo, bool) {
41+
func GetResourceInfo(resourceType string) (ResourceInfo, bool) {
4342
resourceType = strings.ToLower(resourceType)
4443
if info, exists := resourceMap[resourceType]; exists {
4544
return info, true
4645
}
47-
return metadata.ResourceInfo{}, false
46+
return ResourceInfo{}, false
4847
}
4948

5049
// IsNamespaced 判断资源是否为命名空间级别
@@ -57,7 +56,7 @@ func IsNamespaced(resourceType string) bool {
5756
}
5857

5958
// ParseFromRequest 从请求中解析资源元数据
60-
func ParseFromRequest(ctx context.Context, request mcp.CallToolRequest, serverConfig *metadata.ServerConfig) (context.Context, *metadata.ResourceMetadata, error) {
59+
func ParseFromRequest(ctx context.Context, request mcp.CallToolRequest, serverConfig *ServerConfig) (context.Context, *ResourceMetadata, error) {
6160
if serverConfig != nil {
6261
if authVal, ok := ctx.Value(serverConfig.AuthKey).(string); ok {
6362
ctx = context.WithValue(ctx, serverConfig.AuthKey, authVal)
@@ -106,7 +105,7 @@ func ParseFromRequest(ctx context.Context, request mcp.CallToolRequest, serverCo
106105
kind = getStringParam(request, "kind", "")
107106
}
108107

109-
return ctx, &metadata.ResourceMetadata{
108+
return ctx, &ResourceMetadata{
110109
Cluster: cluster,
111110
Namespace: namespace,
112111
Name: name,

‎mcp/server.go

+51-9
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import (
44
"fmt"
55

66
"github.com/mark3labs/mcp-go/server"
7+
"github.com/weibaohui/kom/mcp/metadata"
78
"github.com/weibaohui/kom/mcp/tools/cluster"
89
"github.com/weibaohui/kom/mcp/tools/deployment"
910
"github.com/weibaohui/kom/mcp/tools/dynamic"
@@ -17,6 +18,9 @@ import (
1718
)
1819

1920
func RunMCPServer(name, version string, port int) {
21+
config.Name = name
22+
config.Version = version
23+
config.Port = port
2024
// 创建一个新的 MCP 服务器
2125
s := server.NewMCPServer(
2226
name,
@@ -27,15 +31,15 @@ func RunMCPServer(name, version string, port int) {
2731
)
2832

2933
// 注册工具
30-
dynamic.RegisterTools(s)
31-
pod.RegisterTools(s)
32-
cluster.RegisterTools(s)
33-
event.RegisterTools(s)
34-
deployment.RegisterTools(s)
35-
node.RegisterTools(s)
36-
storageclass.RegisterTools(s)
37-
ingressclass.RegisterTools(s)
38-
yaml.RegisterTools(s)
34+
dynamic.RegisterTools(s, config)
35+
pod.RegisterTools(s, config)
36+
cluster.RegisterTools(s, config)
37+
event.RegisterTools(s, config)
38+
deployment.RegisterTools(s, config)
39+
node.RegisterTools(s, config)
40+
storageclass.RegisterTools(s, config)
41+
ingressclass.RegisterTools(s, config)
42+
yaml.RegisterTools(s, config)
3943
// 创建 SSE 服务器
4044
sseServer := server.NewSSEServer(s)
4145

@@ -45,3 +49,41 @@ func RunMCPServer(name, version string, port int) {
4549
klog.Errorf("MCP Server error: %v\n", err)
4650
}
4751
}
52+
53+
var config *metadata.ServerConfig
54+
55+
func GetServerConfig() *metadata.ServerConfig {
56+
return config
57+
}
58+
func RunMCPServerWithOption(cfg *metadata.ServerConfig) {
59+
if cfg == nil {
60+
klog.Errorf("MCP Server error: config is nil\n")
61+
return
62+
}
63+
config = cfg
64+
// 创建一个新的 MCP 服务器
65+
s := server.NewMCPServer(
66+
config.Name,
67+
config.Version,
68+
config.ServerOptions...,
69+
)
70+
71+
// 注册工具
72+
dynamic.RegisterTools(s, config)
73+
pod.RegisterTools(s, config)
74+
cluster.RegisterTools(s, config)
75+
event.RegisterTools(s, config)
76+
deployment.RegisterTools(s, config)
77+
node.RegisterTools(s, config)
78+
storageclass.RegisterTools(s, config)
79+
ingressclass.RegisterTools(s, config)
80+
yaml.RegisterTools(s, config)
81+
// 创建 SSE 服务器
82+
sseServer := server.NewSSEServer(s, config.SSEOption...)
83+
84+
// 启动服务器
85+
err := sseServer.Start(fmt.Sprintf(":%d", config.Port))
86+
if err != nil {
87+
klog.Errorf("MCP Server error: %v\n", err)
88+
}
89+
}

‎mcp/tools/cluster/cluster_list.go

+3-2
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ import (
55

66
"github.com/mark3labs/mcp-go/mcp"
77
"github.com/weibaohui/kom/kom"
8-
"github.com/weibaohui/kom/mcp/tools"
8+
"github.com/weibaohui/kom/utils"
99
)
1010

1111
func ListClusters() mcp.Tool {
@@ -16,6 +16,7 @@ func ListClusters() mcp.Tool {
1616
}
1717

1818
func ListClustersHandler(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
19+
1920
// 获取所有已注册的集群名称
2021
clusters := kom.Clusters().AllClusters()
2122

@@ -27,5 +28,5 @@ func ListClustersHandler(ctx context.Context, request mcp.CallToolRequest) (*mcp
2728
})
2829
}
2930

30-
return tools.TextResult(result, nil)
31+
return utils.TextResult(result, nil)
3132
}

‎mcp/tools/cluster/reg.go

+5-1
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,13 @@ package cluster
22

33
import (
44
"github.com/mark3labs/mcp-go/server"
5+
"github.com/weibaohui/kom/mcp/metadata"
56
)
67

7-
func RegisterTools(s *server.MCPServer) {
8+
var config *metadata.ServerConfig
9+
10+
func RegisterTools(s *server.MCPServer, cfg *metadata.ServerConfig) {
11+
config = cfg
812
s.AddTool(
913
ListClusters(),
1014
ListClustersHandler,

‎mcp/tools/deployment/hpa.go

+6-4
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,10 @@ import (
55

66
"github.com/mark3labs/mcp-go/mcp"
77
"github.com/weibaohui/kom/kom"
8-
"github.com/weibaohui/kom/mcp/tools"
9-
"github.com/weibaohui/kom/mcp/tools/metadata"
8+
"github.com/weibaohui/kom/mcp/metadata"
9+
10+
"github.com/weibaohui/kom/utils"
11+
1012
appsv1 "k8s.io/api/apps/v1"
1113
"k8s.io/klog/v2"
1214
)
@@ -25,7 +27,7 @@ func HPAListDeploymentTool() mcp.Tool {
2527
// HPAListDeploymentHandler 处理查询Deployment的HPA列表的请求
2628
func HPAListDeploymentHandler(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
2729
// 获取参数
28-
meta, err := metadata.ParseFromRequest(request)
30+
ctx, meta, err := metadata.ParseFromRequest(ctx, request, config)
2931
if err != nil {
3032
return nil, err
3133
}
@@ -42,5 +44,5 @@ func HPAListDeploymentHandler(ctx context.Context, request mcp.CallToolRequest)
4244
result += "HPA " + item.Name + "\n"
4345
}
4446

45-
return tools.TextResult(result, meta)
47+
return utils.TextResult(result, meta)
4648
}

‎mcp/tools/deployment/managed_pods.go

+5-4
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,8 @@ import (
55

66
"github.com/mark3labs/mcp-go/mcp"
77
"github.com/weibaohui/kom/kom"
8-
"github.com/weibaohui/kom/mcp/tools"
9-
"github.com/weibaohui/kom/mcp/tools/metadata"
8+
"github.com/weibaohui/kom/mcp/metadata"
9+
"github.com/weibaohui/kom/utils"
1010
)
1111

1212
// ManagedPodsDeploymentTool 创建一个获取Deployment管理的Pod列表的工具
@@ -23,7 +23,8 @@ func ManagedPodsDeploymentTool() mcp.Tool {
2323
// ManagedPodsDeploymentHandler 处理获取Deployment管理的Pod列表的请求
2424
func ManagedPodsDeploymentHandler(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
2525
// 获取参数
26-
meta, err := metadata.ParseFromRequest(request)
26+
ctx, meta, err := metadata.ParseFromRequest(ctx, request, config)
27+
2728
if err != nil {
2829
return nil, err
2930
}
@@ -39,5 +40,5 @@ func ManagedPodsDeploymentHandler(ctx context.Context, request mcp.CallToolReque
3940
podNames = append(podNames, pod.Name)
4041
}
4142

42-
return tools.TextResult(podNames, meta)
43+
return utils.TextResult(podNames, meta)
4344
}

‎mcp/tools/deployment/reg.go

+5-1
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,13 @@ package deployment
22

33
import (
44
"github.com/mark3labs/mcp-go/server"
5+
"github.com/weibaohui/kom/mcp/metadata"
56
)
67

7-
func RegisterTools(s *server.MCPServer) {
8+
var config *metadata.ServerConfig
9+
10+
func RegisterTools(s *server.MCPServer, cfg *metadata.ServerConfig) {
11+
config = cfg
812
s.AddTool(
913
ManagedPodsDeploymentTool(),
1014
ManagedPodsDeploymentHandler,

‎mcp/tools/deployment/restart.go

+6-4
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,9 @@ import (
55

66
"github.com/mark3labs/mcp-go/mcp"
77
"github.com/weibaohui/kom/kom"
8-
"github.com/weibaohui/kom/mcp/tools"
9-
"github.com/weibaohui/kom/mcp/tools/metadata"
8+
"github.com/weibaohui/kom/mcp/metadata"
9+
"github.com/weibaohui/kom/utils"
10+
1011
appsv1 "k8s.io/api/apps/v1"
1112
)
1213

@@ -24,7 +25,8 @@ func RestartDeploymentTool() mcp.Tool {
2425
// RestartDeploymentHandler 处理重启Deployment的请求
2526
func RestartDeploymentHandler(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
2627
// 获取参数
27-
meta, err := metadata.ParseFromRequest(request)
28+
ctx, meta, err := metadata.ParseFromRequest(ctx, request, config)
29+
2830
if err != nil {
2931
return nil, err
3032
}
@@ -34,5 +36,5 @@ func RestartDeploymentHandler(ctx context.Context, request mcp.CallToolRequest)
3436
return nil, err
3537
}
3638

37-
return tools.TextResult("Successfully restarted deployment", meta)
39+
return utils.TextResult("Successfully restarted deployment", meta)
3840
}

0 commit comments

Comments
 (0)