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

Author: jankeromnes

components/dashboard/src/components/BillingAccountSelector.tsx

@@ -0,0 +1,82 @@
+/**
+ * Copyright (c) 2022 Gitpod GmbH. All rights reserved.
+ * Licensed under the GNU Affero General Public License (AGPL).
+ * See License-AGPL.txt in the project root for license information.
+ */
+
+import { Team } from "@gitpod/gitpod-protocol";
+import { useContext, useEffect, useState } from "react";
+import { Link } from "react-router-dom";
+import { getGitpodService } from "../service/service";
+import { TeamsContext } from "../teams/teams-context";
+import { UserContext } from "../user-context";
+import DropDown from "../components/DropDown";
+import { ReactComponent as Spinner } from "../icons/Spinner.svg";
+
+export function BillingAccountSelector() {
+    const { user } = useContext(UserContext);
+    const { teams } = useContext(TeamsContext);
+    const [teamsWithBillingEnabled, setTeamsWithBillingEnabled] = useState<Team[] | undefined>();
+
+    useEffect(() => {
+        if (!teams) {
+            setTeamsWithBillingEnabled(undefined);
+            return;
+        }
+        const teamsWithBilling: Team[] = [];
+        Promise.all(
+            teams.map(async (t) => {
+                const subscriptionId = await getGitpodService().server.findStripeSubscriptionIdForTeam(t.id);
+                if (subscriptionId) {
+                    teamsWithBilling.push(t);
+                }
+            }),
+        ).then(() => setTeamsWithBillingEnabled(teamsWithBilling));
+    }, [teams]);
+
+    const setUsageAttributionTeam = async (team?: Team) => {
+        if (!user) {
+            return;
+        }
+        const usageAttributionId = team ? `team:${team.id}` : `user:${user.id}`;
+        await getGitpodService().server.setUsageAttribution(usageAttributionId);
+    };
+    return (
+        <>
+            {teamsWithBillingEnabled === undefined && <Spinner className="m-2 h-5 w-5 animate-spin" />}
+            {teamsWithBillingEnabled && teamsWithBillingEnabled.length === 0 && (
+                <div className="flex space-x-2">
+                    <span>
+                        <Link className="gp-link" to="/teams/new">
+                            Create a team
+                        </Link>{" "}
+                        to set up usage-based billing.
+                    </span>
+                </div>
+            )}
+            {teamsWithBillingEnabled && teamsWithBillingEnabled.length > 0 && (
+                <div className="flex space-x-2">
+                    <span>Bill all my usage to:</span>
+                    <DropDown
+                        activeEntry={
+                            teamsWithBillingEnabled.find((t) => `team:${t.id}` === user?.usageAttributionId)?.name
+                        }
+                        customClasses="w-32"
+                        renderAsLink={true}
+                        entries={[
+                            {
+                                title: "(myself)",
+                                onClick: () => setUsageAttributionTeam(undefined),
+                            },
+                        ].concat(
+                            teamsWithBillingEnabled.map((t) => ({
+                                title: t.name,
+                                onClick: () => setUsageAttributionTeam(t),
+                            })),
+                        )}
+                    />
+                </div>
+            )}
+        </>
+    );
+}

components/dashboard/src/settings/Billing.tsx

@@ -4,85 +4,17 @@
  * See License-AGPL.txt in the project root for license information.
  */
 
-import { Team } from "@gitpod/gitpod-protocol";
-import { useContext, useEffect, useState } from "react";
-import { Link } from "react-router-dom";
-import { ReactComponent as Spinner } from "../icons/Spinner.svg";
-import DropDown from "../components/DropDown";
-import { getGitpodService } from "../service/service";
-import { TeamsContext } from "../teams/teams-context";
-import { UserContext } from "../user-context";
 import { PageWithSettingsSubMenu } from "./PageWithSettingsSubMenu";
