Bugfix/Rate limit updating (#3311)

fix rate limit updating

packages/server/src/controllers/chatflows/index.ts

@@ -2,7 +2,7 @@ import { NextFunction, Request, Response } from 'express'
 import { StatusCodes } from 'http-status-codes'
 import apiKeyService from '../../services/apikey'
 import { ChatFlow } from '../../database/entities/ChatFlow'
-import { createRateLimiter } from '../../utils/rateLimit'
+import { updateRateLimiter } from '../../utils/rateLimit'
 import { InternalFlowiseError } from '../../errors/internalFlowiseError'
 import { ChatflowType } from '../../Interface'
 import chatflowsService from '../../services/chatflows'
@@ -130,7 +130,7 @@ const updateChatflow = async (req: Request, res: Response, next: NextFunction) =
         Object.assign(updateChatFlow, body)
 
         updateChatFlow.id = chatflow.id
-        createRateLimiter(updateChatFlow)
+        updateRateLimiter(updateChatFlow)
 
         const apiResponse = await chatflowsService.updateChatflow(chatflow, updateChatFlow)
         return res.json(apiResponse)

packages/server/src/utils/rateLimit.ts

@@ -21,6 +21,12 @@ async function addRateLimiter(id: string, duration: number, limit: number, messa
     }
 }
 
+function removeRateLimit(id: string) {
+    if (rateLimiters[id]) {
+        delete rateLimiters[id]
+    }
+}
+
 export function getRateLimiter(req: Request, res: Response, next: NextFunction) {
     const id = req.params.id
     if (!rateLimiters[id]) return next()
@@ -28,21 +34,22 @@ export function getRateLimiter(req: Request, res: Response, next: NextFunction)
     return idRateLimiter(req, res, next)
 }
 
-export async function createRateLimiter(chatFlow: IChatFlow) {
+export async function updateRateLimiter(chatFlow: IChatFlow) {
     if (!chatFlow.apiConfig) return
     const apiConfig = JSON.parse(chatFlow.apiConfig)
 
-    const rateLimit: { limitDuration: number; limitMax: number; limitMsg: string } = apiConfig.rateLimit
+    const rateLimit: { limitDuration: number; limitMax: number; limitMsg: string; status?: boolean } = apiConfig.rateLimit
     if (!rateLimit) return
 
-    const { limitDuration, limitMax, limitMsg } = rateLimit
-    if (limitMax && limitDuration && limitMsg) await addRateLimiter(chatFlow.id, limitDuration, limitMax, limitMsg)
+    const { limitDuration, limitMax, limitMsg, status } = rateLimit
+    if (status === false) removeRateLimit(chatFlow.id)
+    else if (limitMax && limitDuration && limitMsg) await addRateLimiter(chatFlow.id, limitDuration, limitMax, limitMsg)
 }
 
 export async function initializeRateLimiter(chatFlowPool: IChatFlow[]) {
     await Promise.all(
         chatFlowPool.map(async (chatFlow) => {
-            await createRateLimiter(chatFlow)
+            await updateRateLimiter(chatFlow)
         })
     )
 }

packages/ui/src/ui-component/extended/RateLimit.jsx

@@ -8,6 +8,7 @@ import { Box, Typography, Button, OutlinedInput } from '@mui/material'
 // Project import
 import { StyledButton } from '@/ui-component/button/StyledButton'
 import { TooltipWithParser } from '@/ui-component/tooltip/TooltipWithParser'
+import { SwitchInput } from '@/ui-component/switch/Switch'
 
 // Icons
 import { IconX } from '@tabler/icons-react'
@@ -29,29 +30,47 @@ const RateLimit = () => {
     const enqueueSnackbar = (...args) => dispatch(enqueueSnackbarAction(...args))
     const closeSnackbar = (...args) => dispatch(closeSnackbarAction(...args))
 
+    const [rateLimitStatus, setRateLimitStatus] = useState(false)
     const [limitMax, setLimitMax] = useState(apiConfig?.rateLimit?.limitMax ?? '')
     const [limitDuration, setLimitDuration] = useState(apiConfig?.rateLimit?.limitDuration ?? '')
     const [limitMsg, setLimitMsg] = useState(apiConfig?.rateLimit?.limitMsg ?? '')
 
     const formatObj = () => {
         const obj = {
-            rateLimit: {}
+            rateLimit: { status: rateLimitStatus }
         }
+
+        if (rateLimitStatus) {
             const rateLimitValuesBoolean = [!limitMax, !limitDuration, !limitMsg]
             const rateLimitFilledValues = rateLimitValuesBoolean.filter((value) => value === false)
             if (rateLimitFilledValues.length >= 1 && rateLimitFilledValues.length <= 2) {
                 throw new Error('Need to fill all rate limit input fields')
             } else if (rateLimitFilledValues.length === 3) {
                 obj.rateLimit = {
+                    ...obj.rateLimit,
                     limitMax,
                     limitDuration,
                     limitMsg
                 }
             }
+        }
 
         return obj
     }
 
+    const handleChange = (value) => {
+        setRateLimitStatus(value)
+    }
+
+    const checkDisabled = () => {
+        if (rateLimitStatus) {
+            if (limitMax === '' || limitDuration === '' || limitMsg === '') {
+                return true
+            }
+        }
+        return false
+    }
+
     const onSave = async () => {
         try {
             const saveResp = await chatflowsApi.updateChatflow(chatflowid, {
@@ -139,11 +158,20 @@ const RateLimit = () => {
                     }
                 />
             </Typography>
+            <SwitchInput label='Enable Rate Limit' onChange={handleChange} value={rateLimitStatus} />
+            {rateLimitStatus && (
+                <>
                     {textField(limitMax, 'limitMax', 'Message Limit per Duration', 'number', '5')}
                     {textField(limitDuration, 'limitDuration', 'Duration in Second', 'number', '60')}
                     {textField(limitMsg, 'limitMsg', 'Limit Message', 'string', 'You have reached the quota')}
-
-            <StyledButton style={{ marginBottom: 10, marginTop: 10 }} variant='contained' onClick={() => onSave()}>
+                </>
+            )}
+            <StyledButton
+                disabled={checkDisabled()}
+                style={{ marginBottom: 10, marginTop: 10 }}
+                variant='contained'
+                onClick={() => onSave()}
+            >
                 Save Changes
             </StyledButton>
         </>