Refactor OAuth token refresh handling to improve concurrency management

- Introduce a refresh gate mechanism to prevent duplicate refresh attempts for the same token, allowing concurrent requests to safely handle expired tokens.
- Update the guard and refresh token logic to utilize the new error handling for in-progress refresh operations, enhancing the reliability of the authentication flow.
- Remove the previous caching mechanism for refresh status, streamlining the code and improving clarity in the refresh process.

openapi/oauth/guard.go

@@ -1,9 +1,11 @@
 package oauth
 
 import (
+	"errors"
 	"fmt"
 	"net/http"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/gin-gonic/gin"
@@ -14,6 +16,16 @@ import (
 	"github.com/yaoapp/yao/openapi/response"
 )
 
+var (
+	errRefreshInProgress  = errors.New("refresh in progress")
+	errRefreshAlreadyDone = errors.New("refresh already done")
+	refreshGates          sync.Map // refreshToken → *refreshGate
+)
+
+type refreshGate struct {
+	done chan struct{} // closed when rotation completes
+}
+
 // Guard is the OAuth guard middleware
 func (s *Service) Guard(c *gin.Context) {
 	// Authenticate first (validates token and sets authorized info)
@@ -67,19 +79,17 @@ func (s *Service) Authenticate(c *gin.Context) bool {
 
 		// Signature valid but expired — attempt auto refresh
 		if !expiredClaims.ExpiresAt.IsZero() && expiredClaims.ExpiresAt.Before(time.Now()) {
-			if refreshToken := s.getRefreshToken(c); s.IsRefreshing(refreshToken) {
-				// Another request is already rotating this refresh token.
-				// The signature has been verified, so we can safely let
-				// this request through with the expired-but-authentic claims.
-				claims = expiredClaims
-			} else {
-				newClaims, refreshErr := s.TryRefreshToken(c, expiredClaims)
-				if refreshErr != nil {
+			newClaims, refreshErr := s.TryRefreshToken(c, expiredClaims)
+			if refreshErr != nil {
+				if errors.Is(refreshErr, errRefreshInProgress) || errors.Is(refreshErr, errRefreshAlreadyDone) {
+					claims = expiredClaims
+				} else {
 					log.Error("[OAuth] Token refresh failed: %v", refreshErr)
 					response.RespondWithError(c, http.StatusUnauthorized, types.ErrInvalidRefreshToken)
 					c.Abort()
 					return false
 				}
+			} else {
 				claims = newClaims
 			}
 		} else {
@@ -111,12 +121,28 @@ func (s *Service) TryRefreshToken(c *gin.Context, expiredClaims *types.TokenClai
 		return nil, fmt.Errorf("refresh token missing")
 	}
 
-	// Mark this refresh token as being rotated so concurrent requests
-	// can detect the in-flight rotation and skip duplicate refresh attempts.
-	// TTL acts as a safety net in case defer doesn't run (e.g. process crash).
-	refreshingKey := s.refreshingCacheKey(refreshToken)
-	s.cache.Set(refreshingKey, true, 30*time.Second)
-	defer s.cache.Del(refreshingKey)
+	gate := &refreshGate{done: make(chan struct{})}
+	if actual, loaded := refreshGates.LoadOrStore(refreshToken, gate); loaded {
+		// Another goroutine owns the rotation for this refresh token.
+		// It may still be running or already finished.
+		existing := actual.(*refreshGate)
+		select {
+		case <-existing.done:
+			return nil, errRefreshAlreadyDone
+		default:
+			return nil, errRefreshInProgress
+		}
+	}
+
+	// We own the gate — clean up when finished.
+	defer func() {
+		close(gate.done)
+		// Keep the gate in the map for 30 s so late arrivals see "done"
+		// instead of starting a new rotation with the now-revoked token.
+		time.AfterFunc(30*time.Second, func() {
+			refreshGates.CompareAndDelete(refreshToken, gate)
+		})
+	}()
 
 	refreshClaims, err := s.VerifyRefreshToken(refreshToken)
 	if err != nil {
@@ -258,19 +284,10 @@ func (s *Service) GetRefreshToken(c *gin.Context) string {
 	return s.getRefreshToken(c)
 }
 
-// IsRefreshing reports whether the given refresh token is currently being
-// rotated by another request. Callers can use this to avoid duplicate
-// refresh attempts that would fail because the old token has been revoked.
-func (s *Service) IsRefreshing(refreshToken string) bool {
-	if refreshToken == "" {
-		return false
-	}
-	_, ok := s.cache.Get(s.refreshingCacheKey(refreshToken))
-	return ok
-}
-
-func (s *Service) refreshingCacheKey(refreshToken string) string {
-	return fmt.Sprintf("%soauth:refreshing:%s", s.prefix, refreshToken)
+// IsRefreshInProgress checks whether an error signals that another goroutine
+// is already rotating (or has just rotated) the same refresh token.
+func IsRefreshInProgress(err error) bool {
+	return errors.Is(err, errRefreshInProgress) || errors.Is(err, errRefreshAlreadyDone)
 }
 
 // GetSessionID gets the session ID from the request (public method)

sui/api/guards.go

@@ -122,16 +122,14 @@ func guardOAuth(r *Request) error {
 		expiredClaims, expErr := oauth.OAuth.VerifyTokenAllowExpired(token)
 		if expErr == nil && expiredClaims != nil &&
 			!expiredClaims.ExpiresAt.IsZero() && expiredClaims.ExpiresAt.Before(time.Now()) {
-			if refreshToken := oauth.OAuth.GetRefreshToken(c); oauth.OAuth.IsRefreshing(refreshToken) {
-				// Another request is already rotating this refresh token.
-				// The signature has been verified, so we can safely let
-				// this request through with the expired-but-authentic claims.
-				claims = expiredClaims
-			} else {
-				refreshed, refreshErr := oauth.OAuth.TryRefreshToken(c, expiredClaims)
-				if refreshErr != nil {
+			refreshed, refreshErr := oauth.OAuth.TryRefreshToken(c, expiredClaims)
+			if refreshErr != nil {
+				if oauth.IsRefreshInProgress(refreshErr) {
+					claims = expiredClaims
+				} else {
 					return fmt.Errorf("Exception|401:Token expired and refresh failed")
 				}
+			} else {
 				claims = refreshed
 			}
 		} else {