chore: split trial models to a single API (#36796)

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>

web/__tests__/utils/mock-system-features.tsx

@@ -6,28 +6,39 @@ import { render, renderHook } from '@testing-library/react'
 import { consoleQuery } from '@/service/client'
 import { defaultSystemFeatures } from '@/types/feature'
 
-type DeepPartial<T> = T extends Array<infer U>
-  ? Array<U>
-  : T extends object
-    ? { [K in keyof T]?: DeepPartial<T[K]> }
-    : T
-
 type QueryKeyProvider = {
   queryKey: () => readonly unknown[]
 }
 
+type TrialModelsQueryProvider = {
+  get?: QueryKeyProvider
+}
+
 type AppDslVersionQueryProvider = {
   get?: QueryKeyProvider
 }
 
+const fallbackTrialModelsQueryKey = ['console', 'trialModels', 'get'] as const
 const fallbackAppDslVersionQueryKey = ['console', 'appDslVersion', 'get'] as const
 
+const getTrialModelsQueryKey = () => {
+  const trialModelsQuery = (consoleQuery as { trialModels?: TrialModelsQueryProvider }).trialModels
+
+  return trialModelsQuery?.get?.queryKey() ?? fallbackTrialModelsQueryKey
+}
+
 const getAppDslVersionQueryKey = () => {
   const appDslVersionQuery = (consoleQuery as { appDslVersion?: AppDslVersionQueryProvider }).appDslVersion
 
   return appDslVersionQuery?.get?.queryKey() ?? fallbackAppDslVersionQueryKey
 }
 
+type DeepPartial<T> = T extends Array<infer U>
+  ? Array<U>
+  : T extends object
+    ? { [K in keyof T]?: DeepPartial<T[K]> }
+    : T
+
 const buildSystemFeatures = (
   overrides: DeepPartial<SystemFeatures> = {},
 ): SystemFeatures => {
@@ -86,6 +97,13 @@ export const seedSystemFeatures = (
   return data
 }
 
+const seedTrialModels = (
+  queryClient: QueryClient,
+  trialModels: readonly string[] = [],
+) => {
+  queryClient.setQueryData(getTrialModelsQueryKey(), { trial_models: [...trialModels] })
+}
+
 export const seedAppDslVersion = (
   queryClient: QueryClient,
   appDslVersion = '0.6.0',
@@ -101,6 +119,7 @@ type SystemFeaturesTestOptions = {
    * keep the systemFeatures query in the pending state.
    */
   systemFeatures?: DeepPartial<SystemFeatures> | null
+  trialModels?: readonly string[] | null
   /**
    * Seed the workflow clipboard DSL version query only for tests that need it.
    * Omit or pass `null` to leave it unseeded.
@@ -122,6 +141,8 @@ export const createSystemFeaturesWrapper = (
   const systemFeatures = options.systemFeatures === null
     ? null
     : seedSystemFeatures(queryClient, options.systemFeatures)
+  if (options.trialModels !== undefined && options.trialModels !== null)
+    seedTrialModels(queryClient, options.trialModels)
   if (options.appDslVersion !== undefined && options.appDslVersion !== null)
     seedAppDslVersion(queryClient, options.appDslVersion)
   const wrapper = ({ children }: { children: ReactNode }) => (
@@ -134,9 +155,10 @@ export const renderWithSystemFeatures = (
   ui: ReactElement,
   options: SystemFeaturesTestOptions & Omit<RenderOptions, 'wrapper'> = {},
 ): RenderResult & { queryClient: QueryClient, systemFeatures: SystemFeatures | null } => {
-  const { systemFeatures: sf, appDslVersion, queryClient: qc, ...renderOptions } = options
+  const { systemFeatures: sf, trialModels, appDslVersion, queryClient: qc, ...renderOptions } = options
   const { wrapper, queryClient, systemFeatures } = createSystemFeaturesWrapper({
     systemFeatures: sf,
+    trialModels,
     appDslVersion,
     queryClient: qc,
   })
@@ -148,9 +170,10 @@ export const renderHookWithSystemFeatures = <Result, Props = void>(
   callback: (props: Props) => Result,
   options: SystemFeaturesTestOptions & Omit<RenderHookOptions<Props>, 'wrapper'> = {},
 ): RenderHookResult<Result, Props> & { queryClient: QueryClient, systemFeatures: SystemFeatures | null } => {
-  const { systemFeatures: sf, appDslVersion, queryClient: qc, ...hookOptions } = options
+  const { systemFeatures: sf, trialModels, appDslVersion, queryClient: qc, ...hookOptions } = options
   const { wrapper, queryClient, systemFeatures } = createSystemFeaturesWrapper({
     systemFeatures: sf,
+    trialModels,
     appDslVersion,
     queryClient: qc,
   })

web/app/components/header/account-setting/model-provider-page/model-selector/__tests__/popup.spec.tsx

@@ -1,7 +1,6 @@
 import type { ReactElement } from 'react'
 import type { Model, ModelItem, ModelProvider } from '../../declarations'
 import type { PopupProps } from '../popup'
-import type { SystemFeatures } from '@/types/feature'
 import { Combobox } from '@langgenius/dify-ui/combobox'
 import { fireEvent, screen, waitFor } from '@testing-library/react'
 import userEvent from '@testing-library/user-event'
@@ -105,7 +104,7 @@ function PopupHarness(props: PopupTestProps) {
 }
 
 const renderPopup = (ui: ReactElement<PopupTestProps>) => renderWithSystemFeatures(ui, {
-  systemFeatures: { trial_models: mockTrialModels.current as unknown as SystemFeatures['trial_models'] },
+  trialModels: mockTrialModels.current,
 })
 
 const mockTrialCredits = vi.hoisted(() => ({

web/app/components/header/account-setting/model-provider-page/model-selector/popup.tsx

@@ -3,17 +3,18 @@ import type { ModelSelectorPreviewPayload } from './popup-item'
 import type { ModelProviderQuotaGetPaid } from '@/types/model-provider'
 import { ComboboxList } from '@langgenius/dify-ui/combobox'
 import { createPreviewCardHandle, PreviewCard, PreviewCardContent } from '@langgenius/dify-ui/preview-card'
-import { useSuspenseQuery } from '@tanstack/react-query'
+import { useQuery } from '@tanstack/react-query'
 import { useTheme } from 'next-themes'
 import { useCallback, useMemo, useState } from 'react'
 import { useTranslation } from 'react-i18next'
 import { ACCOUNT_SETTING_MODAL_ACTION, ACCOUNT_SETTING_TAB } from '@/app/components/header/account-setting/constants'
 import checkTaskStatus from '@/app/components/plugins/install-plugin/base/check-task-status'
 import useRefreshPluginList from '@/app/components/plugins/install-plugin/hooks/use-refresh-plugin-list'
+import { IS_CLOUD_EDITION } from '@/config'
 import { useModalContext } from '@/context/modal-context'
 import { useProviderContext } from '@/context/provider-context'
 import { useSearchParams } from '@/next/navigation'
-import { systemFeaturesQueryOptions } from '@/service/system-features'
+import { consoleQuery } from '@/service/client'
 import { useInstallPackageFromMarketPlace } from '@/service/use-plugins'
 import { CustomConfigurationStatusEnum, ModelFeatureEnum, ModelStatusEnum, ModelTypeEnum } from '../declarations'
 import { useLanguage, useMarketplaceAllPlugins } from '../hooks'
@@ -62,8 +63,10 @@ function Popup({
   const { refreshPluginList } = useRefreshPluginList()
   const [installingProvider, setInstallingProvider] = useState<ModelProviderQuotaGetPaid | null>(null)
   const { isExhausted: isCreditsExhausted } = useTrialCredits()
-  const { data: systemFeatures } = useSuspenseQuery(systemFeaturesQueryOptions())
-  const trialModels = systemFeatures.trial_models
+  const { data: trialModels = [] } = useQuery(consoleQuery.trialModels.get.queryOptions({
+    enabled: IS_CLOUD_EDITION,
+    select: data => data.trial_models,
+  }))
   const installedProviderMap = useMemo(() => new Map(
     modelProviders.map(provider => [provider.provider, provider]),
   ), [modelProviders])

web/app/components/header/account-setting/model-provider-page/provider-added-card/__tests__/credential-panel.spec.tsx

@@ -111,7 +111,7 @@ const createProvider = (overrides: Partial<ModelProvider> = {}): ModelProvider =
 
 const renderWithQueryClient = (provider: ModelProvider) => {
   return renderWithSystemFeatures(<CredentialPanel provider={provider} />, {
-    systemFeatures: { trial_models: ['langgenius/openai/openai'] as never },
+    trialModels: ['langgenius/openai/openai'],
   })
 }
 

web/app/components/header/account-setting/model-provider-page/provider-added-card/__tests__/quota-panel.spec.tsx

@@ -47,7 +47,7 @@ vi.mock('../use-trial-credits', () => ({
 }))
 
 const renderQuotaPanel = (ui: ReactElement) => renderWithSystemFeatures(ui, {
-  systemFeatures: mockTrialModels === undefined ? null : { trial_models: mockTrialModels as never },
+  trialModels: mockTrialModels ?? [],
 })
 
 vi.mock('../../hooks', () => ({

web/app/components/header/account-setting/model-provider-page/provider-added-card/__tests__/use-credential-panel-state.spec.ts

@@ -22,7 +22,7 @@ vi.mock('@/config', async (importOriginal) => {
 
 const renderPanelHook = (provider: ModelProvider | undefined) =>
   renderHookWithSystemFeatures(() => useCredentialPanelState(provider), {
-    systemFeatures: { trial_models: mockTrialModels as never },
+    trialModels: mockTrialModels,
   })
 
 const createProvider = (overrides: Partial<ModelProvider> = {}): ModelProvider => ({

web/app/components/header/account-setting/model-provider-page/provider-added-card/quota-panel.tsx

@@ -4,7 +4,7 @@ import type { Plugin } from '@/app/components/plugins/types'
 import type { ModelProviderQuotaGetPaid } from '@/types/model-provider'
 import { cn } from '@langgenius/dify-ui/cn'
 import { Tooltip, TooltipContent, TooltipTrigger } from '@langgenius/dify-ui/tooltip'
-import { useSuspenseQuery } from '@tanstack/react-query'
+import { useQuery } from '@tanstack/react-query'
 import { useBoolean } from 'ahooks'
 import * as React from 'react'
 import { useCallback, useEffect, useMemo, useRef, useState } from 'react'
@@ -12,8 +12,9 @@ import { useTranslation } from 'react-i18next'
 import { Infotip } from '@/app/components/base/infotip'
 import Loading from '@/app/components/base/loading'
 import InstallFromMarketplace from '@/app/components/plugins/install-plugin/install-from-marketplace'
+import { IS_CLOUD_EDITION } from '@/config'
 import useTimestamp from '@/hooks/use-timestamp'
-import { systemFeaturesQueryOptions } from '@/service/system-features'
+import { consoleQuery } from '@/service/client'
 import { formatNumber } from '@/utils/format'
 import { PreferredProviderTypeEnum } from '../declarations'
 import { useMarketplaceAllPlugins } from '../hooks'
@@ -34,8 +35,10 @@ const QuotaPanel: FC<QuotaPanelProps> = ({
 }) => {
   const { t } = useTranslation()
   const { credits, isExhausted, isLoading, nextCreditResetDate } = useTrialCredits()
-  const { data: systemFeatures } = useSuspenseQuery(systemFeaturesQueryOptions())
-  const trialModels = systemFeatures.trial_models
+  const { data: trialModels = [] } = useQuery(consoleQuery.trialModels.get.queryOptions({
+    enabled: IS_CLOUD_EDITION,
+    select: data => data.trial_models,
+  }))
   const providerMap = useMemo(() => new Map(
     providers.map(p => [p.provider, p.preferred_provider_type]),
   ), [providers])

web/app/components/header/account-setting/model-provider-page/provider-added-card/use-credential-panel-state.ts

@@ -1,7 +1,8 @@
 import type { ModelProvider } from '../declarations'
-import { useSuspenseQuery } from '@tanstack/react-query'
+import { useQuery } from '@tanstack/react-query'
 import { useCredentialStatus } from '@/app/components/header/account-setting/model-provider-page/model-auth/hooks'
-import { systemFeaturesQueryOptions } from '@/service/system-features'
+import { IS_CLOUD_EDITION } from '@/config'
+import { consoleQuery } from '@/service/client'
 import {
   PreferredProviderTypeEnum,
 } from '../declarations'
@@ -80,8 +81,10 @@ export function useCredentialPanelState(provider: ModelProvider | undefined): Cr
     current_credential_name,
   } = useCredentialStatus(provider)
 
-  const { data: systemFeatures } = useSuspenseQuery(systemFeaturesQueryOptions())
-  const trialModels = systemFeatures.trial_models
+  const { data: trialModels = [] } = useQuery(consoleQuery.trialModels.get.queryOptions({
+    enabled: IS_CLOUD_EDITION,
+    select: data => data.trial_models,
+  }))
 
   const preferredType = provider?.preferred_provider_type
 

web/types/feature.ts

@@ -1,26 +1,30 @@
-import type { ModelProviderQuotaGetPaid } from './model-provider'
+export const SSOProtocol = {
+  SAML: 'saml',
+  OIDC: 'oidc',
+  OAuth2: 'oauth2',
+} as const
 
-export enum SSOProtocol {
-  SAML = 'saml',
-  OIDC = 'oidc',
-  OAuth2 = 'oauth2',
-}
+export type SSOProtocol = typeof SSOProtocol[keyof typeof SSOProtocol]
 
-export enum LicenseStatus {
-  NONE = 'none',
-  INACTIVE = 'inactive',
-  ACTIVE = 'active',
-  EXPIRING = 'expiring',
-  EXPIRED = 'expired',
-  LOST = 'lost',
-}
+export const LicenseStatus = {
+  NONE: 'none',
+  INACTIVE: 'inactive',
+  ACTIVE: 'active',
+  EXPIRING: 'expiring',
+  EXPIRED: 'expired',
+  LOST: 'lost',
+} as const
 
-export enum InstallationScope {
-  ALL = 'all',
-  NONE = 'none',
-  OFFICIAL_ONLY = 'official_only',
-  OFFICIAL_AND_PARTNER = 'official_and_specific_partners',
-}
+export type LicenseStatus = typeof LicenseStatus[keyof typeof LicenseStatus]
+
+export const InstallationScope = {
+  ALL: 'all',
+  NONE: 'none',
+  OFFICIAL_ONLY: 'official_only',
+  OFFICIAL_AND_PARTNER: 'official_and_specific_partners',
+} as const
+
+export type InstallationScope = typeof InstallationScope[keyof typeof InstallationScope]
 
 type License = {
   status: LicenseStatus
@@ -28,7 +32,6 @@ type License = {
 }
 
 export type SystemFeatures = {
-  trial_models: ModelProviderQuotaGetPaid[]
   plugin_installation_permission: {
     plugin_installation_scope: InstallationScope
     restrict_to_marketplace_only: boolean
@@ -69,7 +72,6 @@ export type SystemFeatures = {
 }
 
 export const defaultSystemFeatures: SystemFeatures = {
-  trial_models: [],
   plugin_installation_permission: {
     plugin_installation_scope: InstallationScope.ALL,
     restrict_to_marketplace_only: false,