Add files via upload

internal/database/batch_task.go

@@ -12,6 +12,7 @@ import (
 type BatchTaskQueueRow struct {
 	ID           string
 	Title        sql.NullString
+	Role         sql.NullString
 	Status       string
 	CreatedAt    time.Time
 	StartedAt    sql.NullTime
@@ -33,7 +34,7 @@ type BatchTaskRow struct {
 }
 
 // CreateBatchQueue 创建批量任务队列
-func (db *DB) CreateBatchQueue(queueID string, title string, tasks []map[string]interface{}) error {
+func (db *DB) CreateBatchQueue(queueID string, title string, role string, tasks []map[string]interface{}) error {
 	tx, err := db.Begin()
 	if err != nil {
 		return fmt.Errorf("开始事务失败: %w", err)
@@ -42,8 +43,8 @@ func (db *DB) CreateBatchQueue(queueID string, title string, tasks []map[string]
 
 	now := time.Now()
 	_, err = tx.Exec(
-		"INSERT INTO batch_task_queues (id, title, status, created_at, current_index) VALUES (?, ?, ?, ?, ?)",
+		"INSERT INTO batch_task_queues (id, title, role, status, created_at, current_index) VALUES (?, ?, ?, ?, ?, ?)",
-		queueID, title, "pending", now, 0,
+		queueID, title, role, "pending", now, 0,
 	)
 	if err != nil {
 		return fmt.Errorf("创建批量任务队列失败: %w", err)
@@ -77,9 +78,9 @@ func (db *DB) GetBatchQueue(queueID string) (*BatchTaskQueueRow, error) {
 	var row BatchTaskQueueRow
 	var createdAt string
 	err := db.QueryRow(
-		"SELECT id, title, status, created_at, started_at, completed_at, current_index FROM batch_task_queues WHERE id = ?",
+		"SELECT id, title, role, status, created_at, started_at, completed_at, current_index FROM batch_task_queues WHERE id = ?",
 		queueID,
-	).Scan(&row.ID, &row.Title, &row.Status, &createdAt, &row.StartedAt, &row.CompletedAt, &row.CurrentIndex)
+	).Scan(&row.ID, &row.Title, &row.Role, &row.Status, &createdAt, &row.StartedAt, &row.CompletedAt, &row.CurrentIndex)
 	if err == sql.ErrNoRows {
 		return nil, nil
 	}
@@ -103,7 +104,7 @@ func (db *DB) GetBatchQueue(queueID string) (*BatchTaskQueueRow, error) {
 // GetAllBatchQueues 获取所有批量任务队列
 func (db *DB) GetAllBatchQueues() ([]*BatchTaskQueueRow, error) {
 	rows, err := db.Query(
-		"SELECT id, title, status, created_at, started_at, completed_at, current_index FROM batch_task_queues ORDER BY created_at DESC",
+		"SELECT id, title, role, status, created_at, started_at, completed_at, current_index FROM batch_task_queues ORDER BY created_at DESC",
 	)
 	if err != nil {
 		return nil, fmt.Errorf("查询批量任务队列列表失败: %w", err)
@@ -114,7 +115,7 @@ func (db *DB) GetAllBatchQueues() ([]*BatchTaskQueueRow, error) {
 	for rows.Next() {
 		var row BatchTaskQueueRow
 		var createdAt string
-		if err := rows.Scan(&row.ID, &row.Title, &row.Status, &createdAt, &row.StartedAt, &row.CompletedAt, &row.CurrentIndex); err != nil {
+		if err := rows.Scan(&row.ID, &row.Title, &row.Role, &row.Status, &createdAt, &row.StartedAt, &row.CompletedAt, &row.CurrentIndex); err != nil {
 			return nil, fmt.Errorf("扫描批量任务队列失败: %w", err)
 		}
 		parsedTime, parseErr := time.Parse("2006-01-02 15:04:05", createdAt)
@@ -134,7 +135,7 @@ func (db *DB) GetAllBatchQueues() ([]*BatchTaskQueueRow, error) {
 
 // ListBatchQueues 列出批量任务队列（支持筛选和分页）
 func (db *DB) ListBatchQueues(limit, offset int, status, keyword string) ([]*BatchTaskQueueRow, error) {
-	query := "SELECT id, title, status, created_at, started_at, completed_at, current_index FROM batch_task_queues WHERE 1=1"
+	query := "SELECT id, title, role, status, created_at, started_at, completed_at, current_index FROM batch_task_queues WHERE 1=1"
 	args := []interface{}{}
 
 	// 状态筛选
@@ -162,7 +163,7 @@ func (db *DB) ListBatchQueues(limit, offset int, status, keyword string) ([]*Bat
 	for rows.Next() {
 		var row BatchTaskQueueRow
 		var createdAt string
-		if err := rows.Scan(&row.ID, &row.Title, &row.Status, &createdAt, &row.StartedAt, &row.CompletedAt, &row.CurrentIndex); err != nil {
+		if err := rows.Scan(&row.ID, &row.Title, &row.Role, &row.Status, &createdAt, &row.StartedAt, &row.CompletedAt, &row.CurrentIndex); err != nil {
 			return nil, fmt.Errorf("扫描批量任务队列失败: %w", err)
 		}
 		parsedTime, parseErr := time.Parse("2006-01-02 15:04:05", createdAt)

internal/database/database.go

@@ -433,7 +433,7 @@ func (db *DB) migrateConversationGroupMappingsTable() error {
 	return nil
 }
 
-// migrateBatchTaskQueuesTable 迁移batch_task_queues表，添加title字段
+// migrateBatchTaskQueuesTable 迁移batch_task_queues表，添加title和role字段
 func (db *DB) migrateBatchTaskQueuesTable() error {
 	// 检查title字段是否存在
 	var count int
@@ -454,6 +454,25 @@ func (db *DB) migrateBatchTaskQueuesTable() error {
 		}
 	}
 
+	// 检查role字段是否存在
+	var roleCount int
+	err = db.QueryRow("SELECT COUNT(*) FROM pragma_table_info('batch_task_queues') WHERE name='role'").Scan(&roleCount)
+	if err != nil {
+		// 如果查询失败，尝试添加字段
+		if _, addErr := db.Exec("ALTER TABLE batch_task_queues ADD COLUMN role TEXT"); addErr != nil {
+			// 如果字段已存在，忽略错误
+			errMsg := strings.ToLower(addErr.Error())
+			if !strings.Contains(errMsg, "duplicate column") && !strings.Contains(errMsg, "already exists") {
+				db.logger.Warn("添加role字段失败", zap.Error(addErr))
+			}
+		}
+	} else if roleCount == 0 {
+		// 字段不存在，添加它
+		if _, err := db.Exec("ALTER TABLE batch_task_queues ADD COLUMN role TEXT"); err != nil {
+			db.logger.Warn("添加role字段失败", zap.Error(err))
+		}
+	}
+
 	return nil
 }
 

internal/handler/agent.go

@@ -811,6 +811,7 @@ func (h *AgentHandler) ListCompletedTasks(c *gin.Context) {
 type BatchTaskRequest struct {
 	Title string   `json:"title"`                    // 任务标题（可选）
 	Tasks []string `json:"tasks" binding:"required"` // 任务列表，每行一个任务
+	Role  string   `json:"role,omitempty"`           // 角色名称（可选，空字符串表示默认角色）
 }
 
 // CreateBatchQueue 创建批量任务队列
@@ -839,7 +840,7 @@ func (h *AgentHandler) CreateBatchQueue(c *gin.Context) {
 		return
 	}
 
-	queue := h.batchTaskManager.CreateBatchQueue(req.Title, validTasks)
+	queue := h.batchTaskManager.CreateBatchQueue(req.Title, req.Role, validTasks)
 	c.JSON(http.StatusOK, gin.H{
 		"queueId": queue.ID,
 		"queue":   queue,
@@ -1095,7 +1096,27 @@ func (h *AgentHandler) executeBatchQueue(queueID string) {
 		// 保存conversationId到任务中（即使是运行中状态也要保存，以便查看对话）
 		h.batchTaskManager.UpdateTaskStatusWithConversationID(queueID, task.ID, "running", "", "", conversationID)
 
-		// 保存用户消息
+		// 应用角色用户提示词和工具配置
+		finalMessage := task.Message
+		var roleTools []string // 角色配置的工具列表
+		if queue.Role != "" && queue.Role != "默认" {
+			if h.config.Roles != nil {
+				if role, exists := h.config.Roles[queue.Role]; exists && role.Enabled {
+					// 应用用户提示词
+					if role.UserPrompt != "" {
+						finalMessage = role.UserPrompt + "\n\n" + task.Message
+						h.logger.Info("应用角色用户提示词", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.String("role", queue.Role))
+					}
+					// 获取角色配置的工具列表（优先使用tools字段，向后兼容mcps字段）
+					if len(role.Tools) > 0 {
+						roleTools = role.Tools
+						h.logger.Info("使用角色配置的工具列表", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.String("role", queue.Role), zap.Int("toolCount", len(roleTools)))
+					}
+				}
+			}
+		}
+
+		// 保存用户消息（保存原始消息，不包含角色提示词）
 		_, err = h.db.AddMessage(conversationID, "user", task.Message, nil)
 		if err != nil {
 			h.logger.Error("保存用户消息失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.String("conversationId", conversationID), zap.Error(err))
@@ -1116,14 +1137,14 @@ func (h *AgentHandler) executeBatchQueue(queueID string) {
 		}
 		progressCallback := h.createProgressCallback(conversationID, assistantMessageID, nil)
 
-		// 执行任务
+		// 执行任务（使用包含角色提示词的finalMessage和角色工具列表）
-		h.logger.Info("执行批量任务", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.String("message", task.Message), zap.String("conversationId", conversationID))
+		h.logger.Info("执行批量任务", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.String("message", task.Message), zap.String("role", queue.Role), zap.String("conversationId", conversationID))
 
 		ctx, cancel := context.WithTimeout(context.Background(), 30*time.Minute)
 		// 存储取消函数，以便在取消队列时能够取消当前任务
 		h.batchTaskManager.SetTaskCancel(queueID, cancel)
-		// 批量任务暂时不支持角色工具过滤，使用所有工具（传入nil）
+		// 使用队列配置的角色工具列表（如果为空，表示使用所有工具）
-		result, err := h.agent.AgentLoopWithProgress(ctx, task.Message, []agent.ChatMessage{}, conversationID, progressCallback, nil)
+		result, err := h.agent.AgentLoopWithProgress(ctx, finalMessage, []agent.ChatMessage{}, conversationID, progressCallback, roleTools)
 		// 任务执行完成，清理取消函数
 		h.batchTaskManager.SetTaskCancel(queueID, nil)
 		cancel()

internal/handler/batch_task_manager.go

@@ -29,6 +29,7 @@ type BatchTask struct {
 type BatchTaskQueue struct {
 	ID           string       `json:"id"`
 	Title        string       `json:"title,omitempty"`
+	Role         string       `json:"role,omitempty"` // 角色名称（空字符串表示默认角色）
 	Tasks        []*BatchTask `json:"tasks"`
 	Status       string       `json:"status"` // pending, running, paused, completed, cancelled
 	CreatedAt    time.Time    `json:"createdAt"`
@@ -62,7 +63,7 @@ func (m *BatchTaskManager) SetDB(db *database.DB) {
 }
 
 // CreateBatchQueue 创建批量任务队列
-func (m *BatchTaskManager) CreateBatchQueue(title string, tasks []string) *BatchTaskQueue {
+func (m *BatchTaskManager) CreateBatchQueue(title, role string, tasks []string) *BatchTaskQueue {
 	m.mu.Lock()
 	defer m.mu.Unlock()
 
@@ -70,6 +71,7 @@ func (m *BatchTaskManager) CreateBatchQueue(title string, tasks []string) *Batch
 	queue := &BatchTaskQueue{
 		ID:           queueID,
 		Title:        title,
+		Role:         role,
 		Tasks:        make([]*BatchTask, 0, len(tasks)),
 		Status:       "pending",
 		CreatedAt:    time.Now(),
@@ -98,7 +100,7 @@ func (m *BatchTaskManager) CreateBatchQueue(title string, tasks []string) *Batch
 
 	// 保存到数据库
 	if m.db != nil {
-		if err := m.db.CreateBatchQueue(queueID, title, dbTasks); err != nil {
+		if err := m.db.CreateBatchQueue(queueID, title, role, dbTasks); err != nil {
 			// 如果数据库保存失败，记录错误但继续（使用内存缓存）
 			// 这里可以添加日志记录
 		}
@@ -158,6 +160,9 @@ func (m *BatchTaskManager) loadQueueFromDB(queueID string) *BatchTaskQueue {
 	if queueRow.Title.Valid {
 		queue.Title = queueRow.Title.String
 	}
+	if queueRow.Role.Valid {
+		queue.Role = queueRow.Role.String
+	}
 	if queueRow.StartedAt.Valid {
 		queue.StartedAt = &queueRow.StartedAt.Time
 	}
@@ -351,6 +356,9 @@ func (m *BatchTaskManager) LoadFromDB() error {
 		if queueRow.Title.Valid {
 			queue.Title = queueRow.Title.String
 		}
+		if queueRow.Role.Valid {
+			queue.Role = queueRow.Role.String
+		}
 		if queueRow.StartedAt.Valid {
 			queue.StartedAt = &queueRow.StartedAt.Time
 		}

web/static/js/tasks.js

@@ -716,23 +716,56 @@ const batchQueuesState = {
     totalPages: 1
 };
 
-// 显示批量导入模态框
+// 显示新建任务模态框
-function showBatchImportModal() {
+async function showBatchImportModal() {
     const modal = document.getElementById('batch-import-modal');
     const input = document.getElementById('batch-tasks-input');
     const titleInput = document.getElementById('batch-queue-title');
+    const roleSelect = document.getElementById('batch-queue-role');
     if (modal && input) {
         input.value = '';
         if (titleInput) {
             titleInput.value = '';
         }
+        // 重置角色选择为默认
+        if (roleSelect) {
+            roleSelect.value = '';
+        }
         updateBatchImportStats('');
+        
+        // 加载并填充角色列表
+        if (roleSelect && typeof loadRoles === 'function') {
+            try {
+                const loadedRoles = await loadRoles();
+                // 清空现有选项（除了默认选项）
+                roleSelect.innerHTML = '<option value="">默认</option>';
+                
+                // 添加已启用的角色
+                const sortedRoles = loadedRoles.sort((a, b) => {
+                    if (a.name === '默认') return -1;
+                    if (b.name === '默认') return 1;
+                    return (a.name || '').localeCompare(b.name || '', 'zh-CN');
+                });
+                
+                sortedRoles.forEach(role => {
+                    if (role.name !== '默认' && role.enabled !== false) {
+                        const option = document.createElement('option');
+                        option.value = role.name;
+                        option.textContent = role.name;
+                        roleSelect.appendChild(option);
+                    }
+                });
+            } catch (error) {
+                console.error('加载角色列表失败:', error);
+            }
+        }
+        
         modal.style.display = 'block';
         input.focus();
     }
 }
 
-// 关闭批量导入模态框
+// 关闭新建任务模态框
 function closeBatchImportModal() {
     const modal = document.getElementById('batch-import-modal');
     if (modal) {
@@ -740,7 +773,7 @@ function closeBatchImportModal() {
     }
 }
 
-// 更新批量导入统计
+// 更新新建任务统计
 function updateBatchImportStats(text) {
     const statsEl = document.getElementById('batch-import-stats');
     if (!statsEl) return;
@@ -770,6 +803,7 @@ document.addEventListener('DOMContentLoaded', function() {
 async function createBatchQueue() {
     const input = document.getElementById('batch-tasks-input');
     const titleInput = document.getElementById('batch-queue-title');
+    const roleSelect = document.getElementById('batch-queue-role');
     if (!input) return;
     
     const text = input.value.trim();
@@ -788,13 +822,16 @@ async function createBatchQueue() {
     // 获取标题（可选）
     const title = titleInput ? titleInput.value.trim() : '';
     
+    // 获取角色（可选，空字符串表示默认角色）
+    const role = roleSelect ? roleSelect.value || '' : '';
+    
     try {
         const response = await apiFetch('/api/batch-tasks', {
             method: 'POST',
             headers: {
                 'Content-Type': 'application/json',
             },
-            body: JSON.stringify({ title, tasks }),
+            body: JSON.stringify({ title, tasks, role }),
         });
         
         if (!response.ok) {
@@ -816,6 +853,34 @@ async function createBatchQueue() {
     }
 }
 
+// 获取角色图标（辅助函数）
+function getRoleIconForDisplay(roleName, rolesList) {
+    if (!roleName || roleName === '') {
+        return '🔵'; // 默认角色图标
+    }
+    
+    if (Array.isArray(rolesList) && rolesList.length > 0) {
+        const role = rolesList.find(r => r.name === roleName);
+        if (role && role.icon) {
+            let icon = role.icon;
+            // 检查是否是 Unicode 转义格式（可能包含引号）
+            const unicodeMatch = icon.match(/^"?\\U([0-9A-F]{8})"?$/i);
+            if (unicodeMatch) {
+                try {
+                    const codePoint = parseInt(unicodeMatch[1], 16);
+                    icon = String.fromCodePoint(codePoint);
+                } catch (e) {
+                    // 转换失败，使用默认图标
+                    console.warn('转换 icon Unicode 转义失败:', icon, e);
+                    return '👤';
+                }
+            }
+            return icon;
+        }
+    }
+    return '👤'; // 默认图标
+}
+
 // 加载批量任务队列列表
 async function loadBatchQueues(page) {
     const section = document.getElementById('batch-queues-section');
@@ -826,6 +891,17 @@ async function loadBatchQueues(page) {
         batchQueuesState.currentPage = page;
     }
     
+    // 加载角色列表（用于显示正确的角色图标）
+    let loadedRoles = [];
+    if (typeof loadRoles === 'function') {
+        try {
+            loadedRoles = await loadRoles();
+        } catch (error) {
+            console.warn('加载角色列表失败，将使用默认图标:', error);
+        }
+    }
+    batchQueuesState.loadedRoles = loadedRoles; // 保存到状态中供渲染使用
+    
     // 构建查询参数
     const params = new URLSearchParams();
     params.append('page', batchQueuesState.currentPage.toString());
@@ -933,11 +1009,18 @@ function renderBatchQueues() {
         
         const titleDisplay = queue.title ? `<span class="batch-queue-title" style="font-weight: 600; color: var(--text-primary); margin-right: 8px;">${escapeHtml(queue.title)}</span>` : '';
         
+        // 显示角色信息（使用正确的角色图标）
+        const loadedRoles = batchQueuesState.loadedRoles || [];
+        const roleIcon = getRoleIconForDisplay(queue.role, loadedRoles);
+        const roleName = queue.role && queue.role !== '' ? queue.role : '默认';
+        const roleDisplay = `<span class="batch-queue-role" style="margin-right: 8px;" title="角色: ${escapeHtml(roleName)}">${roleIcon} ${escapeHtml(roleName)}</span>`;
+        
         return `
             <div class="batch-queue-item" data-queue-id="${queue.id}" onclick="showBatchQueueDetail('${queue.id}')">
                 <div class="batch-queue-header">
                     <div class="batch-queue-info" style="flex: 1;">
                         ${titleDisplay}
+                        ${roleDisplay}
                         <span class="batch-queue-status ${status.class}">${status.text}</span>
                         <span class="batch-queue-id">队列ID: ${escapeHtml(queue.id)}</span>
                         <span class="batch-queue-time">创建时间: ${new Date(queue.createdAt).toLocaleString('zh-CN')}</span>
@@ -1110,6 +1193,16 @@ async function showBatchQueueDetail(queueId) {
         if (!modal || !content) return;
         
         try {
+        // 加载角色列表（如果还未加载）
+        let loadedRoles = [];
+        if (typeof loadRoles === 'function') {
+            try {
+                loadedRoles = await loadRoles();
+            } catch (error) {
+                console.warn('加载角色列表失败，将使用默认图标:', error);
+            }
+        }
+        
         const response = await apiFetch(`/api/batch-tasks/${queueId}`);
         if (!response.ok) {
             throw new Error('获取队列详情失败');
@@ -1164,12 +1257,48 @@ async function showBatchQueueDetail(queueId) {
             'cancelled': { text: '已取消', class: 'batch-task-status-cancelled' }
         };
         
+        // 获取角色信息（如果队列有角色配置）
+        let roleDisplay = '';
+        if (queue.role && queue.role !== '') {
+            // 如果有角色配置，尝试获取角色详细信息
+            let roleName = queue.role;
+            let roleIcon = '👤';
+            // 从已加载的角色列表中查找角色图标
+            if (Array.isArray(loadedRoles) && loadedRoles.length > 0) {
+                const role = loadedRoles.find(r => r.name === roleName);
+                if (role && role.icon) {
+                    let icon = role.icon;
+                    const unicodeMatch = icon.match(/^"?\\U([0-9A-F]{8})"?$/i);
+                    if (unicodeMatch) {
+                        try {
+                            const codePoint = parseInt(unicodeMatch[1], 16);
+                            icon = String.fromCodePoint(codePoint);
+                        } catch (e) {
+                            // 转换失败，使用默认图标
+                        }
+                    }
+                    roleIcon = icon;
+                }
+            }
+            roleDisplay = `<div class="detail-item">
+                <span class="detail-label">角色</span>
+                <span class="detail-value">${roleIcon} ${escapeHtml(roleName)}</span>
+            </div>`;
+        } else {
+            // 默认角色
+            roleDisplay = `<div class="detail-item">
+                <span class="detail-label">角色</span>
+                <span class="detail-value">🔵 默认</span>
+            </div>`;
+        }
+        
         content.innerHTML = `
             <div class="batch-queue-detail-info">
                 ${queue.title ? `<div class="detail-item">
                     <span class="detail-label">任务标题</span>
                     <span class="detail-value">${escapeHtml(queue.title)}</span>
                 </div>` : ''}
+                ${roleDisplay}
                 <div class="detail-item">
                     <span class="detail-label">队列ID</span>
                     <span class="detail-value"><code>${escapeHtml(queue.id)}</code></span>

web/templates/index.html

@@ -611,7 +611,7 @@
                     <div class="page-header">
                         <h2>任务管理</h2>
                         <div class="page-header-actions">
-                            <button class="btn-primary" onclick="showBatchImportModal()">批量导入任务</button>
+                            <button class="btn-primary" onclick="showBatchImportModal()">新建任务</button>
                             <label class="auto-refresh-toggle">
                                 <input type="checkbox" id="tasks-auto-refresh" checked onchange="toggleTasksAutoRefresh(this.checked)">
                                 <span>自动刷新</span>
@@ -1238,11 +1238,11 @@
         </div>
     </div>
 
-    <!-- 批量导入任务模态框 -->
+    <!-- 新建任务模态框 -->
     <div id="batch-import-modal" class="modal">
         <div class="modal-content" style="max-width: 800px;">
             <div class="modal-header">
-                <h2>批量导入任务</h2>
+                <h2>新建任务</h2>
                 <span class="modal-close" onclick="closeBatchImportModal()">&times;</span>
             </div>
             <div class="modal-body">
@@ -1253,6 +1253,15 @@
                         为批量任务队列设置一个标题，方便后续查找和管理。
                     </div>
                 </div>
+                <div class="form-group">
+                    <label for="batch-queue-role">角色</label>
+                    <select id="batch-queue-role" style="width: 100%; padding: 8px; border: 1px solid #ddd; border-radius: 4px; font-size: 0.875rem;">
+                        <option value="">默认</option>
+                    </select>
+                    <div class="form-hint" style="margin-top: 4px;">
+                        选择一个角色，所有任务将使用该角色的配置（提示词和工具）执行。
+                    </div>
+                </div>
                 <div class="form-group">
                     <label for="batch-tasks-input">任务列表（每行一个任务）<span style="color: red;">*</span></label>
                     <textarea id="batch-tasks-input" rows="15" placeholder="请输入任务列表，每行一个任务，例如：&#10;扫描 192.168.1.1 的开放端口&#10;检查 https://example.com 是否存在SQL注入&#10;枚举 example.com 的子域名" style="font-family: 'Monaco', 'Menlo', 'Ubuntu Mono', monospace; font-size: 0.875rem; line-height: 1.5;"></textarea>