[PM-31781] skip unpaid automations for exempt orgs

src/Billing/Services/Implementations/InvoiceCreatedHandler.cs

@@ -40,8 +40,8 @@ public async Task HandleAsync(Event parsedEvent)
                     Status: not StripeConstants.InvoiceStatus.Paid,
                     CollectionMethod: "charge_automatically",
                     BillingReason:
-                    "subscription_cycle" or
-                    "automatic_pending_invoice_item_invoice",
+                    StripeConstants.BillingReasons.SubscriptionCycle or
+                    StripeConstants.BillingReasons.AutomaticPendingInvoiceItemInvoice,
                     Parent.SubscriptionDetails.Subscription: not null
                 })
             {

src/Billing/Services/Implementations/PaymentFailedHandler.cs

@@ -50,7 +50,7 @@ invoice is
             AmountDue: > 0,
             Status: not StripeConstants.InvoiceStatus.Paid,
             CollectionMethod: "charge_automatically",
-            BillingReason: "subscription_cycle" or "automatic_pending_invoice_item_invoice",
+            BillingReason: StripeConstants.BillingReasons.SubscriptionCycle or StripeConstants.BillingReasons.AutomaticPendingInvoiceItemInvoice,
             Parent.SubscriptionDetails: not null
         };
 }

src/Billing/Services/Implementations/StripeEventUtilityService.cs

@@ -258,7 +258,7 @@ invoice is
             AmountDue: > 0,
             Status: not StripeConstants.InvoiceStatus.Paid,
             CollectionMethod: "charge_automatically",
-            BillingReason: "subscription_cycle" or "automatic_pending_invoice_item_invoice",
+            BillingReason: StripeConstants.BillingReasons.SubscriptionCycle or StripeConstants.BillingReasons.AutomaticPendingInvoiceItemInvoice,
             Parent.SubscriptionDetails: not null
         };
 

src/Billing/Services/Implementations/SubscriptionUpdatedHandler.cs

@@ -1,4 +1,6 @@
 using Bit.Core;
+using Bit.Core.AdminConsole.Entities;
+using Bit.Core.AdminConsole.Entities.Provider;
 using Bit.Core.AdminConsole.OrganizationFeatures.Organizations.Interfaces;
 using Bit.Core.AdminConsole.Repositories;
 using Bit.Core.AdminConsole.Services;
@@ -8,6 +10,7 @@
 using Bit.Core.Billing.Pricing;
 using Bit.Core.Billing.Services;
 using Bit.Core.Billing.Subscriptions.Models;
+using Bit.Core.Entities;
 using Bit.Core.OrganizationFeatures.OrganizationSponsorships.FamiliesForEnterprise.Interfaces;
 using Bit.Core.Repositories;
 using Bit.Core.Services;
@@ -83,39 +86,73 @@ public async Task HandleAsync(Event parsedEvent)
         var subscription = await _stripeEventService.GetSubscription(parsedEvent, true, ["customer", "discounts", "latest_invoice", "test_clock"]);
         SubscriberId subscriberId = subscription;
 
+        var subscriber = await GetSubscriberAsync(subscriberId);
+        if (subscriber == null)
+        {
+            _logger.LogWarning(
+                "Subscriber not found for subscription ({SubscriptionId}) in event ({EventId}), skipping handler",
+                subscription.Id,
+                parsedEvent.Id);
+            return;
+        }
+
         var currentPeriodEnd = subscription.GetCurrentPeriodEnd();
+        var clearOrgBillingAutomationExemption = false;
 
