feat/dotcom: use Enterprise Portal for Cody Gateway usage (#63653)

Closes https://linear.app/sourcegraph/issue/CORE-211

See https://linear.app/sourcegraph/issue/CORE-100 for a higher-level
view - this is the first proof-of-concept for achieving our migration
strategy to extract Enterprise subscription data out of dotcom while
retaining the existing UI until a future project ships a dedicated
Enterprise Portal UI
(https://linear.app/sourcegraph/project/kr-p-enterprise-portal-user-interface-dadd5ff28bd8).

The integration uses generated ConnectRPC client code + `react-query`,
the latter of which has already been used elsewhere for SSC
integrations. This is partly supported by
https://github.com/connectrpc/connect-query-es which offers
mostly-first-class integration with `react-query`, but I had to do some
fenangling to provide the query clients directly as I can't get the
React provider thing to work. The ConnectRPC clients point to the
proxies introduced in
https://github.com/sourcegraph/sourcegraph/pull/63652 which
authenticates the requests for Enterprise Portal, until we ship
https://linear.app/sourcegraph/project/kr-p1-streamlined-role-assignment-via-sams-and-entitle-2f118b3f9d4c/overview

## Test plan

### Local

First, `sg start dotcom`

Choose a subscription you have locally. Use `psql -d sourcegraph` to
connect to local database, then:

```
sourcegraph=# delete from product_licenses where product_subscription_id = '<local subscription ID>';
DELETE 1
sourcegraph=# update product_subscriptions set id = '58b95c21-c2d0-4b4b-8b15-bf1b926d3557' where id = '<local subscription ID>';
UPDATE 1
```

Now annoyingly the UI will break because there is no license, we need:

```gql
query getGraphQLID {
  dotcom {
    productSubscription(uuid:"58b95c21-c2d0-4b4b-8b15-bf1b926d3557") {
      id # graphQL ID
    }
  }
}

mutation createLicense {
  dotcom {
    generateProductLicenseForSubscription(productSubscriptionID:"<graphQLID>", license:{
      tags:["dev"]
      userCount:100
      expiresAt:1814815397
    }) {
      id
    }
  }
}
```

This effectively lets us have a "pretend S2" subscription locally.
Visiting the subscription page now at
https://sourcegraph.test:3443/site-admin/dotcom/product/subscriptions/58b95c21-c2d0-4b4b-8b15-bf1b926d3557


![image](https://github.com/sourcegraph/sourcegraph/assets/23356519/1e77d77d-8032-436b-ab1d-393b34e8e4b5)

The data matches the "real" data currently at
https://sourcegraph.com/site-admin/dotcom/product/subscriptions/58b95c21-c2d0-4b4b-8b15-bf1b926d3557

### Against dotcom

```
sg start web-standalone
```

follow
https://www.loom.com/share/6cb3b3ca475b4b9392aa4b11938e76e6?sid=6cd1a689-d75d-4133-bcff-b0c7d25b23f1
and then check out some product subscriptions

Author: bobheadxi

.eslintrc.js

@@ -30,6 +30,8 @@ const config = {
     'typedoc.js',
     'client/web/dev/**/*',
     'graphql-schema-linter.config.js',
+    // Generated code
+    'client/web/src/enterprise/site-admin/dotcom/productSubscriptions/enterpriseportalgen/**',
   ],
   extends: ['@sourcegraph/eslint-config', 'plugin:storybook/recommended'],
   env: {

.prettierignore

@@ -70,3 +70,6 @@ dev/linearhooks/internal/lineargql/schema.graphql
 # This is an embedded external minified library and should not be modified
 internal/appliance/web/static/script/htmx.min.js
 internal/appliance/web/static/script/bootstrap.bundle.min.js
+
+# Generated code
+client/web/src/enterprise/site-admin/dotcom/productSubscriptions/enterpriseportalgen/**

client/shared/package.json

@@ -30,4 +30,4 @@
     "@sourcegraph/wildcard": "workspace:*"
   },
   "sideEffects": true
-}
+}

client/web/BUILD.bazel

@@ -1056,6 +1056,9 @@ ts_project(
         "src/enterprise/site-admin/dotcom/productSubscriptions/SiteAdminProductSubscriptionPage.tsx",
         "src/enterprise/site-admin/dotcom/productSubscriptions/SiteAdminProductSubscriptionsPage.tsx",
         "src/enterprise/site-admin/dotcom/productSubscriptions/backend.ts",
+        "src/enterprise/site-admin/dotcom/productSubscriptions/enterpriseportal.ts",
+        "src/enterprise/site-admin/dotcom/productSubscriptions/enterpriseportalgen/codyaccess-CodyAccessService_connectquery.ts",
+        "src/enterprise/site-admin/dotcom/productSubscriptions/enterpriseportalgen/codyaccess_pb.ts",
         "src/enterprise/site-admin/dotcom/productSubscriptions/plandata.ts",
         "src/enterprise/site-admin/dotcom/productSubscriptions/testUtils.ts",
         "src/enterprise/site-admin/dotcom/productSubscriptions/utils.ts",
@@ -1754,6 +1757,7 @@ ts_project(
         ":node_modules/@types/node",
         ":node_modules/mermaid",
         "//:node_modules/@apollo/client",
+        "//:node_modules/@bufbuild/protobuf",
         "//:node_modules/@codemirror/commands",
         "//:node_modules/@codemirror/lang-json",
         "//:node_modules/@codemirror/lang-markdown",
@@ -1762,6 +1766,9 @@ ts_project(
         "//:node_modules/@codemirror/search",
         "//:node_modules/@codemirror/state",
         "//:node_modules/@codemirror/view",
+        "//:node_modules/@connectrpc/connect",
+        "//:node_modules/@connectrpc/connect-query",
+        "//:node_modules/@connectrpc/connect-web",
         "//:node_modules/@date-fns/utc",
         "//:node_modules/@graphiql/react",
         "//:node_modules/@lezer/common",  #keep

client/web/src/enterprise/site-admin/dotcom/productSubscriptions/CodyServicesSection.tsx

@@ -1,14 +1,15 @@
 import React, { useCallback, useState } from 'react'
 
+import type { ConnectError } from '@connectrpc/connect'
 import { mdiPencil, mdiTrashCan } from '@mdi/js'
-import { parseISO } from 'date-fns'
+import type { UseQueryResult } from '@tanstack/react-query'
 import type { GraphQLError } from 'graphql'
 
 import { Toggle } from '@sourcegraph/branded/src/components/Toggle'
 import { logger } from '@sourcegraph/common'
-import { useMutation, useQuery } from '@sourcegraph/http-client'
+import { useMutation } from '@sourcegraph/http-client'
 import { CodyGatewayRateLimitSource } from '@sourcegraph/shared/src/graphql-operations'
-import { TelemetryV2Props } from '@sourcegraph/shared/src/telemetry'
+import type { TelemetryV2Props } from '@sourcegraph/shared/src/telemetry'
 import {
     H3,
     ProductStatusBadge,
@@ -35,26 +36,20 @@ import type {
     Scalars,
     UpdateCodyGatewayConfigResult,
     UpdateCodyGatewayConfigVariables,
-    CodyGatewayRateLimitUsageDatapoint,
     CodyGatewayRateLimitFields,
-    DotComProductSubscriptionCodyGatewayCompletionsUsageResult,
-    DotComProductSubscriptionCodyGatewayCompletionsUsageVariables,
-    DotComProductSubscriptionCodyGatewayEmbeddingsUsageVariables,
-    DotComProductSubscriptionCodyGatewayEmbeddingsUsageResult,
 } from '../../../../graphql-operations'
 import { ChartContainer } from '../../../../site-admin/analytics/components/ChartContainer'
 
-import {
-    DOTCOM_PRODUCT_SUBSCRIPTION_CODY_GATEWAY_COMPLETIONS_USAGE,
-    DOTCOM_PRODUCT_SUBSCRIPTION_CODY_GATEWAY_EMBEDDINGS_USAGE,
-    UPDATE_CODY_GATEWAY_CONFIG,
-} from './backend'
+import { UPDATE_CODY_GATEWAY_CONFIG } from './backend'
 import { CodyGatewayRateLimitModal } from './CodyGatewayRateLimitModal'
+import { useGetCodyGatewayUsage, type EnterprisePortalEnvironment } from './enterpriseportal'
+import type { CodyGatewayUsage_UsageDatapoint, GetCodyGatewayUsageResponse } from './enterpriseportalgen/codyaccess_pb'
 import { numberFormatter, prettyInterval } from './utils'
 
 import styles from './CodyServicesSection.module.scss'
 
 interface Props extends TelemetryV2Props {
+    enterprisePortalEnvironment: EnterprisePortalEnvironment
     productSubscriptionUUID: string
     productSubscriptionID: Scalars['ID']
     currentSourcegraphAccessToken: string | null
@@ -65,6 +60,7 @@ interface Props extends TelemetryV2Props {
 }
 
 export const CodyServicesSection: React.FunctionComponent<Props> = ({
+    enterprisePortalEnvironment,
     productSubscriptionUUID,
     productSubscriptionID,
     viewerCanAdminister,
@@ -74,6 +70,9 @@ export const CodyServicesSection: React.FunctionComponent<Props> = ({
     codyGatewayAccess,
     telemetryRecorder,
 }) => {
+    // TODO: Figure out strategy for what instance to target
+    const codyGatewayUsageQuery = useGetCodyGatewayUsage(enterprisePortalEnvironment, productSubscriptionUUID)
+
     const [updateCodyGatewayConfig, { loading: updateCodyGatewayConfigLoading, error: updateCodyGatewayConfigError }] =
         useMutation<UpdateCodyGatewayConfigResult, UpdateCodyGatewayConfigVariables>(UPDATE_CODY_GATEWAY_CONFIG)
 
@@ -175,7 +174,7 @@ export const CodyServicesSection: React.FunctionComponent<Props> = ({
                                         />
                                     </tbody>
                                 </table>
-                                <RateLimitUsage mode="completions" productSubscriptionUUID={productSubscriptionUUID} />
+                                <RateLimitUsage mode="completions" usageQuery={codyGatewayUsageQuery} />
 
                                 <hr className="my-3" />
 
@@ -201,10 +200,7 @@ export const CodyServicesSection: React.FunctionComponent<Props> = ({
                                         />
                                     </tbody>
                                 </table>
-                                <EmbeddingsRateLimitUsage
-                                    mode="embeddings"
-                                    productSubscriptionUUID={productSubscriptionUUID}
-                                />
+                                <EmbeddingsRateLimitUsage mode="embeddings" usageQuery={codyGatewayUsageQuery} />
                             </>
                         )}
 
@@ -277,8 +273,8 @@ export const CodyGatewayRateLimitSourceBadge: React.FunctionComponent<{
     }
 }
 
-function generateSeries(data: CodyGatewayRateLimitUsageDatapoint[]): [string, CodyGatewayRateLimitUsageDatapoint[]][] {
-    const series: Record<string, CodyGatewayRateLimitUsageDatapoint[]> = {}
+function generateSeries(data: CodyGatewayUsage_UsageDatapoint[]): [string, CodyGatewayUsage_UsageDatapoint[]][] {
+    const series: Record<string, CodyGatewayUsage_UsageDatapoint[]> = {}
     for (const entry of data) {
         if (!series[entry.model]) {
             series[entry.model] = []
@@ -411,17 +407,12 @@ const RateLimitRow: React.FunctionComponent<RateLimitRowProps> = ({
 }
 
 interface RateLimitUsageProps {
-    productSubscriptionUUID: string
+    usageQuery: UseQueryResult<GetCodyGatewayUsageResponse, ConnectError>
     mode: 'completions' | 'embeddings'
 }
 
-const RateLimitUsage: React.FunctionComponent<RateLimitUsageProps> = ({ productSubscriptionUUID }) => {
-    const { data, loading, error } = useQuery<
-        DotComProductSubscriptionCodyGatewayCompletionsUsageResult,
-        DotComProductSubscriptionCodyGatewayCompletionsUsageVariables
-    >(DOTCOM_PRODUCT_SUBSCRIPTION_CODY_GATEWAY_COMPLETIONS_USAGE, { variables: { uuid: productSubscriptionUUID } })
-
-    if (loading && !data) {
+const RateLimitUsage: React.FunctionComponent<RateLimitUsageProps> = ({ usageQuery: { data, isLoading, error } }) => {
+    if (isLoading && !data) {
         return (
             <>
                 <H5 className="mb-2">Usage</H5>
@@ -439,8 +430,7 @@ const RateLimitUsage: React.FunctionComponent<RateLimitUsageProps> = ({ productS
         )
     }
 
-    const { codyGatewayAccess } = data!.dotcom.productSubscription
-
+    const usage = data?.usage
     return (
         <>
             <H5 className="mb-2">Usage</H5>
@@ -450,28 +440,28 @@ const RateLimitUsage: React.FunctionComponent<RateLimitUsageProps> = ({ productS
                         width={width}
                         height={200}
                         series={[
-                            ...generateSeries(codyGatewayAccess.chatCompletionsRateLimit?.usage ?? []).map(
-                                ([model, data]): Series<CodyGatewayRateLimitUsageDatapoint> => ({
+                            ...generateSeries(usage?.chatCompletionsUsage ?? []).map(
+                                ([model, data]): Series<CodyGatewayUsage_UsageDatapoint> => ({
                                     data,
                                     getXValue(datum) {
-                                        return parseISO(datum.date)
+                                        return datum.time?.toDate() || new Date()
                                     },
                                     getYValue(datum) {
-                                        return Number(datum.count)
+                                        return Number(datum.usage)
                                     },
                                     id: 'chat-usage',
                                     name: 'Chat completions: ' + model,
                                     color: 'var(--purple)',
                                 })
                             ),
-                            ...generateSeries(codyGatewayAccess.codeCompletionsRateLimit?.usage ?? []).map(
-                                ([model, data]): Series<CodyGatewayRateLimitUsageDatapoint> => ({
+                            ...generateSeries(data?.usage?.codeCompletionsUsage ?? []).map(
+                                ([model, data]): Series<CodyGatewayUsage_UsageDatapoint> => ({
                                     data,
                                     getXValue(datum) {
-                                        return parseISO(datum.date)
+                                        return datum.time?.toDate() || new Date()
                                     },
                                     getYValue(datum) {
-                                        return Number(datum.count)
+                                        return Number(datum.usage)
                                     },
                                     id: 'code-completions-usage',
                                     name: 'Code completions: ' + model,
@@ -486,13 +476,10 @@ const RateLimitUsage: React.FunctionComponent<RateLimitUsageProps> = ({ productS
     )
 }
 
-const EmbeddingsRateLimitUsage: React.FunctionComponent<RateLimitUsageProps> = ({ productSubscriptionUUID }) => {
-    const { data, loading, error } = useQuery<
-        DotComProductSubscriptionCodyGatewayEmbeddingsUsageResult,
-        DotComProductSubscriptionCodyGatewayEmbeddingsUsageVariables
-    >(DOTCOM_PRODUCT_SUBSCRIPTION_CODY_GATEWAY_EMBEDDINGS_USAGE, { variables: { uuid: productSubscriptionUUID } })
-
-    if (loading && !data) {
+const EmbeddingsRateLimitUsage: React.FunctionComponent<RateLimitUsageProps> = ({
+    usageQuery: { data, isLoading, error },
+}) => {
+    if (isLoading && !data) {
         return (
             <>
                 <H5 className="mb-2">Usage</H5>
@@ -510,8 +497,7 @@ const EmbeddingsRateLimitUsage: React.FunctionComponent<RateLimitUsageProps> = (
         )
     }
 
-    const { codyGatewayAccess } = data!.dotcom.productSubscription
-
+    const usage = data?.usage
     return (
         <>
             <H5 className="mb-2">Usage</H5>
@@ -521,16 +507,16 @@ const EmbeddingsRateLimitUsage: React.FunctionComponent<RateLimitUsageProps> = (
                         width={width}
                         height={200}
                         series={[
-                            ...generateSeries(codyGatewayAccess.embeddingsRateLimit?.usage ?? []).map(
-                                ([model, data]): Series<CodyGatewayRateLimitUsageDatapoint> => ({
+                            ...generateSeries(usage?.embeddingsUsage ?? []).map(
+                                ([model, data]): Series<CodyGatewayUsage_UsageDatapoint> => ({
                                     data,
                                     getXValue(datum) {
-                                        return parseISO(datum.date)
+                                        return datum.time?.toDate() || new Date()
                                     },
                                     getYValue(datum) {
-                                        return Number(datum.count)
+                                        return Number(datum.usage)
                                     },
-                                    id: 'chat-usage',
+                                    id: 'embeddings-usage',
                                     name: 'Embedded tokens: ' + model,
                                     color: 'var(--purple)',
                                 })

client/web/src/enterprise/site-admin/dotcom/productSubscriptions/SiteAdminCreateProductSubscriptionPage.tsx

@@ -7,7 +7,7 @@ import { catchError, concatMap, map, tap } from 'rxjs/operators'
 
 import { asError, type ErrorLike, isErrorLike } from '@sourcegraph/common'
 import { dataOrThrowErrors, gql } from '@sourcegraph/http-client'
-import { TelemetryV2Props } from '@sourcegraph/shared/src/telemetry'
+import type { TelemetryV2Props } from '@sourcegraph/shared/src/telemetry'
 import { Button, useEventObservable, Link, Alert, Icon, Form, Container, PageHeader } from '@sourcegraph/wildcard'
 
 import type { AuthenticatedUser } from '../../../../auth'

client/web/src/enterprise/site-admin/dotcom/productSubscriptions/SiteAdminProductSubscriptionPage.tsx

@@ -6,7 +6,7 @@ import { useLocation, useNavigate, useParams } from 'react-router-dom'
 import { Timestamp } from '@sourcegraph/branded/src/components/Timestamp'
 import { logger } from '@sourcegraph/common'
 import { useMutation, useQuery } from '@sourcegraph/http-client'
-import { TelemetryV2Props } from '@sourcegraph/shared/src/telemetry'
+import type { TelemetryV2Props } from '@sourcegraph/shared/src/telemetry'
 import { Button, LoadingSpinner, Link, Icon, ErrorAlert, PageHeader, Container, H3, Text } from '@sourcegraph/wildcard'
 
 import {
@@ -36,6 +36,7 @@ import {
     useProductSubscriptionLicensesConnection,
 } from './backend'
 import { CodyServicesSection } from './CodyServicesSection'
+import type { EnterprisePortalEnvironment } from './enterpriseportal'
 import { SiteAdminGenerateProductLicenseForSubscriptionForm } from './SiteAdminGenerateProductLicenseForSubscriptionForm'
 import { SiteAdminProductLicenseNode } from './SiteAdminProductLicenseNode'
 import { accessTokenPath, errorForPath, enterprisePortalID } from './utils'
@@ -122,6 +123,20 @@ export const SiteAdminProductSubscriptionPage: React.FunctionComponent<React.Pro
 
     const productSubscription = data!.dotcom.productSubscription
 
+    /**
+     * TODO(@robert): As part of https://linear.app/sourcegraph/issue/CORE-100,
+     * eventually dev subscriptions will only live on Enterprise Portal dev and
+     * prod subscriptions will only live on Enterprise Portal prod. Until we
+     * cut over, we use license tags to determine what Enterprise Portal
+     * environment to target.
+     */
+    const enterprisePortalEnvironment: EnterprisePortalEnvironment =
+        window.context.deployType === 'dev'
+            ? 'local'
+            : productSubscription.activeLicense?.info?.tags?.includes('dev')
+            ? 'dev'
+            : 'prod'
+
     return (
         <>
             <div className="site-admin-product-subscription-page">
@@ -198,6 +213,7 @@ export const SiteAdminProductSubscriptionPage: React.FunctionComponent<React.Pro
                 </Container>
 
                 <CodyServicesSection
+                    enterprisePortalEnvironment={enterprisePortalEnvironment}
                     viewerCanAdminister={true}
                     currentSourcegraphAccessToken={productSubscription.currentSourcegraphAccessToken}
                     accessTokenError={errorForPath(error, accessTokenPath)}

client/web/src/enterprise/site-admin/dotcom/productSubscriptions/backend.ts

@@ -67,73 +67,6 @@ export const CODY_GATEWAY_ACCESS_FIELDS_FRAGMENT = gql`
     }
 `
 
-const CODY_GATEWAY_RATE_LIMIT_USAGE_FIELDS = gql`
-    fragment CodyGatewayRateLimitUsageFields on CodyGatewayRateLimit {
-        usage {
-            ...CodyGatewayRateLimitUsageDatapoint
-        }
-    }
-
-    fragment CodyGatewayRateLimitUsageDatapoint on CodyGatewayUsageDatapoint {
-        date
-        count
-        model
-    }
-`
-
-export const CODY_GATEWAY_ACCESS_COMPLETIONS_USAGE_FIELDS_FRAGMENT = gql`
-    fragment CodyGatewayAccessCompletionsUsageFields on CodyGatewayAccess {
-        codeCompletionsRateLimit {
-            ...CodyGatewayRateLimitUsageFields
-        }
-        chatCompletionsRateLimit {
-            ...CodyGatewayRateLimitUsageFields
-        }
-    }
-
-    ${CODY_GATEWAY_RATE_LIMIT_USAGE_FIELDS}
-`
-
-export const DOTCOM_PRODUCT_SUBSCRIPTION_CODY_GATEWAY_COMPLETIONS_USAGE = gql`
-    query DotComProductSubscriptionCodyGatewayCompletionsUsage($uuid: String!) {
-        dotcom {
-            productSubscription(uuid: $uuid) {
-                id
-                codyGatewayAccess {
-                    ...CodyGatewayAccessCompletionsUsageFields
-                }
-            }
-        }
-    }
-
-    ${CODY_GATEWAY_ACCESS_COMPLETIONS_USAGE_FIELDS_FRAGMENT}
-`
-
-export const CODY_GATEWAY_ACCESS_EMBEDDINGS_USAGE_FIELDS_FRAGMENT = gql`
-    fragment CodyGatewayAccessEmbeddingsUsageFields on CodyGatewayAccess {
-        embeddingsRateLimit {
-            ...CodyGatewayRateLimitUsageFields
-        }
-    }
-
-    ${CODY_GATEWAY_RATE_LIMIT_USAGE_FIELDS}
-`
-
-export const DOTCOM_PRODUCT_SUBSCRIPTION_CODY_GATEWAY_EMBEDDINGS_USAGE = gql`
-    query DotComProductSubscriptionCodyGatewayEmbeddingsUsage($uuid: String!) {
-        dotcom {
-            productSubscription(uuid: $uuid) {
-                id
-                codyGatewayAccess {
-                    ...CodyGatewayAccessEmbeddingsUsageFields
-                }
-            }
-        }
-    }
-
-    ${CODY_GATEWAY_ACCESS_EMBEDDINGS_USAGE_FIELDS_FRAGMENT}
-`
-
 export const DOTCOM_PRODUCT_SUBSCRIPTION = gql`
     query DotComProductSubscription($uuid: String!) {
         dotcom {