admin/TaskPool
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: add agent start code

Browse Source
This commit is contained in:
engigu
2025-12-28 11:42:58 +08:00
parent 6d7997c92a
commit 78c3707cce
19 changed files with 1491 additions and 501 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/controllers/agent_controller.go
+300 -74
View File
@@ -1,72 +1,46 @@
package controllers
import (
"baihu/internal/logger"
"baihu/internal/models"
"baihu/internal/services"
"baihu/internal/utils"
"encoding/json"
"net/http"
"strconv"
"strings"
"time"
"github.com/gin-gonic/gin"
"github.com/gorilla/websocket"
)
var agentUpgrader = websocket.Upgrader{
CheckOrigin: func(r *http.Request) bool {
return true
},
}
// AgentController Agent 控制器
type AgentController struct {
agentService *services.AgentService
wsManager *services.AgentWSManager
}
// NewAgentController 创建 Agent 控制器
func NewAgentController() *AgentController {
return &AgentController{
agentService: services.NewAgentService(),
wsManager: services.GetAgentWSManager(),
}
}
// List 获取已审核的 Agent 列表
// List 获取 Agent 列表
func (c *AgentController) List(ctx *gin.Context) {
agents := c.agentService.List()
utils.Success(ctx, agents)
}
// ListPending 获取待审核的 Agent 列表
func (c *AgentController) ListPending(ctx *gin.Context) {
agents := c.agentService.ListPending()
utils.Success(ctx, agents)
}
// Approve 审核通过 Agent
func (c *AgentController) Approve(ctx *gin.Context) {
id, err := strconv.ParseUint(ctx.Param("id"), 10, 32)
if err != nil {
utils.BadRequest(ctx, "无效的 ID")
return
}
agent, err := c.agentService.Approve(uint(id))
if err != nil {
utils.BadRequest(ctx, err.Error())
return
}
utils.Success(ctx, agent)
}
// Reject 拒绝 Agent
func (c *AgentController) Reject(ctx *gin.Context) {
id, err := strconv.ParseUint(ctx.Param("id"), 10, 32)
if err != nil {
utils.BadRequest(ctx, "无效的 ID")
return
}
if err := c.agentService.Reject(uint(id)); err != nil {
utils.ServerError(ctx, err.Error())
return
}
utils.SuccessMsg(ctx, "已拒绝")
}
// Update 更新 Agent
func (c *AgentController) Update(ctx *gin.Context) {
id, err := strconv.ParseUint(ctx.Param("id"), 10, 32)
@@ -86,11 +60,36 @@ func (c *AgentController) Update(ctx *gin.Context) {
return
}
// 获取旧状态
oldAgent := c.agentService.GetByID(uint(id))
if oldAgent == nil {
utils.NotFound(ctx, "Agent 不存在")
return
}
wasEnabled := oldAgent.Enabled
if err := c.agentService.Update(uint(id), req.Name, req.Description, req.Enabled); err != nil {
utils.ServerError(ctx, err.Error())
return
}
// 如果启用状态发生变化，通知 Agent
if wasEnabled != req.Enabled {
if req.Enabled {
// 启用：发送任务列表
c.wsManager.SendToAgent(uint(id), services.WSTypeEnabled, map[string]interface{}{
"message": "Agent 已启用",
})
// 发送任务列表
c.wsManager.BroadcastTasks(uint(id))
} else {
// 禁用：发送禁用消息，Agent 收到后清空任务
c.wsManager.SendToAgent(uint(id), services.WSTypeDisabled, map[string]interface{}{
"message": "Agent 已禁用",
})
}
}
utils.SuccessMsg(ctx, "更新成功")
}
@@ -143,49 +142,19 @@ func (c *AgentController) Register(ctx *gin.Context) {
}
ip := ctx.ClientIP()
agent, err := c.agentService.Register(&req, ip)
agent, token, err := c.agentService.Register(&req, ip)
if err != nil {
utils.ServerError(ctx, err.Error())
utils.BadRequest(ctx, err.Error())
return
}
utils.Success(ctx, gin.H{
"agent_id": agent.ID,
"status": agent.Status,
"message": "注册成功，等待审核",
"token": token,
"message": "注册成功",
})
}
// CheckStatus Agent 检查状态（用于轮询等待审核结果）
func (c *AgentController) CheckStatus(ctx *gin.Context) {
var req struct {
Name string `json:"name"`
}
if err := ctx.ShouldBindJSON(&req); err != nil {
utils.BadRequest(ctx, "参数错误")
return
}
ip := ctx.ClientIP()
agent, err := c.agentService.CheckPendingAgent(req.Name, ip)
if err != nil {
utils.NotFound(ctx, err.Error())
return
}
response := gin.H{
"agent_id": agent.ID,
"status": agent.Status,
}
// 如果已审核通过，返回 Token
if agent.Status != "pending" && agent.Token != "" {
response["token"] = agent.Token
}
utils.Success(ctx, response)
}
// Heartbeat Agent 心跳
func (c *AgentController) Heartbeat(ctx *gin.Context) {
token := c.getAgentToken(ctx)
@@ -348,3 +317,260 @@ func (c *AgentController) ForceUpdate(ctx *gin.Context) {
utils.SuccessMsg(ctx, "已标记强制更新，Agent 下次心跳时将自动更新")
}
// ========== WebSocket ==========
// WSConnect Agent WebSocket 连接
func (c *AgentController) WSConnect(ctx *gin.Context) {
ip := ctx.ClientIP()
// 检查 IP 限流
if allowed, reason := c.wsManager.CheckRateLimit(ip); !allowed {
logger.Warnf("[AgentWS] IP %s 被限流: %s", ip, reason)
ctx.JSON(http.StatusTooManyRequests, gin.H{"error": reason})
return
}
token := ctx.Query("token")
if token == "" {
c.wsManager.RecordConnectFail(ip)
ctx.JSON(http.StatusUnauthorized, gin.H{"error": "缺少 token"})
return
}
machineID := ctx.Query("machine_id")
isNewAgent := false
// 先尝试用 token 查找已有 Agent
agent := c.agentService.GetByToken(token)
// 如果没找到，尝试用令牌注册（会检查 machine_id 是否已存在）
if agent == nil {
var err error
agent, isNewAgent, err = c.agentService.RegisterByToken(token, machineID, ip)
if err != nil {
c.wsManager.RecordConnectFail(ip)
ctx.JSON(http.StatusUnauthorized, gin.H{"error": err.Error()})
return
}
}
if !agent.Enabled {
c.wsManager.RecordConnectFail(ip)
ctx.JSON(http.StatusForbidden, gin.H{"error": "Agent 已禁用"})
return
}
conn, err := agentUpgrader.Upgrade(ctx.Writer, ctx.Request, nil)
if err != nil {
logger.Errorf("[AgentWS] 升级连接失败: %v", err)
return
}
// 连接成功，重置失败计数
c.wsManager.RecordConnectSuccess(ip)
// 注册连接
ac := c.wsManager.Register(agent.ID, conn, ip)
// 更新 Agent 状态
c.agentService.Heartbeat(token, ip, "", "", "", "", "")
// 发送连接成功消息（包含注册状态）
c.wsManager.SendToAgent(agent.ID, services.WSTypeConnected, map[string]interface{}{
"agent_id": agent.ID,
"name": agent.Name,
"is_new_agent": isNewAgent,
"machine_id": machineID,
})
// 启动读写协程
go c.wsWritePump(ac)
go c.wsReadPump(ac, agent)
}
// wsReadPump 读取消息
func (c *AgentController) wsReadPump(ac *services.AgentConnection, agent *models.Agent) {
defer func() {
c.wsManager.Unregister(agent.ID)
}()
ac.Conn.SetReadDeadline(time.Now().Add(90 * time.Second))
ac.Conn.SetPongHandler(func(string) error {
ac.Conn.SetReadDeadline(time.Now().Add(90 * time.Second))
return nil
})
for {
_, message, err := ac.Conn.ReadMessage()
if err != nil {
if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
logger.Warnf("[AgentWS] Agent #%d 读取错误: %v", agent.ID, err)
}
break
}
var msg services.WSMessage
if err := json.Unmarshal(message, &msg); err != nil {
continue
}
c.handleWSMessage(ac, agent, &msg)
}
}
// wsWritePump 写入消息
func (c *AgentController) wsWritePump(ac *services.AgentConnection) {
ticker := time.NewTicker(30 * time.Second)
defer ticker.Stop()
for {
select {
case message, ok := <-ac.Send:
if !ok {
return
}
if err := ac.WriteMessage(message); err != nil {
return
}
case <-ticker.C:
ac.Conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
if err := ac.Conn.WriteMessage(websocket.PingMessage, nil); err != nil {
return
}
}
}
}
// handleWSMessage 处理 WebSocket 消息
func (c *AgentController) handleWSMessage(ac *services.AgentConnection, agent *models.Agent, msg *services.WSMessage) {
switch msg.Type {
case services.WSTypeHeartbeat:
c.handleHeartbeat(ac, agent, msg.Data)
case services.WSTypeTaskResult:
c.handleTaskResult(agent, msg.Data)
case services.WSTypeFetchTasks:
c.handleFetchTasks(agent)
}
}
// handleFetchTasks 处理 Agent 请求任务列表
func (c *AgentController) handleFetchTasks(agent *models.Agent) {
tasks := c.agentService.GetTasks(agent.ID)
c.wsManager.SendToAgent(agent.ID, services.WSTypeTasks, map[string]interface{}{
"tasks": tasks,
})
logger.Infof("[AgentWS] Agent #%d 请求任务列表，返回 %d 个任务", agent.ID, len(tasks))
}
// handleHeartbeat 处理心跳
func (c *AgentController) handleHeartbeat(ac *services.AgentConnection, agent *models.Agent, data json.RawMessage) {
var req struct {
Version string `json:"version"`
BuildTime string `json:"build_time"`
Hostname string `json:"hostname"`
OS string `json:"os"`
Arch string `json:"arch"`
AutoUpdate bool `json:"auto_update"`
}
json.Unmarshal(data, &req)
ac.UpdatePing()
// 更新 Agent 信息（使用连接时保存的 IP）
c.agentService.Heartbeat(agent.Token, ac.IP, req.Version, req.BuildTime, req.Hostname, req.OS, req.Arch)
// 检查是否需要更新
latestVersion := c.agentService.GetLatestVersion()
needUpdate := latestVersion != "" && req.Version != "" && req.Version != latestVersion
forceUpdate := agent.ForceUpdate
if forceUpdate && needUpdate {
c.agentService.ClearForceUpdate(agent.ID)
}
// 发送心跳响应
response := map[string]interface{}{
"agent_id": agent.ID,
"name": agent.Name,
"need_update": needUpdate,
"force_update": forceUpdate,
"latest_version": latestVersion,
}
c.wsManager.SendToAgent(agent.ID, services.WSTypeHeartbeatAck, response)
}
// handleTaskResult 处理任务结果
func (c *AgentController) handleTaskResult(agent *models.Agent, data json.RawMessage) {
var result models.AgentTaskResult
if err := json.Unmarshal(data, &result); err != nil {
return
}
result.AgentID = agent.ID
c.agentService.ReportResult(&result)
}
// NotifyTaskUpdate 通知 Agent 任务更新
func (c *AgentController) NotifyTaskUpdate(agentID uint) {
c.wsManager.BroadcastTasks(agentID)
}
// ========== 注册码管理 ==========
// ListRegCodes 获取注册码列表
func (c *AgentController) ListRegCodes(ctx *gin.Context) {
codes := c.agentService.ListRegCodes()
utils.Success(ctx, codes)
}
// CreateRegCode 创建注册码
func (c *AgentController) CreateRegCode(ctx *gin.Context) {
var req struct {
Remark string `json:"remark"`
MaxUses int `json:"max_uses"`
ExpiresAt string `json:"expires_at"` // 格式: 2006-01-02 15:04:05
}
if err := ctx.ShouldBindJSON(&req); err != nil {
utils.BadRequest(ctx, "参数错误")
return
}
var expiresAt *time.Time
if req.ExpiresAt != "" {
t, err := time.ParseInLocation("2006-01-02 15:04:05", req.ExpiresAt, time.Local)
if err != nil {
utils.BadRequest(ctx, "过期时间格式错误")
return
}
expiresAt = &t
}
code, err := c.agentService.CreateRegCode(req.Remark, req.MaxUses, expiresAt)
if err != nil {
utils.ServerError(ctx, err.Error())
return
}
utils.Success(ctx, code)
}
// DeleteRegCode 删除注册码
func (c *AgentController) DeleteRegCode(ctx *gin.Context) {
id, err := strconv.ParseUint(ctx.Param("id"), 10, 32)
if err != nil {
utils.BadRequest(ctx, "无效的 ID")
return
}
if err := c.agentService.DeleteRegCode(uint(id)); err != nil {
utils.ServerError(ctx, err.Error())
return
}
utils.SuccessMsg(ctx, "删除成功")
}
internal/controllers/task_controller.go
+55 -10
View File
@@ -12,14 +12,16 @@ import (
)
type TaskController struct {
taskService *services.TaskService
cronService *services.CronService
taskService *services.TaskService
cronService *services.CronService
agentWSManager *services.AgentWSManager
}
func NewTaskController(taskService *services.TaskService, cronService *services.CronService) *TaskController {
return &TaskController{
taskService: taskService,
cronService: cronService,
taskService: taskService,
cronService: cronService,
agentWSManager: services.GetAgentWSManager(),
}
}
@@ -57,6 +59,7 @@ func (tc *TaskController) CreateTask(c *gin.Context) {
WorkDir string `json:"work_dir"`
CleanConfig string `json:"clean_config"`
Envs string `json:"envs"`
AgentID *uint `json:"agent_id"`
}
if err := c.ShouldBindJSON(&req); err != nil {
@@ -78,8 +81,14 @@ func (tc *TaskController) CreateTask(c *gin.Context) {
// 转换为绝对路径
workDir := resolveWorkDir(req.WorkDir)
task := tc.taskService.CreateTask(req.Name, req.Command, req.Schedule, req.Timeout, workDir, req.CleanConfig, req.Envs, req.Type, req.Config)
tc.cronService.AddTask(task)
task := tc.taskService.CreateTask(req.Name, req.Command, req.Schedule, req.Timeout, workDir, req.CleanConfig, req.Envs, req.Type, req.Config, req.AgentID)
// 如果是 Agent 任务，通知 Agent；否则添加到本地 cron
if task.AgentID != nil && *task.AgentID > 0 {
tc.agentWSManager.BroadcastTasks(*task.AgentID)
} else {
tc.cronService.AddTask(task)
}
utils.Success(c, task)
}
@@ -115,6 +124,13 @@ func (tc *TaskController) UpdateTask(c *gin.Context) {
return
}
// 获取旧任务信息（用于判断 agent 变更）
oldTask := tc.taskService.GetTaskByID(id)
var oldAgentID *uint
if oldTask != nil {
oldAgentID = oldTask.AgentID
}
var req struct {
Name string `json:"name"`
Command string `json:"command"`
@@ -126,6 +142,7 @@ func (tc *TaskController) UpdateTask(c *gin.Context) {
CleanConfig string `json:"clean_config"`
Envs string `json:"envs"`
Enabled bool `json:"enabled"`
AgentID *uint `json:"agent_id"`
}
if err := c.ShouldBindJSON(&req); err != nil {
@@ -140,16 +157,32 @@ func (tc *TaskController) UpdateTask(c *gin.Context) {
}
}
task := tc.taskService.UpdateTask(id, req.Name, req.Command, req.Schedule, req.Timeout, resolveWorkDir(req.WorkDir), req.CleanConfig, req.Envs, req.Enabled, req.Type, req.Config)
task := tc.taskService.UpdateTask(id, req.Name, req.Command, req.Schedule, req.Timeout, resolveWorkDir(req.WorkDir), req.CleanConfig, req.Envs, req.Enabled, req.Type, req.Config, req.AgentID)
if task == nil {
utils.NotFound(c, "任务不存在")
return
}
if task.Enabled {
tc.cronService.AddTask(task)
} else {
// 处理任务调度
if task.AgentID != nil && *task.AgentID > 0 {
// Agent 任务：从本地 cron 移除，通知 Agent
tc.cronService.RemoveTask(task.ID)
tc.agentWSManager.BroadcastTasks(*task.AgentID)
// 如果 agent 变更了，也通知旧 agent
if oldAgentID != nil && *oldAgentID > 0 && *oldAgentID != *task.AgentID {
tc.agentWSManager.BroadcastTasks(*oldAgentID)
}
} else {
// 本地任务
if task.Enabled {
tc.cronService.AddTask(task)
} else {
tc.cronService.RemoveTask(task.ID)
}
// 如果之前是 agent 任务，通知旧 agent 移除
if oldAgentID != nil && *oldAgentID > 0 {
tc.agentWSManager.BroadcastTasks(*oldAgentID)
}
}
utils.Success(c, task)
@@ -162,6 +195,13 @@ func (tc *TaskController) DeleteTask(c *gin.Context) {
return
}
// 获取任务信息（用于通知 agent）
task := tc.taskService.GetTaskByID(id)
var agentID *uint
if task != nil {
agentID = task.AgentID
}
tc.cronService.RemoveTask(uint(id))
success := tc.taskService.DeleteTask(id)
@@ -170,5 +210,10 @@ func (tc *TaskController) DeleteTask(c *gin.Context) {
return
}
// 如果是 agent 任务，通知 agent
if agentID != nil && *agentID > 0 {
tc.agentWSManager.BroadcastTasks(*agentID)
}
utils.SuccessMsg(c, "删除成功")
}