-        if (SubscriptionWentUnpaid(parsedEvent, subscription) ||
-            SubscriptionWentIncompleteExpired(parsedEvent, subscription))
+        if (SubscriptionWentUnpaid(parsedEvent, subscription))
+        {
+            if (SkipUnpaidBillingAutomationsForExemptOrganization(subscriber))
+            {
+                _logger.LogInformation(
+                    "Skipping billing automations for exempt organization ({OrganizationId}). Exemption will be cleared after handler completion",
+                    subscriber.Id);
+                clearOrgBillingAutomationExemption = true;
+            }
+            else
+            {
+                await DisableSubscriberAsync(subscriber, currentPeriodEnd);
+                await SetSubscriptionToCancelAsync(subscription);
+            }
+        }
+        else if (SubscriptionWentIncompleteExpired(parsedEvent, subscription))
         {
-            await DisableSubscriberAsync(subscriberId, currentPeriodEnd);
+            await DisableSubscriberAsync(subscriber, currentPeriodEnd);
             await SetSubscriptionToCancelAsync(subscription);
         }
         else if (SubscriptionBecameActive(parsedEvent, subscription))
         {
-            await EnableSubscriberAsync(subscriberId, currentPeriodEnd);
+            await EnableSubscriberAsync(subscriber, currentPeriodEnd);
             await RemovePendingCancellationAsync(subscription);
         }
 
-        await subscriberId.Match(
-            userId => _userService.UpdatePremiumExpirationAsync(userId.Value, currentPeriodEnd),
-            async organizationId =>
-            {
-                if (_featureService.IsEnabled(FeatureFlagKeys.PM32645_DeferPriceMigrationToRenewal))
+        switch (subscriber)
+        {
+            case User user:
+                await _userService.UpdatePremiumExpirationAsync(user.Id, currentPeriodEnd);
+                break;
+            case Organization organization:
                 {
-                    await HandleScheduleTriggeredFamiliesMigrationAsync(parsedEvent, subscription, organizationId.Value);
-                }
+                    if (_featureService.IsEnabled(FeatureFlagKeys.PM32645_DeferPriceMigrationToRenewal))
+                    {
+                        await HandleScheduleTriggeredFamiliesMigrationAsync(parsedEvent, subscription, organization);
+                    }
 
-                await _organizationService.UpdateExpirationDateAsync(organizationId.Value, currentPeriodEnd);
+                    await _organizationService.UpdateExpirationDateAsync(organization.Id, currentPeriodEnd);
 
-                if (_stripeEventUtilityService.IsSponsoredSubscription(subscription) && currentPeriodEnd.HasValue)
-                {
-                    await _organizationSponsorshipRenewCommand.UpdateExpirationDateAsync(organizationId.Value, currentPeriodEnd.Value);
-                }
+                    if (_stripeEventUtilityService.IsSponsoredSubscription(subscription) && currentPeriodEnd.HasValue)
+                    {
+                        await _organizationSponsorshipRenewCommand.UpdateExpirationDateAsync(organization.Id, currentPeriodEnd.Value);
+                    }
 
-                await RemovePasswordManagerCouponIfRemovingSecretsManagerTrialAsync(parsedEvent, subscription);
-            },
-            _ => Task.CompletedTask);
+                    await RemovePasswordManagerCouponIfRemovingSecretsManagerTrialAsync(parsedEvent, subscription, organization);
+
+                    if (clearOrgBillingAutomationExemption)
+                    {
+                        await ClearBillingAutomationExemptionAsync(organization);
+                    }
+                    break;
+                }
+        }
     }
 
     private static bool SubscriptionWentUnpaid(
@@ -130,7 +167,7 @@ SubscriptionStatus.Active or
         } && currentSubscription is
         {
             Status: SubscriptionStatus.Unpaid,
-            LatestInvoice.BillingReason: BillingReasons.SubscriptionCreate or BillingReasons.SubscriptionCycle
+            LatestInvoice.BillingReason: BillingReasons.SubscriptionCycle
         };
 
     private static bool SubscriptionWentIncompleteExpired(
@@ -159,65 +196,79 @@ SubscriptionStatus.Incomplete or
             LatestInvoice.BillingReason: BillingReasons.SubscriptionCreate or BillingReasons.SubscriptionCycle
         };
 
