diff --git a/apps/web/src/app/(app)/organizations/[id]/custom-modes/page.tsx b/apps/web/src/app/(app)/organizations/[id]/custom-modes/page.tsx index e28b0ebbde..86b0a54b0c 100644 --- a/apps/web/src/app/(app)/organizations/[id]/custom-modes/page.tsx +++ b/apps/web/src/app/(app)/organizations/[id]/custom-modes/page.tsx @@ -9,7 +9,13 @@ export default async function OrganizationCustomModesPage({ return ( } + render={({ organization, role, isGlobalAdmin }) => ( + + )} /> ); } 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 7914abe091..265cc9724a 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 7119240a7a..308305c462 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'; @@ -48,6 +49,7 @@ import { modelNotAllowedResponse, extractHeaderAndLimitLength, noFreeModelsAvailableResponse, + organizationAutoConfigurationResponse, temporarilyUnavailableResponse, usageLimitExceededResponse, wrapInSafeNextResponse, @@ -95,6 +97,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'; @@ -242,6 +245,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(); @@ -271,6 +281,7 @@ export async function POST(request: NextRequest): Promise res.user), @@ -319,6 +331,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..1d971ed64f 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,17 @@ 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)) { + console.warn('organization_auto_invalid_default', { organizationId: organization.id }); + 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..dea9edaccb 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 @@ -11,7 +11,7 @@ const mockedGetAuthorizedOrgContext = jest.mocked(getAuthorizedOrgContext); const mockedGetAllOrganizationModes = jest.mocked(getAllOrganizationModes); describe('GET /api/organizations/[id]/modes', () => { - test('returns defaultModel as part of the additive mode payload', async () => { + test('returns the direct mode payload without Organization Auto route projection', async () => { mockedGetAuthorizedOrgContext.mockResolvedValue({ success: true, data: { @@ -30,7 +30,6 @@ describe('GET /api/organizations/[id]/modes', () => { config: { roleDefinition: 'You are a coding assistant', groups: ['read'], - defaultModel: 'openai/gpt-4o', }, }, ]); @@ -53,57 +52,10 @@ describe('GET /api/organizations/[id]/modes', () => { config: { roleDefinition: 'You are a coding assistant', groups: ['read'], - defaultModel: 'openai/gpt-4o', }, }, ], }); expect(mockedGetAllOrganizationModes).toHaveBeenCalledWith('org-1'); }); - - test('returns a legacy mode row without defaultModel unchanged', async () => { - mockedGetAuthorizedOrgContext.mockResolvedValue({ - success: true, - data: { - organization: { id: 'org-1' }, - }, - } 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'], - }, - }, - ]); - - 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'], - }, - }, - ], - }); - }); }); 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..5d483dded1 100644 --- a/apps/web/src/app/api/organizations/[id]/modes/route.ts +++ b/apps/web/src/app/api/organizations/[id]/modes/route.ts @@ -15,9 +15,8 @@ export async function GET( } const { organization } = data; - const modes = await getAllOrganizationModes(organization.id); return NextResponse.json({ - modes, + modes: await getAllOrganizationModes(organization.id), }); } diff --git a/apps/web/src/app/api/organizations/hooks.ts b/apps/web/src/app/api/organizations/hooks.ts index 2093d3affd..5d74d17645 100644 --- a/apps/web/src/app/api/organizations/hooks.ts +++ b/apps/web/src/app/api/organizations/hooks.ts @@ -173,28 +173,97 @@ export function useUpdateCompanyDomain() { ); } -export function useUpdateOrganizationSettings() { +const useInvalidateOrganizationDataAndDefaults = () => { const trpc = useTRPC(); const queryClient = useQueryClient(); + return function (_: unknown, { organizationId }: { organizationId: Organization['id'] }) { + void queryClient.invalidateQueries({ queryKey: trpc.organizations.pathKey() }); + void queryClient.invalidateQueries({ queryKey: ['organization-defaults', organizationId] }); + }; +}; + +export function useUpdateOrganizationSettings() { + const trpc = useTRPC(); + const onSuccess = useInvalidateOrganizationDataAndDefaults(); return useMutation( trpc.organizations.settings.updateAllowLists.mutationOptions({ + onSuccess, + }) + ); +} + +export function useUpdateDefaultModel() { + const trpc = useTRPC(); + const onSuccess = useInvalidateOrganizationDataAndDefaults(); + + return useMutation( + trpc.organizations.settings.updateDefaultModel.mutationOptions({ + onSuccess, + }) + ); +} + +export function useEnableOrganizationAuto() { + const trpc = useTRPC(); + const onSuccess = useInvalidateOrganizationDataAndDefaults(); + return useMutation( + trpc.organizations.settings.enableOrganizationAuto.mutationOptions({ + onSuccess, + }) + ); +} + +export function useDisableOrganizationAuto() { + const trpc = useTRPC(); + const onSuccess = useInvalidateOrganizationDataAndDefaults(); + return useMutation( + trpc.organizations.settings.disableOrganizationAuto.mutationOptions({ + onSuccess, + }) + ); +} + +export function useConfigureOrganizationDefaultBehavior() { + const trpc = useTRPC(); + const onSuccess = useInvalidateOrganizationDataAndDefaults(); + return useMutation( + trpc.organizations.settings.configureOrganizationDefaultBehavior.mutationOptions({ + onSuccess, + }) + ); +} + +export function useSetOrganizationAutoFallback() { + const trpc = useTRPC(); + const queryClient = useQueryClient(); + return useMutation( + trpc.organizations.settings.setOrganizationAutoFallback.mutationOptions({ onSuccess: () => { - // lazy-mode invalidate everything related to an org if settings change void queryClient.invalidateQueries({ queryKey: trpc.organizations.pathKey() }); }, }) ); } -export function useUpdateDefaultModel() { +export function useSetOrganizationAutoRoute() { const trpc = useTRPC(); const queryClient = useQueryClient(); + return useMutation( + trpc.organizations.settings.setOrganizationAutoRoute.mutationOptions({ + onSuccess: () => { + void queryClient.invalidateQueries({ queryKey: trpc.organizations.pathKey() }); + }, + }) + ); +} +export function useClearOrganizationAutoRoute() { + const trpc = useTRPC(); + const queryClient = useQueryClient(); return useMutation( - trpc.organizations.settings.updateDefaultModel.mutationOptions({ + trpc.organizations.settings.clearOrganizationAutoRoute.mutationOptions({ onSuccess: () => { - // lazy-mode invalidate everything related to an org if settings change void queryClient.invalidateQueries({ queryKey: trpc.organizations.pathKey() }); }, }) @@ -492,56 +561,61 @@ export function useOrganizationModeById(organizationId: string, modeId: string) ); } -export function useCreateOrganizationMode() { +function useInvalidateOrganizationModes() { const trpc = useTRPC(); const queryClient = useQueryClient(); + return async (_: unknown, variables: { organizationId: Organization['id']; modeId?: string }) => { + await Promise.all([ + queryClient.invalidateQueries({ + queryKey: trpc.organizations.modes.list.queryKey({ + organizationId: variables.organizationId, + }), + }), + queryClient.invalidateQueries({ + queryKey: trpc.organizations.withMembers.queryKey({ + organizationId: variables.organizationId, + }), + }), + ...(variables.modeId + ? [ + queryClient.invalidateQueries({ + queryKey: trpc.organizations.modes.getById.queryKey({ + organizationId: variables.organizationId, + modeId: variables.modeId, + }), + }), + ] + : []), + ]); + }; +} + +export function useCreateOrganizationMode() { + const trpc = useTRPC(); + const onSuccess = useInvalidateOrganizationModes(); return useMutation( trpc.organizations.modes.create.mutationOptions({ - onSuccess: (_, variables) => { - void queryClient.invalidateQueries({ - queryKey: trpc.organizations.modes.list.queryKey({ - organizationId: variables.organizationId, - }), - }); - }, + onSuccess, }) ); } export function useUpdateOrganizationMode() { const trpc = useTRPC(); - const queryClient = useQueryClient(); + const onSuccess = useInvalidateOrganizationModes(); return useMutation( trpc.organizations.modes.update.mutationOptions({ - onSuccess: (_, variables) => { - void queryClient.invalidateQueries({ - queryKey: trpc.organizations.modes.list.queryKey({ - organizationId: variables.organizationId, - }), - }); - void queryClient.invalidateQueries({ - queryKey: trpc.organizations.modes.getById.queryKey({ - organizationId: variables.organizationId, - modeId: variables.modeId, - }), - }); - }, + onSuccess, }) ); } export function useDeleteOrganizationMode() { const trpc = useTRPC(); - const queryClient = useQueryClient(); + const onSuccess = useInvalidateOrganizationModes(); return useMutation( trpc.organizations.modes.delete.mutationOptions({ - onSuccess: (_, variables) => { - void queryClient.invalidateQueries({ - queryKey: trpc.organizations.modes.list.queryKey({ - organizationId: variables.organizationId, - }), - }); - }, + onSuccess, }) ); } diff --git a/apps/web/src/components/models/CondensedProviderAndModelsList.tsx b/apps/web/src/components/models/CondensedProviderAndModelsList.tsx index 2bb1aa9cf3..667296ce74 100644 --- a/apps/web/src/components/models/CondensedProviderAndModelsList.tsx +++ b/apps/web/src/components/models/CondensedProviderAndModelsList.tsx @@ -57,10 +57,6 @@ export function CondensedProviderAndModelsList({ return
Loading providers...
; } - if (!selections || providersWithSelections.length === 0) { - return null; - } - return (
{/* Default model row */} @@ -97,7 +93,7 @@ export function CondensedProviderAndModelsList({

You don't have permission to change the default model. Contact your organization - owner to update this setting. + owner or billing manager to update this setting.

diff --git a/apps/web/src/components/organizations/OrganizationDashboard.tsx b/apps/web/src/components/organizations/OrganizationDashboard.tsx index 5b69d93365..2adc4ccb8f 100644 --- a/apps/web/src/components/organizations/OrganizationDashboard.tsx +++ b/apps/web/src/components/organizations/OrganizationDashboard.tsx @@ -149,7 +149,9 @@ export function OrganizationDashboard({ ) : ( )} 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} /> void; onEditClick: (mode: DisplayMode) => void; isDefaultModelConfigEnabled: boolean; + isOrganizationAutoDefaultActive: boolean; }; function ModesList({ modes, readonly, + canManageRoutes, onDeleteClick, onEditClick, isDefaultModelConfigEnabled, + isOrganizationAutoDefaultActive, }: ModesListProps) { return (
- {modes.map((mode, index) => ( - - - -
-
- {mode.name} -

- {mode.config.description} -

-
- {!readonly && ( -
- - {mode.isOverridden ? ( - - ) : ( - !mode.isDefault && ( + {mode.isOverridden ? ( - ) - )} -
- )} -
- - -
- {isDefaultModelConfigEnabled && mode.config.defaultModel && ( -
- Default model - - {mode.config.defaultModel} - -
- )} - {mode.config?.groups && mode.config.groups.length > 0 && ( -
-

Available Tools

-
- {mode.config.groups.map((group, idx) => { - const groupName = Array.isArray(group) ? group[0] : group; - const groupConfig = Array.isArray(group) ? group[1] : null; - const hasRestriction = !!groupConfig; - const tooltipText = groupConfig - ? `Restricted file access: ${groupConfig.description || ''} (${groupConfig.fileRegex})`.trim() - : undefined; - - const badgeContent = ( - onDeleteClick(mode)} + disabled={lifecycleRouteLocked} + title={ + lifecycleRouteLocked + ? 'Route managers must delete a routed mode.' + : undefined + } + className="text-red-400 hover:text-red-500" > - {groupName} - {hasRestriction && ' *'} - - ); + + + ) + )} +
+ )} +
+ + +
+ {isDefaultModelConfigEnabled && mode.routeModel && ( +
+ + {isOrganizationAutoDefaultActive + ? 'Organization Auto route' + : 'Saved Organization Auto route (inactive)'} + + + {mode.routeModel} + +
+ )} + {mode.config?.groups && mode.config.groups.length > 0 && ( +
+

Available Tools

+
+ {mode.config.groups.map((group, idx) => { + const groupName = Array.isArray(group) ? group[0] : group; + const groupConfig = Array.isArray(group) ? group[1] : null; + const hasRestriction = !!groupConfig; + const tooltipText = groupConfig + ? `Restricted file access: ${groupConfig.description || ''} (${groupConfig.fileRegex})`.trim() + : undefined; - return hasRestriction ? ( - - {badgeContent} - {tooltipText} - - ) : ( - badgeContent - ); - })} + const badgeContent = ( + + {groupName} + {hasRestriction && ' *'} + + ); + + return hasRestriction ? ( + + {badgeContent} + {tooltipText} + + ) : ( + badgeContent + ); + })} +
-
- )} -
- - - - ))} + )} +
+ + + + ); + })}
); } -export function CustomModesLayout({ organizationId }: CustomModesLayoutProps) { +export function CustomModesLayout({ organizationId, role, isGlobalAdmin }: CustomModesLayoutProps) { const { data, isLoading, error } = useOrganizationModes(organizationId); const { data: organizationData } = useOrganizationWithMembers(organizationId); const deleteMutation = useDeleteOrganizationMode(); @@ -174,10 +206,14 @@ 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_FLAG); const isDevelopment = process.env.NODE_ENV === 'development'; const isDefaultModelConfigEnabled = isDevelopment || isDefaultModelFeatureEnabled === true; - const canSetDefaultModel = organizationData?.plan === 'enterprise'; + const canMaintainRoutedMode = role === 'owner' || role === 'billing_manager' || isGlobalAdmin; + const canSetDefaultModel = + organizationData?.plan === 'enterprise' && isDefaultModelConfigEnabled && canMaintainRoutedMode; + const isOrganizationAutoDefaultActive = + organizationData?.settings.default_model === ORG_AUTO_MODEL.id; const readonly = isReadOnly; // Separate built-in modes and custom modes @@ -197,6 +233,9 @@ export function CustomModesLayout({ organizationId }: CustomModesLayoutProps) { ...customMode, isDefault: true, isOverridden: true, + routeModel: organizationData + ? getOrganizationAutoRoute(organizationData.settings, customMode.slug) + : undefined, }; } else { // This default mode is not overridden @@ -207,6 +246,9 @@ export function CustomModesLayout({ organizationId }: CustomModesLayoutProps) { slug: defaultMode.slug, name: defaultMode.name, config: defaultMode.config, + routeModel: organizationData + ? getOrganizationAutoRoute(organizationData.settings, defaultMode.slug) + : undefined, created_by: '', created_at: new Date().toISOString(), updated_at: new Date().toISOString(), @@ -224,6 +266,9 @@ export function CustomModesLayout({ organizationId }: CustomModesLayoutProps) { ...mode, isDefault: false, isOverridden: false, + routeModel: organizationData + ? getOrganizationAutoRoute(organizationData.settings, mode.slug) + : undefined, })) .sort((a, b) => a.name.localeCompare(b.name)); @@ -231,7 +276,7 @@ export function CustomModesLayout({ organizationId }: CustomModesLayoutProps) { builtInModes: builtInDisplayModes, customModes: customDisplayModes, }; - }, [data?.modes, organizationId]); + }, [data?.modes, organizationData?.settings, organizationId]); const handleDelete = async () => { if (!modeToDelete) return; @@ -242,6 +287,7 @@ export function CustomModesLayout({ organizationId }: CustomModesLayoutProps) { await deleteMutation.mutateAsync({ organizationId, modeId: modeToDelete.id, + ...(modeToDelete.isDefault && modeToDelete.isOverridden ? { preserve_route: true } : {}), }); toast.success(`Mode "${modeToDelete.name}" ${action} successfully`); setDeleteDialogOpen(false); @@ -319,9 +365,11 @@ export function CustomModesLayout({ organizationId }: CustomModesLayoutProps) { organizationId={organizationId} modes={builtInModes} readonly={readonly} + canManageRoutes={canMaintainRoutedMode} onDeleteClick={openDeleteDialog} onEditClick={handleEditMode} isDefaultModelConfigEnabled={isDefaultModelConfigEnabled} + isOrganizationAutoDefaultActive={isOrganizationAutoDefaultActive} /> @@ -333,9 +381,11 @@ export function CustomModesLayout({ organizationId }: CustomModesLayoutProps) { organizationId={organizationId} modes={customModes} readonly={readonly} + canManageRoutes={canMaintainRoutedMode} onDeleteClick={openDeleteDialog} onEditClick={handleEditMode} isDefaultModelConfigEnabled={isDefaultModelConfigEnabled} + isOrganizationAutoDefaultActive={isOrganizationAutoDefaultActive} /> )} @@ -395,10 +445,7 @@ export function CustomModesLayout({ organizationId }: CustomModesLayoutProps) { footer={
-
} @@ -421,6 +468,7 @@ export function CustomModesLayout({ organizationId }: CustomModesLayoutProps) { defaultModeSlug={editingMode.isDefault ? editingMode.slug : undefined} isDefaultModelConfigEnabled={isDefaultModelConfigEnabled} canSetDefaultModel={canSetDefaultModel} + canMaintainRoutedMode={canMaintainRoutedMode} onSuccess={handleDrawerClose} onCancel={handleDrawerClose} /> diff --git a/apps/web/src/components/organizations/custom-modes/EditModeForm.test.ts b/apps/web/src/components/organizations/custom-modes/EditModeForm.test.ts index 089bc093f7..4d5e900ffa 100644 --- a/apps/web/src/components/organizations/custom-modes/EditModeForm.test.ts +++ b/apps/web/src/components/organizations/custom-modes/EditModeForm.test.ts @@ -3,21 +3,25 @@ import { matchesBuiltInModeState } from './EditModeForm'; import { DEFAULT_MODES } from './default-modes'; import type { ModeFormData } from './ModeForm'; -function codeModeForm(overrides: Partial = {}): ModeFormData { - const codeMode = DEFAULT_MODES.find(mode => mode.slug === 'code')!; +function builtInModeForm(slug: string, overrides: Partial = {}): ModeFormData { + const mode = DEFAULT_MODES.find(defaultMode => defaultMode.slug === slug)!; return { - name: codeMode.name, - slug: codeMode.slug, - roleDefinition: codeMode.config.roleDefinition || '', - description: codeMode.config.description || '', - whenToUse: codeMode.config.whenToUse || '', - groups: [...(codeMode.config.groups || [])], - customInstructions: codeMode.config.customInstructions || '', + name: mode.name, + slug: mode.slug, + roleDefinition: mode.config.roleDefinition || '', + description: mode.config.description || '', + whenToUse: mode.config.whenToUse || '', + groups: [...(mode.config.groups || [])], + customInstructions: mode.config.customInstructions || '', ...overrides, }; } +function codeModeForm(overrides: Partial = {}): ModeFormData { + return builtInModeForm('code', overrides); +} + describe('matchesBuiltInModeState', () => { test('returns true for a built-in mode state with reordered groups', () => { const formData = codeModeForm({ groups: [...codeModeForm().groups].reverse() }); @@ -25,6 +29,14 @@ describe('matchesBuiltInModeState', () => { expect(matchesBuiltInModeState(formData, 'code')).toBe(true); }); + test('returns true for a built-in mode state with reordered configured groups', () => { + const formData = builtInModeForm('architect', { + groups: [...builtInModeForm('architect').groups].reverse(), + }); + + expect(matchesBuiltInModeState(formData, 'architect')).toBe(true); + }); + test('returns false when another customization remains', () => { const formData = codeModeForm({ description: 'Customized description' }); diff --git a/apps/web/src/components/organizations/custom-modes/EditModeForm.tsx b/apps/web/src/components/organizations/custom-modes/EditModeForm.tsx index f4c9c17033..362ebfe479 100644 --- a/apps/web/src/components/organizations/custom-modes/EditModeForm.tsx +++ b/apps/web/src/components/organizations/custom-modes/EditModeForm.tsx @@ -5,12 +5,14 @@ import { useUpdateOrganizationMode, useOrganizationModes, useDeleteOrganizationMode, + useOrganizationWithMembers, } from '@/app/api/organizations/hooks'; import { ModeForm, type ModeFormData } from './ModeForm'; import { LoadingCard } from '@/components/LoadingCard'; import { ErrorCard } from '@/components/ErrorCard'; import { toast } from 'sonner'; import { DEFAULT_MODES } from './default-modes'; +import { ORG_AUTO_MODEL } from '@/lib/ai-gateway/auto-model'; type EditModeFormProps = { organizationId: string; @@ -18,6 +20,7 @@ type EditModeFormProps = { defaultModeSlug?: string; isDefaultModelConfigEnabled?: boolean; canSetDefaultModel?: boolean; + canMaintainRoutedMode?: boolean; onSuccess?: () => void; onCancel?: () => void; }; @@ -31,7 +34,17 @@ function normalizeGroups(groups: unknown): string[] | undefined { return undefined; } - return groups.map(group => JSON.stringify(group)).sort(); + return groups + .map(group => { + if (Array.isArray(group) && group[0] === 'edit') { + return JSON.stringify([ + 'edit', + { fileRegex: group[1]?.fileRegex ?? '', description: group[1]?.description ?? '' }, + ]); + } + return JSON.stringify(group); + }) + .sort(); } export function matchesBuiltInModeState(formData: ModeFormData, defaultModeSlug: string): boolean { @@ -58,6 +71,7 @@ export function EditModeForm({ defaultModeSlug, isDefaultModelConfigEnabled = false, canSetDefaultModel = true, + canMaintainRoutedMode = true, onSuccess, onCancel, }: EditModeFormProps) { @@ -65,29 +79,35 @@ export function EditModeForm({ const { data: modesData } = useOrganizationModes(organizationId); const updateMutation = useUpdateOrganizationMode(); const deleteMutation = useDeleteOrganizationMode(); + const { data: organizationData, isLoading: isOrganizationLoading } = + useOrganizationWithMembers(organizationId); + const currentRouteModel = defaultModeSlug + ? organizationData?.settings.org_auto_model?.routes[defaultModeSlug] + : data?.mode + ? organizationData?.settings.org_auto_model?.routes[data.mode.slug] + : undefined; + const isOrganizationAutoDefaultActive = + organizationData?.settings.default_model === ORG_AUTO_MODEL.id; const handleSubmit = async (formData: ModeFormData) => { try { - if ( - defaultModeSlug && - !formData.defaultModel && - matchesBuiltInModeState(formData, defaultModeSlug) - ) { + const nextRouteModel = formData.defaultModel || undefined; + if (defaultModeSlug && matchesBuiltInModeState(formData, defaultModeSlug)) { + if (currentRouteModel && !canMaintainRoutedMode) { + toast.error('Route managers must revert a routed built-in mode.'); + return; + } await deleteMutation.mutateAsync({ organizationId, modeId, + preserve_route: true, + ...(nextRouteModel === currentRouteModel ? {} : { route_model: nextRouteModel ?? null }), }); toast.success(`Mode "${formData.name}" reverted successfully`); onSuccess?.(); return; } - const persistedDefaultModel = data?.mode?.config.defaultModel ?? ''; - const defaultModelUpdate = - formData.defaultModel === persistedDefaultModel - ? {} - : { defaultModel: formData.defaultModel || null }; - await updateMutation.mutateAsync({ organizationId, modeId, @@ -99,8 +119,8 @@ export function EditModeForm({ whenToUse: formData.whenToUse, groups: formData.groups as ('read' | 'edit' | 'browser' | 'command' | 'mcp')[], customInstructions: formData.customInstructions, - ...defaultModelUpdate, }, + ...(nextRouteModel === currentRouteModel ? {} : { route_model: nextRouteModel ?? null }), }); toast.success(`Mode "${formData.name}" updated successfully`); onSuccess?.(); @@ -130,14 +150,16 @@ export function EditModeForm({ null} /> ); } diff --git a/apps/web/src/components/organizations/custom-modes/ModeForm.tsx b/apps/web/src/components/organizations/custom-modes/ModeForm.tsx index 58b9e1e2d4..181d74ebd0 100644 --- a/apps/web/src/components/organizations/custom-modes/ModeForm.tsx +++ b/apps/web/src/components/organizations/custom-modes/ModeForm.tsx @@ -1,7 +1,7 @@ 'use client'; import type { FormEvent } from 'react'; -import { useMemo, useState, useEffect } from 'react'; +import { useEffect, useMemo, useRef, useState } from 'react'; import * as z from 'zod'; import { Button } from '@/components/ui/button'; import { Input } from '@/components/ui/input'; @@ -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-shared'; +import { CUSTOM_LLM_PREFIX } from '@/lib/ai-gateway/model-utils'; const availableGroups = [ { value: 'read', label: 'Read Files' }, @@ -54,11 +56,14 @@ export type ModeFormData = z.infer; type ModeFormProps = { organizationId: string; mode?: OrganizationMode; + routeModel?: string; onSubmit: (data: ModeFormData) => Promise; isSubmitting: boolean; isEditingBuiltIn?: boolean; isDefaultModelConfigEnabled?: boolean; + isOrganizationAutoDefaultActive?: boolean; canSetDefaultModel?: boolean; + disableSlug?: boolean; existingModes?: OrganizationMode[]; onCancel?: () => void; renderButtons?: (props: { isDirty: boolean; isSubmitting: boolean }) => React.ReactNode; @@ -100,11 +105,14 @@ function denormalizeGroups( export function ModeForm({ organizationId, mode, + routeModel, onSubmit, isSubmitting, isEditingBuiltIn = false, isDefaultModelConfigEnabled = false, + isOrganizationAutoDefaultActive = false, canSetDefaultModel = true, + disableSlug = false, existingModes = [], onCancel, renderButtons, @@ -116,7 +124,7 @@ export function ModeForm({ description: mode?.config?.description || '', whenToUse: mode?.config?.whenToUse || '', customInstructions: mode?.config?.customInstructions || '', - defaultModel: mode?.config?.defaultModel || '', + defaultModel: routeModel || '', }); const [selectedGroups, setSelectedGroups] = useState(() => { const { simpleGroups } = normalizeGroups(mode?.config?.groups || []); @@ -134,7 +142,7 @@ export function ModeForm({ description: mode?.config?.description || '', whenToUse: mode?.config?.whenToUse || '', customInstructions: mode?.config?.customInstructions || '', - defaultModel: mode?.config?.defaultModel || '', + defaultModel: routeModel || '', }); const [initialGroups, setInitialGroups] = useState(() => { const { simpleGroups } = normalizeGroups(mode?.config?.groups || []); @@ -145,6 +153,8 @@ export function ModeForm({ return editConfig || { fileRegex: '', description: '' }; }); const [selectedTemplate, setSelectedTemplate] = useState(''); + const routeFieldDirtyRef = useRef(false); + routeFieldDirtyRef.current = formData.defaultModel !== initialFormData.defaultModel; // Fetch mode templates const { data: templates, isLoading: templatesLoading } = useModeTemplates(); @@ -153,7 +163,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 && @@ -164,8 +186,17 @@ export function ModeForm({ const shouldShowDefaultModelControl = isDefaultModelConfigEnabled && (canSetDefaultModel || !!formData.defaultModel); const defaultModelChanged = formData.defaultModel !== initialFormData.defaultModel; - - // Update form data when mode prop changes + const hasStoredInactiveRoute = !isOrganizationAutoDefaultActive && !!formData.defaultModel; + const routeLabel = isOrganizationAutoDefaultActive + ? 'Route Organization Auto to' + : hasStoredInactiveRoute + ? 'Saved Organization Auto route' + : 'Organization Auto route'; + const routeFallbackLabel = isOrganizationAutoDefaultActive + ? 'Use Organization Auto fallback' + : 'Use Organization Auto fallback when enabled'; + + // Re-seed the full form only when switching modes, not when settings refetch. useEffect(() => { if (mode) { const newFormData = { @@ -175,7 +206,7 @@ export function ModeForm({ description: mode.config?.description || '', whenToUse: mode.config?.whenToUse || '', customInstructions: mode.config?.customInstructions || '', - defaultModel: mode.config?.defaultModel || '', + defaultModel: routeModel || '', }; const { simpleGroups, editConfig } = normalizeGroups(mode.config?.groups || []); const newEditConfig = editConfig || { fileRegex: '', description: '' }; @@ -187,7 +218,16 @@ export function ModeForm({ setInitialGroups(simpleGroups); setInitialEditConfig(newEditConfig); } - }, [mode]); + }, [mode?.id]); + + useEffect(() => { + if (routeFieldDirtyRef.current) { + return; + } + const nextRouteModel = routeModel || ''; + setFormData(previous => ({ ...previous, defaultModel: nextRouteModel })); + setInitialFormData(previous => ({ ...previous, defaultModel: nextRouteModel })); + }, [routeModel]); // Check if form is dirty (has changes) const isDirty = @@ -198,7 +238,7 @@ export function ModeForm({ formData.whenToUse !== initialFormData.whenToUse || formData.customInstructions !== initialFormData.customInstructions || formData.defaultModel !== initialFormData.defaultModel || - JSON.stringify(selectedGroups.sort()) !== JSON.stringify(initialGroups.sort()) || + JSON.stringify([...selectedGroups].sort()) !== JSON.stringify([...initialGroups].sort()) || editGroupConfig.fileRegex !== initialEditConfig.fileRegex || editGroupConfig.description !== initialEditConfig.description; @@ -378,12 +418,14 @@ export function ModeForm({ value={formData.slug} onChange={e => setFormData(prev => ({ ...prev, slug: e.target.value }))} placeholder="e.g., code" - disabled={isSubmitting || isEditingBuiltIn} + disabled={isSubmitting || isEditingBuiltIn || disableSlug} />

{isEditingBuiltIn ? 'Built-in mode slugs cannot be changed' - : 'Unique identifier for this mode.'} + : disableSlug + ? 'Route managers must rename routed modes.' + : 'Unique identifier for this mode.'}

{errors.slug &&

{errors.slug}

} @@ -456,7 +498,7 @@ export function ModeForm({ {shouldShowDefaultModelControl && (
- +

{!canSetDefaultModel - ? 'This organization must be on Enterprise to set mode defaults. Existing defaults 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.'} + ? 'Organization Auto routes are read-only for your role or plan.' + : !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 6f5cd9aed0..c7ff426daa 100644 --- a/apps/web/src/components/organizations/custom-modes/NewModeForm.tsx +++ b/apps/web/src/components/organizations/custom-modes/NewModeForm.tsx @@ -1,11 +1,20 @@ 'use client'; import { useSearchParams } from 'next/navigation'; -import { useCreateOrganizationMode, useOrganizationModes } from '@/app/api/organizations/hooks'; +import { + useClearOrganizationAutoRoute, + useCreateOrganizationMode, + useOrganizationModes, + useOrganizationWithMembers, + useSetOrganizationAutoRoute, +} from '@/app/api/organizations/hooks'; import { ModeForm, type ModeFormData } from './ModeForm'; +import { matchesBuiltInModeState } from './EditModeForm'; 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; @@ -26,7 +35,10 @@ export function NewModeForm({ }: NewModeFormProps) { const searchParams = useSearchParams(); const createMutation = useCreateOrganizationMode(); + const setRouteMutation = useSetOrganizationAutoRoute(); + const clearRouteMutation = useClearOrganizationAutoRoute(); const { data: modesData } = useOrganizationModes(organizationId); + const { data: organizationData } = useOrganizationWithMembers(organizationId); // Check if we're editing a default mode (from prop or search params) const defaultModeSlug = propDefaultModeSlug || searchParams.get('defaultMode'); @@ -34,6 +46,11 @@ export function NewModeForm({ if (!defaultModeSlug) return undefined; return DEFAULT_MODES.find(m => m.slug === defaultModeSlug); }, [defaultModeSlug]); + 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(() => { @@ -50,8 +67,28 @@ export function NewModeForm({ }; }, [defaultMode, organizationId]); + const persistRoute = async (modeSlug: string, targetModelId: string | undefined) => { + if (targetModelId) { + await setRouteMutation.mutateAsync({ + organizationId, + mode_slug: modeSlug, + model_id: targetModelId, + }); + } else { + await clearRouteMutation.mutateAsync({ organizationId, mode_slug: modeSlug }); + } + }; + const handleSubmit = async (data: ModeFormData) => { try { + if (defaultModeSlug && matchesBuiltInModeState(data, defaultModeSlug)) { + await persistRoute(defaultModeSlug, data.defaultModel); + toast.success(`Mode "${data.name}" route updated successfully`); + onSuccess?.(); + return; + } + + const nextRouteModel = data.defaultModel || undefined; await createMutation.mutateAsync({ organizationId, name: data.name, @@ -62,8 +99,8 @@ export function NewModeForm({ whenToUse: data.whenToUse, groups: data.groups as ('read' | 'edit' | 'browser' | 'command' | 'mcp')[], customInstructions: data.customInstructions, - ...(data.defaultModel ? { defaultModel: data.defaultModel } : {}), }, + ...(nextRouteModel === routeModel ? {} : { route_model: nextRouteModel ?? null }), }); toast.success(`Mode "${data.name}" created successfully`); onSuccess?.(); @@ -78,14 +115,17 @@ export function NewModeForm({ 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 5434c68c21..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,18 +1,15 @@ 'use client'; -import { useState } from 'react'; -import { useQueryClient } from '@tanstack/react-query'; -import { LockableContainer } from '../LockableContainer'; -import { useUpdateDefaultModel } from '@/app/api/organizations/hooks'; +import { useEffect, useMemo, useRef, 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, @@ -21,9 +18,22 @@ 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 { cn } from '@/lib/utils'; type DefaultModelDialogProps = { open: boolean; @@ -31,65 +41,207 @@ type DefaultModelDialogProps = { organizationId: string; organizationSettings?: OrganizationSettings; currentDefaultModel?: string; + organizationPlan?: 'teams' | 'enterprise'; }; +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, organizationId, organizationSettings, currentDefaultModel, + organizationPlan, }: DefaultModelDialogProps) { - const queryClient = useQueryClient(); - const [selectedModel, setSelectedModel] = useState(''); - const { data: openRouterModels, isLoading: modelsLoading } = useModelSelectorList(organizationId); - const updateDefaultModelMutation = useUpdateDefaultModel(); - + 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 organizationDefaultModel = organizationSettings?.default_model; - const availableModels = openRouterModels?.data ?? []; - - const handleUpdateDefaultModel = async () => { - if (!selectedModel) return; - - try { - await updateDefaultModelMutation.mutateAsync({ - organizationId, - default_model: selectedModel, - }); + 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(''); - // Invalidate the defaults query to refresh the display - await queryClient.invalidateQueries({ - queryKey: ['organization-defaults', organizationId], - }); + 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, organizationAutoEnabled]); + + const handleSave = async () => { + try { + 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, - }); - - await queryClient.invalidateQueries({ - queryKey: ['organization-defaults', organizationId], - }); - - 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'); + toast.error(error instanceof Error ? error.message : 'Failed to reset default model'); } }; @@ -98,85 +250,159 @@ 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'} +
-
- - - - {availableModels.length === 0 && ( -
- No models available. Configure model access first. + {showOrganizationAutoBehavior && ( + + )} + + {behavior === 'auto' ? ( +
+
+ +

+ Used when a mode has no explicit route or the request uses an unknown mode. + {canConfigureOrganizationAuto && organizationAutoEnabled && ( + <> + {' '} + onOpenChange(false)} + > + Configure mode routes + + + )} +

- )} -
+ + {fallbackNeedsReplacement && ( +

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

+ )} +
+ ) : ( +
+
+ +

+ Every mode uses this model unless a local selection overrides it. +

+
+ + {!modelsLoading && availableModels.length === 0 && ( +

+ No models available. Configure model access first. +

+ )} +
+ )}
- + {organizationDefaultModel && ( )} - +
+ + +
diff --git a/apps/web/src/components/organizations/providers-and-models/ModelsTab.tsx b/apps/web/src/components/organizations/providers-and-models/ModelsTab.tsx index 561e78ebf6..fb76190c82 100644 --- a/apps/web/src/components/organizations/providers-and-models/ModelsTab.tsx +++ b/apps/web/src/components/organizations/providers-and-models/ModelsTab.tsx @@ -132,7 +132,7 @@ export function ModelsTab({ {!canEdit ? (
- Only organization owners can edit model access. + Only organization owners and billing managers can edit model access.
) : null} diff --git a/apps/web/src/components/organizations/providers-and-models/OrganizationProvidersAndModelsPage.tsx b/apps/web/src/components/organizations/providers-and-models/OrganizationProvidersAndModelsPage.tsx index a5639598d8..0ec489438a 100644 --- a/apps/web/src/components/organizations/providers-and-models/OrganizationProvidersAndModelsPage.tsx +++ b/apps/web/src/components/organizations/providers-and-models/OrganizationProvidersAndModelsPage.tsx @@ -90,7 +90,7 @@ export function OrganizationProvidersAndModelsPage({ organizationId, role }: Pro const { assumedRole } = useRoleTesting(); const isKiloAdmin = assumedRole === 'KILO ADMIN'; const currentRole = (isKiloAdmin ? 'owner' : assumedRole) ?? role; - const canEdit = isKiloAdmin || currentRole === 'owner'; + const canEdit = isKiloAdmin || currentRole === 'owner' || currentRole === 'billing_manager'; const updateOrganizationSettings = useUpdateOrganizationSettings(); diff --git a/apps/web/src/components/organizations/providers-and-models/ProvidersTab.tsx b/apps/web/src/components/organizations/providers-and-models/ProvidersTab.tsx index 0955573be9..a42fc431cc 100644 --- a/apps/web/src/components/organizations/providers-and-models/ProvidersTab.tsx +++ b/apps/web/src/components/organizations/providers-and-models/ProvidersTab.tsx @@ -232,7 +232,7 @@ export function ProvidersTab({ {!canEdit ? (
- Only organization owners can edit provider access. + Only organization owners and billing managers can edit provider access.
) : null} diff --git a/apps/web/src/lib/ai-gateway/auto-model/index.ts b/apps/web/src/lib/ai-gateway/auto-model/index.ts index 6596f4b3dd..da73a553d2 100644 --- a/apps/web/src/lib/ai-gateway/auto-model/index.ts +++ b/apps/web/src/lib/ai-gateway/auto-model/index.ts @@ -9,7 +9,7 @@ import type { OpenCodeSettings, Verbosity } from '@kilocode/db/schema-types'; import { QWEN37_PLUS_MODEL_ID } from '@/lib/ai-gateway/providers/qwen'; import { NVIDIA_TRIAL_TOS } from '@/lib/ai-gateway/providers/nvidia'; -type AutoModel = { +export type AutoModel = { id: string; name: string; description: string; @@ -22,6 +22,7 @@ type AutoModel = { supports_images: boolean; supports_pdf: boolean; opencode_settings: OpenCodeSettings | undefined; + status?: 'public' | 'hidden'; }; export type ResolvedAutoModel = { @@ -162,6 +163,21 @@ export const KILO_AUTO_EFFICIENT_MODEL: AutoModel = { 'Routes each request to the cheapest model that gets the job done, based on continuously benchmarked accuracy and cost.', }; +export const ORG_AUTO_MODEL: AutoModel = { + ...KILO_AUTO_BALANCED_MODEL, + id: 'kilo-auto/org', + name: 'Organization Auto', + description: "Routes requests using your organization's mode-specific model settings.", + status: 'hidden', +}; + +export const ORGANIZATION_AUTO_TARGET_MODELS = [ + KILO_AUTO_FREE_MODEL.id, + KILO_AUTO_SMALL_MODEL.id, + KILO_AUTO_BALANCED_MODEL.id, + KILO_AUTO_FRONTIER_MODEL.id, +] as const; + export const AUTO_MODELS = [ KILO_AUTO_FRONTIER_MODEL, KILO_AUTO_BALANCED_MODEL, @@ -171,5 +187,9 @@ export const AUTO_MODELS = [ ]; export function isKiloAutoModel(model: string) { - return AUTO_MODELS.some(m => m.id === model) || model === KILO_AUTO_LEGACY_MODEL; + return ( + AUTO_MODELS.some(m => m.id === model) || + model === ORG_AUTO_MODEL.id || + model === KILO_AUTO_LEGACY_MODEL + ); } diff --git a/apps/web/src/lib/ai-gateway/auto-model/resolution.test.ts b/apps/web/src/lib/ai-gateway/auto-model/resolution.test.ts index 15235c3730..16978784bb 100644 --- a/apps/web/src/lib/ai-gateway/auto-model/resolution.test.ts +++ b/apps/web/src/lib/ai-gateway/auto-model/resolution.test.ts @@ -9,7 +9,12 @@ jest.mock('@/lib/kiloclaw/setup-promo', () => ({ })); import { resolveAutoModel } from './resolution'; -import { BALANCED_QWEN_MODEL, KILO_AUTO_EFFICIENT_MODEL } from '@/lib/ai-gateway/auto-model'; +import { + BALANCED_QWEN_MODEL, + FRONTIER_MODE_TO_MODEL, + KILO_AUTO_EFFICIENT_MODEL, + ORG_AUTO_MODEL, +} from '@/lib/ai-gateway/auto-model'; import type { AutoRoutingDecision } from '@kilocode/auto-routing-contracts'; const baseParams = { @@ -141,3 +146,174 @@ describe('resolveAutoModel — kilo-auto/efficient branch', () => { expect(thunk).toHaveBeenCalledTimes(1); }); }); + +describe('resolveAutoModel — Organization Auto branch', () => { + it('uses exact built-in alias routes before canonical fallback routes', async () => { + const result = await resolveAutoModel( + { + ...baseParams, + model: ORG_AUTO_MODEL.id, + modeHeader: 'build', + apiKind: 'chat_completions', + organizationContext: Promise.resolve({ + organizationId: 'org-1', + plan: 'enterprise', + settings: { + default_model: ORG_AUTO_MODEL.id, + org_auto_model: { + routes: { + code: 'kilo-auto/frontier', + build: 'kilo-auto/small', + }, + fallback_model: 'kilo-auto/balanced', + }, + }, + }), + }, + nullUserPromise, + zeroBalancePromise + ); + + expect(result).toMatchObject({ + kind: 'ok', + routingTarget: 'kilo-auto/small', + }); + }); + + it('uses exact plan routes before architect fallback routes', async () => { + const result = await resolveAutoModel( + { + ...baseParams, + model: ORG_AUTO_MODEL.id, + modeHeader: 'plan', + apiKind: 'chat_completions', + organizationContext: Promise.resolve({ + organizationId: 'org-1', + plan: 'enterprise', + settings: { + default_model: ORG_AUTO_MODEL.id, + org_auto_model: { + routes: { + architect: 'kilo-auto/balanced', + plan: 'kilo-auto/frontier', + }, + fallback_model: 'kilo-auto/balanced', + }, + }, + }), + }, + nullUserPromise, + zeroBalancePromise + ); + + expect(result).toEqual({ + kind: 'ok', + resolved: FRONTIER_MODE_TO_MODEL.plan, + routingTarget: 'kilo-auto/frontier', + }); + }); + + it('uses the configured fallback when no mode route exists', async () => { + const result = await resolveAutoModel( + { + ...baseParams, + model: ORG_AUTO_MODEL.id, + modeHeader: 'custom-mode', + apiKind: 'chat_completions', + organizationContext: Promise.resolve({ + organizationId: 'org-1', + plan: 'enterprise', + settings: { + default_model: ORG_AUTO_MODEL.id, + org_auto_model: { + routes: {}, + fallback_model: 'kilo-auto/balanced', + }, + }, + }), + }, + nullUserPromise, + zeroBalancePromise + ); + + expect(result).toEqual({ + kind: 'ok', + resolved: BALANCED_QWEN_MODEL, + routingTarget: 'kilo-auto/balanced', + }); + }); + + it('rejects Organization Auto without an organization context', async () => { + const result = await resolveAutoModel( + { + ...baseParams, + model: ORG_AUTO_MODEL.id, + apiKind: 'chat_completions', + }, + nullUserPromise, + zeroBalancePromise + ); + + expect(result).toEqual({ + kind: 'organization_auto_configuration_error', + message: 'Organization Auto is not available for this account.', + }); + }); + + it('rejects direct Organization Auto requests after it is disabled', async () => { + const result = await resolveAutoModel( + { + ...baseParams, + model: ORG_AUTO_MODEL.id, + apiKind: 'chat_completions', + organizationContext: Promise.resolve({ + organizationId: 'org-1', + plan: 'enterprise', + settings: { + default_model: 'kilo-auto/balanced', + org_auto_model: { + routes: {}, + fallback_model: 'kilo-auto/balanced', + }, + }, + }), + }, + nullUserPromise, + zeroBalancePromise + ); + + expect(result).toEqual({ + kind: 'organization_auto_configuration_error', + message: 'Organization Auto is not enabled for this organization.', + }); + }); + + it('rejects self-referential route targets at runtime', async () => { + const result = await resolveAutoModel( + { + ...baseParams, + model: ORG_AUTO_MODEL.id, + modeHeader: 'code', + apiKind: 'chat_completions', + organizationContext: Promise.resolve({ + organizationId: 'org-1', + plan: 'enterprise', + settings: { + default_model: ORG_AUTO_MODEL.id, + org_auto_model: { + routes: { code: ORG_AUTO_MODEL.id }, + fallback_model: 'kilo-auto/balanced', + }, + }, + }), + }, + nullUserPromise, + zeroBalancePromise + ); + + expect(result).toEqual({ + kind: 'organization_auto_configuration_error', + message: 'Organization Auto cannot target itself.', + }); + }); +}); diff --git a/apps/web/src/lib/ai-gateway/auto-model/resolution.ts b/apps/web/src/lib/ai-gateway/auto-model/resolution.ts index b45bf1617a..71358e9665 100644 --- a/apps/web/src/lib/ai-gateway/auto-model/resolution.ts +++ b/apps/web/src/lib/ai-gateway/auto-model/resolution.ts @@ -9,6 +9,10 @@ import type { } from '@/lib/ai-gateway/providers/openrouter/types'; import type OpenAI from 'openai'; import type { User } from '@kilocode/db'; +import type { + OrganizationPlan, + OrganizationSettings, +} from '@/lib/organizations/organization-types'; import type { AutoRoutingDecision } from '@kilocode/auto-routing-contracts'; import { KILO_AUTO_FREE_MODEL, @@ -22,6 +26,8 @@ import { FRONTIER_CODE_MODEL, type ResolvedAutoModel, KILO_AUTO_LEGACY_MODEL, + ORG_AUTO_MODEL, + ORGANIZATION_AUTO_TARGET_MODELS, } from '@/lib/ai-gateway/auto-model'; import { userIsWithinFirstKiloClawInstanceWindow } from '@/lib/kiloclaw/setup-promo'; import { getRandomNumber } from '@/lib/ai-gateway/getRandomNumber'; @@ -33,6 +39,7 @@ import { import { getOpenRouterModels } from '@/lib/ai-gateway/providers/gateway-models-cache'; import PROVIDERS from '@/lib/ai-gateway/providers/provider-definitions'; import type { ProviderId } from '@/lib/ai-gateway/providers/types'; +import { validateOrganizationAutoTarget } from '@/lib/organizations/organization-auto-model'; type ResolveAutoModelParams = { model: string; @@ -44,6 +51,11 @@ type ResolveAutoModelParams = { // Lazily fetches the auto-routing worker's decision; only set for // kilo-auto/efficient requests (route.ts owns the request-body capture). efficientDecision?: () => Promise; + organizationContext?: Promise<{ + organizationId?: string; + settings?: OrganizationSettings; + plan?: OrganizationPlan; + }>; }; function resolveMode(modeHeader: string | null, featureHeader: FeatureValue | null) { @@ -88,9 +100,142 @@ function gatewaySupportsApiKind( return provider?.supportedChatApis.some(k => k === apiKind) ?? false; } +type OrganizationAutoContext = { + organizationId?: string; + settings?: OrganizationSettings; + plan?: OrganizationPlan; +}; + +function hasOrganizationAutoRoute(routes: Record, slug: string): boolean { + return Object.prototype.hasOwnProperty.call(routes, slug); +} + +function resolveOrganizationAutoRouteTarget( + settings: OrganizationSettings, + modeHeader: string | null +): string | undefined { + const routes = settings.org_auto_model?.routes ?? {}; + const mode = modeHeader?.trim() ?? ''; + const normalizedMode = mode.toLowerCase(); + + if (normalizedMode && hasOrganizationAutoRoute(routes, normalizedMode)) { + return routes[normalizedMode]; + } + + if (normalizedMode === 'code' || normalizedMode === 'build') { + if (hasOrganizationAutoRoute(routes, 'code')) return routes.code; + if (hasOrganizationAutoRoute(routes, 'build')) return routes.build; + return settings.org_auto_model?.fallback_model; + } + + if (normalizedMode === 'plan' || normalizedMode === 'architect') { + if (hasOrganizationAutoRoute(routes, 'architect')) return routes.architect; + if (hasOrganizationAutoRoute(routes, 'plan')) return routes.plan; + return settings.org_auto_model?.fallback_model; + } + + return settings.org_auto_model?.fallback_model; +} + +function isSupportedNestedAutoTarget(modelId: string): boolean { + return (ORGANIZATION_AUTO_TARGET_MODELS as readonly string[]).includes(modelId); +} + export type ResolveAutoModelResult = - | { kind: 'ok'; resolved: ResolvedAutoModel } - | { kind: 'no_free_models_available' }; + | { kind: 'ok'; resolved: ResolvedAutoModel; routingTarget?: string } + | { kind: 'no_free_models_available' } + | { kind: 'organization_auto_configuration_error'; message: string }; + +async function resolveOrganizationAutoModel( + params: ResolveAutoModelParams, + userPromise: Promise, + balancePromise: Promise +): Promise { + const organizationContext: OrganizationAutoContext = await (params.organizationContext ?? + Promise.resolve({})); + + if (!organizationContext.organizationId || organizationContext.plan !== 'enterprise') { + return { + kind: 'organization_auto_configuration_error', + message: 'Organization Auto is not available for this account.', + }; + } + + if (!organizationContext.settings?.org_auto_model) { + return { + kind: 'organization_auto_configuration_error', + message: 'Organization Auto is not configured for this organization.', + }; + } + + if (organizationContext.settings.default_model !== ORG_AUTO_MODEL.id) { + return { + kind: 'organization_auto_configuration_error', + message: 'Organization Auto is not enabled for this organization.', + }; + } + + const targetModelId = resolveOrganizationAutoRouteTarget( + organizationContext.settings, + params.modeHeader + ); + if (!targetModelId) { + return { + kind: 'organization_auto_configuration_error', + message: 'Organization Auto has no configured fallback model.', + }; + } + + let validation: Awaited>; + try { + validation = await validateOrganizationAutoTarget( + { + id: organizationContext.organizationId, + plan: organizationContext.plan, + settings: organizationContext.settings, + }, + targetModelId, + { apiKind: params.apiKind ?? undefined } + ); + } catch { + return { + kind: 'organization_auto_configuration_error', + message: + 'Organization Auto could not validate this route target against the current model catalog.', + }; + } + if (validation.kind === 'error') { + return { kind: 'organization_auto_configuration_error', message: validation.message }; + } + + if (validation.modelId === ORG_AUTO_MODEL.id) { + return { + kind: 'organization_auto_configuration_error', + message: 'Organization Auto cannot target itself.', + }; + } + + if (isSupportedNestedAutoTarget(validation.modelId)) { + const nestedResult = await resolveAutoModel( + { + ...params, + model: validation.modelId, + }, + userPromise, + balancePromise + ); + if (nestedResult.kind === 'ok') { + return { ...nestedResult, routingTarget: validation.modelId }; + } + return nestedResult; + } + + return { + kind: 'ok', + resolved: { model: validation.modelId }, + routingTarget: validation.modelId, + }; +} export async function resolveAutoModel( params: ResolveAutoModelParams, @@ -98,6 +243,9 @@ export async function resolveAutoModel( balancePromise: Promise ): Promise { const { model, modeHeader, featureHeader, sessionId, apiKind, clientIp } = params; + if (model === ORG_AUTO_MODEL.id) { + return await resolveOrganizationAutoModel(params, userPromise, balancePromise); + } if (model === KILO_AUTO_FREE_MODEL.id) { const candidates = await getAutoFreeCandidates(apiKind); if (candidates.length === 0) { diff --git a/apps/web/src/lib/ai-gateway/byok/index.ts b/apps/web/src/lib/ai-gateway/byok/index.ts index c2d3a15cc0..2add2646fb 100644 --- a/apps/web/src/lib/ai-gateway/byok/index.ts +++ b/apps/web/src/lib/ai-gateway/byok/index.ts @@ -1,4 +1,4 @@ -import { type db } from '@/lib/drizzle'; +import { type db, type DrizzleTransaction } from '@/lib/drizzle'; import { byok_api_keys } from '@kilocode/db/schema'; import { and, eq, inArray } from 'drizzle-orm'; import type { EncryptedData } from '@/lib/ai-gateway/byok/encryption'; @@ -75,7 +75,7 @@ export async function getBYOKforUser( } export async function getBYOKforOrganization( - fromDb: typeof db, + fromDb: typeof db | DrizzleTransaction, organizationId: string, providerIds: UserByokProviderId[] ): Promise { diff --git a/apps/web/src/lib/ai-gateway/llm-proxy-helpers.ts b/apps/web/src/lib/ai-gateway/llm-proxy-helpers.ts index 14ae12088d..ff2690d658 100644 --- a/apps/web/src/lib/ai-gateway/llm-proxy-helpers.ts +++ b/apps/web/src/lib/ai-gateway/llm-proxy-helpers.ts @@ -264,6 +264,17 @@ export function noFreeModelsAvailableResponse() { ); } +export function organizationAutoConfigurationResponse(message: string) { + return NextResponse.json( + { + error: 'Organization Auto configuration error', + error_type: ProxyErrorType.organization_auto_configuration, + message, + }, + { status: 400 } + ); +} + export function featureExclusiveModelResponse(modelId: string) { const exclusiveTo = findKiloExclusiveModel(modelId)?.exclusive_to ?? []; const error = `${modelId} is only available for ${exclusiveTo.join(', ')}. Use ${KILO_AUTO_FREE_MODEL.id} as a free alternative.`; diff --git a/apps/web/src/lib/ai-gateway/providers/openrouter/index.ts b/apps/web/src/lib/ai-gateway/providers/openrouter/index.ts index 07da1860a7..9ca8f481ad 100644 --- a/apps/web/src/lib/ai-gateway/providers/openrouter/index.ts +++ b/apps/web/src/lib/ai-gateway/providers/openrouter/index.ts @@ -19,7 +19,7 @@ import { } from '@/lib/ai-gateway/providers/kilo-exclusive-model'; import { isForbiddenFreeModel } from '@/lib/ai-gateway/forbidden-free-models'; import { getOpenCodeSettings } from '@/lib/ai-gateway/providers/model-settings'; -import { AUTO_MODELS } from '@/lib/ai-gateway/auto-model'; +import { AUTO_MODELS, type AutoModel } from '@/lib/ai-gateway/auto-model'; import { ATTRIBUTION_HEADERS } from '@/lib/ai-gateway/providers/openrouter/attribution-headers'; import { getOpenRouterModelsMetadata } from '@/lib/ai-gateway/providers/gateway-models-cache'; import { getPreferredProviderOrder } from '@/lib/ai-gateway/providers/apply-provider-specific-logic'; @@ -33,51 +33,47 @@ import { addMonths } from 'date-fns'; // Re-export from shared module for backwards compatibility export { normalizeModelId } from '@/lib/ai-gateway/model-utils'; -function buildAutoModels(): OpenRouterModel[] { - return AUTO_MODELS.map(m => { - const input_modalities = ['text']; - if (m.supports_images) { - input_modalities.push('image'); - } - if (m.supports_pdf) { - input_modalities.push('pdf'); - } - return { - id: m.id, - name: m.name, - created: 0, - description: m.description, - architecture: { - input_modalities: input_modalities, - output_modalities: ['text'], - tokenizer: 'Other', - }, - top_provider: { - is_moderated: false, - context_length: m.context_length, - max_completion_tokens: m.max_completion_tokens, - }, - pricing: { - prompt: m.prompt_price, - completion: m.completion_price, - input_cache_read: m.input_cache_read_price, - input_cache_write: m.input_cache_write_price, - request: '0', - image: '0', - web_search: '0', - internal_reasoning: '0', - }, +export function buildAutoModelCatalogEntry(m: AutoModel): OpenRouterModel { + const input_modalities = ['text']; + if (m.supports_images) { + input_modalities.push('image'); + } + if (m.supports_pdf) { + input_modalities.push('pdf'); + } + return { + id: m.id, + name: m.name, + created: 0, + description: m.description, + architecture: { + input_modalities: input_modalities, + output_modalities: ['text'], + tokenizer: 'Other', + }, + top_provider: { + is_moderated: false, context_length: m.context_length, - supported_parameters: [ - 'max_tokens', - 'temperature', - 'tools', - 'reasoning', - 'include_reasoning', - ], - opencode: m.opencode_settings, - }; - }); + max_completion_tokens: m.max_completion_tokens, + }, + pricing: { + prompt: m.prompt_price, + completion: m.completion_price, + input_cache_read: m.input_cache_read_price, + input_cache_write: m.input_cache_write_price, + request: '0', + image: '0', + web_search: '0', + internal_reasoning: '0', + }, + context_length: m.context_length, + supported_parameters: ['max_tokens', 'temperature', 'tools', 'reasoning', 'include_reasoning'], + opencode: m.opencode_settings, + }; +} + +function buildAutoModels(): OpenRouterModel[] { + return AUTO_MODELS.map(buildAutoModelCatalogEntry); } export function formatName(model: OpenRouterModel, preferredIndex: number) { diff --git a/apps/web/src/lib/organizations/organization-auto-model-shared.ts b/apps/web/src/lib/organizations/organization-auto-model-shared.ts new file mode 100644 index 0000000000..2db7cb1d07 --- /dev/null +++ b/apps/web/src/lib/organizations/organization-auto-model-shared.ts @@ -0,0 +1,22 @@ +import type { OrganizationSettings } from '@/lib/organizations/organization-types'; +import { ORGANIZATION_AUTO_TARGET_MODELS } from '@/lib/ai-gateway/auto-model'; + +export const ORGANIZATION_AUTO_MODEL_FLAG = 'organization-auto-model-routing'; +export const MAX_ORGANIZATION_AUTO_ROUTES = 100; + +export function getOrganizationAutoRoute( + settings: OrganizationSettings | undefined, + slug: string +): string | undefined { + if (!settings?.org_auto_model) { + return undefined; + } + if (!Object.prototype.hasOwnProperty.call(settings.org_auto_model.routes, slug)) { + return undefined; + } + return settings.org_auto_model.routes[slug]; +} + +export function isOrganizationAutoTargetModel(modelId: string): boolean { + return (ORGANIZATION_AUTO_TARGET_MODELS as readonly string[]).includes(modelId); +} diff --git a/apps/web/src/lib/organizations/organization-auto-model.ts b/apps/web/src/lib/organizations/organization-auto-model.ts new file mode 100644 index 0000000000..7190873bf8 --- /dev/null +++ b/apps/web/src/lib/organizations/organization-auto-model.ts @@ -0,0 +1,155 @@ +import type { Organization } from '@kilocode/db/schema'; +import type { OrganizationAutoModelSettings } from '@/lib/organizations/organization-types'; +import { getEnhancedOpenRouterModels } from '@/lib/ai-gateway/providers/openrouter'; +import { createAllowPredicateFromRestrictions } from '@/lib/model-allow.server'; +import { CUSTOM_LLM_PREFIX, normalizeModelId } from '@/lib/ai-gateway/model-utils'; +import { + formatDirectByokModelId, + getDirectByokModel, +} from '@/lib/ai-gateway/providers/direct-byok'; +import { getBYOKforOrganization } from '@/lib/ai-gateway/byok'; +import { db, type DrizzleTransaction } from '@/lib/drizzle'; +import { KILO_AUTO_BALANCED_MODEL, ORG_AUTO_MODEL } from '@/lib/ai-gateway/auto-model'; +export { + getOrganizationAutoRoute, + isOrganizationAutoTargetModel, + MAX_ORGANIZATION_AUTO_ROUTES, + ORGANIZATION_AUTO_MODEL_FLAG, +} from '@/lib/organizations/organization-auto-model-shared'; +import { isOrganizationAutoTargetModel } from '@/lib/organizations/organization-auto-model-shared'; + +type OrganizationAutoPolicyOrganization = Pick; + +export const DEFAULT_ORGANIZATION_AUTO_MODEL_SETTINGS: OrganizationAutoModelSettings = { + routes: {}, + fallback_model: KILO_AUTO_BALANCED_MODEL.id, +}; + +export function isOrganizationAutoEligible(organization: Organization): boolean { + return organization.plan === 'enterprise'; +} + +export function isOrganizationAutoConfigured(organization: Organization): boolean { + return isOrganizationAutoEligible(organization) && !!organization.settings.org_auto_model; +} + +export function isOrganizationAutoEnabled(organization: Organization): boolean { + return ( + isOrganizationAutoConfigured(organization) && + organization.settings.default_model === ORG_AUTO_MODEL.id + ); +} + +export function getOrganizationAutoSettings( + organization: Organization +): OrganizationAutoModelSettings | undefined { + return organization.settings.org_auto_model; +} + +export type OrganizationAutoTargetValidationResult = + | { kind: 'ok'; modelId: string } + | { kind: 'error'; message: string }; + +export async function validateOrganizationAutoTarget( + organization: OrganizationAutoPolicyOrganization, + targetModelId: string, + options: { + apiKind?: 'chat_completions' | 'responses' | 'messages'; + dbClient?: typeof db | DrizzleTransaction; + } = {} +): Promise { + const rawModelId = targetModelId.trim().toLowerCase(); + const normalizedModelId = normalizeModelId(rawModelId); + + if (!rawModelId) { + return { kind: 'error', message: 'Organization Auto route target is required.' }; + } + + if (normalizedModelId.endsWith('/*')) { + return { + kind: 'error', + message: `Organization Auto route target '${targetModelId}' is not a concrete model identifier.`, + }; + } + + if (normalizedModelId === ORG_AUTO_MODEL.id) { + return { kind: 'error', message: 'Organization Auto cannot target itself.' }; + } + + if (rawModelId.startsWith(CUSTOM_LLM_PREFIX)) { + return { + kind: 'error', + message: `Organization Auto route target '${targetModelId}' must be a Kilo-hosted model, supported auto tier, or organization-owned BYOK model.`, + }; + } + + const directByokTarget = await getDirectByokModel(rawModelId); + if (directByokTarget.provider && directByokTarget.model) { + const byok = await getBYOKforOrganization(options.dbClient ?? db, organization.id, [ + directByokTarget.provider.id, + ]); + if (!byok || byok.length === 0) { + return { + kind: 'error', + message: `Organization Auto route target '${targetModelId}' is unavailable because this organization does not have an enabled BYOK credential for ${directByokTarget.provider.id}.`, + }; + } + if ( + options.apiKind && + !directByokTarget.provider.supported_chat_apis.includes(options.apiKind) + ) { + return { + kind: 'error', + message: `Organization Auto route target '${targetModelId}' does not support the ${options.apiKind} API.`, + }; + } + return { + kind: 'ok', + modelId: formatDirectByokModelId(directByokTarget.provider, directByokTarget.model), + }; + } + + if (normalizedModelId.startsWith('kilo-auto/')) { + if (!isOrganizationAutoTargetModel(normalizedModelId)) { + return { + kind: 'error', + message: `Organization Auto route target '${targetModelId}' is not a supported auto tier.`, + }; + } + + return { kind: 'ok', modelId: normalizedModelId }; + } + + let models; + try { + models = await getEnhancedOpenRouterModels(); + } catch { + return { + kind: 'error', + message: + 'Organization Auto could not validate this route target against the current model catalog.', + }; + } + const catalogModel = models.data.find(model => model.id.toLowerCase() === rawModelId); + if (!catalogModel) { + return { + kind: 'error', + message: `Organization Auto route target '${targetModelId}' is unavailable.`, + }; + } + + const isAllowed = createAllowPredicateFromRestrictions({ + providerAllowList: + organization.plan === 'enterprise' ? organization.settings.provider_allow_list : undefined, + modelDenyList: + organization.plan === 'enterprise' ? (organization.settings.model_deny_list ?? []) : [], + }); + if (!(await isAllowed(normalizedModelId))) { + return { + kind: 'error', + message: `Organization Auto route target '${targetModelId}' is not allowed by the organization's model policy.`, + }; + } + + return { kind: 'ok', modelId: catalogModel.id }; +} diff --git a/apps/web/src/lib/organizations/organization-base-types.ts b/apps/web/src/lib/organizations/organization-base-types.ts index 9f0cfe57d7..b59594cbcb 100644 --- a/apps/web/src/lib/organizations/organization-base-types.ts +++ b/apps/web/src/lib/organizations/organization-base-types.ts @@ -3,11 +3,13 @@ export { OrganizationPlanSchema, OrganizationModeConfigSchema, OrganizationSettingsSchema, + OrganizationAutoModelSettingsSchema, } from '@kilocode/db/schema-types'; export type { OrganizationRole, OrganizationPlan, OrganizationSettings, + OrganizationAutoModelSettings, OrganizationModeConfig, EditGroupConfig, } from '@kilocode/db/schema-types'; diff --git a/apps/web/src/lib/organizations/organization-models.ts b/apps/web/src/lib/organizations/organization-models.ts index fd604b91fd..3c47559ebf 100644 --- a/apps/web/src/lib/organizations/organization-models.ts +++ b/apps/web/src/lib/organizations/organization-models.ts @@ -1,5 +1,8 @@ import type { OpenRouterModelsResponse } from '@/lib/organizations/organization-types'; -import { getEnhancedOpenRouterModels } from '@/lib/ai-gateway/providers/openrouter'; +import { + buildAutoModelCatalogEntry, + getEnhancedOpenRouterModels, +} from '@/lib/ai-gateway/providers/openrouter'; import { createAllowPredicateFromRestrictions, hasActiveModelRestrictions, @@ -10,6 +13,8 @@ import { getDirectByokModelsForOrganization } from '@/lib/ai-gateway/providers/d import { getOrganizationById } from '@/lib/organizations/organizations'; import { getEffectiveModelRestrictions } from '@/lib/organizations/model-restrictions'; import { listAvailableExperimentModels } from '@/lib/ai-gateway/experiments/list-available-experiment-models'; +import { ORG_AUTO_MODEL } from '@/lib/ai-gateway/auto-model'; +import { isOrganizationAutoEnabled } from '@/lib/organizations/organization-auto-model'; export async function getAvailableModelsForOrganization( organizationId: string @@ -44,6 +49,10 @@ export async function getAvailableModelsForOrganization( filteredModels.push(...(await listAvailableExperimentModels())); } + if (isOrganizationAutoEnabled(organization)) { + filteredModels.push(buildAutoModelCatalogEntry(ORG_AUTO_MODEL)); + } + filteredModels.push(...(await getDirectByokModelsForOrganization(organizationId))); filteredModels.push(...(await listAvailableCustomLlms(organizationId))); diff --git a/apps/web/src/lib/organizations/organization-modes.test.ts b/apps/web/src/lib/organizations/organization-modes.test.ts index 057e0e879d..7887a3adac 100644 --- a/apps/web/src/lib/organizations/organization-modes.test.ts +++ b/apps/web/src/lib/organizations/organization-modes.test.ts @@ -139,19 +139,6 @@ describe('createOrganizationMode', () => { expect(mode).not.toBeNull(); expect(mode?.config).toEqual(config); }); - - test('should preserve an optional default model', async () => { - const user = await insertTestUser(); - const organization = await createOrganization('Test Org', user.id); - - const mode = await createOrganizationMode(organization.id, user.id, 'Code Mode', 'code', { - roleDefinition: 'You are a coding assistant', - groups: ['read', 'edit'], - defaultModel: 'openai/gpt-4o', - }); - - expect(mode?.config.defaultModel).toBe('openai/gpt-4o'); - }); }); describe('updateOrganizationMode', () => { @@ -160,48 +147,6 @@ describe('updateOrganizationMode', () => { await db.delete(organizations); }); - test('should preserve existing config when updating only defaultModel', async () => { - const user = await insertTestUser(); - const organization = await createOrganization('Test Org', user.id); - const mode = await createOrganizationMode(organization.id, user.id, 'Code Mode', 'code', { - roleDefinition: 'You are a coding assistant', - description: 'Write code', - groups: ['read', 'edit'], - }); - - const updatedMode = await updateOrganizationMode(organization.id, mode!.id, { - config: { defaultModel: 'openai/gpt-4o' }, - }); - - expect(updatedMode?.config).toEqual({ - roleDefinition: 'You are a coding assistant', - description: 'Write code', - groups: ['read', 'edit'], - defaultModel: 'openai/gpt-4o', - }); - }); - - test('should clear defaultModel without dropping the rest of config', async () => { - const user = await insertTestUser(); - const organization = await createOrganization('Test Org', user.id); - const mode = await createOrganizationMode(organization.id, user.id, 'Code Mode', 'code', { - roleDefinition: 'You are a coding assistant', - description: 'Write code', - groups: ['read', 'edit'], - defaultModel: 'openai/gpt-4o', - }); - - const updatedMode = await updateOrganizationMode(organization.id, mode!.id, { - config: { defaultModel: undefined }, - }); - - expect(updatedMode?.config).toEqual({ - roleDefinition: 'You are a coding assistant', - description: 'Write code', - groups: ['read', 'edit'], - }); - }); - test('should not lose concurrent partial config updates', async () => { const user = await insertTestUser(); const organization = await createOrganization('Test Org', user.id); @@ -215,7 +160,7 @@ describe('updateOrganizationMode', () => { config: { description: 'Write code' }, }), updateOrganizationMode(organization.id, mode!.id, { - config: { defaultModel: 'openai/gpt-4o' }, + config: { customInstructions: 'Focus on security' }, }), ]); @@ -225,7 +170,7 @@ describe('updateOrganizationMode', () => { roleDefinition: 'You are a coding assistant', groups: ['read'], description: 'Write code', - defaultModel: 'openai/gpt-4o', + customInstructions: 'Focus on security', }); }); @@ -239,7 +184,7 @@ describe('updateOrganizationMode', () => { }); const updatedMode = await updateOrganizationMode(otherOrganization.id, mode!.id, { - config: { defaultModel: 'openai/gpt-4o' }, + config: { description: 'Write code' }, }); expect(updatedMode).toBeNull(); diff --git a/apps/web/src/lib/organizations/organization-modes.ts b/apps/web/src/lib/organizations/organization-modes.ts index f8bc687ff5..c78c368d20 100644 --- a/apps/web/src/lib/organizations/organization-modes.ts +++ b/apps/web/src/lib/organizations/organization-modes.ts @@ -1,4 +1,4 @@ -import { db } from '@/lib/drizzle'; +import { db, type DrizzleTransaction } from '@/lib/drizzle'; import { orgnaization_modes, ORGANIZATION_MODES_ORG_SLUG_CONSTRAINT } from '@kilocode/db/schema'; import { and, eq, sql } from 'drizzle-orm'; import type { OrganizationModeConfig } from '@/lib/organizations/organization-types'; @@ -10,10 +10,19 @@ const defaultConfig: OrganizationModeConfig = { roleDefinition: 'default', }; -function mergeToSatisfy(config: Partial): OrganizationModeConfig { +type OrganizationModeConfigWithLegacyDefaultModel = Partial & { + defaultModel?: string; +}; + +function mergeToSatisfy( + config: OrganizationModeConfigWithLegacyDefaultModel +): OrganizationModeConfig { + const configWithoutLegacyDefaultModel = { ...config }; + delete configWithoutLegacyDefaultModel.defaultModel; + return { ...defaultConfig, - ...config, + ...configWithoutLegacyDefaultModel, }; } @@ -22,9 +31,10 @@ export async function createOrganizationMode( createdBy: string, name: string, slug: string, - config: Partial = {} + config: Partial = {}, + txn?: DrizzleTransaction ): Promise { - const [mode] = await db + const [mode] = await (txn ?? db) .insert(orgnaization_modes) .values({ organization_id: organizationId, @@ -39,8 +49,11 @@ export async function createOrganizationMode( return mode || null; } -export async function getAllOrganizationModes(organizationId: string): Promise { - const modes = await db +export async function getAllOrganizationModes( + organizationId: string, + txn?: DrizzleTransaction +): Promise { + const modes = await (txn ?? db) .select() .from(orgnaization_modes) .where(eq(orgnaization_modes.organization_id, organizationId)); @@ -50,14 +63,17 @@ export async function getAllOrganizationModes(organizationId: string): Promise { - const [mode] = await db + const query = (txn ?? db) .select() .from(orgnaization_modes) .where( and(eq(orgnaization_modes.id, modeId), eq(orgnaization_modes.organization_id, organizationId)) ); + const [mode] = await (lockForUpdate ? query.for('update') : query); return mode ? { ...mode, config: mergeToSatisfy(mode.config) } : null; } @@ -69,7 +85,8 @@ export async function updateOrganizationMode( name?: string; slug?: string; config?: Partial; - } + }, + txn?: DrizzleTransaction ): Promise { const updateData: Record = {}; @@ -81,17 +98,16 @@ export async function updateOrganizationMode( } if (updates.config !== undefined) { const configPatch = Object.fromEntries( - Object.entries(updates.config).map(([key, value]) => [ - key, - value === undefined ? null : value, - ]) + Object.entries(updates.config) + .filter(([key]) => key !== 'defaultModel') + .map(([key, value]) => [key, value === undefined ? null : value]) ); - updateData.config = sql`jsonb_strip_nulls(((${JSON.stringify(defaultConfig)}::jsonb || COALESCE(${orgnaization_modes.config}, '{}'::jsonb)) || ${JSON.stringify(configPatch)}::jsonb))`; + updateData.config = sql`jsonb_strip_nulls(((${JSON.stringify(defaultConfig)}::jsonb || (COALESCE(${orgnaization_modes.config}, '{}'::jsonb) - 'defaultModel')) || ${JSON.stringify(configPatch)}::jsonb))`; } try { - const [mode] = await db + const [mode] = await (txn ?? db) .update(orgnaization_modes) .set(updateData) .where( @@ -112,6 +128,9 @@ export async function updateOrganizationMode( } } -export async function deleteOrganizationMode(modeId: string): Promise { - await db.delete(orgnaization_modes).where(eq(orgnaization_modes.id, modeId)); +export async function deleteOrganizationMode( + modeId: string, + txn?: DrizzleTransaction +): Promise { + await (txn ?? db).delete(orgnaization_modes).where(eq(orgnaization_modes.id, modeId)); } diff --git a/apps/web/src/lib/organizations/organization-seats.ts b/apps/web/src/lib/organizations/organization-seats.ts index fec9045b46..437c610310 100644 --- a/apps/web/src/lib/organizations/organization-seats.ts +++ b/apps/web/src/lib/organizations/organization-seats.ts @@ -311,7 +311,6 @@ async function handleSubscriptionEventInternal( return; } - // Update organization plan from subscription metadata for ALL events (Org Plans 5) const plan = getPlanTypeFromSubscription(subscription); if (plan !== null) { await tx.update(organizations).set({ plan }).where(eq(organizations.id, meta.organizationId)); diff --git a/apps/web/src/lib/organizations/organization-types.ts b/apps/web/src/lib/organizations/organization-types.ts index 1888fe9c73..0a69106f7f 100644 --- a/apps/web/src/lib/organizations/organization-types.ts +++ b/apps/web/src/lib/organizations/organization-types.ts @@ -8,6 +8,7 @@ export type { OrganizationRole, OrganizationPlan, OrganizationSettings, + OrganizationAutoModelSettings, OrganizationModeConfig, EditGroupConfig, } from './organization-base-types'; @@ -15,6 +16,7 @@ export { OrganizationPlanSchema, OrganizationModeConfigSchema, OrganizationSettingsSchema, + OrganizationAutoModelSettingsSchema, } from './organization-base-types'; import type { OrganizationRole, OrganizationPlan } from './organization-base-types'; diff --git a/apps/web/src/lib/organizations/organizations.ts b/apps/web/src/lib/organizations/organizations.ts index f2041355fb..8dbfdbbfd2 100644 --- a/apps/web/src/lib/organizations/organizations.ts +++ b/apps/web/src/lib/organizations/organizations.ts @@ -6,6 +6,7 @@ import { type AcceptInviteResult, type OrganizationSettings, OrganizationSSODomainSchema, + OrganizationSettingsSchema, } from '@/lib/organizations/organization-types'; import { kilocode_users, @@ -706,6 +707,33 @@ export function getAcceptInviteUrl(inviteToken: OrganizationInvitation['token']) return acceptInviteUrl; } +export async function mutateOrganizationSettings( + organizationId: Organization['id'], + mutate: (organization: Organization) => Promise | OrganizationSettings, + txn?: DrizzleTransaction +): Promise { + const run = async (tx: DrizzleTransaction): Promise => { + const [organization] = await tx + .select() + .from(organizations) + .where(eq(organizations.id, organizationId)) + .for('update'); + if (!organization) { + throw new Error(`Organization ${organizationId} not found`); + } + const nextSettings = await mutate(organization); + // Returning the locked settings object is the explicit no-op signal. + if (nextSettings === organization.settings) { + return organization.settings; + } + const settings = OrganizationSettingsSchema.parse(nextSettings); + await tx.update(organizations).set({ settings }).where(eq(organizations.id, organizationId)); + return settings; + }; + + return txn ? run(txn) : db.transaction(run); +} + export async function updateOrganizationSettings( organizationId: Organization['id'], settings: OrganizationSettings, diff --git a/apps/web/src/lib/proxy-error-types.ts b/apps/web/src/lib/proxy-error-types.ts index 8290e5e663..fe0db0cc48 100644 --- a/apps/web/src/lib/proxy-error-types.ts +++ b/apps/web/src/lib/proxy-error-types.ts @@ -28,6 +28,7 @@ export const proxyErrorTypeSchema = z.enum([ 'upstream_error', 'no_free_models_available', 'abuse_blocked', + 'organization_auto_configuration', ]); export type ProxyErrorType = z.infer; diff --git a/apps/web/src/routers/organizations/organization-modes-router.test.ts b/apps/web/src/routers/organizations/organization-modes-router.test.ts index 748132366e..536e6baea3 100644 --- a/apps/web/src/routers/organizations/organization-modes-router.test.ts +++ b/apps/web/src/routers/organizations/organization-modes-router.test.ts @@ -1,12 +1,12 @@ import { createCallerForUser } from '@/routers/test-utils'; import { insertTestUser } from '@/tests/helpers/user.helper'; import { createTestOrganization } from '@/tests/helpers/organization.helper'; -import { addUserToOrganization } from '@/lib/organizations/organizations'; +import { addUserToOrganization, getOrganizationById } from '@/lib/organizations/organizations'; import { getAllOrganizationModes } from '@/lib/organizations/organization-modes'; -import { db } from '@/lib/drizzle'; -import { organizations } from '@kilocode/db/schema'; import type { User, Organization } from '@kilocode/db/schema'; -import { eq } from 'drizzle-orm'; +import { organization_audit_logs } from '@kilocode/db/schema'; +import { db } from '@/lib/drizzle'; +import { desc, eq } from 'drizzle-orm'; import { randomUUID } from 'crypto'; jest.mock('@/lib/posthog-feature-flags', () => ({ @@ -97,224 +97,162 @@ describe('organization modes tRPC router', () => { expect(result.mode.created_by).toBe(member.id); }); - it('should throw error for duplicate slug', async () => { + it('does not clear a canonical route when create omits route_model', async () => { const caller = await createCallerForUser(owner.id); + const freshOrg = await createTestOrganization('Create Route Org', owner.id, 0, {}, false); + await caller.organizations.settings.setOrganizationAutoRoute({ + organizationId: freshOrg.id, + mode_slug: 'create-route-mode', + model_id: 'kilo-auto/frontier', + }); - // Create first mode - await caller.organizations.modes.create({ - organizationId: testOrganization.id, - name: 'First Mode', - slug: 'duplicate-slug', + const result = await caller.organizations.modes.create({ + organizationId: freshOrg.id, + name: 'Create Route Mode', + slug: 'create-route-mode', }); - // Try to create second mode with same slug - await expect( - caller.organizations.modes.create({ - organizationId: testOrganization.id, - name: 'Second Mode', - slug: 'duplicate-slug', - }) - ).rejects.toThrow(); + expect(result.mode.config).not.toHaveProperty('defaultModel'); }); - it('should validate slug format', async () => { + it('records route details in mode create audit messages', async () => { const caller = await createCallerForUser(owner.id); + const result = await caller.organizations.modes.create({ + organizationId: testOrganization.id, + name: 'Routed Mode', + slug: 'routed-mode', + route_model: 'kilo-auto/frontier', + }); - await expect( - caller.organizations.modes.create({ - organizationId: testOrganization.id, - name: 'Invalid Slug Mode', - slug: 'Invalid Slug!', - }) - ).rejects.toThrow(); - }); - - it('should throw error for non-existent organization', async () => { - const caller = await createCallerForUser(owner.id); - const nonExistentId = randomUUID(); + const [audit] = await db + .select() + .from(organization_audit_logs) + .where(eq(organization_audit_logs.organization_id, testOrganization.id)) + .orderBy(desc(organization_audit_logs.created_at)); - await expect( - caller.organizations.modes.create({ - organizationId: nonExistentId, - name: 'Test Mode', - slug: 'test', - }) - ).rejects.toThrow(); + expect(result.mode.slug).toBe('routed-mode'); + expect(audit?.message).toContain('Organization Auto route set'); + expect(audit?.message).toContain('routed-mode'); }); - it('should allow an organization mode default that is not denied', async () => { + it('stores route_model in canonical Organization Auto routes', async () => { const caller = await createCallerForUser(owner.id); - const organization = await createTestOrganization( - 'Allowed Default Model Org', + const freshOrg = await createTestOrganization( + 'Create With Route Org', owner.id, 0, - { model_deny_list: ['anthropic/claude-3-opus'] }, + {}, false ); const result = await caller.organizations.modes.create({ - organizationId: organization.id, - name: 'Code Mode', - slug: 'code', - config: { - roleDefinition: 'You are a coding assistant', - groups: ['read'], - defaultModel: 'openai/gpt-4o', - }, + organizationId: freshOrg.id, + name: 'Create With Route Mode', + slug: 'create-with-route-mode', + route_model: 'kilo-auto/balanced', }); - expect(result.mode.config.defaultModel).toBe('openai/gpt-4o'); + expect(result.mode.config).not.toHaveProperty('defaultModel'); + const updatedOrganization = await getOrganizationById(freshOrg.id); + expect(updatedOrganization?.settings.org_auto_model?.routes['create-with-route-mode']).toBe( + 'kilo-auto/balanced' + ); }); - it('should reject an organization mode default for a non-enterprise organization', async () => { + it('records a clear message when create explicitly clears an existing route', async () => { const caller = await createCallerForUser(owner.id); - const organization = await createTestOrganization( - 'Teams Default Model Org', + const freshOrg = await createTestOrganization( + 'Create Clear Route Org', owner.id, 0, {}, - true + false ); + await caller.organizations.settings.setOrganizationAutoRoute({ + organizationId: freshOrg.id, + mode_slug: 'create-clear-route-mode', + model_id: 'kilo-auto/frontier', + }); - await expect( - caller.organizations.modes.create({ - organizationId: organization.id, - name: 'Code Mode', - slug: 'code', - config: { - roleDefinition: 'You are a coding assistant', - groups: ['read'], - defaultModel: 'openai/gpt-4o', - }, - }) - ).rejects.toThrow('Model access configuration is not available for this organization.'); - }); + await caller.organizations.modes.create({ + organizationId: freshOrg.id, + name: 'Create Clear Route Mode', + slug: 'create-clear-route-mode', + route_model: null, + }); - it('should reject an organization mode default that is denied', async () => { - const caller = await createCallerForUser(owner.id); - const organization = await createTestOrganization( - 'Denied Default Model Org', - owner.id, - 0, - { model_deny_list: ['openai/gpt-4o'] }, - false + const updatedOrganization = await getOrganizationById(freshOrg.id); + expect(updatedOrganization?.settings.org_auto_model?.routes['create-clear-route-mode']).toBe( + undefined ); - await expect( - caller.organizations.modes.create({ - organizationId: organization.id, - name: 'Code Mode', - slug: 'code', - config: { - roleDefinition: 'You are a coding assistant', - groups: ['read'], - defaultModel: 'openai/gpt-4o', - }, - }) - ).rejects.toThrow( - "Default model 'openai/gpt-4o' is not in the organization's allowed models list" - ); + const [audit] = await db + .select() + .from(organization_audit_logs) + .where(eq(organization_audit_logs.organization_id, freshOrg.id)) + .orderBy(desc(organization_audit_logs.created_at)); + + expect(audit?.message).toContain('Organization Auto route cleared'); + expect(audit?.message).toContain('create-clear-route-mode'); }); - it('should reject an empty organization mode default', async () => { + it('rejects route_model for non-enterprise organizations', async () => { const caller = await createCallerForUser(owner.id); + const teamsOrg = await createTestOrganization('Teams Route Org', owner.id, 0, {}, true); await expect( caller.organizations.modes.create({ - organizationId: testOrganization.id, - name: 'Code Mode', - slug: 'empty-default-model', - config: { - roleDefinition: 'You are a coding assistant', - groups: ['read'], - defaultModel: '', - }, + organizationId: teamsOrg.id, + name: 'Teams Route Mode', + slug: 'teams-route-mode', + route_model: 'kilo-auto/balanced', }) - ).rejects.toThrow(); + ).rejects.toThrow('Organization Auto is only available for Enterprise organizations.'); }); - it('should reject mode default writes when the release flag is disabled', async () => { - mockedIsReleaseToggleEnabled.mockResolvedValueOnce(false); + it('should throw error for duplicate slug', async () => { const caller = await createCallerForUser(owner.id); + // Create first mode + await caller.organizations.modes.create({ + organizationId: testOrganization.id, + name: 'First Mode', + slug: 'duplicate-slug', + }); + + // Try to create second mode with same slug await expect( caller.organizations.modes.create({ organizationId: testOrganization.id, - name: 'Code Mode', - slug: 'flag-disabled-default-model', - config: { - roleDefinition: 'You are a coding assistant', - groups: ['read'], - defaultModel: 'openai/gpt-4o', - }, + name: 'Second Mode', + slug: 'duplicate-slug', }) - ).rejects.toThrow('Mode default model configuration is not available'); - }); - - it('should allow mode default writes in development when the release flag is disabled', async () => { - mockedIsReleaseToggleEnabled.mockResolvedValue(false); - const replacedEnv = jest.replaceProperty(process, 'env', { - ...process.env, - NODE_ENV: 'development', - }); - const caller = await createCallerForUser(owner.id); - - try { - await expect( - caller.organizations.modes.create({ - organizationId: testOrganization.id, - name: 'Code Mode', - slug: 'development-default-model', - config: { - roleDefinition: 'You are a coding assistant', - groups: ['read'], - defaultModel: 'openai/gpt-4o', - }, - }) - ).resolves.toMatchObject({ - mode: { - config: { - defaultModel: 'openai/gpt-4o', - }, - }, - }); - } finally { - replacedEnv.restore(); - } + ).rejects.toThrow(); }); - it('should reject a wildcard organization mode default', async () => { + it('should validate slug format', async () => { const caller = await createCallerForUser(owner.id); await expect( caller.organizations.modes.create({ organizationId: testOrganization.id, - name: 'Code Mode', - slug: 'wildcard-default-model', - config: { - roleDefinition: 'You are a coding assistant', - groups: ['read'], - defaultModel: 'openai/*', - }, + name: 'Invalid Slug Mode', + slug: 'Invalid Slug!', }) - ).rejects.toThrow("Default model 'openai/*' is not a concrete model identifier"); + ).rejects.toThrow(); }); - it('should reject a wildcard organization mode default with a variant suffix', async () => { + it('should throw error for non-existent organization', async () => { const caller = await createCallerForUser(owner.id); + const nonExistentId = randomUUID(); await expect( caller.organizations.modes.create({ - organizationId: testOrganization.id, - name: 'Code Mode', - slug: 'wildcard-variant-default-model', - config: { - roleDefinition: 'You are a coding assistant', - groups: ['read'], - defaultModel: 'openai/*:free', - }, + organizationId: nonExistentId, + name: 'Test Mode', + slug: 'test', }) - ).rejects.toThrow("Default model 'openai/*:free' is not a concrete model identifier"); + ).rejects.toThrow(); }); }); @@ -346,6 +284,33 @@ describe('organization modes tRPC router', () => { expect(result.modes.map(m => m.slug).sort()).toEqual(['mode-1', 'mode-2']); }); + it('does not project canonical Organization Auto routes into mode responses', async () => { + const caller = await createCallerForUser(owner.id); + const created = await caller.organizations.modes.create({ + organizationId: testOrganization.id, + name: 'Projected Mode', + slug: 'projected-mode', + }); + await caller.organizations.settings.setOrganizationAutoRoute({ + organizationId: testOrganization.id, + mode_slug: 'projected-mode', + model_id: 'kilo-auto/frontier', + }); + + const result = await caller.organizations.modes.getById({ + organizationId: testOrganization.id, + modeId: created.mode.id, + }); + const listed = await caller.organizations.modes.list({ + organizationId: testOrganization.id, + }); + + expect(result.mode.config).not.toHaveProperty('defaultModel'); + expect(listed.modes.find(mode => mode.id === created.mode.id)?.config).not.toHaveProperty( + 'defaultModel' + ); + }); + it('should return empty array for organization with no modes', async () => { const caller = await createCallerForUser(owner.id); const emptyOrg = await createTestOrganization('Empty Org', owner.id, 0, {}, false); @@ -500,6 +465,34 @@ describe('organization modes tRPC router', () => { expect(result.mode.name).toBe('Member Update'); }); + it('stores update route_model in canonical Organization Auto routes', async () => { + const caller = await createCallerForUser(owner.id); + const freshOrg = await createTestOrganization( + 'Update With Route Org', + owner.id, + 0, + {}, + false + ); + const created = await caller.organizations.modes.create({ + organizationId: freshOrg.id, + name: 'Update With Route Mode', + slug: 'update-with-route-mode', + }); + + const result = await caller.organizations.modes.update({ + organizationId: freshOrg.id, + modeId: created.mode.id, + route_model: 'kilo-auto/frontier', + }); + + expect(result.mode.config).not.toHaveProperty('defaultModel'); + const updatedOrganization = await getOrganizationById(freshOrg.id); + expect(updatedOrganization?.settings.org_auto_model?.routes['update-with-route-mode']).toBe( + 'kilo-auto/frontier' + ); + }); + it('should throw error for non-existent mode', async () => { const caller = await createCallerForUser(owner.id); const nonExistentId = randomUUID(); @@ -539,19 +532,13 @@ describe('organization modes tRPC router', () => { ).rejects.toThrow(); }); - it('should reject an organization mode default on update for a non-enterprise organization', async () => { + it('should allow ordinary mode edits when the release flag is disabled', async () => { + mockedIsReleaseToggleEnabled.mockResolvedValue(false); const caller = await createCallerForUser(owner.id); - const organization = await createTestOrganization( - 'Teams Update Default Model Org', - owner.id, - 0, - {}, - true - ); const created = await caller.organizations.modes.create({ - organizationId: organization.id, + organizationId: testOrganization.id, name: 'Code Mode', - slug: 'code', + slug: 'flag-disabled-normal-update', config: { roleDefinition: 'You are a coding assistant', groups: ['read'], @@ -560,302 +547,420 @@ describe('organization modes tRPC router', () => { await expect( caller.organizations.modes.update({ - organizationId: organization.id, + organizationId: testOrganization.id, modeId: created.mode.id, config: { - defaultModel: 'openai/gpt-4o', + description: 'Updated description', }, }) - ).rejects.toThrow('Model access configuration is not available for this organization.'); + ).resolves.toMatchObject({ + mode: { + config: { + description: 'Updated description', + }, + }, + }); }); - it('should reject a denied organization mode default on update', async () => { + it('allows owners to rename routed modes when the release flag is disabled', async () => { const caller = await createCallerForUser(owner.id); - const organization = await createTestOrganization( - 'Denied Update Default Model Org', - owner.id, - 0, - { model_deny_list: ['openai/gpt-4o'] }, - false - ); const created = await caller.organizations.modes.create({ - organizationId: organization.id, - name: 'Code Mode', - slug: 'code', - config: { - roleDefinition: 'You are a coding assistant', - groups: ['read'], - }, + organizationId: testOrganization.id, + name: 'Flag Disabled Routed Rename', + slug: 'flag-disabled-routed-rename', }); + await caller.organizations.settings.setOrganizationAutoRoute({ + organizationId: testOrganization.id, + mode_slug: 'flag-disabled-routed-rename', + model_id: 'kilo-auto/balanced', + }); + mockedIsReleaseToggleEnabled.mockResolvedValue(false); - await expect( - caller.organizations.modes.update({ - organizationId: organization.id, - modeId: created.mode.id, - config: { - defaultModel: 'openai/gpt-4o', - }, - }) - ).rejects.toThrow( - "Default model 'openai/gpt-4o' is not in the organization's allowed models list" - ); + await caller.organizations.modes.update({ + organizationId: testOrganization.id, + modeId: created.mode.id, + slug: 'flag-disabled-routed-renamed', + }); + + const updatedOrganization = await getOrganizationById(testOrganization.id); + expect( + updatedOrganization?.settings.org_auto_model?.routes['flag-disabled-routed-rename'] + ).toBeUndefined(); + expect( + updatedOrganization?.settings.org_auto_model?.routes['flag-disabled-routed-renamed'] + ).toBe('kilo-auto/balanced'); + }); + }); + + it('migrates an Organization Auto route when a custom mode slug changes', async () => { + const caller = await createCallerForUser(owner.id); + const created = await caller.organizations.modes.create({ + organizationId: testOrganization.id, + name: 'Route Mode', + slug: 'route-mode', + }); + await caller.organizations.settings.setOrganizationAutoRoute({ + organizationId: testOrganization.id, + mode_slug: 'route-mode', + model_id: 'openai/gpt-4o', + }); + + await caller.organizations.modes.update({ + organizationId: testOrganization.id, + modeId: created.mode.id, + slug: 'renamed-route-mode', + }); + + const updatedOrganization = await getOrganizationById(testOrganization.id); + expect(updatedOrganization?.settings.org_auto_model?.routes['route-mode']).toBeUndefined(); + expect(updatedOrganization?.settings.org_auto_model?.routes['renamed-route-mode']).toBe( + 'openai/gpt-4o' + ); + }); + + it('rejects a rename when the destination already has an Organization Auto route', async () => { + const caller = await createCallerForUser(owner.id); + const created = await caller.organizations.modes.create({ + organizationId: testOrganization.id, + name: 'Source Route Mode', + slug: 'source-route-conflict', + }); + await caller.organizations.settings.setOrganizationAutoRoute({ + organizationId: testOrganization.id, + mode_slug: 'source-route-conflict', + model_id: 'kilo-auto/balanced', + }); + await caller.organizations.settings.setOrganizationAutoRoute({ + organizationId: testOrganization.id, + mode_slug: 'destination-route-conflict', + model_id: 'kilo-auto/frontier', + }); + + await expect( + caller.organizations.modes.update({ + organizationId: testOrganization.id, + modeId: created.mode.id, + slug: 'destination-route-conflict', + }) + ).rejects.toThrow( + 'Organization Auto route already exists for mode "destination-route-conflict"' + ); + + const updatedOrganization = await getOrganizationById(testOrganization.id); + expect(updatedOrganization?.settings.org_auto_model?.routes['source-route-conflict']).toBe( + 'kilo-auto/balanced' + ); + expect(updatedOrganization?.settings.org_auto_model?.routes['destination-route-conflict']).toBe( + 'kilo-auto/frontier' + ); + }); + + it('prevents members from renaming a mode with an Organization Auto route', async () => { + const caller = await createCallerForUser(owner.id); + const created = await caller.organizations.modes.create({ + organizationId: testOrganization.id, + name: 'Routed Rename Mode', + slug: 'routed-rename-test', + }); + await caller.organizations.settings.setOrganizationAutoRoute({ + organizationId: testOrganization.id, + mode_slug: 'routed-rename-test', + model_id: 'kilo-auto/balanced', + }); + + const memberCaller = await createCallerForUser(member.id); + await expect( + memberCaller.organizations.modes.update({ + organizationId: testOrganization.id, + modeId: created.mode.id, + slug: 'member-renamed-mode', + }) + ).rejects.toThrow('You do not have the required organizational role to access this feature'); + + const updatedOrganization = await getOrganizationById(testOrganization.id); + expect(updatedOrganization?.settings.org_auto_model?.routes['routed-rename-test']).toBe( + 'kilo-auto/balanced' + ); + }); + + it('rolls back route migration when a renamed mode collides', async () => { + const caller = await createCallerForUser(owner.id); + const routedMode = await caller.organizations.modes.create({ + organizationId: testOrganization.id, + name: 'Routed Collision Mode', + slug: 'routed-collision-test', + }); + await caller.organizations.settings.setOrganizationAutoRoute({ + organizationId: testOrganization.id, + mode_slug: 'routed-collision-test', + model_id: 'kilo-auto/balanced', + }); + await caller.organizations.modes.create({ + organizationId: testOrganization.id, + name: 'Existing Collision Mode', + slug: 'existing-collision-test', + }); + + await expect( + caller.organizations.modes.update({ + organizationId: testOrganization.id, + modeId: routedMode.mode.id, + slug: 'existing-collision-test', + }) + ).rejects.toThrow(); + + const updatedOrganization = await getOrganizationById(testOrganization.id); + expect(updatedOrganization?.settings.org_auto_model?.routes['routed-collision-test']).toBe( + 'kilo-auto/balanced' + ); + expect(updatedOrganization?.settings.org_auto_model?.routes['existing-collision-test']).toBe( + undefined + ); + }); + + describe('delete procedure', () => { + it('should delete a mode', async () => { + const caller = await createCallerForUser(owner.id); + + // Create a mode + const created = await caller.organizations.modes.create({ + organizationId: testOrganization.id, + name: 'To Be Deleted', + slug: 'to-be-deleted', + }); + + const result = await caller.organizations.modes.delete({ + organizationId: testOrganization.id, + modeId: created.mode.id, + }); + + expect(result.success).toBe(true); + + // Verify it's actually deleted + const modes = await getAllOrganizationModes(testOrganization.id); + expect(modes.find(m => m.id === created.mode.id)).toBeUndefined(); }); - it('should reject a wildcard organization mode default on update', async () => { + it('removes an Organization Auto route when a custom mode is deleted', async () => { const caller = await createCallerForUser(owner.id); const created = await caller.organizations.modes.create({ organizationId: testOrganization.id, - name: 'Code Mode', - slug: 'wildcard-update-default-model', - config: { - roleDefinition: 'You are a coding assistant', - groups: ['read'], - }, + name: 'Delete Route Mode', + slug: 'delete-route-mode', + }); + await caller.organizations.settings.setOrganizationAutoRoute({ + organizationId: testOrganization.id, + mode_slug: 'delete-route-mode', + model_id: 'openai/gpt-4o', }); - await expect( - caller.organizations.modes.update({ - organizationId: testOrganization.id, - modeId: created.mode.id, - config: { - defaultModel: 'openai/*', - }, - }) - ).rejects.toThrow("Default model 'openai/*' is not a concrete model identifier"); + await caller.organizations.modes.delete({ + organizationId: testOrganization.id, + modeId: created.mode.id, + }); + + const updatedOrganization = await getOrganizationById(testOrganization.id); + expect( + updatedOrganization?.settings.org_auto_model?.routes['delete-route-mode'] + ).toBeUndefined(); }); - it('should reject a wildcard organization mode default with a variant suffix on update', async () => { + it('allows owners to delete routed modes when the release flag is disabled', async () => { const caller = await createCallerForUser(owner.id); const created = await caller.organizations.modes.create({ organizationId: testOrganization.id, - name: 'Code Mode', - slug: 'wildcard-variant-update-default-model', - config: { - roleDefinition: 'You are a coding assistant', - groups: ['read'], - }, + name: 'Flag Disabled Routed Delete', + slug: 'flag-disabled-routed-delete', + }); + await caller.organizations.settings.setOrganizationAutoRoute({ + organizationId: testOrganization.id, + mode_slug: 'flag-disabled-routed-delete', + model_id: 'kilo-auto/balanced', }); + mockedIsReleaseToggleEnabled.mockResolvedValue(false); - await expect( - caller.organizations.modes.update({ - organizationId: testOrganization.id, - modeId: created.mode.id, - config: { - defaultModel: 'openai/*:free', - }, - }) - ).rejects.toThrow("Default model 'openai/*:free' is not a concrete model identifier"); + await caller.organizations.modes.delete({ + organizationId: testOrganization.id, + modeId: created.mode.id, + }); + + const updatedOrganization = await getOrganizationById(testOrganization.id); + expect( + updatedOrganization?.settings.org_auto_model?.routes['flag-disabled-routed-delete'] + ).toBeUndefined(); }); - it('should allow unrelated mode edits after a stored default becomes denied', async () => { + it('preserves a canonical route when a built-in override is reverted', async () => { const caller = await createCallerForUser(owner.id); - const organization = await createTestOrganization( - 'Stale Default Model Org', - owner.id, - 0, - {}, - false - ); + const freshOrg = await createTestOrganization('Revert Route Org', owner.id, 0, {}, false); const created = await caller.organizations.modes.create({ - organizationId: organization.id, - name: 'Code Mode', - slug: 'stale-default-model', - config: { - roleDefinition: 'You are a coding assistant', - groups: ['read'], - defaultModel: 'openai/gpt-4o', - }, + organizationId: freshOrg.id, + name: 'Custom Code', + slug: 'code', + }); + await caller.organizations.settings.setOrganizationAutoRoute({ + organizationId: freshOrg.id, + mode_slug: 'code', + model_id: 'kilo-auto/frontier', }); - await db - .update(organizations) - .set({ settings: { model_deny_list: ['openai/gpt-4o'] } }) - .where(eq(organizations.id, organization.id)); - await expect( - caller.organizations.modes.update({ - organizationId: organization.id, - modeId: created.mode.id, - config: { - description: 'Updated description', - }, - }) - ).resolves.toMatchObject({ - mode: { - config: { - description: 'Updated description', - defaultModel: 'openai/gpt-4o', - }, - }, + await caller.organizations.modes.delete({ + organizationId: freshOrg.id, + modeId: created.mode.id, + preserve_route: true, }); + + const updatedOrganization = await getOrganizationById(freshOrg.id); + expect(updatedOrganization?.settings.org_auto_model?.routes.code).toBe('kilo-auto/frontier'); }); - it('should allow clearing a mode default after an enterprise organization downgrades', async () => { + it('updates a preserved route when a built-in override is reverted', async () => { const caller = await createCallerForUser(owner.id); - const organization = await createTestOrganization( - 'Downgraded Default Model Org', + const freshOrg = await createTestOrganization( + 'Revert Changed Route Org', owner.id, 0, {}, false ); const created = await caller.organizations.modes.create({ - organizationId: organization.id, - name: 'Code Mode', - slug: 'downgraded-clear-default-model', - config: { - roleDefinition: 'You are a coding assistant', - groups: ['read'], - defaultModel: 'openai/gpt-4o', - }, + organizationId: freshOrg.id, + name: 'Custom Code', + slug: 'code', + }); + await caller.organizations.settings.setOrganizationAutoRoute({ + organizationId: freshOrg.id, + mode_slug: 'code', + model_id: 'kilo-auto/frontier', }); - await db - .update(organizations) - .set({ plan: 'teams' }) - .where(eq(organizations.id, organization.id)); - await expect( - caller.organizations.modes.update({ - organizationId: organization.id, - modeId: created.mode.id, - config: { - defaultModel: null, - }, - }) - ).resolves.toMatchObject({ - mode: { - config: { - roleDefinition: 'You are a coding assistant', - }, - }, + await caller.organizations.modes.delete({ + organizationId: freshOrg.id, + modeId: created.mode.id, + preserve_route: true, + route_model: 'kilo-auto/balanced', }); + + const updatedOrganization = await getOrganizationById(freshOrg.id); + expect(updatedOrganization?.settings.org_auto_model?.routes.code).toBe('kilo-auto/balanced'); }); - it('should reject clearing a mode default when the release flag is disabled', async () => { + it('rejects route_model when deleting a custom mode', async () => { const caller = await createCallerForUser(owner.id); + const freshOrg = await createTestOrganization( + 'Delete Custom Route Org', + owner.id, + 0, + {}, + false + ); const created = await caller.organizations.modes.create({ - organizationId: testOrganization.id, - name: 'Code Mode', - slug: 'flag-disabled-clear-default-model', - config: { - roleDefinition: 'You are a coding assistant', - groups: ['read'], - defaultModel: 'openai/gpt-4o', - }, + organizationId: freshOrg.id, + name: 'Custom Route Mode', + slug: 'custom-route-mode', + }); + await caller.organizations.settings.setOrganizationAutoRoute({ + organizationId: freshOrg.id, + mode_slug: 'custom-route-mode', + model_id: 'kilo-auto/frontier', }); - mockedIsReleaseToggleEnabled.mockResolvedValueOnce(false); await expect( - caller.organizations.modes.update({ - organizationId: testOrganization.id, + caller.organizations.modes.delete({ + organizationId: freshOrg.id, modeId: created.mode.id, - config: { - defaultModel: null, - }, + route_model: 'kilo-auto/balanced', }) - ).rejects.toThrow('Mode default model configuration is not available'); + ).rejects.toThrow('Route updates can only be preserved when reverting a built-in mode.'); + + const updatedOrganization = await getOrganizationById(freshOrg.id); + expect(updatedOrganization?.settings.org_auto_model?.routes['custom-route-mode']).toBe( + 'kilo-auto/frontier' + ); + expect( + (await getAllOrganizationModes(freshOrg.id)).find(mode => mode.id === created.mode.id) + ).toBeDefined(); }); - it('should allow ordinary mode edits when the release flag is disabled', async () => { - mockedIsReleaseToggleEnabled.mockResolvedValue(false); + it('prevents members from preserving a routed built-in override', async () => { const caller = await createCallerForUser(owner.id); + const freshOrg = await createTestOrganization( + 'Member Revert Route Org', + owner.id, + 0, + {}, + false + ); + await addUserToOrganization(freshOrg.id, member.id, 'member'); const created = await caller.organizations.modes.create({ - organizationId: testOrganization.id, - name: 'Code Mode', - slug: 'flag-disabled-normal-update', - config: { - roleDefinition: 'You are a coding assistant', - groups: ['read'], - }, + organizationId: freshOrg.id, + name: 'Custom Code', + slug: 'code', + }); + await caller.organizations.settings.setOrganizationAutoRoute({ + organizationId: freshOrg.id, + mode_slug: 'code', + model_id: 'kilo-auto/frontier', }); + const memberCaller = await createCallerForUser(member.id); await expect( - caller.organizations.modes.update({ - organizationId: testOrganization.id, + memberCaller.organizations.modes.delete({ + organizationId: freshOrg.id, modeId: created.mode.id, - config: { - description: 'Updated description', - }, + preserve_route: true, }) - ).resolves.toMatchObject({ - mode: { - config: { - description: 'Updated description', - }, - }, - }); - }); - - it('should clear an organization mode default on update', async () => { - const caller = await createCallerForUser(owner.id); - const created = await caller.organizations.modes.create({ - organizationId: testOrganization.id, - name: 'Code Mode', - slug: 'clear-default-model', - config: { - roleDefinition: 'You are a coding assistant', - description: 'Write code', - groups: ['read'], - defaultModel: 'openai/gpt-4o', - }, - }); + ).rejects.toThrow('You do not have the required organizational role to access this feature'); - const result = await caller.organizations.modes.update({ - organizationId: testOrganization.id, - modeId: created.mode.id, - config: { - defaultModel: null, - }, - }); - - expect(result.mode.config).toEqual({ - roleDefinition: 'You are a coding assistant', - description: 'Write code', - groups: ['read'], - }); + const updatedOrganization = await getOrganizationById(freshOrg.id); + expect(updatedOrganization?.settings.org_auto_model?.routes.code).toBe('kilo-auto/frontier'); }); - }); - describe('delete procedure', () => { - it('should delete a mode', async () => { + it('should allow members to delete modes', async () => { const caller = await createCallerForUser(owner.id); // Create a mode const created = await caller.organizations.modes.create({ organizationId: testOrganization.id, - name: 'To Be Deleted', - slug: 'to-be-deleted', + name: 'Delete Test Mode', + slug: 'delete-test', }); - const result = await caller.organizations.modes.delete({ + const memberCaller = await createCallerForUser(member.id); + + const result = await memberCaller.organizations.modes.delete({ organizationId: testOrganization.id, modeId: created.mode.id, }); expect(result.success).toBe(true); - - // Verify it's actually deleted - const modes = await getAllOrganizationModes(testOrganization.id); - expect(modes.find(m => m.id === created.mode.id)).toBeUndefined(); }); - it('should allow members to delete modes', async () => { + it('prevents members from deleting a mode with an Organization Auto route', async () => { const caller = await createCallerForUser(owner.id); - - // Create a mode const created = await caller.organizations.modes.create({ organizationId: testOrganization.id, - name: 'Delete Test Mode', - slug: 'delete-test', + name: 'Routed Mode', + slug: 'routed-delete-test', }); - - const memberCaller = await createCallerForUser(member.id); - - const result = await memberCaller.organizations.modes.delete({ + await caller.organizations.settings.setOrganizationAutoRoute({ organizationId: testOrganization.id, - modeId: created.mode.id, + mode_slug: 'routed-delete-test', + model_id: 'kilo-auto/balanced', }); - expect(result.success).toBe(true); + const memberCaller = await createCallerForUser(member.id); + await expect( + memberCaller.organizations.modes.delete({ + organizationId: testOrganization.id, + modeId: created.mode.id, + }) + ).rejects.toThrow('You do not have the required organizational role to access this feature'); + + const updatedOrganization = await getOrganizationById(testOrganization.id); + expect(updatedOrganization?.settings.org_auto_model?.routes['routed-delete-test']).toBe( + 'kilo-auto/balanced' + ); }); it('should throw error for non-existent mode', async () => { diff --git a/apps/web/src/routers/organizations/organization-modes-router.ts b/apps/web/src/routers/organizations/organization-modes-router.ts index 4a78227681..9d42dfce98 100644 --- a/apps/web/src/routers/organizations/organization-modes-router.ts +++ b/apps/web/src/routers/organizations/organization-modes-router.ts @@ -1,5 +1,6 @@ import { createTRPCRouter } from '@/lib/trpc/init'; import { + ensureOrganizationAccess, OrganizationIdInputSchema, organizationMemberProcedure, organizationMemberMutationProcedure, @@ -8,35 +9,37 @@ import { TRPCError } from '@trpc/server'; import * as z from 'zod'; import { createOrganizationMode, + deleteOrganizationMode, getAllOrganizationModes, getOrganizationModeById, + type OrganizationMode, updateOrganizationMode, - deleteOrganizationMode, } from '@/lib/organizations/organization-modes'; import { OrganizationModeConfigSchema, type OrganizationModeConfig, + type OrganizationSettings, } from '@/lib/organizations/organization-types'; import { createAuditLog } from '@/lib/organizations/organization-audit-logs'; -import { getOrganizationById } from '@/lib/organizations/organizations'; +import { getOrganizationById, mutateOrganizationSettings } from '@/lib/organizations/organizations'; import { successResult } from '@/lib/maybe-result'; -import { createAllowPredicateFromRestrictions } from '@/lib/model-allow.server'; -import { getEffectiveModelRestrictions } from '@/lib/organizations/model-restrictions'; -import { normalizeModelId } from '@/lib/ai-gateway/model-utils'; import { isReleaseToggleEnabled } from '@/lib/posthog-feature-flags'; - -const ORGANIZATION_MODE_DEFAULT_MODEL_FLAG = 'org-default-model-config'; +import { db, type DrizzleTransaction } from '@/lib/drizzle'; +import type { Organization } from '@kilocode/db/schema'; +import { + DEFAULT_ORGANIZATION_AUTO_MODEL_SETTINGS, + ORGANIZATION_AUTO_MODEL_FLAG, + MAX_ORGANIZATION_AUTO_ROUTES, + validateOrganizationAutoTarget, +} from '@/lib/organizations/organization-auto-model'; const ModeConfigInputSchema = OrganizationModeConfigSchema.partial(); - -const ModeUpdateConfigInputSchema = ModeConfigInputSchema.extend({ - defaultModel: z.string().min(1, 'Default model cannot be empty').nullable().optional(), -}); - -type ModeUpdateConfigInput = z.infer; -type DefaultModelConfig = { - defaultModel?: string | null; -}; +const RouteModelInputSchema = z + .string() + .trim() + .nullable() + .transform(value => (value === '' ? null : value)) + .optional(); const CreateModeInputSchema = OrganizationIdInputSchema.extend({ name: z @@ -49,6 +52,7 @@ const CreateModeInputSchema = OrganizationIdInputSchema.extend({ .max(50, 'Mode slug must be less than 50 characters') .regex(/^[a-z0-9-]+$/, 'Mode slug must contain only lowercase letters, numbers, and hyphens'), config: ModeConfigInputSchema.optional(), + route_model: RouteModelInputSchema, }); const UpdateModeInputSchema = OrganizationIdInputSchema.extend({ @@ -60,121 +64,180 @@ const UpdateModeInputSchema = OrganizationIdInputSchema.extend({ .max(50) .regex(/^[a-z0-9-]+$/) .optional(), - config: ModeUpdateConfigInputSchema.optional(), + config: ModeConfigInputSchema.optional(), + route_model: RouteModelInputSchema, }); const DeleteModeInputSchema = OrganizationIdInputSchema.extend({ modeId: z.uuid(), + preserve_route: z.boolean().optional(), + route_model: RouteModelInputSchema, }); const ModeIdInputSchema = OrganizationIdInputSchema.extend({ modeId: z.uuid(), }); -type DefaultModelChange = - | { kind: 'none' } - | { kind: 'clear' } - | { kind: 'set'; defaultModel: string }; - -function getDefaultModelChange(config: DefaultModelConfig | undefined): DefaultModelChange { - if (!config || !Object.prototype.hasOwnProperty.call(config, 'defaultModel')) { - return { kind: 'none' }; - } - - if (config.defaultModel === null) { - return { kind: 'clear' }; - } - - if (typeof config.defaultModel === 'string') { - return { kind: 'set', defaultModel: config.defaultModel }; - } +const BUILT_IN_MODE_SLUGS = new Set(['architect', 'code', 'ask', 'debug', 'orchestrator']); - return { kind: 'none' }; +function hasRoute(routes: Record, slug: string): boolean { + return Object.prototype.hasOwnProperty.call(routes, slug); } -function assertDefaultModelCanBeSet( - organization: NonNullable>>, - change: DefaultModelChange -): void { - if (change.kind !== 'set') { - return; - } - - if (organization.plan !== 'enterprise') { +async function assertOrganizationAutoWriteEnabled(userId: string): Promise { + if ( + process.env.NODE_ENV !== 'development' && + !(await isReleaseToggleEnabled(ORGANIZATION_AUTO_MODEL_FLAG, userId)) + ) { throw new TRPCError({ code: 'FORBIDDEN', - message: 'Model access configuration is not available for this organization.', + message: 'Organization Auto routing configuration is not available', }); } } -async function assertDefaultModelConfigEnabled( - userId: string, - change: DefaultModelChange -): Promise { - if (change.kind === 'none') { - return; - } +function getOrganizationAutoSettings( + settings: OrganizationSettings +): typeof DEFAULT_ORGANIZATION_AUTO_MODEL_SETTINGS { + return settings.org_auto_model ?? DEFAULT_ORGANIZATION_AUTO_MODEL_SETTINGS; +} - if ( - process.env.NODE_ENV !== 'development' && - !(await isReleaseToggleEnabled(ORGANIZATION_MODE_DEFAULT_MODEL_FLAG, userId)) - ) { +function assertOrganizationAutoEligible(organization: Pick): void { + if (organization.plan !== 'enterprise') { throw new TRPCError({ code: 'FORBIDDEN', - message: 'Mode default model configuration is not available', + message: 'Organization Auto is only available for Enterprise organizations.', }); } } -function normalizeModeConfig( - config: ModeUpdateConfigInput | undefined -): Partial | undefined { - if (!config) { - return undefined; +type OrganizationAccessContext = Parameters[0]; + +async function applyOrganizationAutoRouteChange( + organization: Pick, + modeSlug: string, + routeModel: string | null, + ctx: OrganizationAccessContext, + tx?: DrizzleTransaction +): Promise { + await assertOrganizationAutoWriteEnabled(ctx.user.id); + assertOrganizationAutoEligible(organization); + await ensureOrganizationAccess(ctx, organization.id, ['owner', 'billing_manager']); + const orgAutoModel = getOrganizationAutoSettings(organization.settings); + const currentRoute = orgAutoModel.routes[modeSlug]; + if (routeModel === null && !currentRoute) { + return organization.settings; } - const { defaultModel, ...rest } = config; - if (defaultModel === null) { - return { ...rest, defaultModel: undefined }; - } - if (defaultModel === undefined) { - return rest; - } + const routes = { ...orgAutoModel.routes }; + if (routeModel === null) { + delete routes[modeSlug]; + } else { + const validation = await validateOrganizationAutoTarget(organization, routeModel, { + dbClient: tx, + }); - return { ...rest, defaultModel }; -} + if (validation.kind === 'error') { + throw new TRPCError({ code: 'BAD_REQUEST', message: validation.message }); + } + if (currentRoute === validation.modelId) { + return organization.settings; + } + routes[modeSlug] = validation.modelId; + } -async function validateDefaultModel( - organization: NonNullable>>, - defaultModel: string -): Promise { - const normalizedDefaultModel = normalizeModelId(defaultModel); - if (normalizedDefaultModel.endsWith('/*')) { + if (Object.keys(routes).length > MAX_ORGANIZATION_AUTO_ROUTES) { throw new TRPCError({ code: 'BAD_REQUEST', - message: `Default model '${defaultModel}' is not a concrete model identifier`, + message: `Organization Auto supports at most ${MAX_ORGANIZATION_AUTO_ROUTES} routes.`, }); } - const isAllowed = createAllowPredicateFromRestrictions( - getEffectiveModelRestrictions(organization) - ); + return { + ...organization.settings, + org_auto_model: { ...orgAutoModel, routes }, + }; +} - if (!(await isAllowed(normalizedDefaultModel))) { - throw new TRPCError({ - code: 'BAD_REQUEST', - message: `Default model '${defaultModel}' is not in the organization's allowed models list`, - }); +function createModeUpdateAuditMessage( + existingMode: OrganizationMode, + updates: { name?: string; slug?: string; config?: Partial } +): string { + const changes: string[] = []; + if (updates.name && updates.name !== existingMode.name) { + changes.push(`name: "${existingMode.name}" → "${updates.name}"`); + } + if (updates.slug && updates.slug !== existingMode.slug) { + changes.push(`slug: "${existingMode.slug}" → "${updates.slug}"`); + } + if (updates.config) { + const auditConfig = updates.config; + const configChanges: string[] = []; + + if ( + 'roleDefinition' in auditConfig && + auditConfig.roleDefinition !== existingMode.config.roleDefinition + ) { + const oldValue = existingMode.config.roleDefinition || '(empty)'; + const newValue = auditConfig.roleDefinition || '(empty)'; + configChanges.push( + `roleDefinition: "${oldValue.substring(0, 50)}${oldValue.length > 50 ? '...' : ''}" → "${newValue.substring(0, 50)}${newValue.length > 50 ? '...' : ''}"` + ); + } + if ('whenToUse' in auditConfig && auditConfig.whenToUse !== existingMode.config.whenToUse) { + const oldValue = existingMode.config.whenToUse || '(empty)'; + const newValue = auditConfig.whenToUse || '(empty)'; + configChanges.push( + `whenToUse: "${oldValue.substring(0, 50)}${oldValue.length > 50 ? '...' : ''}" → "${newValue.substring(0, 50)}${newValue.length > 50 ? '...' : ''}"` + ); + } + if ( + 'description' in auditConfig && + auditConfig.description !== existingMode.config.description + ) { + const oldValue = existingMode.config.description || '(empty)'; + const newValue = auditConfig.description || '(empty)'; + configChanges.push( + `description: "${oldValue.substring(0, 50)}${oldValue.length > 50 ? '...' : ''}" → "${newValue.substring(0, 50)}${newValue.length > 50 ? '...' : ''}"` + ); + } + if ( + 'customInstructions' in auditConfig && + auditConfig.customInstructions !== existingMode.config.customInstructions + ) { + const oldValue = existingMode.config.customInstructions || '(empty)'; + const newValue = auditConfig.customInstructions || '(empty)'; + configChanges.push( + `customInstructions: "${oldValue.substring(0, 50)}${oldValue.length > 50 ? '...' : ''}" → "${newValue.substring(0, 50)}${newValue.length > 50 ? '...' : ''}"` + ); + } + if ( + auditConfig.groups !== undefined && + existingMode.config.groups !== undefined && + JSON.stringify(auditConfig.groups) !== JSON.stringify(existingMode.config.groups) + ) { + const oldValue = JSON.stringify(existingMode.config.groups); + const newValue = JSON.stringify(auditConfig.groups); + configChanges.push( + `groups: ${oldValue.substring(0, 50)}${oldValue.length > 50 ? '...' : ''} → ${newValue.substring(0, 50)}${newValue.length > 50 ? '...' : ''}` + ); + } + + if (configChanges.length > 0) { + changes.push(...configChanges); + } else { + changes.push('config updated (no property changes detected)'); + } } + + return `Updated mode "${existingMode.name}"${changes.length > 0 ? `: ${changes.join(', ')}` : ''}`; } export const organizationModesRouter = createTRPCRouter({ create: organizationMemberMutationProcedure .input(CreateModeInputSchema) .mutation(async ({ input, ctx }) => { - const { organizationId, name, slug, config } = input; - + const { organizationId, name, slug, config, route_model } = input; const organization = await getOrganizationById(organizationId); if (!organization) { throw new TRPCError({ @@ -183,46 +246,85 @@ export const organizationModesRouter = createTRPCRouter({ }); } - const defaultModelChange = getDefaultModelChange(config); - assertDefaultModelCanBeSet(organization, defaultModelChange); - await assertDefaultModelConfigEnabled(ctx.user.id, defaultModelChange); - if (defaultModelChange.kind === 'set') { - await validateDefaultModel(organization, defaultModelChange.defaultModel); + if (route_model !== undefined) { + await assertOrganizationAutoWriteEnabled(ctx.user.id); + assertOrganizationAutoEligible(organization); + await ensureOrganizationAccess(ctx, organizationId, ['owner', 'billing_manager']); } - const mode = await createOrganizationMode( - organizationId, - ctx.user.id, - name, - slug, - normalizeModeConfig(config) - ); + let createdMode: OrganizationMode | null | undefined; + let routeAuditMessage: string | undefined; + await db.transaction(async tx => { + await mutateOrganizationSettings( + organizationId, + async lockedOrganization => { + createdMode = await createOrganizationMode( + organizationId, + ctx.user.id, + name, + slug, + config, + tx + ); - if (!mode) { - throw new TRPCError({ - code: 'CONFLICT', - message: `A mode with slug "${slug}" already exists in this organization`, - }); - } + if (!createdMode) { + throw new TRPCError({ + code: 'CONFLICT', + message: `A mode with slug "${slug}" already exists in this organization`, + }); + } + + if (route_model === undefined) { + return lockedOrganization.settings; + } + const previousRoute = + lockedOrganization.settings.org_auto_model?.routes[createdMode.slug]; + const nextSettings = await applyOrganizationAutoRouteChange( + lockedOrganization, + createdMode.slug, + route_model, + ctx, + tx + ); + + const nextRoute = nextSettings.org_auto_model?.routes[createdMode.slug]; + if (previousRoute !== nextRoute) { + routeAuditMessage = nextRoute + ? `Organization Auto route set: "${createdMode.slug}" → "${nextRoute}"` + : `Organization Auto route cleared: "${createdMode.slug}"${previousRoute ? ` (was "${previousRoute}")` : ''}`; + } - await createAuditLog({ - action: 'organization.mode.create', - actor_email: ctx.user.google_user_email, - actor_id: ctx.user.id, - actor_name: ctx.user.google_user_name, - message: `Created mode "${name}" with slug "${slug}": ${JSON.stringify(config)}`, - organization_id: organizationId, + return nextSettings; + }, + tx + ); + + if (!createdMode) { + throw new TRPCError({ code: 'INTERNAL_SERVER_ERROR', message: 'Mode creation failed' }); + } + + await createAuditLog({ + action: 'organization.mode.create', + actor_email: ctx.user.google_user_email, + actor_id: ctx.user.id, + actor_name: ctx.user.google_user_name, + message: `Created mode "${name}" with slug "${slug}": ${JSON.stringify(config)}${routeAuditMessage ? `, ${routeAuditMessage}` : ''}`, + organization_id: organizationId, + tx, + }); }); - return { mode }; + if (!createdMode) { + throw new TRPCError({ code: 'INTERNAL_SERVER_ERROR', message: 'Mode creation failed' }); + } + + return { mode: createdMode }; }), list: organizationMemberProcedure.input(OrganizationIdInputSchema).query(async ({ input }) => { const { organizationId } = input; - const modes = await getAllOrganizationModes(organizationId); - - return { modes }; + return { modes: await getAllOrganizationModes(organizationId) }; }), getById: organizationMemberProcedure.input(ModeIdInputSchema).query(async ({ input }) => { @@ -243,16 +345,7 @@ export const organizationModesRouter = createTRPCRouter({ update: organizationMemberMutationProcedure .input(UpdateModeInputSchema) .mutation(async ({ input, ctx }) => { - const { modeId, organizationId, ...updates } = input; - - const existingMode = await getOrganizationModeById(organizationId, modeId); - - if (!existingMode) { - throw new TRPCError({ - code: 'NOT_FOUND', - message: 'Mode not found', - }); - } + const { modeId, organizationId, route_model, ...updates } = input; const organization = await getOrganizationById(organizationId); if (!organization) { @@ -262,142 +355,239 @@ export const organizationModesRouter = createTRPCRouter({ }); } - const defaultModelChange = getDefaultModelChange(updates.config); - assertDefaultModelCanBeSet(organization, defaultModelChange); - await assertDefaultModelConfigEnabled(ctx.user.id, defaultModelChange); - if (defaultModelChange.kind === 'set') { - await validateDefaultModel(organization, defaultModelChange.defaultModel); - } - const normalizedConfig = normalizeModeConfig(updates.config); - - const mode = await updateOrganizationMode(organizationId, modeId, { - ...updates, - config: normalizedConfig, - }); - - if (!mode) { - throw new TRPCError({ - code: 'CONFLICT', - message: `A mode with slug "${updates.slug}" already exists in this organization`, - }); + if (route_model !== undefined) { + await assertOrganizationAutoWriteEnabled(ctx.user.id); + assertOrganizationAutoEligible(organization); + await ensureOrganizationAccess(ctx, organizationId, ['owner', 'billing_manager']); } - const changes: string[] = []; - if (updates.name && updates.name !== existingMode.name) { - changes.push(`name: "${existingMode.name}" → "${updates.name}"`); - } - if (updates.slug && updates.slug !== existingMode.slug) { - changes.push(`slug: "${existingMode.slug}" → "${updates.slug}"`); - } - if (updates.config) { - const auditConfig = normalizedConfig ?? updates.config; - const configChanges: string[] = []; - - if ( - 'roleDefinition' in auditConfig && - auditConfig.roleDefinition !== existingMode.config.roleDefinition - ) { - const oldValue = existingMode.config.roleDefinition || '(empty)'; - const newValue = auditConfig.roleDefinition || '(empty)'; - configChanges.push( - `roleDefinition: "${oldValue.substring(0, 50)}${oldValue.length > 50 ? '...' : ''}" → "${newValue.substring(0, 50)}${newValue.length > 50 ? '...' : ''}"` - ); - } - if ('whenToUse' in auditConfig && auditConfig.whenToUse !== existingMode.config.whenToUse) { - const oldValue = existingMode.config.whenToUse || '(empty)'; - const newValue = auditConfig.whenToUse || '(empty)'; - configChanges.push( - `whenToUse: "${oldValue.substring(0, 50)}${oldValue.length > 50 ? '...' : ''}" → "${newValue.substring(0, 50)}${newValue.length > 50 ? '...' : ''}"` - ); - } - if ( - 'description' in auditConfig && - auditConfig.description !== existingMode.config.description - ) { - const oldValue = existingMode.config.description || '(empty)'; - const newValue = auditConfig.description || '(empty)'; - configChanges.push( - `description: "${oldValue.substring(0, 50)}${oldValue.length > 50 ? '...' : ''}" → "${newValue.substring(0, 50)}${newValue.length > 50 ? '...' : ''}"` - ); - } - if ( - 'customInstructions' in auditConfig && - auditConfig.customInstructions !== existingMode.config.customInstructions - ) { - const oldValue = existingMode.config.customInstructions || '(empty)'; - const newValue = auditConfig.customInstructions || '(empty)'; - configChanges.push( - `customInstructions: "${oldValue.substring(0, 50)}${oldValue.length > 50 ? '...' : ''}" → "${newValue.substring(0, 50)}${newValue.length > 50 ? '...' : ''}"` - ); - } - if ( - 'defaultModel' in auditConfig && - auditConfig.defaultModel !== existingMode.config.defaultModel - ) { - if (existingMode.config.defaultModel && auditConfig.defaultModel) { - configChanges.push( - `defaultModel: "${existingMode.config.defaultModel}" → "${auditConfig.defaultModel}"` - ); - } else if (auditConfig.defaultModel) { - configChanges.push(`defaultModel: set to "${auditConfig.defaultModel}"`); - } else if (existingMode.config.defaultModel) { - configChanges.push(`defaultModel: cleared "${existingMode.config.defaultModel}"`); - } - } - if ( - auditConfig.groups !== undefined && - existingMode.config.groups !== undefined && - JSON.stringify(auditConfig.groups) !== JSON.stringify(existingMode.config.groups) - ) { - const oldValue = JSON.stringify(existingMode.config.groups); - const newValue = JSON.stringify(auditConfig.groups); - configChanges.push( - `groups: ${oldValue.substring(0, 50)}${oldValue.length > 50 ? '...' : ''} → ${newValue.substring(0, 50)}${newValue.length > 50 ? '...' : ''}` - ); + const hasModeUpdates = + updates.name !== undefined || updates.slug !== undefined || updates.config !== undefined; + let existingMode: OrganizationMode | undefined; + let updatedMode: OrganizationMode | null | undefined; + const routeAuditChanges: string[] = []; + await db.transaction(async tx => { + await mutateOrganizationSettings( + organizationId, + async lockedOrganization => { + const lockedMode = await getOrganizationModeById(organizationId, modeId, tx, true); + if (!lockedMode) { + throw new TRPCError({ + code: 'NOT_FOUND', + message: 'Mode not found', + }); + } + existingMode = lockedMode; + + const orgAutoModel = getOrganizationAutoSettings(lockedOrganization.settings); + const routes = { ...orgAutoModel.routes }; + const nextSlug = updates.slug ?? lockedMode.slug; + const slugChanged = nextSlug !== lockedMode.slug; + const sourceHasRoute = hasRoute(routes, lockedMode.slug); + const initialRoute = routes[lockedMode.slug]; + + let nextSettings = lockedOrganization.settings; + + if (sourceHasRoute && slugChanged) { + await ensureOrganizationAccess(ctx, organizationId, ['owner', 'billing_manager']); + if (hasRoute(routes, nextSlug)) { + throw new TRPCError({ + code: 'CONFLICT', + message: `Organization Auto route already exists for mode "${nextSlug}"`, + }); + } + + const targetModelId = routes[lockedMode.slug]; + delete routes[lockedMode.slug]; + routes[nextSlug] = targetModelId; + routeAuditChanges.push( + `Organization Auto route migrated: "${lockedMode.slug}" → "${nextSlug}" (${targetModelId})` + ); + + nextSettings = { + ...lockedOrganization.settings, + org_auto_model: { + ...orgAutoModel, + routes, + }, + }; + } + + if (hasModeUpdates) { + updatedMode = await updateOrganizationMode(organizationId, modeId, updates, tx); + + if (!updatedMode) { + throw new TRPCError({ + code: 'CONFLICT', + message: `A mode with slug "${updates.slug}" already exists in this organization`, + }); + } + } else if (route_model !== undefined) { + updatedMode = lockedMode; + } else { + throw new TRPCError({ code: 'BAD_REQUEST', message: 'No mode updates provided' }); + } + + if (route_model !== undefined) { + const previousRoute = nextSettings.org_auto_model?.routes[nextSlug]; + nextSettings = await applyOrganizationAutoRouteChange( + { ...lockedOrganization, settings: nextSettings }, + nextSlug, + route_model, + ctx, + tx + ); + const nextRoute = nextSettings.org_auto_model?.routes[nextSlug]; + if (previousRoute !== nextRoute) { + if (slugChanged) { + routeAuditChanges.length = 0; + routeAuditChanges.push( + nextRoute + ? `Organization Auto route: "${lockedMode.slug}"${initialRoute ? ` "${initialRoute}"` : ''} → "${nextSlug}" "${nextRoute}"` + : `Organization Auto route removed: "${lockedMode.slug}"${initialRoute ? ` (was "${initialRoute}")` : ''}` + ); + } else { + routeAuditChanges.push( + nextRoute + ? `Organization Auto route: "${nextSlug}" ${previousRoute ? `"${previousRoute}" → ` : 'set to '}"${nextRoute}"` + : `Organization Auto route cleared: "${nextSlug}"${previousRoute ? ` (was "${previousRoute}")` : ''}` + ); + } + } + } + + return nextSettings; + }, + tx + ); + + if (!existingMode) { + throw new TRPCError({ code: 'INTERNAL_SERVER_ERROR', message: 'Mode update failed' }); } - if (configChanges.length > 0) { - changes.push(...configChanges); - } else { - changes.push('config updated (no property changes detected)'); + if (hasModeUpdates || routeAuditChanges.length > 0) { + await createAuditLog({ + action: 'organization.mode.update', + actor_email: ctx.user.google_user_email, + actor_id: ctx.user.id, + actor_name: ctx.user.google_user_name, + message: `${createModeUpdateAuditMessage(existingMode, updates)}${routeAuditChanges.length > 0 ? `, ${routeAuditChanges.join(', ')}` : ''}`, + organization_id: existingMode.organization_id, + tx, + }); } - } - - await createAuditLog({ - action: 'organization.mode.update', - actor_email: ctx.user.google_user_email, - actor_id: ctx.user.id, - actor_name: ctx.user.google_user_name, - message: `Updated mode "${existingMode.name}"${changes.length > 0 ? `: ${changes.join(', ')}` : ''}`, - organization_id: existingMode.organization_id, }); - return { mode }; + if (!updatedMode) { + throw new TRPCError({ code: 'INTERNAL_SERVER_ERROR', message: 'Mode update failed' }); + } + + return { mode: updatedMode }; }), delete: organizationMemberMutationProcedure .input(DeleteModeInputSchema) .mutation(async ({ input, ctx }) => { - const { modeId, organizationId } = input; - - const mode = await getOrganizationModeById(organizationId, modeId); + const { modeId, organizationId, preserve_route = false, route_model } = input; - if (!mode) { + const organization = await getOrganizationById(organizationId); + if (!organization) { throw new TRPCError({ code: 'NOT_FOUND', - message: 'Mode not found', + message: 'Organization not found', }); } - await deleteOrganizationMode(modeId); + if (route_model !== undefined) { + await assertOrganizationAutoWriteEnabled(ctx.user.id); + assertOrganizationAutoEligible(organization); + await ensureOrganizationAccess(ctx, organizationId, ['owner', 'billing_manager']); + } - await createAuditLog({ - action: 'organization.mode.delete', - actor_email: ctx.user.google_user_email, - actor_id: ctx.user.id, - actor_name: ctx.user.google_user_name, - message: `Deleted mode "${mode.name}" (slug: "${mode.slug}")`, - organization_id: mode.organization_id, + const routeAuditChanges: string[] = []; + await db.transaction(async tx => { + const settings = await mutateOrganizationSettings( + organizationId, + async lockedOrganization => { + const lockedMode = await getOrganizationModeById(organizationId, modeId, tx, true); + if (!lockedMode) { + throw new TRPCError({ + code: 'NOT_FOUND', + message: 'Mode not found', + }); + } + + const orgAutoModel = getOrganizationAutoSettings(lockedOrganization.settings); + let nextSettings = lockedOrganization.settings; + + const preserveBuiltInRoute = preserve_route && BUILT_IN_MODE_SLUGS.has(lockedMode.slug); + const hasExistingRoute = hasRoute(orgAutoModel.routes, lockedMode.slug); + + if (route_model !== undefined && !preserveBuiltInRoute) { + throw new TRPCError({ + code: 'BAD_REQUEST', + message: 'Route updates can only be preserved when reverting a built-in mode.', + }); + } + + if (route_model !== undefined) { + const previousRoute = + lockedOrganization.settings.org_auto_model?.routes[lockedMode.slug]; + nextSettings = await applyOrganizationAutoRouteChange( + lockedOrganization, + lockedMode.slug, + route_model, + ctx, + tx + ); + const nextRoute = nextSettings.org_auto_model?.routes[lockedMode.slug]; + if (previousRoute !== nextRoute) { + routeAuditChanges.push( + nextRoute + ? `Organization Auto route: "${lockedMode.slug}" ${previousRoute ? `"${previousRoute}" → ` : 'set to '}"${nextRoute}"` + : `Organization Auto route cleared: "${lockedMode.slug}"${previousRoute ? ` (was "${previousRoute}")` : ''}` + ); + } + } else if (hasExistingRoute) { + await ensureOrganizationAccess(ctx, organizationId, ['owner', 'billing_manager']); + + if (!preserveBuiltInRoute) { + routeAuditChanges.push( + `Organization Auto route removed: "${lockedMode.slug}" (was "${orgAutoModel.routes[lockedMode.slug]}")` + ); + const routes = { ...orgAutoModel.routes }; + delete routes[lockedMode.slug]; + nextSettings = { + ...lockedOrganization.settings, + org_auto_model: { + ...orgAutoModel, + routes, + }, + }; + } else { + routeAuditChanges.push( + `Organization Auto route preserved: "${lockedMode.slug}" (${orgAutoModel.routes[lockedMode.slug]})` + ); + } + } + + await deleteOrganizationMode(modeId, tx); + await createAuditLog({ + action: 'organization.mode.delete', + actor_email: ctx.user.google_user_email, + actor_id: ctx.user.id, + actor_name: ctx.user.google_user_name, + message: `Deleted mode "${lockedMode.name}" (slug: "${lockedMode.slug}")${routeAuditChanges.length > 0 ? `, ${routeAuditChanges.join(', ')}` : ''}`, + + organization_id: lockedMode.organization_id, + tx, + }); + return nextSettings; + }, + tx + ); + + return settings; }); return successResult(); 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 c14dcfd112..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,14 +10,33 @@ 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'; +jest.mock('@/lib/posthog-feature-flags', () => ({ + isReleaseToggleEnabled: jest.fn(async () => true), +})); + jest.mock('@/lib/ai-gateway/providers/openrouter', () => { return { getEnhancedOpenRouterModels: jest.fn(), + buildAutoModelCatalogEntry: jest.fn(model => ({ + id: model.id, + name: model.name, + created: 0, + description: model.description, + architecture: { input_modalities: ['text'], output_modalities: ['text'], tokenizer: 'test' }, + top_provider: { is_moderated: false }, + pricing: { prompt: '0', completion: '0' }, + context_length: 8192, + })), }; }); @@ -30,6 +49,19 @@ jest.mock('@/lib/ai-gateway/providers/openrouter/models-by-provider-index.server import { getEnhancedOpenRouterModels } from '@/lib/ai-gateway/providers/openrouter'; import { getProviderSlugsForModel } from '@/lib/ai-gateway/providers/openrouter/models-by-provider-index.server'; +function makeTestOpenRouterModel(id: string): OpenRouterModel { + return { + id, + name: id, + created: 0, + description: '', + architecture: { input_modalities: [], output_modalities: [], tokenizer: 'test' }, + top_provider: { is_moderated: false }, + pricing: { prompt: '0', completion: '0' }, + context_length: 8192, + }; +} + let owner: User; let member: User; let testOrganization: Organization; @@ -45,6 +77,15 @@ describe('organizations settings trpc router', () => { beforeEach(() => { mockedGetProviderSlugsForModel.mockReset(); mockedGetEnhancedOpenRouterModels.mockReset(); + mockedGetEnhancedOpenRouterModels.mockResolvedValue({ + data: [ + makeTestOpenRouterModel('gpt-4'), + makeTestOpenRouterModel('gpt-3.5-turbo'), + makeTestOpenRouterModel('openai/gpt-4o'), + makeTestOpenRouterModel('kilo-auto/balanced'), + makeTestOpenRouterModel('kilo-auto/frontier'), + ], + } satisfies OpenRouterModelsResponse); }); beforeAll(async () => { @@ -357,6 +398,35 @@ describe('organizations settings trpc router', () => { expect(result.data.map(model => model.id)).toEqual(['openai/gpt-4o']); }); + it('should include Organization Auto only for enabled enterprise organizations', async () => { + const openRouterModelsResponse = { + data: [makeOpenRouterModel('openai/gpt-4o')], + } satisfies OpenRouterModelsResponse; + + mockedGetEnhancedOpenRouterModels.mockResolvedValue(openRouterModelsResponse); + const organization = await createTestOrganization( + 'Organization Auto Catalog', + owner.id, + 0, + { + default_model: 'kilo-auto/org', + org_auto_model: { + routes: {}, + fallback_model: 'kilo-auto/balanced', + }, + }, + false + ); + await addUserToOrganization(organization.id, member.id, 'member'); + + const caller = await createCallerForUser(member.id); + const result = await caller.organizations.settings.listAvailableModels({ + organizationId: organization.id, + }); + + expect(result.data.map(model => model.id)).toEqual(['openai/gpt-4o', 'kilo-auto/org']); + }); + it('should return all models for a non-enterprise org even if access settings are set', async () => { const openRouterModelsResponse = { data: [ @@ -407,6 +477,24 @@ describe('organizations settings trpc router', () => { expect(updatedOrg?.settings?.default_model).toBe('gpt-4'); }); + it('preserves an exact catalog variant when setting the default model', async () => { + const caller = await createCallerForUser(owner.id); + const freshOrg = await createTestOrganization('Variant Default Org', owner.id, 0, {}, false); + mockedGetEnhancedOpenRouterModels.mockResolvedValue({ + data: [ + makeTestOpenRouterModel('openai/gpt-4o'), + makeTestOpenRouterModel('openai/gpt-4o:free'), + ], + } satisfies OpenRouterModelsResponse); + + const result = await caller.organizations.settings.updateDefaultModel({ + organizationId: freshOrg.id, + default_model: 'openai/gpt-4o:free', + }); + + expect(result.settings.default_model).toBe('openai/gpt-4o:free'); + }); + it('should reject default_model if it is in the deny list', async () => { const caller = await createCallerForUser(owner.id); @@ -447,6 +535,204 @@ describe('organizations settings trpc router', () => { }); }); + describe('Organization Auto procedures', () => { + it('enables Organization Auto and preserves its default route settings', async () => { + const caller = await createCallerForUser(owner.id); + + const result = await caller.organizations.settings.enableOrganizationAuto({ + organizationId: testOrganization.id, + }); + + expect(result.settings.default_model).toBe('kilo-auto/org'); + expect(result.settings.org_auto_model).toEqual({ + routes: {}, + fallback_model: 'kilo-auto/balanced', + }); + }); + + 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', + }); + 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 () => { + 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); + + await caller.organizations.settings.setOrganizationAutoRoute({ + organizationId: testOrganization.id, + mode_slug: 'code', + model_id: 'kilo-auto/frontier', + }); + + let updatedOrg = await getOrganizationById(testOrganization.id); + expect(updatedOrg?.settings.org_auto_model?.routes.code).toBe('kilo-auto/frontier'); + + await caller.organizations.settings.clearOrganizationAutoRoute({ + organizationId: testOrganization.id, + mode_slug: 'code', + }); + + updatedOrg = await getOrganizationById(testOrganization.id); + expect(updatedOrg?.settings.org_auto_model?.routes.code).toBeUndefined(); + }); + + it('requires a replacement model when disabling Organization Auto', async () => { + const caller = await createCallerForUser(owner.id); + + await caller.organizations.settings.enableOrganizationAuto({ + organizationId: testOrganization.id, + }); + const result = await caller.organizations.settings.disableOrganizationAuto({ + organizationId: testOrganization.id, + replacement_model: 'kilo-auto/balanced', + }); + + expect(result.settings.default_model).toBe('kilo-auto/balanced'); + expect(result.settings.org_auto_model).toEqual({ + routes: {}, + fallback_model: 'kilo-auto/balanced', + }); + }); + + it('rejects disableOrganizationAuto when Organization Auto is not enabled', async () => { + const caller = await createCallerForUser(owner.id); + + await expect( + caller.organizations.settings.disableOrganizationAuto({ + organizationId: testOrganization.id, + replacement_model: 'kilo-auto/balanced', + }) + ).rejects.toThrow('Organization Auto is not enabled for this organization.'); + }); + + it('does not allow updateDefaultModel to clear an active Organization Auto default', async () => { + const caller = await createCallerForUser(owner.id); + + await caller.organizations.settings.enableOrganizationAuto({ + organizationId: testOrganization.id, + }); + + await expect( + caller.organizations.settings.updateDefaultModel({ + organizationId: testOrganization.id, + default_model: null, + }) + ).rejects.toThrow('Disable Organization Auto through its dedicated controls.'); + }); + + it('preserves Organization Auto routes when unrelated settings change', async () => { + const caller = await createCallerForUser(owner.id); + + await caller.organizations.settings.enableOrganizationAuto({ + organizationId: testOrganization.id, + }); + await caller.organizations.settings.setOrganizationAutoRoute({ + organizationId: testOrganization.id, + mode_slug: 'code', + model_id: 'kilo-auto/frontier', + }); + await caller.organizations.settings.updateDataCollection({ + organizationId: testOrganization.id, + dataCollection: 'deny', + }); + + const updatedOrg = await getOrganizationById(testOrganization.id); + expect(updatedOrg?.settings.org_auto_model?.routes.code).toBe('kilo-auto/frontier'); + expect(updatedOrg?.settings.data_collection).toBe('deny'); + }); + + it('keeps Organization Auto enabled when access lists change', async () => { + const caller = await createCallerForUser(owner.id); + + await caller.organizations.settings.enableOrganizationAuto({ + organizationId: testOrganization.id, + }); + const result = await caller.organizations.settings.updateAllowLists({ + organizationId: testOrganization.id, + provider_allow_list: ['anthropic'], + }); + + expect(result.settings.default_model).toBe('kilo-auto/org'); + }); + }); + describe('updateMinimumBalanceAlert procedure', () => { it('should enable minimum balance alert with valid settings', 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 69a77b60ad..4d7413a200 100644 --- a/apps/web/src/routers/organizations/organization-settings-router.ts +++ b/apps/web/src/routers/organizations/organization-settings-router.ts @@ -1,4 +1,4 @@ -import { getOrganizationById, updateOrganizationSettings } from '@/lib/organizations/organizations'; +import { getOrganizationById, mutateOrganizationSettings } from '@/lib/organizations/organizations'; import type { OpenRouterModelsResponse, OrganizationSettings, @@ -17,6 +17,16 @@ import { createAllowPredicateFromRestrictions } from '@/lib/model-allow.server'; import { getAvailableModelsForOrganization } from '@/lib/organizations/organization-models'; import { getEffectiveModelRestrictions } from '@/lib/organizations/model-restrictions'; import { normalizeModelId } from '@/lib/ai-gateway/model-utils'; +import { isReleaseToggleEnabled } from '@/lib/posthog-feature-flags'; +import { db } from '@/lib/drizzle'; +import { ORG_AUTO_MODEL } from '@/lib/ai-gateway/auto-model'; +import { + DEFAULT_ORGANIZATION_AUTO_MODEL_SETTINGS, + ORGANIZATION_AUTO_MODEL_FLAG, + MAX_ORGANIZATION_AUTO_ROUTES, + isOrganizationAutoEligible, + validateOrganizationAutoTarget, +} from '@/lib/organizations/organization-auto-model'; /** * Allowlist of organization IDs that are allowed to modify experimental settings @@ -91,6 +101,15 @@ function createDefaultModelDiffMessage( return 'Updated default model'; } +function assertOrganizationAutoRouteCount(routes: Record): void { + if (Object.keys(routes).length > MAX_ORGANIZATION_AUTO_ROUTES) { + throw new TRPCError({ + code: 'BAD_REQUEST', + message: `Organization Auto supports at most ${MAX_ORGANIZATION_AUTO_ROUTES} routes`, + }); + } +} + const UpdateAllowListsInputSchema = OrganizationIdInputSchema.extend({ provider_allow_list: z.array(z.string()).optional(), model_deny_list: z.array(z.string()).optional(), @@ -104,10 +123,140 @@ function dedupeStrings(values: string[]): string[] { return [...new Set(values)]; } +// Organization Auto rollout is intentionally actor-scoped so flagged admins can +// configure eligible organizations before the broader rollout is enabled. +async function assertOrganizationAutoWriteEnabled(userId: string): Promise { + if ( + process.env.NODE_ENV !== 'development' && + !(await isReleaseToggleEnabled(ORGANIZATION_AUTO_MODEL_FLAG, userId)) + ) { + throw new TRPCError({ + code: 'FORBIDDEN', + message: 'Organization Auto configuration is not available', + }); + } +} + +function assertOrganizationAutoEligible( + organization: NonNullable>> +): void { + if (!isOrganizationAutoEligible(organization)) { + throw new TRPCError({ + code: 'FORBIDDEN', + message: 'Organization Auto is only available for Enterprise organizations.', + }); + } +} + +async function validateOrganizationDefaultModel( + organization: NonNullable>>, + defaultModel: string +): Promise { + const requestedDefaultModel = defaultModel.trim().toLowerCase(); + const normalizedDefaultModel = normalizeModelId(requestedDefaultModel); + if (!normalizedDefaultModel || normalizedDefaultModel.endsWith('/*')) { + throw new TRPCError({ + code: 'BAD_REQUEST', + message: `Default model '${defaultModel}' is not a concrete model identifier`, + }); + } + + if (normalizedDefaultModel === ORG_AUTO_MODEL.id) { + throw new TRPCError({ + code: 'BAD_REQUEST', + message: 'Enable Organization Auto through its dedicated controls.', + }); + } + + const isAllowed = createAllowPredicateFromRestrictions( + getEffectiveModelRestrictions(organization) + ); + if (!(await isAllowed(requestedDefaultModel))) { + throw new TRPCError({ + code: 'BAD_REQUEST', + message: `Default model '${defaultModel}' is not in the organization's allowed models list`, + }); + } + + return defaultModel.trim(); +} + +async function validateOrganizationDefaultReplacement( + organization: NonNullable>>, + replacementModel: string +): Promise { + const normalizedReplacementModel = normalizeModelId(replacementModel.trim().toLowerCase()); + if (normalizedReplacementModel === ORG_AUTO_MODEL.id) { + throw new TRPCError({ + code: 'BAD_REQUEST', + message: 'Choose a replacement model other than Organization Auto.', + }); + } + + const validatedReplacementModel = await validateOrganizationDefaultModel( + organization, + replacementModel + ); + let availableModels: Awaited>; + try { + availableModels = await getAvailableModelsForOrganization(organization.id); + } catch { + throw new TRPCError({ + code: 'BAD_REQUEST', + message: + 'Replacement default model could not be validated against the current model catalog.', + }); + } + const availableModel = availableModels?.data.find( + model => model.id.trim().toLowerCase() === validatedReplacementModel.toLowerCase() + ); + if (!availableModel) { + throw new TRPCError({ + code: 'BAD_REQUEST', + message: `Default model '${replacementModel}' is unavailable for this organization.`, + }); + } + + return availableModel.id; +} + const UpdateDefaultModelInputSchema = OrganizationIdInputSchema.extend({ default_model: z.string().or(z.null()), }); +const EnableOrganizationAutoInputSchema = OrganizationIdInputSchema; + +const DisableOrganizationAutoInputSchema = OrganizationIdInputSchema.extend({ + replacement_model: z.string().min(1), +}); + +const SetOrganizationAutoRouteInputSchema = OrganizationIdInputSchema.extend({ + mode_slug: z + .string() + .min(1) + .max(50) + .regex(/^[a-z0-9-]+$/, 'Mode slug must contain only lowercase letters, numbers, and hyphens'), + model_id: z.string().min(1).max(200), +}); + +const ClearOrganizationAutoRouteInputSchema = OrganizationIdInputSchema.extend({ + mode_slug: z + .string() + .min(1) + .max(50) + .regex(/^[a-z0-9-]+$/, 'Mode slug must contain only lowercase letters, numbers, and hyphens'), +}); + +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(), }); @@ -167,121 +316,577 @@ export const organizationsSettingsRouter = createTRPCRouter({ .output(SettingsResponseSchema) .mutation(async ({ input, ctx }) => { const { organizationId, provider_allow_list, model_deny_list } = input; - const existingOrg = await getOrganizationById(organizationId); if (!existingOrg) { + throw new TRPCError({ code: 'NOT_FOUND', message: 'Organization not found' }); + } + if (existingOrg.plan !== 'enterprise') { throw new TRPCError({ - code: 'NOT_FOUND', - message: 'Organization not found', + code: 'FORBIDDEN', + message: 'Model access configuration is not available for this organization.', }); } - // enterprise only feature + let previousSettings: OrganizationSettings | undefined; + const updatedSettings = await db.transaction(async tx => { + const settings = await mutateOrganizationSettings( + organizationId, + async organization => { + if (organization.plan !== 'enterprise') { + throw new TRPCError({ + code: 'FORBIDDEN', + message: 'Model access configuration is not available for this organization.', + }); + } + previousSettings = organization.settings; + const currentSettings = organization.settings || {}; + const settingsUpdate: OrganizationSettings = { ...currentSettings }; + if (provider_allow_list !== undefined) { + settingsUpdate.provider_allow_list = dedupeStrings(provider_allow_list); + } + if (model_deny_list !== undefined) { + settingsUpdate.model_deny_list = dedupeModels(model_deny_list); + } + if ( + (provider_allow_list !== undefined || model_deny_list !== undefined) && + currentSettings.default_model && + currentSettings.default_model !== ORG_AUTO_MODEL.id + ) { + const isAllowed = createAllowPredicateFromRestrictions({ + providerAllowList: settingsUpdate.provider_allow_list, + modelDenyList: settingsUpdate.model_deny_list ?? [], + }); + if (!(await isAllowed(currentSettings.default_model))) { + settingsUpdate.default_model = undefined; + } + } + return settingsUpdate; + }, + 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: createAccessListsDiffMessage(previousSettings, settings), + organization_id: organizationId, + tx, + }); + return settings; + }); + return { settings: updatedSettings }; + }), + + updateDefaultModel: organizationBillingMutationProcedure + .input(UpdateDefaultModelInputSchema) + .output(SettingsResponseSchema) + .mutation(async ({ input, ctx }) => { + const { organizationId, default_model } = input; + const existingOrg = await getOrganizationById(organizationId); + if (!existingOrg) { + throw new TRPCError({ code: 'NOT_FOUND', message: 'Organization not found' }); + } if (existingOrg.plan !== 'enterprise') { throw new TRPCError({ code: 'FORBIDDEN', message: 'Model access configuration is not available for this organization.', }); } + if (existingOrg.settings.default_model === ORG_AUTO_MODEL.id) { + throw new TRPCError({ + code: 'BAD_REQUEST', + message: 'Disable Organization Auto through its dedicated controls.', + }); + } - // Merge with existing settings - const currentSettings = existingOrg.settings || {}; - const settingsUpdate: OrganizationSettings = { - ...currentSettings, - }; + let validatedDefaultModel: string | undefined; + let previousSettings: OrganizationSettings | undefined; + const updatedSettings = await db.transaction(async tx => { + const settings = await mutateOrganizationSettings( + organizationId, + async organization => { + if (organization.plan !== 'enterprise') { + throw new TRPCError({ + code: 'FORBIDDEN', + message: 'Model access configuration is not available for this organization.', + }); + } + if (organization.settings.default_model === ORG_AUTO_MODEL.id) { + throw new TRPCError({ + code: 'BAD_REQUEST', + message: 'Disable Organization Auto through its dedicated controls.', + }); + } + if (default_model) { + validatedDefaultModel = await validateOrganizationDefaultModel( + organization, + default_model + ); + } else { + validatedDefaultModel = undefined; + } + previousSettings = organization.settings; + return { + ...organization.settings, + default_model: validatedDefaultModel || undefined, + }; + }, + 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: createDefaultModelDiffMessage(previousSettings, settings), + organization_id: organizationId, + tx, + }); + return settings; + }); + return { settings: updatedSettings }; + }), - if (provider_allow_list !== undefined) { - settingsUpdate.provider_allow_list = dedupeStrings(provider_allow_list); - } + enableOrganizationAuto: organizationBillingMutationProcedure + .input(EnableOrganizationAutoInputSchema) + .output(SettingsResponseSchema) + .mutation(async ({ input, ctx }) => { + const { organizationId } = input; + await assertOrganizationAutoWriteEnabled(ctx.user.id); - if (model_deny_list !== undefined) { - settingsUpdate.model_deny_list = dedupeModels(model_deny_list); + const existingOrg = await getOrganizationById(organizationId); + if (!existingOrg) { + throw new TRPCError({ code: 'NOT_FOUND', message: 'Organization not found' }); } + assertOrganizationAutoEligible(existingOrg); - // Check if default_model needs to be cleared when access lists change - if ( - (provider_allow_list !== undefined || model_deny_list !== undefined) && - currentSettings.default_model - ) { - const isAllowed = createAllowPredicateFromRestrictions({ - providerAllowList: settingsUpdate.provider_allow_list, - modelDenyList: settingsUpdate.model_deny_list ?? [], - }); + let didEnableChange = false; + const updatedSettings = await db.transaction(async tx => { + const settings = await mutateOrganizationSettings( + organizationId, + async organization => { + assertOrganizationAutoEligible(organization); + if ( + organization.settings.default_model === ORG_AUTO_MODEL.id && + organization.settings.org_auto_model + ) { + return organization.settings; + } + const existingOrgAutoModel = organization.settings.org_auto_model; + const freshOrgAutoModel = + existingOrgAutoModel ?? DEFAULT_ORGANIZATION_AUTO_MODEL_SETTINGS; + const seededRoutes = { ...freshOrgAutoModel.routes }; - if (!(await isAllowed(currentSettings.default_model))) { - // Clear default_model if it's no longer allowed - settingsUpdate.default_model = undefined; - } - } + assertOrganizationAutoRouteCount(seededRoutes); + for (const [slug, targetModelId] of Object.entries(seededRoutes)) { + const validation = await validateOrganizationAutoTarget(organization, targetModelId, { + dbClient: tx, + }); + if (validation.kind === 'error') { + throw new TRPCError({ + code: 'BAD_REQUEST', + message: `Cannot enable Organization Auto because route "${slug}" is invalid: ${validation.message}`, + }); + } + seededRoutes[slug] = validation.modelId; + } + const fallbackValidation = await validateOrganizationAutoTarget( + organization, + freshOrgAutoModel.fallback_model, + { dbClient: tx } + ); - const updatedSettings = await updateOrganizationSettings(organizationId, settingsUpdate); + if (fallbackValidation.kind === 'error') { + throw new TRPCError({ + code: 'BAD_REQUEST', + message: `Cannot enable Organization Auto because fallback model is invalid: ${fallbackValidation.message}`, + }); + } - 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: createAccessListsDiffMessage(existingOrg.settings, updatedSettings), - organization_id: organizationId, + const nextSettings = { + ...organization.settings, + default_model: ORG_AUTO_MODEL.id, + org_auto_model: { + ...freshOrgAutoModel, + routes: seededRoutes, + fallback_model: fallbackValidation.modelId, + }, + }; + if ( + organization.settings.default_model === ORG_AUTO_MODEL.id && + JSON.stringify(organization.settings.org_auto_model) === + JSON.stringify(nextSettings.org_auto_model) + ) { + return organization.settings; + } + didEnableChange = true; + return nextSettings; + }, + tx + ); + if (didEnableChange) { + 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: 'Enabled Organization Auto and set it as the organization default model.', + organization_id: organizationId, + tx, + }); + } + return settings; }); - return { - settings: updatedSettings, - }; + return { settings: updatedSettings }; }), - updateDefaultModel: organizationBillingMutationProcedure - .input(UpdateDefaultModelInputSchema) + disableOrganizationAuto: organizationBillingMutationProcedure + .input(DisableOrganizationAutoInputSchema) .output(SettingsResponseSchema) .mutation(async ({ input, ctx }) => { - const { organizationId, default_model } = input; + const { organizationId, replacement_model } = input; const existingOrg = await getOrganizationById(organizationId); if (!existingOrg) { + throw new TRPCError({ code: 'NOT_FOUND', message: 'Organization not found' }); + } + assertOrganizationAutoEligible(existingOrg); + if (existingOrg.settings.default_model !== ORG_AUTO_MODEL.id) { throw new TRPCError({ - code: 'NOT_FOUND', - message: 'Organization not found', + code: 'BAD_REQUEST', + message: 'Organization Auto is not enabled for this organization.', }); } - // enterprise only feature - if (existingOrg.plan !== 'enterprise') { - throw new TRPCError({ - code: 'FORBIDDEN', - message: 'Model access configuration is not available for this organization.', + let normalizedReplacementModel: string | undefined; + const updatedSettings = await db.transaction(async tx => { + const settings = await mutateOrganizationSettings( + organizationId, + async organization => { + assertOrganizationAutoEligible(organization); + if (organization.settings.default_model !== ORG_AUTO_MODEL.id) { + throw new TRPCError({ + code: 'BAD_REQUEST', + message: 'Organization Auto is not enabled for this organization.', + }); + } + normalizedReplacementModel = await validateOrganizationDefaultReplacement( + organization, + replacement_model + ); + return { + ...organization.settings, + default_model: normalizedReplacementModel, + }; + }, + 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: `Disabled Organization Auto and set replacement default model: ${normalizedReplacementModel}`, + organization_id: organizationId, + tx, }); + return settings; + }); + + return { settings: updatedSettings }; + }), + + setOrganizationAutoRoute: organizationBillingMutationProcedure + .input(SetOrganizationAutoRouteInputSchema) + .output(SettingsResponseSchema) + .mutation(async ({ input, ctx }) => { + const { organizationId, mode_slug, model_id } = input; + 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 isAllowed = createAllowPredicateFromRestrictions( - getEffectiveModelRestrictions(existingOrg) - ); + let validatedModelId: string | undefined; + let previousRoute: string | undefined; + const updatedSettings = await db.transaction(async tx => { + const settings = await mutateOrganizationSettings( + organizationId, + async organization => { + assertOrganizationAutoEligible(organization); + const validation = await validateOrganizationAutoTarget(organization, model_id, { + dbClient: tx, + }); + if (validation.kind === 'error') { + throw new TRPCError({ code: 'BAD_REQUEST', message: validation.message }); + } + validatedModelId = validation.modelId; + const orgAutoModel = + organization.settings.org_auto_model ?? DEFAULT_ORGANIZATION_AUTO_MODEL_SETTINGS; + previousRoute = orgAutoModel.routes[mode_slug]; + if (previousRoute === validation.modelId) { + return organization.settings; + } + const routes = { + ...orgAutoModel.routes, + [mode_slug]: validation.modelId, + }; + assertOrganizationAutoRouteCount(routes); + return { + ...organization.settings, + org_auto_model: { + ...orgAutoModel, + routes, + }, + }; + }, + tx + ); + if (previousRoute !== validatedModelId) { + 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: + previousRoute === undefined + ? `Set Organization Auto route for mode "${mode_slug}" to "${validatedModelId}"` + : `Updated Organization Auto route for mode "${mode_slug}": "${previousRoute}" → "${validatedModelId}"`, + organization_id: organizationId, + tx, + }); + } + return settings; + }); - if (default_model && !(await isAllowed(default_model))) { - throw new TRPCError({ - code: 'BAD_REQUEST', - message: `Default model '${default_model}' is not in the organization's allowed models list`, - }); + return { settings: updatedSettings }; + }), + + clearOrganizationAutoRoute: organizationBillingMutationProcedure + .input(ClearOrganizationAutoRouteInputSchema) + .output(SettingsResponseSchema) + .mutation(async ({ input, ctx }) => { + const { organizationId, mode_slug } = input; + await assertOrganizationAutoWriteEnabled(ctx.user.id); + + const existingOrg = await getOrganizationById(organizationId); + if (!existingOrg) { + throw new TRPCError({ code: 'NOT_FOUND', message: 'Organization not found' }); } + assertOrganizationAutoEligible(existingOrg); - // Merge with existing settings - const currentSettings = existingOrg.settings || {}; - const updatedSettings = await updateOrganizationSettings(organizationId, { - ...currentSettings, - default_model: default_model ? default_model : undefined, + let removedRoute: string | undefined; + const updatedSettings = await db.transaction(async tx => { + const settings = await mutateOrganizationSettings( + organizationId, + organization => { + assertOrganizationAutoEligible(organization); + const orgAutoModel = + organization.settings.org_auto_model ?? DEFAULT_ORGANIZATION_AUTO_MODEL_SETTINGS; + removedRoute = orgAutoModel.routes[mode_slug]; + if (removedRoute === undefined) { + return organization.settings; + } + const routes = { ...orgAutoModel.routes }; + delete routes[mode_slug]; + return { + ...organization.settings, + org_auto_model: { + ...orgAutoModel, + routes, + }, + }; + }, + tx + ); + if (removedRoute !== undefined) { + 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: `Cleared Organization Auto route for mode "${mode_slug}" (was "${removedRoute}")`, + organization_id: organizationId, + tx, + }); + } + return settings; }); - 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: createDefaultModelDiffMessage(existingOrg.settings, updatedSettings), - organization_id: organizationId, + return { settings: updatedSettings }; + }), + + setOrganizationAutoFallback: organizationBillingMutationProcedure + .input(SetOrganizationAutoFallbackInputSchema) + .output(SettingsResponseSchema) + .mutation(async ({ input, ctx }) => { + const { organizationId, model_id } = input; + await assertOrganizationAutoWriteEnabled(ctx.user.id); + + const existingOrg = await getOrganizationById(organizationId); + if (!existingOrg) { + throw new TRPCError({ code: 'NOT_FOUND', message: 'Organization not found' }); + } + assertOrganizationAutoEligible(existingOrg); + + let validatedModelId: string | undefined; + let previousFallback: string | undefined; + const updatedSettings = await db.transaction(async tx => { + const settings = await mutateOrganizationSettings( + organizationId, + async organization => { + assertOrganizationAutoEligible(organization); + const validation = await validateOrganizationAutoTarget(organization, model_id, { + dbClient: tx, + }); + if (validation.kind === 'error') { + throw new TRPCError({ code: 'BAD_REQUEST', message: validation.message }); + } + validatedModelId = validation.modelId; + const orgAutoModel = + organization.settings.org_auto_model ?? DEFAULT_ORGANIZATION_AUTO_MODEL_SETTINGS; + previousFallback = orgAutoModel.fallback_model; + if (previousFallback === validation.modelId) { + return organization.settings; + } + return { + ...organization.settings, + org_auto_model: { + ...orgAutoModel, + fallback_model: validation.modelId, + }, + }; + }, + tx + ); + if (previousFallback !== validatedModelId) { + 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: `Updated Organization Auto fallback model: "${previousFallback}" → "${validatedModelId}"`, + organization_id: organizationId, + tx, + }); + } + return settings; }); - return { - settings: updatedSettings, - }; + 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); + + 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 }; + } + + 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}` + : 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, + tx, + }); + return settings; + }); + + return { settings: updatedSettings }; }), updateDataCollection: organizationBillingMutationProcedure @@ -298,11 +903,16 @@ export const organizationsSettingsRouter = createTRPCRouter({ }); } - // Update the data collection setting - const updatedSettings = await updateOrganizationSettings(organizationId, { - ...existingOrg.settings, - data_collection: dataCollection, - }); + const updatedSettings = await db.transaction(async tx => + mutateOrganizationSettings( + organizationId, + organization => ({ + ...organization.settings, + data_collection: dataCollection, + }), + tx + ) + ); return { settings: updatedSettings, @@ -335,24 +945,34 @@ export const organizationsSettingsRouter = createTRPCRouter({ }); } - // Merge with existing settings - const currentSettings = existingOrg.settings || {}; - const updatedSettings = await updateOrganizationSettings(organizationId, { - ...currentSettings, - projects_ui_enabled, - }); + let previousSettings: OrganizationSettings | undefined; + const updatedSettings = await db.transaction(async tx => { + const settings = await mutateOrganizationSettings( + organizationId, + organization => { + previousSettings = organization.settings; + return { + ...organization.settings, + projects_ui_enabled, + }; + }, + tx + ); - // Create audit log if the value changed - if (currentSettings.projects_ui_enabled !== projects_ui_enabled) { - 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: `Projects UI: ${projects_ui_enabled ? 'enabled' : 'disabled'}`, - organization_id: organizationId, - }); - } + if (previousSettings?.projects_ui_enabled !== projects_ui_enabled) { + 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: `Projects UI: ${projects_ui_enabled ? 'enabled' : 'disabled'}`, + organization_id: organizationId, + tx, + }); + } + + return settings; + }); return { settings: updatedSettings, @@ -373,24 +993,34 @@ export const organizationsSettingsRouter = createTRPCRouter({ }); } - // Merge with existing settings - const currentSettings = existingOrg.settings || {}; - const updatedSettings = await updateOrganizationSettings(organizationId, { - ...currentSettings, - code_indexing_enabled, - }); + let previousSettings: OrganizationSettings | undefined; + const updatedSettings = await db.transaction(async tx => { + const settings = await mutateOrganizationSettings( + organizationId, + organization => { + previousSettings = organization.settings; + return { + ...organization.settings, + code_indexing_enabled, + }; + }, + tx + ); - // Create audit log if the value changed - if (currentSettings.code_indexing_enabled !== code_indexing_enabled) { - 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: `[Admin] Code indexing: ${code_indexing_enabled ? 'enabled' : 'disabled'}`, - organization_id: organizationId, - }); - } + if (previousSettings?.code_indexing_enabled !== code_indexing_enabled) { + 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: `[Admin] Code indexing: ${code_indexing_enabled ? 'enabled' : 'disabled'}`, + organization_id: organizationId, + tx, + }); + } + + return settings; + }); return { settings: updatedSettings, @@ -411,41 +1041,47 @@ export const organizationsSettingsRouter = createTRPCRouter({ }); } - const currentSettings = existingOrg.settings || {}; - let updatedSettings: OrganizationSettings; + let previousSettings: OrganizationSettings | undefined; + const updatedSettings = await db.transaction(async tx => { + const settings = await mutateOrganizationSettings( + organizationId, + organization => { + previousSettings = organization.settings; + if (enabled) { + return { + ...organization.settings, + minimum_balance, + minimum_balance_alert_email, + }; + } - if (enabled) { - updatedSettings = await updateOrganizationSettings(organizationId, { - ...currentSettings, - minimum_balance, - minimum_balance_alert_email, - }); - } else { - // Remove the fields when disabled - const { - minimum_balance: _mb, - minimum_balance_alert_email: _mbae, - ...rest - } = currentSettings; - updatedSettings = await updateOrganizationSettings(organizationId, rest); - } + const settings = { ...organization.settings }; + delete settings.minimum_balance; + delete settings.minimum_balance_alert_email; + return settings; + }, + tx + ); - // Create audit log - const wasEnabled = - currentSettings.minimum_balance !== undefined && - currentSettings.minimum_balance_alert_email !== undefined; - if (enabled !== wasEnabled || enabled) { - 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: enabled - ? `Minimum balance alert: enabled (threshold: $${minimum_balance}, emails: ${minimum_balance_alert_email?.join(', ')})` - : 'Minimum balance alert: disabled', - organization_id: organizationId, - }); - } + const wasEnabled = + previousSettings?.minimum_balance !== undefined && + previousSettings?.minimum_balance_alert_email !== undefined; + if (enabled !== wasEnabled || enabled) { + 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: enabled + ? `Minimum balance alert: enabled (threshold: $${minimum_balance}, emails: ${minimum_balance_alert_email?.join(', ')})` + : 'Minimum balance alert: disabled', + organization_id: organizationId, + tx, + }); + } + + return settings; + }); return { settings: updatedSettings, diff --git a/packages/db/src/schema-types.test.ts b/packages/db/src/schema-types.test.ts new file mode 100644 index 0000000000..8778acf2da --- /dev/null +++ b/packages/db/src/schema-types.test.ts @@ -0,0 +1,43 @@ +import { describe, expect, it } from '@jest/globals'; +import { OrganizationSettingsSchema } from './schema-types'; + +describe('OrganizationSettingsSchema org_auto_model', () => { + it('accepts bounded route maps and a fallback model', () => { + const result = OrganizationSettingsSchema.safeParse({ + org_auto_model: { + routes: { + code: 'anthropic/claude-sonnet-4.5', + plan: 'kilo-auto/frontier', + }, + fallback_model: 'kilo-auto/balanced', + }, + }); + + expect(result.success).toBe(true); + }); + + it('rejects Organization Auto self-targets', () => { + const result = OrganizationSettingsSchema.safeParse({ + org_auto_model: { + routes: { code: 'kilo-auto/org' }, + fallback_model: 'kilo-auto/balanced', + }, + }); + + expect(result.success).toBe(false); + }); + + it('rejects route maps with more than 100 routes', () => { + const routes = Object.fromEntries( + Array.from({ length: 101 }, (_, index) => [`mode-${index}`, 'kilo-auto/balanced']) + ); + const result = OrganizationSettingsSchema.safeParse({ + org_auto_model: { + routes, + fallback_model: 'kilo-auto/balanced', + }, + }); + + expect(result.success).toBe(false); + }); +}); diff --git a/packages/db/src/schema-types.ts b/packages/db/src/schema-types.ts index 7e77545a0a..9866dd8df7 100644 --- a/packages/db/src/schema-types.ts +++ b/packages/db/src/schema-types.ts @@ -715,12 +715,44 @@ export const OrganizationPlanSchema = z.enum(['teams', 'enterprise']); export type OrganizationPlan = z.infer; +const OrganizationAutoModelRouteSlugSchema = z + .string() + .min(1, 'Organization Auto route slug is required') + .max(50, 'Organization Auto route slug must be less than 50 characters') + .regex( + /^[a-z0-9-]+$/, + 'Organization Auto route slug must contain only lowercase letters, numbers, and hyphens' + ); + +const OrganizationAutoModelTargetSchema = z + .string() + .min(1, 'Organization Auto route target is required') + .max(200, 'Organization Auto route target must be less than 200 characters') + .refine(value => !value.endsWith('/*'), { + message: 'Organization Auto route target must be a concrete model identifier', + }) + .refine(value => value !== 'kilo-auto/org', { + message: 'Organization Auto cannot target itself', + }); + +export const OrganizationAutoModelSettingsSchema = z.object({ + routes: z + .record(OrganizationAutoModelRouteSlugSchema, OrganizationAutoModelTargetSchema) + .refine(routes => Object.keys(routes).length <= 100, { + message: 'Organization Auto supports at most 100 routes', + }), + fallback_model: OrganizationAutoModelTargetSchema, +}); + +export type OrganizationAutoModelSettings = z.infer; + const OrganizationSettingsSchema = z.object({ provider_allow_list: z.array(z.string()).optional(), model_deny_list: z.array(z.string()).optional(), default_model: z.string().optional(), + org_auto_model: OrganizationAutoModelSettingsSchema.optional(), data_collection: z.enum(['allow', 'deny']).nullable().optional(), // null means they were grandfathered in and so they have usage limits enabled enable_usage_limits: z.boolean().optional(), @@ -765,7 +797,6 @@ export const OrganizationModeConfigSchema = z.object({ description: z.string().optional(), customInstructions: z.string().optional(), groups: z.array(GroupEntrySchema), - defaultModel: z.string().min(1, 'Default model cannot be empty').optional(), }); export type OrganizationModeConfig = z.infer;