fix(api): prune the stale team-index entry when the expiration sweep finds an orphan (rebase of #3567)

packages/api/internal/sandbox/storage/redis/expiration_index_test.go

@@ -170,17 +170,60 @@ func TestExpiredItems_RemovesOrphanMember(t *testing.T) {
 
 	storage, client := setupTestStorage(t)
 
-	member := expirationMember(uuid.NewString(), "sbx-orphan", uuid.NewString())
+	teamID := uuid.New()
+	const sandboxID = "sbx-orphan"
+	member := expirationMember(teamID.String(), sandboxID, uuid.NewString())
 	require.NoError(t, client.ZAdd(t.Context(), globalExpirationSet, redis.Z{
 		Score:  float64(time.Now().Add(-time.Minute).UnixMilli()),
 		Member: member,
 	}).Err())
+	// A leftover team-index entry with no sandbox key — the stale state this
+	// sweep must also prune (the sandbox key is absent, so MGET returns nil).
+	teamIndexKey := GetSandboxStorageTeamIndexKey(teamID.String())
+	require.NoError(t, client.SAdd(t.Context(), teamIndexKey, sandboxID).Err())
 
 	items, err := storage.ExpiredItems(t.Context())
 	require.NoError(t, err)
 	require.Empty(t, items)
 
 	requireMemberAbsent(t, client, member)
+
+	isMember, err := client.SIsMember(t.Context(), teamIndexKey, sandboxID).Result()
+	require.NoError(t, err)
+	require.False(t, isMember, "orphan sandbox ID must be pruned from the team index")
+}
+
+// A sandbox whose key is present must never be unindexed by the orphan sweep,
+// even if a stale expired ZSET member names a different (dead) execution of it:
+// the EXISTS-guarded prune only fires when the key is truly gone, so a resume
+// that re-created the key under the same (reused) sandbox ID keeps its index.
+func TestExpiredItems_OrphanPruneSparesLiveKey(t *testing.T) {
+	t.Parallel()
+
+	storage, client := setupTestStorage(t)
+
+	teamID := uuid.New()
+	const sandboxID = "sbx-live-again"
+
+	// Live sandbox: real key + team-index entry (as Add writes them).
+	live := makeIndexedSandbox(teamID, sandboxID, uuid.NewString(), time.Now(), time.Now().Add(time.Hour))
+	require.NoError(t, storage.Add(t.Context(), live))
+
+	// A dead execution of the SAME sandbox ID left an expired orphan-looking member.
+	deadMember := expirationMember(teamID.String(), sandboxID, uuid.NewString())
+	require.NoError(t, client.ZAdd(t.Context(), globalExpirationSet, redis.Z{
+		Score:  float64(time.Now().Add(-time.Minute).UnixMilli()),
+		Member: deadMember,
+	}).Err())
+
+	_, err := storage.ExpiredItems(t.Context())
+	require.NoError(t, err)
+
+	// The dead member is swept, but the live sandbox keeps its team-index entry.
+	requireMemberAbsent(t, client, deadMember)
+	isMember, err := client.SIsMember(t.Context(), GetSandboxStorageTeamIndexKey(teamID.String()), sandboxID).Result()
+	require.NoError(t, err)
+	require.True(t, isMember, "a live sandbox key must not be unindexed by the orphan sweep")
 }
 
 // TestExpiredItems_SweepsInvalidMember: unparseable members must be removed, not skipped,

packages/api/internal/sandbox/storage/redis/items.go

@@ -3,6 +3,7 @@ package redis
 import (
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"time"
 
@@ -107,13 +108,22 @@ func (s *Storage) ExpiredItems(ctx context.Context) ([]sandboxtypes.Sandbox, err
 		for i, raw := range batch.cmd.Val() {
 			ref := batch.refs[i]
 
-			// Sandbox key gone but ZSET entry remains — orphaned. Members are
-			// execution-scoped, so this removal can never unindex a fresh
-			// execution concurrently re-added under the same sandbox ID.
+			// Sandbox key gone but ZSET entry remains — orphaned. Remove the
+			// stale ZSET member, and also prune the leftover team-index entry so
+			// it does not accumulate indefinitely (the healer never repairs the
+			// team index). The prune is an atomic EXISTS-guarded SREM so a resume
+			// that re-creates the key under the same sandbox ID between the MGET
+			// above and here is never unindexed.
 			if raw == nil {
 				staleMembers = append(staleMembers, ref.member)
 				orphanCount++
 
+				teamIndexKey := GetSandboxStorageTeamIndexKey(batch.teamID)
+				sandboxKey := getSandboxKey(batch.teamID, ref.sandboxID)
+				if err := pruneOrphanIndexScript.Run(ctx, s.redisClient, []string{sandboxKey, teamIndexKey}, ref.sandboxID).Err(); err != nil && !errors.Is(err, redis.Nil) {
+					logger.L().Warn(ctx, "Failed to prune orphan from team index", zap.Error(err), logger.WithSandboxID(ref.sandboxID))
+				}
+
 				continue
 			}
 

packages/api/internal/sandbox/storage/redis/scripts.go

@@ -82,4 +82,18 @@ var (
 		redis.call('SET', KEYS[3], '', 'EX', ARGV[4])
 		return 1
 	`)
+
+	// pruneOrphanIndexScript removes a sandbox ID from the team index ONLY if the
+	// sandbox key is still absent, atomically. The orphan sweep found the key gone
+	// via MGET, but between that read and this call a resume can re-create the key
+	// under the same (reused) sandbox ID and re-add it to the index; the EXISTS
+	// guard makes sure we never unindex such a live sandbox. The team index is not
+	// repaired by the healer, so a wrong SREM here would be permanent.
+	// KEYS[1] = sandbox key, KEYS[2] = team index key, ARGV[1] = sandbox ID.
+	pruneOrphanIndexScript = redis.NewScript(`
+		if redis.call('EXISTS', KEYS[1]) == 0 then
+			return redis.call('SREM', KEYS[2], ARGV[1])
+		end
+		return 0
+	`)
 )