245 lines
6.7 KiB
Go
245 lines
6.7 KiB
Go
package mcpclient
|
||
|
||
import (
|
||
"context"
|
||
"crypto/sha256"
|
||
"encoding/json"
|
||
"fmt"
|
||
"github.com/google/jsonschema-go/jsonschema"
|
||
"strings"
|
||
"time"
|
||
|
||
"github.com/modelcontextprotocol/go-sdk/mcp"
|
||
)
|
||
|
||
// AgentTool keeps provider function names unique across MCP servers.
|
||
type AgentTool struct {
|
||
Name string
|
||
Server string
|
||
Tool string
|
||
Description string
|
||
Schema json.RawMessage
|
||
validator *jsonschema.Resolved
|
||
}
|
||
|
||
type liveServer struct {
|
||
config ServerConfig
|
||
secrets Secrets
|
||
session *mcp.ClientSession
|
||
cancel context.CancelFunc
|
||
done chan struct{}
|
||
}
|
||
|
||
type Toolset struct {
|
||
service *Service
|
||
servers []*liveServer
|
||
tools []AgentTool
|
||
targets map[string]*liveServer
|
||
gate chan struct{}
|
||
}
|
||
|
||
// OpenTools gives each chat its own sessions. It discovers only enabled servers;
|
||
// configuration/credentials are snapshotted under lock, never during network IO.
|
||
func OpenTools(ctx context.Context, service *Service, warning func(string), filters ...func(ServerConfig) bool) (*Toolset, error) {
|
||
set := &Toolset{service: service, targets: map[string]*liveServer{}}
|
||
if service == nil {
|
||
return set, nil
|
||
}
|
||
service.mu.Lock()
|
||
if service.loadErr != nil {
|
||
service.mu.Unlock()
|
||
return nil, fmt.Errorf("MCP 配置读取失败，请检查设置")
|
||
}
|
||
var candidates []*liveServer
|
||
for _, config := range service.servers {
|
||
if !config.Enabled {
|
||
continue
|
||
}
|
||
allowed := true
|
||
for _, filter := range filters {
|
||
if !filter(config) {
|
||
allowed = false
|
||
break
|
||
}
|
||
}
|
||
if !allowed {
|
||
continue
|
||
}
|
||
_, secret, _, err := service.resolve(ServerInput{Config: config})
|
||
if err != nil {
|
||
if warning != nil {
|
||
warning("MCP「" + config.Name + "」配置不可用")
|
||
}
|
||
continue
|
||
}
|
||
candidates = append(candidates, &liveServer{config: config, secrets: secret})
|
||
}
|
||
service.mu.Unlock()
|
||
for _, live := range candidates {
|
||
if ctx.Err() != nil {
|
||
set.Close()
|
||
return nil, ctx.Err()
|
||
}
|
||
child, cancel := context.WithCancel(ctx)
|
||
live.cancel, live.done = cancel, make(chan struct{})
|
||
ready := make(chan error, 1)
|
||
var found []AgentTool
|
||
go func() {
|
||
defer close(live.done)
|
||
err := withSession(child, live.config, live.secrets, func(session *mcp.ClientSession) error {
|
||
live.session = session
|
||
for tool, err := range session.Tools(child, nil) {
|
||
if err != nil {
|
||
return err
|
||
}
|
||
if !supportedTool(live.config, tool.Name) {
|
||
continue
|
||
}
|
||
if len(found) >= 128 {
|
||
return fmt.Errorf("服务工具超过 128 个")
|
||
}
|
||
schema, err := json.Marshal(tool.InputSchema)
|
||
if err != nil || len(schema) > 64*1024 {
|
||
return fmt.Errorf("工具定义过大或无效")
|
||
}
|
||
var definition jsonschema.Schema
|
||
if err := json.Unmarshal(schema, &definition); err != nil {
|
||
return fmt.Errorf("工具参数定义无效")
|
||
}
|
||
validator, err := definition.Resolve(nil)
|
||
if err != nil {
|
||
return fmt.Errorf("工具参数定义无法校验")
|
||
}
|
||
hash := fmt.Sprintf("%x", sha256.Sum256([]byte(live.config.ID+":"+tool.Name)))[:12]
|
||
name := strings.Map(func(r rune) rune {
|
||
if r >= 'a' && r <= 'z' || r >= 'A' && r <= 'Z' || r >= '0' && r <= '9' || r == '_' || r == '-' {
|
||
return r
|
||
}
|
||
return '_'
|
||
}, tool.Name)
|
||
if len(name) > 45 {
|
||
name = name[:45]
|
||
}
|
||
found = append(found, AgentTool{Name: "mcp_" + hash + "_" + name, Server: live.config.Name, Tool: tool.Name, Description: toolDescription(live.config, tool.Name, tool.Description), Schema: schema, validator: validator})
|
||
}
|
||
ready <- nil
|
||
<-child.Done()
|
||
return nil
|
||
})
|
||
if err != nil {
|
||
select {
|
||
case ready <- err:
|
||
default:
|
||
}
|
||
}
|
||
}()
|
||
timer := time.NewTimer(time.Duration(live.config.TimeoutSeconds) * time.Second)
|
||
var err error
|
||
select {
|
||
case err = <-ready:
|
||
case <-ctx.Done():
|
||
err = ctx.Err()
|
||
case <-timer.C:
|
||
err = context.DeadlineExceeded
|
||
}
|
||
timer.Stop()
|
||
if err != nil || len(set.tools)+len(found) > 128 {
|
||
cancel()
|
||
<-live.done
|
||
if ctx.Err() != nil {
|
||
set.Close()
|
||
return nil, ctx.Err()
|
||
}
|
||
if warning != nil {
|
||
warning("MCP「" + live.config.Name + "」连接或工具发现失败（最多支持 128 个工具）")
|
||
}
|
||
continue
|
||
}
|
||
set.servers = append(set.servers, live)
|
||
for _, tool := range found {
|
||
set.tools = append(set.tools, tool)
|
||
set.targets[tool.Name] = live
|
||
}
|
||
}
|
||
return set, nil
|
||
}
|
||
|
||
func (s *Toolset) Close() {
|
||
for _, server := range s.servers {
|
||
server.cancel()
|
||
}
|
||
for _, server := range s.servers {
|
||
<-server.done
|
||
}
|
||
if s.gate != nil {
|
||
<-s.gate
|
||
s.gate = nil
|
||
}
|
||
}
|
||
func (s *Toolset) Tools() []AgentTool { return s.tools }
|
||
|
||
func (s *Toolset) Call(ctx context.Context, name, arguments string) (*mcp.CallToolResult, error) {
|
||
live := s.targets[name]
|
||
if live == nil {
|
||
return nil, fmt.Errorf("未知的 MCP 工具")
|
||
}
|
||
// Serialize action sequences across chats. One lease for all MCP servers
|
||
// avoids deadlocks when two chats need the same servers in opposite order.
|
||
// Text-only generation remains concurrent; waiting for the lease can cancel.
|
||
if s.gate == nil {
|
||
s.service.mu.Lock()
|
||
gate := s.service.agentGate
|
||
if gate == nil {
|
||
gate = make(chan struct{}, 1)
|
||
s.service.agentGate = gate
|
||
}
|
||
s.service.mu.Unlock()
|
||
select {
|
||
case gate <- struct{}{}:
|
||
s.gate = gate
|
||
case <-ctx.Done():
|
||
return nil, ctx.Err()
|
||
}
|
||
}
|
||
var args map[string]any
|
||
if len(arguments) > 64*1024 || json.Unmarshal([]byte(arguments), &args) != nil || args == nil {
|
||
return nil, fmt.Errorf("工具参数必须是 JSON 对象（最多 64 KB）")
|
||
}
|
||
// An already-open session must not bypass a later disable/delete/target edit.
|
||
s.service.mu.Lock()
|
||
i := s.service.index(live.config.ID)
|
||
valid := i >= 0 && s.service.servers[i].Enabled && sameTarget(s.service.servers[i], live.config)
|
||
s.service.mu.Unlock()
|
||
if !valid {
|
||
return nil, fmt.Errorf("MCP 服务已停用、删除或修改，请重新发起对话")
|
||
}
|
||
var original string
|
||
for _, tool := range s.tools {
|
||
if tool.Name == name {
|
||
if err := tool.validator.Validate(args); err != nil {
|
||
return nil, fmt.Errorf("工具参数不符合定义: %.1000s", err.Error())
|
||
}
|
||
original = tool.Tool
|
||
break
|
||
}
|
||
}
|
||
callCtx, cancel := context.WithTimeout(ctx, time.Duration(live.config.TimeoutSeconds)*time.Second)
|
||
defer cancel()
|
||
result, err := live.session.CallTool(callCtx, &mcp.CallToolParams{Name: original, Arguments: args})
|
||
if err != nil {
|
||
if ctx.Err() != nil {
|
||
return nil, ctx.Err()
|
||
}
|
||
message := err.Error()
|
||
for _, values := range []map[string]string{live.secrets.Env, live.secrets.Headers} {
|
||
for _, value := range values {
|
||
if value != "" {
|
||
message = strings.ReplaceAll(message, value, "[已隐藏凭据]")
|
||
}
|
||
}
|
||
}
|
||
return nil, fmt.Errorf("MCP 工具执行失败: %.1000s", message)
|
||
}
|
||
return result, nil
|
||
}
|