From d5cc70560da9cdcca19194525f0c251f87bd9f16 Mon Sep 17 00:00:00 2001 From: Ignacio Dobronich Date: Thu, 19 Feb 2026 14:28:31 -0300 Subject: [PATCH] chore: ai assistant model entitlement (#42955) ### Changes Replace hard-coded plan checks with `assistant.advance_model` entitlement for AI model access control. - `AIAssistant.tsx` & `ModelSelector.tsx`: Use `useCheckEntitlements('assistant.advance_model')` instead of `plan.id !== 'free'` **Server-side:** - `org-ai-details.ts`: Check `assistant.advance_model` entitlement via new `checkEntitlement()` helper - `entitlements-query.ts`: Add `checkEntitlement()` function with auth headers support for server-side calls ### Test 1: Free Plan (No Entitlement) - Open AI Assistant with a Free plan org - Check model selector defaults to gpt-5-mini - Click dropdown - verify gpt-5 shows "Upgrade" badge - Click gpt-5 - should redirect to billing page - Assert that you can send a message to the Assistant and that you get a response image ### Test 2: Paid Plan (Has Entitlement) - Switch to Pro/Team/Enterprise org - Check model selector defaults to gpt-5 - Click dropdown - verify gpt-5 shows checkmark (no badge) - Switch to gpt-5-mini, then back to gpt-5 - should work without redirect - Assert that you can send a message to the Assistant and that you get a response --- .../ui/AIAssistantPanel/AIAssistant.tsx | 15 +- .../ui/AIAssistantPanel/ModelSelector.tsx | 10 +- .../data/entitlements/entitlements-query.ts | 27 ++- apps/studio/lib/ai/org-ai-details.test.ts | 12 +- apps/studio/lib/ai/org-ai-details.ts | 17 +- packages/api-types/types/api.d.ts | 110 +++++++++++- packages/api-types/types/platform.d.ts | 161 +++++++++++++++++- 7 files changed, 329 insertions(+), 23 deletions(-) diff --git a/apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx b/apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx index 191c52a3da8..930fbb33a4b 100644 --- a/apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx +++ b/apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx @@ -42,6 +42,7 @@ import { import { Message } from './Message' import AlertError from '../AlertError' import { ASSISTANT_ERRORS } from './AiAssistant.constants' +import { useCheckEntitlements } from '@/hooks/misc/useCheckEntitlements' interface AIAssistantProps { initialMessages?: MessageType[] | undefined @@ -65,18 +66,24 @@ export const AIAssistant = ({ className }: AIAssistantProps) => { const state = useAiAssistantState() const { activeSidebar, closeSidebar } = useSidebarManagerSnapshot() - const isPaidPlan = selectedOrganization?.plan?.id !== 'free' + const { hasAccess: hasAccessToAdvanceModel, isLoading: isLoadingEntitlements } = + useCheckEntitlements('assistant.advance_model') const selectedModel = useMemo(() => { - const defaultModel: AssistantModel = isPaidPlan ? 'gpt-5' : 'gpt-5-mini' + // While entitlements are loading, use the stored model without enforcing access + if (isLoadingEntitlements) { + return snap.model ?? 'gpt-5-mini' + } + + const defaultModel: AssistantModel = hasAccessToAdvanceModel ? 'gpt-5' : 'gpt-5-mini' const model = snap.model ?? defaultModel - if (!isPaidPlan && model === 'gpt-5') { + if (!hasAccessToAdvanceModel && model === 'gpt-5') { return 'gpt-5-mini' } return model - }, [isPaidPlan, snap.model]) + }, [isLoadingEntitlements, hasAccessToAdvanceModel, snap.model]) const [updatedOptInSinceMCP] = useLocalStorageQuery( LOCAL_STORAGE_KEYS.AI_ASSISTANT_MCP_OPT_IN, diff --git a/apps/studio/components/ui/AIAssistantPanel/ModelSelector.tsx b/apps/studio/components/ui/AIAssistantPanel/ModelSelector.tsx index aa050cb9c96..d7c66b64d13 100644 --- a/apps/studio/components/ui/AIAssistantPanel/ModelSelector.tsx +++ b/apps/studio/components/ui/AIAssistantPanel/ModelSelector.tsx @@ -1,7 +1,5 @@ import { Check, ChevronsUpDown } from 'lucide-react' import { useState } from 'react' - -import { useIsFeatureEnabled } from 'hooks/misc/useIsFeatureEnabled' import { useSelectedOrganizationQuery } from 'hooks/misc/useSelectedOrganization' import { useRouter } from 'next/router' import { @@ -18,6 +16,7 @@ import { TooltipTrigger, Tooltip, } from 'ui' +import { useCheckEntitlements } from '@/hooks/misc/useCheckEntitlements' interface ModelSelectorProps { selectedModel: 'gpt-5' | 'gpt-5-mini' @@ -27,16 +26,17 @@ interface ModelSelectorProps { export const ModelSelector = ({ selectedModel, onSelectModel }: ModelSelectorProps) => { const router = useRouter() const { data: organization } = useSelectedOrganizationQuery() + const { hasAccess: hasAccessToAdvanceModel, isLoading: isLoadingEntitlements } = + useCheckEntitlements('assistant.advance_model') const [open, setOpen] = useState(false) - const canAccessProModels = organization?.plan?.id !== 'free' const slug = organization?.slug ?? '_' const upgradeHref = `/org/${slug ?? '_'}/billing?panel=subscriptionPlan&source=ai-assistant-model` const handleSelectModel = (model: 'gpt-5' | 'gpt-5-mini') => { - if (model === 'gpt-5' && !canAccessProModels) { + if (model === 'gpt-5' && !hasAccessToAdvanceModel) { setOpen(false) void router.push(upgradeHref) return @@ -75,7 +75,7 @@ export const ModelSelector = ({ selectedModel, onSelectModel }: ModelSelectorPro className="flex justify-between" > gpt-5 - {canAccessProModels ? ( + {hasAccessToAdvanceModel ? ( selectedModel === 'gpt-5' ? ( ) : null diff --git a/apps/studio/data/entitlements/entitlements-query.ts b/apps/studio/data/entitlements/entitlements-query.ts index a023c4cba64..bf703f6763e 100644 --- a/apps/studio/data/entitlements/entitlements-query.ts +++ b/apps/studio/data/entitlements/entitlements-query.ts @@ -17,12 +17,17 @@ export type EntitlementConfig = export type Entitlement = components['schemas']['ListEntitlementsResponse']['entitlements'][0] export type EntitlementType = Entitlement['type'] -export async function getEntitlements({ slug }: EntitlementsVariables, signal?: AbortSignal) { +export async function getEntitlements( + { slug }: EntitlementsVariables, + signal?: AbortSignal, + headers?: HeadersInit +) { if (!slug) throw new Error('slug is required') const { data, error } = await get('/platform/organizations/{slug}/entitlements', { params: { path: { slug } }, signal, + ...(headers && { headers }), }) if (error) handleError(error) @@ -32,6 +37,26 @@ export async function getEntitlements({ slug }: EntitlementsVariables, signal?: export type EntitlementsData = Awaited> export type EntitlementsError = ResponseError +/** + * Helper to check a specific entitlement for an organization + * For client-side components, please using useCheckEntitlements hook instead + */ +export async function checkEntitlement( + slug: string, + featureKey: string, + signal?: AbortSignal, + headers?: HeadersInit +): Promise<{ hasAccess: boolean; entitlement?: Entitlement }> { + const entitlements = await getEntitlements({ slug }, signal, headers) + + const entitlement = entitlements.entitlements.find((e) => e.feature.key === (featureKey as any)) + + return { + hasAccess: entitlement?.hasAccess ?? false, + entitlement, + } +} + export const useEntitlementsQuery = ( { slug }: EntitlementsVariables, { diff --git a/apps/studio/lib/ai/org-ai-details.test.ts b/apps/studio/lib/ai/org-ai-details.test.ts index 6098b80b08f..d36c4168085 100644 --- a/apps/studio/lib/ai/org-ai-details.test.ts +++ b/apps/studio/lib/ai/org-ai-details.test.ts @@ -26,6 +26,10 @@ vi.mock('components/interfaces/Billing/Subscription/Subscription.utils', () => ( subscriptionHasHipaaAddon: vi.fn(), })) +vi.mock('data/entitlements/entitlements-query', () => ({ + checkEntitlement: vi.fn(), +})) + describe('ai/org-ai-details', () => { let mockGetOrganizations: ReturnType let mockGetProjectDetail: ReturnType @@ -33,6 +37,7 @@ describe('ai/org-ai-details', () => { let mockGetProjectSettings: ReturnType let mockGetAiOptInLevel: ReturnType let mockSubscriptionHasHipaaAddon: ReturnType + let mockCheckEntitlement: ReturnType beforeEach(async () => { vi.clearAllMocks() @@ -45,6 +50,7 @@ describe('ai/org-ai-details', () => { const subscriptionUtils = await import( 'components/interfaces/Billing/Subscription/Subscription.utils' ) + const entitlementsQuery = await import('data/entitlements/entitlements-query') mockGetOrganizations = vi.mocked(orgsQuery.getOrganizations) mockGetProjectDetail = vi.mocked(projectQuery.getProjectDetail) @@ -52,11 +58,13 @@ describe('ai/org-ai-details', () => { mockGetProjectSettings = vi.mocked(settingsQuery.getProjectSettings) mockGetAiOptInLevel = vi.mocked(aiHook.getAiOptInLevel) mockSubscriptionHasHipaaAddon = vi.mocked(subscriptionUtils.subscriptionHasHipaaAddon) + mockCheckEntitlement = vi.mocked(entitlementsQuery.checkEntitlement) // Default mocks for subscription/settings (no HIPAA) mockGetOrgSubscription.mockResolvedValue({ addons: [] }) mockGetProjectSettings.mockResolvedValue({ is_sensitive: false }) mockSubscriptionHasHipaaAddon.mockReturnValue(false) + mockCheckEntitlement.mockResolvedValue({ hasAccess: false }) }) describe('getOrgAIDetails', () => { @@ -140,6 +148,7 @@ describe('ai/org-ai-details', () => { mockGetOrganizations.mockResolvedValue([mockOrg]) mockGetProjectDetail.mockResolvedValue(mockProject) mockGetAiOptInLevel.mockReturnValue('full') + mockCheckEntitlement.mockResolvedValue({ hasAccess: true }) const result = await getOrgAIDetails({ orgSlug: 'test-org', @@ -261,6 +270,7 @@ describe('ai/org-ai-details', () => { mockGetOrganizations.mockResolvedValue(mockOrgs) mockGetProjectDetail.mockResolvedValue(mockProject) mockGetAiOptInLevel.mockReturnValue('full') + mockCheckEntitlement.mockResolvedValue({ hasAccess: true }) const result = await getOrgAIDetails({ orgSlug: 'test-org', @@ -268,7 +278,7 @@ describe('ai/org-ai-details', () => { projectRef: 'test-project', }) - expect(result.isLimited).toBe(false) // Pro plan + expect(result.isLimited).toBe(false) // Has advance model entitlement }) it('should return isHipaaEnabled true when subscription has HIPAA addon and project is sensitive', async () => { diff --git a/apps/studio/lib/ai/org-ai-details.ts b/apps/studio/lib/ai/org-ai-details.ts index 5ce932d45a6..659c0d12aec 100644 --- a/apps/studio/lib/ai/org-ai-details.ts +++ b/apps/studio/lib/ai/org-ai-details.ts @@ -4,6 +4,7 @@ import { getOrganizations } from 'data/organizations/organizations-query' import { getProjectDetail } from 'data/projects/project-detail-query' import { getOrgSubscription } from 'data/subscriptions/org-subscription-query' import { getAiOptInLevel } from 'hooks/misc/useOrgOptedIntoAi' +import { checkEntitlement } from 'data/entitlements/entitlements-query' export const getOrgAIDetails = async ({ orgSlug, @@ -19,12 +20,14 @@ export const getOrgAIDetails = async ({ ...(authorization && { Authorization: authorization }), } - const [organizations, selectedProject, subscription, projectSettings] = await Promise.all([ - getOrganizations({ headers }), - getProjectDetail({ ref: projectRef }, undefined, headers), - getOrgSubscription({ orgSlug }, undefined, headers), - getProjectSettings({ projectRef }, undefined, headers), - ]) + const [organizations, selectedProject, subscription, projectSettings, advanceModelAccess] = + await Promise.all([ + getOrganizations({ headers }), + getProjectDetail({ ref: projectRef }, undefined, headers), + getOrgSubscription({ orgSlug }, undefined, headers), + getProjectSettings({ projectRef }, undefined, headers), + checkEntitlement(orgSlug, 'assistant.advance_model', undefined, headers), + ]) const selectedOrg = organizations.find((org) => org.slug === orgSlug) @@ -34,7 +37,7 @@ export const getOrgAIDetails = async ({ } const aiOptInLevel = getAiOptInLevel(selectedOrg?.opt_in_tags) - const isLimited = selectedOrg?.plan.id === 'free' + const isLimited = !advanceModelAccess.hasAccess const isHipaaEnabled = subscriptionHasHipaaAddon(subscription) && !!projectSettings?.is_sensitive return { diff --git a/packages/api-types/types/api.d.ts b/packages/api-types/types/api.d.ts index cfaf84f2f8b..32207dabcdb 100644 --- a/packages/api-types/types/api.d.ts +++ b/packages/api-types/types/api.d.ts @@ -996,6 +996,23 @@ export interface paths { patch: operations['v1-update-realtime-config'] trace?: never } + '/v1/projects/{ref}/config/realtime/shutdown': { + parameters: { + query?: never + header?: never + path?: never + cookie?: never + } + get?: never + put?: never + /** Shutdowns realtime connections for a project */ + post: operations['v1-shutdown-realtime'] + delete?: never + options?: never + head?: never + patch?: never + trace?: never + } '/v1/projects/{ref}/config/storage': { parameters: { query?: never @@ -2275,6 +2292,7 @@ export interface components { security_captcha_secret: string | null security_manual_linking_enabled: boolean | null security_refresh_token_reuse_interval: number | null + security_sb_forwarded_for_enabled: boolean | null security_update_password_require_reauthentication: boolean | null sessions_inactivity_timeout: number | null sessions_single_per_user: boolean | null @@ -3512,7 +3530,26 @@ export interface components { type: 'smartGroup' }[] specific: { - code: string + /** @enum {string} */ + code: + | 'us-east-1' + | 'us-east-2' + | 'us-west-1' + | 'us-west-2' + | 'ap-southeast-1' + | 'ap-northeast-1' + | 'ap-northeast-2' + | 'ap-east-1' + | 'ap-southeast-2' + | 'eu-west-1' + | 'eu-west-2' + | 'eu-west-3' + | 'eu-north-1' + | 'eu-central-1' + | 'eu-central-2' + | 'ca-central-1' + | 'ap-south-1' + | 'sa-east-1' name: string /** @enum {string} */ provider: 'AWS' | 'FLY' | 'AWS_K8S' | 'AWS_NIMBUS' @@ -3531,7 +3568,26 @@ export interface components { type: 'smartGroup' } specific: { - code: string + /** @enum {string} */ + code: + | 'us-east-1' + | 'us-east-2' + | 'us-west-1' + | 'us-west-2' + | 'ap-southeast-1' + | 'ap-northeast-1' + | 'ap-northeast-2' + | 'ap-east-1' + | 'ap-southeast-2' + | 'eu-west-1' + | 'eu-west-2' + | 'eu-west-3' + | 'eu-north-1' + | 'eu-central-1' + | 'eu-central-2' + | 'ca-central-1' + | 'ap-south-1' + | 'sa-east-1' name: string /** @enum {string} */ provider: 'AWS' | 'FLY' | 'AWS_K8S' | 'AWS_NIMBUS' @@ -3962,6 +4018,7 @@ export interface components { security_captcha_secret?: string | null security_manual_linking_enabled?: boolean | null security_refresh_token_reuse_interval?: number | null + security_sb_forwarded_for_enabled?: boolean | null security_update_password_require_reauthentication?: boolean | null sessions_inactivity_timeout?: number | null sessions_single_per_user?: boolean | null @@ -8475,6 +8532,55 @@ export interface operations { } } } + 'v1-shutdown-realtime': { + parameters: { + query?: never + header?: never + path: { + /** @description Project ref */ + ref: string + } + cookie?: never + } + requestBody?: never + responses: { + /** @description Realtime connections shutdown successfully */ + 204: { + headers: { + [name: string]: unknown + } + content?: never + } + /** @description Unauthorized */ + 401: { + headers: { + [name: string]: unknown + } + content?: never + } + /** @description Forbidden action */ + 403: { + headers: { + [name: string]: unknown + } + content?: never + } + /** @description Tenant not found */ + 404: { + headers: { + [name: string]: unknown + } + content?: never + } + /** @description Rate limit exceeded */ + 429: { + headers: { + [name: string]: unknown + } + content?: never + } + } + } 'v1-get-storage-config': { parameters: { query?: never diff --git a/packages/api-types/types/platform.d.ts b/packages/api-types/types/platform.d.ts index a8faeb2caae..f695cc39e5b 100644 --- a/packages/api-types/types/platform.d.ts +++ b/packages/api-types/types/platform.d.ts @@ -2662,6 +2662,23 @@ export interface paths { patch: operations['RealtimeConfigController_updateConfig'] trace?: never } + '/platform/projects/{ref}/config/realtime/shutdown': { + parameters: { + query?: never + header?: never + path?: never + cookie?: never + } + get?: never + put?: never + /** Shutdowns realtime connections for a project */ + post: operations['RealtimeConfigController_shutdown'] + delete?: never + options?: never + head?: never + patch?: never + trace?: never + } '/platform/projects/{ref}/config/secrets': { parameters: { query?: never @@ -4456,6 +4473,23 @@ export interface paths { patch?: never trace?: never } + '/platform/telemetry/stream': { + parameters: { + query?: never + header?: never + path?: never + cookie?: never + } + /** Stream telemetry events (local dev only) */ + get: operations['TelemetryStreamController_streamEvents'] + put?: never + post?: never + delete?: never + options?: never + head?: never + patch?: never + trace?: never + } '/platform/tos/fly': { parameters: { query?: never @@ -5776,7 +5810,6 @@ export interface components { amount_cents: number /** Format: date-time */ credits_expire_at: string | null - credits_expiry_days: number | null } CreditsTopUpRequest: { address?: { @@ -6854,6 +6887,7 @@ export interface components { SECURITY_CAPTCHA_SECRET: string SECURITY_MANUAL_LINKING_ENABLED: boolean SECURITY_REFRESH_TOKEN_REUSE_INTERVAL: number + SECURITY_SB_FORWARDED_FOR_ENABLED: boolean SECURITY_UPDATE_PASSWORD_REQUIRE_REAUTHENTICATION: boolean SESSIONS_INACTIVITY_TIMEOUT: number SESSIONS_SINGLE_PER_USER: boolean @@ -7073,8 +7107,10 @@ export interface components { | 'instances.compute_update_available_sizes' | 'instances.read_replicas' | 'instances.disk_modifications' + | 'instances.high_availability' | 'replication.etl' | 'storage.max_file_size' + | 'storage.max_file_size.configurable' | 'storage.image_transformations' | 'storage.vector_buckets' | 'storage.iceberg_catalog' @@ -7115,6 +7151,7 @@ export interface components { | 'security.member_roles' | 'project_pausing' | 'project_cloning' + | 'assistant.advance_model' /** @enum {string} */ type: 'boolean' | 'numeric' | 'set' } @@ -8349,7 +8386,26 @@ export interface components { type: 'smartGroup' }[] specific: { - code: string + /** @enum {string} */ + code: + | 'us-east-1' + | 'us-east-2' + | 'us-west-1' + | 'us-west-2' + | 'ap-southeast-1' + | 'ap-northeast-1' + | 'ap-northeast-2' + | 'ap-east-1' + | 'ap-southeast-2' + | 'eu-west-1' + | 'eu-west-2' + | 'eu-west-3' + | 'eu-north-1' + | 'eu-central-1' + | 'eu-central-2' + | 'ca-central-1' + | 'ap-south-1' + | 'sa-east-1' name: string /** @enum {string} */ provider: 'AWS' | 'FLY' | 'AWS_K8S' | 'AWS_NIMBUS' @@ -8368,7 +8424,26 @@ export interface components { type: 'smartGroup' } specific: { - code: string + /** @enum {string} */ + code: + | 'us-east-1' + | 'us-east-2' + | 'us-west-1' + | 'us-west-2' + | 'ap-southeast-1' + | 'ap-northeast-1' + | 'ap-northeast-2' + | 'ap-east-1' + | 'ap-southeast-2' + | 'eu-west-1' + | 'eu-west-2' + | 'eu-west-3' + | 'eu-north-1' + | 'eu-central-1' + | 'eu-central-2' + | 'ca-central-1' + | 'ap-south-1' + | 'sa-east-1' name: string /** @enum {string} */ provider: 'AWS' | 'FLY' | 'AWS_K8S' | 'AWS_NIMBUS' @@ -9933,6 +10008,7 @@ export interface components { SECURITY_CAPTCHA_SECRET?: string | null SECURITY_MANUAL_LINKING_ENABLED?: boolean | null SECURITY_REFRESH_TOKEN_REUSE_INTERVAL?: number | null + SECURITY_SB_FORWARDED_FOR_ENABLED?: boolean | null SECURITY_UPDATE_PASSWORD_REQUIRE_REAUTHENTICATION?: boolean | null SESSIONS_INACTIVITY_TIMEOUT?: number | null SESSIONS_SINGLE_PER_USER?: boolean | null @@ -17046,8 +17122,10 @@ export interface operations { | 'instances.compute_update_available_sizes' | 'instances.read_replicas' | 'instances.disk_modifications' + | 'instances.high_availability' | 'replication.etl' | 'storage.max_file_size' + | 'storage.max_file_size.configurable' | 'storage.image_transformations' | 'storage.vector_buckets' | 'storage.iceberg_catalog' @@ -17088,6 +17166,7 @@ export interface operations { | 'security.member_roles' | 'project_pausing' | 'project_cloning' + | 'assistant.advance_model' } header?: never path?: never @@ -19845,6 +19924,55 @@ export interface operations { } } } + RealtimeConfigController_shutdown: { + parameters: { + query?: never + header?: never + path: { + /** @description Project ref */ + ref: string + } + cookie?: never + } + requestBody?: never + responses: { + /** @description Realtime connections shutdown successfully */ + 204: { + headers: { + [name: string]: unknown + } + content?: never + } + /** @description Unauthorized */ + 401: { + headers: { + [name: string]: unknown + } + content?: never + } + /** @description Forbidden action */ + 403: { + headers: { + [name: string]: unknown + } + content?: never + } + /** @description Tenant not found */ + 404: { + headers: { + [name: string]: unknown + } + content?: never + } + /** @description Rate limit exceeded */ + 429: { + headers: { + [name: string]: unknown + } + content?: never + } + } + } SecretsConfigController_updateConfig: { parameters: { query?: never @@ -26436,6 +26564,33 @@ export interface operations { } } } + TelemetryStreamController_streamEvents: { + parameters: { + query?: { + session_id?: string + } + header?: never + path?: never + cookie?: never + } + requestBody?: never + responses: { + /** @description SSE stream of telemetry events */ + 200: { + headers: { + [name: string]: unknown + } + content?: never + } + /** @description Only available in local development */ + 404: { + headers: { + [name: string]: unknown + } + content?: never + } + } + } TermsOfServiceController_flyTosAccepted: { parameters: { query: {