feat: add configurable tool call blocking and monitoring

internal/mcp/server.go

@@ -16,6 +16,7 @@ import (
 
 	"cyberstrike-ai/internal/authctx"
 	"cyberstrike-ai/internal/mcp/builtin"
+	"cyberstrike-ai/internal/toolguard"
 
 	"github.com/google/uuid"
 	"go.uber.org/zap"
@@ -53,6 +54,7 @@ type Server struct {
 	httpToolTimeoutMinutes *int
 	httpToolTimeoutMu      sync.RWMutex
 	toolAuthorizer         func(context.Context, string, map[string]interface{}) error
+	toolGuard              *toolguard.Manager
 	executionService       *ExecutionService
 	toolWaitTimeout        time.Duration
 	toolResultMaxBytes     int
@@ -72,6 +74,23 @@ func (s *Server) SetToolAuthorizer(authorizer func(context.Context, string, map[
 	s.mu.Unlock()
 }
 
+// SetToolGuard installs the runtime safety rules shared by HTTP and internal calls.
+func (s *Server) SetToolGuard(guard *toolguard.Manager) {
+	if s == nil {
+		return
+	}
+	s.mu.Lock()
+	s.toolGuard = guard
+	s.mu.Unlock()
+}
+
+func (s *Server) checkToolGuard(toolName string, args map[string]interface{}) *ToolResult {
+	s.mu.RLock()
+	guard := s.toolGuard
+	s.mu.RUnlock()
+	return toolGuardBlockedResult(guard, toolName, args)
+}
+
 type sseClient struct {
 	id   string
 	send chan []byte
@@ -566,7 +585,7 @@ func (s *Server) handleCallTool(requestCtx context.Context, msg *Message) *Messa
 			s.mu.Unlock()
 		}
 
-		s.updateStats(req.Name, true)
+		s.updateStats(req.Name, ToolExecutionStatusFailed)
 
 		return &Message{
 			ID:      msg.ID,
@@ -590,10 +609,13 @@ func (s *Server) handleCallTool(requestCtx context.Context, msg *Message) *Messa
 		zap.Any("arguments", req.Arguments),
 	)
 
-	result, err := handler(execCtx, req.Arguments)
+	result := s.checkToolGuard(req.Name, req.Arguments)
+	var err error
+	if result == nil {
+		result, err = handler(execCtx, req.Arguments)
+	}
 	cancelledWithUserNote := s.applyAbortUserNoteToCancelledToolResult(executionID, &result, &err)
 	now := time.Now()
-	var failed bool
 	var finalResult *ToolResult
 
 	s.mu.Lock()
@@ -604,13 +626,15 @@ func (s *Server) handleCallTool(requestCtx context.Context, msg *Message) *Messa
 		st, msg := executionStatusAndMessage(err)
 		execution.Status = st
 		execution.Error = msg
-		failed = st != "cancelled"
+	} else if result != nil && result.Blocked {
+		execution.Status = ToolExecutionStatusBlocked
+		execution.Error = firstToolResultText(result, toolGuardBlockedPrefix)
+		execution.Result = result
 	} else if result != nil && result.IsError {
 		if cancelledWithUserNote {
 			execution.Status = "cancelled"
 			execution.Error = ""
 			execution.Result = result
-			failed = false
 		} else {
 			execution.Status = "failed"
 			if len(result.Content) > 0 {
@@ -619,7 +643,6 @@ func (s *Server) handleCallTool(requestCtx context.Context, msg *Message) *Messa
 				execution.Error = "工具执行返回错误结果"
 			}
 			execution.Result = result
-			failed = true
 		}
 	} else {
 		execution.Status = "completed"
@@ -631,7 +654,6 @@ func (s *Server) handleCallTool(requestCtx context.Context, msg *Message) *Messa
 			}
 		}
 		execution.Result = result
-		failed = false
 	}
 
 	finalResult = execution.Result
@@ -643,7 +665,7 @@ func (s *Server) handleCallTool(requestCtx context.Context, msg *Message) *Messa
 		}
 	}
 
