[PM-37084] Business Aware Schedule Recovery and Cancellation

src/Billing/Services/Implementations/SubscriptionUpdatedHandler.cs

@@ -83,7 +83,7 @@ public SubscriptionUpdatedHandler(
 
     public async Task HandleAsync(Event parsedEvent)
     {
-        var subscription = await _stripeEventService.GetSubscription(parsedEvent, true, ["customer", "discounts", "latest_invoice", "test_clock"]);
+        var subscription = await _stripeEventService.GetSubscription(parsedEvent, true, ["customer.discount", "discounts", "latest_invoice", "test_clock"]);
         SubscriberId subscriberId = subscription;
 
         var subscriber = await GetSubscriberAsync(subscriberId);
@@ -338,8 +338,6 @@ private async Task SetSubscriptionToCancelAsync(Subscription subscription)
 
     private async Task RemovePendingCancellationAsync(Subscription subscription)
     {
-        await _priceIncreaseScheduler.SchedulePersonalPriceIncrease(subscription);
-
         await _stripeAdapter.UpdateSubscriptionAsync(subscription.Id, new SubscriptionUpdateOptions
         {
             CancelAtPeriodEnd = false,
@@ -352,6 +350,7 @@ private async Task RemovePendingCancellationAsync(Subscription subscription)
                 [MetadataKeys.CancellationOrigin] = string.Empty
             }
         });
+        await _priceIncreaseScheduler.ScheduleForSubscription(subscription);
     }
 
     /// <summary>

src/Billing/Services/Implementations/UpcomingInvoiceHandler.cs

@@ -22,6 +22,7 @@
 using Bit.Core.Repositories;
 using Bit.Core.Services;
 using Stripe;
+using Stripe.TestHelpers;
 using Event = Stripe.Event;
 using Plan = Bit.Core.Models.StaticStore.Plan;
 using PremiumPlan = Bit.Core.Billing.Pricing.Premium.Plan;
@@ -356,48 +357,29 @@ private async Task<bool> ScheduleBusinessPlanPriceMigrationAsync(
             }
 
             var assignment = await assignmentRepository.GetByOrganizationIdAsync(organization.Id);
-
             if (assignment is null || assignment.ScheduledDate is not null)
             {
                 return false;
             }
 
-            var cohort = await cohortRepository.GetByIdAsync(assignment.CohortId);
-
-            if (cohort is null || !cohort.IsActive)
-            {
-                return false;
-            }
-
-            if (cohort.MigrationPathId is null)
+            if (subscription.TestClock != null)
             {
-                // Churn-only cohort — no migration to schedule.
-                return false;
-            }
-
-            var migrationPath = MigrationPaths.FromId(cohort.MigrationPathId.Value);
-            if (migrationPath is null)
-            {
-                logger.LogError(
-                    "Unknown MigrationPathId ({MigrationPathId}) on cohort ({CohortId}) for Organization ({OrganizationId})",
-                    cohort.MigrationPathId, cohort.Id, organization.Id);
-                return false;
+                await WaitForTestClockToAdvanceAsync(subscription.TestClock);
             }
 
-            if (organization.PlanType != migrationPath.FromPlan)
-            {
-                logger.LogWarning(
-                    "Skipping business price migration for Organization ({OrganizationId}); PlanType {ActualPlan} does not match cohort {CohortName} source {ExpectedPlan}",
-                    organization.Id, organization.PlanType, cohort.Name, migrationPath.FromPlan);
-                return false;
-            }
+            var scheduled = await priceIncreaseScheduler.ScheduleForSubscription(subscription);
 
-            var scheduled = await priceIncreaseScheduler.ScheduleBusinessPriceIncrease(subscription, cohort);
             if (!scheduled)
             {
                 return true;
             }
 
+            var cohort = await cohortRepository.GetByIdAsync(assignment.CohortId);
+            if (cohort?.MigrationPathId is null) return true;
+
+            var migrationPath = MigrationPaths.FromId(cohort.MigrationPathId.Value);
+            if (migrationPath is null) return true;
+
             var sourcePlan = await pricingClient.GetPlanOrThrow(migrationPath.FromPlan);
             var targetPlan = await pricingClient.GetPlanOrThrow(migrationPath.ToPlan);
 
@@ -428,6 +410,19 @@ private Task SendBusinessRenewalEmailAsync(
         return Task.CompletedTask;
     }
 
+    private async Task WaitForTestClockToAdvanceAsync(TestClock testClock)
+    {
+        while (testClock.Status != "ready")
+        {
+            await Task.Delay(TimeSpan.FromSeconds(2));
+            testClock = await stripeAdapter.GetTestClockAsync(testClock.Id);
+            if (testClock.Status == "internal_failure")
+            {
+                throw new Exception("Stripe Test Clock encountered an internal failure");
+            }
+        }
+    }
+
     #endregion
 
     #region Premium Users

src/Core/Billing/Constants/StripeConstants.cs

@@ -99,6 +99,7 @@ public static class MetadataKeys
         public const string CancelledDuringDeferredPriceIncrease = "cancelled_during_deferred_price_increase";
         public const string MigrationCohortId = "migration_cohort_id";
         public const string MigrationCohortName = "migration_cohort_name";
+        public const string CancellingUserId = "cancellingUserId";
     }
 
     public static class CancellationOrigins

