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
<img width="335" height="156" alt="image"
src="https://github.com/user-attachments/assets/2458cbd1-46ab-46cd-babb-09a47e163fd0"
/>


### 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

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<AssistantModel>(() => {
-    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,

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"
               >
                 <span>gpt-5</span>
-                {canAccessProModels ? (
+                {hasAccessToAdvanceModel ? (
                   selectedModel === 'gpt-5' ? (
                     <Check className="h-3.5 w-3.5" />
                   ) : null

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<ReturnType<typeof getEntitlements>>
 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 = <TData = EntitlementsData>(
   { slug }: EntitlementsVariables,
   {

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<typeof vi.fn>
   let mockGetProjectDetail: ReturnType<typeof vi.fn>
@@ -33,6 +37,7 @@ describe('ai/org-ai-details', () => {
   let mockGetProjectSettings: ReturnType<typeof vi.fn>
   let mockGetAiOptInLevel: ReturnType<typeof vi.fn>
   let mockSubscriptionHasHipaaAddon: ReturnType<typeof vi.fn>
+  let mockCheckEntitlement: ReturnType<typeof vi.fn>
 
   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 () => {

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 {

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

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: {