From 4d934f50a9cc1dce901ef4a08c9be2b6229b1789 Mon Sep 17 00:00:00 2001 From: syn Date: Tue, 16 Jun 2026 16:55:52 -0500 Subject: [PATCH 1/5] feat(ai-gateway): add organization auto model routing --- .../api/openrouter/[...path]/route.test.ts | 49 ++ .../src/app/api/openrouter/[...path]/route.ts | 32 + .../organizations/[id]/defaults/route.test.ts | 32 + .../api/organizations/[id]/defaults/route.ts | 15 +- .../organizations/[id]/modes/route.test.ts | 56 ++ .../app/api/organizations/[id]/modes/route.ts | 6 +- apps/web/src/app/api/organizations/hooks.ts | 36 + ...ionProvidersAndModelsConfigurationCard.tsx | 1 + .../custom-modes/CustomModesLayout.tsx | 16 +- .../organizations/custom-modes/ModeForm.tsx | 32 +- .../DefaultModelDialog.tsx | 184 ++++- .../src/lib/ai-gateway/auto-model/index.ts | 23 +- .../ai-gateway/auto-model/resolution.test.ts | 146 +++- .../lib/ai-gateway/auto-model/resolution.ts | 145 +++- .../src/lib/ai-gateway/llm-proxy-helpers.ts | 11 + .../ai-gateway/providers/openrouter/index.ts | 86 +-- .../organizations/organization-auto-model.ts | 184 +++++ .../organizations/organization-base-types.ts | 2 + .../lib/organizations/organization-models.ts | 11 +- .../lib/organizations/organization-modes.ts | 33 +- .../lib/organizations/organization-types.ts | 2 + .../src/lib/organizations/organizations.ts | 22 + apps/web/src/lib/proxy-error-types.ts | 1 + .../organization-modes-router.test.ts | 245 +++++- .../organization-modes-router.ts | 655 ++++++++++------ .../organization-settings-router.test.ts | 190 +++++ .../organization-settings-router.ts | 722 ++++++++++++++---- packages/db/src/schema-types.test.ts | 43 ++ packages/db/src/schema-types.ts | 32 + 29 files changed, 2586 insertions(+), 426 deletions(-) create mode 100644 apps/web/src/lib/organizations/organization-auto-model.ts create mode 100644 packages/db/src/schema-types.test.ts diff --git a/apps/web/src/app/api/openrouter/[...path]/route.test.ts b/apps/web/src/app/api/openrouter/[...path]/route.test.ts index bb7b22ded3..c46f75a7bf 100644 --- a/apps/web/src/app/api/openrouter/[...path]/route.test.ts +++ b/apps/web/src/app/api/openrouter/[...path]/route.test.ts @@ -13,6 +13,7 @@ import type { Provider } from '@/lib/ai-gateway/providers/types'; import { fetchEfficientAutoDecision } from '@/lib/ai-gateway/auto-routing-decision'; import { logMicrodollarUsage } from '@/lib/ai-gateway/processUsage'; import { applyResolvedAutoModel } from '@/lib/ai-gateway/auto-model/resolution'; +import { getDirectByokModel } from '@/lib/ai-gateway/providers/direct-byok'; jest.mock('next/server', () => { return { @@ -40,6 +41,9 @@ jest.mock('@/lib/ai-gateway/abuse-service', () => { }; }); jest.mock('@/lib/ai-gateway/providers/get-provider'); +jest.mock('@/lib/ai-gateway/providers/direct-byok', () => ({ + getDirectByokModel: jest.fn(async () => ({ provider: null, model: null })), +})); jest.mock('@/lib/ai-gateway/providers/upstream-request'); jest.mock('@/lib/ai-gateway/providers/gateway-models-cache'); jest.mock('@/lib/redis', () => ({ @@ -90,6 +94,7 @@ const mockedRedisSet = jest.mocked(redisClient.set); const mockedFetchEfficientAutoDecision = jest.mocked(fetchEfficientAutoDecision); const mockedLogMicrodollarUsage = jest.mocked(logMicrodollarUsage); const mockedApplyResolvedAutoModel = jest.mocked(applyResolvedAutoModel); +const mockedGetDirectByokModel = jest.mocked(getDirectByokModel); const provider = { id: 'openrouter', @@ -413,7 +418,9 @@ describe('POST /api/openrouter/v1/chat/completions rules-engine actions', () => describe('kilo-auto/efficient classifier billing', () => { beforeEach(() => { jest.clearAllMocks(); + mockedGetDirectByokModel.mockResolvedValue({ provider: null, model: null }); setUserAuth(); + mockedGetProvider.mockResolvedValue({ kind: 'provider', provider, @@ -443,6 +450,48 @@ describe('kilo-auto/efficient classifier billing', () => { }); }); + it('rejects Organization Auto direct-BYOK routes when provider selection falls through', async () => { + mockedGetUserFromAuth.mockResolvedValue({ + user: { + id: 'user-123', + google_user_email: 'test@example.com', + microdollars_used: 0, + } as User, + authFailedResponse: null, + organizationId: 'org-1', + }); + mockedGetBalanceAndOrgSettings.mockResolvedValue({ + balance: 1000, + settings: { + default_model: 'kilo-auto/org', + org_auto_model: { routes: {}, fallback_model: 'kilo-auto/balanced' }, + }, + plan: 'enterprise', + }); + mockedApplyResolvedAutoModel.mockImplementation(async (_params, request) => { + request.body.model = 'martian/moonshotai/kimi-k2.6'; + return { + kind: 'ok', + resolved: { model: 'martian/moonshotai/kimi-k2.6' }, + routingTarget: 'martian/moonshotai/kimi-k2.6', + }; + }); + mockedGetDirectByokModel.mockResolvedValue({ + provider: { id: 'martian' } as never, + model: {} as never, + }); + + const { POST } = await import('./route'); + const response = await POST(makeRequest(makeBody('kilo-auto/org')) as never); + + expect(response.status).toBe(400); + expect(await response.json()).toMatchObject({ + error_type: 'organization_auto_configuration', + message: expect.stringContaining('does not have an enabled BYOK credential for martian'), + }); + expect(mockedUpstreamRequest).not.toHaveBeenCalled(); + }); + it('bills classifier cost when cost > 0 and user is non-BYOK', async () => { mockedFetchEfficientAutoDecision.mockResolvedValue({ decision: { diff --git a/apps/web/src/app/api/openrouter/[...path]/route.ts b/apps/web/src/app/api/openrouter/[...path]/route.ts index d2b1bf7b09..a52067662b 100644 --- a/apps/web/src/app/api/openrouter/[...path]/route.ts +++ b/apps/web/src/app/api/openrouter/[...path]/route.ts @@ -17,6 +17,7 @@ import type { } from '@/lib/ai-gateway/providers/openrouter/types'; import { applyProviderSpecificLogic } from '@/lib/ai-gateway/providers/apply-provider-specific-logic'; import { getProvider } from '@/lib/ai-gateway/providers/get-provider'; +import { getDirectByokModel } from '@/lib/ai-gateway/providers/direct-byok'; import { buildExperimentPromptCapture } from '@/lib/ai-gateway/experiments/persist'; import { isPublicIdExperimented } from '@/lib/ai-gateway/experiments/membership'; import { upstreamRequest } from '@/lib/ai-gateway/providers/upstream-request'; @@ -46,6 +47,7 @@ import { modelNotAllowedResponse, extractHeaderAndLimitLength, noFreeModelsAvailableResponse, + organizationAutoConfigurationResponse, temporarilyUnavailableResponse, usageLimitExceededResponse, wrapInSafeNextResponse, @@ -93,6 +95,7 @@ import { isKiloAutoModel, KILO_AUTO_FREE_MODEL, KILO_AUTO_EFFICIENT_MODEL, + ORG_AUTO_MODEL, } from '@/lib/ai-gateway/auto-model'; import { applyResolvedAutoModel } from '@/lib/ai-gateway/auto-model/resolution'; import { fetchEfficientAutoDecision } from '@/lib/ai-gateway/auto-routing-decision'; @@ -241,6 +244,13 @@ export async function POST(request: NextRequest): Promise ({ + organizationId: auth.organizationId, + settings: balanceAndSettings.settings, + plan: balanceAndSettings.plan, + }) + ); // Extract IP early (needed for free model routing fallback and rate limiting) const ipAddress = request.headers.get('x-forwarded-for')?.split(',')[0]?.trim(); @@ -270,6 +280,7 @@ export async function POST(request: NextRequest): Promise res.user), @@ -318,6 +330,10 @@ export async function POST(request: NextRequest): Promise { expect(mockedGetEnhancedOpenRouterModels).not.toHaveBeenCalled(); }); + test('returns Organization Auto when it is configured as the organization default', async () => { + const user = await insertTestUser(); + const organization = await createOrganization('Test Org', user.id); + + mockedGetEnhancedOpenRouterModels.mockRejectedValue(new Error('should not be called')); + mockedGetAuthorizedOrgContext.mockResolvedValue({ + success: true, + data: { + user: { ...user, role: 'owner' }, + organization: { + ...organization, + plan: 'enterprise' as const, + settings: { + default_model: 'kilo-auto/org', + org_auto_model: { + routes: {}, + fallback_model: 'kilo-auto/balanced', + }, + }, + }, + }, + }); + + const response = await GET(new NextRequest('http://localhost:3000'), { + params: Promise.resolve({ id: organization.id }), + }); + + expect(response.status).toBe(200); + const body = await response.json(); + expect(body.defaultModel).toBe('kilo-auto/org'); + }); + test('returns 409 when all available models are blocked by policy', async () => { const user = await insertTestUser(); const organization = await createOrganization('Test Org', user.id); diff --git a/apps/web/src/app/api/organizations/[id]/defaults/route.ts b/apps/web/src/app/api/organizations/[id]/defaults/route.ts index 90f6744af0..2e92f75d96 100644 --- a/apps/web/src/app/api/organizations/[id]/defaults/route.ts +++ b/apps/web/src/app/api/organizations/[id]/defaults/route.ts @@ -8,8 +8,9 @@ import { hasActiveModelRestrictions, } from '@/lib/model-allow.server'; import { getModelIdToProviderSlugsIndex } from '@/lib/ai-gateway/providers/openrouter/models-by-provider-index.server'; -import { KILO_AUTO_FREE_MODEL } from '@/lib/ai-gateway/auto-model'; +import { KILO_AUTO_FREE_MODEL, ORG_AUTO_MODEL } from '@/lib/ai-gateway/auto-model'; import { getEffectiveModelRestrictions } from '@/lib/organizations/model-restrictions'; +import { isOrganizationAutoConfigured } from '@/lib/organizations/organization-auto-model'; type DefaultsResponse = { defaultModel: string; @@ -67,8 +68,16 @@ export async function GET( return undefined; }; - // If organization has a default model set, validate it against allowed models - if (defaultModel && (defaultModel.endsWith('/*') || !(await isAllowed(defaultModel)))) { + // If organization has a default model set, validate it against allowed models. + // Organization Auto is a virtual organization-only default, so its eligibility + // is validated from persisted organization settings rather than provider policy. + if (defaultModel === ORG_AUTO_MODEL.id && !isOrganizationAutoConfigured(organization)) { + defaultModel = undefined; + } else if ( + defaultModel && + defaultModel !== ORG_AUTO_MODEL.id && + (defaultModel.endsWith('/*') || !(await isAllowed(defaultModel))) + ) { // Organization's configured default model is not permitted; fall back to a safe default. defaultModel = undefined; } diff --git a/apps/web/src/app/api/organizations/[id]/modes/route.test.ts b/apps/web/src/app/api/organizations/[id]/modes/route.test.ts index 1d0a38868d..8d777526ea 100644 --- a/apps/web/src/app/api/organizations/[id]/modes/route.test.ts +++ b/apps/web/src/app/api/organizations/[id]/modes/route.test.ts @@ -61,6 +61,62 @@ describe('GET /api/organizations/[id]/modes', () => { expect(mockedGetAllOrganizationModes).toHaveBeenCalledWith('org-1'); }); + test('projects canonical Organization Auto routes over legacy mode defaults', async () => { + mockedGetAuthorizedOrgContext.mockResolvedValue({ + success: true, + data: { + organization: { + id: 'org-1', + settings: { + org_auto_model: { + routes: { code: 'kilo-auto/frontier' }, + fallback_model: 'kilo-auto/balanced', + }, + }, + }, + }, + } as never); + mockedGetAllOrganizationModes.mockResolvedValue([ + { + id: 'mode-1', + organization_id: 'org-1', + name: 'Code', + slug: 'code', + created_by: 'user-1', + created_at: '2026-01-01T00:00:00.000Z', + updated_at: '2026-01-01T00:00:00.000Z', + config: { + roleDefinition: 'You are a coding assistant', + groups: ['read'], + defaultModel: 'openai/gpt-4o', + }, + }, + ]); + + const response = await GET(new NextRequest('http://localhost:3000'), { + params: Promise.resolve({ id: 'org-1' }), + }); + + await expect(response.json()).resolves.toEqual({ + modes: [ + { + id: 'mode-1', + organization_id: 'org-1', + name: 'Code', + slug: 'code', + created_by: 'user-1', + created_at: '2026-01-01T00:00:00.000Z', + updated_at: '2026-01-01T00:00:00.000Z', + config: { + roleDefinition: 'You are a coding assistant', + groups: ['read'], + defaultModel: 'kilo-auto/frontier', + }, + }, + ], + }); + }); + test('returns a legacy mode row without defaultModel unchanged', async () => { mockedGetAuthorizedOrgContext.mockResolvedValue({ success: true, diff --git a/apps/web/src/app/api/organizations/[id]/modes/route.ts b/apps/web/src/app/api/organizations/[id]/modes/route.ts index 305d406e3b..7f5b97b5df 100644 --- a/apps/web/src/app/api/organizations/[id]/modes/route.ts +++ b/apps/web/src/app/api/organizations/[id]/modes/route.ts @@ -3,6 +3,7 @@ import { getAuthorizedOrgContext } from '@/lib/organizations/organization-auth'; import type { NextRequest } from 'next/server'; import type { OrganizationMode } from '@/lib/organizations/organization-modes'; import { getAllOrganizationModes } from '@/lib/organizations/organization-modes'; +import { projectOrganizationAutoRoutesIntoModes } from '@/lib/organizations/organization-auto-model'; export async function GET( _request: NextRequest, @@ -15,7 +16,10 @@ export async function GET( } const { organization } = data; - const modes = await getAllOrganizationModes(organization.id); + const modes = projectOrganizationAutoRoutesIntoModes( + await getAllOrganizationModes(organization.id), + organization.settings + ); return NextResponse.json({ modes, diff --git a/apps/web/src/app/api/organizations/hooks.ts b/apps/web/src/app/api/organizations/hooks.ts index 2093d3affd..e938c997c4 100644 --- a/apps/web/src/app/api/organizations/hooks.ts +++ b/apps/web/src/app/api/organizations/hooks.ts @@ -201,6 +201,42 @@ export function useUpdateDefaultModel() { ); } +export function useEnableOrganizationAuto() { + const trpc = useTRPC(); + const queryClient = useQueryClient(); + return useMutation( + trpc.organizations.settings.enableOrganizationAuto.mutationOptions({ + onSuccess: () => { + void queryClient.invalidateQueries({ queryKey: trpc.organizations.pathKey() }); + }, + }) + ); +} + +export function useDisableOrganizationAuto() { + const trpc = useTRPC(); + const queryClient = useQueryClient(); + return useMutation( + trpc.organizations.settings.disableOrganizationAuto.mutationOptions({ + onSuccess: () => { + void queryClient.invalidateQueries({ queryKey: trpc.organizations.pathKey() }); + }, + }) + ); +} + +export function useSetOrganizationAutoFallback() { + const trpc = useTRPC(); + const queryClient = useQueryClient(); + return useMutation( + trpc.organizations.settings.setOrganizationAutoFallback.mutationOptions({ + onSuccess: () => { + void queryClient.invalidateQueries({ queryKey: trpc.organizations.pathKey() }); + }, + }) + ); +} + export function useUpdateOrganizationSeatsRequired() { const trpc = useTRPC(); const invalidate = useInvalidateAllOrganizationData(); diff --git a/apps/web/src/components/organizations/OrganizationProvidersAndModelsConfigurationCard.tsx b/apps/web/src/components/organizations/OrganizationProvidersAndModelsConfigurationCard.tsx index 4c9863ca4b..e643973348 100644 --- a/apps/web/src/components/organizations/OrganizationProvidersAndModelsConfigurationCard.tsx +++ b/apps/web/src/components/organizations/OrganizationProvidersAndModelsConfigurationCard.tsx @@ -175,6 +175,7 @@ export function OrganizationProvidersAndModelsConfigurationCard({ organizationId={organizationId} organizationSettings={organizationData?.settings} currentDefaultModel={organizationData?.settings?.default_model} + organizationPlan={organizationData?.plan} /> {isDefaultModelConfigEnabled && mode.config.defaultModel && (
- Default model + + Organization Auto route + {mode.config.defaultModel} @@ -174,7 +177,7 @@ export function CustomModesLayout({ organizationId }: CustomModesLayoutProps) { const [drawerMode, setDrawerMode] = useState<'create' | 'edit'>('create'); const [editingMode, setEditingMode] = useState(null); const isReadOnly = useOrganizationReadOnly(organizationId); - const isDefaultModelFeatureEnabled = useFeatureFlagEnabled('org-default-model-config'); + const isDefaultModelFeatureEnabled = useFeatureFlagEnabled('organization-auto-model-routing'); const isDevelopment = process.env.NODE_ENV === 'development'; const isDefaultModelConfigEnabled = isDevelopment || isDefaultModelFeatureEnabled === true; const canSetDefaultModel = organizationData?.plan === 'enterprise'; @@ -206,7 +209,12 @@ export function CustomModesLayout({ organizationId }: CustomModesLayoutProps) { organization_id: organizationId, slug: defaultMode.slug, name: defaultMode.name, - config: defaultMode.config, + config: { + ...defaultMode.config, + defaultModel: organizationData + ? getOrganizationAutoRoute(organizationData.settings, defaultMode.slug) + : defaultMode.config.defaultModel, + }, created_by: '', created_at: new Date().toISOString(), updated_at: new Date().toISOString(), @@ -231,7 +239,7 @@ export function CustomModesLayout({ organizationId }: CustomModesLayoutProps) { builtInModes: builtInDisplayModes, customModes: customDisplayModes, }; - }, [data?.modes, organizationId]); + }, [data?.modes, organizationData?.settings, organizationId]); const handleDelete = async () => { if (!modeToDelete) return; diff --git a/apps/web/src/components/organizations/custom-modes/ModeForm.tsx b/apps/web/src/components/organizations/custom-modes/ModeForm.tsx index 58b9e1e2d4..b9c836910b 100644 --- a/apps/web/src/components/organizations/custom-modes/ModeForm.tsx +++ b/apps/web/src/components/organizations/custom-modes/ModeForm.tsx @@ -20,6 +20,8 @@ import type { EditGroupConfig } from '@/lib/organizations/organization-types'; import { Save, FileText } from 'lucide-react'; import { useModeTemplates } from './useModeTemplates'; import { useModelSelectorList } from '@/app/api/openrouter/hooks'; +import { isOrganizationAutoTargetModel } from '@/lib/organizations/organization-auto-model'; +import { CUSTOM_LLM_PREFIX } from '@/lib/ai-gateway/model-utils'; const availableGroups = [ { value: 'read', label: 'Read Files' }, @@ -153,7 +155,19 @@ export function ModeForm({ isLoading: modelsLoading, error: modelsError, } = useModelSelectorList(organizationId, isDefaultModelConfigEnabled && canSetDefaultModel); - const modelOptions = useMemo(() => modelsData?.data || [], [modelsData?.data]); + const modelOptions = useMemo( + () => + (modelsData?.data || []).filter(model => { + if (model.id.startsWith(CUSTOM_LLM_PREFIX)) { + return false; + } + if (model.id.startsWith('kilo-auto/')) { + return isOrganizationAutoTargetModel(model.id); + } + return true; + }), + [modelsData?.data] + ); const hasCurrentDefaultModelOption = canSetDefaultModel && !!formData.defaultModel && @@ -456,7 +470,7 @@ export function ModeForm({ {shouldShowDefaultModelControl && (
- +

{!canSetDefaultModel - ? 'This organization must be on Enterprise to set mode defaults. Existing defaults can still be cleared.' + ? 'This organization must be on Enterprise to configure Organization Auto routes. Existing routes can still be cleared.' : modelsLoading ? 'Loading organization-allowed models...' : modelsError ? 'Unable to load organization models.' : modelOptions.length === 0 ? 'No organization-allowed models are available.' - : 'Members can still override this locally in Kilo Code.'} + : 'Members can still override Organization Auto locally in Kilo Code.'}

{hasUnavailableDefaultModel && (

diff --git a/apps/web/src/components/organizations/providers-and-models/DefaultModelDialog.tsx b/apps/web/src/components/organizations/providers-and-models/DefaultModelDialog.tsx index 5434c68c21..78ad2ec28c 100644 --- a/apps/web/src/components/organizations/providers-and-models/DefaultModelDialog.tsx +++ b/apps/web/src/components/organizations/providers-and-models/DefaultModelDialog.tsx @@ -3,7 +3,12 @@ import { useState } from 'react'; import { useQueryClient } from '@tanstack/react-query'; import { LockableContainer } from '../LockableContainer'; -import { useUpdateDefaultModel } from '@/app/api/organizations/hooks'; +import { + useDisableOrganizationAuto, + useEnableOrganizationAuto, + useSetOrganizationAutoFallback, + useUpdateDefaultModel, +} from '@/app/api/organizations/hooks'; import { useModelSelectorList } from '@/app/api/openrouter/hooks'; import { Button } from '@/components/ui/button'; import { @@ -24,6 +29,10 @@ import { import type { OrganizationSettings } from '@/lib/organizations/organization-types'; import { toast } from 'sonner'; import { Settings2 } from 'lucide-react'; +import { ORG_AUTO_MODEL } from '@/lib/ai-gateway/auto-model'; +import { isOrganizationAutoTargetModel } from '@/lib/organizations/organization-auto-model'; +import { CUSTOM_LLM_PREFIX } from '@/lib/ai-gateway/model-utils'; +import { useFeatureFlagEnabled } from 'posthog-js/react'; type DefaultModelDialogProps = { open: boolean; @@ -31,6 +40,7 @@ type DefaultModelDialogProps = { organizationId: string; organizationSettings?: OrganizationSettings; currentDefaultModel?: string; + organizationPlan?: 'teams' | 'enterprise'; }; export function DefaultModelDialog({ @@ -39,15 +49,38 @@ export function DefaultModelDialog({ organizationId, organizationSettings, currentDefaultModel, + organizationPlan, }: DefaultModelDialogProps) { const queryClient = useQueryClient(); const [selectedModel, setSelectedModel] = useState(''); + const [selectedFallbackModel, setSelectedFallbackModel] = useState(''); const { data: openRouterModels, isLoading: modelsLoading } = useModelSelectorList(organizationId); const updateDefaultModelMutation = useUpdateDefaultModel(); + const enableOrganizationAutoMutation = useEnableOrganizationAuto(); + const disableOrganizationAutoMutation = useDisableOrganizationAuto(); + const setOrganizationAutoFallbackMutation = useSetOrganizationAutoFallback(); const organizationDefaultModel = organizationSettings?.default_model; - const availableModels = openRouterModels?.data ?? []; + const organizationAutoFallbackModel = organizationSettings?.org_auto_model?.fallback_model; + const organizationAutoEnabled = organizationDefaultModel === ORG_AUTO_MODEL.id; + const organizationAutoFeatureEnabled = useFeatureFlagEnabled('organization-auto-model-routing'); + const isDevelopment = process.env.NODE_ENV === 'development'; + const canConfigureOrganizationAuto = + organizationPlan === 'enterprise' && (isDevelopment || organizationAutoFeatureEnabled === true); + const showOrganizationAutoSection = organizationAutoEnabled || canConfigureOrganizationAuto; + const availableModels = (openRouterModels?.data ?? []).filter( + model => model.id !== ORG_AUTO_MODEL.id + ); + const organizationAutoTargetModels = availableModels.filter(model => { + if (model.id.startsWith(CUSTOM_LLM_PREFIX)) { + return false; + } + if (model.id.startsWith('kilo-auto/')) { + return isOrganizationAutoTargetModel(model.id); + } + return true; + }); const handleUpdateDefaultModel = async () => { if (!selectedModel) return; @@ -93,6 +126,56 @@ export function DefaultModelDialog({ } }; + const handleEnableOrganizationAuto = async () => { + try { + await enableOrganizationAutoMutation.mutateAsync({ organizationId }); + await queryClient.invalidateQueries({ queryKey: ['organization-defaults', organizationId] }); + setSelectedModel(''); + onOpenChange(false); + toast.success('Organization Auto enabled'); + } catch (error) { + console.error('Failed to enable Organization Auto:', error); + toast.error(error instanceof Error ? error.message : 'Failed to enable Organization Auto'); + } + }; + + const handleDisableOrganizationAuto = async () => { + if (!selectedModel) return; + + try { + await disableOrganizationAutoMutation.mutateAsync({ + organizationId, + replacement_model: selectedModel, + }); + await queryClient.invalidateQueries({ queryKey: ['organization-defaults', organizationId] }); + setSelectedModel(''); + onOpenChange(false); + toast.success('Organization Auto disabled'); + } catch (error) { + console.error('Failed to disable Organization Auto:', error); + toast.error(error instanceof Error ? error.message : 'Failed to disable Organization Auto'); + } + }; + + const handleSetOrganizationAutoFallback = async () => { + const fallbackModel = selectedFallbackModel || organizationAutoFallbackModel; + if (!fallbackModel) return; + + try { + await setOrganizationAutoFallbackMutation.mutateAsync({ + organizationId, + model_id: fallbackModel, + }); + setSelectedFallbackModel(''); + toast.success('Organization Auto fallback updated'); + } catch (error) { + console.error('Failed to update Organization Auto fallback:', error); + toast.error( + error instanceof Error ? error.message : 'Failed to update Organization Auto fallback' + ); + } + }; + return (

@@ -129,6 +212,86 @@ export function DefaultModelDialog({
+ {showOrganizationAutoSection && ( +
+
+ +

+ Route the organization default by mode. Local model selections still override + it. +

+ {canConfigureOrganizationAuto && ( + + Configure mode routes + + )} +
+ {organizationAutoEnabled ? ( +

+ Enabled. Choose a replacement model below to disable it. +

+ ) : ( + + )} + {canConfigureOrganizationAuto && organizationSettings?.org_auto_model && ( +
+ +
+ + +
+
+ )} +
+ )} +

{!canSetDefaultModel - ? 'This organization must be on Enterprise to configure Organization Auto routes. Existing routes can still be cleared.' + ? 'Organization Auto routes are read-only for your role or plan.' : modelsLoading ? 'Loading organization-allowed models...' : modelsError diff --git a/apps/web/src/components/organizations/custom-modes/NewModeForm.tsx b/apps/web/src/components/organizations/custom-modes/NewModeForm.tsx index 8d08034ee1..673eb2aea4 100644 --- a/apps/web/src/components/organizations/custom-modes/NewModeForm.tsx +++ b/apps/web/src/components/organizations/custom-modes/NewModeForm.tsx @@ -5,6 +5,7 @@ import { useClearOrganizationAutoRoute, useCreateOrganizationMode, useOrganizationModes, + useOrganizationWithMembers, useSetOrganizationAutoRoute, } from '@/app/api/organizations/hooks'; import { ModeForm, type ModeFormData } from './ModeForm'; @@ -12,13 +13,11 @@ import { matchesBuiltInModeState } from './EditModeForm'; import { toast } from 'sonner'; import { DEFAULT_MODES } from './default-modes'; import { useMemo } from 'react'; -import { useOrganizationWithMembers } from '@/app/api/organizations/hooks'; -import { getOrganizationAutoRoute } from '@/lib/organizations/organization-auto-model'; +import { getOrganizationAutoRoute } from '@/lib/organizations/organization-auto-model-shared'; type NewModeFormProps = { organizationId: string; defaultModeSlug?: string; - routeModel?: string; isDefaultModelConfigEnabled?: boolean; canSetDefaultModel?: boolean; onSuccess?: () => void; @@ -28,7 +27,6 @@ type NewModeFormProps = { export function NewModeForm({ organizationId, defaultModeSlug: propDefaultModeSlug, - routeModel: propRouteModel, isDefaultModelConfigEnabled = false, canSetDefaultModel = true, onSuccess, @@ -47,11 +45,9 @@ export function NewModeForm({ if (!defaultModeSlug) return undefined; return DEFAULT_MODES.find(m => m.slug === defaultModeSlug); }, [defaultModeSlug]); - const routeModel = - propRouteModel ?? - (defaultModeSlug - ? getOrganizationAutoRoute(organizationData?.settings, defaultModeSlug) - : undefined); + const routeModel = defaultModeSlug + ? getOrganizationAutoRoute(organizationData?.settings, defaultModeSlug) + : undefined; // Convert default mode to the format expected by ModeForm const initialMode = useMemo(() => { @@ -89,7 +85,8 @@ export function NewModeForm({ return; } - const created = await createMutation.mutateAsync({ + const nextRouteModel = data.defaultModel || undefined; + await createMutation.mutateAsync({ organizationId, name: data.name, slug: data.slug, @@ -100,8 +97,8 @@ export function NewModeForm({ groups: data.groups as ('read' | 'edit' | 'browser' | 'command' | 'mcp')[], customInstructions: data.customInstructions, }, + ...(nextRouteModel === routeModel ? {} : { route_model: nextRouteModel ?? null }), }); - await persistRoute(created.mode.slug, data.defaultModel); toast.success(`Mode "${data.name}" created successfully`); onSuccess?.(); } catch (error) { @@ -125,7 +122,6 @@ export function NewModeForm({ canSetDefaultModel={canSetDefaultModel} existingModes={modesData?.modes || []} onCancel={onCancel} - renderButtons={() => null} /> ); } diff --git a/apps/web/src/components/organizations/providers-and-models/DefaultModelDialog.tsx b/apps/web/src/components/organizations/providers-and-models/DefaultModelDialog.tsx index 78ad2ec28c..3e1813d5f6 100644 --- a/apps/web/src/components/organizations/providers-and-models/DefaultModelDialog.tsx +++ b/apps/web/src/components/organizations/providers-and-models/DefaultModelDialog.tsx @@ -1,7 +1,6 @@ 'use client'; -import { useState } from 'react'; -import { useQueryClient } from '@tanstack/react-query'; +import { useEffect, useState } from 'react'; import { LockableContainer } from '../LockableContainer'; import { useDisableOrganizationAuto, @@ -30,9 +29,13 @@ import type { OrganizationSettings } from '@/lib/organizations/organization-type import { toast } from 'sonner'; import { Settings2 } from 'lucide-react'; import { ORG_AUTO_MODEL } from '@/lib/ai-gateway/auto-model'; -import { isOrganizationAutoTargetModel } from '@/lib/organizations/organization-auto-model'; +import { + isOrganizationAutoTargetModel, + ORGANIZATION_AUTO_MODEL_FLAG, +} from '@/lib/organizations/organization-auto-model-shared'; import { CUSTOM_LLM_PREFIX } from '@/lib/ai-gateway/model-utils'; import { useFeatureFlagEnabled } from 'posthog-js/react'; +import Link from 'next/link'; type DefaultModelDialogProps = { open: boolean; @@ -51,7 +54,6 @@ export function DefaultModelDialog({ currentDefaultModel, organizationPlan, }: DefaultModelDialogProps) { - const queryClient = useQueryClient(); const [selectedModel, setSelectedModel] = useState(''); const [selectedFallbackModel, setSelectedFallbackModel] = useState(''); @@ -64,7 +66,7 @@ export function DefaultModelDialog({ const organizationDefaultModel = organizationSettings?.default_model; const organizationAutoFallbackModel = organizationSettings?.org_auto_model?.fallback_model; const organizationAutoEnabled = organizationDefaultModel === ORG_AUTO_MODEL.id; - const organizationAutoFeatureEnabled = useFeatureFlagEnabled('organization-auto-model-routing'); + const organizationAutoFeatureEnabled = useFeatureFlagEnabled(ORGANIZATION_AUTO_MODEL_FLAG); const isDevelopment = process.env.NODE_ENV === 'development'; const canConfigureOrganizationAuto = organizationPlan === 'enterprise' && (isDevelopment || organizationAutoFeatureEnabled === true); @@ -72,6 +74,13 @@ export function DefaultModelDialog({ const availableModels = (openRouterModels?.data ?? []).filter( model => model.id !== ORG_AUTO_MODEL.id ); + useEffect(() => { + if (!open) { + setSelectedModel(''); + setSelectedFallbackModel(''); + } + }, [open]); + const organizationAutoTargetModels = availableModels.filter(model => { if (model.id.startsWith(CUSTOM_LLM_PREFIX)) { return false; @@ -81,6 +90,9 @@ export function DefaultModelDialog({ } return true; }); + const organizationAutoFallbackUnavailable = + !!organizationAutoFallbackModel && + !organizationAutoTargetModels.some(model => model.id === organizationAutoFallbackModel); const handleUpdateDefaultModel = async () => { if (!selectedModel) return; @@ -91,11 +103,6 @@ export function DefaultModelDialog({ default_model: selectedModel, }); - // Invalidate the defaults query to refresh the display - await queryClient.invalidateQueries({ - queryKey: ['organization-defaults', organizationId], - }); - setSelectedModel(''); onOpenChange(false); toast.success('Default model updated successfully'); @@ -113,10 +120,6 @@ export function DefaultModelDialog({ default_model: null, }); - await queryClient.invalidateQueries({ - queryKey: ['organization-defaults', organizationId], - }); - setSelectedModel(''); onOpenChange(false); toast.success('Default model cleared - will use global default'); @@ -129,7 +132,6 @@ export function DefaultModelDialog({ const handleEnableOrganizationAuto = async () => { try { await enableOrganizationAutoMutation.mutateAsync({ organizationId }); - await queryClient.invalidateQueries({ queryKey: ['organization-defaults', organizationId] }); setSelectedModel(''); onOpenChange(false); toast.success('Organization Auto enabled'); @@ -147,7 +149,6 @@ export function DefaultModelDialog({ organizationId, replacement_model: selectedModel, }); - await queryClient.invalidateQueries({ queryKey: ['organization-defaults', organizationId] }); setSelectedModel(''); onOpenChange(false); toast.success('Organization Auto disabled'); @@ -221,12 +222,13 @@ export function DefaultModelDialog({ it.

{canConfigureOrganizationAuto && ( - onOpenChange(false)} > Configure mode routes - + )}
{organizationAutoEnabled ? ( @@ -260,6 +262,18 @@ export function DefaultModelDialog({ + {organizationAutoFallbackUnavailable && organizationAutoFallbackModel && ( + +
+ + {organizationAutoFallbackModel} + + + Unavailable current fallback + +
+ + )} {organizationAutoTargetModels.map(model => (
@@ -287,13 +301,21 @@ export function DefaultModelDialog({ {setOrganizationAutoFallbackMutation.isPending ? 'Saving...' : 'Save'}
+ {organizationAutoFallbackUnavailable && ( +

+ This fallback is no longer available. Modes without explicit routes will + fail until you replace it. +

+ )} )} )}
- +