src/Core/Billing/Pricing/OrganizationPriceIncreaseOptions.cs

@@ -0,0 +1,15 @@
+namespace Bit.Core.Billing.Pricing;
+
+/// <summary>
+/// Controls optional guard behavior when scheduling an organization price increase via
+/// <see cref="IPriceIncreaseScheduler.ScheduleForSubscription"/>. Guards are applied
+/// during cohort validation before any Stripe calls are made.
+/// </summary>
+public record OrganizationPriceIncreaseOptions
+{
+    /// <summary>
+    /// Skip scheduling if a price increase has already been scheduled for this
+    /// organization (i.e. <c>assignment.ScheduledDate</c> is set).
+    /// </summary>
+    public bool SkipIfAlreadyScheduled { get; init; }
+}

src/Core/Billing/Pricing/PriceIncreaseScheduler.cs

@@ -1,11 +1,13 @@
-using Bit.Core.Billing.Enums;
+using Bit.Core.AdminConsole.Entities;
+using Bit.Core.Billing.Enums;
 using Bit.Core.Billing.Extensions;
 using Bit.Core.Billing.Organizations.PlanMigration;
 using Bit.Core.Billing.Organizations.PlanMigration.Entities;
 using Bit.Core.Billing.Organizations.PlanMigration.Repositories;
 using Bit.Core.Billing.Organizations.PlanMigration.ValueObjects;
 using Bit.Core.Billing.Services;
 using Bit.Core.Billing.Subscriptions.Models;
+using Bit.Core.Repositories;
 using Bit.Core.Services;
 using Microsoft.Extensions.Logging;
 using Stripe;
@@ -45,6 +47,19 @@ public interface IPriceIncreaseScheduler
     /// <returns>True if a new schedule was created; false if skipped.</returns>
     Task<bool> ScheduleBusinessPriceIncrease(Subscription subscription, OrganizationPlanMigrationCohort cohort);
 
+    /// <summary>
+    /// Creates a deferred price-increase schedule for the given subscription,
+    /// dispatching to the correct path based on the subscription owner.
+    /// </summary>
+    /// <param name="subscription">The Stripe subscription to schedule a price increase for.</param>
+    /// <param name="options">
+    /// Optional guards applied before scheduling.
+    /// </param>
+    /// <returns>True if a new schedule was created; false if skipped.</returns>
+    Task<bool> ScheduleForSubscription(
+        Subscription subscription,
+        OrganizationPriceIncreaseOptions? options = null);
+
     /// <summary>
     /// Releases any active subscription schedule for the given subscription, cancelling a pending
     /// deferred price increase. Use when the subscription operation makes the scheduled migration
