heicode-mananger/heicode/controller/agent_async_deploy_test.go

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")
}