From 8ce30c829965961ca04d1fdac4da2f6698f83437 Mon Sep 17 00:00:00 2001
From: Taras Yemets <yemets.taras@gmail.com>
Date: Thu, 20 Jun 2024 23:08:38 +0300
Subject: [PATCH 1/8] feat(plg): remove seats from subscription

---
 client/web/src/cody/invites/InviteUsers.tsx | 43 +++++++++++++++++----
 1 file changed, 36 insertions(+), 7 deletions(-)

diff --git a/client/web/src/cody/invites/InviteUsers.tsx b/client/web/src/cody/invites/InviteUsers.tsx
index e3cf2c5e5dc1..b609d9c51896 100644
--- a/client/web/src/cody/invites/InviteUsers.tsx
+++ b/client/web/src/cody/invites/InviteUsers.tsx
@@ -2,13 +2,13 @@ import React, { useState, useCallback, useMemo } from 'react'
 
 import { pluralize } from '@sourcegraph/common'
 import type { TelemetryV2Props } from '@sourcegraph/shared/src/telemetry'
-import { ButtonLink, H2, Link, Text, H3, TextArea } from '@sourcegraph/wildcard'
+import { H2, Link, Text, H3, TextArea, Button, H1 } from '@sourcegraph/wildcard'
 
 import { CodyAlert } from '../components/CodyAlert'
 import { CodyContainer } from '../components/CodyContainer'
 import { CodyProBadgeDeck } from '../components/CodyProBadgeDeck'
 import { useSendInvite, useTeamInvites } from '../management/api/react-query/invites'
-import { useCurrentSubscription } from '../management/api/react-query/subscriptions'
+import { useCurrentSubscription, useUpdateCurrentSubscription } from '../management/api/react-query/subscriptions'
 import { useTeamMembers } from '../management/api/react-query/teams'
 import type { SubscriptionSummary } from '../management/api/teamSubscriptions'
 import { isValidEmailAddress } from '../util'
