From bfe1d28650122d80e601368bf88a6ad7663dd28e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=85=AC=E6=98=8E?= <83812544+Ed1s0nZ@users.noreply.github.com> Date: Tue, 21 Jul 2026 20:52:43 +0800 Subject: [PATCH] Add files via upload --- internal/mcp/builtin/constants.go | 11 + internal/mcp/execution_control_tools.go | 224 ++++++++ internal/mcp/execution_service.go | 557 ++++++++++++++++++++ internal/mcp/execution_service_test.go | 41 ++ internal/mcp/external_manager.go | 489 ++++++++++++----- internal/mcp/external_manager_async_test.go | 230 ++++++++ internal/mcp/server.go | 221 ++++---- internal/mcp/server_authorization_test.go | 146 +++++ internal/mcp/tool_result_guard.go | 99 ++++ internal/mcp/tool_result_guard_test.go | 125 +++++ 10 files changed, 1916 insertions(+), 227 deletions(-) create mode 100644 internal/mcp/execution_control_tools.go create mode 100644 internal/mcp/execution_service.go create mode 100644 internal/mcp/execution_service_test.go create mode 100644 internal/mcp/external_manager_async_test.go create mode 100644 internal/mcp/tool_result_guard.go create mode 100644 internal/mcp/tool_result_guard_test.go diff --git a/internal/mcp/builtin/constants.go b/internal/mcp/builtin/constants.go index cdc23b7f..d38bcaae 100644 --- a/internal/mcp/builtin/constants.go +++ b/internal/mcp/builtin/constants.go @@ -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, diff --git a/internal/mcp/execution_control_tools.go b/internal/mcp/execution_control_tools.go new file mode 100644 index 00000000..f8b56783 --- /dev/null +++ b/internal/mcp/execution_control_tools.go @@ -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 +} diff --git a/internal/mcp/execution_service.go b/internal/mcp/execution_service.go new file mode 100644 index 00000000..8de9c376 --- /dev/null +++ b/internal/mcp/execution_service.go @@ -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 +} diff --git a/internal/mcp/execution_service_test.go b/internal/mcp/execution_service_test.go new file mode 100644 index 00000000..29f981c5 --- /dev/null +++ b/internal/mcp/execution_service_test.go @@ -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") + } +} diff --git a/internal/mcp/external_manager.go b/internal/mcp/external_manager.go index 9f69ade0..ed6516ff 100644 --- a/internal/mcp/external_manager.go +++ b/internal/mcp/external_manager.go @@ -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() } diff --git a/internal/mcp/external_manager_async_test.go b/internal/mcp/external_manager_async_test.go new file mode 100644 index 00000000..6c280ed9 --- /dev/null +++ b/internal/mcp/external_manager_async_test.go @@ -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) + } +} diff --git a/internal/mcp/server.go b/internal/mcp/server.go index 6d97d999..9251c667 100644 --- a/internal/mcp/server.go +++ b/internal/mcp/server.go @@ -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{}{} } diff --git a/internal/mcp/server_authorization_test.go b/internal/mcp/server_authorization_test.go index 51c44283..ecfcf3cf 100644 --- a/internal/mcp/server_authorization_test.go +++ b/internal/mcp/server_authorization_test.go @@ -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) +} diff --git a/internal/mcp/tool_result_guard.go b/internal/mcp/tool_result_guard.go new file mode 100644 index 00000000..70fdbb9b --- /dev/null +++ b/internal/mcp/tool_result_guard.go @@ -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] +} diff --git a/internal/mcp/tool_result_guard_test.go b/internal/mcp/tool_result_guard_test.go new file mode 100644 index 00000000..4a0ff55c --- /dev/null +++ b/internal/mcp/tool_result_guard_test.go @@ -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) + } +}