[server] When starting a workspace but usage attribution is unclear, prompt for explicit user choice

Author: jankeromnes

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

@@ -2008,20 +2008,6 @@ export class GitpodServerEEImpl extends GitpodServerImpl {
                 id: attributionId,
                 spendingLimit: this.defaultSpendingLimit,
             });
-
-            // 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.usageAttributionId) {
-                        u.usageAttributionId = attributionId;
-                        await this.userDB.storeUser(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`);

components/server/src/user/user-service.ts

@@ -15,6 +15,7 @@ import {
     WORKSPACE_TIMEOUT_DEFAULT_LONG,
     WORKSPACE_TIMEOUT_EXTENDED,
     WORKSPACE_TIMEOUT_EXTENDED_ALT,
+    Team,
 } from "@gitpod/gitpod-protocol";
 import { CostCenterDB, ProjectDB, TeamDB, TermsAcceptanceDB, UserDB } from "@gitpod/gitpod-db/lib";
 import { HostContextProvider } from "../auth/host-context-provider";
@@ -211,6 +212,57 @@ export class UserService {
         return false;
     }
 
+    protected async findTeamUsageBasedSubscriptionId(team: Team): Promise<string | undefined> {
+        const customer = await this.stripeService.findCustomerByTeamId(team.id);
+        if (!customer) {
+            return;
+        }
+        const subscription = await this.stripeService.findUncancelledSubscriptionByCustomer(customer.id);
+        return subscription?.id;
+    }
+
+    protected async validateUsageAttributionId(user: User, usageAttributionId: string): Promise<void> {
+        const attribution = AttributionId.parse(usageAttributionId);
+        if (attribution?.kind === "team") {
+            const team = await this.teamDB.findTeamById(attribution.teamId);
+            if (!team) {
+                throw new Error("Selected team does not exist!");
+            }
+            const members = await this.teamDB.findMembersByTeam(team.id);
+            if (!members.find((m) => m.userId === user.id)) {
+                throw new Error("User is not part of selected team!");
+            }
+            const subscriptionId = await this.findTeamUsageBasedSubscriptionId(team);
+            if (!subscriptionId) {
+                throw new Error("Selected team has no subscription!");
+            }
+        }
+    }
+
+    protected async findSingleTeamWithUsageBasedBilling(user: User): Promise<Team | undefined> {
+        // Find all the user's teams with usage-based billing enabled.
+        const teams = await this.teamDB.findTeamsByUser(user.id);
+        const teamsWithBilling: Team[] = [];
+        await Promise.all(
+            teams.map(async (team) => {
+                const subscriptionId = await this.findTeamUsageBasedSubscriptionId(team);
+                if (subscriptionId) {
+                    teamsWithBilling.push(team);
+                }
+            }),
+        );
+        if (teamsWithBilling.length > 1) {
+            // Multiple teams with usage-based billing enabled -- ask the user to make an explicit choice.
+            throw new Error("Multiple billing teams! Please choose one");
+        }
+        if (teamsWithBilling.length === 1) {
+            // Single team with usage-based billing enabled -- attribute all usage to it.
+            return teamsWithBilling[0];
+        }
+        // No team with usage-based billing enabled.
+        return undefined;
+    }
+
     /**
      * Identifies the team or user to which a workspace instance's running time should be attributed to
      * (e.g. for usage analytics or billing purposes).
@@ -229,12 +281,18 @@ export class UserService {
     async getWorkspaceUsageAttributionId(user: User, projectId?: string): Promise<string | undefined> {
         // A. Billing-based attribution
         if (this.config.enablePayment) {
-            if (!user.usageAttributionId) {
-                // No explicit user attribution ID yet -- attribute all usage to the user by default (regardless of project/team).
-                return AttributionId.render({ kind: "user", userId: user.id });
+            if (user.usageAttributionId) {
+                await this.validateUsageAttributionId(user, user.usageAttributionId);
+                // Return the user's explicit attribution ID.
+                return user.usageAttributionId;
             }
-            // Return the user's explicit attribution ID.
-            return user.usageAttributionId;
+            const billingTeam = await this.findSingleTeamWithUsageBasedBilling(user);
+            if (billingTeam) {
+                // Single team with usage-based billing enabled -- attribute all usage to it.
+                return AttributionId.render({ kind: "team", teamId: billingTeam.id });
+            }
+            // Attribute all usage to the user by default (regardless of project/team).
+            return AttributionId.render({ kind: "user", userId: user.id });
         }
 
         // B. Project-based attribution