-    private Task DisableSubscriberAsync(SubscriberId subscriberId, DateTime? currentPeriodEnd) =>
-        subscriberId.Match(
-            async userId =>
-            {
-                await _userService.DisablePremiumAsync(userId.Value, currentPeriodEnd);
-                var user = await _userRepository.GetByIdAsync(userId.Value);
-                if (user != null)
-                {
-                    await _pushNotificationAdapter.NotifyPremiumStatusChangedAsync(user);
-                }
-            },
-            async organizationId =>
-            {
-                await _organizationDisableCommand.DisableAsync(organizationId.Value, currentPeriodEnd);
-                var organization = await _organizationRepository.GetByIdAsync(organizationId.Value);
-                if (organization != null)
-                {
-                    await _pushNotificationAdapter.NotifyEnabledChangedAsync(organization);
-                }
-            },
-            async providerId =>
-            {
-                var provider = await _providerRepository.GetByIdAsync(providerId.Value);
-                if (provider != null)
-                {
-                    provider.Enabled = false;
-                    await _providerService.UpdateAsync(provider);
-                }
-            });
+    private static bool SkipUnpaidBillingAutomationsForExemptOrganization(ISubscriber subscriber) =>
+        subscriber is Organization { Enabled: true, ExemptFromBillingAutomation: true };
 
