feat: tao-mcp-go core features
This commit is contained in:
OpenClaw Agent
295
main.go
Normal file
295
main.go
Normal file
@@ -0,0 +1,295 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"net/http"
|
||||
"os"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// --- MCP Protocol Types ---
|
||||
type MCPRequest struct {
|
||||
JSONRPC string `json:"jsonrpc"`
|
||||
ID any `json:"id,omitempty"`
|
||||
Method string `json:"method"`
|
||||
Params json.RawMessage `json:"params,omitempty"`
|
||||
}
|
||||
|
||||
type MCPResponse struct {
|
||||
JSONRPC string `json:"jsonrpc"`
|
||||
ID any `json:"id,omitempty"`
|
||||
Result any `json:"result,omitempty"`
|
||||
Error interface{} `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
type MCPContent struct {
|
||||
Type string `json:"type"`
|
||||
Text string `json:"text"`
|
||||
}
|
||||
|
||||
func sendMCPResponse(id any, result any) MCPResponse {
|
||||
return MCPResponse{JSONRPC: "2.0", ID: id, Result: result}
|
||||
}
|
||||
|
||||
func (s *TaoServer) dispatchMCP(token string, req MCPRequest) {
|
||||
var resp MCPResponse
|
||||
resp.JSONRPC = "2.0"
|
||||
resp.ID = req.ID
|
||||
|
||||
switch req.Method {
|
||||
case "initialize":
|
||||
resp.Result = map[string]interface{}{
|
||||
"protocolVersion": "2025-06-18",
|
||||
"capabilities": map[string]interface{}{
|
||||
"tools": map[string]interface{}{"listChanged": false},
|
||||
"resources": map[string]interface{}{"listChanged": false},
|
||||
"prompts": map[string]interface{}{"listChanged": false},
|
||||
"logging": map[string]interface{}{},
|
||||
},
|
||||
"serverInfo": map[string]string{
|
||||
"name": "Tao-Memory-Server",
|
||||
"version": "1.2.0",
|
||||
},
|
||||
}
|
||||
case "notifications/initialized":
|
||||
log.Printf("[MCP Notify] initialized from %s", token)
|
||||
return
|
||||
case "tools/list":
|
||||
resp.Result = map[string]interface{}{
|
||||
"tools": buildToolList(),
|
||||
}
|
||||
case "tools/call":
|
||||
var params struct {
|
||||
Name string `json:"name"`
|
||||
Arguments map[string]interface{} `json:"arguments"`
|
||||
}
|
||||
_ = json.Unmarshal(req.Params, ¶ms)
|
||||
if tool, ok := ToolRegistry[params.Name]; ok {
|
||||
result, err := tool.Handler(params.Arguments)
|
||||
if err != nil {
|
||||
resp.Result = map[string]interface{}{
|
||||
"content": []MCPContent{{Type: "text", Text: "error: " + err.Error()}},
|
||||
}
|
||||
} else {
|
||||
resp.Result = map[string]interface{}{
|
||||
"content": []MCPContent{{Type: "text", Text: result}},
|
||||
}
|
||||
}
|
||||
} else {
|
||||
resp.Result = map[string]interface{}{
|
||||
"content": []MCPContent{{Type: "text", Text: "error: tool not found"}},
|
||||
}
|
||||
}
|
||||
default:
|
||||
_ = s.Record("agent_action", fmt.Sprintf("执行指令: %+v", req), 2)
|
||||
return
|
||||
}
|
||||
|
||||
if token == "" {
|
||||
log.Printf("[MCP Response] missing token for method=%s", req.Method)
|
||||
return
|
||||
}
|
||||
if ch, ok := s.conns.Load(token); ok {
|
||||
if b, err := json.Marshal(resp); err == nil {
|
||||
ch.(chan string) <- string(b)
|
||||
log.Printf("[MCP Response] sent via SSE method=%s", req.Method)
|
||||
}
|
||||
} else {
|
||||
log.Printf("[MCP Response] no SSE channel for token=%s method=%s", token, req.Method)
|
||||
}
|
||||
}
|
||||
|
||||
func getEnv(key, def string) string {
|
||||
if v := os.Getenv(key); v != "" {
|
||||
return v
|
||||
}
|
||||
return def
|
||||
}
|
||||
|
||||
// --- 以简御繁：鉴权 ---
|
||||
func (s *TaoServer) checkAuth(r *http.Request) bool {
|
||||
token := getEnv("TAO_AUTH_TOKEN", "")
|
||||
if token == "" {
|
||||
return true // 未配置则不启用鉴权
|
||||
}
|
||||
// Header Bearer
|
||||
h := r.Header.Get("Authorization")
|
||||
if h == "Bearer "+token {
|
||||
return true
|
||||
}
|
||||
// Query token
|
||||
if q := r.URL.Query().Get("token"); q != "" && q == token {
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func (s *TaoServer) requireAuth(next http.HandlerFunc) http.HandlerFunc {
|
||||
return func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method == "OPTIONS" {
|
||||
w.Header().Set("Access-Control-Allow-Origin", "*")
|
||||
w.Header().Set("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
|
||||
w.Header().Set("Access-Control-Allow-Headers", "Content-Type, Authorization")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
return
|
||||
}
|
||||
if !s.checkAuth(r) {
|
||||
http.Error(w, "Unauthorized", http.StatusUnauthorized)
|
||||
return
|
||||
}
|
||||
next(w, r)
|
||||
}
|
||||
}
|
||||
|
||||
// --- 感官 (Webhook Adapters) ---
|
||||
// 适配 Gitea 的 Push Webhook
|
||||
func (s *TaoServer) GiteaHandler(w http.ResponseWriter, r *http.Request) {
|
||||
var payload struct {
|
||||
Repository struct {
|
||||
Name string `json:"name"`
|
||||
} `json:"repository"`
|
||||
Commits []struct {
|
||||
Message string `json:"message"`
|
||||
} `json:"commits"`
|
||||
}
|
||||
|
||||
if err := json.NewDecoder(r.Body).Decode(&payload); err != nil {
|
||||
http.Error(w, "Bad Request", 400)
|
||||
return
|
||||
}
|
||||
|
||||
if len(payload.Commits) > 0 {
|
||||
msg := payload.Commits[0].Message
|
||||
summary := fmt.Sprintf("代码演化于 [%s]: %s", payload.Repository.Name, msg)
|
||||
_ = s.Record("code", summary, 4)
|
||||
}
|
||||
w.WriteHeader(200)
|
||||
}
|
||||
|
||||
// 适配 SmsReceiver-go 的短信推送
|
||||
func (s *TaoServer) SmsHandler(w http.ResponseWriter, r *http.Request) {
|
||||
var payload struct {
|
||||
From string `json:"from"`
|
||||
Content string `json:"content"`
|
||||
}
|
||||
if err := json.NewDecoder(r.Body).Decode(&payload); err != nil {
|
||||
http.Error(w, "Bad Request", 400)
|
||||
return
|
||||
}
|
||||
|
||||
summary := fmt.Sprintf("收到信号 [%s]: %s", payload.From, payload.Content)
|
||||
_ = s.Record("sms", summary, 3)
|
||||
w.WriteHeader(200)
|
||||
}
|
||||
|
||||
// --- MCP SSE ---
|
||||
func (s *TaoServer) SSEHandler(w http.ResponseWriter, r *http.Request) {
|
||||
log.Printf("[SSE Connect] Remote=%s URL=%s", r.RemoteAddr, r.URL.String())
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
w.Header().Set("Cache-Control", "no-cache")
|
||||
w.Header().Set("Connection", "keep-alive")
|
||||
w.Header().Set("Access-Control-Allow-Origin", "*")
|
||||
w.Header().Set("X-Accel-Buffering", "no")
|
||||
|
||||
flusher, ok := w.(http.Flusher)
|
||||
if !ok {
|
||||
http.Error(w, "Streaming unsupported", http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
|
||||
// 告知客户端 POST 入口（按客户端拼接习惯输出）
|
||||
style := getEnv("TAO_ENDPOINT_STYLE", "message")
|
||||
endpoint := "mcp/message"
|
||||
if style == "message" {
|
||||
endpoint = "message"
|
||||
}
|
||||
// 若通过 query token 访问，也把 token 拼到 endpoint（便于客户端无 Header）
|
||||
token := r.URL.Query().Get("token")
|
||||
if token != "" {
|
||||
if strings.Contains(endpoint, "?") {
|
||||
endpoint = endpoint + "&token=" + token
|
||||
} else {
|
||||
endpoint = endpoint + "?token=" + token
|
||||
}
|
||||
}
|
||||
fmt.Fprintf(w, "event: endpoint\ndata: %s\n\n", endpoint)
|
||||
flusher.Flush()
|
||||
|
||||
var msgChan chan string
|
||||
if token != "" {
|
||||
msgChan = make(chan string, 50)
|
||||
s.conns.Store(token, msgChan)
|
||||
defer s.conns.Delete(token)
|
||||
}
|
||||
|
||||
ticker := time.NewTicker(5 * time.Second)
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-r.Context().Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
fmt.Fprintf(w, ":ping\n\n")
|
||||
flusher.Flush()
|
||||
case msg := <-msgChan:
|
||||
fmt.Fprintf(w, "event: message\ndata: %s\n\n", msg)
|
||||
flusher.Flush()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// --- MCP Message ---
|
||||
func (s *TaoServer) MessageHandler(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Access-Control-Allow-Origin", "*")
|
||||
w.Header().Set("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
|
||||
w.Header().Set("Access-Control-Allow-Headers", "Content-Type, Authorization")
|
||||
|
||||
if r.Method == "OPTIONS" {
|
||||
w.WriteHeader(http.StatusOK)
|
||||
return
|
||||
}
|
||||
|
||||
bodyBytes, _ := io.ReadAll(r.Body)
|
||||
r.Body = io.NopCloser(bytes.NewBuffer(bodyBytes))
|
||||
log.Printf("[MCP POST] From=%s URL=%s Body=%s", r.RemoteAddr, r.URL.String(), string(bodyBytes))
|
||||
|
||||
var req MCPRequest
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
http.Error(w, "Bad Request", 400)
|
||||
return
|
||||
}
|
||||
|
||||
token := r.URL.Query().Get("token")
|
||||
w.WriteHeader(http.StatusAccepted)
|
||||
|
||||
go s.dispatchMCP(token, req)
|
||||
}
|
||||
|
||||
// --- 主程序 (Main) ---
|
||||
func main() {
|
||||
server := &TaoServer{
|
||||
config: Config{
|
||||
MemoryRoot: getEnv("MEMORY_ROOT", "./knowledge_ocean"),
|
||||
Port: getEnv("PORT", "5001"),
|
||||
},
|
||||
}
|
||||
|
||||
server.RegisterTools()
|
||||
|
||||
// 启动 Webhook 监听 (感知层)
|
||||
http.HandleFunc("/ingest/gitea", server.requireAuth(server.GiteaHandler))
|
||||
http.HandleFunc("/ingest/sms", server.requireAuth(server.SmsHandler))
|
||||
|
||||
// MCP SSE + Message
|
||||
http.HandleFunc("/mcp/sse", server.requireAuth(server.SSEHandler))
|
||||
http.HandleFunc("/mcp/message", server.requireAuth(server.MessageHandler))
|
||||
|
||||
fmt.Printf("Tao Memory Server 启动。道场地址: :%s\n", server.config.Port)
|
||||
log.Fatal(http.ListenAndServe(":"+server.config.Port, nil))
|
||||
}
|
||||
Reference in New Issue
Block a user