Add files via upload

2026-06-23 22:40:05 +02:00 · 2026-06-23 15:35:12 +08:00
parent a52452ceea
commit b366dc0287
3 changed files with 214 additions and 0 deletions
@@ -27,6 +27,7 @@ type Config struct {
 	Database    DatabaseConfig        `yaml:"database"`
 	Auth        AuthConfig            `yaml:"auth"`
 	Audit       AuditConfig           `yaml:"audit,omitempty" json:"audit,omitempty"`
+	Monitor     MonitorConfig         `yaml:"monitor,omitempty" json:"monitor,omitempty"`
 	ExternalMCP ExternalMCPConfig     `yaml:"external_mcp,omitempty"`
 	Knowledge   KnowledgeConfig       `yaml:"knowledge,omitempty"`
 	C2          C2Config              `yaml:"c2,omitempty" json:"c2,omitempty"` // 内置 C2 总开关；未配置时默认启用
@@ -623,6 +624,23 @@ type AuthConfig struct {
 	GeneratedPasswordPersistErr string `yaml:"-" json:"-"`
 }

+// MonitorConfig MCP 状态监控（tool_executions）保留策略。
+type MonitorConfig struct {
+	// RetentionDays 执行记录保留天数；省略时默认 90；0 表示不自动清理。
+	RetentionDays *int `yaml:"retention_days,omitempty" json:"retention_days,omitempty"`
+}
+
+// RetentionDaysEffective returns retention; 0 means keep forever; omitted defaults to 90.
+func (m MonitorConfig) RetentionDaysEffective() int {
+	if m.RetentionDays == nil {
+		return 90
+	}
+	if *m.RetentionDays < 0 {
+		return 0
+	}
+	return *m.RetentionDays
+}
+
 // AuditConfig platform operation audit log settings (not chat/tool execution bodies).
 type AuditConfig struct {
 	// Enabled nil or true enables persistence; explicit false disables.
@@ -1274,6 +1292,10 @@ func Default() *Config {
 				Enabled:        &on,
 			}
 		}(),
+		Monitor: func() MonitorConfig {
+			days := 90
+			return MonitorConfig{RetentionDays: &days}
+		}(),
 		Robots: RobotsConfig{
 			Session: RobotSessionConfig{
 				StrictUserIdentity: &strictRobotIdentity,
@@ -410,6 +410,76 @@ func (db *DB) GetToolExecutionsByIds(ids []string) ([]*mcp.ToolExecution, error)
 	return executions, nil
 }

+type toolExecutionStatDelta struct {
+	totalCalls   int
+	successCalls int
+	failedCalls  int
+}
+
+// PurgeToolExecutionsBefore deletes executions older than cutoff and adjusts tool_stats.
+func (db *DB) PurgeToolExecutionsBefore(cutoff time.Time) (int64, error) {
+	query := `
+		SELECT tool_name, status, COUNT(*) AS cnt
+		FROM tool_executions
+		WHERE ` + sqliteEpochGE("start_time", "<") + `
+		GROUP BY tool_name, status
+	`
+	rows, err := db.Query(query, formatSQLiteUTC(cutoff))
+	if err != nil {
+		return 0, err
+	}
+	defer rows.Close()
+
+	deltas := make(map[string]*toolExecutionStatDelta)
+	for rows.Next() {
+		var toolName, status string
+		var count int
+		if err := rows.Scan(&toolName, &status, &count); err != nil {
+			db.logger.Warn("读取待清理执行记录统计失败", zap.Error(err))
+			continue
+		}
+		toolName = strings.TrimSpace(toolName)
+		if toolName == "" || count <= 0 {
+			continue
+		}
+		delta := deltas[toolName]
+		if delta == nil {
+			delta = &toolExecutionStatDelta{}
+			deltas[toolName] = delta
+		}
+		delta.totalCalls += count
+		switch status {
+		case "failed", "cancelled":
+			delta.failedCalls += count
+		case "completed":
+			delta.successCalls += count
+		}
+	}
+	if err := rows.Err(); err != nil {
+		return 0, err
+	}
+
+	res, err := db.Exec(`DELETE FROM tool_executions WHERE `+sqliteEpochGE("start_time", "<"), formatSQLiteUTC(cutoff))
+	if err != nil {
+		return 0, err
+	}
+	deleted, err := res.RowsAffected()
+	if err != nil {
+		return 0, err
+	}
+
+	for toolName, delta := range deltas {
+		if err := db.DecreaseToolStats(toolName, delta.totalCalls, delta.successCalls, delta.failedCalls); err != nil {
+			db.logger.Warn("清理过期执行记录后更新统计失败",
+				zap.Error(err),
+				zap.String("toolName", toolName),
+			)
+		}
+	}
+
+	return deleted, nil
+}
+
 // SaveToolStats 保存工具统计信息
 func (db *DB) SaveToolStats(toolName string, stats *mcp.ToolStats) error {
 	var lastCallTime sql.NullTime
@@ -0,0 +1,122 @@
+package database
+
+import (
+	"path/filepath"
+	"testing"
+	"time"
+
+	"cyberstrike-ai/internal/mcp"
+
+	"go.uber.org/zap"
+)
+
+func TestPurgeToolExecutionsBefore(t *testing.T) {
+	dbPath := filepath.Join(t.TempDir(), "monitor.db")
+	db, err := NewDB(dbPath, zap.NewNop())
+	if err != nil {
+		t.Fatalf("NewDB: %v", err)
+	}
+	defer db.Close()
+
+	oldStart := time.Now().AddDate(0, 0, -100)
+	newStart := time.Now().AddDate(0, 0, -1)
+
+	oldExec := &mcp.ToolExecution{
+		ID:        "old-completed",
+		ToolName:  "nmap::scan",
+		Arguments: map[string]interface{}{"target": "127.0.0.1"},
+		Status:    "completed",
+		StartTime: oldStart,
+	}
+	oldFailed := &mcp.ToolExecution{
+		ID:        "old-failed",
+		ToolName:  "nmap::scan",
+		Arguments: map[string]interface{}{"target": "127.0.0.1"},
+		Status:    "failed",
+		Error:     "timeout",
+		StartTime: oldStart,
+	}
+	newExec := &mcp.ToolExecution{
+		ID:        "new-completed",
+		ToolName:  "nmap::scan",
+		Arguments: map[string]interface{}{"target": "127.0.0.1"},
+		Status:    "completed",
+		StartTime: newStart,
+	}
+	for _, exec := range []*mcp.ToolExecution{oldExec, oldFailed, newExec} {
+		if err := db.SaveToolExecution(exec); err != nil {
+			t.Fatalf("SaveToolExecution(%s): %v", exec.ID, err)
+		}
+	}
+	if err := db.UpdateToolStats("nmap::scan", 3, 2, 1, &newStart); err != nil {
+		t.Fatalf("UpdateToolStats: %v", err)
+	}
+
+	cutoff := time.Now().AddDate(0, 0, -90)
+	deleted, err := db.PurgeToolExecutionsBefore(cutoff)
+	if err != nil {
+		t.Fatalf("PurgeToolExecutionsBefore: %v", err)
+	}
+	if deleted != 2 {
+		t.Fatalf("deleted = %d, want 2", deleted)
+	}
+
+	if _, err := db.GetToolExecution("old-completed"); err == nil {
+		t.Fatal("old-completed should be deleted")
+	}
+	if _, err := db.GetToolExecution("old-failed"); err == nil {
+		t.Fatal("old-failed should be deleted")
+	}
+	if _, err := db.GetToolExecution("new-completed"); err != nil {
+		t.Fatalf("new-completed should remain: %v", err)
+	}
+
+	stats, err := db.LoadToolStats()
+	if err != nil {
+		t.Fatalf("LoadToolStats: %v", err)
+	}
+	stat := stats["nmap::scan"]
+	if stat == nil {
+		t.Fatal("expected stats for nmap::scan")
+	}
+	if stat.TotalCalls != 1 || stat.SuccessCalls != 1 || stat.FailedCalls != 0 {
+		t.Fatalf("stats after purge = %+v, want total=1 success=1 failed=0", stat)
+	}
+
+	total, err := db.CountToolExecutions("", "")
+	if err != nil {
+		t.Fatalf("CountToolExecutions: %v", err)
+	}
+	if total != 1 {
+		t.Fatalf("remaining executions = %d, want 1", total)
+	}
+}
+
+func TestPurgeToolExecutionsBefore_zeroRetentionSkipsViaService(t *testing.T) {
+	// RetentionDaysEffective: 0 means no purge at service layer; DB method still works when called directly.
+	dbPath := filepath.Join(t.TempDir(), "monitor.db")
+	db, err := NewDB(dbPath, zap.NewNop())
+	if err != nil {
+		t.Fatalf("NewDB: %v", err)
+	}
+	defer db.Close()
+
+	exec := &mcp.ToolExecution{
+		ID:        "ancient",
+		ToolName:  "curl::get",
+		Arguments: map[string]interface{}{},
+		Status:    "completed",
+		StartTime: time.Now().AddDate(-1, 0, 0),
+	}
+	if err := db.SaveToolExecution(exec); err != nil {
+		t.Fatalf("SaveToolExecution: %v", err)
+	}
+
+	deleted, err := db.PurgeToolExecutionsBefore(time.Now())
+	if err != nil {
+		t.Fatalf("PurgeToolExecutionsBefore: %v", err)
+	}
+	if deleted != 1 {
+		t.Fatalf("deleted = %d, want 1", deleted)
+	}
+}