Automatically set a user's usageAttributionId in cases where there is no ambiguity

components/server/ee/src/workspace/gitpod-server-impl.ts

@@ -1428,29 +1428,44 @@ export class GitpodServerEEImpl extends GitpodServerImpl {
 
     protected async onTeamMemberAdded(userId: string, teamId: string): Promise<void> {
         const now = new Date();
-        const teamSubscription = await this.teamSubscription2DB.findForTeam(teamId, now.toISOString());
-        if (!teamSubscription) {
-            // No team subscription, nothing to do 🌴
-            return;
+        const ts2 = await this.teamSubscription2DB.findForTeam(teamId, now.toISOString());
+        if (ts2) {
+            await this.updateTeamSubscriptionQuantity(ts2);
+            await this.teamSubscription2Service.addTeamMemberSubscription(ts2, userId);
+        }
+        const [teamCustomer, user] = await Promise.all([
+            this.stripeService.findCustomerByTeamId(teamId),
+            this.userDB.findUserById(userId),
+        ]);
+        if (teamCustomer && user && !user.additionalData?.usageAttributionId) {
+            // If the user didn't explicitly choose yet where their usage should be attributed to, and
+            // they join a team which accepts usage attribution (i.e. with usage-based billing enabled),
+            // then we simplify the UX by automatically attributing the user's usage to that team.
+            // Note: This default choice can be changed at any time by the user in their billing settings.
+            const subscription = await this.stripeService.findUncancelledSubscriptionByCustomer(teamCustomer.id);
+            if (subscription) {
+                user.additionalData = user.additionalData || {};
+                user.additionalData.usageAttributionId = `team:${teamId}`;
+                await this.userDB.updateUserPartial(user);
+            }
         }
-        await this.updateTeamSubscriptionQuantity(teamSubscription);
-        await this.teamSubscription2Service.addTeamMemberSubscription(teamSubscription, userId);
     }
 
     protected async onTeamMemberRemoved(userId: string, teamId: string, teamMembershipId: string): Promise<void> {
         const now = new Date();
-        const teamSubscription = await this.teamSubscription2DB.findForTeam(teamId, now.toISOString());
-        if (!teamSubscription) {
-            // No team subscription, nothing to do 🌴
-            return;
+        const ts2 = await this.teamSubscription2DB.findForTeam(teamId, now.toISOString());
+        if (ts2) {
+            await this.updateTeamSubscriptionQuantity(ts2);
+            await this.teamSubscription2Service.cancelTeamMemberSubscription(ts2, userId, teamMembershipId, now);
+        }
+        const user = await this.userDB.findUserById(userId);
+        if (user && user.additionalData?.usageAttributionId === `team:${teamId}`) {
+            // If the user previously attributed all their usage to a given team, but they are now leaving this
+            // team, then the currently selected usage attribution ID is no longer valid. In this case, we must
+            // reset this ID to the default value.
+            user.additionalData.usageAttributionId = undefined;
+            await this.userDB.updateUserPartial(user);
         }
-        await this.updateTeamSubscriptionQuantity(teamSubscription);
-        await this.teamSubscription2Service.cancelTeamMemberSubscription(
-            teamSubscription,
-            userId,
-            teamMembershipId,
-            now,
-        );
     }
 
     protected async onTeamDeleted(teamId: string): Promise<void> {
@@ -1945,6 +1960,20 @@ export class GitpodServerEEImpl extends GitpodServerImpl {
                 customer = await this.stripeService.createCustomerForTeam(user, team!, setupIntentId);
             }
             await this.stripeService.createSubscriptionForCustomer(customer.id, currency);
+            // For all team members that didn't explicitly choose yet where their usage should be attributed to,
+            // we simplify the UX by automatically attributing their usage to this recently-upgraded team.
+            // Note: This default choice can be changed at any time by members in their personal billing settings.
+            const members = await this.teamDB.findMembersByTeam(teamId);
+            await Promise.all(
+                members.map(async (m) => {
+                    const u = await this.userDB.findUserById(m.userId);
+                    if (u && !u.additionalData?.usageAttributionId) {
+                        u.additionalData = u.additionalData || {};
+                        u.additionalData.usageAttributionId = `team:${teamId}`;
+                        await this.userDB.updateUserPartial(u);
+                    }
+                }),
+            );
         } catch (error) {
             log.error(`Failed to subscribe team '${teamId}' to Stripe`, error);
             throw new ResponseError(ErrorCodes.INTERNAL_SERVER_ERROR, `Failed to subscribe team '${teamId}' to Stripe`);