From 69da61b9a0e67c3dd8273cf6214518f28b2db391 Mon Sep 17 00:00:00 2001 From: engigu Date: Sat, 3 Jan 2026 12:21:10 +0800 Subject: [PATCH] fix: agent exec envs --- agent/agent.go | 58 +++++ internal/services/agent_service.go | 75 +++--- internal/services/executor_service.go | 183 +++---------- internal/services/task_execution_service.go | 275 ++++++++++++++++++++ internal/services/task_log_service.go | 161 ++++++++++++ 5 files changed, 566 insertions(+), 186 deletions(-) create mode 100644 internal/services/task_execution_service.go create mode 100644 internal/services/task_log_service.go diff --git a/agent/agent.go b/agent/agent.go index a4d8792..6fd0b0d 100644 --- a/agent/agent.go +++ b/agent/agent.go @@ -262,6 +262,8 @@ func (a *Agent) handleWSMessage(msg *WSMessage) { case WSTypeEnabled: log.Info("Agent 已被启用，主动拉取任务") a.fetchTasks() + case "execute": + a.handleExecute(msg.Data) } } @@ -320,6 +322,31 @@ func (a *Agent) handleTasks(data json.RawMessage) { a.updateTasks(resp.Tasks) } +func (a *Agent) handleExecute(data json.RawMessage) { + var req struct { + TaskID uint `json:"task_id"` + } + if err := json.Unmarshal(data, &req); err != nil { + log.Errorf("解析立即执行请求失败: %v", err) + return + } + + log.Infof("收到立即执行命令: 任务 #%d", req.TaskID) + + // 查找任务 + a.mu.RLock() + task, exists := a.tasks[req.TaskID] + a.mu.RUnlock() + + if !exists { + log.Warnf("任务 #%d 不存在，无法执行", req.TaskID) + return + } + + // 立即执行任务 + go a.executeTask(task) +} + func (a *Agent) sendWSMessage(msgType string, data interface{}) error { a.wsMu.Lock() defer a.wsMu.Unlock() @@ -496,6 +523,15 @@ func (a *Agent) executeTask(task *AgentTask) { cmd = exec.CommandContext(ctx, "sh", "-c", finalCommand) } + // 处理环境变量 + if task.Envs != "" { + envVars := a.parseEnvVars(task.Envs) + if len(envVars) > 0 { + cmd.Env = append(os.Environ(), envVars...) + log.Infof("任务 #%d 设置了 %d 个环境变量", task.ID, len(envVars)) + } + } + var stdout, stderr bytes.Buffer cmd.Stdout = &stdout cmd.Stderr = &stderr @@ -536,6 +572,28 @@ func (a *Agent) executeTask(task *AgentTask) { log.Infof("任务 #%d 执行完成 (%s)", result.TaskID, result.Status) } +// parseEnvVars 解析环境变量字符串 "KEY1=VALUE1,KEY2=VALUE2" +func (a *Agent) parseEnvVars(envStr string) []string { + if envStr == "" { + return nil + } + + pairs := strings.Split(envStr, ",") + result := make([]string, 0, len(pairs)) + + for _, pair := range pairs { + if pair == "" { + continue + } + // 解码特殊字符 + pair = strings.ReplaceAll(pair, "{{COMMA}}", ",") + pair = strings.ReplaceAll(pair, "{{EQUAL}}", "=") + result = append(result, pair) + } + + return result +} + func (a *Agent) doRequest(method, path string, body interface{}) (*http.Response, error) { var bodyReader io.Reader if body != nil { diff --git a/internal/services/agent_service.go b/internal/services/agent_service.go index 695faf9..8c4f76c 100644 --- a/internal/services/agent_service.go +++ b/internal/services/agent_service.go @@ -5,7 +5,6 @@ import ( "baihu/internal/database" "baihu/internal/logger" "baihu/internal/models" - "baihu/internal/utils" "crypto/rand" "encoding/hex" "fmt" @@ -306,6 +305,9 @@ func (s *AgentService) GetTasks(agentID uint) []models.AgentTask { result := make([]models.AgentTask, len(tasks)) for i, task := range tasks { + // 将环境变量 ID 转换为实际的环境变量键值对 + envVarsStr := s.buildEnvVarsString(task.Envs) + result[i] = models.AgentTask{ ID: task.ID, Name: task.Name, @@ -313,7 +315,7 @@ func (s *AgentService) GetTasks(agentID uint) []models.AgentTask { Schedule: task.Schedule, Timeout: task.Timeout, WorkDir: task.WorkDir, - Envs: task.Envs, + Envs: envVarsStr, // 传递 "KEY1=VALUE1,KEY2=VALUE2" 格式 Enabled: task.Enabled, } } @@ -321,48 +323,37 @@ func (s *AgentService) GetTasks(agentID uint) []models.AgentTask { return result } +// buildEnvVarsString 将环境变量 ID 列表转换为键值对字符串 +func (s *AgentService) buildEnvVarsString(envIDs string) string { + if envIDs == "" { + return "" + } + + var envVars []models.EnvironmentVariable + ids := strings.Split(envIDs, ",") + database.DB.Where("id IN ?", ids).Find(&envVars) + + if len(envVars) == 0 { + return "" + } + + // 构建 "KEY1=VALUE1,KEY2=VALUE2" 格式 + pairs := make([]string, 0, len(envVars)) + for _, env := range envVars { + // 对值进行转义，避免特殊字符问题 + encodedValue := strings.ReplaceAll(env.Value, ",", "{{COMMA}}") + encodedValue = strings.ReplaceAll(encodedValue, "=", "{{EQUAL}}") + pairs = append(pairs, fmt.Sprintf("%s=%s", env.Name, encodedValue)) + } + return strings.Join(pairs, ",") +} + // ReportResult Agent 上报执行结果 func (s *AgentService) ReportResult(result *models.AgentTaskResult) error { - // 压缩输出 - compressed, err := utils.CompressToBase64(result.Output) - if err != nil { - logger.Errorf("[Agent] 压缩日志失败: %v", err) - compressed = "" - } - - taskLog := &models.TaskLog{ - TaskID: result.TaskID, - AgentID: &result.AgentID, - Command: result.Command, - Output: compressed, - Status: result.Status, - Duration: result.Duration, - ExitCode: result.ExitCode, - } - - // 处理开始和结束时间 - if result.StartTime > 0 { - startTime := models.LocalTime(time.Unix(result.StartTime, 0)) - taskLog.StartTime = &startTime - } - if result.EndTime > 0 { - endTime := models.LocalTime(time.Unix(result.EndTime, 0)) - taskLog.EndTime = &endTime - } - - if err := database.DB.Create(taskLog).Error; err != nil { - return err - } - - // 更新任务的 last_run - database.DB.Model(&models.Task{}).Where("id = ?", result.TaskID).Update("last_run", time.Now()) - - // 更新统计 - sendStatsService := NewSendStatsService() - sendStatsService.IncrementStats(result.TaskID, result.Status) - - logger.Infof("[Agent] 收到任务结果 #%d (agent=%d, status=%s)", result.TaskID, result.AgentID, result.Status) - return nil + taskExecutionService := NewTaskExecutionService() + + // 使用统一的结果处理流程 + return taskExecutionService.ProcessAgentResult(result) } // UpdateOfflineAgents 更新离线 Agent 状态（超过 2 分钟无心跳） diff --git a/internal/services/executor_service.go b/internal/services/executor_service.go index 599547f..4510692 100644 --- a/internal/services/executor_service.go +++ b/internal/services/executor_service.go @@ -2,7 +2,6 @@ package services import ( "baihu/internal/constant" - "baihu/internal/database" "baihu/internal/logger" "baihu/internal/models" "baihu/internal/utils" @@ -28,9 +27,6 @@ type ExecutionResult struct { End time.Time } -// ExecutionCallback 任务执行完成后的回调函数类型 -type ExecutionCallback func(taskID uint, command string, result *ExecutionResult) - // taskJob 任务队列项 type taskJob struct { taskID int @@ -38,12 +34,12 @@ type taskJob struct { // ExecutorService handles task execution type ExecutorService struct { - taskService *TaskService - results []ExecutionResult - runningTasks map[int]bool - callbacks []ExecutionCallback - mu sync.RWMutex - resultsMu sync.RWMutex + taskService *TaskService + taskExecutionService *TaskExecutionService + results []ExecutionResult + runningTasks map[int]bool + mu sync.RWMutex + resultsMu sync.RWMutex // 任务队列和 worker pool taskQueue chan taskJob @@ -64,21 +60,16 @@ func NewExecutorService(taskService *TaskService) *ExecutorService { logger.Infof("[Executor] 配置: workers=%d, queue=%d, rate=%dms", workerCount, queueSize, rateInterval) es := &ExecutorService{ - taskService: taskService, - results: make([]ExecutionResult, 0, 100), - runningTasks: make(map[int]bool), - callbacks: make([]ExecutionCallback, 0), - taskQueue: make(chan taskJob, queueSize), - workerCount: workerCount, - rateLimiter: time.Tick(time.Duration(rateInterval) * time.Millisecond), - stopCh: make(chan struct{}), + taskService: taskService, + taskExecutionService: NewTaskExecutionService(), + results: make([]ExecutionResult, 0, 100), + runningTasks: make(map[int]bool), + taskQueue: make(chan taskJob, queueSize), + workerCount: workerCount, + rateLimiter: time.Tick(time.Duration(rateInterval) * time.Millisecond), + stopCh: make(chan struct{}), } - // 注册默认回调 - es.RegisterCallback(es.saveTaskLogCallback) - es.RegisterCallback(es.updateStatsCallback) - es.RegisterCallback(es.cleanLogsCallback) - // 启动 worker pool es.startWorkers() @@ -156,118 +147,6 @@ func (es *ExecutorService) Reload() { logger.Infof("[Executor] 配置已重载: workers=%d, queue=%d, rate=%dms", workerCount, queueSize, rateInterval) } -// RegisterCallback 注册执行完成回调 -func (es *ExecutorService) RegisterCallback(cb ExecutionCallback) { - es.mu.Lock() - es.callbacks = append(es.callbacks, cb) - es.mu.Unlock() -} - -// executeCallbacksAsync 异步执行所有回调 -func (es *ExecutorService) executeCallbacksAsync(taskID uint, command string, result *ExecutionResult) { - es.mu.RLock() - callbacks := make([]ExecutionCallback, len(es.callbacks)) - copy(callbacks, es.callbacks) - es.mu.RUnlock() - - go func() { - for _, cb := range callbacks { - cb(taskID, command, result) - } - }() -} - -// saveTaskLogCallback 保存任务日志的回调（异步执行，压缩在此处进行） -func (es *ExecutorService) saveTaskLogCallback(taskID uint, command string, result *ExecutionResult) { - output := result.Output - if result.Error != "" { - output += "\n[ERROR]\n" + result.Error - } - - compressed, err := utils.CompressToBase64(output) - if err != nil { - logger.Errorf("[Executor] 压缩日志失败: %v", err) - compressed = "" - } - - status := "success" - if !result.Success { - status = "failed" - } - - startTime := models.LocalTime(result.Start) - endTime := models.LocalTime(result.End) - - taskLog := &models.TaskLog{ - TaskID: taskID, - Command: command, - Output: compressed, - Status: status, - Duration: result.End.Sub(result.Start).Milliseconds(), - StartTime: &startTime, - EndTime: &endTime, - } - - if err := database.DB.Create(taskLog).Error; err != nil { - logger.Errorf("[Executor] 保存任务日志失败: %v", err) - } -} - -// updateStatsCallback 更新统计数据的回调 -func (es *ExecutorService) updateStatsCallback(taskID uint, _ string, result *ExecutionResult) { - status := "success" - if !result.Success { - status = "failed" - } - sendStatsService := NewSendStatsService() - if err := sendStatsService.IncrementStats(taskID, status); err != nil { - logger.Errorf("[Executor] 更新统计失败: %v", err) - } -} - -// CleanConfig 清理配置结构 -type CleanConfig struct { - Type string `json:"type"` // "day" 或 "count" - Keep int `json:"keep"` // 保留天数或条数 -} - -// cleanLogsCallback 清理日志的回调 -func (es *ExecutorService) cleanLogsCallback(taskID uint, _ string, _ *ExecutionResult) { - task := es.taskService.GetTaskByID(int(taskID)) - if task == nil || task.CleanConfig == "" { - return - } - - var config CleanConfig - if err := json.Unmarshal([]byte(task.CleanConfig), &config); err != nil { - logger.Errorf("[Executor] 解析清理配置失败: %v", err) - return - } - - if config.Keep <= 0 { - return - } - - var deleted int64 - switch config.Type { - case "day": - cutoff := time.Now().AddDate(0, 0, -config.Keep) - result := database.DB.Where("task_id = ? AND created_at < ?", taskID, cutoff).Delete(&models.TaskLog{}) - deleted = result.RowsAffected - case "count": - var boundaryLog models.TaskLog - err := database.DB.Where("task_id = ?", taskID).Order("id DESC").Offset(config.Keep - 1).Limit(1).First(&boundaryLog).Error - if err == nil { - result := database.DB.Where("task_id = ? AND id < ?", taskID, boundaryLog.ID).Delete(&models.TaskLog{}) - deleted = result.RowsAffected - } - } - - if deleted > 0 { - //logger.Infof("Cleaned %d logs for task %d", deleted, taskID) - } -} - // EnqueueTask 将任务加入队列（供 cron 调度器调用） func (es *ExecutorService) EnqueueTask(taskID int) { select { @@ -305,23 +184,39 @@ func (es *ExecutorService) executeTaskInternal(taskID int) *ExecutionResult { var result *ExecutionResult - // 根据任务类型执行不同逻辑 - if task.Type == "repo" { - result = es.executeRepoTask(task) - } else { - result = es.executeNormalTask(task) + // 使用统一的任务执行服务 + req := &TaskExecutionRequest{ + TaskID: uint(taskID), + Task: task, } - result.TaskID = taskID + start := time.Now() + err := es.taskExecutionService.ExecuteTask(req) + end := time.Now() + + if err != nil { + result = &ExecutionResult{ + TaskID: taskID, + Success: false, + Error: err.Error(), + Start: start, + End: end, + } + } else { + result = &ExecutionResult{ + TaskID: taskID, + Success: true, + Output: "任务已提交执行", + Start: start, + End: end, + } + } // 标记任务结束 es.mu.Lock() delete(es.runningTasks, taskID) es.mu.Unlock() - // 异步执行回调（日志压缩、统计更新、日志清理） - es.executeCallbacksAsync(uint(taskID), task.Command, result) - return result } diff --git a/internal/services/task_execution_service.go b/internal/services/task_execution_service.go new file mode 100644 index 0000000..15f4da2 --- /dev/null +++ b/internal/services/task_execution_service.go @@ -0,0 +1,275 @@ +package services + +import ( + "baihu/internal/database" + "baihu/internal/logger" + "baihu/internal/models" + "bytes" + "context" + "fmt" + "os" + "os/exec" + "path/filepath" + "runtime" + "strings" + "time" +) + +// TaskExecutionService 统一的任务执行服务 +type TaskExecutionService struct { + taskLogService *TaskLogService +} + +// NewTaskExecutionService 创建任务执行服务 +func NewTaskExecutionService() *TaskExecutionService { + return &TaskExecutionService{ + taskLogService: NewTaskLogService(), + } +} + +// TaskExecutionRequest 任务执行请求 +type TaskExecutionRequest struct { + TaskID uint + Task *models.Task + AgentID *uint // nil 表示本地执行 +} + +// TaskExecutionResult 任务执行结果 +type TaskExecutionResult struct { + TaskID uint + AgentID *uint + Command string + Output string + Status string // success, failed + Duration int64 // milliseconds + ExitCode int + Start time.Time + End time.Time +} + +// ExecuteTask 执行任务（统一入口） +func (s *TaskExecutionService) ExecuteTask(req *TaskExecutionRequest) error { + if req.Task.AgentID != nil && *req.Task.AgentID > 0 { + // 远程执行：通过 Agent + return s.executeRemote(req) + } + // 本地执行 + return s.executeLocal(req) +} + +// executeLocal 本地执行任务 +func (s *TaskExecutionService) executeLocal(req *TaskExecutionRequest) error { + task := req.Task + logger.Infof("[TaskExecution] 本地执行任务 #%d: %s", task.ID, task.Name) + + start := time.Now() + + // 准备命令 + ctx, cancel := s.createContext(task.Timeout) + defer cancel() + + cmd, err := s.prepareCommand(ctx, task) + if err != nil { + return s.handleExecutionError(task.ID, task.Command, start, err) + } + + // 执行命令 + var stdout, stderr bytes.Buffer + cmd.Stdout = &stdout + cmd.Stderr = &stderr + + execErr := cmd.Run() + end := time.Now() + + // 构建结果 + result := &TaskExecutionResult{ + TaskID: task.ID, + AgentID: nil, + Command: task.Command, + Output: stdout.String(), + Start: start, + End: end, + Duration: end.Sub(start).Milliseconds(), + } + + if execErr != nil { + result.Status = "failed" + result.Output += "\n[ERROR]\n" + stderr.String() + "\n" + execErr.Error() + if exitErr, ok := execErr.(*exec.ExitError); ok { + result.ExitCode = exitErr.ExitCode() + } else { + result.ExitCode = 1 + } + } else { + result.Status = "success" + result.ExitCode = 0 + } + + // 处理执行结果 + return s.processExecutionResult(result) +} + +// executeRemote 远程执行任务（通过 Agent） +func (s *TaskExecutionService) executeRemote(req *TaskExecutionRequest) error { + task := req.Task + agentID := *task.AgentID + + logger.Infof("[TaskExecution] 远程执行任务 #%d: %s (Agent #%d)", task.ID, task.Name, agentID) + + // 检查 Agent 是否在线 + var agent models.Agent + if err := database.DB.First(&agent, agentID).Error; err != nil { + return fmt.Errorf("Agent #%d 不存在", agentID) + } + + if !agent.Enabled { + return fmt.Errorf("Agent #%d 已禁用", agentID) + } + + // 通过 WebSocket 发送立即执行命令给 Agent + manager := GetAgentWSManager() + err := manager.SendToAgent(agentID, "execute", map[string]interface{}{ + "task_id": task.ID, + }) + if err != nil { + return fmt.Errorf("发送执行命令失败: %v", err) + } + + logger.Infof("[TaskExecution] 已发送立即执行命令给 Agent #%d，任务 #%d", agentID, task.ID) + return nil +} + +// prepareCommand 准备执行命令 +func (s *TaskExecutionService) prepareCommand(ctx context.Context, task *models.Task) (*exec.Cmd, error) { + command := task.Command + + // 处理工作目录 + if task.WorkDir != "" { + // 验证工作目录 + if _, err := os.Stat(task.WorkDir); err != nil { + return nil, fmt.Errorf("工作目录不存在或无法访问: %s", task.WorkDir) + } + } + + // 处理环境变量 + envVars := s.loadEnvVars(task.Envs) + + // 根据操作系统创建命令 + var cmd *exec.Cmd + if runtime.GOOS == "windows" { + cmd = exec.CommandContext(ctx, "cmd", "/c", command) + } else { + // 如果有工作目录，在命令前加 cd + if task.WorkDir != "" { + command = fmt.Sprintf("cd %s && %s", task.WorkDir, command) + } + cmd = exec.CommandContext(ctx, "sh", "-c", command) + } + + // 设置环境变量 + if len(envVars) > 0 { + cmd.Env = append(os.Environ(), envVars...) + } + + return cmd, nil +} + +// createContext 创建带超时的上下文 +func (s *TaskExecutionService) createContext(timeout int) (context.Context, context.CancelFunc) { + if timeout <= 0 { + timeout = 30 // 默认 30 分钟 + } + return context.WithTimeout(context.Background(), time.Duration(timeout)*time.Minute) +} + +// loadEnvVars 加载环境变量 +func (s *TaskExecutionService) loadEnvVars(envIDs string) []string { + if envIDs == "" { + return nil + } + + var envVars []models.EnvironmentVariable + ids := strings.Split(envIDs, ",") + database.DB.Where("id IN ?", ids).Find(&envVars) + + result := make([]string, 0, len(envVars)) + for _, env := range envVars { + result = append(result, fmt.Sprintf("%s=%s", env.Name, env.Value)) + } + return result +} + +// handleExecutionError 处理执行错误 +func (s *TaskExecutionService) handleExecutionError(taskID uint, command string, start time.Time, err error) error { + end := time.Now() + result := &TaskExecutionResult{ + TaskID: taskID, + Command: command, + Output: fmt.Sprintf("[ERROR] 任务执行失败: %v", err), + Status: "failed", + Duration: end.Sub(start).Milliseconds(), + ExitCode: 1, + Start: start, + End: end, + } + return s.processExecutionResult(result) +} + +// processExecutionResult 处理执行结果（统一的结果处理） +func (s *TaskExecutionService) processExecutionResult(result *TaskExecutionResult) error { + // 创建任务日志 + taskLog, err := s.taskLogService.CreateTaskLogFromLocalExecution( + result.TaskID, + result.Command, + result.Output, + result.Status, + result.Duration, + result.ExitCode, + result.Start, + result.End, + ) + if err != nil { + logger.Errorf("[TaskExecution] 创建任务日志失败: %v", err) + return err + } + + // 如果是 Agent 执行的，设置 AgentID + if result.AgentID != nil { + taskLog.AgentID = result.AgentID + } + + // 处理任务完成（保存日志、更新统计、清理旧日志） + if err := s.taskLogService.ProcessTaskCompletion(taskLog); err != nil { + logger.Errorf("[TaskExecution] 处理任务完成失败: %v", err) + return err + } + + logger.Infof("[TaskExecution] 任务 #%d 执行完成 (%s)", result.TaskID, result.Status) + return nil +} + +// ProcessAgentResult 处理 Agent 上报的结果（统一入口） +func (s *TaskExecutionService) ProcessAgentResult(agentResult *models.AgentTaskResult) error { + logger.Infof("[TaskExecution] 处理 Agent #%d 上报的任务 #%d 结果", agentResult.AgentID, agentResult.TaskID) + + // 转换为统一的执行结果 + result := &TaskExecutionResult{ + TaskID: agentResult.TaskID, + AgentID: &agentResult.AgentID, + Command: agentResult.Command, + Output: agentResult.Output, + Status: agentResult.Status, + Duration: agentResult.Duration, + ExitCode: agentResult.ExitCode, + Start: time.Unix(agentResult.StartTime, 0), + End: time.Unix(agentResult.EndTime, 0), + } + + // 使用统一的结果处理流程 + return s.processExecutionResult(result) +} + +// GetScriptPath 获取脚本路径 +func (s *TaskExecutionService) GetScriptPath(scriptName string) string { + return filepath.Join("data", "scripts", scriptName) +} diff --git a/internal/services/task_log_service.go b/internal/services/task_log_service.go new file mode 100644 index 0000000..27adaa8 --- /dev/null +++ b/internal/services/task_log_service.go @@ -0,0 +1,161 @@ +package services + +import ( + "baihu/internal/database" + "baihu/internal/logger" + "baihu/internal/models" + "baihu/internal/utils" + "encoding/json" + "time" +) + +// TaskLogService 任务日志服务 +type TaskLogService struct{} + +// NewTaskLogService 创建任务日志服务 +func NewTaskLogService() *TaskLogService { + return &TaskLogService{} +} + +// CleanConfig 清理配置 +type CleanConfig struct { + Type string `json:"type"` // day 或 count + Keep int `json:"keep"` // 保留天数或条数 +} + +// SaveTaskLog 保存任务日志（通用方法） +func (s *TaskLogService) SaveTaskLog(taskLog *models.TaskLog) error { + if err := database.DB.Create(taskLog).Error; err != nil { + return err + } + + // 更新任务的 last_run + database.DB.Model(&models.Task{}).Where("id = ?", taskLog.TaskID).Update("last_run", time.Now()) + + return nil +} + +// UpdateTaskStats 更新任务统计 +func (s *TaskLogService) UpdateTaskStats(taskID uint, status string) { + sendStatsService := NewSendStatsService() + err := sendStatsService.IncrementStats(taskID, status) + if err != nil { + logger.Errorf("UpdateTaskStats err: %v", err) + return + } +} + +// CleanTaskLogs 清理任务日志 +func (s *TaskLogService) CleanTaskLogs(taskID uint) { + var task models.Task + if err := database.DB.First(&task, taskID).Error; err != nil { + return + } + + if task.CleanConfig == "" { + return + } + + var config CleanConfig + if err := json.Unmarshal([]byte(task.CleanConfig), &config); err != nil { + logger.Errorf("[TaskLog] 解析清理配置失败: %v", err) + return + } + + if config.Keep <= 0 { + return + } + + var deleted int64 + switch config.Type { + case "day": + cutoff := time.Now().AddDate(0, 0, -config.Keep) + result := database.DB.Where("task_id = ? AND created_at < ?", taskID, cutoff).Delete(&models.TaskLog{}) + deleted = result.RowsAffected + case "count": + var boundaryLog models.TaskLog + err := database.DB.Where("task_id = ?", taskID).Order("id DESC").Offset(config.Keep - 1).Limit(1).First(&boundaryLog).Error + if err == nil { + result := database.DB.Where("task_id = ? AND id < ?", taskID, boundaryLog.ID).Delete(&models.TaskLog{}) + deleted = result.RowsAffected + } + } + + if deleted > 0 { + logger.Infof("[TaskLog] 清理任务 #%d 的 %d 条日志", taskID, deleted) + } +} + +// ProcessTaskCompletion 处理任务完成后的所有操作（保存日志、更新统计、清理旧日志） +func (s *TaskLogService) ProcessTaskCompletion(taskLog *models.TaskLog) error { + // 1. 保存日志 + if err := s.SaveTaskLog(taskLog); err != nil { + return err + } + + // 2. 更新统计 + s.UpdateTaskStats(taskLog.TaskID, taskLog.Status) + + // 3. 异步清理旧日志 + go s.CleanTaskLogs(taskLog.TaskID) + + return nil +} + +// CreateTaskLogFromAgentResult 从 Agent 结果创建任务日志 +func (s *TaskLogService) CreateTaskLogFromAgentResult(result *models.AgentTaskResult) (*models.TaskLog, error) { + // 压缩输出 + compressed, err := utils.CompressToBase64(result.Output) + if err != nil { + logger.Errorf("[TaskLog] 压缩日志失败: %v", err) + compressed = "" + } + + taskLog := &models.TaskLog{ + TaskID: result.TaskID, + AgentID: &result.AgentID, + Command: result.Command, + Output: compressed, + Status: result.Status, + Duration: result.Duration, + ExitCode: result.ExitCode, + } + + // 处理开始和结束时间 + if result.StartTime > 0 { + startTime := models.LocalTime(time.Unix(result.StartTime, 0)) + taskLog.StartTime = &startTime + } + if result.EndTime > 0 { + endTime := models.LocalTime(time.Unix(result.EndTime, 0)) + taskLog.EndTime = &endTime + } + + return taskLog, nil +} + +// CreateTaskLogFromLocalExecution 从本地执行结果创建任务日志 +func (s *TaskLogService) CreateTaskLogFromLocalExecution(taskID uint, command, output, status string, duration int64, exitCode int, start, end time.Time) (*models.TaskLog, error) { + // 压缩输出 + compressed, err := utils.CompressToBase64(output) + if err != nil { + logger.Errorf("[TaskLog] 压缩日志失败: %v", err) + compressed = "" + } + + startTime := models.LocalTime(start) + endTime := models.LocalTime(end) + + taskLog := &models.TaskLog{ + TaskID: taskID, + Command: command, + Output: compressed, + Status: status, + Duration: duration, + ExitCode: exitCode, + StartTime: &startTime, + EndTime: &endTime, + } + + return taskLog, nil +}