feat: 完成token功能

Yi.Abp.Net8/module/ai-hub/Yi.Framework.AiHub.Domain/Managers/TokenManager.cs

@@ -1,64 +1,134 @@
-using Volo.Abp.Domain.Services;
-using Volo.Abp.Users;
+using SqlSugar;
+using Volo.Abp.Domain.Services;
+using Yi.Framework.AiHub.Domain.Entities;
 using Yi.Framework.AiHub.Domain.Entities.OpenApi;
+using Yi.Framework.AiHub.Domain.Shared.Consts;
 using Yi.Framework.SqlSugarCore.Abstractions;
 
 namespace Yi.Framework.AiHub.Domain.Managers;
 
+/// <summary>
+/// Token验证结果
+/// </summary>
+public class TokenValidationResult
+{
+    /// <summary>
+    /// 用户Id
+    /// </summary>
+    public Guid UserId { get; set; }
+
+    /// <summary>
+    /// Token Id
+    /// </summary>
+    public Guid TokenId { get; set; }
+}
+
 public class TokenManager : DomainService
 {
     private readonly ISqlSugarRepository<TokenAggregateRoot> _tokenRepository;
+    private readonly ISqlSugarRepository<UsageStatisticsAggregateRoot> _usageStatisticsRepository;
 
-    public TokenManager(ISqlSugarRepository<TokenAggregateRoot> tokenRepository)
+    public TokenManager(
+        ISqlSugarRepository<TokenAggregateRoot> tokenRepository,
+        ISqlSugarRepository<UsageStatisticsAggregateRoot> usageStatisticsRepository)
     {
         _tokenRepository = tokenRepository;
+        _usageStatisticsRepository = usageStatisticsRepository;
     }
 
-    public async Task<string?> GetAsync(Guid userId)
-    {
-        var entity = await _tokenRepository._DbQueryable.FirstAsync(x => x.UserId == userId);
-        if (entity is not null)
-        {
-            return entity.Token;
-        }
-        else
-        {
-            return null;
-        }
-    }
-
-    public async Task CreateAsync(Guid userId)
-    {
-        var entity = await _tokenRepository._DbQueryable.FirstAsync(x => x.UserId == userId);
-        if (entity is not null)
-        {
-            entity.ResetToken();
-            await _tokenRepository.UpdateAsync(entity);
-        }
-        else
-        {
-            var token = new TokenAggregateRoot(userId);
-            await _tokenRepository.InsertAsync(token);
-        }
-    }
-
-    public async Task<Guid> GetUserIdAsync(string? token)
+    /// <summary>
+    /// 验证Token并返回用户Id和TokenId
+    /// </summary>
+    /// <param name="token">Token密钥</param>
+    /// <param name="modelId">模型Id（用于判断是否是尊享模型需要检查额度）</param>
+    /// <returns>Token验证结果</returns>
+    public async Task<TokenValidationResult> ValidateTokenAsync(string? token, string? modelId = null)
     {
         if (token is null)
         {
             throw new UserFriendlyException("当前请求未包含token", "401");
         }
 
-        if (token.StartsWith("yi-"))
+        if (!token.StartsWith("yi-"))
         {
-            var entity = await _tokenRepository._DbQueryable.Where(x => x.Token == token).FirstAsync();
-            if (entity is null)
-            {
-                throw new UserFriendlyException("当前请求token无效", "401");
-            }
-
-            return entity.UserId;
+            throw new UserFriendlyException("当前请求token非法", "401");
         }
-        throw new UserFriendlyException("当前请求token非法", "401");
+
+        var entity = await _tokenRepository._DbQueryable
+            .Where(x => x.Token == token)
+            .FirstAsync();
+
+        if (entity is null)
+        {
+            throw new UserFriendlyException("当前请求token无效", "401");
+        }
+
+        // 检查Token是否被禁用
+        if (entity.IsDisabled)
+        {
+            throw new UserFriendlyException("当前Token已被禁用，请启用后再使用", "403");
+        }
+
+        // 检查Token是否过期
+        if (entity.ExpireTime.HasValue && entity.ExpireTime.Value < DateTime.Now)
+        {
+            throw new UserFriendlyException("当前Token已过期，请更新过期时间或创建新的Token", "403");
+        }
+
+        // 如果是尊享模型且Token设置了额度限制，检查是否超限
+        if (!string.IsNullOrEmpty(modelId) &&
+            PremiumPackageConst.ModeIds.Contains(modelId) &&
+            entity.PremiumQuotaLimit.HasValue)
+        {
+            var usedQuota = await GetTokenPremiumUsedQuotaAsync(entity.UserId, entity.Id);
+            if (usedQuota >= entity.PremiumQuotaLimit.Value)
+            {
+                throw new UserFriendlyException($"当前Token的尊享包额度已用完（已使用：{usedQuota}，限制：{entity.PremiumQuotaLimit.Value}），请调整额度限制或使用其他Token", "403");
+            }
+        }
+
+        return new TokenValidationResult
+        {
+            UserId = entity.UserId,
+            TokenId = entity.Id
+        };
     }
-}
+
+    /// <summary>
+    /// 获取Token的尊享包已使用额度
+    /// </summary>
+    private async Task<long> GetTokenPremiumUsedQuotaAsync(Guid userId, Guid tokenId)
+    {
+        var premiumModelIds = PremiumPackageConst.ModeIds;
+
+        var usedQuota = await _usageStatisticsRepository._DbQueryable
+            .Where(x => x.UserId == userId && x.TokenId == tokenId && premiumModelIds.Contains(x.ModelId))
+            .SumAsync(x => x.TotalTokenCount);
+
+        return usedQuota;
+    }
+
+    /// <summary>
+    /// 获取用户的Token（兼容旧接口，返回第一个可用的Token）
+    /// </summary>
+    [Obsolete("请使用 ValidateTokenAsync 方法")]
+    public async Task<string?> GetAsync(Guid userId)
+    {
+        var entity = await _tokenRepository._DbQueryable
+            .Where(x => x.UserId == userId && !x.IsDisabled)
+            .OrderBy(x => x.CreationTime)
+            .FirstAsync();
+
+        return entity?.Token;
+    }
+
+    /// <summary>
+    /// 获取用户Id（兼容旧接口）
+    /// </summary>
+    [Obsolete("请使用 ValidateTokenAsync 方法")]
+    public async Task<Guid> GetUserIdAsync(string? token)
+    {
+        var result = await ValidateTokenAsync(token);
+        return result.UserId;
+    }
+}