Merge pull request #827 from chaitin/fix/vm-idle-clickhouse-guard

修复虚拟机空闲回收误判

Author: yokowu

backend/biz/host/repo/host.go

@@ -21,6 +21,7 @@ import (
 	"github.com/chaitin/MonkeyCode/backend/db/model"
 	"github.com/chaitin/MonkeyCode/backend/db/predicate"
 	"github.com/chaitin/MonkeyCode/backend/db/projecttask"
+	"github.com/chaitin/MonkeyCode/backend/db/task"
 	"github.com/chaitin/MonkeyCode/backend/db/taskvirtualmachine"
 	"github.com/chaitin/MonkeyCode/backend/db/teamgroup"
 	"github.com/chaitin/MonkeyCode/backend/db/user"
@@ -610,6 +611,7 @@ func (h *HostRepo) CompleteCreateVirtualMachine(ctx context.Context, id string,
 func (h *HostRepo) DeleteVirtualMachine(ctx context.Context, uid uuid.UUID, hostID, id string, fn func(*db.VirtualMachine) error) error {
 	return entx.WithTx2(ctx, h.db, func(tx *db.Tx) error {
 		vm, err := tx.VirtualMachine.Query().
+			WithTasks().
 			Where(virtualmachine.ID(id)).
 			Where(virtualmachine.UserID(uid)).
 			First(ctx)
@@ -627,6 +629,27 @@ func (h *HostRepo) DeleteVirtualMachine(ctx context.Context, uid uuid.UUID, host
 			return err
 		}
 
+		taskIDs := make([]uuid.UUID, 0, len(vm.Edges.Tasks))
+		for _, tk := range vm.Edges.Tasks {
+			if tk == nil {
+				continue
+			}
+			taskIDs = append(taskIDs, tk.ID)
+		}
+		if len(taskIDs) > 0 {
+			_, err := tx.Task.Update().
+				Where(
+					task.IDIn(taskIDs...),
+					task.StatusNotIn(consts.TaskStatusFinished, consts.TaskStatusError),
+				).
+				SetStatus(consts.TaskStatusFinished).
+				SetCompletedAt(time.Now()).
+				Save(ctx)
+			if err != nil {
+				return err
+			}
+		}
+
 		_, err = tx.VirtualMachine.Delete().Where(virtualmachine.ID(id)).Exec(ctx)
 		if err != nil {
 			return err

backend/biz/host/usecase/host.go

@@ -533,7 +533,7 @@ func (h *HostUsecase) CreateVM(ctx context.Context, user *domain.User, req *doma
 
 // DeleteVM 删除虚拟机
 func (h *HostUsecase) DeleteVM(ctx context.Context, uid uuid.UUID, hostID, vmID string) error {
-	h.logger.InfoContext(ctx, "delete vm", "vmID", vmID)
+	h.logger.InfoContext(ctx, "delete vm", "vmID", vmID, "user_id", uid, "host_id", hostID)
 	return h.repo.DeleteVirtualMachine(ctx, uid, hostID, vmID, func(vm *db.VirtualMachine) error {
 		if err := h.taskflow.VirtualMachiner().Delete(ctx, &taskflow.DeleteVirtualMachineReq{
 			UserID: uid.String(),

backend/biz/host/usecase/host_test.go

@@ -23,6 +23,7 @@ import (
 	"github.com/chaitin/MonkeyCode/backend/db"
 	"github.com/chaitin/MonkeyCode/backend/db/enttest"
 	"github.com/chaitin/MonkeyCode/backend/domain"
+	"github.com/chaitin/MonkeyCode/backend/pkg/delayqueue"
 	"github.com/chaitin/MonkeyCode/backend/pkg/taskflow"
 )
 
@@ -337,6 +338,97 @@ func TestHostUsecase_markRecycledTasksFinished(t *testing.T) {
 	}
 }
 
+func TestHostUsecase_DeleteVMFinishesBoundTasks(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+	client := enttest.Open(t, "sqlite3", "file:host-usecase-delete-vm-finish-task-test?mode=memory&cache=shared&_fk=1")
+	defer client.Close()
+
+	userID := uuid.New()
+	if _, err := client.User.Create().
+		SetID(userID).
+		SetName("tester").
+		SetRole(consts.UserRoleIndividual).
+		SetStatus(consts.UserStatusActive).
+		Save(ctx); err != nil {
+		t.Fatalf("create user: %v", err)
+	}
+
+	hostID := "host-1"
+	if _, err := client.Host.Create().
+		SetID(hostID).
+		SetUserID(userID).
+		SetHostname("host").
+		Save(ctx); err != nil {
+		t.Fatalf("create host: %v", err)
+	}
+
+	vmID := "vm-1"
+	if _, err := client.VirtualMachine.Create().
+		SetID(vmID).
+		SetHostID(hostID).
+		SetUserID(userID).
+		SetName("vm").
+		SetEnvironmentID("env-1").
+		Save(ctx); err != nil {
+		t.Fatalf("create vm: %v", err)
+	}
+
+	taskID := uuid.New()
+	if _, err := client.Task.Create().
+		SetID(taskID).
+		SetUserID(userID).
+		SetKind(consts.TaskTypeDevelop).
+		SetContent("content").
+		SetStatus(consts.TaskStatusProcessing).
+		Save(ctx); err != nil {
+		t.Fatalf("create task: %v", err)
+	}
+	if _, err := client.TaskVirtualMachine.Create().
+		SetID(uuid.New()).
+		SetTaskID(taskID).
+		SetVirtualmachineID(vmID).
+		Save(ctx); err != nil {
+		t.Fatalf("bind task vm: %v", err)
+	}
+
+	i := do.New()
+	do.ProvideValue(i, client)
+	redisServer := miniredis.RunT(t)
+	redisClient := redis.NewClient(&redis.Options{Addr: redisServer.Addr()})
+	t.Cleanup(func() { _ = redisClient.Close() })
+	logger := slog.New(slog.NewTextHandler(io.Discard, nil))
+	do.ProvideValue(i, &config.Config{})
+	do.ProvideValue(i, logger)
+	do.ProvideValue(i, redisClient)
+	hostRepo, err := repo.NewHostRepo(i)
+	if err != nil {
+		t.Fatalf("new host repo: %v", err)
+	}
+	u := &HostUsecase{
+		repo:          hostRepo,
+		taskflow:      &preinsertTaskflowStub{vm: &preinsertVMCreateStub{db: client}},
+		logger:        logger,
+		vmexpireQueue: delayqueue.NewVMExpireQueue(redisClient, logger),
+	}
+
+	if err := u.DeleteVM(ctx, userID, hostID, vmID); err != nil {
+		t.Fatalf("DeleteVM() error = %v", err)
+	}
+
+	gotTask, err := client.Task.Get(ctx, taskID)
+	if err != nil {
+		t.Fatalf("query task: %v", err)
+	}
+	if gotTask.Status != consts.TaskStatusFinished {
+		t.Fatalf("task status = %s, want %s", gotTask.Status, consts.TaskStatusFinished)
+	}
+	if gotTask.CompletedAt.IsZero() {
+		t.Fatal("expected task completed_at to be set")
+	}
+}
+
 func TestHostUsecase_CreateVMPreinsertsVirtualMachineBeforeTaskflowCreate(t *testing.T) {
 	t.Parallel()
 

backend/biz/vmidle/usecase/policy_test.go

@@ -14,10 +14,12 @@ import (
 	"github.com/redis/go-redis/v9"
 
 	"github.com/chaitin/MonkeyCode/backend/config"
+	"github.com/chaitin/MonkeyCode/backend/consts"
 	"github.com/chaitin/MonkeyCode/backend/db"
 	"github.com/chaitin/MonkeyCode/backend/db/enttest"
 	"github.com/chaitin/MonkeyCode/backend/db/virtualmachine"
 	"github.com/chaitin/MonkeyCode/backend/domain"
+	"github.com/chaitin/MonkeyCode/backend/pkg/delayqueue"
 	"github.com/chaitin/MonkeyCode/backend/pkg/taskflow"
 )
 
@@ -107,6 +109,104 @@ func TestRefreshCachesNotFoundVM(t *testing.T) {
 	}
 }
 
+func TestShouldSkipRecycleForRecentTaskLastActiveAtSkipsClickHouse(t *testing.T) {
+	ctx := context.Background()
+	now := time.Date(2026, 7, 7, 10, 34, 38, 0, time.UTC)
+	taskID := uuid.New()
+	vm := &db.VirtualMachine{ID: "vm-recent-pg-activity", UserID: uuid.New()}
+	vm.Edges.Tasks = []*db.Task{{
+		ID:           taskID,
+		Status:       consts.TaskStatusProcessing,
+		LastActiveAt: now.Add(-time.Minute),
+	}}
+	redisClient := newTestRedis(t)
+	logger := slog.Default()
+	repo := &refreshHostRepoStub{vm: vm}
+	activity := &taskLogActivityStub{err: errors.New("unexpected clickhouse query")}
+	r := &vmIdleRefresher{
+		cfg:             &config.Config{VMIdle: config.VMIdle{SleepSeconds: 600, RecycleSeconds: 3600}},
+		redis:           redisClient,
+		logger:          logger,
+		hostRepo:        repo,
+		taskLogActivity: activity,
+		sleepQueue:      delayqueue.NewVMSleepQueue(redisClient, logger),
+		notifyQueue:     delayqueue.NewVMNotifyQueue(redisClient, logger),
+		recycleQueue:    delayqueue.NewVMRecycleQueue(redisClient, logger),
+	}
+
+	skip, err := r.shouldSkipRecycleForRecentActivity(ctx, vm, now)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !skip {
+		t.Fatal("expected recent task last_active_at to skip recycle")
+	}
+	if activity.calls != 0 {
+		t.Fatalf("clickhouse calls = %d, want 0", activity.calls)
+	}
+	if repo.getVirtualMachineCalls != 1 {
+		t.Fatalf("GetVirtualMachine calls = %d, want 1", repo.getVirtualMachineCalls)
+	}
+}
+
+func TestShouldSkipRecycleForRecentTaskLogRefreshesVM(t *testing.T) {
+	ctx := context.Background()
+	now := time.Date(2026, 7, 7, 10, 34, 38, 0, time.UTC)
+	taskID := uuid.New()
+	vm := &db.VirtualMachine{ID: "vm-recent-log", UserID: uuid.New()}
+	vm.Edges.Tasks = []*db.Task{{ID: taskID, Status: consts.TaskStatusProcessing}}
+	redisClient := newTestRedis(t)
+	logger := slog.Default()
+	repo := &refreshHostRepoStub{vm: vm}
+	r := &vmIdleRefresher{
+		cfg:             &config.Config{VMIdle: config.VMIdle{SleepSeconds: 600, RecycleSeconds: 3600}},
+		redis:           redisClient,
+		logger:          logger,
+		hostRepo:        repo,
+		taskLogActivity: &taskLogActivityStub{latest: map[uuid.UUID]time.Time{taskID: now.Add(-time.Minute)}},
+		sleepQueue:      delayqueue.NewVMSleepQueue(redisClient, logger),
+		notifyQueue:     delayqueue.NewVMNotifyQueue(redisClient, logger),
+		recycleQueue:    delayqueue.NewVMRecycleQueue(redisClient, logger),
+	}
+
+	skip, err := r.shouldSkipRecycleForRecentActivity(ctx, vm, now)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !skip {
+		t.Fatal("expected recent task log to skip recycle")
+	}
+	if repo.getVirtualMachineCalls != 1 {
+		t.Fatalf("GetVirtualMachine calls = %d, want 1", repo.getVirtualMachineCalls)
+	}
+}
+
+func TestShouldSkipRecycleForRecentTaskLogAllowsStaleVM(t *testing.T) {
+	ctx := context.Background()
+	now := time.Date(2026, 7, 7, 10, 34, 38, 0, time.UTC)
+	taskID := uuid.New()
+	vm := &db.VirtualMachine{ID: "vm-stale-log", UserID: uuid.New()}
+	vm.Edges.Tasks = []*db.Task{{ID: taskID, Status: consts.TaskStatusProcessing}}
+	repo := &refreshHostRepoStub{vm: vm}
+	r := &vmIdleRefresher{
+		cfg:             &config.Config{VMIdle: config.VMIdle{SleepSeconds: 600, RecycleSeconds: 3600}},
+		logger:          slog.Default(),
+		hostRepo:        repo,
+		taskLogActivity: &taskLogActivityStub{latest: map[uuid.UUID]time.Time{taskID: now.Add(-2 * time.Hour)}},
+	}
+
+	skip, err := r.shouldSkipRecycleForRecentActivity(ctx, vm, now)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if skip {
+		t.Fatal("expected stale task log to allow recycle")
+	}
+	if repo.getVirtualMachineCalls != 0 {
+		t.Fatalf("GetVirtualMachine calls = %d, want 0", repo.getVirtualMachineCalls)
+	}
+}
+
 func newTestRedis(t *testing.T) *redis.Client {
 	t.Helper()
 	srv := miniredis.RunT(t)
@@ -139,6 +239,21 @@ type refreshHostRepoStub struct {
 	getVirtualMachineCalls int
 }
 
+type taskLogActivityStub struct {
+	latest map[uuid.UUID]time.Time
+	err    error
+	calls  int
+}
+
+func (s *taskLogActivityStub) LatestTaskLogTime(_ context.Context, taskID uuid.UUID) (time.Time, bool, error) {
+	s.calls++
+	if s.err != nil {
+		return time.Time{}, false, s.err
+	}
+	latest, ok := s.latest[taskID]
+	return latest, ok, nil
+}
+
 func (s *refreshHostRepoStub) List(context.Context, uuid.UUID) ([]*db.Host, error) {
 	return nil, errors.New("not implemented")
 }