merge: sync upstream changes

backend/internal/repository/dashboard_aggregation_repo.go

@@ -8,6 +8,7 @@ import (
 	"strings"
 	"time"
 
+	"github.com/Wei-Shaw/sub2api/internal/pkg/timezone"
 	"github.com/Wei-Shaw/sub2api/internal/service"
 	"github.com/lib/pq"
 )
@@ -41,21 +42,22 @@ func isPostgresDriver(db *sql.DB) bool {
 }
 
 func (r *dashboardAggregationRepository) AggregateRange(ctx context.Context, start, end time.Time) error {
-	startUTC := start.UTC()
-	endUTC := end.UTC()
-	if !endUTC.After(startUTC) {
+	loc := timezone.Location()
+	startLocal := start.In(loc)
+	endLocal := end.In(loc)
+	if !endLocal.After(startLocal) {
 		return nil
 	}
 
-	hourStart := startUTC.Truncate(time.Hour)
-	hourEnd := endUTC.Truncate(time.Hour)
-	if endUTC.After(hourEnd) {
+	hourStart := startLocal.Truncate(time.Hour)
+	hourEnd := endLocal.Truncate(time.Hour)
+	if endLocal.After(hourEnd) {
 		hourEnd = hourEnd.Add(time.Hour)
 	}
 
-	dayStart := truncateToDayUTC(startUTC)
-	dayEnd := truncateToDayUTC(endUTC)
-	if endUTC.After(dayEnd) {
+	dayStart := truncateToDay(startLocal)
+	dayEnd := truncateToDay(endLocal)
+	if endLocal.After(dayEnd) {
 		dayEnd = dayEnd.Add(24 * time.Hour)
 	}
 
@@ -146,38 +148,41 @@ func (r *dashboardAggregationRepository) EnsureUsageLogsPartitions(ctx context.C
 }
 
 func (r *dashboardAggregationRepository) insertHourlyActiveUsers(ctx context.Context, start, end time.Time) error {
+	tzName := timezone.Name()
 	query := `
 		INSERT INTO usage_dashboard_hourly_users (bucket_start, user_id)
 		SELECT DISTINCT
-			date_trunc('hour', created_at AT TIME ZONE 'UTC') AT TIME ZONE 'UTC' AS bucket_start,
+			date_trunc('hour', created_at AT TIME ZONE $3) AT TIME ZONE $3 AS bucket_start,
 			user_id
 		FROM usage_logs
 		WHERE created_at >= $1 AND created_at < $2
 		ON CONFLICT DO NOTHING
 	`
-	_, err := r.sql.ExecContext(ctx, query, start.UTC(), end.UTC())
+	_, err := r.sql.ExecContext(ctx, query, start, end, tzName)
 	return err
 }
 
 func (r *dashboardAggregationRepository) insertDailyActiveUsers(ctx context.Context, start, end time.Time) error {
+	tzName := timezone.Name()
 	query := `
 		INSERT INTO usage_dashboard_daily_users (bucket_date, user_id)
 		SELECT DISTINCT
-			(bucket_start AT TIME ZONE 'UTC')::date AS bucket_date,
+			(bucket_start AT TIME ZONE $3)::date AS bucket_date,
 			user_id
 		FROM usage_dashboard_hourly_users
 		WHERE bucket_start >= $1 AND bucket_start < $2
 		ON CONFLICT DO NOTHING
 	`
-	_, err := r.sql.ExecContext(ctx, query, start.UTC(), end.UTC())
+	_, err := r.sql.ExecContext(ctx, query, start, end, tzName)
 	return err
 }
 
 func (r *dashboardAggregationRepository) upsertHourlyAggregates(ctx context.Context, start, end time.Time) error {
+	tzName := timezone.Name()
 	query := `
 		WITH hourly AS (
 			SELECT
-				date_trunc('hour', created_at AT TIME ZONE 'UTC') AT TIME ZONE 'UTC' AS bucket_start,
+				date_trunc('hour', created_at AT TIME ZONE $3) AT TIME ZONE $3 AS bucket_start,
 				COUNT(*) AS total_requests,
 				COALESCE(SUM(input_tokens), 0) AS input_tokens,
 				COALESCE(SUM(output_tokens), 0) AS output_tokens,
@@ -236,15 +241,16 @@ func (r *dashboardAggregationRepository) upsertHourlyAggregates(ctx context.Cont
 			active_users = EXCLUDED.active_users,
 			computed_at = EXCLUDED.computed_at
 	`
-	_, err := r.sql.ExecContext(ctx, query, start.UTC(), end.UTC())
+	_, err := r.sql.ExecContext(ctx, query, start, end, tzName)
 	return err
 }
 
 func (r *dashboardAggregationRepository) upsertDailyAggregates(ctx context.Context, start, end time.Time) error {
+	tzName := timezone.Name()
 	query := `
 		WITH daily AS (
 			SELECT
-				(bucket_start AT TIME ZONE 'UTC')::date AS bucket_date,
+				(bucket_start AT TIME ZONE $5)::date AS bucket_date,
 				COALESCE(SUM(total_requests), 0) AS total_requests,
 				COALESCE(SUM(input_tokens), 0) AS input_tokens,
 				COALESCE(SUM(output_tokens), 0) AS output_tokens,
@@ -255,7 +261,7 @@ func (r *dashboardAggregationRepository) upsertDailyAggregates(ctx context.Conte
 				COALESCE(SUM(total_duration_ms), 0) AS total_duration_ms
 			FROM usage_dashboard_hourly
 			WHERE bucket_start >= $1 AND bucket_start < $2
-			GROUP BY (bucket_start AT TIME ZONE 'UTC')::date
+			GROUP BY (bucket_start AT TIME ZONE $5)::date
 		),
 		user_counts AS (
 			SELECT bucket_date, COUNT(*) AS active_users
@@ -303,7 +309,7 @@ func (r *dashboardAggregationRepository) upsertDailyAggregates(ctx context.Conte
 			active_users = EXCLUDED.active_users,
 			computed_at = EXCLUDED.computed_at
 	`
-	_, err := r.sql.ExecContext(ctx, query, start.UTC(), end.UTC(), start.UTC(), end.UTC())
+	_, err := r.sql.ExecContext(ctx, query, start, end, start, end, tzName)
 	return err
 }
 
@@ -376,9 +382,8 @@ func (r *dashboardAggregationRepository) createUsageLogsPartition(ctx context.Co
 	return err
 }
 
-func truncateToDayUTC(t time.Time) time.Time {
-	t = t.UTC()
-	return time.Date(t.Year(), t.Month(), t.Day(), 0, 0, 0, 0, time.UTC)
+func truncateToDay(t time.Time) time.Time {
+	return timezone.StartOfDay(t)
 }
 
 func truncateToMonthUTC(t time.Time) time.Time {

backend/internal/repository/gemini_token_cache.go

@@ -11,8 +11,8 @@ import (
 )
 
 const (
-	geminiTokenKeyPrefix       = "gemini:token:"
-	geminiRefreshLockKeyPrefix = "gemini:refresh_lock:"
+	oauthTokenKeyPrefix       = "oauth:token:"
+	oauthRefreshLockKeyPrefix = "oauth:refresh_lock:"
 )
 
 type geminiTokenCache struct {
@@ -24,21 +24,26 @@ func NewGeminiTokenCache(rdb *redis.Client) service.GeminiTokenCache {
 }
 
 func (c *geminiTokenCache) GetAccessToken(ctx context.Context, cacheKey string) (string, error) {
-	key := fmt.Sprintf("%s%s", geminiTokenKeyPrefix, cacheKey)
+	key := fmt.Sprintf("%s%s", oauthTokenKeyPrefix, cacheKey)
 	return c.rdb.Get(ctx, key).Result()
 }
 
 func (c *geminiTokenCache) SetAccessToken(ctx context.Context, cacheKey string, token string, ttl time.Duration) error {
-	key := fmt.Sprintf("%s%s", geminiTokenKeyPrefix, cacheKey)
+	key := fmt.Sprintf("%s%s", oauthTokenKeyPrefix, cacheKey)
 	return c.rdb.Set(ctx, key, token, ttl).Err()
 }
 
+func (c *geminiTokenCache) DeleteAccessToken(ctx context.Context, cacheKey string) error {
+	key := fmt.Sprintf("%s%s", oauthTokenKeyPrefix, cacheKey)
+	return c.rdb.Del(ctx, key).Err()
+}
+
 func (c *geminiTokenCache) AcquireRefreshLock(ctx context.Context, cacheKey string, ttl time.Duration) (bool, error) {
-	key := fmt.Sprintf("%s%s", geminiRefreshLockKeyPrefix, cacheKey)
+	key := fmt.Sprintf("%s%s", oauthRefreshLockKeyPrefix, cacheKey)
 	return c.rdb.SetNX(ctx, key, 1, ttl).Result()
 }
 
 func (c *geminiTokenCache) ReleaseRefreshLock(ctx context.Context, cacheKey string) error {
-	key := fmt.Sprintf("%s%s", geminiRefreshLockKeyPrefix, cacheKey)
+	key := fmt.Sprintf("%s%s", oauthRefreshLockKeyPrefix, cacheKey)
 	return c.rdb.Del(ctx, key).Err()
 }

backend/internal/repository/gemini_token_cache_integration_test.go

+//go:build integration
+
+package repository
+
+import (
+	"errors"
+	"testing"
+	"time"
+
+	"github.com/Wei-Shaw/sub2api/internal/service"
+	"github.com/redis/go-redis/v9"
+	"github.com/stretchr/testify/require"
+	"github.com/stretchr/testify/suite"
+)
+
+type GeminiTokenCacheSuite struct {
+	IntegrationRedisSuite
+	cache service.GeminiTokenCache
+}
+
+func (s *GeminiTokenCacheSuite) SetupTest() {
+	s.IntegrationRedisSuite.SetupTest()
+	s.cache = NewGeminiTokenCache(s.rdb)
+}
+
+func (s *GeminiTokenCacheSuite) TestDeleteAccessToken() {
+	cacheKey := "project-123"
+	token := "token-value"
+	require.NoError(s.T(), s.cache.SetAccessToken(s.ctx, cacheKey, token, time.Minute))
+
+	got, err := s.cache.GetAccessToken(s.ctx, cacheKey)
+	require.NoError(s.T(), err)
+	require.Equal(s.T(), token, got)
+
+	require.NoError(s.T(), s.cache.DeleteAccessToken(s.ctx, cacheKey))
+
+	_, err = s.cache.GetAccessToken(s.ctx, cacheKey)
+	require.True(s.T(), errors.Is(err, redis.Nil), "expected redis.Nil after delete")
+}
+
+func (s *GeminiTokenCacheSuite) TestDeleteAccessToken_MissingKey() {
+	require.NoError(s.T(), s.cache.DeleteAccessToken(s.ctx, "missing-key"))
+}
+
+func TestGeminiTokenCacheSuite(t *testing.T) {
+	suite.Run(t, new(GeminiTokenCacheSuite))
+}

backend/internal/repository/gemini_token_cache_test.go

+//go:build unit
+
+package repository
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/redis/go-redis/v9"
+	"github.com/stretchr/testify/require"
+)
+
+func TestGeminiTokenCache_DeleteAccessToken_RedisError(t *testing.T) {
+	rdb := redis.NewClient(&redis.Options{
+		Addr:         "127.0.0.1:1",
+		DialTimeout:  50 * time.Millisecond,
+		ReadTimeout:  50 * time.Millisecond,
+		WriteTimeout: 50 * time.Millisecond,
+	})
+	t.Cleanup(func() {
+		_ = rdb.Close()
+	})
+
+	cache := NewGeminiTokenCache(rdb)
+	err := cache.DeleteAccessToken(context.Background(), "broken")
+	require.Error(t, err)
+}

backend/internal/repository/group_repo.go

@@ -49,7 +49,13 @@ func (r *groupRepository) Create(ctx context.Context, groupIn *service.Group) er
 		SetNillableImagePrice4k(groupIn.ImagePrice4K).
 		SetDefaultValidityDays(groupIn.DefaultValidityDays).
 		SetClaudeCodeOnly(groupIn.ClaudeCodeOnly).
-		SetNillableFallbackGroupID(groupIn.FallbackGroupID)
+		SetNillableFallbackGroupID(groupIn.FallbackGroupID).
+		SetModelRoutingEnabled(groupIn.ModelRoutingEnabled)
+
+	// 设置模型路由配置
+	if groupIn.ModelRouting != nil {
+		builder = builder.SetModelRouting(groupIn.ModelRouting)
+	}
 
 	created, err := builder.Save(ctx)
 	if err == nil {
@@ -101,7 +107,8 @@ func (r *groupRepository) Update(ctx context.Context, groupIn *service.Group) er
 		SetNillableImagePrice2k(groupIn.ImagePrice2K).
 		SetNillableImagePrice4k(groupIn.ImagePrice4K).
 		SetDefaultValidityDays(groupIn.DefaultValidityDays).
-		SetClaudeCodeOnly(groupIn.ClaudeCodeOnly)
+		SetClaudeCodeOnly(groupIn.ClaudeCodeOnly).
+		SetModelRoutingEnabled(groupIn.ModelRoutingEnabled)
 
 	// 处理 FallbackGroupID：nil 时清除，否则设置
 	if groupIn.FallbackGroupID != nil {
@@ -110,6 +117,13 @@ func (r *groupRepository) Update(ctx context.Context, groupIn *service.Group) er
 		builder = builder.ClearFallbackGroupID()
 	}
 
+	// 处理 ModelRouting：nil 时清除，否则设置
+	if groupIn.ModelRouting != nil {
+		builder = builder.SetModelRouting(groupIn.ModelRouting)
+	} else {
+		builder = builder.ClearModelRouting()
+	}
+
 	updated, err := builder.Save(ctx)
 	if err != nil {
 		return translatePersistenceError(err, service.ErrGroupNotFound, service.ErrGroupExists)

backend/internal/repository/proxy_probe_service_test.go

@@ -21,7 +21,7 @@ type ProxyProbeServiceSuite struct {
 func (s *ProxyProbeServiceSuite) SetupTest() {
 	s.ctx = context.Background()
 	s.prober = &proxyProbeService{
-		ipInfoURL:         "http://ipinfo.test/json",
+		ipInfoURL:         "http://ip-api.test/json/?lang=zh-CN",
 		allowPrivateHosts: true,
 	}
 }
@@ -54,7 +54,7 @@ func (s *ProxyProbeServiceSuite) TestProbeProxy_Success() {
 	s.setupProxyServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		seen <- r.RequestURI
 		w.Header().Set("Content-Type", "application/json")
-		_, _ = io.WriteString(w, `{"ip":"1.2.3.4","city":"c","region":"r","country":"cc"}`)
+		_, _ = io.WriteString(w, `{"status":"success","query":"1.2.3.4","city":"c","regionName":"r","country":"cc","countryCode":"CC"}`)
 	}))
 
 	info, latencyMs, err := s.prober.ProbeProxy(s.ctx, s.proxySrv.URL)
@@ -64,11 +64,12 @@ func (s *ProxyProbeServiceSuite) TestProbeProxy_Success() {
 	require.Equal(s.T(), "c", info.City)
 	require.Equal(s.T(), "r", info.Region)
 	require.Equal(s.T(), "cc", info.Country)
+	require.Equal(s.T(), "CC", info.CountryCode)
 
 	// Verify proxy received the request
 	select {
 	case uri := <-seen:
-		require.Contains(s.T(), uri, "ipinfo.test", "expected request to go through proxy")
+		require.Contains(s.T(), uri, "ip-api.test", "expected request to go through proxy")
 	default:
 		require.Fail(s.T(), "expected proxy to receive request")
 	}