[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,45 +357,38 @@ 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)
+            if (subscription.TestClock != null)
             {
-                return false;
+                await WaitForTestClockToAdvanceAsync(subscription.TestClock);
             }
 
-            if (cohort.MigrationPathId is null)
-            {
-                // Churn-only cohort — no migration to schedule.
-                return false;
-            }
+            var migrationScheduled = await priceIncreaseScheduler.ScheduleForSubscription(subscription);
 
-            var migrationPath = MigrationPaths.FromId(cohort.MigrationPathId.Value);
-            if (migrationPath is null)
+            if (!migrationScheduled)
             {
-                logger.LogError(
-                    "Unknown MigrationPathId ({MigrationPathId}) on cohort ({CohortId}) for Organization ({OrganizationId})",
-                    cohort.MigrationPathId, cohort.Id, organization.Id);
                 return false;
             }
 
-            if (organization.PlanType != migrationPath.FromPlan)
+            var cohort = await cohortRepository.GetByIdAsync(assignment.CohortId);
+            if (cohort?.MigrationPathId is null)
             {
                 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;
+                    "Cohort ({CohortId}) missing or has no MigrationPathId; skipping renewal email for Organization ({OrganizationId})",
+                    assignment.CohortId, organization.Id);
+                return true;
             }
 
-            var scheduled = await priceIncreaseScheduler.ScheduleBusinessPriceIncrease(subscription, cohort);
-            if (!scheduled)
+            var migrationPath = MigrationPaths.FromId(cohort.MigrationPathId.Value);
+            if (migrationPath is null)
             {
+                logger.LogWarning(
+                    "Unknown MigrationPathId ({MigrationPathId}) on cohort ({CohortId}); skipping renewal email for Organization ({OrganizationId})",
+                    cohort.MigrationPathId, cohort.Id, organization.Id);
                 return true;
             }
 
@@ -428,6 +422,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 => DispatchOrganizationScheduleAsync(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,90 @@ 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> DispatchOrganizationScheduleAsync(
+        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 (organization.PlanType.GetProductTier() is not (ProductTierType.Teams or ProductTierType.Enterprise))
+        {
+            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.
+    /// </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;
+    }
+
+
 }

src/Core/Billing/Services/Implementations/SubscriberService.cs

@@ -67,7 +67,7 @@ public async Task CancelSubscription(
         ];
 
         // Build once from survey — null when survey is absent (system-initiated cancellation)
-        var cancellationDetails = offboardingSurveyResponse != null
+        var cancellationDetails = offboardingSurveyResponse is not null
             ? new SubscriptionCancellationDetailsOptions
             {
                 Comment = offboardingSurveyResponse.Feedback,
@@ -77,21 +77,16 @@ public async Task CancelSubscription(
             }
             : null;
 
-        var cancellingUserMetadata = offboardingSurveyResponse != null
+        var cancellingUserMetadata = offboardingSurveyResponse is not null
             ? new Dictionary<string, string>
             {
-                { "cancellingUserId", offboardingSurveyResponse.UserId.ToString() }
+                { MetadataKeys.CancellingUserId, offboardingSurveyResponse.UserId.ToString() }
             }
             : null;
 
-        if (cancelImmediately)
-        {
-            await CancelSubscriptionImmediatelyAsync(subscription, cancellationDetails, cancellingUserMetadata);
-        }
-        else
-        {
-            await CancelSubscriptionAtPeriodEndAsync(subscription, cancellationDetails, cancellingUserMetadata);
-        }
+        await (cancelImmediately
+            ? CancelSubscriptionImmediatelyAsync(subscription, cancellationDetails, cancellingUserMetadata)
+            : CancelSubscriptionAtPeriodEndAsync(subscription, cancellationDetails, cancellingUserMetadata));
     }
 
     public async Task<string> CreateBraintreeCustomer(
@@ -233,7 +228,8 @@ private async Task CancelSubscriptionImmediatelyAsync(
         SubscriptionCancellationDetailsOptions? cancellationDetails,
         Dictionary<string, string>? cancellingUserMetadata)
     {
-        if (featureService.IsEnabled(FeatureFlagKeys.PM32645_DeferPriceMigrationToRenewal))
+        if (featureService.IsEnabled(FeatureFlagKeys.PM32645_DeferPriceMigrationToRenewal) ||
+            featureService.IsEnabled(FeatureFlagKeys.PM35215_BusinessPlanPriceMigration))
         {
             var activeSchedule = await GetActiveScheduleAsync(subscription);
             if (activeSchedule != null)
@@ -268,7 +264,8 @@ private async Task CancelSubscriptionAtPeriodEndAsync(
             Metadata = cancellingUserMetadata
         };
 
-        if (featureService.IsEnabled(FeatureFlagKeys.PM32645_DeferPriceMigrationToRenewal))
+        if (featureService.IsEnabled(FeatureFlagKeys.PM32645_DeferPriceMigrationToRenewal) ||
+            featureService.IsEnabled(FeatureFlagKeys.PM35215_BusinessPlanPriceMigration))
         {
             var activeSchedule = await GetActiveScheduleAsync(subscription);
 
@@ -557,7 +554,8 @@ await braintreeGateway.Customer.UpdateAsync(
 
     public async Task ResumeFromUnpaidCancellationAsync(ISubscriber subscriber)
     {
-        var subscription = await GetSubscription(subscriber);
+        var subscription = await GetSubscription(subscriber,
+            new SubscriptionGetOptions { Expand = ["customer.discount", "discounts"] });
 
         if (subscription is null ||
             subscription.Status != SubscriptionStatus.Unpaid ||
@@ -578,6 +576,8 @@ subscription.Metadata is null ||
             }
         });
 
+        await priceIncreaseScheduler.ScheduleForSubscription(subscription);
+
         logger.LogInformation(
             "Cleared pending unpaid-lifecycle cancellation for subscription ({SubscriptionId}) after subscriber re-enable",
             subscription.Id);
@@ -602,6 +602,8 @@ public async Task ScheduleUnpaidCancellationAsync(ISubscriber subscriber)
 
         var now = subscription.TestClock?.FrozenTime ?? DateTime.UtcNow;
 
+        await priceIncreaseScheduler.Release(subscription.CustomerId, subscription.Id);
+
         await stripeAdapter.UpdateSubscriptionAsync(subscription.Id, new SubscriptionUpdateOptions
         {
             CancelAt = now.AddDays(7),