476 lines
12 KiB
Go
476 lines
12 KiB
Go
package main
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"log"
|
|
"net/http"
|
|
"os"
|
|
"path/filepath"
|
|
"strconv"
|
|
"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, client 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 == "" && !getEnvBool("TAO_ALLOW_ANON", false) {
|
|
log.Printf("[MCP Response] missing token for method=%s", req.Method)
|
|
return
|
|
}
|
|
connKey := buildConnKey(token, client)
|
|
if ch, ok := s.conns.Load(connKey); 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 client=%s method=%s", token, client, req.Method)
|
|
}
|
|
}
|
|
|
|
func getEnv(key, def string) string {
|
|
if v := os.Getenv(key); v != "" {
|
|
return v
|
|
}
|
|
return def
|
|
}
|
|
|
|
func getEnvBool(key string, def bool) bool {
|
|
v := strings.ToLower(strings.TrimSpace(os.Getenv(key)))
|
|
if v == "" {
|
|
return def
|
|
}
|
|
switch v {
|
|
case "1", "true", "yes", "on":
|
|
return true
|
|
case "0", "false", "no", "off":
|
|
return false
|
|
default:
|
|
return def
|
|
}
|
|
}
|
|
|
|
func getEnvInt(key string, def int) int {
|
|
if v := os.Getenv(key); v != "" {
|
|
if n, err := strconv.Atoi(v); err == nil {
|
|
return n
|
|
}
|
|
}
|
|
return def
|
|
}
|
|
|
|
func extractToken(r *http.Request) (string, bool) {
|
|
if q := r.URL.Query().Get("token"); q != "" {
|
|
return q, true
|
|
}
|
|
h := r.Header.Get("Authorization")
|
|
if strings.HasPrefix(h, "Bearer ") {
|
|
return strings.TrimSpace(strings.TrimPrefix(h, "Bearer ")), false
|
|
}
|
|
return "", false
|
|
}
|
|
|
|
func buildConnKey(token string, client string) string {
|
|
if token == "" {
|
|
token = "anon"
|
|
}
|
|
if client != "" {
|
|
return token + "_" + client
|
|
}
|
|
return token
|
|
}
|
|
|
|
func generateClientID() string {
|
|
return fmt.Sprintf("c%d", time.Now().UnixNano())
|
|
}
|
|
|
|
func parseCORSOrigins() (bool, []string) {
|
|
raw := strings.TrimSpace(os.Getenv("TAO_CORS_ORIGINS"))
|
|
if raw == "" {
|
|
return false, nil
|
|
}
|
|
if raw == "*" {
|
|
return true, nil
|
|
}
|
|
parts := strings.Split(raw, ",")
|
|
var origins []string
|
|
for _, p := range parts {
|
|
p = strings.TrimSpace(p)
|
|
if p != "" {
|
|
origins = append(origins, p)
|
|
}
|
|
}
|
|
return false, origins
|
|
}
|
|
|
|
func setCORSHeaders(w http.ResponseWriter, r *http.Request) {
|
|
origin := r.Header.Get("Origin")
|
|
if origin == "" {
|
|
return
|
|
}
|
|
allowAll, origins := parseCORSOrigins()
|
|
if allowAll {
|
|
w.Header().Set("Access-Control-Allow-Origin", "*")
|
|
} else {
|
|
allowed := false
|
|
for _, o := range origins {
|
|
if o == origin {
|
|
allowed = true
|
|
break
|
|
}
|
|
}
|
|
if !allowed {
|
|
return
|
|
}
|
|
w.Header().Set("Access-Control-Allow-Origin", origin)
|
|
w.Header().Set("Vary", "Origin")
|
|
}
|
|
w.Header().Set("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
|
|
w.Header().Set("Access-Control-Allow-Headers", "Content-Type, Authorization")
|
|
}
|
|
|
|
func isSubpath(path string, base string) bool {
|
|
absPath, err1 := filepath.Abs(path)
|
|
absBase, err2 := filepath.Abs(base)
|
|
if err1 != nil || err2 != nil {
|
|
return false
|
|
}
|
|
if absPath == absBase {
|
|
return true
|
|
}
|
|
return strings.HasPrefix(absPath, absBase+string(filepath.Separator))
|
|
}
|
|
|
|
// --- 以简御繁：鉴权 ---
|
|
func (s *TaoServer) checkAuth(r *http.Request) bool {
|
|
token := getEnv("TAO_AUTH_TOKEN", "")
|
|
if token == "" {
|
|
return getEnvBool("TAO_ALLOW_ANON", false)
|
|
}
|
|
reqToken, _ := extractToken(r)
|
|
if reqToken == 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" {
|
|
setCORSHeaders(w, r)
|
|
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("X-Accel-Buffering", "no")
|
|
setCORSHeaders(w, r)
|
|
|
|
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"
|
|
}
|
|
|
|
queryToken := r.URL.Query().Get("token")
|
|
token, _ := extractToken(r)
|
|
if token == "" && !getEnvBool("TAO_ALLOW_ANON", false) {
|
|
http.Error(w, "Unauthorized", http.StatusUnauthorized)
|
|
return
|
|
}
|
|
|
|
client := r.URL.Query().Get("client")
|
|
if client == "" {
|
|
client = generateClientID()
|
|
}
|
|
|
|
endpointToken := queryToken
|
|
if endpointToken == "" && getEnvBool("TAO_ENDPOINT_TOKEN_ECHO", false) {
|
|
endpointToken = token
|
|
}
|
|
|
|
if endpointToken != "" {
|
|
if strings.Contains(endpoint, "?") {
|
|
endpoint = endpoint + "&token=" + endpointToken
|
|
} else {
|
|
endpoint = endpoint + "?token=" + endpointToken
|
|
}
|
|
}
|
|
if client != "" {
|
|
if strings.Contains(endpoint, "?") {
|
|
endpoint = endpoint + "&client=" + client
|
|
} else {
|
|
endpoint = endpoint + "?client=" + client
|
|
}
|
|
}
|
|
fmt.Fprintf(w, "event: endpoint\ndata: %s\n\n", endpoint)
|
|
flusher.Flush()
|
|
|
|
msgChan := make(chan string, 50)
|
|
connKey := buildConnKey(token, client)
|
|
s.conns.Store(connKey, msgChan)
|
|
defer s.conns.Delete(connKey)
|
|
|
|
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) {
|
|
setCORSHeaders(w, r)
|
|
|
|
if r.Method == "OPTIONS" {
|
|
w.WriteHeader(http.StatusOK)
|
|
return
|
|
}
|
|
|
|
bodyBytes, _ := io.ReadAll(r.Body)
|
|
r.Body = io.NopCloser(bytes.NewBuffer(bodyBytes))
|
|
if getEnvBool("TAO_DEBUG", false) {
|
|
log.Printf("[MCP POST] From=%s URL=%s Body=%s", r.RemoteAddr, r.URL.String(), string(bodyBytes))
|
|
} else {
|
|
log.Printf("[MCP POST] From=%s URL=%s", r.RemoteAddr, r.URL.String())
|
|
}
|
|
|
|
var req MCPRequest
|
|
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
|
http.Error(w, "Bad Request", 400)
|
|
return
|
|
}
|
|
|
|
token, _ := extractToken(r)
|
|
client := r.URL.Query().Get("client")
|
|
w.WriteHeader(http.StatusAccepted)
|
|
|
|
go s.dispatchMCP(token, client, req)
|
|
}
|
|
|
|
// --- MCP Unified Gateway (/mcp) ---
|
|
func (s *TaoServer) MCPUnifiedHandler(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 r.Method == http.MethodGet {
|
|
s.SSEHandler(w, r)
|
|
return
|
|
}
|
|
if r.Method == http.MethodPost {
|
|
accept := r.Header.Get("Accept")
|
|
if strings.Contains(accept, "text/event-stream") {
|
|
s.SSEHandler(w, r)
|
|
return
|
|
}
|
|
if r.ContentLength == 0 {
|
|
s.SSEHandler(w, r)
|
|
return
|
|
}
|
|
bodyBytes, _ := io.ReadAll(r.Body)
|
|
r.Body = io.NopCloser(bytes.NewBuffer(bodyBytes))
|
|
if len(bodyBytes) == 0 {
|
|
s.SSEHandler(w, r)
|
|
return
|
|
}
|
|
if strings.Contains(string(bodyBytes), "\"jsonrpc\":\"2.0\"") {
|
|
s.MessageHandler(w, r)
|
|
return
|
|
}
|
|
http.Error(w, "Bad Request", 400)
|
|
return
|
|
}
|
|
http.Error(w, "Method Not Allowed", 405)
|
|
}
|
|
|
|
// --- 主程序 (Main) ---
|
|
func main() {
|
|
if getEnv("TAO_AUTH_TOKEN", "") == "" && !getEnvBool("TAO_ALLOW_ANON", false) {
|
|
log.Fatal("TAO_AUTH_TOKEN is required unless TAO_ALLOW_ANON=true")
|
|
}
|
|
|
|
memoryRoot := getEnv("MEMORY_ROOT", "./knowledge_ocean")
|
|
searchRoot := getEnv("TAO_SEARCH_ROOT", memoryRoot)
|
|
if !isSubpath(searchRoot, memoryRoot) {
|
|
log.Printf("TAO_SEARCH_ROOT must be under MEMORY_ROOT, fallback to MEMORY_ROOT")
|
|
searchRoot = memoryRoot
|
|
}
|
|
|
|
server := &TaoServer{
|
|
config: Config{
|
|
MemoryRoot: memoryRoot,
|
|
Port: getEnv("PORT", "5001"),
|
|
SearchRoot: searchRoot,
|
|
MaxSearchFiles: getEnvInt("TAO_SEARCH_MAX_FILES", 2000),
|
|
},
|
|
}
|
|
|
|
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", server.requireAuth(server.MCPUnifiedHandler))
|
|
http.HandleFunc("/mcp/", server.requireAuth(server.MCPUnifiedHandler))
|
|
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))
|
|
}
|