diff --git a/heicode/controller/agent_async_deploy_test.go b/heicode/controller/agent_async_deploy_test.go index 15dae80..620382c 100644 --- a/heicode/controller/agent_async_deploy_test.go +++ b/heicode/controller/agent_async_deploy_test.go @@ -3,6 +3,8 @@ package controller import ( "net/http" "net/http/httptest" + "strings" + "sync" "testing" "time" @@ -76,3 +78,61 @@ func TestStartTemplateAgentAsync_FailureMarksFailed(t *testing.T) { }, 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") +} diff --git a/heicode/controller/agent_template_handlers.go b/heicode/controller/agent_template_handlers.go index bd1e7b8..7c92a58 100644 --- a/heicode/controller/agent_template_handlers.go +++ b/heicode/controller/agent_template_handlers.go @@ -240,20 +240,32 @@ func startTemplateAgentAsync(deploymentID string, modelTokenID int, args amStart if len(reason) > 480 { reason = strings.ToValidUTF8(reason[:480], "") } - // Mark failed and revoke the minted model key (the agent never started). - _ = model.DB.Model(&model.AgentDeployment{}). - Where("deployment_id = ?", deploymentID). + // 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, - }).Error - revokeAgentModelToken(modelTokenID) + }) + // 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 = ?", deploymentID). + Where("deployment_id = ? AND LOWER(status) = ?", deploymentID, "pending"). Updates(map[string]any{ "subdomain": result.Subdomain, "runtime_deployment_id": result.RuntimeID, @@ -265,11 +277,22 @@ func startTemplateAgentAsync(deploymentID string, modelTokenID int, args amStart common.SysLog("startTemplateAgentAsync: persist result failed for " + deploymentID + ": " + res.Error.Error()) return } - if res.RowsAffected == 0 && strings.TrimSpace(result.RuntimeID) != "" { - // Record was deleted (user cancelled) before AM finished starting — clean - // up the orphan agent in AM so it doesn't keep running untracked. - if delErr := amDeleteTemplateAgent(context.Background(), result.RuntimeID); delErr != nil { - common.SysLog("startTemplateAgentAsync orphan cleanup failed: " + delErr.Error()) + 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()) + } } } }()