ORB-91: Reserve AI quota atomically

architecture.json

@@ -2884,7 +2884,7 @@
       "Challenges": 14,
       "Chat": 59,
       "ChecklistTemplates": 9,
-      "Common": 161,
+      "Common": 163,
       "Gamification": 33,
       "Goals": 38,
       "Habits": 81,
@@ -4177,8 +4177,7 @@
         "file": "tests/Orbit.Application.Tests/Commands/Habits/SuggestHabitSetupCommandHandlerTests.cs",
         "references": [
           "SuggestHabitSetupCommand",
-          "SuggestHabitSetupCommandHandler",
-          "User"
+          "SuggestHabitSetupCommandHandler"
         ]
       },
       {
@@ -4735,6 +4734,14 @@
         "file": "tests/Orbit.Application.Tests/Common/LocaleHelperTests.cs",
         "references": []
       },
+      {
+        "testClass": "PayGateServiceExpiredCycleTests",
+        "file": "tests/Orbit.Application.Tests/Common/PayGateServiceExpiredCycleTests.cs",
+        "references": [
+          "Habit",
+          "User"
+        ]
+      },
       {
         "testClass": "PayGateServiceTests",
         "file": "tests/Orbit.Application.Tests/Common/PayGateServiceTests.cs",
@@ -5304,8 +5311,7 @@
         "references": [
           "SuggestTagsQuery",
           "SuggestTagsQueryHandler",
-          "Tag",
-          "User"
+          "Tag"
         ]
       },
       {

src/Orbit.Application/Chat/Commands/ProcessUserChatCommand.Ai.cs

@@ -66,6 +66,13 @@ private async Task<Result<AiResponse>> RequestInitialAiResponseAsync(
 
         LogCallingAiIntentService(logger, toolDeclarations.Count);
 
+        var reservation = await execution.PayGateService.TryConsumeAiMessage(
+            request.UserId,
+            execution.UnitOfWork,
+            cancellationToken);
+        if (reservation.IsFailure)
+            return reservation.PropagateError<AiResponse>();
+
         return await ai.IntentService.SendWithToolsAsync(
             new AiToolRequest(
                 request.Message,

src/Orbit.Application/Chat/Commands/ProcessUserChatCommand.Context.cs

@@ -1,6 +1,5 @@
 using Microsoft.EntityFrameworkCore;
 using Microsoft.Extensions.Logging;
-using Orbit.Application.Common;
 using Orbit.Domain.Common;
 using Orbit.Domain.Entities;
 using Orbit.Domain.Enums;
@@ -43,10 +42,6 @@ private async Task<Result<ChatContext>> LoadChatContextAsync(
             activeGoals = loadedGoals;
         }
 
-        var messageGate = await execution.PayGateService.CanSendAiMessage(request.UserId, cancellationToken);
-        if (messageGate.IsFailure)
-            return messageGate.PropagateError<ChatContext>();
-
         IReadOnlyList<UserFact> userFacts = [];
         if (aiMemoryEnabled)
         {

src/Orbit.Application/Chat/Commands/ProcessUserChatCommand.Persistence.cs

@@ -1,7 +1,6 @@
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Logging;
 using Orbit.Application.Common;
-using Orbit.Domain.Common;
 using Orbit.Domain.Entities;
 using Orbit.Domain.Enums;
 using Orbit.Domain.Interfaces;
@@ -49,7 +48,7 @@ private static bool RequiresStreakRecalculation(IEnumerable<ActionResult> action
     }
 
     /// <summary>
-    /// Fires off background work for fact extraction and AI message counter increment.
+    /// Fires off background work for fact extraction.
     /// Runs in a separate DI scope so it doesn't block the response.
     /// </summary>
     private void RunBackgroundPostResponseWork(
@@ -59,19 +58,15 @@ private void RunBackgroundPostResponseWork(
         bool shouldExtractFacts,
         IReadOnlyList<UserFact> existingFacts)
     {
+        if (!shouldExtractFacts)
+            return;
+
         _ = Task.Run(async () =>
         {
             try
             {
                 using var scope = execution.ServiceScopeFactory.CreateScope();
-                var bgUnitOfWork = scope.ServiceProvider.GetRequiredService<IUnitOfWork>();
-                var bgUserRepo = scope.ServiceProvider.GetRequiredService<IGenericRepository<User>>();
-                var bgLogger = scope.ServiceProvider.GetRequiredService<ILogger<ProcessUserChatCommandHandler>>();
-
-                if (shouldExtractFacts)
-                    await SubmitFactExtractionBatchAsync(scope, userId, userMessage, aiMessage, existingFacts);
-
-                await IncrementAiMessageCountAsync(bgUserRepo, bgUnitOfWork, userId, bgLogger);
+                await SubmitFactExtractionBatchAsync(scope, userId, userMessage, aiMessage, existingFacts);
             }
             catch (Exception ex)
             {
@@ -91,30 +86,4 @@ private static async Task SubmitFactExtractionBatchAsync(
         await bgFactService.SubmitBatchAsync(userMessage: userMessage, aiResponse: aiMessage,
             existingFacts: existingFacts, userId: userId, cancellationToken: CancellationToken.None);
     }
-
-    private static async Task IncrementAiMessageCountAsync(
-        IGenericRepository<User> bgUserRepo,
-        IUnitOfWork bgUnitOfWork,
-        Guid userId,
-        ILogger bgLogger)
-    {
-        try
-        {
-            await ConcurrencyRetry.ExecuteAsync(
-                bgUserRepo,
-                bgUnitOfWork,
-                ct => bgUserRepo.FindOneTrackedAsync(u => u.Id == userId, cancellationToken: ct),
-                user =>
-                {
-                    user.IncrementAiMessageCount();
-                    return Task.FromResult(Result.Success());
-                },
-                ErrorMessages.UserNotFound,
-                CancellationToken.None);
-        }
-        catch (Exception ex)
-        {
-            LogBackgroundMessageCounterFailed(bgLogger, ex);
-        }
-    }
 }

src/Orbit.Application/Chat/Commands/ProcessUserChatCommand.cs

@@ -266,9 +266,6 @@ internal static bool IsShareableFaqTurn(ToolExecutionAccumulator results) =>
     [LoggerMessage(EventId = 20, Level = LogLevel.Debug, Message = "Saving changes to database...")]
     private static partial void LogSavingChanges(ILogger logger);
 
-    [LoggerMessage(EventId = 22, Level = LogLevel.Warning, Message = "Background message counter increment failed")]
-    private static partial void LogBackgroundMessageCounterFailed(ILogger logger, Exception ex);
-
     [LoggerMessage(EventId = 24, Level = LogLevel.Information, Message = "Tool {Name} requested clarification (operationId={OperationId}, missing={MissingKey})")]
     private static partial void LogClarificationRequested(ILogger logger, string name, Guid operationId, string missingKey);
 

src/Orbit.Application/Common/PayGateService.cs

@@ -68,6 +68,44 @@ public async Task<Result> CanSendAiMessage(Guid userId, CancellationToken ct = d
         return Result.Success();
     }
 
+    public async Task<Result> TryConsumeAiMessage(
+        Guid userId,
+        IUnitOfWork unitOfWork,
+        CancellationToken ct = default)
+    {
+        var freeLimit = await appConfig.GetAsync(AppConfigKeys.FreeAiMessagesPerMonth, AppConstants.DefaultFreeAiMessages, ct);
+        var proLimit = await appConfig.GetAsync(AppConfigKeys.ProAiMessagesPerMonth, AppConstants.DefaultProAiMessages, ct);
+
+        var consumption = await ConcurrencyRetry.ExecuteAsync(
+            userRepository,
+            unitOfWork,
+            token => userRepository.FindOneTrackedAsync(user => user.Id == userId, cancellationToken: token),
+            user =>
+            {
+                var currentAtUtc = DateTime.UtcNow;
+                if (!IsProductionSmokeAccount(user.Email))
+                {
+                    var messageLimit = (user.HasProAccess ? proLimit : freeLimit) + user.AdRewardBonusMessages;
+                    var cycleIsActive = user.AiMessagesResetAt.HasValue && user.AiMessagesResetAt.Value > currentAtUtc;
+                    if (cycleIsActive && user.AiMessagesUsedThisMonth >= messageLimit)
+                    {
+                        var errorMessage = user.HasProAccess
+                            ? $"You've reached your monthly AI message limit ({messageLimit})."
+                            : $"You've reached your monthly AI message limit ({messageLimit}). Upgrade to Pro for {proLimit} messages per month.";
+
+                        return Task.FromResult(Result.PayGateFailure(errorMessage));
+                    }
+                }
+
+                user.IncrementAiMessageCount(currentAtUtc);
+                return Task.FromResult(Result.Success());
+            },
+            ErrorMessages.UserNotFound,
+            ct);
+
+        return consumption.IsSuccess ? Result.Success() : consumption.PropagateError();
+    }
+
     public async Task<Result> CanUseDailySummary(Guid userId, CancellationToken ct = default)
     {
         var user = await userRepository.GetByIdAsync(userId, ct);

src/Orbit.Application/Habits/Commands/SuggestHabitSetupCommand.cs

@@ -2,10 +2,8 @@
 using System.Text;
 using MediatR;
 using Microsoft.Extensions.Caching.Memory;
-using Microsoft.Extensions.Logging;
 using Orbit.Application.Common;
 using Orbit.Domain.Common;
-using Orbit.Domain.Entities;
 using Orbit.Domain.Interfaces;
 using Orbit.Domain.Models;
 
@@ -19,64 +17,42 @@ public record SuggestHabitSetupCommand(
 public partial class SuggestHabitSetupCommandHandler(
     IPayGateService payGate,
     IHabitSuggestionService suggestionService,
-    IGenericRepository<User> userRepository,
     IUnitOfWork unitOfWork,
-    IMemoryCache cache,
-    ILogger<SuggestHabitSetupCommandHandler> logger)
+    IMemoryCache cache)
     : IRequestHandler<SuggestHabitSetupCommand, Result<HabitSetupSuggestion>>
 {
     private static readonly TimeSpan CacheTtl = TimeSpan.FromHours(1);
 
     public async Task<Result<HabitSetupSuggestion>> Handle(
         SuggestHabitSetupCommand request, CancellationToken cancellationToken)
     {
-        var gateCheck = await payGate.CanSendAiMessage(request.UserId, cancellationToken);
-        if (gateCheck.IsFailure)
-            return gateCheck.PropagateError<HabitSetupSuggestion>();
-
         var language = string.IsNullOrWhiteSpace(request.Language) ? "en" : request.Language;
         var cacheKey = BuildCacheKey(request.UserId, request.Title, language);
 
         if (cache.TryGetValue(cacheKey, out HabitSetupSuggestion? cached) && cached is not null)
             return Result.Success(cached);
 
+        var reservation = await payGate.TryConsumeAiMessage(
+            request.UserId,
+            unitOfWork,
+            cancellationToken);
+        if (reservation.IsFailure)
+            return reservation.PropagateError<HabitSetupSuggestion>();
+
         var suggestionResult = await suggestionService.SuggestSetupAsync(
             request.Title, language, cancellationToken);
         if (suggestionResult.IsFailure)
             return suggestionResult;
 
-        await IncrementUsageAsync(request.UserId, cancellationToken);
-
         cache.Set(cacheKey, suggestionResult.Value, CacheTtl);
 
         return suggestionResult;
     }
 
-    private async Task IncrementUsageAsync(Guid userId, CancellationToken cancellationToken)
-    {
-        var increment = await ConcurrencyRetry.ExecuteAsync(
-            userRepository,
-            unitOfWork,
-            ct => userRepository.FindOneTrackedAsync(user => user.Id == userId, cancellationToken: ct),
-            user =>
-            {
-                user.IncrementAiMessageCount();
-                return Task.FromResult(Result.Success());
-            },
-            ErrorMessages.UserNotFound,
-            cancellationToken);
-
-        if (increment.IsFailure)
-            LogUsageIncrementFailed(logger, userId);
-    }
-
     private static string BuildCacheKey(Guid userId, string title, string language)
     {
         var normalizedTitle = title.Trim().ToLowerInvariant();
         var titleHash = Convert.ToHexString(SHA256.HashData(Encoding.UTF8.GetBytes(normalizedTitle)));
         return $"suggest-setup:{userId}:{titleHash}:{language.ToLowerInvariant()}";
     }
-
-    [LoggerMessage(EventId = 1, Level = LogLevel.Warning, Message = "Failed to increment AI message usage after a habit suggestion for user {UserId}")]
-    private static partial void LogUsageIncrementFailed(ILogger logger, Guid userId);
 }

src/Orbit.Application/Tags/Queries/SuggestTagsQuery.cs

@@ -17,7 +17,6 @@ public class SuggestTagsQueryHandler(
     IPayGateService payGate,
     ITagSuggestionService tagSuggestionService,
     IGenericRepository<Tag> tagRepository,
-    IGenericRepository<User> userRepository,
     IUnitOfWork unitOfWork) : IRequestHandler<SuggestTagsQuery, Result<SuggestTagsResponse>>
 {
     private const string NewTagColor = "#7c3aed";
@@ -26,16 +25,19 @@ public async Task<Result<SuggestTagsResponse>> Handle(
         SuggestTagsQuery request,
         CancellationToken cancellationToken)
     {
-        var gateCheck = await payGate.CanSendAiMessage(request.UserId, cancellationToken);
-        if (gateCheck.IsFailure)
-            return gateCheck.PropagateError<SuggestTagsResponse>();
-
         var existingTags = await tagRepository.FindAsync(
             tag => tag.UserId == request.UserId,
             cancellationToken);
 
         var existingNames = existingTags.Select(tag => tag.Name).ToList();
 
+        var reservation = await payGate.TryConsumeAiMessage(
+            request.UserId,
+            unitOfWork,
+            cancellationToken);
+        if (reservation.IsFailure)
+            return reservation.PropagateError<SuggestTagsResponse>();
+
         var suggestionResult = await tagSuggestionService.SuggestTagsAsync(
             request.Title,
             request.Description,
@@ -48,8 +50,6 @@ public async Task<Result<SuggestTagsResponse>> Handle(
 
         var suggestions = MapSuggestions(suggestionResult.Value, existingTags);
 
-        await MeterAiMessageAsync(request.UserId, cancellationToken);
-
         return Result.Success(new SuggestTagsResponse(suggestions));
     }
 
@@ -81,18 +81,6 @@ private static List<SuggestedTag> MapSuggestions(
         return mapped;
     }
 
-    private async Task MeterAiMessageAsync(Guid userId, CancellationToken cancellationToken)
-    {
-        var user = await userRepository.FindOneTrackedAsync(
-            candidate => candidate.Id == userId,
-            cancellationToken: cancellationToken);
-        if (user is null)
-            return;
-
-        user.IncrementAiMessageCount();
-        await unitOfWork.SaveChangesAsync(cancellationToken);
-    }
-
     private static string Capitalize(string value) =>
         string.IsNullOrEmpty(value) ? value : char.ToUpper(value[0]) + value[1..].ToLower();
 }

src/Orbit.Domain/Interfaces/IPayGateService.cs

@@ -19,6 +19,14 @@ public interface IPayGateService
     /// </summary>
     Task<Result> CanSendAiMessage(Guid userId, CancellationToken ct = default);
 
+    /// <summary>
+    /// Atomically consumes one AI message when the user has quota remaining.
+    /// </summary>
+    Task<Result> TryConsumeAiMessage(
+        Guid userId,
+        IUnitOfWork unitOfWork,
+        CancellationToken ct = default);
+
     /// <summary>
     /// Checks if the user can use daily AI summaries (Pro-only feature).
     /// </summary>