chore(plg): provide user Cody plan on a route level

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",

client/web/src/cody/codyProRoutes.tsx

@@ -1,10 +1,12 @@
-import type { RouteObject } from 'react-router-dom'
+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'
 
 import { QueryClientProvider } from './management/api/react-query/QueryClientProvider'
+import { useUserCodySubscription } from './subscription/useUserCodySubscription'
 import { isEmbeddedCodyProUIEnabled } from './util'
 
 export enum CodyProRoutes {
@@ -77,10 +79,29 @@ 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) {
+        // 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
+    }
+
     const Component = routeComponents[props.path]
     return (
         <QueryClientProvider>
-            <Component authenticatedUser={props.authenticatedUser} telemetryRecorder={props.telemetryRecorder} />
+            <Component
+                authenticatedUser={props.authenticatedUser}
+                telemetryRecorder={props.telemetryRecorder}
+                codySubscription={data.currentUser.codySubscription}
+            />
         </QueryClientProvider>
     )
 }

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 }
                                         )
                                     }
                                 >

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,25 +101,21 @@ 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 (
         <>
             <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>
@@ -191,7 +174,7 @@ export const CodyManagementPage: React.FunctionComponent<CodyManagementPageProps
                             </div>
                         )}
                     </div>
-                    <SubscriptionStats {...{ subscription, usageData }} />
+                    <SubscriptionStats subscription={codySubscription} usageData={usageData} />
                 </div>
 
                 <UseCodyInEditorSection

client/web/src/cody/management/subscription/manage/CodySubscriptionManagePage.tsx

@@ -4,23 +4,18 @@ 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 { useCurrentSubscription } from '../../api/react-query/subscriptions'
+import type { UserCodySubscription } from '../../../subscription/useUserCodySubscription'
+import { useCurrentSubscription, useSubscriptionSummary } from '../../api/react-query/subscriptions'
 
 import { InvoiceHistory } from './InvoiceHistory'
 import { PaymentDetails } from './PaymentDetails'
@@ -30,14 +25,11 @@ 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(
         function recordViewEvent() {
@@ -46,24 +38,23 @@ const AuthenticatedCodySubscriptionManagePage: React.FC<Props> = ({ telemetryRec
         [telemetryRecorder]
     )
 
-    if (userCodyPlanLoading) {
+    if (subscriptionSummaryQuery.isLoading) {
         return <LoadingSpinner />
     }
 
-    if (useCodyPlanError) {
-        logger.error('Failed to fetch Cody subscription data', useCodyPlanError)
+    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.')
+    if (!subscriptionSummaryQuery.data) {
+        logger.error('Cody subscription summary is not available.')
         return null
     }
 
     // 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} />
     }
 

client/web/src/cody/management/subscription/new/NewCodyProSubscriptionPage.tsx

@@ -9,22 +9,17 @@ 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'
 
 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,24 +36,20 @@ 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')
     const isTeam = addSeats || parseInt(urlSearchParams.get('seats') || '', 10) > 1
 
     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,21 @@ 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} />
     }
 
+    // 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" />
+    }
+
     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,22 +90,14 @@ const AuthenticatedNewCodyProSubscriptionPage: FunctionComponent<NewCodyProSubsc
         </Page>
     )
 
-    if (userCodyPlanLoading || subscriptionQueryResult.isLoading) {
+    if (subscriptionQueryResult.isLoading) {
         return (
             <PageWithHeader>
                 <LoadingSpinner className="mx-auto" />
             </PageWithHeader>
         )
     }
 
-    if (userCodyPlanError) {
-        return (
-            <PageWithHeader>
-                <Alert variant="danger">Failed to fetch user Cody plan data</Alert>
-            </PageWithHeader>
-        )
-    }
-
     if (addSeats && subscriptionQueryResult.isError) {
         return (
             <PageWithHeader>