mirror of
https://github.com/RuoJi6/ARTEX.git
synced 2026-10-27 17:03:49 +08:00
91 lines
2.7 KiB
Go
91 lines
2.7 KiB
Go
package server
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"log"
|
|
"time"
|
|
|
|
"github.com/Autumn-27/artex/db"
|
|
"github.com/Autumn-27/artex/mcphttp"
|
|
"github.com/Autumn-27/norma/mcp"
|
|
actool "github.com/Autumn-27/norma/tool"
|
|
)
|
|
|
|
// mcpClient is the shared surface of a connected MCP server (stdio, Streamable HTTP,
|
|
// or legacy SSE),
|
|
// so tools/list and cleanup are handled uniformly regardless of transport.
|
|
type mcpClient interface {
|
|
Tools(context.Context) ([]actool.CoreTool, error)
|
|
Close() error
|
|
}
|
|
|
|
// connectMCP dials one MCP server per its transport. Callers must Close the client.
|
|
func connectMCP(ctx context.Context, m *db.MCPServer) (mcpClient, error) {
|
|
switch m.Transport {
|
|
case "stdio":
|
|
if m.Command == "" {
|
|
return nil, fmt.Errorf("stdio 传输缺少命令")
|
|
}
|
|
return mcp.NewStdioClient(ctx, m.Name, m.Command, jsonStrMap(m.Env), jsonStrSlice(m.Args)...)
|
|
case "http":
|
|
if m.URL == "" {
|
|
return nil, fmt.Errorf("http 传输缺少 URL")
|
|
}
|
|
// env map doubles as HTTP headers (e.g. Authorization).
|
|
return mcphttp.New(ctx, m.Name, m.URL, jsonStrMap(m.Env), m.Insecure)
|
|
case "sse":
|
|
if m.URL == "" {
|
|
return nil, fmt.Errorf("sse 传输缺少 URL")
|
|
}
|
|
return mcphttp.NewSSE(ctx, m.Name, m.URL, jsonStrMap(m.Env), m.Insecure)
|
|
default:
|
|
return nil, fmt.Errorf("未知传输方式 %q", m.Transport)
|
|
}
|
|
}
|
|
|
|
// discoverAndCacheMCP connects to one MCP, lists its tools, and persists the tool
|
|
// names to mcp_tools_cache so the UI shows them without a live connection.
|
|
func (s *Server) discoverAndCacheMCP(ctx context.Context, m *db.MCPServer) error {
|
|
cl, err := connectMCP(ctx, m)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer cl.Close()
|
|
ts, err := cl.Tools(ctx)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
tools := make([]db.MCPTool, 0, len(ts))
|
|
for _, t := range ts {
|
|
tools = append(tools, db.MCPTool{Name: t.Name(), Description: t.Description()})
|
|
}
|
|
if err := s.m.pg.SaveMCPTools(m.ID, tools); err != nil {
|
|
return err
|
|
}
|
|
log.Printf("[mcp] %s 发现 %d 个工具并已缓存", m.Name, len(tools))
|
|
return nil
|
|
}
|
|
|
|
// discoverEmptyMCPsOnStartup fills the tool cache for any enabled MCP that has none
|
|
// yet (notably the seeded browser MCP on first run). Runs sequentially in one
|
|
// goroutine so we never spawn many stdio servers (npx) at once, and never blocks
|
|
// startup. Best-effort: a failure leaves the cache empty to retry next start.
|
|
func (s *Server) discoverEmptyMCPsOnStartup() {
|
|
servers, err := s.m.pg.ListMCP()
|
|
if err != nil {
|
|
log.Printf("[mcp] 启动自动发现: 读取列表失败: %v", err)
|
|
return
|
|
}
|
|
for _, m := range servers {
|
|
if !m.Enabled || len(m.Tools) > 0 {
|
|
continue
|
|
}
|
|
ctx, cancel := context.WithTimeout(s.ctx, 90*time.Second)
|
|
if err := s.discoverAndCacheMCP(ctx, m); err != nil {
|
|
log.Printf("[mcp] 启动自动发现 %s 失败: %v", m.Name, err)
|
|
}
|
|
cancel()
|
|
}
|
|
}
|