@@ -61,7 +76,9 @@ public class PriceIncreaseScheduler(
     IStripeAdapter stripeAdapter,
     IFeatureService featureService,
     IPricingClient pricingClient,
+    IOrganizationRepository organizationRepository,
     IOrganizationPlanMigrationCohortAssignmentRepository assignmentRepository,
+    IOrganizationPlanMigrationCohortRepository cohortRepository,
     ILogger<PriceIncreaseScheduler> logger) : IPriceIncreaseScheduler
 {
     public async Task<bool> SchedulePersonalPriceIncrease(Subscription subscription)
@@ -169,6 +186,33 @@ public async Task<bool> ScheduleBusinessPriceIncrease(
         return true;
     }
 
+    public async Task<bool> ScheduleForSubscription(
+        Subscription subscription,
+        OrganizationPriceIncreaseOptions? options = null)
+    {
+        try
+        {
+            SubscriberId subscriberId = subscription;
+            return await subscriberId.Match(
+                _ => SchedulePersonalPriceIncrease(subscription),
+                orgId => ScheduleForOrganizationAsync(subscription, orgId.Value, options),
+                _ =>
+                {
+                    logger.LogWarning(
+                        "Provider subscriptions do not support schedule recovery ({SubscriptionId})",
+                        subscription.Id);
+                    return Task.FromResult(false);
+                });
+        }
+        catch (Exception ex)
+        {
+            logger.LogError(ex,
+                "Failed to resolve subscriber type for subscription ({SubscriptionId}), cannot recover schedule",
+                subscription.Id);
+            return false;
+        }
+    }
+
     public async Task Release(string customerId, string subscriptionId)
     {
         try
@@ -547,4 +591,101 @@ await stripeAdapter.UpdateSubscriptionScheduleAsync(schedule.Id,
         };
     }
 
+    /// <summary>
+    /// Coordinates the full organization scheduling flow. Resolves the organization, routes
+    /// non-business plan types (personal, family, and 2019-era plans) to the personal scheduling
+    /// path, then validates cohort eligibility before scheduling.
+    /// </summary>
+    private async Task<bool> ScheduleForOrganizationAsync(
+        Subscription subscription,
+        Guid organizationId,
+        OrganizationPriceIncreaseOptions? options)
+    {
+        var organization = await organizationRepository.GetByIdAsync(organizationId);
+        if (organization is null)
+        {
+            logger.LogError(
+                "Organization ({OrganizationId}) not found; cannot recover schedule for subscription ({SubscriptionId})",
+                organizationId, subscription.Id);
+            return false;
+        }
+
+        if (!IsTrackABusinessPlanType(organization.PlanType))
+        {
+            return await SchedulePersonalPriceIncrease(subscription);
+        }
+
+        var assignment = await assignmentRepository.GetByOrganizationIdAsync(organization.Id);
+        if (assignment is null)
+        {
+            return false;
+        }
+
+        var cohort = await cohortRepository.GetByIdAsync(assignment.CohortId);
+        if (cohort is null)
+        {
+            return false;
+        }
+
+        if (!IsEligibleForScheduling(organization, assignment, cohort, options))
+        {
+            return false;
+        }
+
+        return await ScheduleBusinessPriceIncrease(subscription, cohort);
+    }
+
+    /// <summary>
+    /// Returns true if the organization is eligible for a business plan price increase.
+    /// Guards from <paramref name="options"/> are evaluated first against all resolved data;
+    /// structural eligibility checks follow. Add new option guards here —
+    /// <see cref="ScheduleForOrganizationAsync"/> never needs to change.
+    /// Does not make any database or Stripe calls.
+    /// </summary>
+    private bool IsEligibleForScheduling(
+        Organization organization,
+        OrganizationPlanMigrationCohortAssignment assignment,
+        OrganizationPlanMigrationCohort cohort,
+        OrganizationPriceIncreaseOptions? options)
+    {
+        if (options?.SkipIfAlreadyScheduled == true && assignment.ScheduledDate is not null)
+        {
+            return false;
+        }
+
+        if (!cohort.IsActive || cohort.MigrationPathId is null)
+        {
+            return false;
+        }
+
+        var migrationPath = MigrationPaths.FromId(cohort.MigrationPathId.Value);
+        if (migrationPath is null)
+        {
+            logger.LogError(
+                "Unknown MigrationPathId ({MigrationPathId}) on cohort ({CohortId}); skipping schedule for organization ({OrganizationId})",
+                cohort.MigrationPathId, cohort.Id, organization.Id);
+            return false;
+        }
+
+        if (organization.PlanType != migrationPath.FromPlan)
+        {
+            logger.LogWarning(
+                "Skipping schedule for Organization ({OrganizationId}); PlanType {ActualPlan} does not match cohort {CohortName} source {ExpectedPlan}",
+                organization.Id, organization.PlanType, cohort.Name, migrationPath.FromPlan);
+            return false;
+        }
+
+        return true;
+    }
+
+    /// <summary>
+    /// Returns true if the plan type is a Track A business plan type.
+    /// </summary>
+    /// <remarks>Expand to include additional business plan types as new tracks are added.</remarks>
+    private static bool IsTrackABusinessPlanType(PlanType planType) => planType is
+        PlanType.TeamsMonthly2020 or
+        PlanType.TeamsAnnually2020 or
+        PlanType.EnterpriseMonthly2020 or
+        PlanType.EnterpriseAnnually2020;
+
 }