diff --git a/app/(chat)/api/models/route.ts b/app/(chat)/api/models/route.ts index c1ada40aa..8df54ecda 100644 --- a/app/(chat)/api/models/route.ts +++ b/app/(chat)/api/models/route.ts @@ -9,6 +9,7 @@ import { getUserById, } from "@/lib/db/queries"; import { getAiModelFrontendAvailability } from "@/lib/ai/model-frontend-availability"; +import { canAccessHiddenAiModels } from "@/lib/ai/model-visibility"; import { buildVideoToolCapability } from "@/lib/ai/video-tool-capability"; import { isImageGenerationModelDescriptor, @@ -146,8 +147,11 @@ export async function GET(request: Request) { } const currentUser = await getUserById({ id: session.user.id }); - const allowHiddenModels = - session.user.type === "admin" || Boolean(currentUser?.isHighVip); + const allowHiddenModels = canAccessHiddenAiModels({ + sessionUserType: session.user.type, + databaseUserRole: currentUser?.role, + isHighVip: currentUser?.isHighVip, + }); const includeUnavailable = new URL(request.url).searchParams.get("includeUnavailable") === "1"; diff --git a/app/api/style-pro/model-options/route.ts b/app/api/style-pro/model-options/route.ts index 26b5f2b13..44d82d1b2 100644 --- a/app/api/style-pro/model-options/route.ts +++ b/app/api/style-pro/model-options/route.ts @@ -1,6 +1,7 @@ import { NextResponse } from "next/server"; import { auth } from "@/app/(auth)/auth"; import { getModelGenerationCreditCost } from "@/lib/ai/image-model-pricing"; +import { canAccessHiddenAiModels } from "@/lib/ai/model-visibility"; import { getSystemSetting, getUserById } from "@/lib/db/queries"; import { findStyleProImageModelByRef, @@ -27,8 +28,11 @@ export async function GET() { } const currentUser = await getUserById({ id: session.user.id }); - const allowHiddenModels = - session.user.type === "admin" || Boolean(currentUser?.isHighVip); + const allowHiddenModels = canAccessHiddenAiModels({ + sessionUserType: session.user.type, + databaseUserRole: currentUser?.role, + isHighVip: currentUser?.isHighVip, + }); const [displayEntries, entries, realisticModelEntries, setting] = await Promise.all([ @@ -67,14 +71,21 @@ export async function GET() { }) : 0; - return NextResponse.json({ - success: true, - defaultModelId: defaultEntry?.model.id || null, - realisticUnitCost, - models: sortStyleProImageModelEntries(displayEntries).map( - toPublicStyleProModelOption - ), - }); + return NextResponse.json( + { + success: true, + defaultModelId: defaultEntry?.model.id || null, + realisticUnitCost, + models: sortStyleProImageModelEntries(displayEntries).map( + toPublicStyleProModelOption + ), + }, + { + headers: { + "Cache-Control": "private, no-store, max-age=0", + }, + } + ); } catch (error) { console.error("[Style Pro Model Options] Failed:", error); return NextResponse.json( diff --git a/components/style-pro/StyleProWorkspace.tsx b/components/style-pro/StyleProWorkspace.tsx index 5fc2d065f..dd7cfcc4a 100644 --- a/components/style-pro/StyleProWorkspace.tsx +++ b/components/style-pro/StyleProWorkspace.tsx @@ -1001,7 +1001,12 @@ export function StyleProWorkspace({ const { data: configData, isLoading: configLoading } = useSWR("/api/style-pro/config", fetcher); const { data: modelsData, isLoading: modelsLoading } = - useSWR("/api/style-pro/model-options", fetcher); + useSWR("/api/style-pro/model-options", fetcher, { + dedupingInterval: 5_000, + refreshInterval: 30_000, + revalidateOnFocus: true, + revalidateOnReconnect: true, + }); const config = configData?.config; const modelOptions = useMemo( () => (modelsData?.models || []).filter((model) => Boolean(model.id)), diff --git a/lib/ai/model-classification.ts b/lib/ai/model-classification.ts index b8928bfc7..d3137b0c1 100644 --- a/lib/ai/model-classification.ts +++ b/lib/ai/model-classification.ts @@ -178,6 +178,18 @@ export function isImageGenerationModelDescriptor(params: { return STATIC_IMAGE_PROVIDER_TYPES.has(providerType); } +export function isStillImageGenerationModelDescriptor(params: { + providerType?: string | null; + modelId?: string | null; + capabilities?: unknown; + supportsImageGen?: boolean | null; +}): boolean { + return ( + isImageGenerationModelDescriptor(params) && + !isVideoModelDescriptor(params) + ); +} + export function modelSupportsTextToImageDescriptor(params: { providerType?: string | null; modelId?: string | null; diff --git a/lib/ai/model-visibility.ts b/lib/ai/model-visibility.ts new file mode 100644 index 000000000..7eb4e15be --- /dev/null +++ b/lib/ai/model-visibility.ts @@ -0,0 +1,19 @@ +export type HiddenModelAccessInput = { + sessionUserType?: string | null; + databaseUserRole?: string | null; + isHighVip?: boolean | null; +}; + +const HIDDEN_MODEL_ADMIN_ROLES = new Set(["admin", "agent_admin"]); + +export function canAccessHiddenAiModels({ + sessionUserType, + databaseUserRole, + isHighVip, +}: HiddenModelAccessInput) { + return ( + isHighVip === true || + HIDDEN_MODEL_ADMIN_ROLES.has(String(databaseUserRole || "").trim()) || + HIDDEN_MODEL_ADMIN_ROLES.has(String(sessionUserType || "").trim()) + ); +} diff --git a/lib/style-pro/model-options.ts b/lib/style-pro/model-options.ts index dc39b7cac..800ac3fde 100644 --- a/lib/style-pro/model-options.ts +++ b/lib/style-pro/model-options.ts @@ -9,6 +9,7 @@ import { getAiModelFrontendAvailability, type ModelFrontendAvailabilityStatus, } from "@/lib/ai/model-frontend-availability"; +import { isStillImageGenerationModelDescriptor } from "@/lib/ai/model-classification"; import { getEnabledAiModelsWithProvider, getFrontendDisplayAiModelsWithProvider, @@ -37,127 +38,15 @@ export type PublicStyleProModelOption = { disabledReason: string | null; }; -const VIDEO_PROVIDER_TYPES = new Set([ - "grsai-video", - "qianwen-video", - "doubao-video", - "ruxa-video", - "volcano-video", - "fal-video", - "runway", -]); - -const VIDEO_MODEL_KEYWORDS = [ - "/generate/video/", - "video", - "veo", - "sora", - "kling", - "hailuo", - "seedance", - "runway", - "gen-3", - "gen-4", - "gen 3", - "gen 4", - "wan-2.", - "wan 2.", - "t2v", - "i2v", -]; - -const EXPLICIT_IMAGE_MODEL_KEYWORDS = [ - "/generate/image/", - "gpt-image", - "sora_image", - "qwen-image", - "seedream", - "nano-banana", - "banana", - "flux", - "z-image", - "image-edit", - "gen4_image", - "gen-4 image", -]; - -function hasVideoGenerationConfig(value: unknown) { - if (!value || typeof value !== "object") return false; - const config = value as Record; - return ( - "creditCostPerVideo" in config || - "creditCostPerVideoByDuration" in config || - "creditCostPerVideoByResolution" in config || - "videoDurationOptions" in config || - "videoResolutionOptions" in config - ); -} - -function hasVideoCapabilities(value: unknown) { - if (!value || typeof value !== "object") return false; - const serialized = JSON.stringify(value).toLowerCase(); - return ( - serialized.includes("text-to-video") || - serialized.includes("image-to-video") - ); -} - -function isExplicitImageGenerationModel( - entry: StyleProModelEntry | StyleProDisplayModelEntry -) { - const searchable = [ - entry.model.modelId, - entry.model.displayName, - entry.model.description, - ] - .filter(Boolean) - .join(" ") - .toLowerCase(); - - if ( - searchable.includes("image-to-video") || - searchable.includes("text-to-video") || - searchable.includes("i2v") || - searchable.includes("t2v") - ) { - return false; - } - - return EXPLICIT_IMAGE_MODEL_KEYWORDS.some((keyword) => - searchable.includes(keyword) - ); -} - -function isVideoLikeModel( - entry: StyleProModelEntry | StyleProDisplayModelEntry -) { - if (VIDEO_PROVIDER_TYPES.has(entry.provider.type)) return true; - if ( - hasVideoGenerationConfig(entry.model.imageGenConfig) || - hasVideoCapabilities(entry.model.capabilities) || - hasVideoCapabilities(entry.model.defaultParams) - ) { - return true; - } - const searchable = [ - entry.provider.type, - entry.provider.name, - entry.model.modelId, - entry.model.displayName, - entry.model.description, - ] - .filter(Boolean) - .join(" ") - .toLowerCase(); - return VIDEO_MODEL_KEYWORDS.some((keyword) => searchable.includes(keyword)); -} - export function isStyleProImageModel( entry: StyleProModelEntry | StyleProDisplayModelEntry ) { - if (!entry.model.supportsImageGen) return false; - if (isExplicitImageGenerationModel(entry)) return true; - return !isVideoLikeModel(entry); + return isStillImageGenerationModelDescriptor({ + providerType: entry.provider.type, + modelId: entry.model.modelId, + capabilities: entry.model.capabilities, + supportsImageGen: entry.model.supportsImageGen, + }); } export async function getStyleProImageModelEntries(options?: { diff --git a/playwright.helpers.config.ts b/playwright.helpers.config.ts index f1999b5a7..06cb08527 100644 --- a/playwright.helpers.config.ts +++ b/playwright.helpers.config.ts @@ -13,6 +13,7 @@ export default defineConfig({ /routes\/canvas-scale-adjust-intent\.test\.ts/, /routes\/canvas-runtime-internal-endpoint\.test\.ts/, /routes\/model-classification\.test\.ts/, + /routes\/model-visibility\.test\.ts/, /routes\/model3d-style-pro-regressions\.test\.ts/, ], fullyParallel: true, diff --git a/tests/routes/model-classification.test.ts b/tests/routes/model-classification.test.ts index 04cc84f1a..ad3dadcc9 100644 --- a/tests/routes/model-classification.test.ts +++ b/tests/routes/model-classification.test.ts @@ -1,6 +1,7 @@ import { expect, test } from "@playwright/test"; import { isImageGenerationModelDescriptor, + isStillImageGenerationModelDescriptor, isVideoModelDescriptor, } from "@/lib/ai/model-classification"; @@ -32,4 +33,22 @@ test.describe("model classification", () => { }) ).toBe(true); }); + + test("uses one image-only rule for canvas and Style Pro model pools", () => { + expect( + isStillImageGenerationModelDescriptor({ + providerType: "ruxa-image", + modelId: "google/nano-banana-2", + supportsImageGen: true, + }) + ).toBe(true); + + expect( + isStillImageGenerationModelDescriptor({ + providerType: "ruxa-video", + modelId: "kling-video", + supportsImageGen: true, + }) + ).toBe(false); + }); }); diff --git a/tests/routes/model-visibility.test.ts b/tests/routes/model-visibility.test.ts new file mode 100644 index 000000000..c97112ec0 --- /dev/null +++ b/tests/routes/model-visibility.test.ts @@ -0,0 +1,31 @@ +import { expect, test } from "@playwright/test"; +import { canAccessHiddenAiModels } from "@/lib/ai/model-visibility"; + +test.describe("hidden AI model visibility", () => { + test("uses the current database role when a JWT still contains an old role", () => { + expect( + canAccessHiddenAiModels({ + sessionUserType: "regular", + databaseUserRole: "admin", + isHighVip: false, + }) + ).toBe(true); + }); + + test("allows Style Pro administrators and high VIP users", () => { + expect( + canAccessHiddenAiModels({ databaseUserRole: "agent_admin" }) + ).toBe(true); + expect(canAccessHiddenAiModels({ isHighVip: true })).toBe(true); + }); + + test("keeps hidden provider variants out of regular user selectors", () => { + expect( + canAccessHiddenAiModels({ + sessionUserType: "regular", + databaseUserRole: "regular", + isHighVip: false, + }) + ).toBe(false); + }); +}); diff --git a/tests/routes/model3d-style-pro-regressions.test.ts b/tests/routes/model3d-style-pro-regressions.test.ts index 3b4bf3d72..56624c323 100644 --- a/tests/routes/model3d-style-pro-regressions.test.ts +++ b/tests/routes/model3d-style-pro-regressions.test.ts @@ -35,11 +35,25 @@ test.describe("3D workspace and Style Pro regressions", () => { test("uses the same hidden-model access rule in Style Pro as canvas", () => { const route = readSource("app/api/style-pro/model-options/route.ts"); + const canvasModelsRoute = readSource("app/(chat)/api/models/route.ts"); expect(route).toContain("getUserById"); - expect(route).toContain('session.user.type === "admin"'); - expect(route).toContain("Boolean(currentUser?.isHighVip)"); + expect(route).toContain("canAccessHiddenAiModels"); + expect(route).toContain("databaseUserRole: currentUser?.role"); expect(route).toContain("includeHidden: allowHiddenModels"); + expect(route).toContain('"Cache-Control": "private, no-store, max-age=0"'); + expect(canvasModelsRoute).toContain("canAccessHiddenAiModels"); + expect(canvasModelsRoute).toContain("databaseUserRole: currentUser?.role"); + }); + + test("refreshes Style Pro models after an administrator changes the catalog", () => { + const workspace = readSource( + "components/style-pro/StyleProWorkspace.tsx" + ); + + expect(workspace).toContain("refreshInterval: 30_000"); + expect(workspace).toContain("revalidateOnFocus: true"); + expect(workspace).toContain("revalidateOnReconnect: true"); }); test("opens the canvas conversation panel by default", () => {