Add files via upload

internal/database/vulnerability.go

@@ -191,6 +191,7 @@ func (db *DB) CreateVulnerability(vuln *Vulnerability) (*Vulnerability, error) {
 	if err != nil {
 		return nil, fmt.Errorf("创建漏洞失败: %w", err)
 	}
+	db.refreshAssetRiskCacheForConversationsBestEffort(vuln.ConversationID)
 	return vuln, nil
 }
 
@@ -299,6 +300,8 @@ func (db *DB) CountVulnerabilitiesForAccess(filter VulnerabilityListFilter, acce
 // UpdateVulnerability 更新漏洞
 func (db *DB) UpdateVulnerability(id string, vuln *Vulnerability) error {
 	vuln.UpdatedAt = time.Now()
+	var oldConversationID string
+	_ = db.QueryRow(`SELECT COALESCE(conversation_id,'') FROM vulnerabilities WHERE id = ?`, id).Scan(&oldConversationID)
 
 	query := `
 		UPDATE vulnerabilities
@@ -318,6 +321,7 @@ func (db *DB) UpdateVulnerability(id string, vuln *Vulnerability) error {
 		return fmt.Errorf("更新漏洞失败: %w", err)
 	}
 
+	db.refreshAssetRiskCacheForConversationsBestEffort(oldConversationID, vuln.ConversationID)
 	return nil
 }
 
@@ -337,6 +341,10 @@ func (db *DB) DeleteVulnerabilitiesByFilterForAccess(filter VulnerabilityListFil
 	args := []interface{}{}
 	where, args = filter.appendWhere(where, args)
 	where, args = appendVulnerabilityAccessFilter(where, args, access)
+	affectedConversations, err := collectVulnerabilityConversationIDs(tx, where, args)
+	if err != nil {
+		return 0, err
+	}
 
 	clearQuery := `UPDATE project_facts SET related_vulnerability_id = NULL
 		WHERE related_vulnerability_id IN (SELECT id FROM vulnerabilities ` + where + `)`
@@ -356,6 +364,7 @@ func (db *DB) DeleteVulnerabilitiesByFilterForAccess(filter VulnerabilityListFil
 	if err := tx.Commit(); err != nil {
 		return 0, fmt.Errorf("提交事务失败: %w", err)
 	}
+	db.refreshAssetRiskCacheForConversationsBestEffort(affectedConversations...)
 	return deleted, nil
 }
 
@@ -366,6 +375,8 @@ func (db *DB) DeleteVulnerability(id string) error {
 		return fmt.Errorf("开启事务失败: %w", err)
 	}
 	defer func() { _ = tx.Rollback() }()
+	var conversationID string
+	_ = tx.QueryRow(`SELECT COALESCE(conversation_id,'') FROM vulnerabilities WHERE id = ?`, id).Scan(&conversationID)
 
 	// 删除漏洞前先解除项目事实中的关联，避免前端继续显示已删除漏洞的短 ID。
 	if _, err := tx.Exec("UPDATE project_facts SET related_vulnerability_id = NULL WHERE related_vulnerability_id = ?", id); err != nil {
@@ -377,9 +388,29 @@ func (db *DB) DeleteVulnerability(id string) error {
 	if err := tx.Commit(); err != nil {
 		return fmt.Errorf("提交事务失败: %w", err)
 	}
+	db.refreshAssetRiskCacheForConversationsBestEffort(conversationID)
 	return nil
 }
 
+func collectVulnerabilityConversationIDs(tx *sql.Tx, where string, args []interface{}) ([]string, error) {
+	rows, err := tx.Query(`SELECT DISTINCT COALESCE(conversation_id,'') FROM vulnerabilities `+where, args...)
+	if err != nil {
+		return nil, fmt.Errorf("查询受影响漏洞会话失败: %w", err)
+	}
+	defer rows.Close()
+	ids := []string{}
+	for rows.Next() {
+		var id string
+		if err := rows.Scan(&id); err != nil {
+			return nil, err
+		}
+		if strings.TrimSpace(id) != "" {
+			ids = append(ids, id)
+		}
+	}
+	return ids, rows.Err()
+}
+
 // GetVulnerabilityStats 获取漏洞统计（筛选条件与 ListVulnerabilities / CountVulnerabilities 一致）
 func (db *DB) GetVulnerabilityStats(filter VulnerabilityListFilter) (map[string]interface{}, error) {
 	return db.GetVulnerabilityStatsForAccess(filter, RBACListAccess{})