-    private Task EnableSubscriberAsync(SubscriberId subscriberId, DateTime? currentPeriodEnd) =>
-        subscriberId.Match(
-            async userId =>
-            {
-                await _userService.EnablePremiumAsync(userId.Value, currentPeriodEnd);
-                var user = await _userRepository.GetByIdAsync(userId.Value);
-                if (user != null)
-                {
-                    await _pushNotificationAdapter.NotifyPremiumStatusChangedAsync(user!);
-                }
-            },
-            async organizationId =>
-            {
-                await _organizationEnableCommand.EnableAsync(organizationId.Value, currentPeriodEnd);
-                var organization = await _organizationRepository.GetByIdAsync(organizationId.Value);
-                if (organization != null)
-                {
-                    await _pushNotificationAdapter.NotifyEnabledChangedAsync(organization);
-                }
-            },
-            async providerId =>
-            {
-                var provider = await _providerRepository.GetByIdAsync(providerId.Value);
-                if (provider != null)
-                {
-                    provider.Enabled = true;
-                    await _providerService.UpdateAsync(provider);
-                }
-            });
+    private Task<ISubscriber?> GetSubscriberAsync(SubscriberId subscriberId) =>
+        subscriberId.Match<Task<ISubscriber?>>(
+            async userId => await _userRepository.GetByIdAsync(userId.Value),
+            async organizationId => await _organizationRepository.GetByIdAsync(organizationId.Value),
+            async providerId => await _providerRepository.GetByIdAsync(providerId.Value));
+
+    private Task DisableSubscriberAsync(ISubscriber subscriber, DateTime? currentPeriodEnd) =>
+        subscriber switch
+        {
+            User user => DisableUserAsync(user, currentPeriodEnd),
+            Organization organization => DisableOrganizationAsync(organization, currentPeriodEnd),
+            Provider provider => DisableProviderAsync(provider),
+            _ => Task.CompletedTask
+        };
+
+    private async Task DisableUserAsync(User user, DateTime? currentPeriodEnd)
+    {
+        await _userService.DisablePremiumAsync(user.Id, currentPeriodEnd);
+        await _pushNotificationAdapter.NotifyPremiumStatusChangedAsync(user);
+    }
+
+    private async Task DisableOrganizationAsync(Organization organization, DateTime? currentPeriodEnd)
+    {
+        await _organizationDisableCommand.DisableAsync(organization.Id, currentPeriodEnd);
+        await _pushNotificationAdapter.NotifyEnabledChangedAsync(organization);
+    }
+
+    private async Task DisableProviderAsync(Provider provider)
+    {
+        provider.Enabled = false;
+        await _providerService.UpdateAsync(provider);
+    }
+
+    private async Task ClearBillingAutomationExemptionAsync(Organization organization)
+    {
+        organization.ExemptFromBillingAutomation = false;
+        organization.RevisionDate = DateTime.UtcNow;
+        await _organizationRepository.ReplaceAsync(organization);
+
+        _logger.LogInformation(
+            "Exemption has been cleared for organization ({OrganizationId})",
+            organization.Id);
+    }
+
+    private Task EnableSubscriberAsync(ISubscriber subscriber, DateTime? currentPeriodEnd) =>
+        subscriber switch
+        {
+            User user => EnableUserAsync(user, currentPeriodEnd),
+            Organization organization => EnableOrganizationAsync(organization, currentPeriodEnd),
+            Provider provider => EnableProviderAsync(provider),
+            _ => Task.CompletedTask
+        };
+
+    private async Task EnableUserAsync(User user, DateTime? currentPeriodEnd)
+    {
+        await _userService.EnablePremiumAsync(user.Id, currentPeriodEnd);
+        await _pushNotificationAdapter.NotifyPremiumStatusChangedAsync(user);
+    }
+
+    private async Task EnableOrganizationAsync(Organization organization, DateTime? currentPeriodEnd)
+    {
+        await _organizationEnableCommand.EnableAsync(organization.Id, currentPeriodEnd);
+        await _pushNotificationAdapter.NotifyEnabledChangedAsync(organization);
+    }
+
+    private async Task EnableProviderAsync(Provider provider)
+    {
+        provider.Enabled = true;
+        await _providerService.UpdateAsync(provider);
+    }
 
     private async Task SetSubscriptionToCancelAsync(Subscription subscription)
     {
@@ -258,24 +309,17 @@ private async Task RemovePendingCancellationAsync(Subscription subscription)
     /// </summary>
     /// <param name="parsedEvent"></param>
     /// <param name="subscription"></param>
+    /// <param name="organization"></param>
     private async Task RemovePasswordManagerCouponIfRemovingSecretsManagerTrialAsync(
         Event parsedEvent,
-        Subscription subscription)
+        Subscription subscription,
+        Organization organization)
     {
         if (parsedEvent.Data.PreviousAttributes?.items is null)
         {
             return;
         }
 
-        var organization = subscription.Metadata.TryGetValue("organizationId", out var organizationId)
-            ? await _organizationRepository.GetByIdAsync(Guid.Parse(organizationId))
-            : null;
-
-        if (organization == null)
-        {
-            return;
-        }
-
         var plan = await _pricingClient.GetPlanOrThrow(organization.PlanType);
 
         if (!plan.SupportsSecretsManager)
@@ -346,7 +390,7 @@ private async Task WaitForTestClockToAdvanceAsync(TestClock testClock)
     private async Task HandleScheduleTriggeredFamiliesMigrationAsync(
         Event parsedEvent,
         Subscription subscription,
-        Guid organizationId)
+        Organization organization)
     {
         try
         {
@@ -388,15 +432,6 @@ private async Task HandleScheduleTriggeredFamiliesMigrationAsync(
             }
 
             // Sync org DB to match the plan Stripe already transitioned to
-            var organization = await _organizationRepository.GetByIdAsync(organizationId);
-            if (organization == null)
-            {
-                _logger.LogWarning(
-                    "Organization ({OrganizationId}) not found for schedule-triggered Families migration",
-                    organizationId);
-                return;
-            }
-
             organization.PlanType = familiesPlan.Type;
             organization.Plan = familiesPlan.Name;
             organization.UsersGetPremium = familiesPlan.UsersGetPremium;
@@ -406,7 +441,7 @@ private async Task HandleScheduleTriggeredFamiliesMigrationAsync(
 
             _logger.LogInformation(
                 "Updated organization ({OrganizationId}) to FamiliesAnnually plan after schedule transition",
-                organizationId);
+                organization.Id);
         }
         catch (BillingException)
         {
@@ -421,7 +456,7 @@ private async Task HandleScheduleTriggeredFamiliesMigrationAsync(
             _logger.LogError(
                 exception,
                 "Failed to handle schedule-triggered Families migration for organization ({OrganizationId})",
-                organizationId);
+                organization.Id);
         }
     }
 }

src/Core/Billing/Constants/StripeConstants.cs

@@ -14,8 +14,10 @@ public static class AutomaticTaxStatus
 
     public static class BillingReasons
     {
+        public const string AutomaticPendingInvoiceItemInvoice = "automatic_pending_invoice_item_invoice";
         public const string SubscriptionCreate = "subscription_create";
         public const string SubscriptionCycle = "subscription_cycle";
+        public const string SubscriptionUpdate = "subscription_update";
     }
 
     public static class CollectionMethod