+import { BillingAccountSelector } from "../components/BillingAccountSelector";
 
 export default function Billing() {
-    const { user } = useContext(UserContext);
-    const { teams } = useContext(TeamsContext);
-    const [teamsWithBillingEnabled, setTeamsWithBillingEnabled] = useState<Team[] | undefined>();
-
-    useEffect(() => {
-        if (!teams) {
-            setTeamsWithBillingEnabled(undefined);
-            return;
-        }
-        const teamsWithBilling: Team[] = [];
-        Promise.all(
-            teams.map(async (t) => {
-                const subscriptionId = await getGitpodService().server.findStripeSubscriptionIdForTeam(t.id);
-                if (subscriptionId) {
-                    teamsWithBilling.push(t);
-                }
-            }),
-        ).then(() => setTeamsWithBillingEnabled(teamsWithBilling));
-    }, [teams]);
-
-    const setUsageAttributionTeam = async (team?: Team) => {
-        if (!user) {
-            return;
-        }
-        const usageAttributionId = team ? `team:${team.id}` : `user:${user.id}`;
-        await getGitpodService().server.setUsageAttribution(usageAttributionId);
-    };
-
     return (
         <PageWithSettingsSubMenu title="Billing" subtitle="Usage-Based Billing.">
             <h3>Usage-Based Billing</h3>
             <h2 className="text-gray-500">Manage usage-based billing, spending limit, and payment method.</h2>
             <div className="mt-8">
                 <h3>Billing Account</h3>
-                {teamsWithBillingEnabled === undefined && <Spinner className="m-2 h-5 w-5 animate-spin" />}
-                {teamsWithBillingEnabled && teamsWithBillingEnabled.length === 0 && (
-                    <div className="flex space-x-2">
-                        <span>
-                            <Link className="gp-link" to="/teams/new">
-                                Create a team
-                            </Link>{" "}
-                            to set up usage-based billing.
-                        </span>
-                    </div>
-                )}
-                {teamsWithBillingEnabled && teamsWithBillingEnabled.length > 0 && (
-                    <div className="flex space-x-2">
-                        <span>Bill all my usage to:</span>
-                        <DropDown
-                            activeEntry={
-                                teamsWithBillingEnabled.find((t) => `team:${t.id}` === user?.usageAttributionId)?.name
-                            }
-                            customClasses="w-32"
-                            renderAsLink={true}
-                            entries={[
-                                {
-                                    title: "(myself)",
-                                    onClick: () => setUsageAttributionTeam(undefined),
-                                },
-                            ].concat(
-                                teamsWithBillingEnabled.map((t) => ({
-                                    title: t.name,
-                                    onClick: () => setUsageAttributionTeam(t),
-                                })),
-                            )}
-                        />
-                    </div>
-                )}
+                <BillingAccountSelector />
             </div>
         </PageWithSettingsSubMenu>
     );

components/dashboard/src/start/CreateWorkspace.tsx

@@ -25,6 +25,7 @@ import PrebuildLogs from "../components/PrebuildLogs";
 import CodeText from "../components/CodeText";
 import FeedbackComponent from "../feedback-form/FeedbackComponent";
 import { isGitpodIo } from "../utils";
+import { BillingAccountSelector } from "../components/BillingAccountSelector";
 
 export interface CreateWorkspaceProps {
     contextUrl: string;
@@ -185,6 +186,12 @@ export default class CreateWorkspace extends React.Component<CreateWorkspaceProp
                     phase = StartPhase.Stopped;
                     statusMessage = <LimitReachedOutOfHours />;
                     break;
+                case ErrorCodes.INVALID_COST_CENTER:
+                    // HACK: Hide the error (behind the modal)
+                    error = undefined;
+                    phase = StartPhase.Stopped;
+                    statusMessage = <SelectCostCenterModal />;
+                    break;
                 default:
                     statusMessage = (
                         <p className="text-base text-gitpod-red w-96">
@@ -290,6 +297,15 @@ export default class CreateWorkspace extends React.Component<CreateWorkspaceProp
     }
 }
 
+function SelectCostCenterModal() {
+    return (
+        <Modal visible={true} closeable={false} onClose={() => {}}>
+            <h3>Choose Billing Team</h3>
+            <BillingAccountSelector />
+        </Modal>
+    );
+}
+
 function LimitReachedModal(p: { children: React.ReactNode }) {
     const { user } = useContext(UserContext);
     return (

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