-	s.updateStats(req.Name, failed)
+	s.updateStats(req.Name, execution.Status)
 
 	if s.storage != nil {
 		s.mu.Lock()
@@ -683,6 +705,8 @@ func (s *Server) handleCallTool(requestCtx context.Context, msg *Message) *Messa
 		errorResult, _ := json.Marshal(CallToolResponse{
 			Content: finalResult.Content,
 			IsError: true,
+			Blocked: finalResult.Blocked,
+			Meta:    toolResultProtocolMeta(finalResult),
 		})
 		return &Message{
 			ID:      msg.ID,
@@ -719,15 +743,15 @@ func (s *Server) handleCallTool(requestCtx context.Context, msg *Message) *Messa
 }
 
 // updateStats 更新统计信息
-func (s *Server) updateStats(toolName string, failed bool) {
+func (s *Server) updateStats(toolName string, status string) {
 	now := time.Now()
 	if s.storage != nil {
 		totalCalls := 1
 		successCalls := 0
 		failedCalls := 0
-		if failed {
+		if executionStatusCountsAsFailed(status) {
 			failedCalls = 1
-		} else {
+		} else if status == ToolExecutionStatusCompleted {
 			successCalls = 1
 		}
 		if err := s.storage.UpdateToolStats(toolName, totalCalls, successCalls, failedCalls, &now); err != nil {
@@ -749,10 +773,12 @@ func (s *Server) updateStats(toolName string, failed bool) {
 	stats.TotalCalls++
 	stats.LastCallTime = &now
 
-	if failed {
+	if executionStatusCountsAsFailed(status) {
 		stats.FailedCalls++
-	} else {
+	} else if status == ToolExecutionStatusCompleted {
 		stats.SuccessCalls++
+	} else if status == ToolExecutionStatusBlocked {
+		stats.BlockedCalls++
 	}
 }
 
@@ -925,11 +951,15 @@ func (s *Server) CallTool(ctx context.Context, toolName string, args map[string]
 			if !exists {
 				return nil, fmt.Errorf("工具 %s 未找到", toolName)
 			}
+			if blocked := s.checkToolGuard(toolName, args); blocked != nil {
+				return blocked, nil
+			}
 			return handler(runCtx, args)
 		},
 		OnDone: func(exec *ToolExecution) {
-			failed := exec != nil && exec.Status != ToolExecutionStatusCompleted && exec.Status != ToolExecutionStatusCancelled
-			s.updateStats(toolName, failed)
+			if exec != nil {
+				s.updateStats(toolName, exec.Status)
+			}
 		},
 	})
 	if err != nil {
@@ -1111,7 +1141,7 @@ func (s *Server) FinishToolExecution(ctx context.Context, executionID, toolName
 		}
 	}
 
-	s.updateStats(exec.ToolName, failed)
+	s.updateStats(exec.ToolName, exec.Status)
 
 	if s.storage != nil {
 		s.mu.Lock()
@@ -1155,6 +1185,11 @@ func (s *Server) UpdateToolExecutionResult(executionID string, result *ToolResul
 	if executionID == "" || result == nil {
 		return nil
 	}
+	if previous, ok := s.GetExecution(executionID); ok && previous != nil &&
+		(previous.Status == ToolExecutionStatusBlocked || previous.Result != nil && previous.Result.Blocked) {
+		result = cloneToolResult(result)
+		result.Blocked, result.IsError = true, true
+	}
 	s.mu.Lock()
 	spill := ToolResultSpillConfig{
 		RootDir:     s.spillRootDir,
@@ -1270,6 +1305,9 @@ func (s *Server) applyAbortUserNoteToCancelledToolResult(executionID string, res
 	}
 	hasErr := err != nil && *err != nil
 	hasRes := result != nil && *result != nil
+	if hasRes && (*result).Blocked {
+		return false
+	}
 	if !hasErr && !hasRes {
 		return false
 	}