From cc7d59dd7ef97ebd21cd7aa6ee9a760e947f6e1e Mon Sep 17 00:00:00 2001 From: thomasluizon Date: Fri, 7 Aug 2026 00:30:50 -0300 Subject: [PATCH 1/6] chore: start ORB-91 From 93f70f513acff890616f504552a71a5de8393ffc Mon Sep 17 00:00:00 2001 From: thomasluizon Date: Fri, 7 Aug 2026 00:39:57 -0300 Subject: [PATCH 2/6] fix: reserve chat AI quota atomically --- .../ProcessUserChatCommand.Context.cs | 2 +- .../ProcessUserChatCommand.Persistence.cs | 41 ++------- .../Chat/Commands/ProcessUserChatCommand.cs | 3 - .../Common/PayGateService.cs | 36 +++++++- .../Interfaces/IPayGateService.cs | 5 ++ .../ProcessUserChatCommandHandlerTests.cs | 87 ++++++++++++++++++- .../Common/PayGateServiceTests.cs | 77 +++++++++++++++- 7 files changed, 205 insertions(+), 46 deletions(-) diff --git a/src/Orbit.Application/Chat/Commands/ProcessUserChatCommand.Context.cs b/src/Orbit.Application/Chat/Commands/ProcessUserChatCommand.Context.cs index dd9e6000..63960e74 100644 --- a/src/Orbit.Application/Chat/Commands/ProcessUserChatCommand.Context.cs +++ b/src/Orbit.Application/Chat/Commands/ProcessUserChatCommand.Context.cs @@ -43,7 +43,7 @@ private async Task> LoadChatContextAsync( activeGoals = loadedGoals; } - var messageGate = await execution.PayGateService.CanSendAiMessage(request.UserId, cancellationToken); + var messageGate = await execution.PayGateService.TryConsumeAiMessage(request.UserId, cancellationToken); if (messageGate.IsFailure) return messageGate.PropagateError(); diff --git a/src/Orbit.Application/Chat/Commands/ProcessUserChatCommand.Persistence.cs b/src/Orbit.Application/Chat/Commands/ProcessUserChatCommand.Persistence.cs index e1150204..7922f10d 100644 --- a/src/Orbit.Application/Chat/Commands/ProcessUserChatCommand.Persistence.cs +++ b/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 action } /// - /// 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. /// private void RunBackgroundPostResponseWork( @@ -59,19 +58,15 @@ private void RunBackgroundPostResponseWork( bool shouldExtractFacts, IReadOnlyList existingFacts) { + if (!shouldExtractFacts) + return; + _ = Task.Run(async () => { try { using var scope = execution.ServiceScopeFactory.CreateScope(); - var bgUnitOfWork = scope.ServiceProvider.GetRequiredService(); - var bgUserRepo = scope.ServiceProvider.GetRequiredService>(); - var bgLogger = scope.ServiceProvider.GetRequiredService>(); - - 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 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); - } - } } diff --git a/src/Orbit.Application/Chat/Commands/ProcessUserChatCommand.cs b/src/Orbit.Application/Chat/Commands/ProcessUserChatCommand.cs index cea3cbe8..dad314e9 100644 --- a/src/Orbit.Application/Chat/Commands/ProcessUserChatCommand.cs +++ b/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); diff --git a/src/Orbit.Application/Common/PayGateService.cs b/src/Orbit.Application/Common/PayGateService.cs index ced219e6..30abf18a 100644 --- a/src/Orbit.Application/Common/PayGateService.cs +++ b/src/Orbit.Application/Common/PayGateService.cs @@ -7,7 +7,8 @@ namespace Orbit.Application.Common; public class PayGateService( IGenericRepository habitRepository, IGenericRepository userRepository, - IAppConfigService appConfig) : IPayGateService + IAppConfigService appConfig, + IUnitOfWork unitOfWork) : IPayGateService { public async Task CanCreateHabits(Guid userId, int count = 1, CancellationToken ct = default) { @@ -67,6 +68,39 @@ public async Task CanSendAiMessage(Guid userId, CancellationToken ct = d return Result.Success(); } + public async Task TryConsumeAiMessage(Guid userId, 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 => + { + if (!IsProductionSmokeAccount(user.Email)) + { + var messageLimit = (user.HasProAccess ? proLimit : freeLimit) + user.AdRewardBonusMessages; + if (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(); + return Task.FromResult(Result.Success()); + }, + ErrorMessages.UserNotFound, + ct); + + return consumption.IsSuccess ? Result.Success() : consumption.PropagateError(); + } + public async Task CanUseDailySummary(Guid userId, CancellationToken ct = default) { var user = await userRepository.GetByIdAsync(userId, ct); diff --git a/src/Orbit.Domain/Interfaces/IPayGateService.cs b/src/Orbit.Domain/Interfaces/IPayGateService.cs index 66cf2a5d..2b2b5282 100644 --- a/src/Orbit.Domain/Interfaces/IPayGateService.cs +++ b/src/Orbit.Domain/Interfaces/IPayGateService.cs @@ -19,6 +19,11 @@ public interface IPayGateService /// Task CanSendAiMessage(Guid userId, CancellationToken ct = default); + /// + /// Atomically consumes one AI message when the user has quota remaining. + /// + Task TryConsumeAiMessage(Guid userId, CancellationToken ct = default); + /// /// Checks if the user can use daily AI summaries (Pro-only feature). /// diff --git a/tests/Orbit.Application.Tests/Commands/Chat/ProcessUserChatCommandHandlerTests.cs b/tests/Orbit.Application.Tests/Commands/Chat/ProcessUserChatCommandHandlerTests.cs index 3ea8bbb5..5fe5f386 100644 --- a/tests/Orbit.Application.Tests/Commands/Chat/ProcessUserChatCommandHandlerTests.cs +++ b/tests/Orbit.Application.Tests/Commands/Chat/ProcessUserChatCommandHandlerTests.cs @@ -7,6 +7,7 @@ using Orbit.Application.Chat.Commands; using Orbit.Application.Chat.Models; using Orbit.Application.Chat.Tools; +using Orbit.Application.Common; using Orbit.Application.Goals.Services; using Orbit.Domain.Common; using Orbit.Domain.Entities; @@ -62,13 +63,20 @@ private static Habit CreateHabit(string title, bool isCompleted = false) } private ProcessUserChatCommandHandler CreateHandler(params IAiTool[] tools) + { + return CreateHandler(_payGate, tools); + } + + private ProcessUserChatCommandHandler CreateHandler( + IPayGateService payGate, + params IAiTool[] tools) { var toolRegistry = new AiToolRegistry(tools); SetupOperationExecutor(toolRegistry); var aiDeps = new ChatAiDependencies(_aiIntentService, toolRegistry, _promptBuilder, _catalogService); var dataDeps = new ChatDataDependencies(_habitRepo, _goalRepo, _userRepo, _userFactRepo, _tagRepo, _checklistTemplateRepo, _featureFlagService); var executionDeps = new ChatExecutionDependencies( - _userDateService, _userStreakService, _payGate, _unitOfWork, _scopeFactory, _operationExecutor, _pendingClarificationStore, _streakGoalReadSyncer, _gamificationService); + _userDateService, _userStreakService, payGate, _unitOfWork, _scopeFactory, _operationExecutor, _pendingClarificationStore, _streakGoalReadSyncer, _gamificationService); return new ProcessUserChatCommandHandler( dataDeps, aiDeps, executionDeps, _logger); @@ -208,10 +216,28 @@ private void SetupUserAndPayGate(User? user = null, bool payGatePass = true) { user ??= User.Create("Thomas", "thomas@test.com").Value; _userRepo.GetByIdAsync(UserId, Arg.Any()).Returns(user); - _payGate.CanSendAiMessage(UserId, Arg.Any()) + _payGate.TryConsumeAiMessage(UserId, Arg.Any()) .Returns(payGatePass ? Result.Success() : Result.PayGateFailure("AI message limit reached.")); } + private PayGateService CreateRealPayGate(User user) + { + _userRepo.GetByIdAsync(UserId, Arg.Any()).Returns(user); + _userRepo.FindOneTrackedAsync( + Arg.Any>>(), + Arg.Any, IQueryable>?>(), + Arg.Any()) + .Returns(user); + + var appConfig = Substitute.For(); + appConfig.GetAsync(AppConfigKeys.FreeAiMessagesPerMonth, AppConstants.DefaultFreeAiMessages, Arg.Any()) + .Returns(20); + appConfig.GetAsync(AppConfigKeys.ProAiMessagesPerMonth, AppConstants.DefaultProAiMessages, Arg.Any()) + .Returns(500); + + return new PayGateService(_habitRepo, _userRepo, appConfig, _unitOfWork); + } + private static readonly AiConversationContext TestConversationContext = new() { Messages = new List(), @@ -258,6 +284,31 @@ public async Task Handle_PayGateBlocks_ReturnsPayGateError() result.ErrorCode.Should().Be("PAY_GATE"); } + [Fact] + public async Task Handle_TenConcurrentRequestsWithOneMessageRemaining_OnlyOneCallsAi() + { + var user = User.Create("Thomas", "thomas@test.com").Value; + user.StartTrial(DateTime.UtcNow.AddDays(-1)); + for (var i = 0; i < 19; i++) + user.IncrementAiMessageCount(); + + SetupAiResponse(new AiResponse { TextMessage = "Reserved response", ToolCalls = null }); + var handler = CreateHandler(CreateRealPayGate(user)); + var command = new ProcessUserChatCommand(UserId, "Hello AI"); + + var results = await Task.WhenAll( + Enumerable.Range(0, 10) + .Select(_ => handler.Handle(command, CancellationToken.None))); + + results.Should().ContainSingle(result => result.IsSuccess); + results.Count(result => result.IsFailure && result.ErrorCode == Result.PayGateErrorCode).Should().Be(9); + user.AiMessagesUsedThisMonth.Should().Be(20); + await _aiIntentService.Received(1).SendWithToolsAsync( + Arg.Any(), + Arg.Any?>(), + Arg.Any()); + } + [Fact] public async Task Handle_AiServiceFails_ReturnsFailure() { @@ -272,6 +323,34 @@ public async Task Handle_AiServiceFails_ReturnsFailure() result.Error.Should().Be("AI service unavailable"); } + [Fact] + public async Task Handle_AiServiceFails_RetainsConsumedQuota() + { + var user = User.Create("Thomas", "thomas@test.com").Value; + user.StartTrial(DateTime.UtcNow.AddDays(-1)); + for (var i = 0; i < 19; i++) + user.IncrementAiMessageCount(); + + _aiIntentService.SendWithToolsAsync( + Arg.Any(), + Arg.Any?>(), + Arg.Any()) + .Returns(_ => + { + user.AiMessagesUsedThisMonth.Should().Be(20); + return Result.Failure("AI service unavailable"); + }); + var handler = CreateHandler(CreateRealPayGate(user)); + + var result = await handler.Handle( + new ProcessUserChatCommand(UserId, "Hello AI"), + CancellationToken.None); + + result.IsFailure.Should().BeTrue(); + result.Error.Should().Be("AI service unavailable"); + user.AiMessagesUsedThisMonth.Should().Be(20); + } + [Fact] public async Task Handle_SuccessfulResponse_ReturnsChatResponse() { @@ -503,7 +582,7 @@ public async Task Handle_AiResponseWithJsonWrapper_StripsWrapper() public async Task Handle_NullUser_StillSucceeds() { _userRepo.GetByIdAsync(UserId, Arg.Any()).Returns((User?)null); - _payGate.CanSendAiMessage(UserId, Arg.Any()).Returns(Result.Success()); + _payGate.TryConsumeAiMessage(UserId, Arg.Any()).Returns(Result.Success()); SetupAiResponse(new AiResponse { TextMessage = "Response", ToolCalls = null }); var handler = CreateHandler(); @@ -1272,7 +1351,7 @@ public async Task Handle_AiMemoryDisabled_DoesNotLoadFacts() var user = User.Create("Thomas", "thomas@test.com").Value; user.SetAiMemory(false); _userRepo.GetByIdAsync(UserId, Arg.Any()).Returns(user); - _payGate.CanSendAiMessage(UserId, Arg.Any()).Returns(Result.Success()); + _payGate.TryConsumeAiMessage(UserId, Arg.Any()).Returns(Result.Success()); SetupAiResponse(new AiResponse { TextMessage = "Hi!", ToolCalls = null }); var handler = CreateHandler(); diff --git a/tests/Orbit.Application.Tests/Common/PayGateServiceTests.cs b/tests/Orbit.Application.Tests/Common/PayGateServiceTests.cs index 210739cb..a7362559 100644 --- a/tests/Orbit.Application.Tests/Common/PayGateServiceTests.cs +++ b/tests/Orbit.Application.Tests/Common/PayGateServiceTests.cs @@ -13,13 +13,14 @@ public class PayGateServiceTests private readonly IGenericRepository _habitRepo = Substitute.For>(); private readonly IGenericRepository _userRepo = Substitute.For>(); private readonly IAppConfigService _appConfig = Substitute.For(); + private readonly IUnitOfWork _unitOfWork = Substitute.For(); private readonly PayGateService _sut; private static readonly Guid UserId = Guid.NewGuid(); public PayGateServiceTests() { - _sut = new PayGateService(_habitRepo, _userRepo, _appConfig); + _sut = new PayGateService(_habitRepo, _userRepo, _appConfig, _unitOfWork); _appConfig.GetAsync("FreeMaxHabits", 10, Arg.Any()).Returns(10); _appConfig.GetAsync("SubHabitsProOnly", true, Arg.Any()).Returns(true); @@ -209,6 +210,80 @@ await WithEnvironment("Development", user.Email, async () => }); } + [Fact] + public async Task TryConsumeAiMessage_UnderLimit_IncrementsAndSaves() + { + var user = CreateFreeUser(); + user.StartTrial(DateTime.UtcNow.AddDays(-1)); + _userRepo.FindOneTrackedAsync( + Arg.Any>>(), + Arg.Any, IQueryable>?>(), + Arg.Any()) + .Returns(user); + + var result = await _sut.TryConsumeAiMessage(UserId); + + result.IsSuccess.Should().BeTrue(); + user.AiMessagesUsedThisMonth.Should().Be(1); + await _unitOfWork.Received(1).SaveChangesAsync(Arg.Any()); + } + + [Fact] + public async Task TryConsumeAiMessage_AtLimit_ReturnsPayGateFailureWithoutSaving() + { + var user = CreateFreeUser(); + user.StartTrial(DateTime.UtcNow.AddDays(-1)); + for (var i = 0; i < 20; i++) + user.IncrementAiMessageCount(); + _userRepo.FindOneTrackedAsync( + Arg.Any>>(), + Arg.Any, IQueryable>?>(), + Arg.Any()) + .Returns(user); + + var result = await _sut.TryConsumeAiMessage(UserId); + + result.IsFailure.Should().BeTrue(); + result.ErrorCode.Should().Be("PAY_GATE"); + user.AiMessagesUsedThisMonth.Should().Be(20); + await _unitOfWork.DidNotReceive().SaveChangesAsync(Arg.Any()); + } + + [Fact] + public async Task TryConsumeAiMessage_ProUserWithHeadroom_IncrementsAndSaves() + { + var user = CreateProUser(); + for (var i = 0; i < 20; i++) + user.IncrementAiMessageCount(); + _userRepo.FindOneTrackedAsync( + Arg.Any>>(), + Arg.Any, IQueryable>?>(), + Arg.Any()) + .Returns(user); + + var result = await _sut.TryConsumeAiMessage(UserId); + + result.IsSuccess.Should().BeTrue(); + user.AiMessagesUsedThisMonth.Should().Be(21); + await _unitOfWork.Received(1).SaveChangesAsync(Arg.Any()); + } + + [Fact] + public async Task TryConsumeAiMessage_UserNotFound_ReturnsFailureWithoutSaving() + { + _userRepo.FindOneTrackedAsync( + Arg.Any>>(), + Arg.Any, IQueryable>?>(), + Arg.Any()) + .Returns((User?)null); + + var result = await _sut.TryConsumeAiMessage(UserId); + + result.IsFailure.Should().BeTrue(); + result.Error.Should().Contain("User not found"); + await _unitOfWork.DidNotReceive().SaveChangesAsync(Arg.Any()); + } + [Fact] public async Task CanUseDailySummary_ProUser_Success() { From 44d9440ad16298952ab0910982d110bb9f4ee882 Mon Sep 17 00:00:00 2001 From: thomasluizon Date: Fri, 7 Aug 2026 00:44:16 -0300 Subject: [PATCH 3/6] fix: satisfy atomic quota delivery guards --- architecture.html | 2 +- architecture.json | 2 +- .../ProcessUserChatCommand.Context.cs | 5 +- .../Common/PayGateService.cs | 8 +- .../Interfaces/IPayGateService.cs | 5 +- .../ProcessUserChatCommandHandlerTests.cs | 8 +- .../Common/PayGateServiceTests.cs | 77 +------------------ 7 files changed, 20 insertions(+), 87 deletions(-) diff --git a/architecture.html b/architecture.html index 2033cc4b..f9317d28 100644 --- a/architecture.html +++ b/architecture.html @@ -47,7 +47,7 @@

Handlers with no endpoint

RequestHandler file

Entities

EntityDomain file
- + + +