From 4d934f50a9cc1dce901ef4a08c9be2b6229b1789 Mon Sep 17 00:00:00 2001 From: syn Date: Tue, 16 Jun 2026 16:55:52 -0500 Subject: [PATCH 1/8] 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. +

+ )} )} )}
- + @@ -511,15 +522,11 @@ export function ModeForm({ aria-invalid={Boolean(errors.defaultModel)} > - - Use Organization Auto fallback - + {routeFallbackLabel} {shouldRenderCurrentDefaultModel && (
@@ -552,13 +559,17 @@ export function ModeForm({

{!canSetDefaultModel ? 'Organization Auto routes are read-only for your role or plan.' - : modelsLoading - ? 'Loading organization-allowed models...' - : modelsError - ? 'Unable to load organization models.' - : modelOptions.length === 0 - ? 'No organization-allowed models are available.' - : 'Members can still override Organization Auto locally in Kilo Code.'} + : !isOrganizationAutoDefaultActive + ? hasStoredInactiveRoute + ? 'Organization Auto is off. This saved route will apply if you enable it.' + : 'Organization Auto is off. Select a route now to use it when enabled.' + : modelsLoading + ? 'Loading organization-allowed models...' + : modelsError + ? 'Unable to load organization models.' + : modelOptions.length === 0 + ? 'No organization-allowed models are available.' + : 'Members can still override Organization Auto locally in Kilo Code.'}

{hasUnavailableDefaultModel && (

diff --git a/apps/web/src/components/organizations/custom-modes/NewModeForm.tsx b/apps/web/src/components/organizations/custom-modes/NewModeForm.tsx index 673eb2aea4..c7ff426daa 100644 --- a/apps/web/src/components/organizations/custom-modes/NewModeForm.tsx +++ b/apps/web/src/components/organizations/custom-modes/NewModeForm.tsx @@ -14,6 +14,7 @@ import { toast } from 'sonner'; import { DEFAULT_MODES } from './default-modes'; import { useMemo } from 'react'; import { getOrganizationAutoRoute } from '@/lib/organizations/organization-auto-model-shared'; +import { ORG_AUTO_MODEL } from '@/lib/ai-gateway/auto-model'; type NewModeFormProps = { organizationId: string; @@ -48,6 +49,8 @@ export function NewModeForm({ const routeModel = defaultModeSlug ? getOrganizationAutoRoute(organizationData?.settings, defaultModeSlug) : undefined; + const isOrganizationAutoDefaultActive = + organizationData?.settings.default_model === ORG_AUTO_MODEL.id; // Convert default mode to the format expected by ModeForm const initialMode = useMemo(() => { @@ -119,6 +122,7 @@ export function NewModeForm({ } isEditingBuiltIn={!!defaultMode} isDefaultModelConfigEnabled={isDefaultModelConfigEnabled} + isOrganizationAutoDefaultActive={isOrganizationAutoDefaultActive} canSetDefaultModel={canSetDefaultModel} existingModes={modesData?.modes || []} onCancel={onCancel} 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 3e1813d5f6..369d84d93a 100644 --- a/apps/web/src/components/organizations/providers-and-models/DefaultModelDialog.tsx +++ b/apps/web/src/components/organizations/providers-and-models/DefaultModelDialog.tsx @@ -1,22 +1,15 @@ 'use client'; -import { useEffect, useState } from 'react'; -import { LockableContainer } from '../LockableContainer'; -import { - useDisableOrganizationAuto, - useEnableOrganizationAuto, - useSetOrganizationAutoFallback, - useUpdateDefaultModel, -} from '@/app/api/organizations/hooks'; +import { useEffect, useMemo, useState } from 'react'; +import Link from 'next/link'; +import { useFeatureFlagEnabled } from 'posthog-js/react'; +import { Settings2 } from 'lucide-react'; +import { toast } from 'sonner'; +import { useConfigureOrganizationDefaultBehavior } from '@/app/api/organizations/hooks'; import { useModelSelectorList } from '@/app/api/openrouter/hooks'; +import { LockableContainer } from '../LockableContainer'; +import { Badge } from '@/components/ui/badge'; import { Button } from '@/components/ui/button'; -import { - Select, - SelectContent, - SelectItem, - SelectTrigger, - SelectValue, -} from '@/components/ui/select'; import { Dialog, DialogContent, @@ -25,17 +18,21 @@ import { DialogHeader, DialogTitle, } from '@/components/ui/dialog'; +import { Label } from '@/components/ui/label'; +import { + Select, + SelectContent, + SelectItem, + SelectTrigger, + SelectValue, +} from '@/components/ui/select'; 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, 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; @@ -46,6 +43,8 @@ type DefaultModelDialogProps = { organizationPlan?: 'teams' | 'enterprise'; }; +type DefaultBehavior = 'auto' | 'specific'; + export function DefaultModelDialog({ open, onOpenChange, @@ -54,126 +53,92 @@ export function DefaultModelDialog({ currentDefaultModel, organizationPlan, }: DefaultModelDialogProps) { - 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 organizationAutoFallbackModel = organizationSettings?.org_auto_model?.fallback_model; - const organizationAutoEnabled = organizationDefaultModel === ORG_AUTO_MODEL.id; + const configureMutation = useConfigureOrganizationDefaultBehavior(); const organizationAutoFeatureEnabled = useFeatureFlagEnabled(ORGANIZATION_AUTO_MODEL_FLAG); const isDevelopment = process.env.NODE_ENV === 'development'; const canConfigureOrganizationAuto = organizationPlan === 'enterprise' && (isDevelopment || organizationAutoFeatureEnabled === true); - const showOrganizationAutoSection = organizationAutoEnabled || canConfigureOrganizationAuto; + const organizationDefaultModel = organizationSettings?.default_model; + const organizationAutoEnabled = organizationDefaultModel === ORG_AUTO_MODEL.id; + const showOrganizationAutoBehavior = canConfigureOrganizationAuto || organizationAutoEnabled; + const organizationAutoFallbackModel = organizationSettings?.org_auto_model?.fallback_model; + const [behavior, setBehavior] = useState( + organizationAutoEnabled ? 'auto' : 'specific' + ); + const [selectedModel, setSelectedModel] = useState(''); + const [selectedFallbackModel, setSelectedFallbackModel] = useState(''); + const availableModels = (openRouterModels?.data ?? []).filter( model => model.id !== ORG_AUTO_MODEL.id ); + const autoTargetModels = useMemo( + () => + availableModels.filter(model => { + if (model.id.startsWith(CUSTOM_LLM_PREFIX)) return false; + if (model.id.startsWith('kilo-auto/')) return isOrganizationAutoTargetModel(model.id); + return true; + }), + [availableModels] + ); + const fallbackUnavailable = + !!organizationAutoFallbackModel && + !autoTargetModels.some(model => model.id === organizationAutoFallbackModel); + const effectiveFallback = + selectedFallbackModel || organizationAutoFallbackModel || 'kilo-auto/balanced'; + const fallbackNeedsReplacement = + fallbackUnavailable && + !modelsLoading && + !autoTargetModels.some(model => model.id === effectiveFallback); + const effectiveSpecificModel = selectedModel || organizationDefaultModel || ''; + const isDirty = + behavior !== (organizationAutoEnabled ? 'auto' : 'specific') || + (behavior === 'auto' && + effectiveFallback !== (organizationAutoFallbackModel || 'kilo-auto/balanced')) || + (behavior === 'specific' && effectiveSpecificModel !== (organizationDefaultModel || '')); + useEffect(() => { if (!open) { setSelectedModel(''); setSelectedFallbackModel(''); + setBehavior(organizationAutoEnabled ? 'auto' : 'specific'); } - }, [open]); - - 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 organizationAutoFallbackUnavailable = - !!organizationAutoFallbackModel && - !organizationAutoTargetModels.some(model => model.id === organizationAutoFallbackModel); - - const handleUpdateDefaultModel = async () => { - if (!selectedModel) return; + }, [open, organizationAutoEnabled]); + const handleSave = async () => { try { - await updateDefaultModelMutation.mutateAsync({ - organizationId, - default_model: selectedModel, - }); - - setSelectedModel(''); + if (behavior === 'auto') { + await configureMutation.mutateAsync({ + organizationId, + behavior: 'auto', + fallback_model: effectiveFallback, + }); + toast.success('Organization Auto default updated'); + } else { + if (!selectedModel) { + toast.error('Choose a specific default model.'); + return; + } + await configureMutation.mutateAsync({ + organizationId, + behavior: 'specific', + specific_model: selectedModel, + }); + toast.success('Default model updated'); + } onOpenChange(false); - toast.success('Default model updated successfully'); } catch (error) { - console.error('Failed to update default model:', error); - toast.error(error instanceof Error ? error.message : 'Failed to update default model'); + toast.error(error instanceof Error ? error.message : 'Failed to update default behavior'); } }; - const handleClearDefaultModel = async () => { + const handleReset = async () => { try { - // Clear only the default model; provider/model access policy stays unchanged. - await updateDefaultModelMutation.mutateAsync({ - organizationId, - default_model: null, - }); - - setSelectedModel(''); + await configureMutation.mutateAsync({ organizationId, behavior: 'global' }); + toast.success('Reset to global default'); onOpenChange(false); - toast.success('Default model cleared - will use global default'); } catch (error) { - console.error('Failed to clear default model:', error); - toast.error(error instanceof Error ? error.message : 'Failed to clear default model'); - } - }; - - const handleEnableOrganizationAuto = async () => { - try { - await enableOrganizationAutoMutation.mutateAsync({ 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, - }); - 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' - ); + toast.error(error instanceof Error ? error.message : 'Failed to reset default model'); } }; @@ -182,199 +147,188 @@ export function DefaultModelDialog({ - - - Set Organization Default Model + + + Default model behavior - Choose a default model for this organization. Members will use this model by default - unless they specify otherwise. + Members use this model by default unless they select another model locally. -

-
- -
- - {currentDefaultModel} - - {organizationDefaultModel ? ( -
- Organization-specific default is set -
- ) : ( -
- Using global default (no organization-specific default set) -
- )} -
+
+
+ Current default + + {currentDefaultModel || 'global default'} +
- {showOrganizationAutoSection && ( -
+ {showOrganizationAutoBehavior && ( +
+ + +
+ )} + + {behavior === 'auto' ? ( +
- +

- Route the organization default by mode. Local model selections still override - it. + Used when a mode has no explicit route or the request uses an unknown mode. + {canConfigureOrganizationAuto && organizationAutoEnabled && ( + <> + {' '} + onOpenChange(false)} + > + Configure mode routes + + + )}

- {canConfigureOrganizationAuto && ( - onOpenChange(false)} - > - Configure mode routes - - )}
- {organizationAutoEnabled ? ( -

- Enabled. Choose a replacement model below to disable it. + + {fallbackNeedsReplacement && ( +

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

- ) : ( - )} - {canConfigureOrganizationAuto && organizationSettings?.org_auto_model && ( -
- -
- + + + + + {availableModels.map(model => ( + +
+ {model.id} + {model.name !== model.id && ( + {model.name} )} - {organizationAutoTargetModels.map(model => ( - -
- {model.id} - {model.name !== model.id && ( - - {model.name} - - )} -
- - ))} - - - -
- {organizationAutoFallbackUnavailable && ( -

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

- )} -
+
+ + ))} + + + {!modelsLoading && availableModels.length === 0 && ( +

+ No models available. Configure model access first. +

)}
)} - -
- - - - {availableModels.length === 0 && ( -
- No models available. Configure model access first. -
- )} -
- - {organizationDefaultModel && !organizationAutoEnabled && ( + + {organizationDefaultModel && ( )} - {organizationAutoEnabled && ( +
+ - )} - +
diff --git a/apps/web/src/routers/organizations/organization-settings-router.test.ts b/apps/web/src/routers/organizations/organization-settings-router.test.ts index 3637c3aa84..335e052787 100644 --- a/apps/web/src/routers/organizations/organization-settings-router.test.ts +++ b/apps/web/src/routers/organizations/organization-settings-router.test.ts @@ -545,6 +545,77 @@ describe('organizations settings trpc router', () => { }); }); + it('resets an active Organization Auto default to the global default', async () => { + const caller = await createCallerForUser(owner.id); + const autoOrg = await createTestOrganization('Active Auto Org', owner.id, 0, {}, false); + + await caller.organizations.settings.configureOrganizationDefaultBehavior({ + organizationId: autoOrg.id, + behavior: 'auto', + fallback_model: 'kilo-auto/balanced', + }); + await caller.organizations.settings.setOrganizationAutoRoute({ + organizationId: autoOrg.id, + mode_slug: 'code', + model_id: 'kilo-auto/frontier', + }); + + const result = await caller.organizations.settings.configureOrganizationDefaultBehavior({ + organizationId: autoOrg.id, + behavior: 'global', + }); + + expect(result.settings.default_model).toBeUndefined(); + expect(result.settings.org_auto_model).toEqual({ + routes: { code: 'kilo-auto/frontier' }, + fallback_model: 'kilo-auto/balanced', + }); + }); + + it('validates stored routes before enabling Organization Auto', async () => { + const caller = await createCallerForUser(owner.id); + const autoOrg = await createTestOrganization( + 'Invalid Auto Route Org', + owner.id, + 0, + { + default_model: 'gpt-4', + org_auto_model: { + routes: { code: 'custom-llm/stale-model' }, + fallback_model: 'kilo-auto/balanced', + }, + }, + false + ); + + await expect( + caller.organizations.settings.configureOrganizationDefaultBehavior({ + organizationId: autoOrg.id, + behavior: 'auto', + fallback_model: 'kilo-auto/balanced', + }) + ).rejects.toThrow('Cannot enable Organization Auto because route "code" is invalid'); + }); + + it('preserves non-auto specific default semantics when configuring a specific model', async () => { + const caller = await createCallerForUser(owner.id); + const specificOrg = await createTestOrganization( + 'Specific Default Org', + owner.id, + 0, + {}, + false + ); + + const result = await caller.organizations.settings.configureOrganizationDefaultBehavior({ + organizationId: specificOrg.id, + behavior: 'specific', + specific_model: 'any-model', + }); + + expect(result.settings.default_model).toBe('any-model'); + }); + it('sets and clears Organization Auto routes', async () => { const caller = await createCallerForUser(owner.id); diff --git a/apps/web/src/routers/organizations/organization-settings-router.ts b/apps/web/src/routers/organizations/organization-settings-router.ts index 1821d042fa..854e9a4632 100644 --- a/apps/web/src/routers/organizations/organization-settings-router.ts +++ b/apps/web/src/routers/organizations/organization-settings-router.ts @@ -251,6 +251,12 @@ const SetOrganizationAutoFallbackInputSchema = OrganizationIdInputSchema.extend( model_id: z.string().min(1).max(200), }); +const ConfigureOrganizationDefaultBehaviorInputSchema = OrganizationIdInputSchema.extend({ + behavior: z.enum(['auto', 'specific', 'global']), + fallback_model: z.string().min(1).max(200).optional(), + specific_model: z.string().min(1).max(200).optional(), +}); + const UpdateDataCollectionInputSchema = OrganizationIdInputSchema.extend({ dataCollection: z.enum(['allow', 'deny']).nullable(), }); @@ -774,6 +780,113 @@ export const organizationsSettingsRouter = createTRPCRouter({ return { settings: updatedSettings }; }), + configureOrganizationDefaultBehavior: organizationBillingMutationProcedure + .input(ConfigureOrganizationDefaultBehaviorInputSchema) + .output(SettingsResponseSchema) + .mutation(async ({ input, ctx }) => { + const { organizationId, behavior, fallback_model, specific_model } = input; + if (behavior === 'auto') { + await assertOrganizationAutoWriteEnabled(ctx.user.id); + } + + const existingOrg = await getOrganizationById(organizationId); + if (!existingOrg) { + throw new TRPCError({ code: 'NOT_FOUND', message: 'Organization not found' }); + } + assertOrganizationAutoEligible(existingOrg); + + const updatedSettings = await db.transaction(async tx => { + const settings = await mutateOrganizationSettings( + organizationId, + async organization => { + assertOrganizationAutoEligible(organization); + if (behavior === 'global') { + return { ...organization.settings, default_model: undefined }; + } + + if (behavior === 'specific') { + if (!specific_model) { + throw new TRPCError({ + code: 'BAD_REQUEST', + message: 'Specific model is required.', + }); + } + const modelId = + organization.settings.default_model === ORG_AUTO_MODEL.id + ? await validateOrganizationDefaultReplacement(organization, specific_model) + : await validateOrganizationDefaultModel(organization, specific_model); + return { ...organization.settings, default_model: modelId }; + } + + const orgAutoModel = + organization.settings.org_auto_model ?? DEFAULT_ORGANIZATION_AUTO_MODEL_SETTINGS; + const isEnablingOrganizationAuto = + organization.settings.default_model !== ORG_AUTO_MODEL.id; + const routes = { ...orgAutoModel.routes }; + + if (isEnablingOrganizationAuto) { + assertOrganizationAutoRouteCount(routes); + for (const [slug, targetModelId] of Object.entries(routes)) { + const routeValidation = await validateOrganizationAutoTarget( + organization, + targetModelId, + { dbClient: tx } + ); + if (routeValidation.kind === 'error') { + throw new TRPCError({ + code: 'BAD_REQUEST', + message: `Cannot enable Organization Auto because route "${slug}" is invalid: ${routeValidation.message}`, + }); + } + routes[slug] = routeValidation.modelId; + } + } + + const requestedFallback = fallback_model ?? orgAutoModel.fallback_model; + const validation = await validateOrganizationAutoTarget( + organization, + requestedFallback, + { + dbClient: tx, + } + ); + if (validation.kind === 'error') { + throw new TRPCError({ code: 'BAD_REQUEST', message: validation.message }); + } + return { + ...organization.settings, + default_model: ORG_AUTO_MODEL.id, + org_auto_model: { + ...orgAutoModel, + routes, + fallback_model: validation.modelId, + }, + }; + }, + tx + ); + await createAuditLog({ + action: 'organization.settings.change', + actor_email: ctx.user.google_user_email, + actor_id: ctx.user.id, + actor_name: ctx.user.google_user_name, + message: + behavior === 'auto' + ? 'Configured Organization Auto default behavior.' + : behavior === 'specific' + ? `Configured specific organization default model: ${settings.default_model}` + : existingOrg.settings.default_model === ORG_AUTO_MODEL.id + ? 'Disabled Organization Auto and reset organization default model to global default.' + : 'Reset organization default model to global default.', + organization_id: organizationId, + tx, + }); + return settings; + }); + + return { settings: updatedSettings }; + }), + updateDataCollection: organizationBillingMutationProcedure .input(UpdateDataCollectionInputSchema) .output(SettingsResponseSchema) From cf3f457ad7b06ab2d95337c9735b269ed1562e53 Mon Sep 17 00:00:00 2001 From: syn Date: Wed, 17 Jun 2026 13:42:21 -0500 Subject: [PATCH 7/8] fix(organizations): improve default behavior chooser --- .../DefaultModelDialog.tsx | 136 ++++++++++++++---- 1 file changed, 105 insertions(+), 31 deletions(-) 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 369d84d93a..32ce1649cf 100644 --- a/apps/web/src/components/organizations/providers-and-models/DefaultModelDialog.tsx +++ b/apps/web/src/components/organizations/providers-and-models/DefaultModelDialog.tsx @@ -1,6 +1,6 @@ 'use client'; -import { useEffect, useMemo, useState } from 'react'; +import { useEffect, useMemo, useRef, useState } from 'react'; import Link from 'next/link'; import { useFeatureFlagEnabled } from 'posthog-js/react'; import { Settings2 } from 'lucide-react'; @@ -33,6 +33,7 @@ import { ORGANIZATION_AUTO_MODEL_FLAG, } from '@/lib/organizations/organization-auto-model-shared'; import { CUSTOM_LLM_PREFIX } from '@/lib/ai-gateway/model-utils'; +import { cn } from '@/lib/utils'; type DefaultModelDialogProps = { open: boolean; @@ -45,6 +46,108 @@ type DefaultModelDialogProps = { type DefaultBehavior = 'auto' | 'specific'; +const BEHAVIOR_OPTIONS: { + value: DefaultBehavior; + title: string; + description: string; + recommended?: boolean; +}[] = [ + { + value: 'auto', + title: 'Organization Auto', + description: 'Route each mode to the right model, with one fallback.', + recommended: true, + }, + { + value: 'specific', + title: 'Specific model', + description: 'Pin a single model as the organization default.', + }, +]; + +const BEHAVIOR_ORDER = BEHAVIOR_OPTIONS.map(option => option.value); + +function BehaviorChooser({ + value, + onChange, +}: { + value: DefaultBehavior; + onChange: (value: DefaultBehavior) => void; +}) { + const itemRefs = useRef>>({}); + + const moveTo = (next: DefaultBehavior) => { + onChange(next); + itemRefs.current[next]?.focus(); + }; + + const handleKeyDown = (event: React.KeyboardEvent, current: DefaultBehavior) => { + const index = BEHAVIOR_ORDER.indexOf(current); + if (event.key === 'ArrowRight' || event.key === 'ArrowDown') { + event.preventDefault(); + moveTo(BEHAVIOR_ORDER[(index + 1) % BEHAVIOR_ORDER.length]); + } else if (event.key === 'ArrowLeft' || event.key === 'ArrowUp') { + event.preventDefault(); + moveTo(BEHAVIOR_ORDER[(index - 1 + BEHAVIOR_ORDER.length) % BEHAVIOR_ORDER.length]); + } else if (event.key === 'Home') { + event.preventDefault(); + moveTo(BEHAVIOR_ORDER[0]); + } else if (event.key === 'End') { + event.preventDefault(); + moveTo(BEHAVIOR_ORDER[BEHAVIOR_ORDER.length - 1]); + } + }; + + return ( +
+ {BEHAVIOR_OPTIONS.map(option => { + const selected = value === option.value; + return ( + + ); + })} +
+ ); +} + export function DefaultModelDialog({ open, onOpenChange, @@ -165,36 +268,7 @@ export function DefaultModelDialog({
{showOrganizationAutoBehavior && ( -
- - -
+ )} {behavior === 'auto' ? ( From f753b11d1541ab67f3c48074feeeb90c31c9132a Mon Sep 17 00:00:00 2001 From: syn Date: Wed, 17 Jun 2026 14:12:05 -0500 Subject: [PATCH 8/8] fix(organizations): capture locked default for audit logs --- .../organization-settings-router.test.ts | 18 +++++++++++++++++- .../organization-settings-router.ts | 4 +++- 2 files changed, 20 insertions(+), 2 deletions(-) diff --git a/apps/web/src/routers/organizations/organization-settings-router.test.ts b/apps/web/src/routers/organizations/organization-settings-router.test.ts index 335e052787..c83544fe7d 100644 --- a/apps/web/src/routers/organizations/organization-settings-router.test.ts +++ b/apps/web/src/routers/organizations/organization-settings-router.test.ts @@ -10,7 +10,12 @@ import type { OpenRouterModel, OpenRouterModelsResponse, } from '@/lib/organizations/organization-types'; -import { type User, type Organization, organizations } from '@kilocode/db/schema'; +import { + type User, + type Organization, + organization_audit_logs, + organizations, +} from '@kilocode/db/schema'; import { eq } from 'drizzle-orm'; import { randomUUID } from 'crypto'; import { db } from '@/lib/drizzle'; @@ -570,6 +575,17 @@ describe('organizations settings trpc router', () => { routes: { code: 'kilo-auto/frontier' }, fallback_model: 'kilo-auto/balanced', }); + const auditLogs = await db.query.organization_audit_logs.findMany({ + where: eq(organization_audit_logs.organization_id, autoOrg.id), + }); + expect(auditLogs).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + message: + 'Disabled Organization Auto and reset organization default model to global default.', + }), + ]) + ); }); it('validates stored routes before enabling Organization Auto', async () => { diff --git a/apps/web/src/routers/organizations/organization-settings-router.ts b/apps/web/src/routers/organizations/organization-settings-router.ts index 854e9a4632..4d7413a200 100644 --- a/apps/web/src/routers/organizations/organization-settings-router.ts +++ b/apps/web/src/routers/organizations/organization-settings-router.ts @@ -795,10 +795,12 @@ export const organizationsSettingsRouter = createTRPCRouter({ } assertOrganizationAutoEligible(existingOrg); + let previousDefaultModel: string | undefined; const updatedSettings = await db.transaction(async tx => { const settings = await mutateOrganizationSettings( organizationId, async organization => { + previousDefaultModel = organization.settings.default_model; assertOrganizationAutoEligible(organization); if (behavior === 'global') { return { ...organization.settings, default_model: undefined }; @@ -875,7 +877,7 @@ export const organizationsSettingsRouter = createTRPCRouter({ ? 'Configured Organization Auto default behavior.' : behavior === 'specific' ? `Configured specific organization default model: ${settings.default_model}` - : existingOrg.settings.default_model === ORG_AUTO_MODEL.id + : previousDefaultModel === ORG_AUTO_MODEL.id ? 'Disabled Organization Auto and reset organization default model to global default.' : 'Reset organization default model to global default.', organization_id: organizationId,