test(auth): preserve session affinity while waiting

Author: DeliciousBuding

auth/session_affinity_test.go

@@ -1,6 +1,10 @@
 package auth
 
-import "testing"
+import (
+	"context"
+	"testing"
+	"time"
+)
 
 func TestNextForSessionPrefersBoundAccountAndProxy(t *testing.T) {
 	store := &Store{
@@ -45,3 +49,77 @@ func TestNextForSessionFallsBackWhenBoundAccountExcluded(t *testing.T) {
 		t.Fatalf("proxyURL = %q, want empty fallback proxy", proxyURL)
 	}
 }
+
+func TestWaitForSessionAvailableReturnsBoundAccount(t *testing.T) {
+	store := &Store{
+		accounts: []*Account{
+			{DBID: 1, AccessToken: "tok-1"},
+			{DBID: 2, AccessToken: "tok-2"},
+		},
+		maxConcurrency: 1,
+	}
+	store.bindSessionAffinity("session-1", store.accounts[1], "http://proxy-2")
+
+	ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
+	defer cancel()
+
+	acc, proxyURL := store.WaitForSessionAvailable(ctx, "session-1", 50*time.Millisecond, nil)
+	if acc == nil {
+		t.Fatal("expected bound account")
+	}
+	if acc.DBID != 2 {
+		t.Fatalf("account DBID = %d, want %d", acc.DBID, 2)
+	}
+	if proxyURL != "http://proxy-2" {
+		t.Fatalf("proxyURL = %q, want %q", proxyURL, "http://proxy-2")
+	}
+}
+
+func TestWaitForSessionAvailableFallsBackWhenBindingExpired(t *testing.T) {
+	store := &Store{
+		accounts: []*Account{
+			{DBID: 1, AccessToken: "tok-1"},
+		},
+		maxConcurrency:  1,
+		sessionBindings: map[string]sessionAffinity{"session-1": {accountID: 99, proxyURL: "http://stale", expiresAt: time.Now().Add(-time.Minute)}},
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
+	defer cancel()
+
+	acc, proxyURL := store.WaitForSessionAvailable(ctx, "session-1", 50*time.Millisecond, nil)
+	if acc == nil {
+		t.Fatal("expected fallback account")
+	}
+	if acc.DBID != 1 {
+		t.Fatalf("account DBID = %d, want %d", acc.DBID, 1)
+	}
+	if proxyURL != "" {
+		t.Fatalf("proxyURL = %q, want empty fallback proxy", proxyURL)
+	}
+}
+
+func TestWaitForSessionAvailableRespectsExcludeSet(t *testing.T) {
+	store := &Store{
+		accounts: []*Account{
+			{DBID: 1, AccessToken: "tok-1"},
+			{DBID: 2, AccessToken: "tok-2"},
+		},
+		maxConcurrency: 1,
+	}
+	store.bindSessionAffinity("session-1", store.accounts[1], "http://proxy-2")
+
+	ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
+	defer cancel()
+
+	acc, proxyURL := store.WaitForSessionAvailable(ctx, "session-1", 50*time.Millisecond, map[int64]bool{2: true})
+	if acc == nil {
+		t.Fatal("expected fallback account")
+	}
+	if acc.DBID != 1 {
+		t.Fatalf("account DBID = %d, want %d", acc.DBID, 1)
+	}
+	if proxyURL != "" {
+		t.Fatalf("proxyURL = %q, want empty fallback proxy", proxyURL)
+	}
+}

auth/store.go

@@ -1378,6 +1378,12 @@ func (s *Store) takeByIDExcluding(id int64, exclude map[int64]bool) *Account {
 
 // WaitForAvailable 等待可用账号（带超时的请求排队）
 func (s *Store) WaitForAvailable(ctx context.Context, timeout time.Duration) *Account {
+	acc, _ := s.WaitForSessionAvailable(ctx, "", timeout, nil)
+	return acc
+}
+
+// WaitForSessionAvailable waits for a session-preferred account and proxy pair.
+func (s *Store) WaitForSessionAvailable(ctx context.Context, key string, timeout time.Duration, exclude map[int64]bool) (*Account, string) {
 	deadline := time.NewTimer(timeout)
 	defer deadline.Stop()
 
@@ -1388,13 +1394,13 @@ func (s *Store) WaitForAvailable(ctx context.Context, timeout time.Duration) *Ac
 	for {
 		select {
 		case <-ctx.Done():
-			return nil
+			return nil, ""
 		case <-deadline.C:
-			return nil
+			return nil, ""
 		default:
-			acc := s.Next()
+			acc, proxyURL := s.NextForSession(key, exclude)
 			if acc != nil {
-				return acc
+				return acc, proxyURL
 			}
 			// 等待一下再重试（指数退避，最大 500ms）
 			backoffTimer.Reset(backoff)
@@ -1404,9 +1410,9 @@ func (s *Store) WaitForAvailable(ctx context.Context, timeout time.Duration) *Ac
 					backoff *= 2
 				}
 			case <-ctx.Done():
-				return nil
+				return nil, ""
 			case <-deadline.C:
-				return nil
+				return nil, ""
 			}
 		}
 	}

proxy/handler.go

@@ -454,7 +454,7 @@ func (h *Handler) Responses(c *gin.Context) {
 		account, stickyProxyURL := h.nextAccountForSession(sessionID, excludeAccounts)
 		if account == nil {
 			// 排队等待可用账号（最多 30s）
-			account = h.store.WaitForAvailable(c.Request.Context(), 30*time.Second)
+			account, stickyProxyURL = h.store.WaitForSessionAvailable(c.Request.Context(), sessionID, 30*time.Second, excludeAccounts)
 			if account == nil {
 				if lastStatusCode == http.StatusTooManyRequests && len(lastBody) > 0 {
 					h.sendFinalUpstreamError(c, lastStatusCode, lastBody)
@@ -823,7 +823,7 @@ func (h *Handler) ChatCompletions(c *gin.Context) {
 		account, stickyProxyURL := h.nextAccountForSession(sessionID, excludeAccounts)
 		if account == nil {
 			// 排队等待可用账号（最多 30s）
-			account = h.store.WaitForAvailable(c.Request.Context(), 30*time.Second)
+			account, stickyProxyURL = h.store.WaitForSessionAvailable(c.Request.Context(), sessionID, 30*time.Second, excludeAccounts)
 			if account == nil {
 				if lastStatusCode == http.StatusTooManyRequests && len(lastBody) > 0 {
 					h.sendFinalUpstreamError(c, lastStatusCode, lastBody)

proxy/handler_anthropic.go

@@ -121,7 +121,7 @@ func (h *Handler) Messages(c *gin.Context) {
 	for attempt := 0; attempt <= maxRetries; attempt++ {
 		account, stickyProxyURL := h.nextAccountForSession(sessionID, excludeAccounts)
 		if account == nil {
-			account = h.store.WaitForAvailable(c.Request.Context(), 30*time.Second)
+			account, stickyProxyURL = h.store.WaitForSessionAvailable(c.Request.Context(), sessionID, 30*time.Second, excludeAccounts)
 			if account == nil {
 				if lastStatusCode == http.StatusTooManyRequests && len(lastBody) > 0 {
 					sendAnthropicError(c, http.StatusTooManyRequests, "rate_limit_error", "All accounts rate limited")