CalvinBackup/CyberStrikeAI
1
0
Fork 0
mirror of https://github.com/Ed1s0nZ/CyberStrikeAI.git synced 2026-09-13 21:29:03 +02:00
Code Issues Packages Projects Releases Wiki Activity

feat: add configurable tool call blocking and monitoring

Browse Source
This commit is contained in:
Ed1s0nZ
2026-09-08 09:44:32 +08:00
parent c70da22de7
commit 6ad9ea2d13
54 changed files with 4635 additions and 152 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/database/blocked_execution_test.go
+120
View File
@@ -0,0 +1,120 @@
package database
import (
"fmt"
"path/filepath"
"testing"
"time"
"cyberstrike-ai/internal/mcp"
"go.uber.org/zap"
)
func TestBlockedExecutionPersistenceStatsAndReconciliation(t *testing.T) {
db, conversationID, _ := setupProcessDetailsSummaryTest(t)
now := time.Now()
for _, status := range []string{"completed", "failed", "blocked", "cancelled"} {
result := &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: "policy message"}}, IsError: status != "completed", Blocked: status == "blocked"}
if err := db.SaveToolExecution(&mcp.ToolExecution{ID: status, ToolName: "test", Status: status, Result: result, StartTime: now.Add(-time.Minute), EndTime: &now, ConversationID: conversationID}); err != nil {
t.Fatal(err)
}
}
if err := db.UpdateToolStats("test", 4, 1, 1, &now); err != nil {
t.Fatal(err)
}
if err := db.UpdateToolExecutionResult("blocked", &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: "reduced output"}}}); err != nil {
t.Fatal(err)
}
reloaded, err := db.GetToolExecution("blocked")
if err != nil || reloaded.Status != "blocked" || !reloaded.Result.Blocked || !reloaded.Result.IsError || reloaded.Result.Content[0].Text != "reduced output" {
t.Fatalf("reduction/storage lost blocked classification: %#v err=%v", reloaded, err)
}
count, err := db.CancelOrphanedRunningToolExecutions(now, "restart")
if err != nil || count != 0 {
t.Fatalf("terminal blocks reclassified as orphaned: count=%d err=%v", count, err)
}
page, err := db.LoadToolExecutionListPage(0, 10, "blocked", "")
if err != nil || len(page) != 1 || page[0].ID != "blocked" {
t.Fatalf("blocked status filter failed: %#v err=%v", page, err)
}
summary, err := db.LoadToolStatsSummary(1)
if err != nil {
t.Fatal(err)
}
if summary.Summary.TotalCalls != 4 || summary.Summary.SuccessCalls != 1 || summary.Summary.FailedCalls != 1 || summary.Summary.BlockedCalls != 1 || summary.TopTools[0].BlockedCalls != 1 {
t.Fatalf("incorrect summary: %#v top=%#v", summary.Summary, summary.TopTools)
}
stats, err := db.LoadToolStats()
if err != nil || stats["test"].BlockedCalls != 1 || stats["test"].FailedCalls != 1 {
t.Fatalf("incorrect legacy stats: %#v err=%v", stats, err)
}
for _, daily := range []bool{false, true} {
buckets, err := db.LoadCallsTimeline(now.Add(-time.Hour), daily)
if err != nil || len(buckets) != 1 || buckets[0].Total != 4 || buckets[0].Failed != 1 || buckets[0].Blocked != 1 {
t.Fatalf("incorrect timeline daily=%v: %#v err=%v", daily, buckets, err)
}
}
}
func TestLegacyToolGuardBlockMigrationIsStrictAndIdempotent(t *testing.T) {
path := filepath.Join(t.TempDir(), "legacy-guard.db")
db, err := NewDB(path, zap.NewNop())
if err != nil {
t.Fatal(err)
}
now := time.Now()
refusal := "工具调用已被安全规则拦截：识别到 example.gov，禁止操作。\n规则: 政府网站保护 (government-domains)\n匹配内容: \"example.gov\""
for i, reason := range []string{
refusal,
"upstream returned: " + refusal,
"工具调用已被安全规则拦截：regular error without the envelope",
"工具调用已被安全规则拦截：malformed match\n规则: Rule (id)\n匹配内容: unquoted",
} {
if err := db.SaveToolExecution(&mcp.ToolExecution{ID: fmt.Sprint(i), ToolName: "test", Status: "failed", Error: reason, StartTime: now, EndTime: &now}); err != nil {
t.Fatal(err)
}
}
if err := db.UpdateToolStats("test", 4, 0, 4, &now); err != nil {
t.Fatal(err)
}
if err := db.Close(); err != nil {
t.Fatal(err)
}
for run := 0; run < 2; run++ {
db, err = NewDB(path, zap.NewNop())
if err != nil {
t.Fatal(err)
}
exec, err := db.GetToolExecution("0")
if err != nil || exec.Status != "blocked" || !exec.Result.Blocked || !exec.Result.IsError || exec.Result.Content[0].Text != refusal {
t.Fatalf("migration did not retain refusal: %#v err=%v", exec, err)
}
stats, err := db.LoadToolStats()
if err != nil || stats["test"].TotalCalls != 4 || stats["test"].FailedCalls != 3 || stats["test"].BlockedCalls != 1 {
t.Fatalf("migration run=%d stats=%#v err=%v", run, stats, err)
}
count, err := db.CountToolExecutions("failed", "")
if err != nil || count != 3 {
t.Fatalf("migration changed unrelated failures: count=%d err=%v", count, err)
}
if err := db.Close(); err != nil {
t.Fatal(err)
}
}
}
func TestToolResultStatusFromPayloadDistinguishesBlocked(t *testing.T) {
for _, tc := range []struct {
payload map[string]interface{}
want string
}{
{map[string]interface{}{"blocked": true, "success": false, "isError": true}, "blocked"},
{map[string]interface{}{"status": "blocked", "success": false}, "blocked"},
{map[string]interface{}{"success": false, "isError": true, "result": "工具调用已被安全规则拦截"}, "failed"},
{map[string]interface{}{"success": true}, "completed"},
} {
if got := toolResultStatusFromPayload(tc.payload, "tool_result"); got != tc.want {
t.Fatalf("payload=%#v status=%s want=%s", tc.payload, got, tc.want)
}
}
}
internal/database/conversation.go
+3
View File
@@ -1637,6 +1637,9 @@ func toolResultStatusFromPayload(payload map[string]interface{}, eventType strin
if eventType != "tool_result" {
return ""
}
if blocked, _ := payload["blocked"].(bool); blocked || strings.EqualFold(processDetailString(payload, "status"), "blocked") {
return "blocked"
}
if status := processDetailString(payload, "status"); strings.EqualFold(status, "background_running") {
return "background_running"
}
internal/database/database.go
+4
View File
@@ -155,6 +155,10 @@ func NewDB(dbPath string, logger *zap.Logger) (*DB, error) {
_ = db.Close()
return nil, fmt.Errorf("初始化表失败: %w", err)
}
if err := database.migrateLegacyToolGuardBlocks(); err != nil {
_ = db.Close()
return nil, fmt.Errorf("迁移历史安全拦截记录失败: %w", err)
}
database.startPassiveCheckpointLoop("conversations")
return database, nil
internal/database/monitor.go
+33 -9
View File
@@ -91,6 +91,15 @@ func (db *DB) UpdateToolExecutionResult(id string, result *mcp.ToolResult) error
if id == "" || result == nil {
return nil
}
var status string
if err := db.QueryRow(`SELECT status FROM tool_executions WHERE id = ?`, id).Scan(&status); err != nil && err != sql.ErrNoRows {
return err
}
if status == mcp.ToolExecutionStatusBlocked {
copy := *result
copy.Blocked, copy.IsError = true, true
result = &copy
}
resultBytes, err := json.Marshal(result)
if err != nil {
return err
@@ -276,6 +285,7 @@ type ToolStatsSummary struct {
TotalCalls int
SuccessCalls int
FailedCalls int
BlockedCalls int
LastCallTime *time.Time
ToolCount int
}
@@ -304,6 +314,7 @@ func (db *DB) LoadToolStatsSummary(topN int) (*ToolStatsSummaryResult, error) {
SELECT COUNT(*),
COALESCE(SUM(CASE WHEN status = 'completed' THEN 1 ELSE 0 END), 0),
COALESCE(SUM(CASE WHEN status IN ('failed', 'hard_timeout', 'orphaned') THEN 1 ELSE 0 END), 0),
COALESCE(SUM(CASE WHEN status = 'blocked' THEN 1 ELSE 0 END), 0),
MAX(start_time),
COUNT(DISTINCT tool_name)
FROM tool_executions
@@ -313,6 +324,7 @@ func (db *DB) LoadToolStatsSummary(topN int) (*ToolStatsSummaryResult, error) {
&result.Summary.TotalCalls,
&result.Summary.SuccessCalls,
&result.Summary.FailedCalls,
&result.Summary.BlockedCalls,
&lastCallRaw,
&result.Summary.ToolCount,
)
@@ -334,6 +346,7 @@ func (db *DB) LoadToolStatsSummary(topN int) (*ToolStatsSummaryResult, error) {
COUNT(*) AS total_calls,
SUM(CASE WHEN status = 'completed' THEN 1 ELSE 0 END) AS success_calls,
SUM(CASE WHEN status IN ('failed', 'hard_timeout', 'orphaned') THEN 1 ELSE 0 END) AS failed_calls,
SUM(CASE WHEN status = 'blocked' THEN 1 ELSE 0 END) AS blocked_calls,
MAX(start_time) AS last_call_time
FROM tool_executions
GROUP BY tool_name
@@ -354,6 +367,7 @@ func (db *DB) LoadToolStatsSummary(topN int) (*ToolStatsSummaryResult, error) {
&stat.TotalCalls,
&stat.SuccessCalls,
&stat.FailedCalls,
&stat.BlockedCalls,
&lastCallTime,
); err != nil {
db.logger.Warn("加载 Top 工具统计失败", zap.Error(err))
@@ -385,8 +399,9 @@ func (db *DB) LoadToolStatsSummaryForAccess(topN int, access RBACListAccess) (*T
err := db.QueryRow(`SELECT COUNT(*),
COALESCE(SUM(CASE WHEN status = 'completed' THEN 1 ELSE 0 END), 0),
COALESCE(SUM(CASE WHEN status IN ('failed', 'hard_timeout', 'orphaned') THEN 1 ELSE 0 END), 0),
COALESCE(SUM(CASE WHEN status = 'blocked' THEN 1 ELSE 0 END), 0),
MAX(start_time), COUNT(DISTINCT tool_name)`+fromSQL, args...).Scan(
&result.Summary.TotalCalls, &result.Summary.SuccessCalls, &result.Summary.FailedCalls,
&result.Summary.TotalCalls, &result.Summary.SuccessCalls, &result.Summary.FailedCalls, &result.Summary.BlockedCalls,
&lastCall, &result.Summary.ToolCount,
)
if err != nil {
@@ -398,7 +413,8 @@ func (db *DB) LoadToolStatsSummaryForAccess(topN int, access RBACListAccess) (*T
}
rows, err := db.Query(`SELECT tool_name, COUNT(*),
SUM(CASE WHEN status = 'completed' THEN 1 ELSE 0 END),
SUM(CASE WHEN status IN ('failed', 'hard_timeout', 'orphaned') THEN 1 ELSE 0 END), MAX(start_time)`+
SUM(CASE WHEN status IN ('failed', 'hard_timeout', 'orphaned') THEN 1 ELSE 0 END),
SUM(CASE WHEN status = 'blocked' THEN 1 ELSE 0 END), MAX(start_time)`+
fromSQL+` GROUP BY tool_name ORDER BY COUNT(*) DESC, tool_name ASC LIMIT ?`, append(args, topN)...)
if err != nil {
return nil, err
@@ -407,7 +423,7 @@ func (db *DB) LoadToolStatsSummaryForAccess(topN int, access RBACListAccess) (*T
for rows.Next() {
var stat mcp.ToolStats
var last sql.NullString
if err := rows.Scan(&stat.ToolName, &stat.TotalCalls, &stat.SuccessCalls, &stat.FailedCalls, &last); err != nil {
if err := rows.Scan(&stat.ToolName, &stat.TotalCalls, &stat.SuccessCalls, &stat.FailedCalls, &stat.BlockedCalls, &last); err != nil {
return nil, err
}
if last.Valid {
@@ -916,8 +932,11 @@ func (db *DB) SaveToolStats(toolName string, stats *mcp.ToolStats) error {
// LoadToolStats 加载所有工具统计信息
func (db *DB) LoadToolStats() (map[string]*mcp.ToolStats, error) {
query := `
SELECT tool_name, total_calls, success_calls, failed_calls, last_call_time
FROM tool_stats
SELECT stats.tool_name, total_calls, success_calls, failed_calls, last_call_time,
COALESCE(blocked.calls, 0)
FROM tool_stats stats
LEFT JOIN (SELECT tool_name, COUNT(*) AS calls FROM tool_executions WHERE status = 'blocked' GROUP BY tool_name) blocked
ON blocked.tool_name = stats.tool_name
`
rows, err := db.Query(query)
@@ -937,6 +956,7 @@ func (db *DB) LoadToolStats() (map[string]*mcp.ToolStats, error) {
&stat.SuccessCalls,
&stat.FailedCalls,
&lastCallTime,
&stat.BlockedCalls,
)
if err != nil {
db.logger.Warn("加载统计信息失败", zap.Error(err))
@@ -989,6 +1009,7 @@ type CallsTimelineBucket struct {
BucketTime time.Time
Total int
Failed int
Blocked int
}
// truncateCallsTimelineBucket 将时间截断到趋势图桶边界（本地时区，与 handler 侧 truncateToBucket 一致）
@@ -1008,7 +1029,8 @@ func (db *DB) LoadCallsTimeline(since time.Time, dailyBuckets bool) ([]CallsTime
query = `
SELECT date(start_time, 'localtime') AS bucket,
COUNT(*) AS total,
SUM(CASE WHEN status IN ('failed', 'hard_timeout', 'orphaned') THEN 1 ELSE 0 END) AS failed
SUM(CASE WHEN status IN ('failed', 'hard_timeout', 'orphaned') THEN 1 ELSE 0 END) AS failed,
SUM(CASE WHEN status = 'blocked' THEN 1 ELSE 0 END) AS blocked
FROM tool_executions
WHERE start_time >= ?
GROUP BY bucket
@@ -1018,7 +1040,8 @@ func (db *DB) LoadCallsTimeline(since time.Time, dailyBuckets bool) ([]CallsTime
query = `
SELECT strftime('%Y-%m-%d %H:00:00', start_time, 'localtime') AS bucket,
COUNT(*) AS total,
SUM(CASE WHEN status IN ('failed', 'hard_timeout', 'orphaned') THEN 1 ELSE 0 END) AS failed
SUM(CASE WHEN status IN ('failed', 'hard_timeout', 'orphaned') THEN 1 ELSE 0 END) AS failed,
SUM(CASE WHEN status = 'blocked' THEN 1 ELSE 0 END) AS blocked
FROM tool_executions
WHERE start_time >= ?
GROUP BY bucket
@@ -1035,8 +1058,8 @@ func (db *DB) LoadCallsTimeline(since time.Time, dailyBuckets bool) ([]CallsTime
buckets := make([]CallsTimelineBucket, 0)
for rows.Next() {
var bucketStr string
var total, failed int
if err := rows.Scan(&bucketStr, &total, &failed); err != nil {
var total, failed, blocked int
if err := rows.Scan(&bucketStr, &total, &failed, &blocked); err != nil {
db.logger.Warn("加载调用趋势失败", zap.Error(err))
continue
}
@@ -1049,6 +1072,7 @@ func (db *DB) LoadCallsTimeline(since time.Time, dailyBuckets bool) ([]CallsTime
BucketTime: bucketTime,
Total: total,
Failed: failed,
Blocked: blocked,
})
}
return buckets, nil
internal/database/tool_guard_migration.go
+84
View File
@@ -0,0 +1,84 @@
package database
import (
"encoding/json"
"strconv"
"strings"
"cyberstrike-ai/internal/mcp"
)
const legacyToolGuardPrefix = "工具调用已被安全规则拦截"
// Only the exact envelope emitted by the old local guard is recognized here.
// New executions use the structured marker and never infer policy from text.
func isLegacyToolGuardRefusal(text string) bool {
if !strings.HasPrefix(text, legacyToolGuardPrefix+"：") && !strings.HasPrefix(text, legacyToolGuardPrefix+"\n规则: ") {
return false
}
matchIndex := strings.LastIndex(text, "\n匹配内容: ")
if matchIndex < 0 {
return false
}
if _, err := strconv.Unquote(text[matchIndex+len("\n匹配内容: "):]); err != nil {
return false
}
ruleIndex := strings.LastIndex(text[:matchIndex], "\n规则: ")
if ruleIndex < 0 {
return false
}
rule := text[ruleIndex+len("\n规则: ") : matchIndex]
idIndex := strings.LastIndex(rule, " (")
return idIndex > 0 && strings.HasSuffix(rule, ")") && len(rule[idIndex+2:len(rule)-1]) > 0 && !strings.Contains(rule, "\n")
}
// migrateLegacyToolGuardBlocks is idempotent because only failed records qualify.
// Keeping status and accumulated failure counts in one transaction makes monitor
// filters, badges and statistics agree immediately after upgrading.
func (db *DB) migrateLegacyToolGuardBlocks() error {
tx, err := db.Begin()
if err != nil {
return err
}
defer tx.Rollback()
rows, err := tx.Query(`SELECT id, tool_name, error, COALESCE(result, '') FROM tool_executions WHERE status = 'failed' AND error LIKE ?`, legacyToolGuardPrefix+"%")
if err != nil {
return err
}
type record struct{ id, tool, reason, result string }
var records []record
for rows.Next() {
var r record
if err := rows.Scan(&r.id, &r.tool, &r.reason, &r.result); err != nil {
rows.Close()
return err
}
if isLegacyToolGuardRefusal(r.reason) {
records = append(records, r)
}
}
err = rows.Err()
rows.Close()
if err != nil {
return err
}
for _, r := range records {
var result mcp.ToolResult
_ = json.Unmarshal([]byte(r.result), &result)
if len(result.Content) == 0 {
result.Content = []mcp.Content{{Type: "text", Text: r.reason}}
}
result.Blocked, result.IsError = true, true
encoded, err := json.Marshal(result)
if err != nil {
return err
}
if _, err := tx.Exec(`UPDATE tool_executions SET status = 'blocked', result = ? WHERE id = ?`, string(encoded), r.id); err != nil {
return err
}
if _, err := tx.Exec(`UPDATE tool_stats SET failed_calls = MAX(0, failed_calls - 1) WHERE tool_name = ?`, r.tool); err != nil {
return err
}
}
return tx.Commit()
}