mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-04-21 18:26:38 +02:00
Add files via upload
This commit is contained in:
@@ -659,6 +659,8 @@ func setupRoutes(
|
||||
protected.GET("/batch-tasks/:queueId", agentHandler.GetBatchQueue)
|
||||
protected.POST("/batch-tasks/:queueId/start", agentHandler.StartBatchQueue)
|
||||
protected.POST("/batch-tasks/:queueId/pause", agentHandler.PauseBatchQueue)
|
||||
protected.PUT("/batch-tasks/:queueId/metadata", agentHandler.UpdateBatchQueueMetadata)
|
||||
protected.PUT("/batch-tasks/:queueId/schedule", agentHandler.UpdateBatchQueueSchedule)
|
||||
protected.PUT("/batch-tasks/:queueId/schedule-enabled", agentHandler.SetBatchQueueScheduleEnabled)
|
||||
protected.DELETE("/batch-tasks/:queueId", agentHandler.DeleteBatchQueue)
|
||||
protected.PUT("/batch-tasks/:queueId/tasks/:taskId", agentHandler.UpdateBatchTask)
|
||||
|
||||
@@ -352,6 +352,18 @@ func (db *DB) UpdateBatchQueueCurrentIndex(queueID string, currentIndex int) err
|
||||
return nil
|
||||
}
|
||||
|
||||
// UpdateBatchQueueMetadata 更新批量任务队列标题和角色
|
||||
func (db *DB) UpdateBatchQueueMetadata(queueID, title, role string) error {
|
||||
_, err := db.Exec(
|
||||
"UPDATE batch_task_queues SET title = ?, role = ? WHERE id = ?",
|
||||
title, role, queueID,
|
||||
)
|
||||
if err != nil {
|
||||
return fmt.Errorf("更新批量任务队列元数据失败: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// UpdateBatchQueueSchedule 更新批量任务队列调度相关信息
|
||||
func (db *DB) UpdateBatchQueueSchedule(queueID, scheduleMode, cronExpr string, nextRunAt *time.Time) error {
|
||||
var nextRunAtValue interface{}
|
||||
@@ -435,7 +447,7 @@ func (db *DB) ResetBatchQueueForRerun(queueID string) error {
|
||||
defer tx.Rollback()
|
||||
|
||||
_, err = tx.Exec(
|
||||
"UPDATE batch_task_queues SET status = ?, current_index = 0, started_at = NULL, completed_at = NULL WHERE id = ?",
|
||||
"UPDATE batch_task_queues SET status = ?, current_index = 0, started_at = NULL, completed_at = NULL, last_run_error = NULL, last_schedule_error = NULL WHERE id = ?",
|
||||
"pending", queueID,
|
||||
)
|
||||
if err != nil {
|
||||
|
||||
@@ -89,7 +89,7 @@ type AgentHandler struct {
|
||||
|
||||
// NewAgentHandler 创建新的Agent处理器
|
||||
func NewAgentHandler(agent *agent.Agent, db *database.DB, cfg *config.Config, logger *zap.Logger) *AgentHandler {
|
||||
batchTaskManager := NewBatchTaskManager()
|
||||
batchTaskManager := NewBatchTaskManager(logger)
|
||||
batchTaskManager.SetDB(db)
|
||||
|
||||
// 从数据库加载所有批量任务队列
|
||||
@@ -1650,7 +1650,11 @@ func (h *AgentHandler) CreateBatchQueue(c *gin.Context) {
|
||||
nextRunAt = &next
|
||||
}
|
||||
|
||||
queue := h.batchTaskManager.CreateBatchQueue(req.Title, req.Role, agentMode, scheduleMode, cronExpr, nextRunAt, validTasks)
|
||||
queue, createErr := h.batchTaskManager.CreateBatchQueue(req.Title, req.Role, agentMode, scheduleMode, cronExpr, nextRunAt, validTasks)
|
||||
if createErr != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": createErr.Error()})
|
||||
return
|
||||
}
|
||||
started := false
|
||||
if req.ExecuteNow {
|
||||
ok, err := h.startBatchQueueExecution(queue.ID, false)
|
||||
@@ -1721,6 +1725,11 @@ func (h *AgentHandler) ListBatchQueues(c *gin.Context) {
|
||||
if offset < 0 {
|
||||
offset = 0
|
||||
}
|
||||
// 防止恶意大 offset 导致 DB 性能问题
|
||||
const maxOffset = 100000
|
||||
if offset > maxOffset {
|
||||
offset = maxOffset
|
||||
}
|
||||
|
||||
// 默认status为"all"
|
||||
if status == "" {
|
||||
@@ -1783,6 +1792,67 @@ func (h *AgentHandler) PauseBatchQueue(c *gin.Context) {
|
||||
c.JSON(http.StatusOK, gin.H{"message": "批量任务已暂停"})
|
||||
}
|
||||
|
||||
// UpdateBatchQueueMetadata 修改批量任务队列的标题和角色
|
||||
func (h *AgentHandler) UpdateBatchQueueMetadata(c *gin.Context) {
|
||||
queueID := c.Param("queueId")
|
||||
var req struct {
|
||||
Title string `json:"title"`
|
||||
Role string `json:"role"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if err := h.batchTaskManager.UpdateQueueMetadata(queueID, req.Title, req.Role); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
updated, _ := h.batchTaskManager.GetBatchQueue(queueID)
|
||||
c.JSON(http.StatusOK, gin.H{"queue": updated})
|
||||
}
|
||||
|
||||
// UpdateBatchQueueSchedule 修改批量任务队列的调度配置(scheduleMode / cronExpr)
|
||||
func (h *AgentHandler) UpdateBatchQueueSchedule(c *gin.Context) {
|
||||
queueID := c.Param("queueId")
|
||||
queue, exists := h.batchTaskManager.GetBatchQueue(queueID)
|
||||
if !exists {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "队列不存在"})
|
||||
return
|
||||
}
|
||||
// 仅在非 running 状态下允许修改调度
|
||||
if queue.Status == "running" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "队列正在运行中,无法修改调度配置"})
|
||||
return
|
||||
}
|
||||
var req struct {
|
||||
ScheduleMode string `json:"scheduleMode"`
|
||||
CronExpr string `json:"cronExpr"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
scheduleMode := normalizeBatchQueueScheduleMode(req.ScheduleMode)
|
||||
cronExpr := strings.TrimSpace(req.CronExpr)
|
||||
var nextRunAt *time.Time
|
||||
if scheduleMode == "cron" {
|
||||
if cronExpr == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "启用 Cron 调度时,调度表达式不能为空"})
|
||||
return
|
||||
}
|
||||
schedule, err := h.batchCronParser.Parse(cronExpr)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "无效的 Cron 表达式: " + err.Error()})
|
||||
return
|
||||
}
|
||||
next := schedule.Next(time.Now())
|
||||
nextRunAt = &next
|
||||
}
|
||||
h.batchTaskManager.UpdateQueueSchedule(queueID, scheduleMode, cronExpr, nextRunAt)
|
||||
updated, _ := h.batchTaskManager.GetBatchQueue(queueID)
|
||||
c.JSON(http.StatusOK, gin.H{"queue": updated})
|
||||
}
|
||||
|
||||
// SetBatchQueueScheduleEnabled 开启/关闭 Cron 自动调度(手工执行不受影响)
|
||||
func (h *AgentHandler) SetBatchQueueScheduleEnabled(c *gin.Context) {
|
||||
queueID := c.Param("queueId")
|
||||
|
||||
@@ -11,6 +11,32 @@ import (
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/database"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// 批量任务状态常量
|
||||
const (
|
||||
BatchQueueStatusPending = "pending"
|
||||
BatchQueueStatusRunning = "running"
|
||||
BatchQueueStatusPaused = "paused"
|
||||
BatchQueueStatusCompleted = "completed"
|
||||
BatchQueueStatusCancelled = "cancelled"
|
||||
|
||||
BatchTaskStatusPending = "pending"
|
||||
BatchTaskStatusRunning = "running"
|
||||
BatchTaskStatusCompleted = "completed"
|
||||
BatchTaskStatusFailed = "failed"
|
||||
BatchTaskStatusCancelled = "cancelled"
|
||||
|
||||
// MaxBatchTasksPerQueue 单个队列最大任务数
|
||||
MaxBatchTasksPerQueue = 10000
|
||||
|
||||
// MaxBatchQueueTitleLen 队列标题最大长度
|
||||
MaxBatchQueueTitleLen = 200
|
||||
|
||||
// MaxBatchQueueRoleLen 角色名最大长度
|
||||
MaxBatchQueueRoleLen = 100
|
||||
)
|
||||
|
||||
// BatchTask 批量任务项
|
||||
@@ -50,14 +76,19 @@ type BatchTaskQueue struct {
|
||||
// BatchTaskManager 批量任务管理器
|
||||
type BatchTaskManager struct {
|
||||
db *database.DB
|
||||
logger *zap.Logger
|
||||
queues map[string]*BatchTaskQueue
|
||||
taskCancels map[string]context.CancelFunc // 存储每个队列当前任务的取消函数
|
||||
mu sync.RWMutex
|
||||
}
|
||||
|
||||
// NewBatchTaskManager 创建批量任务管理器
|
||||
func NewBatchTaskManager() *BatchTaskManager {
|
||||
func NewBatchTaskManager(logger *zap.Logger) *BatchTaskManager {
|
||||
if logger == nil {
|
||||
logger = zap.NewNop()
|
||||
}
|
||||
return &BatchTaskManager{
|
||||
logger: logger,
|
||||
queues: make(map[string]*BatchTaskQueue),
|
||||
taskCancels: make(map[string]context.CancelFunc),
|
||||
}
|
||||
@@ -75,7 +106,18 @@ func (m *BatchTaskManager) CreateBatchQueue(
|
||||
title, role, agentMode, scheduleMode, cronExpr string,
|
||||
nextRunAt *time.Time,
|
||||
tasks []string,
|
||||
) *BatchTaskQueue {
|
||||
) (*BatchTaskQueue, error) {
|
||||
// 输入校验
|
||||
if len(title) > MaxBatchQueueTitleLen {
|
||||
return nil, fmt.Errorf("标题不能超过 %d 个字符", MaxBatchQueueTitleLen)
|
||||
}
|
||||
if len(role) > MaxBatchQueueRoleLen {
|
||||
return nil, fmt.Errorf("角色名不能超过 %d 个字符", MaxBatchQueueRoleLen)
|
||||
}
|
||||
if len(tasks) > MaxBatchTasksPerQueue {
|
||||
return nil, fmt.Errorf("单个队列最多 %d 条任务", MaxBatchTasksPerQueue)
|
||||
}
|
||||
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
|
||||
@@ -131,13 +173,12 @@ func (m *BatchTaskManager) CreateBatchQueue(
|
||||
queue.NextRunAt,
|
||||
dbTasks,
|
||||
); err != nil {
|
||||
// 如果数据库保存失败,记录错误但继续(使用内存缓存)
|
||||
// 这里可以添加日志记录
|
||||
m.logger.Warn("batch queue DB create failed", zap.String("queueId", queueID), zap.Error(err))
|
||||
}
|
||||
}
|
||||
|
||||
m.queues[queueID] = queue
|
||||
return queue
|
||||
return queue, nil
|
||||
}
|
||||
|
||||
// GetBatchQueue 获取批量任务队列
|
||||
@@ -525,7 +566,7 @@ func (m *BatchTaskManager) UpdateTaskStatusWithConversationID(queueID, taskID, s
|
||||
// 同步到数据库
|
||||
if m.db != nil {
|
||||
if err := m.db.UpdateBatchTaskStatus(queueID, taskID, status, conversationID, result, errorMsg); err != nil {
|
||||
// 记录错误但继续(使用内存缓存)
|
||||
m.logger.Warn("batch task DB status update failed", zap.String("queueId", queueID), zap.String("taskId", taskID), zap.Error(err))
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -552,7 +593,7 @@ func (m *BatchTaskManager) UpdateQueueStatus(queueID, status string) {
|
||||
// 同步到数据库
|
||||
if m.db != nil {
|
||||
if err := m.db.UpdateBatchQueueStatus(queueID, status); err != nil {
|
||||
// 记录错误但继续(使用内存缓存)
|
||||
m.logger.Warn("batch queue DB status update failed", zap.String("queueId", queueID), zap.Error(err))
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -578,11 +619,42 @@ func (m *BatchTaskManager) UpdateQueueSchedule(queueID, scheduleMode, cronExpr s
|
||||
|
||||
if m.db != nil {
|
||||
if err := m.db.UpdateBatchQueueSchedule(queueID, queue.ScheduleMode, queue.CronExpr, queue.NextRunAt); err != nil {
|
||||
// 记录错误但继续(使用内存缓存)
|
||||
m.logger.Warn("batch queue DB schedule update failed", zap.String("queueId", queueID), zap.Error(err))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// UpdateQueueMetadata 更新队列标题和角色(非 running 时可用)
|
||||
func (m *BatchTaskManager) UpdateQueueMetadata(queueID, title, role string) error {
|
||||
if len(title) > MaxBatchQueueTitleLen {
|
||||
return fmt.Errorf("标题不能超过 %d 个字符", MaxBatchQueueTitleLen)
|
||||
}
|
||||
if len(role) > MaxBatchQueueRoleLen {
|
||||
return fmt.Errorf("角色名不能超过 %d 个字符", MaxBatchQueueRoleLen)
|
||||
}
|
||||
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
|
||||
queue, exists := m.queues[queueID]
|
||||
if !exists {
|
||||
return fmt.Errorf("队列不存在")
|
||||
}
|
||||
if queue.Status == "running" {
|
||||
return fmt.Errorf("队列正在运行中,无法修改")
|
||||
}
|
||||
|
||||
queue.Title = title
|
||||
queue.Role = role
|
||||
|
||||
if m.db != nil {
|
||||
if err := m.db.UpdateBatchQueueMetadata(queueID, title, role); err != nil {
|
||||
m.logger.Warn("batch queue DB metadata update failed", zap.String("queueId", queueID), zap.Error(err))
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// SetScheduleEnabled 暂停/恢复 Cron 自动调度(不影响手工执行)
|
||||
func (m *BatchTaskManager) SetScheduleEnabled(queueID string, enabled bool) bool {
|
||||
m.mu.Lock()
|
||||
@@ -661,6 +733,8 @@ func (m *BatchTaskManager) ResetQueueForRerun(queueID string) bool {
|
||||
queue.StartedAt = nil
|
||||
queue.CompletedAt = nil
|
||||
queue.NextRunAt = nil
|
||||
queue.LastRunError = ""
|
||||
queue.LastScheduleError = ""
|
||||
for _, task := range queue.Tasks {
|
||||
task.Status = "pending"
|
||||
task.ConversationID = ""
|
||||
@@ -832,8 +906,8 @@ func queueAllowsTaskListMutationLocked(queue *BatchTaskQueue) bool {
|
||||
|
||||
// GetNextTask 获取下一个待执行的任务
|
||||
func (m *BatchTaskManager) GetNextTask(queueID string) (*BatchTask, bool) {
|
||||
m.mu.RLock()
|
||||
defer m.mu.RUnlock()
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
|
||||
queue, exists := m.queues[queueID]
|
||||
if !exists {
|
||||
@@ -866,7 +940,7 @@ func (m *BatchTaskManager) MoveToNextTask(queueID string) {
|
||||
// 同步到数据库
|
||||
if m.db != nil {
|
||||
if err := m.db.UpdateBatchQueueCurrentIndex(queueID, queue.CurrentIndex); err != nil {
|
||||
// 记录错误但继续(使用内存缓存)
|
||||
m.logger.Warn("batch queue DB index update failed", zap.String("queueId", queueID), zap.Error(err))
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -885,15 +959,14 @@ func (m *BatchTaskManager) SetTaskCancel(queueID string, cancel context.CancelFu
|
||||
// PauseQueue 暂停队列
|
||||
func (m *BatchTaskManager) PauseQueue(queueID string) bool {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
|
||||
queue, exists := m.queues[queueID]
|
||||
if !exists {
|
||||
m.mu.Unlock()
|
||||
return false
|
||||
}
|
||||
|
||||
if queue.Status != "running" {
|
||||
m.mu.Unlock()
|
||||
return false
|
||||
}
|
||||
|
||||
@@ -905,12 +978,10 @@ func (m *BatchTaskManager) PauseQueue(queueID string) bool {
|
||||
delete(m.taskCancels, queueID)
|
||||
}
|
||||
|
||||
m.mu.Unlock()
|
||||
|
||||
// 同步队列状态到数据库
|
||||
// 同步队列状态到数据库(在锁内完成,避免竞态)
|
||||
if m.db != nil {
|
||||
if err := m.db.UpdateBatchQueueStatus(queueID, "paused"); err != nil {
|
||||
// 记录错误但继续(使用内存缓存)
|
||||
m.logger.Warn("batch queue DB pause update failed", zap.String("queueId", queueID), zap.Error(err))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -920,15 +991,14 @@ func (m *BatchTaskManager) PauseQueue(queueID string) bool {
|
||||
// CancelQueue 取消队列(保留此方法以保持向后兼容,但建议使用PauseQueue)
|
||||
func (m *BatchTaskManager) CancelQueue(queueID string) bool {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
|
||||
queue, exists := m.queues[queueID]
|
||||
if !exists {
|
||||
m.mu.Unlock()
|
||||
return false
|
||||
}
|
||||
|
||||
if queue.Status == "completed" || queue.Status == "cancelled" {
|
||||
m.mu.Unlock()
|
||||
return false
|
||||
}
|
||||
|
||||
@@ -941,7 +1011,6 @@ func (m *BatchTaskManager) CancelQueue(queueID string) bool {
|
||||
if task.Status == "pending" {
|
||||
task.Status = "cancelled"
|
||||
task.CompletedAt = &now
|
||||
// 同步到数据库
|
||||
if m.db != nil {
|
||||
m.db.UpdateBatchTaskStatus(queueID, task.ID, "cancelled", "", "", "")
|
||||
}
|
||||
@@ -954,35 +1023,38 @@ func (m *BatchTaskManager) CancelQueue(queueID string) bool {
|
||||
delete(m.taskCancels, queueID)
|
||||
}
|
||||
|
||||
m.mu.Unlock()
|
||||
|
||||
// 同步队列状态到数据库
|
||||
// 同步队列状态到数据库(在锁内完成)
|
||||
if m.db != nil {
|
||||
if err := m.db.UpdateBatchQueueStatus(queueID, "cancelled"); err != nil {
|
||||
// 记录错误但继续(使用内存缓存)
|
||||
m.logger.Warn("batch queue DB cancel update failed", zap.String("queueId", queueID), zap.Error(err))
|
||||
}
|
||||
}
|
||||
|
||||
return true
|
||||
}
|
||||
|
||||
// DeleteQueue 删除队列
|
||||
// DeleteQueue 删除队列(运行中的队列不允许删除)
|
||||
func (m *BatchTaskManager) DeleteQueue(queueID string) bool {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
|
||||
_, exists := m.queues[queueID]
|
||||
queue, exists := m.queues[queueID]
|
||||
if !exists {
|
||||
return false
|
||||
}
|
||||
|
||||
// 运行中的队列不允许删除,防止孤儿协程和数据丢失
|
||||
if queue.Status == "running" {
|
||||
return false
|
||||
}
|
||||
|
||||
// 清理取消函数
|
||||
delete(m.taskCancels, queueID)
|
||||
|
||||
// 从数据库删除
|
||||
if m.db != nil {
|
||||
if err := m.db.DeleteBatchQueue(queueID); err != nil {
|
||||
// 记录错误但继续(使用内存缓存)
|
||||
m.logger.Warn("batch queue DB delete failed", zap.String("queueId", queueID), zap.Error(err))
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -161,7 +161,7 @@ agent_mode: single(默认)或 multi(需系统启用多代理)。schedule
|
||||
},
|
||||
"cron_expr": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "schedule_mode 为 cron 时必填",
|
||||
"description": "schedule_mode 为 cron 时必填。标准 5 段格式:分钟 小时 日 月 星期,例如 \"0 */6 * * *\"(每6小时)、\"30 2 * * 1-5\"(工作日凌晨2:30)",
|
||||
},
|
||||
"execute_now": map[string]interface{}{
|
||||
"type": "boolean",
|
||||
@@ -195,7 +195,10 @@ agent_mode: single(默认)或 multi(需系统启用多代理)。schedule
|
||||
if !ok {
|
||||
executeNow = false
|
||||
}
|
||||
queue := h.batchTaskManager.CreateBatchQueue(title, role, agentMode, scheduleMode, cronExpr, nextRunAt, tasks)
|
||||
queue, createErr := h.batchTaskManager.CreateBatchQueue(title, role, agentMode, scheduleMode, cronExpr, nextRunAt, tasks)
|
||||
if createErr != nil {
|
||||
return batchMCPTextResult("创建队列失败: "+createErr.Error(), true), nil
|
||||
}
|
||||
started := false
|
||||
if executeNow {
|
||||
ok, err := h.startBatchQueueExecution(queue.ID, false)
|
||||
@@ -311,6 +314,101 @@ agent_mode: single(默认)或 multi(需系统启用多代理)。schedule
|
||||
return batchMCPTextResult("队列已删除。", false), nil
|
||||
})
|
||||
|
||||
// --- update metadata (title/role) ---
|
||||
reg(mcp.Tool{
|
||||
Name: builtin.ToolBatchTaskUpdateMetadata,
|
||||
Description: "修改批量任务队列的标题和角色。仅在队列非 running 状态下可修改。",
|
||||
ShortDescription: "修改批量任务队列标题/角色",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"queue_id": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "队列 ID",
|
||||
},
|
||||
"title": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "新标题(空字符串清除标题)",
|
||||
},
|
||||
"role": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "新角色名(空字符串使用默认角色)",
|
||||
},
|
||||
},
|
||||
"required": []string{"queue_id"},
|
||||
},
|
||||
}, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
qid := mcpArgString(args, "queue_id")
|
||||
if qid == "" {
|
||||
return batchMCPTextResult("queue_id 不能为空", true), nil
|
||||
}
|
||||
title := mcpArgString(args, "title")
|
||||
role := mcpArgString(args, "role")
|
||||
if err := h.batchTaskManager.UpdateQueueMetadata(qid, title, role); err != nil {
|
||||
return batchMCPTextResult(err.Error(), true), nil
|
||||
}
|
||||
updated, _ := h.batchTaskManager.GetBatchQueue(qid)
|
||||
logger.Info("MCP batch_task_update_metadata", zap.String("queueId", qid))
|
||||
return batchMCPJSONResult(updated)
|
||||
})
|
||||
|
||||
// --- update schedule ---
|
||||
reg(mcp.Tool{
|
||||
Name: builtin.ToolBatchTaskUpdateSchedule,
|
||||
Description: `修改批量任务队列的调度方式和 Cron 表达式。仅在队列非 running 状态下可修改。
|
||||
schedule_mode 为 cron 时必须提供有效 cron_expr;为 manual 时会清除 Cron 配置。`,
|
||||
ShortDescription: "修改批量任务调度配置(Cron 表达式)",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"queue_id": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "队列 ID",
|
||||
},
|
||||
"schedule_mode": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "manual 或 cron",
|
||||
"enum": []string{"manual", "cron"},
|
||||
},
|
||||
"cron_expr": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "Cron 表达式(schedule_mode 为 cron 时必填)。标准 5 段格式:分钟 小时 日 月 星期,如 \"0 */6 * * *\"(每6小时)、\"30 2 * * 1-5\"(工作日凌晨2:30)",
|
||||
},
|
||||
},
|
||||
"required": []string{"queue_id", "schedule_mode"},
|
||||
},
|
||||
}, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
qid := mcpArgString(args, "queue_id")
|
||||
if qid == "" {
|
||||
return batchMCPTextResult("queue_id 不能为空", true), nil
|
||||
}
|
||||
queue, exists := h.batchTaskManager.GetBatchQueue(qid)
|
||||
if !exists {
|
||||
return batchMCPTextResult("队列不存在: "+qid, true), nil
|
||||
}
|
||||
if queue.Status == "running" {
|
||||
return batchMCPTextResult("队列正在运行中,无法修改调度配置", true), nil
|
||||
}
|
||||
scheduleMode := normalizeBatchQueueScheduleMode(mcpArgString(args, "schedule_mode"))
|
||||
cronExpr := strings.TrimSpace(mcpArgString(args, "cron_expr"))
|
||||
var nextRunAt *time.Time
|
||||
if scheduleMode == "cron" {
|
||||
if cronExpr == "" {
|
||||
return batchMCPTextResult("Cron 调度模式下 cron_expr 不能为空", true), nil
|
||||
}
|
||||
sch, err := h.batchCronParser.Parse(cronExpr)
|
||||
if err != nil {
|
||||
return batchMCPTextResult("无效的 Cron 表达式: "+err.Error(), true), nil
|
||||
}
|
||||
n := sch.Next(time.Now())
|
||||
nextRunAt = &n
|
||||
}
|
||||
h.batchTaskManager.UpdateQueueSchedule(qid, scheduleMode, cronExpr, nextRunAt)
|
||||
updated, _ := h.batchTaskManager.GetBatchQueue(qid)
|
||||
logger.Info("MCP batch_task_update_schedule", zap.String("queueId", qid), zap.String("scheduleMode", scheduleMode), zap.String("cronExpr", cronExpr))
|
||||
return batchMCPJSONResult(updated)
|
||||
})
|
||||
|
||||
// --- schedule enabled ---
|
||||
reg(mcp.Tool{
|
||||
Name: builtin.ToolBatchTaskScheduleEnabled,
|
||||
@@ -456,7 +554,7 @@ agent_mode: single(默认)或 multi(需系统启用多代理)。schedule
|
||||
return batchMCPJSONResult(queue)
|
||||
})
|
||||
|
||||
logger.Info("批量任务 MCP 工具已注册", zap.Int("count", 10))
|
||||
logger.Info("批量任务 MCP 工具已注册", zap.Int("count", 12))
|
||||
}
|
||||
|
||||
// --- batch_task_list 精简结构(避免把每条子任务的 result 等大段文本塞进列表上下文) ---
|
||||
@@ -509,6 +607,8 @@ func truncateStringRunes(s string, maxRunes int) string {
|
||||
return s
|
||||
}
|
||||
|
||||
const mcpBatchListMaxTasksPerQueue = 200 // 列表中每个队列最多返回的子任务摘要数
|
||||
|
||||
func toBatchTaskQueueMCPListItem(q *BatchTaskQueue) batchTaskQueueMCPListItem {
|
||||
counts := map[string]int{
|
||||
"pending": 0,
|
||||
@@ -523,11 +623,14 @@ func toBatchTaskQueueMCPListItem(q *BatchTaskQueue) batchTaskQueueMCPListItem {
|
||||
continue
|
||||
}
|
||||
counts[t.Status]++
|
||||
tasks = append(tasks, batchTaskMCPListSummary{
|
||||
ID: t.ID,
|
||||
Status: t.Status,
|
||||
Message: truncateStringRunes(t.Message, mcpBatchListTaskMessageMaxRunes),
|
||||
})
|
||||
// 列表视图限制子任务摘要数量,完整列表通过 batch_task_get 查看
|
||||
if len(tasks) < mcpBatchListMaxTasksPerQueue {
|
||||
tasks = append(tasks, batchTaskMCPListSummary{
|
||||
ID: t.ID,
|
||||
Status: t.Status,
|
||||
Message: truncateStringRunes(t.Message, mcpBatchListTaskMessageMaxRunes),
|
||||
})
|
||||
}
|
||||
}
|
||||
return batchTaskQueueMCPListItem{
|
||||
ID: q.ID,
|
||||
|
||||
@@ -34,6 +34,8 @@ const (
|
||||
ToolBatchTaskStart = "batch_task_start"
|
||||
ToolBatchTaskPause = "batch_task_pause"
|
||||
ToolBatchTaskDelete = "batch_task_delete"
|
||||
ToolBatchTaskUpdateMetadata = "batch_task_update_metadata"
|
||||
ToolBatchTaskUpdateSchedule = "batch_task_update_schedule"
|
||||
ToolBatchTaskScheduleEnabled = "batch_task_schedule_enabled"
|
||||
ToolBatchTaskAdd = "batch_task_add_task"
|
||||
ToolBatchTaskUpdate = "batch_task_update_task"
|
||||
@@ -63,6 +65,8 @@ func IsBuiltinTool(toolName string) bool {
|
||||
ToolBatchTaskStart,
|
||||
ToolBatchTaskPause,
|
||||
ToolBatchTaskDelete,
|
||||
ToolBatchTaskUpdateMetadata,
|
||||
ToolBatchTaskUpdateSchedule,
|
||||
ToolBatchTaskScheduleEnabled,
|
||||
ToolBatchTaskAdd,
|
||||
ToolBatchTaskUpdate,
|
||||
@@ -96,6 +100,8 @@ func GetAllBuiltinTools() []string {
|
||||
ToolBatchTaskStart,
|
||||
ToolBatchTaskPause,
|
||||
ToolBatchTaskDelete,
|
||||
ToolBatchTaskUpdateMetadata,
|
||||
ToolBatchTaskUpdateSchedule,
|
||||
ToolBatchTaskScheduleEnabled,
|
||||
ToolBatchTaskAdd,
|
||||
ToolBatchTaskUpdate,
|
||||
|
||||
Reference in New Issue
Block a user