Merge pull request #18 from xmindlab-heicode/fix/async-agent-deploy

fix(agent): 部署 agent 改异步，避免 Azure 网关 504 + 回滚
2026-06-08 18:08:17 +08:00
parent 206fee1f78 f6da9c541e
commit 03b9b24e45
2 changed files with 246 additions and 36 deletions
@@ -0,0 +1,138 @@
+package controller
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/heicode/manager/model"
+	"github.com/stretchr/testify/require"
+)
+
+// startTemplateAgentAsync must NOT block the caller on the (slow) AM start, and
+// must flip the pre-persisted Pending record to running once AM responds — this
+// is the core of the 504 fix (deploy returns immediately; status lands later).
+func TestStartTemplateAgentAsync_PendingToRunning(t *testing.T) {
+	setupResourceControllerTestDB(t)
+	require.NoError(t, model.DB.AutoMigrate(&model.AgentDeployment{}))
+
+	// Slow AM mock — simulates the ~30s start the handler must not block on.
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		time.Sleep(150 * time.Millisecond)
+		w.Header().Set("Content-Type", "application/json")
+		_, _ = w.Write([]byte(`{"success":true,"data":{"namespace":"rt-async-1","access_info":{"domain":"https://async.agents.example"},"status":"running"}}`))
+	}))
+	defer srv.Close()
+	t.Setenv("AGENT_RUNTIME_BASE_URL", srv.URL)
+
+	// What HeicodeDeployAgent persists before returning.
+	require.NoError(t, model.DB.Create(&model.AgentDeployment{
+		DeploymentID: "dep_async_run",
+		UserID:       "22",
+		TemplateID:   "architect",
+		Status:       "Pending",
+	}).Error)
+
+	start := time.Now()
+	startTemplateAgentAsync("dep_async_run", 0, amStartArgs{ManagerDeploymentID: "dep_async_run", UserID: "22", TemplateKey: "architect"})
+	require.Less(t, time.Since(start), 50*time.Millisecond, "must not block on the AM start call")
+
+	var got model.AgentDeployment
+	require.Eventually(t, func() bool {
+		model.DB.Where("deployment_id = ?", "dep_async_run").First(&got)
+		return got.Status == "running"
+	}, 3*time.Second, 20*time.Millisecond, "Pending should flip to running after async AM start")
+	require.Equal(t, "rt-async-1", got.RuntimeDeploymentID)
+	require.Equal(t, "https://async.agents.example", got.Subdomain)
+}
+
+// On AM failure the record must be marked failed (with a reason), not left
+// stuck on Pending — so the client surfaces a failure instead of a hang.
+func TestStartTemplateAgentAsync_FailureMarksFailed(t *testing.T) {
+	setupResourceControllerTestDB(t)
+	require.NoError(t, model.DB.AutoMigrate(&model.AgentDeployment{}))
+
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusInternalServerError)
+		_, _ = w.Write([]byte(`{"success":false,"message":"AM boom"}`))
+	}))
+	defer srv.Close()
+	t.Setenv("AGENT_RUNTIME_BASE_URL", srv.URL)
+
+	require.NoError(t, model.DB.Create(&model.AgentDeployment{
+		DeploymentID: "dep_async_fail",
+		UserID:       "22",
+		TemplateID:   "architect",
+		Status:       "Pending",
+	}).Error)
+
+	startTemplateAgentAsync("dep_async_fail", 0, amStartArgs{ManagerDeploymentID: "dep_async_fail", UserID: "22", TemplateKey: "architect"})
+
+	var got model.AgentDeployment
+	require.Eventually(t, func() bool {
+		model.DB.Where("deployment_id = ?", "dep_async_fail").First(&got)
+		return got.Status == "failed"
+	}, 3*time.Second, 20*time.Millisecond, "AM failure should mark the record failed")
+	require.NotEmpty(t, got.FailureReason)
+}
+
+// Concurrency boundary: if the user stops the agent while AM is still starting,
+// the async goroutine must NOT resurrect it to running, and must delete the
+// runtime AM just started (orphan cleanup). The AM start is held until the test
+// has marked the record stopped, so the conditional update always runs after.
+func TestStartTemplateAgentAsync_StoppedDuringStart_NoResurrect(t *testing.T) {
+	setupResourceControllerTestDB(t)
+	require.NoError(t, model.DB.AutoMigrate(&model.AgentDeployment{}))
+
+	release := make(chan struct{})
+	var mu sync.Mutex
+	var deletedRuntime string
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		switch {
+		case r.Method == http.MethodPost && r.URL.Path == "/agents":
+			<-release // block the start until the test marks the row stopped
+			w.Header().Set("Content-Type", "application/json")
+			_, _ = w.Write([]byte(`{"success":true,"data":{"namespace":"rt-race-1","access_info":{"domain":"https://race.agents.example"},"status":"running"}}`))
+		case r.Method == http.MethodDelete && strings.HasPrefix(r.URL.Path, "/agents/"):
+			mu.Lock()
+			deletedRuntime = strings.TrimPrefix(r.URL.Path, "/agents/")
+			mu.Unlock()
+			w.Header().Set("Content-Type", "application/json")
+			_, _ = w.Write([]byte(`{"success":true}`))
+		default:
+			w.WriteHeader(http.StatusNotFound)
+		}
+	}))
+	defer srv.Close()
+	t.Setenv("AGENT_RUNTIME_BASE_URL", srv.URL)
+
+	require.NoError(t, model.DB.Create(&model.AgentDeployment{
+		DeploymentID: "dep_race_stop",
+		UserID:       "22",
+		TemplateID:   "architect",
+		Status:       "Pending",
+	}).Error)
+
+	startTemplateAgentAsync("dep_race_stop", 0, amStartArgs{ManagerDeploymentID: "dep_race_stop", UserID: "22", TemplateKey: "architect"})
+
+	// User stops the agent while AM is still starting. Done before releasing AM, so
+	// the goroutine's Pending-guarded update is guaranteed to see "stopped".
+	require.NoError(t, model.DB.Model(&model.AgentDeployment{}).
+		Where("deployment_id = ?", "dep_race_stop").
+		Updates(map[string]any{"status": "stopped"}).Error)
+	close(release)
+
+	require.Eventually(t, func() bool {
+		mu.Lock()
+		defer mu.Unlock()
+		return deletedRuntime == "rt-race-1"
+	}, 3*time.Second, 20*time.Millisecond, "orphan runtime must be deleted in AM")
+
+	var row model.AgentDeployment
+	require.NoError(t, model.DB.Where("deployment_id = ?", "dep_race_stop").First(&row).Error)
+	require.Equal(t, "stopped", row.Status, "stopped record must not be flipped back to running")
+	require.Empty(t, row.RuntimeDeploymentID, "runtime id must not be backfilled onto a stopped record")
+}
@@ -199,8 +199,36 @@ func HeicodeDeployAgent(c *gin.Context) {
 	env["AGENT_ACCESS_TOKEN"] = agentAccessToken
 	env["HEICODE_AGENT_ID"] = deploymentID

-	// Ask AM to start the agent with the template definition (.md) + env injected.
-	result, err := amStartTemplateAgent(c.Request.Context(), amStartArgs{
+	// Persist the agent as Pending and return immediately; start it in AM in the
+	// background. Starting an agent in AM takes ~30s — longer than the Azure
+	// gateway's ~20s timeout — so doing it inside the request returned 504, which
+	// cancelled the request context, aborted the AM call and rolled the deploy
+	// back. Now the client gets the Pending agent at once and polls
+	// GET /api/heicode/agents (which refreshes live status from AM), seeing
+	// Pending -> running once AM is up (or -> failed if AM start failed).
+	now := agentNow()
+	nowMs := time.Now().UnixMilli()
+	row := model.AgentDeployment{
+		DeploymentID:   deploymentID,
+		UserID:         strconv.Itoa(userID),
+		TemplateID:     req.TemplateID,
+		AccessToken:    sealAgentToken(agentAccessToken),
+		BindingIDsJSON: string(bindingIDsJSON),
+		Status:         "Pending",
+		ModelTokenID:   modelTokenID,
+		CreatedAtText:  now,
+		UpdatedAtText:  now,
+		CreatedAtMs:    nowMs,
+		UpdatedAtMs:    nowMs,
+	}
+	if err := model.DB.Create(&row).Error; err != nil {
+		common.SysLog("HeicodeDeployAgent persist: " + err.Error())
+		revokeAgentModelToken(modelTokenID)
+		agentError(c, "DEPLOYMENT_PERSIST_FAILED", "failed to persist agent")
+		return
+	}
+
+	startTemplateAgentAsync(deploymentID, modelTokenID, amStartArgs{
 		ManagerDeploymentID: deploymentID,
 		UserID:              strconv.Itoa(userID),
 		TemplateKey:         tpl.TemplateKey,
@@ -209,43 +237,87 @@ func HeicodeDeployAgent(c *gin.Context) {
 		Env:                 env,
 		CallbackURL:         agentRuntimeCallbackURL(),
 	})
-	if err != nil {
-		revokeAgentModelToken(modelTokenID) // don't leak the minted key
-		agentError(c, "RUNTIME_UNAVAILABLE", "failed to start agent: "+err.Error())
-		return
-	}

-	now := agentNow()
-	nowMs := time.Now().UnixMilli()
-	row := model.AgentDeployment{
-		DeploymentID:        deploymentID,
-		UserID:              strconv.Itoa(userID),
-		TemplateID:          req.TemplateID,
-		Subdomain:           result.Subdomain,
-		AccessToken:         sealAgentToken(agentAccessToken),
-		BindingIDsJSON:      string(bindingIDsJSON),
-		RuntimeDeploymentID: result.RuntimeID,
-		Status:              firstNonEmpty(result.Status, "running"),
-		ModelTokenID:        modelTokenID,
-		CreatedAtText:       now,
-		UpdatedAtText:       now,
-		CreatedAtMs:         nowMs,
-		UpdatedAtMs:         nowMs,
-	}
-	if err := model.DB.Create(&row).Error; err != nil {
-		common.SysLog("HeicodeDeployAgent persist: " + err.Error())
-		revokeAgentModelToken(modelTokenID)
-		// We started an agent in AM but failed to record it: roll back the
-		// orphan so it does not leak/keep running with no Manager record.
-		if strings.TrimSpace(result.RuntimeID) != "" {
-			if delErr := amDeleteTemplateAgent(c.Request.Context(), result.RuntimeID); delErr != nil {
-				common.SysLog("HeicodeDeployAgent orphan cleanup failed: " + delErr.Error())
+	common.ApiSuccess(c, templateAgentResponse(row))
+}
+
+// startTemplateAgentAsync starts the template agent in AM off the request path
+// and records the outcome. The deploy handler returns the Pending agent at once
+// so it never blocks ~30s on AM (which exceeded the Azure gateway timeout → 504
+// + rollback). Uses context.Background() because the request context is gone
+// once the handler returned; the AM HTTP call is still bounded by amTemplateDo's
+// own client timeout. The client observes Pending -> running (or -> failed) by
+// polling the agent list, which refreshes live status from AM.
+func startTemplateAgentAsync(deploymentID string, modelTokenID int, args amStartArgs) {
+	go func() {
+		result, err := amStartTemplateAgent(context.Background(), args)
+		if model.DB == nil {
+			return
+		}
+		now := agentNow()
+		nowMs := time.Now().UnixMilli()
+		if err != nil {
+			common.SysLog("startTemplateAgentAsync: AM start failed for " + deploymentID + ": " + err.Error())
+			reason := err.Error()
+			if len(reason) > 480 {
+				reason = strings.ToValidUTF8(reason[:480], "")
+			}
+			// Only flip Pending->failed. If the user stopped/deleted/cancelled during
+			// the AM start window the status is no longer Pending (or the row is gone),
+			// and we must NOT clobber that terminal intent. The status guard makes the
+			// transition own-or-nothing; LOWER() because the row is stored as "Pending".
+			res := model.DB.Model(&model.AgentDeployment{}).
+				Where("deployment_id = ? AND LOWER(status) = ?", deploymentID, "pending").
+				Updates(map[string]any{
+					"status":          "failed",
+					"failure_reason":  reason,
+					"updated_at_text": now,
+					"updated_at_ms":   nowMs,
+				})
+			// Revoke the minted key ONLY when we owned the Pending->failed transition
+			// (the agent never started). If the user already acted, leave the token to
+			// that path: delete revokes it, stop intentionally retains it. Avoids both
+			// double-revoke and revoking a key the stop path means to keep.
+			if res.Error == nil && res.RowsAffected > 0 {
+				revokeAgentModelToken(modelTokenID)
+			}
+			return
+		}
+		// Success: flip Pending->running and fill in the runtime. Same status guard —
+		// if the user stopped/deleted/cancelled while AM was starting, this updates 0
+		// rows and we must clean up the now-orphaned runtime instead of resurrecting it.
+		res := model.DB.Model(&model.AgentDeployment{}).
+			Where("deployment_id = ? AND LOWER(status) = ?", deploymentID, "pending").
+			Updates(map[string]any{
+				"subdomain":             result.Subdomain,
+				"runtime_deployment_id": result.RuntimeID,
+				"status":                firstNonEmpty(result.Status, "running"),
+				"updated_at_text":       now,
+				"updated_at_ms":         nowMs,
+			})
+		if res.Error != nil {
+			common.SysLog("startTemplateAgentAsync: persist result failed for " + deploymentID + ": " + res.Error.Error())
+			return
+		}
+		if res.RowsAffected == 0 {
+			// No longer Pending: the user stopped/deleted/cancelled while AM was still
+			// starting. Re-read to log which case, then delete the orphan runtime AM
+			// just started — the action is identical either way (clean up, never
+			// resurrect), so the record is never flipped back to running. Token handling
+			// is left to the user's stop/delete path (delete revoked it; stop keeps it).
+			var cur model.AgentDeployment
+			if model.DB.Where("deployment_id = ?", deploymentID).First(&cur).Error != nil {
+				common.SysLog("startTemplateAgentAsync: " + deploymentID + " deleted during AM start; cleaning orphan runtime")
+			} else {
+				common.SysLog("startTemplateAgentAsync: " + deploymentID + " no longer Pending (status=" + cur.Status + ") during AM start; cleaning orphan runtime, preserving user state")
+			}
+			if strings.TrimSpace(result.RuntimeID) != "" {
+				if delErr := amDeleteTemplateAgent(context.Background(), result.RuntimeID); delErr != nil {
+					common.SysLog("startTemplateAgentAsync orphan cleanup failed for " + deploymentID + ": " + delErr.Error())
+				}
 			}
 		}
-		agentError(c, "DEPLOYMENT_PERSIST_FAILED", "failed to persist agent")
-		return
-	}
-	common.ApiSuccess(c, templateAgentResponse(row))
+	}()
 }

 // findUserTemplateAgent loads a deployed template agent owned by the caller.