From 01e6eada6d33761e95f702b79b39f7c1e39953dc Mon Sep 17 00:00:00 2001 From: John Richmond <5629+jr@users.noreply.github.com> Date: Sun, 18 May 2025 08:44:33 -0700 Subject: [PATCH] Connect more of the Default Parameters and Model/Provider whitelist screen (#28) --- src/actions/defaultParameters.ts | 69 +-- src/actions/organizationSettings.ts | 130 +++++- .../DefaultParameters.tsx | 210 ++++++--- .../ProviderWhitelist.tsx | 434 +++++++++++------- src/db/index.ts | 4 + src/lib/server/auditLogs.ts | 20 +- src/schemas/index.ts | 22 +- 7 files changed, 621 insertions(+), 268 deletions(-) diff --git a/src/actions/defaultParameters.ts b/src/actions/defaultParameters.ts index dbb1d324a7..522dd337be 100644 --- a/src/actions/defaultParameters.ts +++ b/src/actions/defaultParameters.ts @@ -2,34 +2,33 @@ import { z } from 'zod'; -import { AuditLogTargetType } from '@/db/schema'; -import { createAuditLog } from '@/lib/server/auditLogs'; - +import { db } from '@/db'; +import { AuditLogTargetType, orgSettings } from '@/db/schema'; +import { insertAuditLog } from '@/lib/server/auditLogs'; import { handleError, isAuthSuccess, validateAuth, type ApiResponse, } from './apiUtils'; +import { sql } from 'drizzle-orm'; +import { ORGANIZATION_ALLOW_ALL } from '@/schemas'; const defaultParametersSchema = z.object({ experimentalPowerSteering: z.boolean().optional(), - terminalOutputLimit: z.number().int().nonnegative().optional(), - compressProgressBar: z.boolean().optional(), + terminalOutputLineLimit: z.number().int().nonnegative().optional(), + terminalCompressProgressBar: z.boolean().optional(), inheritEnvVars: z.boolean().optional(), - disableShellIntegration: z.boolean().optional(), - shellIntegrationTimeout: z.number().int().nonnegative().optional(), - commandDelay: z.number().int().nonnegative().optional(), - enablePowerShellCounter: z.boolean().optional(), - clearZshEol: z.boolean().optional(), - enableOhMyZsh: z.boolean().optional(), + terminalShellIntegrationDisabled: z.boolean().optional(), + terminalShellIntegrationTimeout: z.number().int().nonnegative().optional(), + terminalCommandDelay: z.number().int().nonnegative().optional(), + terminalZshClearEolMark: z.boolean().optional(), enablePowerlevel10k: z.boolean().optional(), - openTabsLimit: z.number().int().nonnegative().optional(), - workspaceFilesLimit: z.number().int().nonnegative().optional(), - showRooignoreFiles: z.boolean().optional(), - fileReadThreshold: z.number().int().optional(), - alwaysReadEntireFile: z.boolean().optional(), - enableAutoCheckpoints: z.boolean().optional(), + maxOpenTabsContext: z.number().int().nonnegative().optional(), + maxWorkspaceFiles: z.number().int().nonnegative().optional(), + showRooIgnoredFiles: z.boolean().optional(), + maxReadFileLine: z.number().int().gte(-1).optional(), + enableCheckpoints: z.boolean().optional(), useCustomTemperature: z.boolean().optional(), temperature: z.number().nonnegative().optional(), rateLimit: z.number().nonnegative().optional(), @@ -62,15 +61,33 @@ export async function updateDefaultParameters( const validatedData = result.data; - // TODO: Audit log description should contain more granular information on - // what changed. - await createAuditLog({ - userId: authResult.userId, - orgId: authResult.orgId, - targetType: AuditLogTargetType.DEFAULT_PARAMETERS, - targetId: 'default-parameters', - newValue: validatedData, - description: 'Updated default parameters', + await db.transaction(async (tx) => { + await tx + .insert(orgSettings) + .values({ + orgId: authResult.orgId, + version: 1, + defaultSettings: validatedData, + allowList: ORGANIZATION_ALLOW_ALL, + }) + .onConflictDoUpdate({ + target: orgSettings.orgId, + set: { + defaultSettings: validatedData, + version: sql`${orgSettings.version} + 1`, + }, + }); + + // TODO: consider trying to capture the changes more granularly, + // although that would prevent upsert + await insertAuditLog(tx, { + userId: authResult.userId, + orgId: authResult.orgId, + targetType: AuditLogTargetType.DEFAULT_PARAMETERS, + targetId: 'default-parameters', + newValue: validatedData, + description: 'Updated default parameters', + }); }); return { diff --git a/src/actions/organizationSettings.ts b/src/actions/organizationSettings.ts index 9dfe923589..2484df065b 100644 --- a/src/actions/organizationSettings.ts +++ b/src/actions/organizationSettings.ts @@ -1,13 +1,27 @@ 'use server'; -import { eq } from 'drizzle-orm'; import { auth } from '@clerk/nextjs/server'; - import { db } from '@/db'; -import { orgSettings, type OrgSettings } from '@/db/schema'; +import { AuditLogTargetType, orgSettings } from '@/db/schema'; +import { eq, sql } from 'drizzle-orm'; +import { + ORGANIZATION_ALLOW_ALL, + ORGANIZATION_DEFAULT, + type OrganizationSettings, + organizationAllowListSchema, + organizationDefaultSettingsSchema, +} from '@/schemas'; +import { z } from 'zod'; +import { + handleError, + isAuthSuccess, + validateAuth, + type ApiResponse, +} from './apiUtils'; +import { insertAuditLog } from '@/lib/server/auditLogs'; export async function getOrganizationSettings(): Promise< - OrgSettings | undefined + OrganizationSettings | undefined > { const { userId, orgId } = await auth(); @@ -25,5 +39,111 @@ export async function getOrganizationSettings(): Promise< .where(eq(orgSettings.orgId, orgId)) .limit(1); - return settings.length === 0 ? undefined : settings[0]; + return settings.length === 0 ? ORGANIZATION_DEFAULT : settings[0]; +} + +/** + * Schema for updating organization settings + */ +const updateOrganizationSchema = z + .object({ + defaultSettings: organizationDefaultSettingsSchema.optional(), + allowList: organizationAllowListSchema.optional(), + }) + .refine( + (data) => + data.defaultSettings !== undefined || data.allowList !== undefined, + { + message: 'At least one of defaultSettings or allowList must be provided', + }, + ); + +type UpdateOrganizationRequest = z.infer; + +export async function updateOrganization( + data: UpdateOrganizationRequest, +): Promise { + try { + const authResult = await validateAuth(); + if (!isAuthSuccess(authResult)) return authResult; + const { userId, orgId } = authResult; + + const result = updateOrganizationSchema.safeParse(data); + if (!result.success) { + return { + success: false, + error: 'Invalid request data', + }; + } + const validatedData = result.data; + + // Perform database update in a transaction + await db.transaction(async (tx) => { + // Get current settings or prepare for insert + const currentSettings = await tx + .select() + .from(orgSettings) + .where(eq(orgSettings.orgId, orgId)) + .limit(1); + + const isNewRecord = currentSettings.length === 0; + + const updateData: Partial = {}; + + if (validatedData.defaultSettings) { + updateData.defaultSettings = validatedData.defaultSettings; + } + + if (validatedData.allowList) { + updateData.allowList = validatedData.allowList; + } + + if (isNewRecord) { + await tx.insert(orgSettings).values({ + orgId, + version: 1, + defaultSettings: validatedData.defaultSettings || {}, + allowList: validatedData.allowList || ORGANIZATION_ALLOW_ALL, + }); + } else { + await tx + .update(orgSettings) + .set({ + ...updateData, + version: sql`${orgSettings.version} + 1`, + updatedAt: new Date(), + }) + .where(eq(orgSettings.orgId, orgId)); + } + + if (validatedData.defaultSettings) { + await insertAuditLog(tx, { + userId, + orgId, + targetType: AuditLogTargetType.DEFAULT_PARAMETERS, + targetId: 'organization-default-settings', + newValue: validatedData.defaultSettings, + description: 'Updated organization default settings', + }); + } + + if (validatedData.allowList) { + await insertAuditLog(tx, { + userId, + orgId, + targetType: AuditLogTargetType.PROVIDER_WHITELIST, + targetId: 'organization-allow-list', + newValue: validatedData.allowList, + description: 'Updated organization allow list', + }); + } + }); + + return { + success: true, + message: 'Organization settings updated successfully', + }; + } catch (error) { + return handleError(error, 'organization_settings'); + } } diff --git a/src/app/(authenticated)/org/[[...organization-profile]]/DefaultParameters.tsx b/src/app/(authenticated)/org/[[...organization-profile]]/DefaultParameters.tsx index b535715618..521c32593a 100644 --- a/src/app/(authenticated)/org/[[...organization-profile]]/DefaultParameters.tsx +++ b/src/app/(authenticated)/org/[[...organization-profile]]/DefaultParameters.tsx @@ -3,11 +3,18 @@ 'use client'; import { useTranslations } from 'next-intl'; -import { useState } from 'react'; +import { useState, useMemo, useRef, useEffect } from 'react'; import { useForm } from 'react-hook-form'; import { toast } from 'sonner'; +import { + useQuery, + useQueryClient, + type QueryClient, +} from '@tanstack/react-query'; import { updateDefaultParameters } from '@/actions/defaultParameters'; +import { getOrganizationSettings } from '@/actions/organizationSettings'; +import { type OrganizationSettings } from '@/schemas'; import { Button, Checkbox, @@ -24,22 +31,19 @@ import { CheckCheck, FlaskConical, SlidersHorizontal } from 'lucide-react'; type DefaultParamsFormValues = { experimentalPowerSteering: boolean; - terminalOutputLimit: number; - compressProgressBar: boolean; + terminalOutputLineLimit: number; + terminalCompressProgressBar: boolean; inheritEnvVars: boolean; - disableShellIntegration: boolean; - shellIntegrationTimeout: number; - commandDelay: number; - enablePowerShellCounter: boolean; - clearZshEol: boolean; - enableOhMyZsh: boolean; + terminalShellIntegrationDisabled: boolean; + terminalShellIntegrationTimeout: number; + terminalCommandDelay: number; + terminalZshClearEolMark: boolean; enablePowerlevel10k: boolean; - openTabsLimit: number; - workspaceFilesLimit: number; - showRooignoreFiles: boolean; - fileReadThreshold: number; - alwaysReadEntireFile: boolean; - enableAutoCheckpoints: boolean; + maxOpenTabsContext: number; + maxWorkspaceFiles: number; + showRooIgnoredFiles: boolean; + maxReadFileLine: number; + enableCheckpoints: boolean; useCustomTemperature: boolean; temperature: number; rateLimit: number; @@ -47,44 +51,102 @@ type DefaultParamsFormValues = { matchPrecision: number; }; -export const DefaultParameters = () => { +const mergeWithDefaultValues = ( + defaultValues: DefaultParamsFormValues, + orgSettings?: OrganizationSettings, +): DefaultParamsFormValues => { + if (!orgSettings || !orgSettings.defaultSettings) { + return defaultValues; + } + + return { + ...defaultValues, + ...orgSettings.defaultSettings, + }; +}; + +const defaultFormValues: DefaultParamsFormValues = { + experimentalPowerSteering: true, + terminalOutputLineLimit: 500, + terminalCompressProgressBar: true, + inheritEnvVars: true, + terminalShellIntegrationDisabled: false, + terminalShellIntegrationTimeout: 5, + terminalCommandDelay: 0, + terminalZshClearEolMark: true, + enablePowerlevel10k: false, + maxOpenTabsContext: 20, + maxWorkspaceFiles: 200, + showRooIgnoredFiles: true, + maxReadFileLine: 500, + enableCheckpoints: true, + useCustomTemperature: true, + temperature: 0, + rateLimit: 0, + enableEditingThroughDiffs: true, + matchPrecision: 100, +} as const; + +type ParametersFormProps = { + orgSettings: OrganizationSettings; + queryClient: QueryClient; +}; + +const ParametersForm = ({ orgSettings, queryClient }: ParametersFormProps) => { const t = useTranslations('ProviderWhitelist'); const [isSaving, setIsSaving] = useState(false); + const [readEntireFile, setReadEntireFile] = useState(false); + + const mergedValues = useMemo( + () => mergeWithDefaultValues(defaultFormValues, orgSettings), + [orgSettings], + ); + + const previousMaxReadFileLine = useRef( + mergedValues.maxReadFileLine === -1 ? 500 : mergedValues.maxReadFileLine, + ); const form = useForm({ - defaultValues: { - experimentalPowerSteering: true, - terminalOutputLimit: 500, - compressProgressBar: true, - inheritEnvVars: true, - disableShellIntegration: false, - shellIntegrationTimeout: 5, - commandDelay: 0, - enablePowerShellCounter: false, - clearZshEol: true, - enableOhMyZsh: false, - enablePowerlevel10k: false, - openTabsLimit: 20, - workspaceFilesLimit: 200, - showRooignoreFiles: true, - fileReadThreshold: 500, - alwaysReadEntireFile: false, - enableAutoCheckpoints: true, - useCustomTemperature: true, - temperature: 0, - rateLimit: 0, - enableEditingThroughDiffs: true, - matchPrecision: 100, - }, + defaultValues: mergedValues, }); + const maxReadFileLineValue = form.watch('maxReadFileLine'); + + useEffect(() => { + previousMaxReadFileLine.current = + mergedValues.maxReadFileLine === -1 ? 500 : mergedValues.maxReadFileLine; + }, [mergedValues]); + + useEffect(() => { + setReadEntireFile(maxReadFileLineValue === -1); + + if (maxReadFileLineValue !== -1) { + previousMaxReadFileLine.current = maxReadFileLineValue; + } + }, [maxReadFileLineValue, form]); + + const handleReadEntireFileChange = (checked: boolean) => { + if (checked) { + const currentValue = form.getValues('maxReadFileLine'); + if (currentValue !== -1) { + previousMaxReadFileLine.current = currentValue; + } + form.setValue('maxReadFileLine', -1); + } else { + form.setValue('maxReadFileLine', previousMaxReadFileLine.current); + } + setReadEntireFile(checked); + }; + const onSubmit = async (data: DefaultParamsFormValues) => { setIsSaving(true); try { const result = await updateDefaultParameters(data); if (result.success) { + queryClient.invalidateQueries({ queryKey: ['organizationSettings'] }); + toast('Settings saved', { description: 'Default parameters have been updated successfully.', }); @@ -124,7 +186,7 @@ export const DefaultParameters = () => {
( @@ -172,7 +234,7 @@ export const DefaultParameters = () => {
(
@@ -202,7 +264,7 @@ export const DefaultParameters = () => { (
@@ -232,7 +294,7 @@ export const DefaultParameters = () => { ( @@ -261,7 +323,7 @@ export const DefaultParameters = () => {
( @@ -274,27 +336,23 @@ export const DefaultParameters = () => { field.onChange(Number(e.target.value)) } className="w-32" + disabled={readEntireFile} + value={readEntireFile ? '' : field.value} /> )} /> lines - ( - - - - - Always read entire file - - )} - /> + + + + + Always read entire file +
Roo reads this number of lines when the model omits @@ -482,7 +540,7 @@ export const DefaultParameters = () => { (
@@ -512,7 +570,7 @@ export const DefaultParameters = () => { ( @@ -570,7 +628,7 @@ export const DefaultParameters = () => { ( @@ -596,7 +654,7 @@ export const DefaultParameters = () => { (
@@ -630,7 +688,7 @@ export const DefaultParameters = () => { (
@@ -668,7 +726,7 @@ export const DefaultParameters = () => { ( @@ -746,3 +804,23 @@ export const DefaultParameters = () => { ); }; + +export const DefaultParameters = () => { + const queryClient = useQueryClient(); + + const { data: orgSettings } = useQuery({ + queryKey: ['organizationSettings'], + queryFn: getOrganizationSettings, + }); + + if (!orgSettings) { + return ( +
+
⟳
+ Loading settings... +
+ ); + } + + return ; +}; diff --git a/src/app/(authenticated)/org/[[...organization-profile]]/ProviderWhitelist.tsx b/src/app/(authenticated)/org/[[...organization-profile]]/ProviderWhitelist.tsx index b0c626e979..626dd97580 100644 --- a/src/app/(authenticated)/org/[[...organization-profile]]/ProviderWhitelist.tsx +++ b/src/app/(authenticated)/org/[[...organization-profile]]/ProviderWhitelist.tsx @@ -2,211 +2,274 @@ import { useTranslations } from 'next-intl'; import { useState } from 'react'; +import { useQuery, useQueryClient } from '@tanstack/react-query'; import { - updateAllowAllProviders, - updateModelStatus, - updateProviderStatus, -} from '@/actions/providerWhitelist'; -import { Badge, Checkbox, Label } from '@/components/ui'; + getOrganizationSettings, + updateOrganization, +} from '@/actions/organizationSettings'; +import { Badge, Button, Checkbox, Label } from '@/components/ui'; import { toast } from 'sonner'; +import { + ORGANIZATION_ALLOW_ALL, + type OrganizationAllowList, + type OrganizationSettings, +} from '@/schemas'; +import { type ProviderName } from '@roo-code/types'; -const initialProviders = [ +type ProviderSetting = { + allowAll: boolean; + models?: string[]; +}; + +// Define a type for providers record +type ProvidersRecord = Record; + +// Provider metadata without state information +const providerMetadata: { + id: ProviderName; + name: string; + models: { + id: string; + name: string; + }[]; +}[] = [ { - id: 'openai', + id: 'openai-native', name: 'OpenAI', - enabled: true, models: [ - { id: 'gpt-4', name: 'GPT-4', enabled: true }, - { id: 'gpt-3.5-turbo', name: 'GPT-3.5 Turbo', enabled: true }, - { id: 'gpt-4o', name: 'GPT-4o', enabled: false }, + { id: 'gpt-4', name: 'GPT-4' }, + { id: 'gpt-3.5-turbo', name: 'GPT-3.5 Turbo' }, + { id: 'gpt-4o', name: 'GPT-4o' }, ], }, { id: 'anthropic', name: 'Anthropic', - enabled: true, models: [ - { id: 'claude-3-opus', name: 'Claude 3 Opus', enabled: true }, - { id: 'claude-3-sonnet', name: 'Claude 3 Sonnet', enabled: true }, - { id: 'claude-3-haiku', name: 'Claude 3 Haiku', enabled: false }, + { id: 'claude-3-opus-20240229', name: 'Claude 3 Opus' }, + { id: 'claude-3-7-sonnet-20250219', name: 'Claude 3 Sonnet' }, + { + id: 'claude-3-7-sonnet-20250219:thinking', + name: 'Claude 3 Sonnet Thinking', + }, + { id: 'claude-3-haiku-20240307', name: 'Claude 3 Haiku' }, + ], + }, + { + id: 'openrouter', + name: 'OpenRouter', + models: [ + { id: 'anthropic/claude-3-opus', name: 'Claude 3 Opus' }, + { id: 'anthropic/claude-3.7-sonnet', name: 'Claude 3.7 Sonnet' }, + { + id: 'anthropic/claude-3.7-sonnet:thinking', + name: 'Claude 3.7 Sonnet Thinking', + }, + { id: 'anthropic/claude-3.5-haiku', name: 'Claude 3.5 Haiku' }, + { id: 'openai/gpt-4', name: 'GPT-4' }, + { id: 'openai/gpt-3.5-turbo', name: 'GPT-3.5 Turbo' }, + { id: 'openai/gpt-4o', name: 'GPT-4o' }, + { + id: 'google/gemini-2.5-flash-preview', + name: 'Gemini 2.5 Flash Preview', + }, ], }, { id: 'mistral', name: 'Mistral AI', - enabled: false, models: [ - { id: 'mistral-large', name: 'Mistral Large', enabled: false }, - { id: 'mistral-medium', name: 'Mistral Medium', enabled: false }, - { id: 'mistral-small', name: 'Mistral Small', enabled: false }, - ], - }, - { - id: 'cohere', - name: 'Cohere', - enabled: false, - models: [ - { id: 'command-r', name: 'Command R', enabled: false }, - { id: 'command-r-plus', name: 'Command R+', enabled: false }, + { id: 'mistral-large-latest', name: 'Mistral Large' }, + { id: 'mistral-small-latest', name: 'Mistral Small' }, ], }, ]; -export const ProviderWhitelist = () => { +type ProviderWhitelistFormProps = { + orgSettings: OrganizationSettings; + queryClient: ReturnType; +}; + +const ProviderWhitelistForm = ({ + orgSettings, + queryClient, +}: ProviderWhitelistFormProps) => { const t = useTranslations('ProviderWhitelist'); - const [providers, setProviders] = useState(initialProviders); - const [policyVersion] = useState(1); - const [allowAllProviders, setAllowAllProviders] = useState(true); - const [isUpdating, setIsUpdating] = useState(false); + const [allowList, setAllowList] = useState( + orgSettings?.allowList || ORGANIZATION_ALLOW_ALL, + ); - // Common error handling function for API calls. - const handleApiError = (error: unknown, errorMessage: string) => { - console.error(errorMessage, error); - toast.error(errorMessage); - setIsUpdating(false); + const [hasChanges, setHasChanges] = useState(false); + const [isSaving, setIsSaving] = useState(false); + + const isAllowAllChecked = () => allowList.allowAll; + + const isProviderChecked = (providerId: ProviderName) => { + if (allowList.allowAll) return true; + return allowList.providers[providerId]?.allowAll || false; }; - const toggleAllowAllProviders = async () => { - const previousAllowAllProviders = allowAllProviders; - const previousProviders = [...providers]; - const newAllowAllProviders = !allowAllProviders; - setAllowAllProviders(newAllowAllProviders); + const isModelChecked = (providerId: ProviderName, modelId: string) => { + if (allowList.allowAll) return true; + const provider = allowList.providers[providerId]; + if (!provider) return false; + if (provider.allowAll) return true; + return provider.models?.includes(modelId) || false; + }; - let updatedProviders = providers; + const toggleAllowAll = () => { + const newAllowAll = !allowList.allowAll; - if (newAllowAllProviders) { - // Enable all providers and their models/ - updatedProviders = providers.map((provider) => ({ - ...provider, - enabled: true, - models: provider.models.map((model) => ({ ...model, enabled: true })), - })); + if (newAllowAll) { + setAllowList({ allowAll: true, providers: {} }); + } else { + const newProviders: Partial = {}; - setProviders(updatedProviders); - } - - // Call API to update backend with the changed allow all providers setting. - setIsUpdating(true); - - try { - const result = await updateAllowAllProviders({ - allowAllProviders: newAllowAllProviders, - policyVersion, + providerMetadata.forEach((provider) => { + newProviders[provider.id] = { allowAll: true }; }); - if (!result.success) { - throw new Error(result.error || 'Failed to update setting'); - } - } catch (error) { - setAllowAllProviders(previousAllowAllProviders); - setProviders(previousProviders); - handleApiError( - error, - 'Failed to update allow all providers setting. Please try again.', - ); - } finally { - setIsUpdating(false); + setAllowList({ + allowAll: false, + providers: newProviders, + }); } + + setHasChanges(true); }; - const toggleProvider = async (providerId: string) => { - if (allowAllProviders) { - return; + const toggleProvider = (providerId: ProviderName) => { + if (allowList.allowAll) { + const newProviders: Partial = {}; + + providerMetadata.forEach((provider) => { + if (provider.id !== providerId) { + newProviders[provider.id] = { allowAll: true }; + } + }); + + setAllowList({ + allowAll: false, + providers: newProviders, + }); + } else { + const newProviders = { ...allowList.providers }; + const providersRecord = newProviders; + const isCurrentlyEnabled = isProviderChecked(providerId); + + if (isCurrentlyEnabled) { + if (providersRecord[providerId]) { + providersRecord[providerId] = { + ...providersRecord[providerId], + allowAll: false, + models: [], + }; + } + } else { + providersRecord[providerId] = { allowAll: true }; + } + + setAllowList({ ...allowList, providers: newProviders }); } - // Find the provider to get its current state. - const provider = providers.find((p) => p.id === providerId); - if (!provider) return; + setHasChanges(true); + }; - const previousProviders = [...providers]; - const newEnabled = !provider.enabled; + const toggleModel = (providerId: ProviderName, modelId: string) => { + if (allowList.allowAll || isProviderChecked(providerId)) { + const newProviders = { ...allowList.providers }; + const providersRecord = newProviders; - const updatedProviders = providers.map((provider) => { - if (provider.id === providerId) { - return { + if (allowList.allowAll) { + providerMetadata.forEach((provider) => { + providersRecord[provider.id] = { allowAll: true }; + }); + + const provider = providerMetadata.find((p) => p.id === providerId); + if (provider) { + const models = provider.models + .filter((m) => m.id !== modelId) + .map((m) => m.id); + + providersRecord[providerId] = { allowAll: false, models }; + } + + setAllowList({ allowAll: false, providers: newProviders }); + } else { + const provider = providerMetadata.find((p) => p.id === providerId); + if (provider) { + const models = provider.models + .filter((m) => m.id !== modelId) + .map((m) => m.id); + + providersRecord[providerId] = { allowAll: false, models }; + setAllowList({ ...allowList, providers: newProviders }); + } + } + } else { + const newProviders = { ...allowList.providers }; + const providersRecord = newProviders; + const provider = providersRecord[providerId] || { + allowAll: false, + models: [], + }; + const models = provider.models || []; + + const isCurrentlyEnabled = models.includes(modelId); + + if (isCurrentlyEnabled) { + providersRecord[providerId] = { ...provider, - enabled: newEnabled, - // If provider is disabled, disable all its models - models: provider.models.map((model) => ({ - ...model, - enabled: newEnabled ? model.enabled : false, - })), + models: models.filter((m: string) => m !== modelId), + }; + } else { + providersRecord[providerId] = { + ...provider, + models: [...models, modelId], }; - } - return provider; - }); - setProviders(updatedProviders); + const allModels = + providerMetadata + .find((p) => p.id === providerId) + ?.models.map((m) => m.id) || []; + const enabledModels = providersRecord[providerId].models || []; - // Call API to update backend with the changed provider status - setIsUpdating(true); - try { - const result = await updateProviderStatus({ - providerId, - enabled: newEnabled, - policyVersion, - }); - if (!result.success) { - throw new Error(result.error || 'Failed to update provider status'); + if ( + allModels.length === enabledModels.length && + allModels.every((m) => enabledModels.includes(m)) + ) { + providersRecord[providerId] = { allowAll: true }; + } } - } catch (error) { - setProviders(previousProviders); - handleApiError( - error, - 'Failed to update provider status. Please try again.', - ); - } finally { - setIsUpdating(false); + + setAllowList({ ...allowList, providers: newProviders }); } + + setHasChanges(true); }; - const toggleModel = async (providerId: string, modelId: string) => { - if (allowAllProviders) { - return; - } - - // Find the model to get its current state. - const provider = providers.find((p) => p.id === providerId); - if (!provider) return; - - const model = provider.models.find((m) => m.id === modelId); - if (!model) return; - - const previousProviders = [...providers]; - const newEnabled = !model.enabled; - - const updatedProviders = providers.map((provider) => { - return provider.id === providerId - ? { - ...provider, - models: provider.models.map((model) => - model.id === modelId ? { ...model, enabled: newEnabled } : model, - ), - } - : provider; - }); - - setProviders(updatedProviders); - setIsUpdating(true); - + const saveChanges = async () => { + setIsSaving(true); try { - const result = await updateModelStatus({ - providerId, - modelId, - enabled: newEnabled, - policyVersion, + const result = await updateOrganization({ + allowList: allowList, }); if (!result.success) { - throw new Error(result.error || 'Failed to update model status'); + throw new Error(result.error || 'Failed to update settings'); } + + queryClient.invalidateQueries({ queryKey: ['organizationSettings'] }); + + toast.success('Provider settings saved successfully'); + setHasChanges(false); } catch (error) { - setProviders(previousProviders); - handleApiError(error, 'Failed to update model status. Please try again.'); + console.error('Failed to save provider settings', error); + toast.error('Failed to save provider settings'); } finally { - setIsUpdating(false); + setIsSaving(false); } }; @@ -227,8 +290,9 @@ export const ProviderWhitelist = () => {
-
- {providers.map((provider) => ( +
+ {providerMetadata.map((provider) => (
toggleProvider(provider.id)} />
); }; + +export const ProviderWhitelist = () => { + const queryClient = useQueryClient(); + + const { data: orgSettings } = useQuery({ + queryKey: ['organizationSettings'], + queryFn: getOrganizationSettings, + }); + + if (!orgSettings) { + return ( +
+
⟳
+ Loading settings... +
+ ); + } + + return ( + + ); +}; diff --git a/src/db/index.ts b/src/db/index.ts index e010c8734a..0f17699ad0 100644 --- a/src/db/index.ts +++ b/src/db/index.ts @@ -26,4 +26,8 @@ const disconnect = async () => { await pool.end(); }; +export type DB_OR_TX = + | typeof db + | Parameters[0]>[0]; + export { db, testDb, disconnect }; diff --git a/src/lib/server/auditLogs.ts b/src/lib/server/auditLogs.ts index 029d71d906..1ec49f5026 100644 --- a/src/lib/server/auditLogs.ts +++ b/src/lib/server/auditLogs.ts @@ -1,12 +1,22 @@ -import { db } from '@/db'; +import { db, type DB_OR_TX } from '@/db'; import { auditLogs, type CreateAuditLog } from '@/db/schema'; import { logger } from '@/lib/server/logger'; -export async function createAuditLog(values: CreateAuditLog) { +export async function insertAuditLog( + db: DB_OR_TX, + values: CreateAuditLog, +): Promise { + await db.insert(auditLogs).values(values); + const { userId, orgId, targetType } = values; + logger.info({ userId, orgId, targetType }); +} + +export async function createAuditLog(values: CreateAuditLog): Promise<{ + success: boolean; + error?: string | Record; +}> { try { - await db.insert(auditLogs).values(values); - const { userId, orgId, targetType } = values; - logger.info({ userId, orgId, targetType }); + await insertAuditLog(db, values); return { success: true }; } catch (e) { const error = diff --git a/src/schemas/index.ts b/src/schemas/index.ts index d7de12938c..dc4314e9bc 100644 --- a/src/schemas/index.ts +++ b/src/schemas/index.ts @@ -2,6 +2,7 @@ * TimePeriod */ +import { providerNames } from '@roo-code/types'; import { z } from 'zod'; export const timePeriods = [7, 30, 90] as const; @@ -11,6 +12,7 @@ export type TimePeriod = (typeof timePeriods)[number]; export const organizationAllowListSchema = z.object({ allowAll: z.boolean(), providers: z.record( + z.enum(providerNames), z.object({ allowAll: z.boolean(), models: z.array(z.string()).optional(), @@ -29,11 +31,17 @@ export const ORGANIZATION_ALLOW_ALL: OrganizationAllowList = { export const organizationDefaultSettingsSchema = z.object({ enableCheckpoints: z.boolean().optional(), - maxOpenTabsContext: z.number().optional(), - maxWorkspaceFiles: z.number().optional(), - showRooIgnoredFiles: z.boolean().optional(), - maxReadFileLine: z.number().optional(), fuzzyMatchThreshold: z.number().optional(), + maxOpenTabsContext: z.number().int().nonnegative().optional(), + maxReadFileLine: z.number().int().gte(-1).optional(), + maxWorkspaceFiles: z.number().int().nonnegative().optional(), + showRooIgnoredFiles: z.boolean().optional(), + terminalCommandDelay: z.number().int().nonnegative().optional(), + terminalCompressProgressBar: z.boolean().optional(), + terminalOutputLineLimit: z.number().int().nonnegative().optional(), + terminalShellIntegrationDisabled: z.boolean().optional(), + terminalShellIntegrationTimeout: z.number().int().nonnegative().optional(), + terminalZshClearEolMark: z.boolean().optional(), }); export type OrganizationDefaultSettings = z.infer< @@ -47,3 +55,9 @@ export const organizationSettingsSchema = z.object({ }); export type OrganizationSettings = z.infer; + +export const ORGANIZATION_DEFAULT: OrganizationSettings = { + version: 0, + defaultSettings: {}, + allowList: ORGANIZATION_ALLOW_ALL, +} as const;