mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-08-01 00:27:35 +02:00
Add files via upload
This commit is contained in:
@@ -31,6 +31,11 @@ const (
|
||||
// 视觉分析(本地图片 → VL 模型 → 文本摘要)
|
||||
ToolAnalyzeImage = "analyze_image"
|
||||
|
||||
// 长耗时工具执行控制(后台 execution 查询/等待/取消)
|
||||
ToolGetToolExecution = "get_tool_execution"
|
||||
ToolWaitToolExecution = "wait_tool_execution"
|
||||
ToolCancelToolExecution = "cancel_tool_execution"
|
||||
|
||||
// WebShell 助手工具(AI 在 WebShell 管理 - AI 助手 中使用)
|
||||
ToolWebshellExec = "webshell_exec"
|
||||
ToolWebshellFileList = "webshell_file_list"
|
||||
@@ -91,6 +96,9 @@ func IsBuiltinTool(toolName string) bool {
|
||||
ToolListKnowledgeRiskTypes,
|
||||
ToolSearchKnowledgeBase,
|
||||
ToolAnalyzeImage,
|
||||
ToolGetToolExecution,
|
||||
ToolWaitToolExecution,
|
||||
ToolCancelToolExecution,
|
||||
ToolWebshellExec,
|
||||
ToolWebshellFileList,
|
||||
ToolWebshellFileRead,
|
||||
@@ -149,6 +157,9 @@ func GetAllBuiltinTools() []string {
|
||||
ToolListKnowledgeRiskTypes,
|
||||
ToolSearchKnowledgeBase,
|
||||
ToolAnalyzeImage,
|
||||
ToolGetToolExecution,
|
||||
ToolWaitToolExecution,
|
||||
ToolCancelToolExecution,
|
||||
ToolWebshellExec,
|
||||
ToolWebshellFileList,
|
||||
ToolWebshellFileRead,
|
||||
|
||||
@@ -0,0 +1,224 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/mcp/builtin"
|
||||
)
|
||||
|
||||
const (
|
||||
defaultExecutionWaitTimeout = 60 * time.Second
|
||||
maxExecutionWaitTimeout = 10 * time.Minute
|
||||
)
|
||||
|
||||
// RegisterExecutionControlTools exposes execution handle operations to Eino as
|
||||
// ordinary MCP tools. This keeps the agent loop native: the model calls a tool,
|
||||
// receives a bounded result, and may call wait_tool_execution again if needed.
|
||||
func RegisterExecutionControlTools(server *Server, external *ExternalMCPManager) {
|
||||
if server == nil {
|
||||
return
|
||||
}
|
||||
|
||||
server.RegisterTool(Tool{
|
||||
Name: builtin.ToolGetToolExecution,
|
||||
Description: "查询后台工具 execution 的当前状态、结果和错误。用于外部 MCP 工具等待超时后,凭 execution_id 继续查看进度。",
|
||||
ShortDescription: "查询后台工具执行状态",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"execution_id": map[string]interface{}{"type": "string", "description": "工具执行 ID"},
|
||||
},
|
||||
"required": []string{"execution_id"},
|
||||
},
|
||||
}, func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) {
|
||||
id := stringArg(args, "execution_id")
|
||||
if id == "" {
|
||||
return textToolResult("execution_id 必填", true), nil
|
||||
}
|
||||
exec := lookupToolExecution(server, external, id)
|
||||
if exec == nil {
|
||||
return textToolResult("未找到该 execution_id: "+id, true), nil
|
||||
}
|
||||
return textToolResult(formatExecutionForModel(exec), false), nil
|
||||
})
|
||||
|
||||
server.RegisterTool(Tool{
|
||||
Name: builtin.ToolWaitToolExecution,
|
||||
Description: "继续等待一个后台工具 execution 完成。每次等待都有 timeout_seconds 上限;若仍未完成,会返回当前状态,模型可稍后再次调用。",
|
||||
ShortDescription: "有界等待后台工具执行",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"execution_id": map[string]interface{}{"type": "string", "description": "工具执行 ID"},
|
||||
"timeout_seconds": map[string]interface{}{"type": "number", "description": "本次最多等待秒数,默认 60,最大 600"},
|
||||
},
|
||||
"required": []string{"execution_id"},
|
||||
},
|
||||
}, func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) {
|
||||
id := stringArg(args, "execution_id")
|
||||
if id == "" {
|
||||
return textToolResult("execution_id 必填", true), nil
|
||||
}
|
||||
wait := durationSecondsArg(args, "timeout_seconds", defaultExecutionWaitTimeout, maxExecutionWaitTimeout)
|
||||
snap, err := waitToolExecutionSnapshot(ctx, server, external, id, wait)
|
||||
if err != nil && !errors.Is(err, ErrExecutionWaitTimeout) {
|
||||
return textToolResult("等待 execution 失败: "+err.Error(), true), nil
|
||||
}
|
||||
if snap == nil || snap.Execution == nil {
|
||||
return textToolResult("未找到该 execution_id: "+id, true), nil
|
||||
}
|
||||
body := formatExecutionForModel(snap.Execution)
|
||||
if errors.Is(err, ErrExecutionWaitTimeout) {
|
||||
body += "\n\n本次等待已到达 timeout_seconds,上述 execution 仍未完成。可继续等待、取消,或采用其他步骤。"
|
||||
}
|
||||
return textToolResult(body, false), nil
|
||||
})
|
||||
|
||||
server.RegisterTool(Tool{
|
||||
Name: builtin.ToolCancelToolExecution,
|
||||
Description: "取消一个后台工具 execution。用于外部 MCP 工具长时间运行、误调用或用户要求停止时。",
|
||||
ShortDescription: "取消后台工具执行",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"execution_id": map[string]interface{}{"type": "string", "description": "工具执行 ID"},
|
||||
"reason": map[string]interface{}{"type": "string", "description": "取消原因,可选,会写入终止说明"},
|
||||
},
|
||||
"required": []string{"execution_id"},
|
||||
},
|
||||
}, func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) {
|
||||
id := stringArg(args, "execution_id")
|
||||
if id == "" {
|
||||
return textToolResult("execution_id 必填", true), nil
|
||||
}
|
||||
reason := stringArg(args, "reason")
|
||||
if server.CancelToolExecutionWithNote(id, reason) {
|
||||
return textToolResult("已请求取消内部工具 execution: "+id, false), nil
|
||||
}
|
||||
if external != nil && external.CancelToolExecutionWithNote(id, reason) {
|
||||
return textToolResult("已请求取消外部 MCP execution: "+id, false), nil
|
||||
}
|
||||
return textToolResult("未找到进行中的 execution,或该 execution 已结束: "+id, true), nil
|
||||
})
|
||||
}
|
||||
|
||||
func waitToolExecutionSnapshot(ctx context.Context, server *Server, external *ExternalMCPManager, id string, wait time.Duration) (*ExecutionSnapshot, error) {
|
||||
if server != nil && server.executionService != nil && server.executionService.getEntry(id) != nil {
|
||||
return server.executionService.Wait(ctx, id, wait)
|
||||
}
|
||||
if external != nil && external.executionService != nil && external.executionService.getEntry(id) != nil {
|
||||
return external.executionService.Wait(ctx, id, wait)
|
||||
}
|
||||
if server != nil && server.executionService != nil {
|
||||
if snap, err := server.executionService.Get(id); err == nil {
|
||||
return snap, nil
|
||||
}
|
||||
}
|
||||
if external != nil && external.executionService != nil {
|
||||
return external.executionService.Get(id)
|
||||
}
|
||||
exec := lookupToolExecution(server, external, id)
|
||||
if exec == nil {
|
||||
return nil, fmt.Errorf("execution not found: %s", id)
|
||||
}
|
||||
return &ExecutionSnapshot{Execution: exec}, nil
|
||||
}
|
||||
|
||||
func lookupToolExecution(server *Server, external *ExternalMCPManager, id string) *ToolExecution {
|
||||
if server != nil {
|
||||
if exec, ok := server.GetExecution(id); ok && exec != nil {
|
||||
return exec
|
||||
}
|
||||
}
|
||||
if external != nil {
|
||||
if exec, ok := external.GetExecution(id); ok && exec != nil {
|
||||
return exec
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func formatExecutionForModel(exec *ToolExecution) string {
|
||||
if exec == nil {
|
||||
return "execution: null"
|
||||
}
|
||||
payload := map[string]interface{}{
|
||||
"execution_id": exec.ID,
|
||||
"tool": exec.ToolName,
|
||||
"status": exec.Status,
|
||||
"started_at": exec.StartTime.Format(time.RFC3339),
|
||||
}
|
||||
if exec.EndTime != nil {
|
||||
payload["ended_at"] = exec.EndTime.Format(time.RFC3339)
|
||||
}
|
||||
if exec.Duration > 0 {
|
||||
payload["duration"] = exec.Duration.String()
|
||||
}
|
||||
if exec.Error != "" {
|
||||
payload["error"] = exec.Error
|
||||
}
|
||||
if exec.Result != nil {
|
||||
payload["result"] = ToolResultPlainText(exec.Result)
|
||||
payload["is_error"] = exec.Result.IsError
|
||||
}
|
||||
b, err := json.MarshalIndent(payload, "", " ")
|
||||
if err != nil {
|
||||
return fmt.Sprintf("execution_id: %s\nstatus: %s\nerror: %s", exec.ID, exec.Status, exec.Error)
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
|
||||
func textToolResult(text string, isErr bool) *ToolResult {
|
||||
return &ToolResult{Content: []Content{{Type: "text", Text: text}}, IsError: isErr}
|
||||
}
|
||||
|
||||
func stringArg(args map[string]interface{}, key string) string {
|
||||
if args == nil {
|
||||
return ""
|
||||
}
|
||||
raw, ok := args[key]
|
||||
if !ok || raw == nil {
|
||||
return ""
|
||||
}
|
||||
switch v := raw.(type) {
|
||||
case string:
|
||||
return strings.TrimSpace(v)
|
||||
default:
|
||||
return strings.TrimSpace(fmt.Sprint(v))
|
||||
}
|
||||
}
|
||||
|
||||
func durationSecondsArg(args map[string]interface{}, key string, def, max time.Duration) time.Duration {
|
||||
if args == nil {
|
||||
return def
|
||||
}
|
||||
var seconds float64
|
||||
switch v := args[key].(type) {
|
||||
case int:
|
||||
seconds = float64(v)
|
||||
case int64:
|
||||
seconds = float64(v)
|
||||
case float64:
|
||||
seconds = v
|
||||
case json.Number:
|
||||
f, _ := v.Float64()
|
||||
seconds = f
|
||||
case string:
|
||||
f, _ := strconv.ParseFloat(strings.TrimSpace(v), 64)
|
||||
seconds = f
|
||||
}
|
||||
if seconds <= 0 {
|
||||
return def
|
||||
}
|
||||
d := time.Duration(seconds * float64(time.Second))
|
||||
if max > 0 && d > max {
|
||||
return max
|
||||
}
|
||||
return d
|
||||
}
|
||||
@@ -0,0 +1,557 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/authctx"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
const (
|
||||
ToolExecutionStatusQueued = "queued"
|
||||
ToolExecutionStatusRunning = "running"
|
||||
ToolExecutionStatusCompleted = "completed"
|
||||
ToolExecutionStatusFailed = "failed"
|
||||
ToolExecutionStatusCancelled = "cancelled"
|
||||
ToolExecutionStatusHardTimeout = "hard_timeout"
|
||||
ToolExecutionStatusOrphaned = "orphaned"
|
||||
)
|
||||
|
||||
var ErrExecutionWaitTimeout = errors.New("tool execution wait timeout")
|
||||
|
||||
// ExecutionRunFunc is the blocking operation owned by a worker.
|
||||
type ExecutionRunFunc func(context.Context) (*ToolResult, error)
|
||||
|
||||
type ExecutionPreRunFunc func(context.Context, *ToolExecution) (func(), error)
|
||||
|
||||
// ExecutionDoneFunc observes the final persisted state. It is invoked once,
|
||||
// including for late completions after an agent has stopped waiting.
|
||||
type ExecutionDoneFunc func(*ToolExecution)
|
||||
|
||||
type ExecutionRequest struct {
|
||||
ID string
|
||||
ToolName string
|
||||
Arguments map[string]interface{}
|
||||
ConversationID string
|
||||
OwnerUserID string
|
||||
HardTimeout time.Duration
|
||||
PreRun ExecutionPreRunFunc
|
||||
Run ExecutionRunFunc
|
||||
OnDone ExecutionDoneFunc
|
||||
}
|
||||
|
||||
type ExecutionHandle struct {
|
||||
ID string
|
||||
}
|
||||
|
||||
type ExecutionSnapshot struct {
|
||||
Execution *ToolExecution
|
||||
}
|
||||
|
||||
type executionEntry struct {
|
||||
exec *ToolExecution
|
||||
cancel context.CancelFunc
|
||||
done chan struct{}
|
||||
preRun ExecutionPreRunFunc
|
||||
run ExecutionRunFunc
|
||||
result *ToolResult
|
||||
err error
|
||||
}
|
||||
|
||||
// ExecutionService keeps Eino-facing tool calls synchronous while moving the
|
||||
// untrusted blocking work into cancellable workers with explicit execution IDs.
|
||||
type ExecutionService struct {
|
||||
storage MonitorStorage
|
||||
logger *zap.Logger
|
||||
|
||||
mu sync.Mutex
|
||||
entries map[string]*executionEntry
|
||||
abortUserNotes map[string]string
|
||||
maxInMemory int
|
||||
resultMaxBytes int
|
||||
}
|
||||
|
||||
func NewExecutionService(storage MonitorStorage, logger *zap.Logger) *ExecutionService {
|
||||
if logger == nil {
|
||||
logger = zap.NewNop()
|
||||
}
|
||||
return &ExecutionService{
|
||||
storage: storage,
|
||||
logger: logger,
|
||||
entries: make(map[string]*executionEntry),
|
||||
abortUserNotes: make(map[string]string),
|
||||
maxInMemory: 1000,
|
||||
resultMaxBytes: DefaultToolResultMaxBytes,
|
||||
}
|
||||
}
|
||||
|
||||
func (s *ExecutionService) ConfigureToolResultMaxBytes(maxBytes int) {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.resultMaxBytes = maxBytes
|
||||
}
|
||||
|
||||
func (s *ExecutionService) Submit(ctx context.Context, req ExecutionRequest) (*ExecutionHandle, error) {
|
||||
if s == nil {
|
||||
return nil, fmt.Errorf("execution service is nil")
|
||||
}
|
||||
if req.Run == nil {
|
||||
return nil, fmt.Errorf("execution run func is nil")
|
||||
}
|
||||
id := strings.TrimSpace(req.ID)
|
||||
if id == "" {
|
||||
id = uuid.New().String()
|
||||
}
|
||||
start := time.Now()
|
||||
exec := &ToolExecution{
|
||||
ID: id,
|
||||
ToolName: strings.TrimSpace(req.ToolName),
|
||||
Arguments: cloneArgsMap(req.Arguments),
|
||||
Status: ToolExecutionStatusQueued,
|
||||
StartTime: start,
|
||||
ConversationID: strings.TrimSpace(req.ConversationID),
|
||||
OwnerUserID: strings.TrimSpace(req.OwnerUserID),
|
||||
}
|
||||
if exec.ConversationID == "" {
|
||||
exec.ConversationID = MCPConversationIDFromContext(ctx)
|
||||
}
|
||||
if exec.OwnerUserID == "" {
|
||||
if principal, ok := authctx.PrincipalFromContext(ctx); ok {
|
||||
exec.OwnerUserID = principal.UserID
|
||||
}
|
||||
}
|
||||
|
||||
runCtx := detachedExecutionContext(ctx)
|
||||
var cancel context.CancelFunc
|
||||
if req.HardTimeout > 0 {
|
||||
runCtx, cancel = context.WithTimeout(runCtx, req.HardTimeout)
|
||||
} else {
|
||||
runCtx, cancel = context.WithCancel(runCtx)
|
||||
}
|
||||
entry := &executionEntry{exec: exec, cancel: cancel, done: make(chan struct{}), preRun: req.PreRun, run: req.Run}
|
||||
|
||||
s.mu.Lock()
|
||||
if _, exists := s.entries[id]; exists {
|
||||
s.mu.Unlock()
|
||||
cancel()
|
||||
return nil, fmt.Errorf("execution already exists: %s", id)
|
||||
}
|
||||
s.entries[id] = entry
|
||||
s.cleanupOldEntriesLocked()
|
||||
s.mu.Unlock()
|
||||
|
||||
if s.storage != nil {
|
||||
if err := s.storage.SaveToolExecution(exec); err != nil {
|
||||
s.logger.Warn("保存执行记录到数据库失败", zap.Error(err), zap.String("executionId", id))
|
||||
}
|
||||
}
|
||||
notifyToolRunBegin(ctx, id)
|
||||
|
||||
go s.runWorker(runCtx, entry, req.OnDone)
|
||||
return &ExecutionHandle{ID: id}, nil
|
||||
}
|
||||
|
||||
func (s *ExecutionService) runWorker(ctx context.Context, entry *executionEntry, onDone ExecutionDoneFunc) {
|
||||
id := entry.exec.ID
|
||||
var release func()
|
||||
defer func() {
|
||||
if release != nil {
|
||||
release()
|
||||
}
|
||||
entry.cancel()
|
||||
notifyToolRunEnd(ctx, id)
|
||||
close(entry.done)
|
||||
}()
|
||||
|
||||
if entry.preRun != nil {
|
||||
var preErr error
|
||||
release, preErr = entry.preRun(ctx, cloneToolExecution(entry.exec))
|
||||
if preErr != nil {
|
||||
s.finishEntry(ctx, entry, nil, preErr, onDone)
|
||||
return
|
||||
}
|
||||
}
|
||||
s.markEntryRunning(entry)
|
||||
|
||||
result, err := entryResultRecover(ctx, entry.exec.ToolName, s.logger, func() (*ToolResult, error) {
|
||||
return nilSafeRun(ctx, entry)
|
||||
})
|
||||
s.finishEntry(ctx, entry, result, err, onDone)
|
||||
}
|
||||
|
||||
func (s *ExecutionService) markEntryRunning(entry *executionEntry) {
|
||||
if s == nil || entry == nil || entry.exec == nil {
|
||||
return
|
||||
}
|
||||
s.mu.Lock()
|
||||
if !isExecutionTerminal(entry.exec.Status) {
|
||||
entry.exec.Status = ToolExecutionStatusRunning
|
||||
}
|
||||
runningExec := cloneToolExecution(entry.exec)
|
||||
s.mu.Unlock()
|
||||
if s.storage != nil {
|
||||
if err := s.storage.SaveToolExecution(runningExec); err != nil {
|
||||
s.logger.Warn("保存执行记录到数据库失败", zap.Error(err), zap.String("executionId", runningExec.ID))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *ExecutionService) finishEntry(ctx context.Context, entry *executionEntry, result *ToolResult, err error, onDone ExecutionDoneFunc) {
|
||||
id := entry.exec.ID
|
||||
cancelledWithUserNote := s.applyAbortUserNoteToCancelledToolResult(id, &result, &err)
|
||||
|
||||
now := time.Now()
|
||||
s.mu.Lock()
|
||||
result = NormalizeToolResultForStorage(result, s.resultMaxBytes)
|
||||
entry.result = result
|
||||
entry.err = err
|
||||
entry.exec.EndTime = &now
|
||||
entry.exec.Duration = now.Sub(entry.exec.StartTime)
|
||||
if err != nil {
|
||||
switch {
|
||||
case errors.Is(err, context.DeadlineExceeded):
|
||||
entry.exec.Status = ToolExecutionStatusHardTimeout
|
||||
entry.exec.Error = "工具执行超过硬超时限制"
|
||||
case errors.Is(err, context.Canceled):
|
||||
entry.exec.Status = ToolExecutionStatusCancelled
|
||||
entry.exec.Error = "已手动终止或任务已取消"
|
||||
default:
|
||||
entry.exec.Status = ToolExecutionStatusFailed
|
||||
entry.exec.Error = err.Error()
|
||||
}
|
||||
} else if result != nil && result.IsError {
|
||||
if cancelledWithUserNote {
|
||||
entry.exec.Status = ToolExecutionStatusCancelled
|
||||
entry.exec.Error = ""
|
||||
} else if isBackgroundWaitToolResult(result) {
|
||||
entry.exec.Status = ToolExecutionStatusCompleted
|
||||
entry.exec.Error = ""
|
||||
} else {
|
||||
entry.exec.Status = ToolExecutionStatusFailed
|
||||
entry.exec.Error = firstToolResultText(result, "工具执行返回错误结果")
|
||||
}
|
||||
entry.exec.Result = result
|
||||
} else {
|
||||
entry.exec.Status = ToolExecutionStatusCompleted
|
||||
if result == nil {
|
||||
result = &ToolResult{Content: []Content{{Type: "text", Text: "工具执行完成,但未返回结果"}}}
|
||||
entry.result = result
|
||||
}
|
||||
entry.exec.Result = result
|
||||
}
|
||||
finalExec := cloneToolExecution(entry.exec)
|
||||
s.mu.Unlock()
|
||||
|
||||
if s.storage != nil {
|
||||
if saveErr := s.storage.SaveToolExecution(finalExec); saveErr != nil {
|
||||
s.logger.Warn("保存执行记录到数据库失败", zap.Error(saveErr), zap.String("executionId", id))
|
||||
}
|
||||
}
|
||||
if onDone != nil {
|
||||
onDone(finalExec)
|
||||
}
|
||||
}
|
||||
|
||||
func nilSafeRun(ctx context.Context, entry *executionEntry) (*ToolResult, error) {
|
||||
if entry == nil {
|
||||
return nil, fmt.Errorf("execution entry is nil")
|
||||
}
|
||||
if entry.run == nil {
|
||||
return nil, fmt.Errorf("execution run func not wired")
|
||||
}
|
||||
return entry.run(ctx)
|
||||
}
|
||||
|
||||
func entryResultRecover(ctx context.Context, toolName string, logger *zap.Logger, fn func() (*ToolResult, error)) (res *ToolResult, err error) {
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
if logger != nil {
|
||||
logger.Error("tool execution worker panic recovered", zap.Any("recover", r), zap.String("toolName", toolName), zap.Stack("stack"))
|
||||
}
|
||||
err = fmt.Errorf("tool execution panic: %v", r)
|
||||
}
|
||||
}()
|
||||
return fn()
|
||||
}
|
||||
|
||||
func (s *ExecutionService) Wait(ctx context.Context, executionID string, timeout time.Duration) (*ExecutionSnapshot, error) {
|
||||
entry := s.getEntry(executionID)
|
||||
if entry == nil {
|
||||
return s.getPersistedSnapshot(executionID)
|
||||
}
|
||||
if isExecutionTerminal(entry.exec.Status) {
|
||||
return &ExecutionSnapshot{Execution: cloneToolExecution(entry.exec)}, nil
|
||||
}
|
||||
|
||||
var timeoutCh <-chan time.Time
|
||||
var timer *time.Timer
|
||||
if timeout > 0 {
|
||||
timer = time.NewTimer(timeout)
|
||||
timeoutCh = timer.C
|
||||
defer timer.Stop()
|
||||
}
|
||||
|
||||
select {
|
||||
case <-entry.done:
|
||||
return &ExecutionSnapshot{Execution: cloneToolExecution(entry.exec)}, nil
|
||||
case <-timeoutCh:
|
||||
return &ExecutionSnapshot{Execution: cloneToolExecution(entry.exec)}, ErrExecutionWaitTimeout
|
||||
case <-ctxDone(ctx):
|
||||
return &ExecutionSnapshot{Execution: cloneToolExecution(entry.exec)}, ctx.Err()
|
||||
}
|
||||
}
|
||||
|
||||
func (s *ExecutionService) Get(executionID string) (*ExecutionSnapshot, error) {
|
||||
entry := s.getEntry(executionID)
|
||||
if entry != nil {
|
||||
return &ExecutionSnapshot{Execution: cloneToolExecution(entry.exec)}, nil
|
||||
}
|
||||
return s.getPersistedSnapshot(executionID)
|
||||
}
|
||||
|
||||
func (s *ExecutionService) Cancel(executionID, note string) bool {
|
||||
id := strings.TrimSpace(executionID)
|
||||
if id == "" || s == nil {
|
||||
return false
|
||||
}
|
||||
s.mu.Lock()
|
||||
entry := s.entries[id]
|
||||
if entry == nil || isExecutionTerminal(entry.exec.Status) {
|
||||
s.mu.Unlock()
|
||||
return false
|
||||
}
|
||||
if strings.TrimSpace(note) != "" {
|
||||
s.abortUserNotes[id] = strings.TrimSpace(note)
|
||||
}
|
||||
cancel := entry.cancel
|
||||
s.mu.Unlock()
|
||||
if cancel != nil {
|
||||
cancel()
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func (s *ExecutionService) ActiveRunningExecutionIDs() map[string]struct{} {
|
||||
if s == nil {
|
||||
return nil
|
||||
}
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
out := make(map[string]struct{})
|
||||
for id, entry := range s.entries {
|
||||
if entry != nil && entry.exec != nil && !isExecutionTerminal(entry.exec.Status) {
|
||||
out[id] = struct{}{}
|
||||
}
|
||||
}
|
||||
if len(out) == 0 {
|
||||
return nil
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (s *ExecutionService) CancelAll(note string) {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
s.mu.Lock()
|
||||
cancels := make([]context.CancelFunc, 0, len(s.entries))
|
||||
for id, entry := range s.entries {
|
||||
if entry == nil || isExecutionTerminal(entry.exec.Status) {
|
||||
continue
|
||||
}
|
||||
if strings.TrimSpace(note) != "" {
|
||||
s.abortUserNotes[id] = strings.TrimSpace(note)
|
||||
}
|
||||
if entry.cancel != nil {
|
||||
cancels = append(cancels, entry.cancel)
|
||||
}
|
||||
}
|
||||
s.mu.Unlock()
|
||||
for _, cancel := range cancels {
|
||||
cancel()
|
||||
}
|
||||
}
|
||||
|
||||
func (s *ExecutionService) getEntry(executionID string) *executionEntry {
|
||||
if s == nil {
|
||||
return nil
|
||||
}
|
||||
id := strings.TrimSpace(executionID)
|
||||
if id == "" {
|
||||
return nil
|
||||
}
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
return s.entries[id]
|
||||
}
|
||||
|
||||
func (s *ExecutionService) getPersistedSnapshot(executionID string) (*ExecutionSnapshot, error) {
|
||||
id := strings.TrimSpace(executionID)
|
||||
if id == "" {
|
||||
return nil, fmt.Errorf("execution_id is required")
|
||||
}
|
||||
if s != nil && s.storage != nil {
|
||||
exec, err := s.storage.GetToolExecution(id)
|
||||
if err == nil && exec != nil {
|
||||
return &ExecutionSnapshot{Execution: exec}, nil
|
||||
}
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
return nil, fmt.Errorf("execution not found: %s", id)
|
||||
}
|
||||
|
||||
func (s *ExecutionService) applyAbortUserNoteToCancelledToolResult(executionID string, result **ToolResult, err *error) (cancelledWithUserNote bool) {
|
||||
note := strings.TrimSpace(s.takeAbortUserNote(executionID))
|
||||
if note == "" {
|
||||
return false
|
||||
}
|
||||
hasErr := err != nil && *err != nil
|
||||
hasRes := result != nil && *result != nil
|
||||
if !hasErr && !hasRes {
|
||||
return false
|
||||
}
|
||||
partial := ""
|
||||
if hasRes {
|
||||
partial = ToolResultPlainText(*result)
|
||||
}
|
||||
if partial == "" && hasErr {
|
||||
partial = (*err).Error()
|
||||
}
|
||||
merged := MergePartialToolOutputAndAbortNote(partial, note)
|
||||
if err != nil {
|
||||
*err = nil
|
||||
}
|
||||
if result != nil {
|
||||
*result = &ToolResult{Content: []Content{{Type: "text", Text: merged}}, IsError: true}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func (s *ExecutionService) takeAbortUserNote(id string) string {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
note := s.abortUserNotes[id]
|
||||
delete(s.abortUserNotes, id)
|
||||
return note
|
||||
}
|
||||
|
||||
func (s *ExecutionService) cleanupOldEntriesLocked() {
|
||||
if s.maxInMemory <= 0 || len(s.entries) <= s.maxInMemory {
|
||||
return
|
||||
}
|
||||
type oldEntry struct {
|
||||
id string
|
||||
startTime time.Time
|
||||
}
|
||||
var terminal []oldEntry
|
||||
for id, entry := range s.entries {
|
||||
if entry != nil && entry.exec != nil && isExecutionTerminal(entry.exec.Status) {
|
||||
terminal = append(terminal, oldEntry{id: id, startTime: entry.exec.StartTime})
|
||||
}
|
||||
}
|
||||
for len(s.entries) > s.maxInMemory && len(terminal) > 0 {
|
||||
oldest := 0
|
||||
for i := 1; i < len(terminal); i++ {
|
||||
if terminal[i].startTime.Before(terminal[oldest].startTime) {
|
||||
oldest = i
|
||||
}
|
||||
}
|
||||
delete(s.entries, terminal[oldest].id)
|
||||
terminal = append(terminal[:oldest], terminal[oldest+1:]...)
|
||||
}
|
||||
}
|
||||
|
||||
func firstToolResultText(result *ToolResult, fallback string) string {
|
||||
if result != nil {
|
||||
for _, c := range result.Content {
|
||||
if strings.TrimSpace(c.Text) != "" {
|
||||
return c.Text
|
||||
}
|
||||
}
|
||||
}
|
||||
return fallback
|
||||
}
|
||||
|
||||
func isBackgroundWaitToolResult(result *ToolResult) bool {
|
||||
text := strings.ToLower(strings.TrimSpace(ToolResultPlainText(result)))
|
||||
if text == "" {
|
||||
return false
|
||||
}
|
||||
hasExecutionID := strings.Contains(text, "execution_id:") || strings.Contains(text, `"execution_id"`)
|
||||
hasRunningStatus := strings.Contains(text, "status: running") || strings.Contains(text, "status: queued") ||
|
||||
strings.Contains(text, `"status": "running"`) || strings.Contains(text, `"status":"running"`) ||
|
||||
strings.Contains(text, `"status": "queued"`) || strings.Contains(text, `"status":"queued"`)
|
||||
hasSoftWaitSignal := strings.Contains(text, "工具已提交到后台执行") ||
|
||||
strings.Contains(text, "本次等待已到达") ||
|
||||
strings.Contains(text, "wait_timeout:") ||
|
||||
strings.Contains(text, "background execution") ||
|
||||
strings.Contains(text, "still running") ||
|
||||
strings.Contains(text, "仍未完成")
|
||||
return hasExecutionID && hasRunningStatus && hasSoftWaitSignal
|
||||
}
|
||||
|
||||
func isExecutionTerminal(status string) bool {
|
||||
switch strings.TrimSpace(strings.ToLower(status)) {
|
||||
case ToolExecutionStatusCompleted, ToolExecutionStatusFailed, ToolExecutionStatusCancelled, ToolExecutionStatusHardTimeout, ToolExecutionStatusOrphaned:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func ctxDone(ctx context.Context) <-chan struct{} {
|
||||
if ctx == nil {
|
||||
return nil
|
||||
}
|
||||
return ctx.Done()
|
||||
}
|
||||
|
||||
func detachedExecutionContext(ctx context.Context) context.Context {
|
||||
if ctx == nil {
|
||||
return context.Background()
|
||||
}
|
||||
return context.WithoutCancel(ctx)
|
||||
}
|
||||
|
||||
func cloneArgsMap(in map[string]interface{}) map[string]interface{} {
|
||||
if in == nil {
|
||||
return map[string]interface{}{}
|
||||
}
|
||||
out := make(map[string]interface{}, len(in))
|
||||
for k, v := range in {
|
||||
out[k] = v
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func cloneToolExecution(in *ToolExecution) *ToolExecution {
|
||||
if in == nil {
|
||||
return nil
|
||||
}
|
||||
out := *in
|
||||
out.Arguments = cloneArgsMap(in.Arguments)
|
||||
if in.Result != nil {
|
||||
res := *in.Result
|
||||
if in.Result.Content != nil {
|
||||
res.Content = append([]Content(nil), in.Result.Content...)
|
||||
}
|
||||
out.Result = &res
|
||||
}
|
||||
if in.EndTime != nil {
|
||||
t := *in.EndTime
|
||||
out.EndTime = &t
|
||||
}
|
||||
return &out
|
||||
}
|
||||
@@ -0,0 +1,41 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestExecutionServiceBackgroundWaitResultCompletesWaitTool(t *testing.T) {
|
||||
service := NewExecutionService(nil, nil)
|
||||
handle, err := service.Submit(context.Background(), ExecutionRequest{
|
||||
ToolName: "wait_tool_execution",
|
||||
Run: func(context.Context) (*ToolResult, error) {
|
||||
return &ToolResult{
|
||||
Content: []Content{{Type: "text", Text: `{
|
||||
"execution_id": "3eaaa391-050b-4be1-a870-48a855923cb7",
|
||||
"tool": "exec",
|
||||
"status": "running"
|
||||
}
|
||||
|
||||
本次等待已到达 timeout_seconds,上述 execution 仍未完成。可继续等待、取消,或采用其他步骤。`}},
|
||||
IsError: true,
|
||||
}, nil
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Submit: %v", err)
|
||||
}
|
||||
snap, err := service.Wait(context.Background(), handle.ID, 0)
|
||||
if err != nil {
|
||||
t.Fatalf("Wait: %v", err)
|
||||
}
|
||||
if snap == nil || snap.Execution == nil {
|
||||
t.Fatal("missing execution snapshot")
|
||||
}
|
||||
if snap.Execution.Status != ToolExecutionStatusCompleted {
|
||||
t.Fatalf("status = %q, want %q", snap.Execution.Status, ToolExecutionStatusCompleted)
|
||||
}
|
||||
if snap.Execution.Result == nil || !snap.Execution.Result.IsError {
|
||||
t.Fatal("model-facing result should remain IsError")
|
||||
}
|
||||
}
|
||||
+359
-130
@@ -2,6 +2,7 @@ package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync"
|
||||
@@ -11,8 +12,6 @@ import (
|
||||
"cyberstrike-ai/internal/authctx"
|
||||
"cyberstrike-ai/internal/config"
|
||||
|
||||
"github.com/google/uuid"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
@@ -36,32 +35,51 @@ type listToolsInflight struct {
|
||||
err error
|
||||
}
|
||||
|
||||
type ExternalMCPResilienceConfig struct {
|
||||
MaxConcurrentPerServer int
|
||||
MaxConcurrentTotal int
|
||||
CircuitFailureThreshold int
|
||||
CircuitCooldown time.Duration
|
||||
}
|
||||
|
||||
type externalMCPServerRuntime struct {
|
||||
semaphore chan struct{}
|
||||
consecutiveFailures int
|
||||
circuitOpenUntil time.Time
|
||||
}
|
||||
|
||||
// ExternalMCPManager 外部MCP管理器
|
||||
type ExternalMCPManager struct {
|
||||
clients map[string]ExternalMCPClient
|
||||
configs map[string]config.ExternalMCPServerConfig
|
||||
logger *zap.Logger
|
||||
storage MonitorStorage // 可选的持久化存储
|
||||
executions map[string]*ToolExecution // 执行记录
|
||||
stats map[string]*ToolStats // 工具统计信息
|
||||
errors map[string]string // 错误信息
|
||||
toolCounts map[string]int // 工具数量缓存
|
||||
toolCountsMu sync.RWMutex // 工具数量缓存的锁
|
||||
toolCache map[string]toolListCacheEntry // 工具列表缓存:MCP名称 -> 工具列表
|
||||
toolCacheMu sync.RWMutex // 工具列表缓存的锁
|
||||
listToolsMu sync.Mutex
|
||||
listToolsInflight map[string]*listToolsInflight
|
||||
stopRefresh chan struct{} // 停止后台刷新的信号
|
||||
refreshWg sync.WaitGroup // 等待后台刷新goroutine完成
|
||||
refreshing atomic.Bool // 防止 refreshToolCounts 并发堆积
|
||||
mu sync.RWMutex
|
||||
runningCancels map[string]context.CancelFunc
|
||||
abortUserNotes map[string]string
|
||||
reconnectMu sync.Mutex
|
||||
reconnecting map[string]bool
|
||||
reconnectLastTry map[string]time.Time
|
||||
reconnectAttempts map[string]int
|
||||
toolAuthorizer func(context.Context, string, map[string]interface{}) error
|
||||
clients map[string]ExternalMCPClient
|
||||
configs map[string]config.ExternalMCPServerConfig
|
||||
logger *zap.Logger
|
||||
storage MonitorStorage // 可选的持久化存储
|
||||
executions map[string]*ToolExecution // 执行记录
|
||||
stats map[string]*ToolStats // 工具统计信息
|
||||
errors map[string]string // 错误信息
|
||||
toolCounts map[string]int // 工具数量缓存
|
||||
toolCountsMu sync.RWMutex // 工具数量缓存的锁
|
||||
toolCache map[string]toolListCacheEntry // 工具列表缓存:MCP名称 -> 工具列表
|
||||
toolCacheMu sync.RWMutex // 工具列表缓存的锁
|
||||
listToolsMu sync.Mutex
|
||||
listToolsInflight map[string]*listToolsInflight
|
||||
stopRefresh chan struct{} // 停止后台刷新的信号
|
||||
refreshWg sync.WaitGroup // 等待后台刷新goroutine完成
|
||||
refreshing atomic.Bool // 防止 refreshToolCounts 并发堆积
|
||||
mu sync.RWMutex
|
||||
runningCancels map[string]context.CancelFunc
|
||||
abortUserNotes map[string]string
|
||||
reconnectMu sync.Mutex
|
||||
reconnecting map[string]bool
|
||||
reconnectLastTry map[string]time.Time
|
||||
reconnectAttempts map[string]int
|
||||
toolAuthorizer func(context.Context, string, map[string]interface{}) error
|
||||
executionService *ExecutionService
|
||||
toolWaitTimeout time.Duration
|
||||
toolResultMaxBytes int
|
||||
resilience ExternalMCPResilienceConfig
|
||||
serverRuntimes map[string]*externalMCPServerRuntime
|
||||
globalSemaphore chan struct{}
|
||||
}
|
||||
|
||||
// NewExternalMCPManager 创建外部MCP管理器
|
||||
@@ -80,28 +98,105 @@ func (m *ExternalMCPManager) SetToolAuthorizer(authorizer func(context.Context,
|
||||
// NewExternalMCPManagerWithStorage 创建外部MCP管理器(带持久化存储)
|
||||
func NewExternalMCPManagerWithStorage(logger *zap.Logger, storage MonitorStorage) *ExternalMCPManager {
|
||||
manager := &ExternalMCPManager{
|
||||
clients: make(map[string]ExternalMCPClient),
|
||||
configs: make(map[string]config.ExternalMCPServerConfig),
|
||||
logger: logger,
|
||||
storage: storage,
|
||||
executions: make(map[string]*ToolExecution),
|
||||
stats: make(map[string]*ToolStats),
|
||||
errors: make(map[string]string),
|
||||
toolCounts: make(map[string]int),
|
||||
toolCache: make(map[string]toolListCacheEntry),
|
||||
listToolsInflight: make(map[string]*listToolsInflight),
|
||||
stopRefresh: make(chan struct{}),
|
||||
runningCancels: make(map[string]context.CancelFunc),
|
||||
abortUserNotes: make(map[string]string),
|
||||
reconnecting: make(map[string]bool),
|
||||
reconnectLastTry: make(map[string]time.Time),
|
||||
reconnectAttempts: make(map[string]int),
|
||||
clients: make(map[string]ExternalMCPClient),
|
||||
configs: make(map[string]config.ExternalMCPServerConfig),
|
||||
logger: logger,
|
||||
storage: storage,
|
||||
executions: make(map[string]*ToolExecution),
|
||||
stats: make(map[string]*ToolStats),
|
||||
errors: make(map[string]string),
|
||||
toolCounts: make(map[string]int),
|
||||
toolCache: make(map[string]toolListCacheEntry),
|
||||
listToolsInflight: make(map[string]*listToolsInflight),
|
||||
stopRefresh: make(chan struct{}),
|
||||
runningCancels: make(map[string]context.CancelFunc),
|
||||
abortUserNotes: make(map[string]string),
|
||||
reconnecting: make(map[string]bool),
|
||||
reconnectLastTry: make(map[string]time.Time),
|
||||
reconnectAttempts: make(map[string]int),
|
||||
toolWaitTimeout: 60 * time.Second,
|
||||
toolResultMaxBytes: DefaultToolResultMaxBytes,
|
||||
resilience: ExternalMCPResilienceConfig{
|
||||
MaxConcurrentPerServer: 2,
|
||||
MaxConcurrentTotal: 16,
|
||||
CircuitFailureThreshold: 3,
|
||||
CircuitCooldown: 60 * time.Second,
|
||||
},
|
||||
serverRuntimes: make(map[string]*externalMCPServerRuntime),
|
||||
globalSemaphore: make(chan struct{}, 16),
|
||||
}
|
||||
manager.executionService = NewExecutionService(storage, logger)
|
||||
// 启动后台刷新工具数量的goroutine
|
||||
manager.startToolCountRefresh()
|
||||
return manager
|
||||
}
|
||||
|
||||
func (m *ExternalMCPManager) ConfigureToolResultMaxBytes(maxBytes int) {
|
||||
if m == nil {
|
||||
return
|
||||
}
|
||||
m.mu.Lock()
|
||||
m.toolResultMaxBytes = maxBytes
|
||||
m.mu.Unlock()
|
||||
if m.executionService != nil {
|
||||
m.executionService.ConfigureToolResultMaxBytes(maxBytes)
|
||||
}
|
||||
}
|
||||
|
||||
// ConfigureToolWaitTimeoutSeconds controls how long an agent-facing tool call
|
||||
// waits for an external MCP execution before returning an execution_id that can
|
||||
// be polled with wait_tool_execution. seconds<=0 waits until completion.
|
||||
func (m *ExternalMCPManager) ConfigureToolWaitTimeoutSeconds(seconds int) {
|
||||
if m == nil {
|
||||
return
|
||||
}
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
if seconds <= 0 {
|
||||
m.toolWaitTimeout = 0
|
||||
return
|
||||
}
|
||||
m.toolWaitTimeout = time.Duration(seconds) * time.Second
|
||||
}
|
||||
|
||||
func (m *ExternalMCPManager) ConfigureResilience(cfg ExternalMCPResilienceConfig) {
|
||||
if m == nil {
|
||||
return
|
||||
}
|
||||
normalized := normalizeExternalMCPResilienceConfig(cfg)
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
m.resilience = normalized
|
||||
m.serverRuntimes = make(map[string]*externalMCPServerRuntime)
|
||||
if normalized.MaxConcurrentTotal > 0 {
|
||||
m.globalSemaphore = make(chan struct{}, normalized.MaxConcurrentTotal)
|
||||
} else {
|
||||
m.globalSemaphore = nil
|
||||
}
|
||||
}
|
||||
|
||||
func normalizeExternalMCPResilienceConfig(cfg ExternalMCPResilienceConfig) ExternalMCPResilienceConfig {
|
||||
if cfg.MaxConcurrentPerServer == 0 {
|
||||
cfg.MaxConcurrentPerServer = 2
|
||||
}
|
||||
if cfg.MaxConcurrentTotal == 0 {
|
||||
cfg.MaxConcurrentTotal = 16
|
||||
}
|
||||
if cfg.CircuitFailureThreshold == 0 {
|
||||
cfg.CircuitFailureThreshold = 3
|
||||
}
|
||||
if cfg.CircuitCooldown <= 0 {
|
||||
cfg.CircuitCooldown = 60 * time.Second
|
||||
}
|
||||
if cfg.MaxConcurrentPerServer < 0 {
|
||||
cfg.MaxConcurrentPerServer = 0
|
||||
}
|
||||
if cfg.MaxConcurrentTotal < 0 {
|
||||
cfg.MaxConcurrentTotal = 0
|
||||
}
|
||||
return cfg
|
||||
}
|
||||
|
||||
// LoadConfigs 加载配置
|
||||
func (m *ExternalMCPManager) LoadConfigs(cfg *config.ExternalMCPConfig) {
|
||||
m.mu.Lock()
|
||||
@@ -588,6 +683,9 @@ func (m *ExternalMCPManager) CallTool(ctx context.Context, toolName string, args
|
||||
if !exists {
|
||||
return nil, "", fmt.Errorf("外部MCP客户端不存在: %s", mcpName)
|
||||
}
|
||||
if err := m.checkExternalMCPCircuit(mcpName); err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
|
||||
// 检查连接状态,如果未连接或状态为error,不允许调用
|
||||
if !client.IsConnected() {
|
||||
@@ -603,107 +701,216 @@ func (m *ExternalMCPManager) CallTool(ctx context.Context, toolName string, args
|
||||
return nil, "", fmt.Errorf("外部MCP客户端未连接: %s (状态: %s)", mcpName, status)
|
||||
}
|
||||
|
||||
// 创建执行记录
|
||||
executionID := uuid.New().String()
|
||||
execution := &ToolExecution{
|
||||
ID: executionID,
|
||||
ToolName: toolName, // 使用完整工具名称(包含MCP名称)
|
||||
Arguments: args,
|
||||
Status: "running",
|
||||
StartTime: time.Now(),
|
||||
if m.executionService == nil {
|
||||
m.executionService = NewExecutionService(m.storage, m.logger)
|
||||
m.executionService.ConfigureToolResultMaxBytes(m.toolResultMaxBytes)
|
||||
}
|
||||
var ownerUserID string
|
||||
if principal, ok := authctx.PrincipalFromContext(ctx); ok {
|
||||
execution.OwnerUserID = principal.UserID
|
||||
ownerUserID = principal.UserID
|
||||
}
|
||||
execution.ConversationID = MCPConversationIDFromContext(ctx)
|
||||
|
||||
m.mu.Lock()
|
||||
m.executions[executionID] = execution
|
||||
// 如果内存中的执行记录超过限制,清理最旧的记录
|
||||
m.cleanupOldExecutions()
|
||||
m.mu.Unlock()
|
||||
|
||||
if m.storage != nil {
|
||||
if err := m.storage.SaveToolExecution(execution); err != nil {
|
||||
m.logger.Warn("保存执行记录到数据库失败", zap.Error(err))
|
||||
}
|
||||
}
|
||||
|
||||
execCtx, runCancel := context.WithCancel(ctx)
|
||||
m.registerRunningCancel(executionID, runCancel)
|
||||
notifyToolRunBegin(ctx, executionID)
|
||||
defer func() {
|
||||
notifyToolRunEnd(ctx, executionID)
|
||||
runCancel()
|
||||
m.unregisterRunningCancel(executionID)
|
||||
}()
|
||||
|
||||
// 调用工具
|
||||
result, err := client.CallTool(execCtx, actualToolName, args)
|
||||
handle, err := m.executionService.Submit(ctx, ExecutionRequest{
|
||||
ToolName: toolName,
|
||||
Arguments: args,
|
||||
ConversationID: MCPConversationIDFromContext(ctx),
|
||||
OwnerUserID: ownerUserID,
|
||||
PreRun: func(runCtx context.Context, exec *ToolExecution) (func(), error) {
|
||||
release, acquireErr := m.acquireExternalMCPCallSlot(runCtx, mcpName)
|
||||
if acquireErr != nil {
|
||||
return nil, acquireErr
|
||||
}
|
||||
return release, nil
|
||||
},
|
||||
Run: func(runCtx context.Context) (*ToolResult, error) {
|
||||
result, callErr := client.CallTool(runCtx, actualToolName, args)
|
||||
if callErr != nil {
|
||||
m.handleConnectionDead(mcpName, client, callErr)
|
||||
}
|
||||
return result, callErr
|
||||
},
|
||||
OnDone: func(exec *ToolExecution) {
|
||||
failed := exec != nil && exec.Status != ToolExecutionStatusCompleted
|
||||
m.recordExternalMCPResult(mcpName, failed)
|
||||
m.updateStats(toolName, failed)
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
m.handleConnectionDead(mcpName, client, err)
|
||||
return nil, "", err
|
||||
}
|
||||
cancelledWithUserNote := m.applyAbortUserNoteToCancelledToolResult(executionID, &result, &err)
|
||||
|
||||
// 更新执行记录
|
||||
m.mu.RLock()
|
||||
waitTimeout := m.toolWaitTimeout
|
||||
m.mu.RUnlock()
|
||||
snapshot, waitErr := m.executionService.Wait(ctx, handle.ID, waitTimeout)
|
||||
if errors.Is(waitErr, ErrExecutionWaitTimeout) {
|
||||
return externalMCPWaitTimeoutResult(snapshot, waitTimeout), handle.ID, nil
|
||||
}
|
||||
if waitErr != nil {
|
||||
return nil, handle.ID, waitErr
|
||||
}
|
||||
if snapshot == nil || snapshot.Execution == nil {
|
||||
return &ToolResult{Content: []Content{{Type: "text", Text: "工具执行完成,但未返回执行快照"}}, IsError: true}, handle.ID, nil
|
||||
}
|
||||
if snapshot.Execution.Result != nil {
|
||||
return snapshot.Execution.Result, handle.ID, nil
|
||||
}
|
||||
if snapshot.Execution.Error != "" {
|
||||
return nil, handle.ID, errors.New(snapshot.Execution.Error)
|
||||
}
|
||||
return &ToolResult{Content: []Content{{Type: "text", Text: "工具执行完成,但未返回结果"}}, IsError: false}, handle.ID, nil
|
||||
}
|
||||
|
||||
func externalMCPWaitTimeoutResult(snapshot *ExecutionSnapshot, waitTimeout time.Duration) *ToolResult {
|
||||
execID := ""
|
||||
status := ToolExecutionStatusRunning
|
||||
toolName := ""
|
||||
elapsed := time.Duration(0)
|
||||
if snapshot != nil && snapshot.Execution != nil {
|
||||
execID = snapshot.Execution.ID
|
||||
status = snapshot.Execution.Status
|
||||
toolName = snapshot.Execution.ToolName
|
||||
elapsed = time.Since(snapshot.Execution.StartTime).Round(time.Second)
|
||||
}
|
||||
waitText := "unbounded"
|
||||
if waitTimeout > 0 {
|
||||
waitText = waitTimeout.Round(time.Second).String()
|
||||
}
|
||||
msg := fmt.Sprintf(`工具已提交到后台执行,但本次等待已到达上限。
|
||||
|
||||
execution_id: %s
|
||||
tool: %s
|
||||
status: %s
|
||||
wait_timeout: %s
|
||||
elapsed: %s
|
||||
|
||||
你可以继续推理、改用其他工具,或调用 wait_tool_execution 继续等待该 execution_id;也可以调用 cancel_tool_execution 取消。`, execID, toolName, status, waitText, elapsed)
|
||||
return &ToolResult{Content: []Content{{Type: "text", Text: msg}}, IsError: true}
|
||||
}
|
||||
|
||||
func (m *ExternalMCPManager) checkExternalMCPCircuit(mcpName string) error {
|
||||
if m == nil {
|
||||
return nil
|
||||
}
|
||||
name := strings.TrimSpace(mcpName)
|
||||
if name == "" {
|
||||
return nil
|
||||
}
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
if m.resilience.CircuitFailureThreshold < 0 {
|
||||
return nil
|
||||
}
|
||||
rt := m.externalMCPRuntimeLocked(name)
|
||||
if rt == nil || rt.circuitOpenUntil.IsZero() {
|
||||
return nil
|
||||
}
|
||||
now := time.Now()
|
||||
execution.EndTime = &now
|
||||
execution.Duration = now.Sub(execution.StartTime)
|
||||
|
||||
if err != nil {
|
||||
st, msg := executionStatusAndMessage(err)
|
||||
execution.Status = st
|
||||
execution.Error = msg
|
||||
} else if result != nil && result.IsError {
|
||||
if cancelledWithUserNote {
|
||||
execution.Status = "cancelled"
|
||||
execution.Error = ""
|
||||
execution.Result = result
|
||||
} else {
|
||||
execution.Status = "failed"
|
||||
if len(result.Content) > 0 {
|
||||
execution.Error = result.Content[0].Text
|
||||
} else {
|
||||
execution.Error = "工具执行返回错误结果"
|
||||
}
|
||||
execution.Result = result
|
||||
}
|
||||
} else {
|
||||
execution.Status = "completed"
|
||||
if result == nil {
|
||||
result = &ToolResult{
|
||||
Content: []Content{
|
||||
{Type: "text", Text: "工具执行完成,但未返回结果"},
|
||||
},
|
||||
}
|
||||
}
|
||||
execution.Result = result
|
||||
if now.Before(rt.circuitOpenUntil) {
|
||||
return fmt.Errorf("外部MCP服务 %s 已临时熔断,预计 %s 后重试", name, time.Until(rt.circuitOpenUntil).Round(time.Second))
|
||||
}
|
||||
rt.circuitOpenUntil = time.Time{}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *ExternalMCPManager) acquireExternalMCPCallSlot(ctx context.Context, mcpName string) (func(), error) {
|
||||
if m == nil {
|
||||
return func() {}, nil
|
||||
}
|
||||
name := strings.TrimSpace(mcpName)
|
||||
m.mu.Lock()
|
||||
rt := m.externalMCPRuntimeLocked(name)
|
||||
serverSem := chan struct{}(nil)
|
||||
if rt != nil {
|
||||
serverSem = rt.semaphore
|
||||
}
|
||||
globalSem := m.globalSemaphore
|
||||
m.mu.Unlock()
|
||||
|
||||
if m.storage != nil {
|
||||
if err := m.storage.SaveToolExecution(execution); err != nil {
|
||||
m.logger.Warn("保存执行记录到数据库失败", zap.Error(err))
|
||||
releaseGlobal := false
|
||||
if globalSem != nil {
|
||||
select {
|
||||
case globalSem <- struct{}{}:
|
||||
releaseGlobal = true
|
||||
case <-ctxDone(ctx):
|
||||
return func() {}, contextErr(ctx)
|
||||
}
|
||||
}
|
||||
|
||||
// 更新统计信息
|
||||
failed := err != nil || (result != nil && result.IsError)
|
||||
m.updateStats(toolName, failed)
|
||||
|
||||
// 如果使用存储,从内存中删除(已持久化)
|
||||
if m.storage != nil {
|
||||
m.mu.Lock()
|
||||
delete(m.executions, executionID)
|
||||
m.mu.Unlock()
|
||||
releaseServer := false
|
||||
if serverSem != nil {
|
||||
select {
|
||||
case serverSem <- struct{}{}:
|
||||
releaseServer = true
|
||||
case <-ctxDone(ctx):
|
||||
if releaseGlobal {
|
||||
<-globalSem
|
||||
}
|
||||
return func() {}, contextErr(ctx)
|
||||
}
|
||||
}
|
||||
return func() {
|
||||
if releaseServer {
|
||||
<-serverSem
|
||||
}
|
||||
if releaseGlobal {
|
||||
<-globalSem
|
||||
}
|
||||
}, nil
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
return nil, executionID, err
|
||||
func contextErr(ctx context.Context) error {
|
||||
if ctx == nil || ctx.Err() == nil {
|
||||
return context.Canceled
|
||||
}
|
||||
return ctx.Err()
|
||||
}
|
||||
|
||||
return result, executionID, nil
|
||||
func (m *ExternalMCPManager) recordExternalMCPResult(mcpName string, failed bool) {
|
||||
if m == nil || strings.TrimSpace(mcpName) == "" {
|
||||
return
|
||||
}
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
rt := m.externalMCPRuntimeLocked(mcpName)
|
||||
if rt == nil {
|
||||
return
|
||||
}
|
||||
if !failed {
|
||||
rt.consecutiveFailures = 0
|
||||
rt.circuitOpenUntil = time.Time{}
|
||||
return
|
||||
}
|
||||
if m.resilience.CircuitFailureThreshold < 0 {
|
||||
return
|
||||
}
|
||||
rt.consecutiveFailures++
|
||||
if rt.consecutiveFailures >= m.resilience.CircuitFailureThreshold {
|
||||
rt.circuitOpenUntil = time.Now().Add(m.resilience.CircuitCooldown)
|
||||
m.logger.Warn("外部MCP服务触发熔断",
|
||||
zap.String("name", mcpName),
|
||||
zap.Int("consecutiveFailures", rt.consecutiveFailures),
|
||||
zap.Duration("cooldown", m.resilience.CircuitCooldown),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
func (m *ExternalMCPManager) externalMCPRuntimeLocked(mcpName string) *externalMCPServerRuntime {
|
||||
if m.serverRuntimes == nil {
|
||||
m.serverRuntimes = make(map[string]*externalMCPServerRuntime)
|
||||
}
|
||||
name := strings.TrimSpace(mcpName)
|
||||
if name == "" {
|
||||
return nil
|
||||
}
|
||||
if rt := m.serverRuntimes[name]; rt != nil {
|
||||
return rt
|
||||
}
|
||||
var sem chan struct{}
|
||||
if m.resilience.MaxConcurrentPerServer > 0 {
|
||||
sem = make(chan struct{}, m.resilience.MaxConcurrentPerServer)
|
||||
}
|
||||
rt := &externalMCPServerRuntime{semaphore: sem}
|
||||
m.serverRuntimes[name] = rt
|
||||
return rt
|
||||
}
|
||||
|
||||
func (m *ExternalMCPManager) applyAbortUserNoteToCancelledToolResult(executionID string, result **ToolResult, err *error) (cancelledWithUserNote bool) {
|
||||
@@ -785,6 +992,11 @@ func (m *ExternalMCPManager) cleanupOldExecutions() {
|
||||
|
||||
// GetExecution 获取执行记录(先从内存查找,再从数据库查找)
|
||||
func (m *ExternalMCPManager) GetExecution(id string) (*ToolExecution, bool) {
|
||||
if m.executionService != nil {
|
||||
if snap, err := m.executionService.Get(id); err == nil && snap != nil && snap.Execution != nil {
|
||||
return snap.Execution, true
|
||||
}
|
||||
}
|
||||
m.mu.RLock()
|
||||
exec, exists := m.executions[id]
|
||||
m.mu.RUnlock()
|
||||
@@ -817,6 +1029,9 @@ func (m *ExternalMCPManager) unregisterRunningCancel(id string) {
|
||||
|
||||
// CancelToolExecutionWithNote 取消外部 MCP 工具;note 非空时与已返回输出合并后交给模型。
|
||||
func (m *ExternalMCPManager) CancelToolExecutionWithNote(id string, note string) bool {
|
||||
if m.executionService != nil && m.executionService.Cancel(id, note) {
|
||||
return true
|
||||
}
|
||||
m.mu.Lock()
|
||||
cancel, ok := m.runningCancels[id]
|
||||
if !ok || cancel == nil {
|
||||
@@ -844,6 +1059,11 @@ func (m *ExternalMCPManager) ActiveRunningExecutionIDs() map[string]struct{} {
|
||||
if m == nil {
|
||||
return nil
|
||||
}
|
||||
if m.executionService != nil {
|
||||
if ids := m.executionService.ActiveRunningExecutionIDs(); len(ids) > 0 {
|
||||
return ids
|
||||
}
|
||||
}
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
if len(m.runningCancels) == 0 {
|
||||
@@ -1335,12 +1555,21 @@ func (m *ExternalMCPManager) StartAllEnabled() {
|
||||
|
||||
// StopAll 停止所有客户端
|
||||
func (m *ExternalMCPManager) StopAll() {
|
||||
if m.executionService != nil {
|
||||
m.executionService.CancelAll("外部 MCP 管理器正在停止")
|
||||
}
|
||||
clients := make(map[string]ExternalMCPClient)
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
|
||||
for name, client := range m.clients {
|
||||
client.Close()
|
||||
clients[name] = client
|
||||
delete(m.clients, name)
|
||||
}
|
||||
m.mu.Unlock()
|
||||
|
||||
for name, client := range clients {
|
||||
if client != nil {
|
||||
_ = client.Close()
|
||||
}
|
||||
m.clearReconnectState(name)
|
||||
}
|
||||
|
||||
@@ -1360,6 +1589,6 @@ func (m *ExternalMCPManager) StopAll() {
|
||||
// 已经关闭,不需要再次关闭
|
||||
default:
|
||||
close(m.stopRefresh)
|
||||
m.refreshWg.Wait()
|
||||
}
|
||||
m.refreshWg.Wait()
|
||||
}
|
||||
|
||||
@@ -0,0 +1,230 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
type blockingExternalMCPClient struct {
|
||||
started chan struct{}
|
||||
calls chan string
|
||||
release chan struct{}
|
||||
result *ToolResult
|
||||
count atomic.Int32
|
||||
}
|
||||
|
||||
func newBlockingExternalMCPClient(resultText string) *blockingExternalMCPClient {
|
||||
return &blockingExternalMCPClient{
|
||||
started: make(chan struct{}),
|
||||
calls: make(chan string, 8),
|
||||
release: make(chan struct{}),
|
||||
result: &ToolResult{Content: []Content{{Type: "text", Text: resultText}}},
|
||||
}
|
||||
}
|
||||
|
||||
func (c *blockingExternalMCPClient) Initialize(ctx context.Context) error { return nil }
|
||||
func (c *blockingExternalMCPClient) ListTools(ctx context.Context) ([]Tool, error) {
|
||||
return []Tool{{Name: "slow_tool"}}, nil
|
||||
}
|
||||
func (c *blockingExternalMCPClient) CallTool(ctx context.Context, name string, args map[string]interface{}) (*ToolResult, error) {
|
||||
c.count.Add(1)
|
||||
select {
|
||||
case c.calls <- name:
|
||||
default:
|
||||
}
|
||||
select {
|
||||
case <-c.started:
|
||||
default:
|
||||
close(c.started)
|
||||
}
|
||||
select {
|
||||
case <-c.release:
|
||||
return c.result, nil
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
}
|
||||
func (c *blockingExternalMCPClient) Close() error { return nil }
|
||||
func (c *blockingExternalMCPClient) IsConnected() bool { return true }
|
||||
func (c *blockingExternalMCPClient) GetStatus() string { return "connected" }
|
||||
|
||||
type failingExternalMCPClient struct{}
|
||||
|
||||
func (c *failingExternalMCPClient) Initialize(ctx context.Context) error { return nil }
|
||||
func (c *failingExternalMCPClient) ListTools(ctx context.Context) ([]Tool, error) {
|
||||
return []Tool{{Name: "fail_tool"}}, nil
|
||||
}
|
||||
func (c *failingExternalMCPClient) CallTool(ctx context.Context, name string, args map[string]interface{}) (*ToolResult, error) {
|
||||
return nil, errors.New("boom")
|
||||
}
|
||||
func (c *failingExternalMCPClient) Close() error { return nil }
|
||||
func (c *failingExternalMCPClient) IsConnected() bool { return true }
|
||||
func (c *failingExternalMCPClient) GetStatus() string { return "connected" }
|
||||
|
||||
func TestExternalMCPManager_CallToolBoundedWaitThenContinue(t *testing.T) {
|
||||
manager := NewExternalMCPManager(zap.NewNop())
|
||||
manager.ConfigureToolWaitTimeoutSeconds(1)
|
||||
manager.toolWaitTimeout = 10 * time.Millisecond
|
||||
client := newBlockingExternalMCPClient("slow result ready")
|
||||
manager.clients["lab"] = client
|
||||
|
||||
callCtx, callCancel := context.WithCancel(context.Background())
|
||||
result, executionID, err := manager.CallTool(callCtx, "lab::slow_tool", map[string]interface{}{"target": "example"})
|
||||
if err != nil {
|
||||
t.Fatalf("CallTool returned error: %v", err)
|
||||
}
|
||||
if executionID == "" {
|
||||
t.Fatal("expected execution id")
|
||||
}
|
||||
if result == nil || !result.IsError {
|
||||
t.Fatalf("expected soft timeout tool result, got %#v", result)
|
||||
}
|
||||
text := ToolResultPlainText(result)
|
||||
if !strings.Contains(text, executionID) || !strings.Contains(text, "wait_tool_execution") {
|
||||
t.Fatalf("timeout result should include execution id and wait guidance, got %q", text)
|
||||
}
|
||||
|
||||
select {
|
||||
case <-client.started:
|
||||
default:
|
||||
t.Fatal("worker did not start")
|
||||
}
|
||||
callCancel()
|
||||
close(client.release)
|
||||
|
||||
snapshot, err := manager.executionService.Wait(context.Background(), executionID, time.Second)
|
||||
if err != nil {
|
||||
t.Fatalf("Wait returned error: %v", err)
|
||||
}
|
||||
if snapshot == nil || snapshot.Execution == nil {
|
||||
t.Fatal("expected execution snapshot")
|
||||
}
|
||||
if snapshot.Execution.Status != ToolExecutionStatusCompleted {
|
||||
t.Fatalf("status = %q, want completed", snapshot.Execution.Status)
|
||||
}
|
||||
if got := ToolResultPlainText(snapshot.Execution.Result); got != "slow result ready" {
|
||||
t.Fatalf("result = %q, want slow result ready", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecutionControlWaitToolReturnsCompletedResult(t *testing.T) {
|
||||
manager := NewExternalMCPManager(zap.NewNop())
|
||||
manager.toolWaitTimeout = 10 * time.Millisecond
|
||||
client := newBlockingExternalMCPClient("control wait result")
|
||||
manager.clients["lab"] = client
|
||||
|
||||
result, executionID, err := manager.CallTool(context.Background(), "lab::slow_tool", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("CallTool returned error: %v", err)
|
||||
}
|
||||
if result == nil || !result.IsError || executionID == "" {
|
||||
t.Fatalf("expected soft timeout and execution id, got result=%#v id=%q", result, executionID)
|
||||
}
|
||||
|
||||
server := NewServer(zap.NewNop())
|
||||
RegisterExecutionControlTools(server, manager)
|
||||
close(client.release)
|
||||
|
||||
waitResult, _, err := server.CallTool(context.Background(), "wait_tool_execution", map[string]interface{}{
|
||||
"execution_id": executionID,
|
||||
"timeout_seconds": 1,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("wait_tool_execution returned error: %v", err)
|
||||
}
|
||||
if waitResult == nil || waitResult.IsError {
|
||||
t.Fatalf("expected successful wait result, got %#v", waitResult)
|
||||
}
|
||||
body := ToolResultPlainText(waitResult)
|
||||
if !strings.Contains(body, `"status": "completed"`) || !strings.Contains(body, "control wait result") {
|
||||
t.Fatalf("wait result body missing completed status/result: %s", body)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExternalMCPManager_PerServerConcurrencyLimitsWorkers(t *testing.T) {
|
||||
manager := NewExternalMCPManager(zap.NewNop())
|
||||
manager.toolWaitTimeout = 10 * time.Millisecond
|
||||
manager.ConfigureResilience(ExternalMCPResilienceConfig{
|
||||
MaxConcurrentPerServer: 1,
|
||||
MaxConcurrentTotal: 4,
|
||||
CircuitFailureThreshold: -1,
|
||||
CircuitCooldown: time.Second,
|
||||
})
|
||||
client := newBlockingExternalMCPClient("ok")
|
||||
manager.clients["lab"] = client
|
||||
|
||||
done1 := make(chan struct{})
|
||||
go func() {
|
||||
_, _, _ = manager.CallTool(context.Background(), "lab::slow_tool", nil)
|
||||
close(done1)
|
||||
}()
|
||||
select {
|
||||
case <-client.calls:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("first worker did not enter client")
|
||||
}
|
||||
|
||||
type callOutcome struct {
|
||||
executionID string
|
||||
err error
|
||||
}
|
||||
done2 := make(chan callOutcome, 1)
|
||||
go func() {
|
||||
_, executionID, err := manager.CallTool(context.Background(), "lab::slow_tool", nil)
|
||||
done2 <- callOutcome{executionID: executionID, err: err}
|
||||
}()
|
||||
select {
|
||||
case <-client.calls:
|
||||
t.Fatal("second worker entered client before per-server slot was released")
|
||||
case <-time.After(50 * time.Millisecond):
|
||||
}
|
||||
var second callOutcome
|
||||
select {
|
||||
case second = <-done2:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("second call did not return after bounded wait")
|
||||
}
|
||||
if second.err != nil || second.executionID == "" {
|
||||
t.Fatalf("second call should return queued execution id after bounded wait, id=%q err=%v", second.executionID, second.err)
|
||||
}
|
||||
snapshot, err := manager.executionService.Get(second.executionID)
|
||||
if err != nil {
|
||||
t.Fatalf("Get queued execution: %v", err)
|
||||
}
|
||||
if snapshot == nil || snapshot.Execution == nil || snapshot.Execution.Status != ToolExecutionStatusQueued {
|
||||
t.Fatalf("second execution status = %#v, want queued", snapshot)
|
||||
}
|
||||
close(client.release)
|
||||
select {
|
||||
case <-client.calls:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("second worker did not enter client after slot release")
|
||||
}
|
||||
<-done1
|
||||
}
|
||||
|
||||
func TestExternalMCPManager_CircuitBreakerOpensAfterFailures(t *testing.T) {
|
||||
manager := NewExternalMCPManager(zap.NewNop())
|
||||
manager.ConfigureResilience(ExternalMCPResilienceConfig{
|
||||
MaxConcurrentPerServer: 2,
|
||||
MaxConcurrentTotal: 4,
|
||||
CircuitFailureThreshold: 1,
|
||||
CircuitCooldown: time.Minute,
|
||||
})
|
||||
manager.clients["lab"] = &failingExternalMCPClient{}
|
||||
|
||||
_, _, err := manager.CallTool(context.Background(), "lab::fail_tool", nil)
|
||||
if err == nil || !strings.Contains(err.Error(), "boom") {
|
||||
t.Fatalf("expected first call to fail with client error, got %v", err)
|
||||
}
|
||||
_, _, err = manager.CallTool(context.Background(), "lab::fail_tool", nil)
|
||||
if err == nil || !strings.Contains(err.Error(), "熔断") {
|
||||
t.Fatalf("expected circuit breaker rejection, got %v", err)
|
||||
}
|
||||
}
|
||||
+124
-97
@@ -15,6 +15,7 @@ import (
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/authctx"
|
||||
"cyberstrike-ai/internal/mcp/builtin"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"go.uber.org/zap"
|
||||
@@ -52,6 +53,9 @@ type Server struct {
|
||||
httpToolTimeoutMinutes *int
|
||||
httpToolTimeoutMu sync.RWMutex
|
||||
toolAuthorizer func(context.Context, string, map[string]interface{}) error
|
||||
executionService *ExecutionService
|
||||
toolWaitTimeout time.Duration
|
||||
toolResultMaxBytes int
|
||||
}
|
||||
|
||||
// SetToolAuthorizer installs the common policy decision point for every
|
||||
@@ -100,7 +104,10 @@ func NewServerWithStorage(logger *zap.Logger, storage MonitorStorage) *Server {
|
||||
sseClients: make(map[string]*sseClient),
|
||||
runningCancels: make(map[string]context.CancelFunc),
|
||||
abortUserNotes: make(map[string]string),
|
||||
toolWaitTimeout: 60 * time.Second,
|
||||
toolResultMaxBytes: DefaultToolResultMaxBytes,
|
||||
}
|
||||
s.executionService = NewExecutionService(storage, logger)
|
||||
|
||||
// 初始化默认提示词和资源
|
||||
s.initDefaultPrompts()
|
||||
@@ -109,6 +116,18 @@ func NewServerWithStorage(logger *zap.Logger, storage MonitorStorage) *Server {
|
||||
return s
|
||||
}
|
||||
|
||||
func (s *Server) ConfigureToolResultMaxBytes(maxBytes int) {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
s.mu.Lock()
|
||||
s.toolResultMaxBytes = maxBytes
|
||||
s.mu.Unlock()
|
||||
if s.executionService != nil {
|
||||
s.executionService.ConfigureToolResultMaxBytes(maxBytes)
|
||||
}
|
||||
}
|
||||
|
||||
// ConfigureHTTPToolCallTimeoutFromAgentMinutes 将 agent.tool_timeout_minutes 同步到经 HTTP POST /api/mcp 触发的 tools/call。
|
||||
// minutes<=0 表示不设置硬性截止时间(与配置「0 不限制」一致);minutes>0 为该次调用的最长等待时间。
|
||||
// 未调用前对 tools/call 使用默认 30 分钟(与历史硬编码一致)。
|
||||
@@ -125,6 +144,19 @@ func (s *Server) ConfigureHTTPToolCallTimeoutFromAgentMinutes(minutes int) {
|
||||
s.httpToolTimeoutMinutes = &v
|
||||
}
|
||||
|
||||
func (s *Server) ConfigureToolWaitTimeoutSeconds(seconds int) {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if seconds <= 0 {
|
||||
s.toolWaitTimeout = 0
|
||||
return
|
||||
}
|
||||
s.toolWaitTimeout = time.Duration(seconds) * time.Second
|
||||
}
|
||||
|
||||
func (s *Server) effectiveHTTPToolCallDeadline(parent context.Context) (context.Context, context.CancelFunc) {
|
||||
const defaultDur = 30 * time.Minute
|
||||
if parent == nil {
|
||||
@@ -709,6 +741,11 @@ func (s *Server) updateStats(toolName string, failed bool) {
|
||||
|
||||
// GetExecution 获取执行记录(先从内存查找,再从数据库查找)
|
||||
func (s *Server) GetExecution(id string) (*ToolExecution, bool) {
|
||||
if s.executionService != nil {
|
||||
if snap, err := s.executionService.Get(id); err == nil && snap != nil && snap.Execution != nil {
|
||||
return snap.Execution, true
|
||||
}
|
||||
}
|
||||
s.mu.RLock()
|
||||
exec, exists := s.executions[id]
|
||||
s.mu.RUnlock()
|
||||
@@ -860,112 +897,90 @@ func (s *Server) CallTool(ctx context.Context, toolName string, args map[string]
|
||||
return nil, "", fmt.Errorf("工具 %s 未找到", toolName)
|
||||
}
|
||||
|
||||
// 创建执行记录
|
||||
executionID := uuid.New().String()
|
||||
execution := &ToolExecution{
|
||||
ID: executionID,
|
||||
ToolName: toolName,
|
||||
Arguments: args,
|
||||
Status: "running",
|
||||
StartTime: time.Now(),
|
||||
if s.executionService == nil {
|
||||
s.executionService = NewExecutionService(s.storage, s.logger)
|
||||
s.executionService.ConfigureToolResultMaxBytes(s.toolResultMaxBytes)
|
||||
}
|
||||
var ownerUserID string
|
||||
if principal, ok := authctx.PrincipalFromContext(ctx); ok {
|
||||
execution.OwnerUserID = principal.UserID
|
||||
ownerUserID = principal.UserID
|
||||
}
|
||||
execution.ConversationID = MCPConversationIDFromContext(ctx)
|
||||
|
||||
s.mu.Lock()
|
||||
s.executions[executionID] = execution
|
||||
// 如果内存中的执行记录超过限制,清理最旧的记录
|
||||
s.cleanupOldExecutions()
|
||||
s.mu.Unlock()
|
||||
|
||||
if s.storage != nil {
|
||||
if err := s.storage.SaveToolExecution(execution); err != nil {
|
||||
s.logger.Warn("保存执行记录到数据库失败", zap.Error(err))
|
||||
}
|
||||
}
|
||||
|
||||
execCtx, runCancel := context.WithCancel(ctx)
|
||||
s.registerRunningCancel(executionID, runCancel)
|
||||
notifyToolRunBegin(ctx, executionID)
|
||||
defer func() {
|
||||
notifyToolRunEnd(ctx, executionID)
|
||||
runCancel()
|
||||
s.unregisterRunningCancel(executionID)
|
||||
}()
|
||||
|
||||
result, err := handler(execCtx, args)
|
||||
cancelledWithUserNote := s.applyAbortUserNoteToCancelledToolResult(executionID, &result, &err)
|
||||
|
||||
s.mu.Lock()
|
||||
now := time.Now()
|
||||
execution.EndTime = &now
|
||||
execution.Duration = now.Sub(execution.StartTime)
|
||||
var failed bool
|
||||
var finalResult *ToolResult
|
||||
|
||||
handle, err := s.executionService.Submit(ctx, ExecutionRequest{
|
||||
ToolName: toolName,
|
||||
Arguments: args,
|
||||
ConversationID: MCPConversationIDFromContext(ctx),
|
||||
OwnerUserID: ownerUserID,
|
||||
Run: func(runCtx context.Context) (*ToolResult, error) {
|
||||
return handler(runCtx, args)
|
||||
},
|
||||
OnDone: func(exec *ToolExecution) {
|
||||
failed := exec != nil && exec.Status != ToolExecutionStatusCompleted
|
||||
s.updateStats(toolName, failed)
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
st, msg := executionStatusAndMessage(err)
|
||||
execution.Status = st
|
||||
execution.Error = msg
|
||||
failed = true
|
||||
} else if result != nil && result.IsError {
|
||||
if cancelledWithUserNote {
|
||||
execution.Status = "cancelled"
|
||||
execution.Error = ""
|
||||
execution.Result = result
|
||||
failed = true
|
||||
finalResult = result
|
||||
} else {
|
||||
execution.Status = "failed"
|
||||
if len(result.Content) > 0 {
|
||||
execution.Error = result.Content[0].Text
|
||||
} else {
|
||||
execution.Error = "工具执行返回错误结果"
|
||||
}
|
||||
execution.Result = result
|
||||
failed = true
|
||||
finalResult = result
|
||||
}
|
||||
} else {
|
||||
execution.Status = "completed"
|
||||
if result == nil {
|
||||
result = &ToolResult{
|
||||
Content: []Content{
|
||||
{Type: "text", Text: "工具执行完成,但未返回结果"},
|
||||
},
|
||||
}
|
||||
}
|
||||
execution.Result = result
|
||||
finalResult = result
|
||||
failed = false
|
||||
return nil, "", err
|
||||
}
|
||||
|
||||
if finalResult == nil {
|
||||
finalResult = execution.Result
|
||||
s.mu.RLock()
|
||||
waitTimeout := s.toolWaitTimeout
|
||||
s.mu.RUnlock()
|
||||
if isExecutionControlTool(toolName) {
|
||||
waitTimeout = 0
|
||||
}
|
||||
s.mu.Unlock()
|
||||
|
||||
if s.storage != nil {
|
||||
if err := s.storage.SaveToolExecution(execution); err != nil {
|
||||
s.logger.Warn("保存执行记录到数据库失败", zap.Error(err))
|
||||
}
|
||||
snapshot, waitErr := s.executionService.Wait(ctx, handle.ID, waitTimeout)
|
||||
if errors.Is(waitErr, ErrExecutionWaitTimeout) {
|
||||
return internalMCPWaitTimeoutResult(snapshot, waitTimeout), handle.ID, nil
|
||||
}
|
||||
|
||||
s.updateStats(toolName, failed)
|
||||
|
||||
if s.storage != nil {
|
||||
s.mu.Lock()
|
||||
delete(s.executions, executionID)
|
||||
s.mu.Unlock()
|
||||
if waitErr != nil {
|
||||
return nil, handle.ID, waitErr
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
return nil, executionID, err
|
||||
if snapshot == nil || snapshot.Execution == nil {
|
||||
return &ToolResult{Content: []Content{{Type: "text", Text: "工具执行完成,但未返回执行快照"}}, IsError: true}, handle.ID, nil
|
||||
}
|
||||
if snapshot.Execution.Result != nil {
|
||||
return snapshot.Execution.Result, handle.ID, nil
|
||||
}
|
||||
if snapshot.Execution.Error != "" {
|
||||
return nil, handle.ID, errors.New(snapshot.Execution.Error)
|
||||
}
|
||||
return &ToolResult{Content: []Content{{Type: "text", Text: "工具执行完成,但未返回结果"}}, IsError: false}, handle.ID, nil
|
||||
}
|
||||
|
||||
return finalResult, executionID, nil
|
||||
func internalMCPWaitTimeoutResult(snapshot *ExecutionSnapshot, waitTimeout time.Duration) *ToolResult {
|
||||
execID := ""
|
||||
status := ToolExecutionStatusRunning
|
||||
toolName := ""
|
||||
elapsed := time.Duration(0)
|
||||
if snapshot != nil && snapshot.Execution != nil {
|
||||
execID = snapshot.Execution.ID
|
||||
status = snapshot.Execution.Status
|
||||
toolName = snapshot.Execution.ToolName
|
||||
elapsed = time.Since(snapshot.Execution.StartTime).Round(time.Second)
|
||||
}
|
||||
waitText := "unbounded"
|
||||
if waitTimeout > 0 {
|
||||
waitText = waitTimeout.Round(time.Second).String()
|
||||
}
|
||||
msg := fmt.Sprintf(`工具已提交到后台执行,但本次等待已到达上限。
|
||||
|
||||
execution_id: %s
|
||||
tool: %s
|
||||
status: %s
|
||||
wait_timeout: %s
|
||||
elapsed: %s
|
||||
|
||||
你可以继续推理、改用其他工具,或调用 wait_tool_execution 继续等待该 execution_id;也可以调用 cancel_tool_execution 取消。`, execID, toolName, status, waitText, elapsed)
|
||||
return &ToolResult{Content: []Content{{Type: "text", Text: msg}}, IsError: true}
|
||||
}
|
||||
|
||||
func isExecutionControlTool(toolName string) bool {
|
||||
switch strings.TrimSpace(toolName) {
|
||||
case builtin.ToolGetToolExecution, builtin.ToolWaitToolExecution, builtin.ToolCancelToolExecution:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
// BeginToolExecution 创建 running 状态的执行记录,供 Eino 等非 CallTool 路径在工具开始时落库。
|
||||
@@ -1020,6 +1035,7 @@ func (s *Server) FinishToolExecution(ctx context.Context, executionID, toolName
|
||||
var finalResult *ToolResult
|
||||
|
||||
s.mu.Lock()
|
||||
maxBytes := s.toolResultMaxBytes
|
||||
exec, inMem := s.executions[id]
|
||||
if !inMem || exec == nil {
|
||||
exec = &ToolExecution{
|
||||
@@ -1053,6 +1069,7 @@ func (s *Server) FinishToolExecution(ctx context.Context, executionID, toolName
|
||||
exec.Error = msg
|
||||
if strings.TrimSpace(resultText) != "" {
|
||||
finalResult = &ToolResult{Content: []Content{{Type: "text", Text: resultText}}}
|
||||
finalResult = NormalizeToolResultForStorage(finalResult, maxBytes)
|
||||
exec.Result = finalResult
|
||||
}
|
||||
} else {
|
||||
@@ -1062,6 +1079,7 @@ func (s *Server) FinishToolExecution(ctx context.Context, executionID, toolName
|
||||
text = "(无输出)"
|
||||
}
|
||||
finalResult = &ToolResult{Content: []Content{{Type: "text", Text: text}}}
|
||||
finalResult = NormalizeToolResultForStorage(finalResult, maxBytes)
|
||||
exec.Result = finalResult
|
||||
}
|
||||
s.mu.Unlock()
|
||||
@@ -1098,6 +1116,7 @@ func (s *Server) UpdateToolExecutionResult(executionID string, result *ToolResul
|
||||
return nil
|
||||
}
|
||||
s.mu.Lock()
|
||||
result = NormalizeToolResultForStorage(result, s.toolResultMaxBytes)
|
||||
if exec, ok := s.executions[executionID]; ok && exec != nil {
|
||||
exec.Result = result
|
||||
}
|
||||
@@ -1205,6 +1224,9 @@ func (s *Server) applyAbortUserNoteToCancelledToolResult(executionID string, res
|
||||
|
||||
// CancelToolExecutionWithNote 取消内部工具;note 非空时与工具已返回文本合并后交给上层模型。
|
||||
func (s *Server) CancelToolExecutionWithNote(id string, note string) bool {
|
||||
if s.executionService != nil && s.executionService.Cancel(id, note) {
|
||||
return true
|
||||
}
|
||||
s.runningCancelsMu.Lock()
|
||||
cancel, ok := s.runningCancels[id]
|
||||
if !ok || cancel == nil {
|
||||
@@ -1232,12 +1254,17 @@ func (s *Server) ActiveRunningExecutionIDs() map[string]struct{} {
|
||||
if s == nil {
|
||||
return nil
|
||||
}
|
||||
out := make(map[string]struct{})
|
||||
if s.executionService != nil {
|
||||
for id := range s.executionService.ActiveRunningExecutionIDs() {
|
||||
out[id] = struct{}{}
|
||||
}
|
||||
}
|
||||
s.runningCancelsMu.Lock()
|
||||
defer s.runningCancelsMu.Unlock()
|
||||
if len(s.runningCancels) == 0 {
|
||||
if len(s.runningCancels) == 0 && len(out) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make(map[string]struct{}, len(s.runningCancels))
|
||||
for id := range s.runningCancels {
|
||||
out[id] = struct{}{}
|
||||
}
|
||||
|
||||
@@ -3,7 +3,9 @@ package mcp
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/authctx"
|
||||
|
||||
@@ -34,3 +36,147 @@ func TestToolAuthorizerIsUniversalAndExecutionKeepsOwner(t *testing.T) {
|
||||
t.Fatalf("execution owner = %#v, want u1", execution)
|
||||
}
|
||||
}
|
||||
|
||||
func TestServerCallToolBoundedWaitForInternalTool(t *testing.T) {
|
||||
server := NewServer(zap.NewNop())
|
||||
server.toolWaitTimeout = 10 * time.Millisecond
|
||||
release := make(chan struct{})
|
||||
started := make(chan struct{})
|
||||
server.RegisterTool(Tool{Name: "slow", InputSchema: map[string]interface{}{"type": "object"}}, func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) {
|
||||
close(started)
|
||||
select {
|
||||
case <-release:
|
||||
return &ToolResult{Content: []Content{{Type: "text", Text: "internal done"}}}, nil
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
})
|
||||
|
||||
callCtx, callCancel := context.WithCancel(context.Background())
|
||||
result, executionID, err := server.CallTool(callCtx, "slow", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("CallTool returned error: %v", err)
|
||||
}
|
||||
if executionID == "" || result == nil || !result.IsError {
|
||||
t.Fatalf("expected soft timeout with execution id, result=%#v id=%q", result, executionID)
|
||||
}
|
||||
if text := ToolResultPlainText(result); !strings.Contains(text, executionID) || !strings.Contains(text, "wait_tool_execution") {
|
||||
t.Fatalf("timeout result missing execution guidance: %q", text)
|
||||
}
|
||||
select {
|
||||
case <-started:
|
||||
default:
|
||||
t.Fatal("internal worker did not start")
|
||||
}
|
||||
callCancel()
|
||||
close(release)
|
||||
|
||||
snapshot, err := server.executionService.Wait(context.Background(), executionID, time.Second)
|
||||
if err != nil {
|
||||
t.Fatalf("wait internal execution: %v", err)
|
||||
}
|
||||
if snapshot == nil || snapshot.Execution == nil || snapshot.Execution.Status != ToolExecutionStatusCompleted {
|
||||
t.Fatalf("snapshot = %#v, want completed", snapshot)
|
||||
}
|
||||
if got := ToolResultPlainText(snapshot.Execution.Result); got != "internal done" {
|
||||
t.Fatalf("result = %q, want internal done", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWaitToolExecutionWaitsForInternalActiveExecution(t *testing.T) {
|
||||
server := NewServer(zap.NewNop())
|
||||
server.toolWaitTimeout = 10 * time.Millisecond
|
||||
release := make(chan struct{})
|
||||
server.RegisterTool(Tool{Name: "slow", InputSchema: map[string]interface{}{"type": "object"}}, func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) {
|
||||
select {
|
||||
case <-release:
|
||||
return &ToolResult{Content: []Content{{Type: "text", Text: "wait saw completion"}}}, nil
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
})
|
||||
RegisterExecutionControlTools(server, nil)
|
||||
|
||||
result, executionID, err := server.CallTool(context.Background(), "slow", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("CallTool returned error: %v", err)
|
||||
}
|
||||
if result == nil || !result.IsError || executionID == "" {
|
||||
t.Fatalf("expected initial bounded wait timeout, result=%#v id=%q", result, executionID)
|
||||
}
|
||||
|
||||
done := make(chan *ToolResult, 1)
|
||||
errCh := make(chan error, 1)
|
||||
go func() {
|
||||
waitResult, _, waitErr := server.CallTool(context.Background(), "wait_tool_execution", map[string]interface{}{
|
||||
"execution_id": executionID,
|
||||
"timeout_seconds": 1,
|
||||
})
|
||||
if waitErr != nil {
|
||||
errCh <- waitErr
|
||||
return
|
||||
}
|
||||
done <- waitResult
|
||||
}()
|
||||
|
||||
select {
|
||||
case <-done:
|
||||
t.Fatal("wait_tool_execution returned before target execution completed")
|
||||
case err := <-errCh:
|
||||
t.Fatalf("wait_tool_execution errored before release: %v", err)
|
||||
case <-time.After(50 * time.Millisecond):
|
||||
}
|
||||
close(release)
|
||||
|
||||
select {
|
||||
case err := <-errCh:
|
||||
t.Fatalf("wait_tool_execution returned error: %v", err)
|
||||
case waitResult := <-done:
|
||||
if waitResult == nil || waitResult.IsError {
|
||||
t.Fatalf("expected successful wait result, got %#v", waitResult)
|
||||
}
|
||||
if body := ToolResultPlainText(waitResult); !strings.Contains(body, "wait saw completion") || !strings.Contains(body, `"status": "completed"`) {
|
||||
t.Fatalf("wait result missing completed target: %s", body)
|
||||
}
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("wait_tool_execution did not return after target completion")
|
||||
}
|
||||
}
|
||||
|
||||
func TestWaitToolExecutionTimeoutIsObservationNotFailure(t *testing.T) {
|
||||
server := NewServer(zap.NewNop())
|
||||
server.toolWaitTimeout = 10 * time.Millisecond
|
||||
release := make(chan struct{})
|
||||
server.RegisterTool(Tool{Name: "slow_observed", InputSchema: map[string]interface{}{"type": "object"}}, func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) {
|
||||
<-release
|
||||
return &ToolResult{Content: []Content{{Type: "text", Text: "done"}}}, nil
|
||||
})
|
||||
RegisterExecutionControlTools(server, nil)
|
||||
|
||||
result, executionID, err := server.CallTool(context.Background(), "slow_observed", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("CallTool returned error: %v", err)
|
||||
}
|
||||
if result == nil || !result.IsError || executionID == "" {
|
||||
t.Fatalf("expected initial bounded wait timeout, result=%#v id=%q", result, executionID)
|
||||
}
|
||||
|
||||
waitResult, _, err := server.CallTool(context.Background(), "wait_tool_execution", map[string]interface{}{
|
||||
"execution_id": executionID,
|
||||
"timeout_seconds": 0.01,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("wait_tool_execution returned error: %v", err)
|
||||
}
|
||||
if waitResult == nil {
|
||||
t.Fatal("missing wait result")
|
||||
}
|
||||
if waitResult.IsError {
|
||||
t.Fatalf("wait timeout should be a successful observation, got %#v", waitResult)
|
||||
}
|
||||
body := ToolResultPlainText(waitResult)
|
||||
if !strings.Contains(body, `"status": "running"`) || !strings.Contains(body, "本次等待已到达") {
|
||||
t.Fatalf("wait timeout body missing running status/guidance: %s", body)
|
||||
}
|
||||
close(release)
|
||||
}
|
||||
|
||||
@@ -0,0 +1,99 @@
|
||||
package mcp
|
||||
|
||||
import "fmt"
|
||||
|
||||
const DefaultToolResultMaxBytes = 12000
|
||||
|
||||
// NormalizeToolResultForStorage returns the canonical result used by both the
|
||||
// agent-facing response and monitor persistence.
|
||||
func NormalizeToolResultForStorage(result *ToolResult, maxBytes int) *ToolResult {
|
||||
if result == nil {
|
||||
return nil
|
||||
}
|
||||
out := cloneToolResult(result)
|
||||
if maxBytes <= 0 {
|
||||
return out
|
||||
}
|
||||
|
||||
total := 0
|
||||
for _, c := range out.Content {
|
||||
if c.Type == "text" {
|
||||
total += len(c.Text)
|
||||
}
|
||||
}
|
||||
if total <= maxBytes {
|
||||
return out
|
||||
}
|
||||
|
||||
remaining := maxBytes
|
||||
truncated := false
|
||||
for i := range out.Content {
|
||||
if out.Content[i].Type != "text" {
|
||||
continue
|
||||
}
|
||||
if remaining <= 0 {
|
||||
out.Content[i].Text = ""
|
||||
truncated = true
|
||||
continue
|
||||
}
|
||||
text := out.Content[i].Text
|
||||
if len(text) <= remaining {
|
||||
remaining -= len(text)
|
||||
continue
|
||||
}
|
||||
out.Content[i].Text = truncateUTF8Bytes(text, remaining)
|
||||
remaining = 0
|
||||
truncated = true
|
||||
}
|
||||
|
||||
if truncated {
|
||||
marker := fmt.Sprintf("\n\n...[tool output truncated: original %d bytes, kept %d bytes]...", total, maxBytes)
|
||||
textBudget := maxBytes - len(marker)
|
||||
if textBudget < 0 {
|
||||
marker = truncateUTF8Bytes(marker, maxBytes)
|
||||
textBudget = 0
|
||||
}
|
||||
for i := range out.Content {
|
||||
if out.Content[i].Type == "text" {
|
||||
out.Content[i].Text = truncateUTF8Bytes(out.Content[i].Text, textBudget) + marker
|
||||
remaining = 0
|
||||
for j := range out.Content {
|
||||
if j != i && out.Content[j].Type == "text" {
|
||||
out.Content[j].Text = ""
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
}
|
||||
out.Content = append(out.Content, Content{Type: "text", Text: marker})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func cloneToolResult(in *ToolResult) *ToolResult {
|
||||
if in == nil {
|
||||
return nil
|
||||
}
|
||||
out := *in
|
||||
if in.Content != nil {
|
||||
out.Content = append([]Content(nil), in.Content...)
|
||||
}
|
||||
return &out
|
||||
}
|
||||
|
||||
func truncateUTF8Bytes(s string, maxBytes int) string {
|
||||
if maxBytes <= 0 {
|
||||
return ""
|
||||
}
|
||||
if len(s) <= maxBytes {
|
||||
return s
|
||||
}
|
||||
cut := maxBytes
|
||||
for cut > 0 && (s[cut]&0xC0) == 0x80 {
|
||||
cut--
|
||||
}
|
||||
if cut <= 0 {
|
||||
return ""
|
||||
}
|
||||
return s[:cut]
|
||||
}
|
||||
@@ -0,0 +1,125 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
type inMemoryMonitorStorage struct {
|
||||
executions map[string]*ToolExecution
|
||||
}
|
||||
|
||||
func newInMemoryMonitorStorage() *inMemoryMonitorStorage {
|
||||
return &inMemoryMonitorStorage{executions: map[string]*ToolExecution{}}
|
||||
}
|
||||
|
||||
func (s *inMemoryMonitorStorage) SaveToolExecution(exec *ToolExecution) error {
|
||||
if exec != nil {
|
||||
s.executions[exec.ID] = cloneToolExecution(exec)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *inMemoryMonitorStorage) UpdateToolExecutionResult(id string, result *ToolResult) error {
|
||||
exec := s.executions[id]
|
||||
if exec == nil {
|
||||
exec = &ToolExecution{ID: id}
|
||||
s.executions[id] = exec
|
||||
}
|
||||
exec.Result = cloneToolResult(result)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *inMemoryMonitorStorage) LoadToolExecutions() ([]*ToolExecution, error) {
|
||||
out := make([]*ToolExecution, 0, len(s.executions))
|
||||
for _, exec := range s.executions {
|
||||
out = append(out, cloneToolExecution(exec))
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (s *inMemoryMonitorStorage) GetToolExecution(id string) (*ToolExecution, error) {
|
||||
if exec := s.executions[id]; exec != nil {
|
||||
return cloneToolExecution(exec), nil
|
||||
}
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func (s *inMemoryMonitorStorage) SaveToolStats(string, *ToolStats) error { return nil }
|
||||
|
||||
func (s *inMemoryMonitorStorage) LoadToolStats() (map[string]*ToolStats, error) {
|
||||
return map[string]*ToolStats{}, nil
|
||||
}
|
||||
|
||||
func (s *inMemoryMonitorStorage) UpdateToolStats(string, int, int, int, *time.Time) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func TestServerCallToolStoresAndReturnsSameGuardedResult(t *testing.T) {
|
||||
storage := newInMemoryMonitorStorage()
|
||||
server := NewServerWithStorage(zap.NewNop(), storage)
|
||||
server.ConfigureToolWaitTimeoutSeconds(0)
|
||||
server.ConfigureToolResultMaxBytes(50)
|
||||
server.RegisterTool(Tool{Name: "big", InputSchema: map[string]interface{}{"type": "object"}}, func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) {
|
||||
return &ToolResult{Content: []Content{{Type: "text", Text: strings.Repeat("x", 100)}}}, nil
|
||||
})
|
||||
|
||||
result, executionID, err := server.CallTool(context.Background(), "big", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("CallTool: %v", err)
|
||||
}
|
||||
if executionID == "" {
|
||||
t.Fatal("missing execution id")
|
||||
}
|
||||
returned := ToolResultPlainText(result)
|
||||
if !strings.Contains(returned, "tool output truncated") || strings.Contains(returned, strings.Repeat("x", 100)) {
|
||||
t.Fatalf("returned result was not guarded: %q", returned)
|
||||
}
|
||||
if len(returned) > 50 {
|
||||
t.Fatalf("returned result exceeded hard limit: len=%d text=%q", len(returned), returned)
|
||||
}
|
||||
|
||||
inMem, ok := server.GetExecution(executionID)
|
||||
if !ok || inMem == nil || inMem.Result == nil {
|
||||
t.Fatalf("missing in-memory execution: %#v", inMem)
|
||||
}
|
||||
stored := storage.executions[executionID]
|
||||
if stored == nil || stored.Result == nil {
|
||||
t.Fatalf("missing stored execution: %#v", stored)
|
||||
}
|
||||
if ToolResultPlainText(inMem.Result) != returned {
|
||||
t.Fatalf("in-memory result != returned\nmem=%q\nret=%q", ToolResultPlainText(inMem.Result), returned)
|
||||
}
|
||||
if ToolResultPlainText(stored.Result) != returned {
|
||||
t.Fatalf("stored result != returned\nstored=%q\nret=%q", ToolResultPlainText(stored.Result), returned)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecutionServiceStoresGuardedResult(t *testing.T) {
|
||||
service := NewExecutionService(nil, zap.NewNop())
|
||||
service.ConfigureToolResultMaxBytes(80)
|
||||
handle, err := service.Submit(context.Background(), ExecutionRequest{
|
||||
ToolName: "big",
|
||||
Run: func(context.Context) (*ToolResult, error) {
|
||||
return &ToolResult{Content: []Content{{Type: "text", Text: strings.Repeat("a", 200)}}}, nil
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Submit: %v", err)
|
||||
}
|
||||
snap, err := service.Wait(context.Background(), handle.ID, time.Second)
|
||||
if err != nil {
|
||||
t.Fatalf("Wait: %v", err)
|
||||
}
|
||||
got := ToolResultPlainText(snap.Execution.Result)
|
||||
if !strings.Contains(got, "tool output truncated") || strings.Contains(got, strings.Repeat("a", 64)) {
|
||||
t.Fatalf("service result was not guarded: %q", got)
|
||||
}
|
||||
if len(got) > 80 {
|
||||
t.Fatalf("service result exceeded hard limit: len=%d text=%q", len(got), got)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user