@@ -39,6 +39,7 @@ export const InviteUsers: React.FunctionComponent<InviteUsersProps> = ({ telemet
     const [emailAddressErrorMessage, setEmailAddressErrorMessage] = useState<string | null>(null)
 
     const sendInviteMutation = useSendInvite()
+    const updateSubscriptionMutation = useUpdateCurrentSubscription()
 
     const verifyEmailList = useCallback((): Error | void => {
         if (emailAddresses.length === 0) {
@@ -104,12 +105,22 @@ export const InviteUsers: React.FunctionComponent<InviteUsersProps> = ({ telemet
         })
     }, [emailAddresses, sendInviteMutation.mutateAsync, teamId, telemetryRecorder, verifyEmailList])
 
-    if (!isAdmin || !remainingInviteCount) {
+    if (!isAdmin || !remainingInviteCount || !subscriptionQueryResult.data) {
         return null
     }
 
+    const { maxSeats } = subscriptionQueryResult.data
+
     return (
         <>
+            {updateSubscriptionMutation.isSuccess && (
+                <CodyAlert variant="greenSuccess">
+                    <H1 as="p" className="mb-2">
+                        Remaining invites removed from plan
+                    </H1>
+                    <Text className="mb-0">You can add more seats at any time with the "Add seats" button.</Text>
+                </CodyAlert>
+            )}
             {sendInviteMutation.status === 'success' && (
                 <CodyAlert variant="greenSuccess">
                     <H3>
@@ -153,15 +164,33 @@ export const InviteUsers: React.FunctionComponent<InviteUsersProps> = ({ telemet
                             isValid={emailAddressErrorMessage ? false : undefined}
                         />
                         {emailAddressErrorMessage ? (
-                            <Text className="text-danger mb-2">{emailAddressErrorMessage}</Text>
+                            <Text className="text-danger">{emailAddressErrorMessage}</Text>
                         ) : (
-                            <Text className="text-muted mb-2">Enter email addresses separated by a comma.</Text>
+                            <Text className="text-muted">Enter email addresses separated by a comma.</Text>
                         )}
 
                         <div>
-                            <ButtonLink variant="success" size="sm" onSelect={onSendInvitesClicked}>
+                            <Button
+                                disabled={updateSubscriptionMutation.isPending || sendInviteMutation.isPending}
+                                variant="success"
+                                onClick={onSendInvitesClicked}
+                                className="mr-2"
+                            >
                                 Send
-                            </ButtonLink>
+                            </Button>
+                            <Button
+                                variant="link"
+                                disabled={updateSubscriptionMutation.isPending || sendInviteMutation.isPending}
+                                onClick={() =>
+                                    updateSubscriptionMutation.mutate({
+                                        subscriptionUpdate: {
+                                            newSeatCount: maxSeats - remainingInviteCount,
+                                        },
+                                    })
+                                }
+                            >
+                                Remove invites from plan
+                            </Button>
                         </div>
                     </div>
                 </div>

From bae2776719b1992cb8c7b2f99de6ddadd93be635 Mon Sep 17 00:00:00 2001
From: Taras Yemets <yemets.taras@gmail.com>
Date: Fri, 21 Jun 2024 01:10:10 +0300
Subject: [PATCH 2/8] chore(plg): provide user cody plan on page level

---
 client/web/src/cody/codyProRoutes.tsx         | 24 +++++++++-
 .../cody/management/CodyManagementPage.tsx    | 37 ++++-----------
 .../manage/CodySubscriptionManagePage.tsx     | 32 +++----------
 .../new/NewCodyProSubscriptionPage.tsx        | 46 +++++--------------
 .../subscription/CodySubscriptionPage.tsx     | 33 ++++---------
 .../subscription/useUserCodySubscription.ts   | 29 ++++++++++++
 .../web/src/cody/team/CodyManageTeamPage.tsx  |  2 +
 7 files changed, 90 insertions(+), 113 deletions(-)
 create mode 100644 client/web/src/cody/subscription/useUserCodySubscription.ts

diff --git a/client/web/src/cody/codyProRoutes.tsx b/client/web/src/cody/codyProRoutes.tsx
index c9a3ebfadf1f..60240e01d95c 100644
--- a/client/web/src/cody/codyProRoutes.tsx
+++ b/client/web/src/cody/codyProRoutes.tsx
@@ -1,9 +1,10 @@
-import type { RouteObject } from 'react-router-dom'
+import { Navigate, type RouteObject } from 'react-router-dom'
 
 import { lazyComponent } from '@sourcegraph/shared/src/util/lazyComponent'
 
 import { type LegacyLayoutRouteContext, LegacyRoute } from '../LegacyRouteContext'
 
+import { useUserCodySubscription } from './subscription/useUserCodySubscription'
 import { isEmbeddedCodyProUIEnabled } from './util'
 
 export enum CodyProRoutes {
@@ -76,6 +77,25 @@ interface CodyProPageProps extends Pick<LegacyLayoutRouteContext, 'authenticated
  * only applies to non-Enterprise users) from the rest of the Sourcegraph UI.
  */
 const CodyProPage: React.FC<CodyProPageProps> = props => {
+    const { data } = useUserCodySubscription()
+    if (!data) {
+        return null
+    }
+
+    if (!data.currentUser) {
+        return <Navigate to={`/sign-in?returnTo=${CodyProRoutes.Manage}`} replace={true} />
+    }
+
+    if (!data.currentUser.codySubscription) {
+        return null
+    }
+
     const Component = routeComponents[props.path]
-    return <Component authenticatedUser={props.authenticatedUser} telemetryRecorder={props.telemetryRecorder} />
+    return (
+        <Component
+            authenticatedUser={props.authenticatedUser}
+            telemetryRecorder={props.telemetryRecorder}
+            codySubscription={data.currentUser.codySubscription}
+        />
+    )
 }
diff --git a/client/web/src/cody/management/CodyManagementPage.tsx b/client/web/src/cody/management/CodyManagementPage.tsx
index 311f5ee64fba..a11ed2a48f0d 100644
--- a/client/web/src/cody/management/CodyManagementPage.tsx
+++ b/client/web/src/cody/management/CodyManagementPage.tsx
@@ -11,13 +11,7 @@ import { ButtonLink, H1, H2, Icon, Link, PageHeader, Text, useSearchParameters,
 import type { AuthenticatedUser } from '../../auth'
 import { Page } from '../../components/Page'
 import { PageTitle } from '../../components/PageTitle'
-import {
-    type UserCodyPlanResult,
-    type UserCodyPlanVariables,
-    type UserCodyUsageResult,
-    type UserCodyUsageVariables,
-    CodySubscriptionPlan,
-} from '../../graphql-operations'
+import { type UserCodyUsageResult, type UserCodyUsageVariables, CodySubscriptionPlan } from '../../graphql-operations'
 import { CodyProRoutes } from '../codyProRoutes'
 import { CodyAlert } from '../components/CodyAlert'
 import { ProIcon } from '../components/CodyIcon'
@@ -26,7 +20,8 @@ import { AcceptInviteBanner } from '../invites/AcceptInviteBanner'
 import { InviteUsers } from '../invites/InviteUsers'
 import { isCodyEnabled } from '../isCodyEnabled'
 import { CodyOnboarding, type IEditor } from '../onboarding/CodyOnboarding'
-import { USER_CODY_PLAN, USER_CODY_USAGE } from '../subscription/queries'
+import { USER_CODY_USAGE } from '../subscription/queries'
+import type { UserCodySubscription } from '../subscription/useUserCodySubscription'
 import { getManageSubscriptionPageURL } from '../util'
 
 import { useSubscriptionSummary } from './api/react-query/subscriptions'
@@ -37,6 +32,7 @@ import styles from './CodyManagementPage.module.scss'
 
 interface CodyManagementPageProps extends TelemetryV2Props {
     authenticatedUser: AuthenticatedUser | null
+    codySubscription: UserCodySubscription
 }
 
 export enum EditorStep {
@@ -47,6 +43,7 @@ export enum EditorStep {
 export const CodyManagementPage: React.FunctionComponent<CodyManagementPageProps> = ({
     authenticatedUser,
     telemetryRecorder,
+    codySubscription,
 }) => {
     const navigate = useNavigate()
     const parameters = useSearchParameters()
@@ -69,8 +66,6 @@ export const CodyManagementPage: React.FunctionComponent<CodyManagementPageProps
 
     const welcomeToPro = parameters.get('welcome') === '1'
 
-    const { data, error: dataError, refetch } = useQuery<UserCodyPlanResult, UserCodyPlanVariables>(USER_CODY_PLAN, {})
-
     const { data: usageData, error: usageDateError } = useQuery<UserCodyUsageResult, UserCodyUsageVariables>(
         USER_CODY_USAGE,
         {}
@@ -82,14 +77,6 @@ export const CodyManagementPage: React.FunctionComponent<CodyManagementPageProps
     const [selectedEditor, setSelectedEditor] = React.useState<IEditor | null>(null)
     const [selectedEditorStep, setSelectedEditorStep] = React.useState<EditorStep | null>(null)
 
-    const subscription = data?.currentUser?.codySubscription
-
-    useEffect(() => {
-        if (!!data && !data?.currentUser) {
-            navigate(`/sign-in?returnTo=${CodyProRoutes.Manage}`)
-        }
-    }, [data, navigate])
-
     const getTeamInviteButton = (): JSX.Element | null => {
         const isSoloUser = subscriptionSummaryQueryResult?.data?.teamMaxMembers === 1
         const hasFreeSeats = subscriptionSummaryQueryResult?.data
@@ -114,19 +101,15 @@ export const CodyManagementPage: React.FunctionComponent<CodyManagementPageProps
         telemetryRecorder.recordEvent('cody.management.upgradeToProCTA', 'click')
     }, [telemetryRecorder])
 
-    if (accountSwitchRequired) {
+    if (accountSwitchRequired || !isCodyEnabled()) {
         return null
     }
 
-    if (dataError || usageDateError) {
-        throw dataError || usageDateError
-    }
-
-    if (!isCodyEnabled() || !subscription) {
-        return null
+    if (usageDateError) {
+        throw usageDateError
     }
 
-    const isUserOnProTier = subscription.plan === CodySubscriptionPlan.PRO
+    const isUserOnProTier = codySubscription.plan === CodySubscriptionPlan.PRO
 
     return (
         <>
@@ -191,7 +174,7 @@ export const CodyManagementPage: React.FunctionComponent<CodyManagementPageProps
                             </div>
                         )}
                     </div>
-                    <SubscriptionStats {...{ subscription, usageData }} />
+                    <SubscriptionStats subscription={codySubscription} usageData={usageData} />
                 </div>
 
                 <UseCodyInEditorSection
diff --git a/client/web/src/cody/management/subscription/manage/CodySubscriptionManagePage.tsx b/client/web/src/cody/management/subscription/manage/CodySubscriptionManagePage.tsx
index 9310665f0f74..cdadaa34f181 100644
--- a/client/web/src/cody/management/subscription/manage/CodySubscriptionManagePage.tsx
+++ b/client/web/src/cody/management/subscription/manage/CodySubscriptionManagePage.tsx
@@ -4,22 +4,17 @@ import classNames from 'classnames'
 import { Navigate } from 'react-router-dom'
 
 import { logger } from '@sourcegraph/common'
-import { useQuery } from '@sourcegraph/http-client'
 import type { AuthenticatedUser } from '@sourcegraph/shared/src/auth'
 import { Alert, Card, LoadingSpinner, PageHeader, Text } from '@sourcegraph/wildcard'
 
 import { withAuthenticatedUser } from '../../../../auth/withAuthenticatedUser'
 import { Page } from '../../../../components/Page'
 import { PageTitle } from '../../../../components/PageTitle'
-import {
-    CodySubscriptionPlan,
-    type UserCodyPlanResult,
-    type UserCodyPlanVariables,
-} from '../../../../graphql-operations'
+import { CodySubscriptionPlan } from '../../../../graphql-operations'
 import type { LegacyLayoutRouteContext } from '../../../../LegacyRouteContext'
 import { CodyProRoutes } from '../../../codyProRoutes'
 import { PageHeaderIcon } from '../../../components/PageHeaderIcon'
-import { USER_CODY_PLAN } from '../../../subscription/queries'
+import type { UserCodySubscription } from '../../../subscription/useUserCodySubscription'
 import { useCurrentSubscription, useSubscriptionSummary } from '../../api/react-query/subscriptions'
 
 import { InvoiceHistory } from './InvoiceHistory'
@@ -30,14 +25,10 @@ import styles from './CodySubscriptionManagePage.module.scss'
 
 interface Props extends Pick<LegacyLayoutRouteContext, 'telemetryRecorder'> {
     authenticatedUser: AuthenticatedUser
+    codySubscription: UserCodySubscription
 }
 
-const AuthenticatedCodySubscriptionManagePage: React.FC<Props> = ({ telemetryRecorder }) => {
-    const {
-        loading: userCodyPlanLoading,
-        error: useCodyPlanError,
-        data: userCodyPlanData,
-    } = useQuery<UserCodyPlanResult, UserCodyPlanVariables>(USER_CODY_PLAN, {})
+const AuthenticatedCodySubscriptionManagePage: React.FC<Props> = ({ telemetryRecorder, codySubscription }) => {
     const subscriptionSummaryQuery = useSubscriptionSummary()
 
     useEffect(
@@ -47,26 +38,15 @@ const AuthenticatedCodySubscriptionManagePage: React.FC<Props> = ({ telemetryRec
         [telemetryRecorder]
     )
 
-    if (userCodyPlanLoading || subscriptionSummaryQuery.isLoading) {
+    if (subscriptionSummaryQuery.isLoading) {
         return <LoadingSpinner />
     }
 
-    if (useCodyPlanError) {
-        logger.error('Failed to fetch Cody subscription data', useCodyPlanError)
-        return null
-    }
-
     if (subscriptionSummaryQuery.isError) {
         logger.error('Failed to fetch Cody subscription summary', subscriptionSummaryQuery.error)
         return null
     }
 
-    const subscriptionData = userCodyPlanData?.currentUser?.codySubscription
-    if (!subscriptionData) {
-        logger.error('Cody subscription data is not available.')
-        return null
-    }
-
     if (!subscriptionSummaryQuery.data) {
         logger.error('Cody subscription summary is not available.')
         return null
@@ -74,7 +54,7 @@ const AuthenticatedCodySubscriptionManagePage: React.FC<Props> = ({ telemetryRec
 
     // This page only applies to users who have a Cody Pro subscription to manage.
     // Otherwise, direct them to the ./new page to sign up.
-    if (subscriptionData.plan !== CodySubscriptionPlan.PRO) {
+    if (codySubscription.plan !== CodySubscriptionPlan.PRO) {
         return <Navigate to={CodyProRoutes.NewProSubscription} replace={true} />
     }
 
diff --git a/client/web/src/cody/management/subscription/new/NewCodyProSubscriptionPage.tsx b/client/web/src/cody/management/subscription/new/NewCodyProSubscriptionPage.tsx
index d922b958c943..844f364a53e5 100644
--- a/client/web/src/cody/management/subscription/new/NewCodyProSubscriptionPage.tsx
+++ b/client/web/src/cody/management/subscription/new/NewCodyProSubscriptionPage.tsx
@@ -9,7 +9,6 @@ import { loadStripe } from '@stripe/stripe-js'
 import classNames from 'classnames'
 import { Navigate, useSearchParams } from 'react-router-dom'
 
-import { useQuery } from '@sourcegraph/http-client'
 import type { TelemetryV2Props } from '@sourcegraph/shared/src/telemetry'
 import { PageHeader, Text, LoadingSpinner, Alert, logger } from '@sourcegraph/wildcard'
 
@@ -17,14 +16,10 @@ import type { AuthenticatedUser } from '../../../../auth'
 import { withAuthenticatedUser } from '../../../../auth/withAuthenticatedUser'
 import { Page } from '../../../../components/Page'
 import { PageTitle } from '../../../../components/PageTitle'
-import {
-    type UserCodyPlanResult,
-    type UserCodyPlanVariables,
-    CodySubscriptionPlan,
-} from '../../../../graphql-operations'
+import { CodySubscriptionPlan } from '../../../../graphql-operations'
 import { CodyProRoutes } from '../../../codyProRoutes'
 import { PageHeaderIcon } from '../../../components/PageHeaderIcon'
-import { USER_CODY_PLAN } from '../../../subscription/queries'
+import type { UserCodySubscription } from '../../../subscription/useUserCodySubscription'
 import { defaultCodyProApiClientContext, CodyProApiClientContext } from '../../api/components/CodyProApiClient'
 import { useCurrentSubscription } from '../../api/react-query/subscriptions'
 import { useBillingAddressStripeElementsOptions } from '../manage/BillingAddress'
@@ -41,11 +36,13 @@ const stripe = await loadStripe(publishableKey || '')
 
 interface NewCodyProSubscriptionPageProps extends TelemetryV2Props {
     authenticatedUser: AuthenticatedUser
+    codySubscription: UserCodySubscription
 }
 
 const AuthenticatedNewCodyProSubscriptionPage: FunctionComponent<NewCodyProSubscriptionPageProps> = ({
     authenticatedUser,
     telemetryRecorder,
+    codySubscription,
 }) => {
     const [urlSearchParams] = useSearchParams()
     const addSeats = !!urlSearchParams.get('addSeats')
@@ -53,12 +50,6 @@ const AuthenticatedNewCodyProSubscriptionPage: FunctionComponent<NewCodyProSubsc
 
     const stripeElementsOptions = useBillingAddressStripeElementsOptions()
 
-    // Load data
-    const {
-        loading: userCodyPlanLoading,
-        data: userCodyPlanData,
-        error: userCodyPlanError,
-    } = useQuery<UserCodyPlanResult, UserCodyPlanVariables>(USER_CODY_PLAN, {})
     const subscriptionQueryResult = useCurrentSubscription()
     const subscription = subscriptionQueryResult?.data
 
@@ -67,25 +58,20 @@ const AuthenticatedNewCodyProSubscriptionPage: FunctionComponent<NewCodyProSubsc
     }, [telemetryRecorder])
 
     useEffect(() => {
-        if (userCodyPlanError) {
-            logger.error('Failed to fetch subscription data', userCodyPlanError)
-        }
-    }, [userCodyPlanError])
-    useEffect(() => {
-        if (subscriptionQueryResult.isError) {
+        if (subscriptionQueryResult.error) {
             logger.error('Failed to fetch subscription data', subscriptionQueryResult.error)
         }
-    }, [subscriptionQueryResult.isError, subscriptionQueryResult.error])
-
-    if (userCodyPlanLoading || subscriptionQueryResult.isLoading) {
-        return <LoadingSpinner className="mx-auto" />
-    }
+    }, [subscriptionQueryResult.error])
 
     // If the user already has a Cody Pro subscription, direct them back to the Cody Management page.
-    if (!addSeats && userCodyPlanData?.currentUser?.codySubscription?.plan === CodySubscriptionPlan.PRO) {
+    if (!addSeats && codySubscription.plan === CodySubscriptionPlan.PRO) {
         return <Navigate to={CodyProRoutes.Manage} replace={true} />
     }
 
+    if (subscriptionQueryResult.isLoading) {
+        return <LoadingSpinner className="mx-auto" />
+    }
+
     const PageWithHeader = ({ children }: { children: React.ReactNode }): React.ReactElement => (
         <Page className={classNames('d-flex flex-column', styles.page)}>
             <PageTitle title={addSeats ? 'Add seats' : 'New subscription'} />
@@ -103,7 +89,7 @@ const AuthenticatedNewCodyProSubscriptionPage: FunctionComponent<NewCodyProSubsc
         </Page>
     )
 
-    if (userCodyPlanLoading || subscriptionQueryResult.isLoading) {
+    if (subscriptionQueryResult.isLoading) {
         return (
             <PageWithHeader>
                 <LoadingSpinner className="mx-auto" />
@@ -111,14 +97,6 @@ const AuthenticatedNewCodyProSubscriptionPage: FunctionComponent<NewCodyProSubsc
         )
     }
 
-    if (userCodyPlanError) {
-        return (
-            <PageWithHeader>
-                <Alert variant="danger">Failed to fetch user Cody plan data</Alert>
-            </PageWithHeader>
-        )
-    }
-
     if (addSeats && subscriptionQueryResult.isError) {
         return (
             <PageWithHeader>
diff --git a/client/web/src/cody/subscription/CodySubscriptionPage.tsx b/client/web/src/cody/subscription/CodySubscriptionPage.tsx
index 2a174cf8a8cf..15e40e3d3c25 100644
--- a/client/web/src/cody/subscription/CodySubscriptionPage.tsx
+++ b/client/web/src/cody/subscription/CodySubscriptionPage.tsx
@@ -4,7 +4,6 @@ import { mdiArrowLeft, mdiInformationOutline, mdiTrendingUp, mdiCreditCardOutlin
 import classNames from 'classnames'
 import { useNavigate } from 'react-router-dom'
 
-import { useQuery } from '@sourcegraph/http-client'
 import type { TelemetryV2Props } from '@sourcegraph/shared/src/telemetry'
 import {
     Badge,
@@ -21,27 +20,27 @@ import {
 } from '@sourcegraph/wildcard'
 
 import type { AuthenticatedUser } from '../../auth'
+import { withAuthenticatedUser } from '../../auth/withAuthenticatedUser'
 import { Page } from '../../components/Page'
 import { PageTitle } from '../../components/PageTitle'
 import { CodySubscriptionPlan } from '../../graphql-operations'
-import type { UserCodyPlanResult, UserCodyPlanVariables } from '../../graphql-operations'
 import { CodyProRoutes } from '../codyProRoutes'
 import { ProIcon } from '../components/CodyIcon'
 import { PageHeaderIcon } from '../components/PageHeaderIcon'
-import { isCodyEnabled } from '../isCodyEnabled'
 import { getManageSubscriptionPageURL, isEmbeddedCodyProUIEnabled, manageSubscriptionRedirectURL } from '../util'
 
-import { USER_CODY_PLAN } from './queries'
+import type { UserCodySubscription } from './useUserCodySubscription'
 
 import styles from './CodySubscriptionPage.module.scss'
 
 interface CodySubscriptionPageProps extends TelemetryV2Props {
-    authenticatedUser?: AuthenticatedUser | null
+    authenticatedUser: AuthenticatedUser
+    codySubscription: UserCodySubscription
 }
 
-export const CodySubscriptionPage: React.FunctionComponent<CodySubscriptionPageProps> = ({
-    authenticatedUser,
+const AuthenticatedCodySubscriptionManagePage: React.FunctionComponent<CodySubscriptionPageProps> = ({
     telemetryRecorder,
+    codySubscription,
 }) => {
     const parameters = useSearchParameters()
 
@@ -50,26 +49,10 @@ export const CodySubscriptionPage: React.FunctionComponent<CodySubscriptionPageP
         telemetryRecorder.recordEvent('cody.planSelection', 'view')
     }, [utm_source, telemetryRecorder])
 
-    const { data, error: dataError } = useQuery<UserCodyPlanResult, UserCodyPlanVariables>(USER_CODY_PLAN, {})
-
     const navigate = useNavigate()
     const useEmbeddedCodyUI = useMemo(() => isEmbeddedCodyProUIEnabled(), [])
 
-    useEffect(() => {
-        if (!!data && !data?.currentUser) {
-            navigate(`/sign-in?returnTo=${CodyProRoutes.Subscription}`)
-        }
-    }, [data, navigate])
-
-    if (dataError) {
-        throw dataError
-    }
-
-    if (!isCodyEnabled() || !data?.currentUser || !authenticatedUser) {
-        return null
-    }
-
-    const isProUser = data.currentUser.codySubscription?.plan === CodySubscriptionPlan.PRO
+    const isProUser = codySubscription.plan === CodySubscriptionPlan.PRO
 
     return (
         <>
@@ -454,6 +437,8 @@ export const CodySubscriptionPage: React.FunctionComponent<CodySubscriptionPageP
     )
 }
 
+export const CodySubscriptionPage = withAuthenticatedUser(AuthenticatedCodySubscriptionManagePage)
+
 export const ProTierIcon = ({ className }: { className?: string }): ReactElement => (
     <svg
         xmlns="http://www.w3.org/2000/svg"
diff --git a/client/web/src/cody/subscription/useUserCodySubscription.ts b/client/web/src/cody/subscription/useUserCodySubscription.ts
new file mode 100644
index 000000000000..6b406099a657
--- /dev/null
+++ b/client/web/src/cody/subscription/useUserCodySubscription.ts
@@ -0,0 +1,29 @@
+import type { QueryResult } from '@apollo/client'
+
+import { gql, useQuery } from '@sourcegraph/http-client'
+
+import type { UserCodyPlanResult, UserCodyPlanVariables } from '../../graphql-operations'
+
+export type UserCodySubscription = Omit<
+    NonNullable<NonNullable<NonNullable<UserCodyPlanResult['currentUser']>>['codySubscription']>,
+    '__typename'
+>
+
+const USER_CODY_PLAN = gql`
+    query UserCodyPlan {
+        currentUser {
+            id
+            codySubscription {
+                status
+                plan
+                applyProRateLimits
+                currentPeriodStartAt
+                currentPeriodEndAt
+                cancelAtPeriodEnd
+            }
+        }
+    }
+`
+
+export const useUserCodySubscription = (): QueryResult<UserCodyPlanResult, UserCodyPlanVariables> =>
+    useQuery<UserCodyPlanResult, UserCodyPlanVariables>(USER_CODY_PLAN, {})
diff --git a/client/web/src/cody/team/CodyManageTeamPage.tsx b/client/web/src/cody/team/CodyManageTeamPage.tsx
index 4fea4ee10369..0c01b141db75 100644
--- a/client/web/src/cody/team/CodyManageTeamPage.tsx
+++ b/client/web/src/cody/team/CodyManageTeamPage.tsx
@@ -18,11 +18,13 @@ import { InviteUsers } from '../invites/InviteUsers'
 import { useTeamInvites } from '../management/api/react-query/invites'
 import { useCurrentSubscription, useSubscriptionSummary } from '../management/api/react-query/subscriptions'
 import { useTeamMembers } from '../management/api/react-query/teams'
+import type { UserCodySubscription } from '../subscription/useUserCodySubscription'
 
 import { TeamMemberList } from './TeamMemberList'
 
 interface CodyManageTeamPageProps extends TelemetryV2Props {
     authenticatedUser: AuthenticatedUser
+    codySubscription: UserCodySubscription
 }
 
 const AuthenticatedCodyManageTeamPage: React.FunctionComponent<CodyManageTeamPageProps> = ({ telemetryRecorder }) => {

From 32090933094f1fd90ea5e5da6f690ecfa0b1864a Mon Sep 17 00:00:00 2001
From: Taras Yemets <yemets.taras@gmail.com>
Date: Fri, 21 Jun 2024 01:39:07 +0300
Subject: [PATCH 3/8] fix banner styles

---
 .../src/cody/components/CodyAlert.module.scss |  7 +++++++
 client/web/src/cody/invites/InviteUsers.tsx   | 19 +++++++++++--------
 2 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/client/web/src/cody/components/CodyAlert.module.scss b/client/web/src/cody/components/CodyAlert.module.scss
index 4437c4ce8556..8b6ed5c1a84c 100644
--- a/client/web/src/cody/components/CodyAlert.module.scss
+++ b/client/web/src/cody/components/CodyAlert.module.scss
@@ -41,8 +41,15 @@
         background-color: #51cf66;
     }
     &::after {
+        top: calc(50% - 16px);
+        left: 1.5rem;
+        /* stylelint-disable-next-line declaration-property-unit-allowed-list */
+        width: 32px;
+        /* stylelint-disable-next-line declaration-property-unit-allowed-list */
+        height: 32px;
         // Green checkmark SVG
         background-image: url("data:image/svg+xml,<svg xmlns='http://www.w3.org/2000/svg' viewBox='0 0 31 25'><path fill='%23319E44' d='M27 .4a3.4 3.4 0 0 0-2.4 1.1L12 15.6 7 9.3a3.4 3.4 0 0 0-5-.6 3.4 3.4 0 0 0-.5 4.9L9.2 23a3.4 3.4 0 0 0 5.3.1l15.2-17a3.4 3.4 0 0 0-.3-4.9 3.4 3.4 0 0 0-2.5-.8Z'/></svg>");
+        background-repeat: no-repeat;
     }
 }
 
diff --git a/client/web/src/cody/invites/InviteUsers.tsx b/client/web/src/cody/invites/InviteUsers.tsx
index b609d9c51896..74a47b45ffec 100644
--- a/client/web/src/cody/invites/InviteUsers.tsx
+++ b/client/web/src/cody/invites/InviteUsers.tsx
@@ -105,6 +105,17 @@ export const InviteUsers: React.FunctionComponent<InviteUsersProps> = ({ telemet
         })
     }, [emailAddresses, sendInviteMutation.mutateAsync, teamId, telemetryRecorder, verifyEmailList])
 
+    if (updateSubscriptionMutation.isSuccess) {
+        return (
+            <CodyAlert variant="greenSuccess">
+                <H1 as="p" className="mb-2">
+                    Remaining invites removed from plan
+                </H1>
+                <Text className="mb-0">You can add more seats at any time with the "Add seats" button.</Text>
+            </CodyAlert>
+        )
+    }
+
     if (!isAdmin || !remainingInviteCount || !subscriptionQueryResult.data) {
         return null
     }
@@ -113,14 +124,6 @@ export const InviteUsers: React.FunctionComponent<InviteUsersProps> = ({ telemet
 
     return (
         <>
-            {updateSubscriptionMutation.isSuccess && (
-                <CodyAlert variant="greenSuccess">
-                    <H1 as="p" className="mb-2">
-                        Remaining invites removed from plan
-                    </H1>
-                    <Text className="mb-0">You can add more seats at any time with the "Add seats" button.</Text>
-                </CodyAlert>
-            )}
             {sendInviteMutation.status === 'success' && (
                 <CodyAlert variant="greenSuccess">
                     <H3>

From 94807802e66d12cbc22c2e9f41343e6541d8dc25 Mon Sep 17 00:00:00 2001
From: Taras Yemets <yemets.taras@gmail.com>
Date: Fri, 21 Jun 2024 01:48:11 +0300
Subject: [PATCH 4/8] add comment and log error

---
 client/web/src/cody/codyProRoutes.tsx | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/client/web/src/cody/codyProRoutes.tsx b/client/web/src/cody/codyProRoutes.tsx
index 60240e01d95c..58772f256df9 100644
--- a/client/web/src/cody/codyProRoutes.tsx
+++ b/client/web/src/cody/codyProRoutes.tsx
@@ -1,5 +1,6 @@
 import { Navigate, type RouteObject } from 'react-router-dom'
 
+import { logger } from '@sourcegraph/common'
 import { lazyComponent } from '@sourcegraph/shared/src/util/lazyComponent'
 
 import { type LegacyLayoutRouteContext, LegacyRoute } from '../LegacyRouteContext'
@@ -83,10 +84,12 @@ const CodyProPage: React.FC<CodyProPageProps> = props => {
     }
 
     if (!data.currentUser) {
+        // Cody plan is not available if the user is not authenticated. Redirecting to the sign-in page.
         return <Navigate to={`/sign-in?returnTo=${CodyProRoutes.Manage}`} replace={true} />
     }
 
     if (!data.currentUser.codySubscription) {
+        logger.error('Cody subscription data is not available.')
         return null
     }
 

From 84349d1a8be5a627863dcd4f9eca33393a7d5f08 Mon Sep 17 00:00:00 2001
From: Taras Yemets <yemets.taras@gmail.com>
Date: Fri, 21 Jun 2024 01:57:02 +0300
Subject: [PATCH 5/8] refetch user cody plan after accepting an invite

---
 .../src/cody/invites/AcceptInviteBanner.tsx    | 18 +++++++++---------
 .../src/cody/management/CodyManagementPage.tsx |  2 +-
 2 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/client/web/src/cody/invites/AcceptInviteBanner.tsx b/client/web/src/cody/invites/AcceptInviteBanner.tsx
index 0af1cda0cdc3..57d8ee4601c3 100644
--- a/client/web/src/cody/invites/AcceptInviteBanner.tsx
+++ b/client/web/src/cody/invites/AcceptInviteBanner.tsx
@@ -3,11 +3,12 @@ import { Button, ButtonLink, H1, Text } from '@sourcegraph/wildcard'
 import { CodyProRoutes } from '../codyProRoutes'
 import { CodyAlert } from '../components/CodyAlert'
 import { useAcceptInvite, useCancelInvite } from '../management/api/react-query/invites'
+import { useUserCodySubscription } from '../subscription/useUserCodySubscription'
 
 import { useInviteParams } from './useInviteParams'
 import { UserInviteStatus, useInviteState } from './useInviteState'
 
-export const AcceptInviteBanner: React.FC<{ onSuccess: () => unknown }> = ({ onSuccess }) => {
+export const AcceptInviteBanner: React.FC = () => {
     const { inviteParams, clearInviteParams } = useInviteParams()
     if (!inviteParams) {
         return null
@@ -16,18 +17,17 @@ export const AcceptInviteBanner: React.FC<{ onSuccess: () => unknown }> = ({ onS
         <AcceptInviteBannerContent
             teamId={inviteParams.teamId}
             inviteId={inviteParams.inviteId}
-            onSuccess={onSuccess}
             clearInviteParams={clearInviteParams}
         />
     )
 }
 
-const AcceptInviteBannerContent: React.FC<{
-    teamId: string
-    inviteId: string
-    onSuccess: () => unknown
-    clearInviteParams: () => void
-}> = ({ teamId, inviteId, onSuccess, clearInviteParams }) => {
+const AcceptInviteBannerContent: React.FC<{ teamId: string; inviteId: string; clearInviteParams: () => void }> = ({
+    teamId,
+    inviteId,
+    clearInviteParams,
+}) => {
+    const { refetch } = useUserCodySubscription()
     const inviteState = useInviteState(teamId, inviteId)
     const acceptInviteMutation = useAcceptInvite()
     const cancelInviteMutation = useCancelInvite()
@@ -108,7 +108,7 @@ const AcceptInviteBannerContent: React.FC<{
                                     onClick={() =>
                                         acceptInviteMutation.mutate(
                                             { teamId, inviteId },
-                                            { onSuccess, onSettled: clearInviteParams }
+                                            { onSuccess: () => refetch(), onSettled: clearInviteParams }
                                         )
                                     }
                                 >
diff --git a/client/web/src/cody/management/CodyManagementPage.tsx b/client/web/src/cody/management/CodyManagementPage.tsx
index a11ed2a48f0d..3dbc21389352 100644
--- a/client/web/src/cody/management/CodyManagementPage.tsx
+++ b/client/web/src/cody/management/CodyManagementPage.tsx
@@ -115,7 +115,7 @@ export const CodyManagementPage: React.FunctionComponent<CodyManagementPageProps
         <>
             <Page className={classNames('d-flex flex-column')}>
                 <PageTitle title="Dashboard" />
-                <AcceptInviteBanner onSuccess={refetch} />
+                <AcceptInviteBanner />
                 {welcomeToPro && (
                     <CodyAlert variant="greenCodyPro">
                         <H2 className="mt-4">Welcome to Cody Pro</H2>

From e29333e1e2cca9df209bf02b1d126b12fe8a560f Mon Sep 17 00:00:00 2001
From: Taras Yemets <yemets.taras@gmail.com>
Date: Fri, 21 Jun 2024 02:03:40 +0300
Subject: [PATCH 6/8] fix page name

---
 client/web/src/cody/subscription/CodySubscriptionPage.tsx | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/client/web/src/cody/subscription/CodySubscriptionPage.tsx b/client/web/src/cody/subscription/CodySubscriptionPage.tsx
index 15e40e3d3c25..b9eb7ea0b346 100644
--- a/client/web/src/cody/subscription/CodySubscriptionPage.tsx
+++ b/client/web/src/cody/subscription/CodySubscriptionPage.tsx
@@ -38,7 +38,7 @@ interface CodySubscriptionPageProps extends TelemetryV2Props {
     codySubscription: UserCodySubscription
 }
 
-const AuthenticatedCodySubscriptionManagePage: React.FunctionComponent<CodySubscriptionPageProps> = ({
+const AuthenticatedCodySubscriptionPage: React.FunctionComponent<CodySubscriptionPageProps> = ({
     telemetryRecorder,
     codySubscription,
 }) => {
@@ -437,7 +437,7 @@ const AuthenticatedCodySubscriptionManagePage: React.FunctionComponent<CodySubsc
     )
 }
 
-export const CodySubscriptionPage = withAuthenticatedUser(AuthenticatedCodySubscriptionManagePage)
+export const CodySubscriptionPage = withAuthenticatedUser(AuthenticatedCodySubscriptionPage)
 
 export const ProTierIcon = ({ className }: { className?: string }): ReactElement => (
     <svg

From a563e81ebadb60da83ffc086e5be3969806825bd Mon Sep 17 00:00:00 2001
From: Taras Yemets <yemets.taras@gmail.com>
Date: Fri, 21 Jun 2024 02:28:01 +0300
Subject: [PATCH 7/8] bazel configure

---
 client/web/BUILD.bazel | 1 +
 1 file changed, 1 insertion(+)

diff --git a/client/web/BUILD.bazel b/client/web/BUILD.bazel
index 382bea2945c9..b8cee9378440 100644
--- a/client/web/BUILD.bazel
+++ b/client/web/BUILD.bazel
@@ -281,6 +281,7 @@ ts_project(
         "src/cody/sidebar/useSidebarSize.tsx",
         "src/cody/subscription/CodySubscriptionPage.tsx",
         "src/cody/subscription/queries.tsx",
+        "src/cody/subscription/useUserCodySubscription.ts",
         "src/cody/switch-account/CodySwitchAccountPage.tsx",
         "src/cody/team/CodyManageTeamPage.tsx",
         "src/cody/team/TeamMemberList.tsx",

From acf25ac9d0ac17405f4d471184e2f3d2a469b9f0 Mon Sep 17 00:00:00 2001
From: Taras Yemets <yemets.taras@gmail.com>
Date: Tue, 25 Jun 2024 15:12:52 +0300
Subject: [PATCH 8/8] Update
 client/web/src/cody/management/subscription/new/NewCodyProSubscriptionPage.tsx

Co-authored-by: David Veszelovszki <veszelovszki@gmail.com>
---
 .../management/subscription/new/NewCodyProSubscriptionPage.tsx   | 1 +
 1 file changed, 1 insertion(+)

diff --git a/client/web/src/cody/management/subscription/new/NewCodyProSubscriptionPage.tsx b/client/web/src/cody/management/subscription/new/NewCodyProSubscriptionPage.tsx
index 844f364a53e5..e2e003e3d730 100644
--- a/client/web/src/cody/management/subscription/new/NewCodyProSubscriptionPage.tsx
+++ b/client/web/src/cody/management/subscription/new/NewCodyProSubscriptionPage.tsx
@@ -68,6 +68,7 @@ const AuthenticatedNewCodyProSubscriptionPage: FunctionComponent<NewCodyProSubsc
         return <Navigate to={CodyProRoutes.Manage} replace={true} />
     }
 
+    // Display spinner without page header because without the subscription, we don't know which header to show
     if (subscriptionQueryResult.isLoading) {
         return <LoadingSpinner className="mx-auto" />
     }