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

Author: jankeromnes

components/dashboard/src/start/CreateWorkspace.tsx

@@ -185,6 +185,10 @@ export default class CreateWorkspace extends React.Component<CreateWorkspaceProp
                     phase = StartPhase.Stopped;
                     statusMessage = <LimitReachedOutOfHours />;
                     break;
+                case ErrorCodes.INVALID_COST_CENTER:
+                    phase = StartPhase.Stopped;
+                    statusMessage = <p>Invalid cost center: {String(error)}</p>;
+                    break;
                 default:
                     statusMessage = (
                         <p className="text-base text-gitpod-red w-96">

components/gitpod-protocol/src/messaging/error.ts

@@ -32,15 +32,18 @@ export namespace ErrorCodes {
     // 430 Repository not whitelisted (custom status code)
     export const REPOSITORY_NOT_WHITELISTED = 430;
 
+    // 450 Payment error
+    export const PAYMENT_ERROR = 450;
+
+    // 455 Invalid cost center (custom status code)
+    export const INVALID_COST_CENTER = 455;
+
     // 460 Context Parse Error (custom status code)
     export const CONTEXT_PARSE_ERROR = 460;
 
-    // 461 Invalid gitpod yml
+    // 461 Invalid gitpod yml (custom status code)
     export const INVALID_GITPOD_YML = 461;
 
-    // 450 Payment error
-    export const PAYMENT_ERROR = 450;
-
     // 470 User Blocked (custom status code)
     export const USER_BLOCKED = 470;
 

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";
@@ -29,6 +30,8 @@ import { EmailAddressAlreadyTakenException, SelectAccountException } from "../au
 import { SelectAccountPayload } from "@gitpod/gitpod-protocol/lib/auth";
 import { AttributionId } from "@gitpod/gitpod-protocol/lib/attribution";
 import { StripeService } from "../../ee/src/user/stripe-service";
+import { ResponseError } from "vscode-ws-jsonrpc";
+import { ErrorCodes } from "@gitpod/gitpod-protocol/lib/messaging/error";
 
 export interface FindUserByIdentityStrResult {
     user: User;
@@ -211,6 +214,66 @@ 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 ResponseError(
+                    ErrorCodes.INVALID_COST_CENTER,
+                    "The billing team you've selected no longer exists.",
+                );
+            }
+            const members = await this.teamDB.findMembersByTeam(team.id);
+            if (!members.find((m) => m.userId === user.id)) {
+                throw new ResponseError(
+                    ErrorCodes.INVALID_COST_CENTER,
+                    "You're no longer a member of the selected billing team.",
+                );
+            }
+            const subscriptionId = await this.findTeamUsageBasedSubscriptionId(team);
+            if (!subscriptionId) {
+                throw new ResponseError(
+                    ErrorCodes.INVALID_COST_CENTER,
+                    "The billing team you've selected has no active 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 ResponseError(ErrorCodes.INVALID_COST_CENTER, "Multiple teams have billing enabled.");
+        }
+        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 +292,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