From f970bd48c98775eca3c4d538b1bfe8c001ae0c14 Mon Sep 17 00:00:00 2001
From: rick147 <w.shao.fei@gmail.com>
Date: Sun, 2 Aug 2026 15:46:50 +0800
Subject: [PATCH] fix: stop scheduler work after request cancellation

---
 .../scheduler_snapshot_cancellation_test.go   | 79 +++++++++++++++++++
 .../service/scheduler_snapshot_service.go     | 18 +++++
 2 files changed, 97 insertions(+)
 create mode 100644 backend/internal/service/scheduler_snapshot_cancellation_test.go

diff --git a/backend/internal/service/scheduler_snapshot_cancellation_test.go b/backend/internal/service/scheduler_snapshot_cancellation_test.go
new file mode 100644
index 0000000000..3f46cb14aa
--- /dev/null
+++ b/backend/internal/service/scheduler_snapshot_cancellation_test.go
@@ -0,0 +1,79 @@
+//go:build unit
+
+package service
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+type schedulerCancellationCache struct {
+	SchedulerCache
+	cancel        context.CancelFunc
+	tokenCaptures int
+}
+
+func (c *schedulerCancellationCache) GetSnapshot(ctx context.Context, _ SchedulerBucket) ([]*Account, bool, error) {
+	c.cancel()
+	return nil, false, ctx.Err()
+}
+
+func (c *schedulerCancellationCache) CaptureBucketWriteToken(ctx context.Context, _ SchedulerBucket) (SchedulerBucketWriteToken, error) {
+	c.tokenCaptures++
+	return SchedulerBucketWriteToken{}, ctx.Err()
+}
+
+func (c *schedulerCancellationCache) GetAccount(ctx context.Context, _ int64) (*Account, error) {
+	c.cancel()
+	return nil, ctx.Err()
+}
+
+type schedulerCancellationAccountRepo struct {
+	AccountRepository
+	listCalls    int
+	getByIDCalls int
+}
+
+func (r *schedulerCancellationAccountRepo) ListSchedulableUngroupedByPlatform(ctx context.Context, _ string) ([]Account, error) {
+	r.listCalls++
+	return nil, ctx.Err()
+}
+
+func (r *schedulerCancellationAccountRepo) GetByID(ctx context.Context, _ int64) (*Account, error) {
+	r.getByIDCalls++
+	return nil, ctx.Err()
+}
+
+func TestSchedulerSnapshotListStopsAfterRequestCancellation(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	t.Cleanup(cancel)
+
+	cache := &schedulerCancellationCache{cancel: cancel}
+	repo := &schedulerCancellationAccountRepo{}
+	svc := NewSchedulerSnapshotService(cache, nil, repo, nil, nil)
+
+	accounts, useMixed, err := svc.ListSchedulableAccounts(ctx, nil, PlatformOpenAI, false)
+
+	require.ErrorIs(t, err, context.Canceled)
+	require.Nil(t, accounts)
+	require.False(t, useMixed)
+	require.Zero(t, cache.tokenCaptures, "canceled requests must not capture a cache publish token")
+	require.Zero(t, repo.listCalls, "canceled requests must not fall back to the database")
+}
+
+func TestSchedulerSnapshotGetAccountStopsAfterRequestCancellation(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	t.Cleanup(cancel)
+
+	cache := &schedulerCancellationCache{cancel: cancel}
+	repo := &schedulerCancellationAccountRepo{}
+	svc := NewSchedulerSnapshotService(cache, nil, repo, nil, nil)
+
+	account, err := svc.GetAccount(ctx, 42)
+
+	require.ErrorIs(t, err, context.Canceled)
+	require.Nil(t, account)
+	require.Zero(t, repo.getByIDCalls, "canceled requests must not fall back to the database")
+}
diff --git a/backend/internal/service/scheduler_snapshot_service.go b/backend/internal/service/scheduler_snapshot_service.go
index 542fc18b97..9655f54f25 100644
--- a/backend/internal/service/scheduler_snapshot_service.go
+++ b/backend/internal/service/scheduler_snapshot_service.go
@@ -213,15 +213,24 @@ func (s *SchedulerSnapshotService) ListSchedulableAccounts(ctx context.Context,
 	bucket := s.bucketFor(groupID, platform, mode)
 	var writeToken SchedulerBucketWriteToken
 	canPublish := false
+	if err := ctx.Err(); err != nil {
+		return nil, useMixed, err
+	}
 
 	if s.cache != nil {
 		cached, hit, err := s.cache.GetSnapshot(ctx, bucket)
+		if ctxErr := ctx.Err(); ctxErr != nil {
+			return nil, useMixed, ctxErr
+		}
 		if err != nil {
 			logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] cache read failed: bucket=%s err=%v", bucket.String(), err)
 		} else if hit {
 			return derefAccounts(cached), useMixed, nil
 		}
 		token, err := s.cache.CaptureBucketWriteToken(ctx, bucket)
+		if ctxErr := ctx.Err(); ctxErr != nil {
+			return nil, useMixed, ctxErr
+		}
 		if err != nil {
 			if errors.Is(err, ErrSchedulerBucketRetired) || errors.Is(err, ErrSchedulerBucketWriteFenced) {
 				slog.Debug("[Scheduler] cache publish fenced", "bucket", bucket.String())
@@ -245,6 +254,9 @@ func (s *SchedulerSnapshotService) ListSchedulableAccounts(ctx context.Context,
 	if err != nil {
 		return nil, useMixed, err
 	}
+	if ctxErr := ctx.Err(); ctxErr != nil {
+		return nil, useMixed, ctxErr
+	}
 
 	if s.cache != nil && canPublish {
 		if err := s.cache.SetSnapshot(fallbackCtx, bucket, writeToken, accounts); err != nil {
@@ -263,8 +275,14 @@ func (s *SchedulerSnapshotService) GetAccount(ctx context.Context, accountID int
 	if accountID <= 0 {
 		return nil, nil
 	}
+	if err := ctx.Err(); err != nil {
+		return nil, err
+	}
 	if s.cache != nil {
 		account, err := s.cache.GetAccount(ctx, accountID)
+		if ctxErr := ctx.Err(); ctxErr != nil {
+			return nil, ctxErr
+		}
 		if err != nil {
 			logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] account cache read failed: id=%d err=%v", accountID, err)
 		} else if account != nil {