Files
higress/plugins/wasm-go/pkg/mcp/server/base_server.go

101 lines
2.8 KiB
Go

// Copyright (c) 2022 Alibaba Group Holding Ltd.
//
// Licensed 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 server
import (
"encoding/base64"
"encoding/json"
"github.com/higress-group/proxy-wasm-go-sdk/proxywasm"
"github.com/higress-group/wasm-go/pkg/log"
)
// BaseMCPServer provides common functionality for MCP servers
type BaseMCPServer struct {
tools map[string]Tool
config []byte
}
// NewBaseMCPServer creates a new BaseMCPServer
func NewBaseMCPServer() BaseMCPServer {
return BaseMCPServer{
tools: make(map[string]Tool),
}
}
// AddMCPTool adds a tool to the server
func (s *BaseMCPServer) AddMCPTool(name string, tool Tool) Server {
if _, exist := s.tools[name]; exist {
log.Errorf("Conflict! There is a tool with the same name:%s", name)
return s
}
s.tools[name] = tool
return s
}
// GetMCPTools returns all tools registered with the server
func (s *BaseMCPServer) GetMCPTools() map[string]Tool {
return s.tools
}
// SetConfig sets the server configuration
func (s *BaseMCPServer) SetConfig(config []byte) {
s.config = config
}
// GetConfig gets the server configuration
// It first tries to get the config from the request header, then falls back to the stored config
func (s *BaseMCPServer) GetConfig(v any) {
var config []byte
serverConfigBase64, _ := proxywasm.GetHttpRequestHeader("x-higress-mcpserver-config")
proxywasm.RemoveHttpRequestHeader("x-higress-mcpserver-config")
if serverConfigBase64 != "" {
serverConfig, err := base64.StdEncoding.DecodeString(serverConfigBase64)
if err != nil {
log.Errorf("base64 decode mcp server config failed:%s, bytes:%s", err, serverConfigBase64)
} else {
config = serverConfig
}
log.Infof("parse server config from request, config:%s", serverConfig)
} else {
config = s.config
}
if len(config) == 0 {
return
}
err := json.Unmarshal(config, v)
if err != nil {
log.Errorf("json unmarshal server config failed:%v, config:%s", err, config)
}
}
// Clone creates a copy of the server
// This method should be overridden by derived types
func (s *BaseMCPServer) Clone() Server {
panic("Clone method must be implemented by derived types")
}
// CloneBase creates a copy of the base server
func (s *BaseMCPServer) CloneBase() BaseMCPServer {
newServer := BaseMCPServer{
tools: make(map[string]Tool),
config: s.config,
}
for k, v := range s.tools {
newServer.tools[k] = v
}
return newServer
}