fix(site): share AI Bridge entitlement/permissions logic (#23834)

Introduces a new `getAIBridgePermissions` method that all AI Bridge
pages can use to restrict access/paywall. Also adds the paywall and
alert to the session threads page bc I totes forgot.

---------

Co-authored-by: Jake Howell <jacob@coder.com>

site/src/pages/AIBridgePage/AIBridgeSetupAlert.stories.tsx

@@ -0,0 +1,14 @@
+import type { Meta, StoryObj } from "@storybook/react-vite";
+import { AIBridgeSetupAlert } from "./AIBridgeSetupAlert";
+
+const meta: Meta<typeof AIBridgeSetupAlert> = {
+	title: "pages/AIBridgePage/AIBridgeSetupAlert",
+	component: AIBridgeSetupAlert,
+};
+
+export default meta;
+type Story = StoryObj<typeof AIBridgeSetupAlert>;
+
+export const Alert: Story = {
+	args: {},
+};

site/src/pages/AIBridgePage/AIBridgeSetupAlert.tsx

@@ -0,0 +1,22 @@
+import type { FC } from "react";
+import { Alert, AlertDescription, AlertTitle } from "#/components/Alert/Alert";
+import { Link } from "#/components/Link/Link";
+import { docs } from "#/utils/docs";
+
+export const AIBridgeSetupAlert: FC = () => {
+	return (
+		<Alert className="mb-12" severity="warning" prominent>
+			<AlertTitle>
+				AI Bridge is included in your license, but not set up yet.
+			</AlertTitle>
+			<AlertDescription>
+				You have access to AI Governance, but it still needs to be setup. Check
+				out the{" "}
+				<Link href={docs("/ai-coder/ai-bridge")} target="_blank">
+					AI Bridge
+				</Link>{" "}
+				documentation to get started.
+			</AlertDescription>
+		</Alert>
+	);
+};

site/src/pages/AIBridgePage/ListSessionsPage/ListSessionsPage.tsx

@@ -8,6 +8,7 @@ import { usePaginatedQuery } from "#/hooks/usePaginatedQuery";
 import { useDashboard } from "#/modules/dashboard/useDashboard";
 import { RequirePermission } from "#/modules/permissions/RequirePermission";
 import { pageTitle } from "#/utils/page";
+import { getAIBridgePermissions } from "../getAIBridgePermissions";
 import { useClientFilterMenu } from "../RequestLogsPage/RequestLogsFilter/ClientFilter";
 import { useProviderFilterMenu } from "../RequestLogsPage/RequestLogsFilter/ProviderFilter";
 import { ListSessionsPageView } from "./ListSessionsPageView";
@@ -17,15 +18,11 @@ const AISessionListPage: FC = () => {
 	const { entitlements } = useDashboard();
 	const navigate = useNavigate();
 
-	// Users are allowed to view their own request logs via the API,
-	// but this page is only visible if the feature is enabled and the user
-	// has the `viewAnyAIBridgeInterception` permission.
-	// (as its defined in the Admin settings dropdown).
-	const isEntitled =
-		entitlements.features.aibridge.entitlement === "entitled" ||
-		entitlements.features.aibridge.entitlement === "grace_period";
-	const isEnabled = entitlements.features.aibridge.enabled;
-	const hasPermission = permissions.viewAnyAIBridgeInterception;
+	const { isEntitled, isEnabled, hasPermission } = getAIBridgePermissions(
+		entitlements,
+		permissions,
+	);
+
 	const canViewSessions = isEntitled && hasPermission;
 
 	const [searchParams, setSearchParams] = useSearchParams();
@@ -33,6 +30,7 @@ const AISessionListPage: FC = () => {
 		...paginatedSessions(searchParams),
 		enabled: canViewSessions,
 	});
+
 	const filter = useFilter({
 		searchParams,
 		onSearchParamsChange: setSearchParams,

site/src/pages/AIBridgePage/ListSessionsPage/ListSessionsPageView.tsx

@@ -1,8 +1,6 @@
 import { InfoIcon } from "lucide-react";
 import type { ComponentProps, FC, PropsWithChildren } from "react";
 import type { AIBridgeSession } from "#/api/typesGenerated";
-import { Alert, AlertDescription, AlertTitle } from "#/components/Alert/Alert";
-import { Link } from "#/components/Link/Link";
 import {
 	PaginationContainer,
 	type PaginationResult,
@@ -23,8 +21,8 @@ import {
 	TooltipProvider,
 	TooltipTrigger,
 } from "#/components/Tooltip/Tooltip";
-import { docs } from "#/utils/docs";
 import { DATE_FORMAT, formatDateTime } from "#/utils/time";
+import { AIBridgeSetupAlert } from "../AIBridgeSetupAlert";
 import { ListSessionsFilter } from "./ListSessionsFilter";
 import { ListSessionsRow } from "./ListSessionsRow";
 
@@ -70,21 +68,7 @@ export const ListSessionsPageView: FC<ListSessionsPageViewProps> = ({
 	}
 
 	if (!isAISessionsEnabled) {
-		return (
-			<Alert className="mb-12" severity="warning" prominent>
-				<AlertTitle>
-					AI Bridge is included in your license, but not set up yet.
-				</AlertTitle>
-				<AlertDescription>
-					You have access to AI Governance, but it still needs to be setup.
-					Check out the{" "}
-					<Link href={docs("/ai-coder/ai-bridge")} target="_blank">
-						AI Bridge
-					</Link>{" "}
-					documentation to get started.
-				</AlertDescription>
-			</Alert>
-		);
+		return <AIBridgeSetupAlert />;
 	}
 
 	const utcOffset = formatDateTime(new Date(), DATE_FORMAT.UTC_OFFSET);

site/src/pages/AIBridgePage/RequestLogsPage/RequestLogsPage.tsx

@@ -8,6 +8,7 @@ import { usePaginatedQuery } from "#/hooks/usePaginatedQuery";
 import { useDashboard } from "#/modules/dashboard/useDashboard";
 import { RequirePermission } from "#/modules/permissions/RequirePermission";
 import { pageTitle } from "#/utils/page";
+import { getAIBridgePermissions } from "../getAIBridgePermissions";
 import { useClientFilterMenu } from "./RequestLogsFilter/ClientFilter";
 import { useModelFilterMenu } from "./RequestLogsFilter/ModelFilter";
 import { useProviderFilterMenu } from "./RequestLogsFilter/ProviderFilter";
@@ -17,15 +18,11 @@ const RequestLogsPage: FC = () => {
 	const { permissions } = useAuthenticated();
 	const { entitlements } = useDashboard();
 
-	// Users are allowed to view their own request logs via the API,
-	// but this page is only visible if the feature is enabled and the user
-	// has the `viewAnyAIBridgeInterception` permission.
-	// (as its defined in the Admin settings dropdown).
-	const isEntitled =
-		entitlements.features.aibridge.entitlement === "entitled" ||
-		entitlements.features.aibridge.entitlement === "grace_period";
-	const isEnabled = entitlements.features.aibridge.enabled;
-	const hasPermission = permissions.viewAnyAIBridgeInterception;
+	const { isEntitled, isEnabled, hasPermission } = getAIBridgePermissions(
+		entitlements,
+		permissions,
+	);
+
 	const canViewRequestLogs = isEntitled && hasPermission;
 
 	const [searchParams, setSearchParams] = useSearchParams();

site/src/pages/AIBridgePage/RequestLogsPage/RequestLogsPageView.tsx

@@ -1,6 +1,6 @@
 import type { ComponentProps, FC } from "react";
 import type { AIBridgeInterception } from "#/api/typesGenerated";
-import { Alert, AlertDescription, AlertTitle } from "#/components/Alert/Alert";
+import { Alert } from "#/components/Alert/Alert";
 import { Link } from "#/components/Link/Link";
 import {
 	PaginationContainer,
@@ -16,7 +16,7 @@ import {
 } from "#/components/Table/Table";
 import { TableEmpty } from "#/components/TableEmpty/TableEmpty";
 import { TableLoader } from "#/components/TableLoader/TableLoader";
-import { docs } from "#/utils/docs";
+import { AIBridgeSetupAlert } from "../AIBridgeSetupAlert";
 import { RequestLogsFilter } from "./RequestLogsFilter/RequestLogsFilter";
 import { RequestLogsRow } from "./RequestLogsRow/RequestLogsRow";
 
@@ -42,21 +42,7 @@ export const RequestLogsPageView: FC<RequestLogsPageViewProps> = ({
 	}
 
 	if (!isRequestLogsEnabled) {
-		return (
-			<Alert className="mb-12" severity="warning" prominent>
-				<AlertTitle>
-					AI Bridge is included in your license, but not set up yet.
-				</AlertTitle>
-				<AlertDescription>
-					You have access to AI Governance, but it still needs to be setup.
-					Check out the{" "}
-					<Link href={docs("/ai-coder/ai-bridge")} target="_blank">
-						AI Bridge
-					</Link>{" "}
-					documentation to get started.
-				</AlertDescription>
-			</Alert>
-		);
+		return <AIBridgeSetupAlert />;
 	}
 
 	return (

site/src/pages/AIBridgePage/SessionThreadsPage/SessionThreadsPage.tsx

@@ -2,14 +2,29 @@ import type { FC } from "react";
 import { useInfiniteQuery } from "react-query";
 import { useParams } from "react-router";
 import { infiniteSessionThreads } from "#/api/queries/aiBridge";
+import { useAuthenticated } from "#/hooks/useAuthenticated";
+import { useDashboard } from "#/modules/dashboard/useDashboard";
+import { RequirePermission } from "#/modules/permissions/RequirePermission";
+import { pageTitle } from "#/utils/page";
+import { getAIBridgePermissions } from "../getAIBridgePermissions";
 import { SessionThreadsPageView } from "./SessionThreadsPageView";
 
 const SessionThreadsPage: FC = () => {
+	const { permissions } = useAuthenticated();
+	const { entitlements } = useDashboard();
+
+	const { isEntitled, isEnabled, hasPermission } = getAIBridgePermissions(
+		entitlements,
+		permissions,
+	);
+
+	const canViewSessionThreads = isEntitled && hasPermission;
+
 	const { sessionId } = useParams() as { sessionId: string };
 
 	const sessionQuery = useInfiniteQuery({
 		...infiniteSessionThreads(sessionId),
-		enabled: !!sessionId,
+		enabled: canViewSessionThreads,
 	});
 
 	const firstPage = sessionQuery.data?.pages[0];
@@ -17,14 +32,20 @@ const SessionThreadsPage: FC = () => {
 		sessionQuery.data?.pages.flatMap((page) => page.threads) ?? [];
 
 	return (
-		<SessionThreadsPageView
-			session={firstPage}
-			threads={allThreads}
-			loading={sessionQuery.isLoading}
-			hasNextPage={sessionQuery.hasNextPage}
-			isFetchingNextPage={sessionQuery.isFetchingNextPage}
-			onFetchNextPage={sessionQuery.fetchNextPage}
-		/>
+		<RequirePermission isFeatureVisible={hasPermission}>
+			<title>{pageTitle("Session Threads", "AI Bridge")}</title>
+
+			<SessionThreadsPageView
+				session={firstPage}
+				threads={allThreads}
+				loading={sessionQuery.isLoading}
+				hasNextPage={sessionQuery.hasNextPage}
+				isFetchingNextPage={sessionQuery.isFetchingNextPage}
+				onFetchNextPage={sessionQuery.fetchNextPage}
+				isAISessionsEnabled={isEnabled}
+				isAISessionsEntitled={isEntitled}
+			/>
+		</RequirePermission>
 	);
 };
 

site/src/pages/AIBridgePage/SessionThreadsPage/SessionThreadsPageView.tsx

@@ -7,12 +7,14 @@ import type {
 } from "#/api/typesGenerated";
 import { Button } from "#/components/Button/Button";
 import { Loader } from "#/components/Loader/Loader";
+import { PaywallAIGovernance } from "#/components/Paywall/PaywallAIGovernance";
 import {
 	Tooltip,
 	TooltipContent,
 	TooltipProvider,
 	TooltipTrigger,
 } from "#/components/Tooltip/Tooltip";
+import { AIBridgeSetupAlert } from "../AIBridgeSetupAlert";
 import { SessionSummaryTable } from "./SessionSummaryTable";
 import { SessionTimeline } from "./SessionTimeline/SessionTimeline";
 
@@ -43,6 +45,8 @@ interface SessionThreadsPageViewProps {
 	hasNextPage: boolean;
 	isFetchingNextPage: boolean;
 	onFetchNextPage: () => void;
+	isAISessionsEnabled: boolean;
+	isAISessionsEntitled: boolean;
 }
 
 export const SessionThreadsPageView: FC<SessionThreadsPageViewProps> = ({
@@ -52,7 +56,17 @@ export const SessionThreadsPageView: FC<SessionThreadsPageViewProps> = ({
 	hasNextPage,
 	isFetchingNextPage,
 	onFetchNextPage,
+	isAISessionsEnabled,
+	isAISessionsEntitled,
 }) => {
+	if (!isAISessionsEntitled) {
+		return <PaywallAIGovernance />;
+	}
+
+	if (!isAISessionsEnabled) {
+		return <AIBridgeSetupAlert />;
+	}
+
 	// calculate the total number of tool calls across all loaded threads
 	const toolCallCount = threads.reduce(
 		(acc, thread) => acc + (thread.agentic_actions?.length ?? 0),

site/src/pages/AIBridgePage/getAIBridgePermissions.ts

@@ -0,0 +1,23 @@
+import type { Entitlements } from "#/api/typesGenerated";
+import type { Permissions } from "#/modules/permissions";
+
+// Users are allowed to view their own request logs via the API,
+// but an AI Bridge page is only visible if the feature is enabled and
+// the user has the `viewAnyAIBridgeInterception` permission. (as it's
+// defined in the Admin settings dropdown).
+export const getAIBridgePermissions = (
+	entitlements: Entitlements,
+	permissions: Permissions,
+) => {
+	// the user is entitled if they have either "entitled" or "grace_period"
+	// status for the AI Bridge feature
+	const isEntitled =
+		entitlements.features.aibridge.entitlement === "entitled" ||
+		entitlements.features.aibridge.entitlement === "grace_period";
+	// the feature is enabled if it's toggled on in the admin settings
+	const isEnabled = entitlements.features.aibridge.enabled;
+	// the user has permission if they have the `viewAnyAIBridgeInterception` permission
+	const hasPermission = permissions.viewAnyAIBridgeInterception;
+
+	return { isEntitled, isEnabled, hasPermission };
+};