From 94c6e1b7fb1005d50b17dd19909d5ee94bead5ad Mon Sep 17 00:00:00 2001
From: engigu <sayheya@qq.com>
Date: Tue, 30 Dec 2025 21:09:36 +0800
Subject: [PATCH] fix:  agent tasks cmd

---
 agent/agent.go                           |  1 +
 internal/controllers/agent_controller.go | 14 ++++++++++++++
 internal/services/agent_service.go       |  9 +++++++++
 3 files changed, 24 insertions(+)

diff --git a/agent/agent.go b/agent/agent.go
index 171b917..34ee0ff 100644
--- a/agent/agent.go
+++ b/agent/agent.go
@@ -490,6 +490,7 @@ func (a *Agent) doRequest(method, path string, body interface{}) (*http.Response
 
 	req.Header.Set("Authorization", "Bearer "+a.config.Token)
 	req.Header.Set("Content-Type", "application/json")
+	req.Header.Set("X-Machine-ID", a.machineID)
 
 	return a.client.Do(req)
 }
diff --git a/internal/controllers/agent_controller.go b/internal/controllers/agent_controller.go
index dc4d127..73ef0ee 100644
--- a/internal/controllers/agent_controller.go
+++ b/internal/controllers/agent_controller.go
@@ -207,7 +207,21 @@ func (c *AgentController) GetTasks(ctx *gin.Context) {
 		return
 	}
 
+	// 先尝试通过 token 查找 Agent
 	agent := c.agentService.GetByToken(token)
+	
+	// 如果找不到，尝试验证令牌并通过 machine_id 查找
+	if agent == nil {
+		machineID := ctx.GetHeader("X-Machine-ID")
+		if machineID != "" {
+			// 验证令牌是否有效
+			if _, err := c.agentService.ValidateToken(token); err == nil {
+				// 令牌有效，尝试通过 machine_id 查找 Agent
+				agent = c.agentService.GetByMachineID(machineID)
+			}
+		}
+	}
+
 	if agent == nil {
 		utils.Unauthorized(ctx, "无效的 Token")
 		return
diff --git a/internal/services/agent_service.go b/internal/services/agent_service.go
index 2a97743..ff02abf 100644
--- a/internal/services/agent_service.go
+++ b/internal/services/agent_service.go
@@ -230,6 +230,15 @@ func (s *AgentService) GetByToken(token string) *models.Agent {
 	return &agent
 }
 
+// GetByMachineID 根据 MachineID 获取 Agent
+func (s *AgentService) GetByMachineID(machineID string) *models.Agent {
+	var agent models.Agent
+	if err := database.DB.Where("machine_id = ?", machineID).First(&agent).Error; err != nil {
+		return nil
+	}
+	return &agent
+}
+
 // List 获取 Agent 列表
 func (s *AgentService) List() []models.Agent {
 	var agents []models.Agent