Skip to content

Commit 864622c

Browse files
icecrasher321icecrasher321
authored
fix(ratelimits): enterprise and team checks should be pooled limit (#1255)
* fix(ratelimits): enterprise and team checks should be pooled limit" * fix * fix dynamic imports * fix tests" ;
1 parent 8668622 commit 864622c

File tree

11 files changed

+6192
-81
lines changed

11 files changed

+6192
-81
lines changed

‎apps/sim/app/api/schedules/execute/route.ts‎

Lines changed: 7 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import { NextResponse } from 'next/server'
44
import { v4 as uuidv4 } from 'uuid'
55
import { z } from 'zod'
66
import { checkServerSideUsageLimits } from '@/lib/billing'
7+
import { getHighestPrioritySubscription } from '@/lib/billing/core/subscription'
78
import { getPersonalAndWorkspaceEnv } from '@/lib/environment/utils'
89
import { createLogger } from '@/lib/logs/console/logger'
910
import { LoggingSession } from '@/lib/logs/execution/logging-session'
@@ -18,7 +19,7 @@ import { decryptSecret } from '@/lib/utils'
1819
import { loadWorkflowFromNormalizedTables } from '@/lib/workflows/db-helpers'
1920
import { updateWorkflowRunCounts } from '@/lib/workflows/utils'
2021
import { db } from '@/db'
21-
import { subscription, userStats, workflow, workflowSchedule } from '@/db/schema'
22+
import { userStats, workflow, workflowSchedule } from '@/db/schema'
2223
import { Executor } from '@/executor'
2324
import { Serializer } from '@/serializer'
2425
import { RateLimiter } from '@/services/queue'
@@ -108,19 +109,15 @@ export async function GET() {
108109
continue
109110
}
110111

111-
// Check rate limits for scheduled execution
112-
const [subscriptionRecord] = await db
113-
.select({ plan: subscription.plan })
114-
.from(subscription)
115-
.where(eq(subscription.referenceId, workflowRecord.userId))
116-
.limit(1)
112+
// Check rate limits for scheduled execution (checks both personal and org subscriptions)
113+
const userSubscription = await getHighestPrioritySubscription(workflowRecord.userId)
117114

118-
const subscriptionPlan = (subscriptionRecord?.plan || 'free') as SubscriptionPlan
115+
const subscriptionPlan = (userSubscription?.plan || 'free') as SubscriptionPlan
119116

120117
const rateLimiter = new RateLimiter()
121-
const rateLimitCheck = await rateLimiter.checkRateLimit(
118+
const rateLimitCheck = await rateLimiter.checkRateLimitWithSubscription(
122119
workflowRecord.userId,
123-
subscriptionPlan,
120+
userSubscription,
124121
'schedule',
125122
false // schedules are always sync
126123
)

‎apps/sim/app/api/users/me/rate-limit/route.ts‎

Lines changed: 8 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,11 @@
11
import { eq } from 'drizzle-orm'
22
import { type NextRequest, NextResponse } from 'next/server'
33
import { getSession } from '@/lib/auth'
4+
import { getHighestPrioritySubscription } from '@/lib/billing/core/subscription'
45
import { createLogger } from '@/lib/logs/console/logger'
56
import { createErrorResponse } from '@/app/api/workflows/utils'
67
import { db } from '@/db'
7-
import { apiKey as apiKeyTable, subscription } from '@/db/schema'
8+
import { apiKey as apiKeyTable } from '@/db/schema'
89
import { RateLimiter } from '@/services/queue'
910

1011
const logger = createLogger('RateLimitAPI')
@@ -33,31 +34,22 @@ export async function GET(request: NextRequest) {
3334
return createErrorResponse('Authentication required', 401)
3435
}
3536

36-
const [subscriptionRecord] = await db
37-
.select({ plan: subscription.plan })
38-
.from(subscription)
39-
.where(eq(subscription.referenceId, authenticatedUserId))
40-
.limit(1)
41-
42-
const subscriptionPlan = (subscriptionRecord?.plan || 'free') as
43-
| 'free'
44-
| 'pro'
45-
| 'team'
46-
| 'enterprise'
37+
// Get user subscription (checks both personal and org subscriptions)
38+
const userSubscription = await getHighestPrioritySubscription(authenticatedUserId)
4739

4840
const rateLimiter = new RateLimiter()
4941
const isApiAuth = !session?.user?.id
5042
const triggerType = isApiAuth ? 'api' : 'manual'
5143

52-
const syncStatus = await rateLimiter.getRateLimitStatus(
44+
const syncStatus = await rateLimiter.getRateLimitStatusWithSubscription(
5345
authenticatedUserId,
54-
subscriptionPlan,
46+
userSubscription,
5547
triggerType,
5648
false
5749
)
58-
const asyncStatus = await rateLimiter.getRateLimitStatus(
50+
const asyncStatus = await rateLimiter.getRateLimitStatusWithSubscription(
5951
authenticatedUserId,
60-
subscriptionPlan,
52+
userSubscription,
6153
triggerType,
6254
true
6355
)

‎apps/sim/app/api/webhooks/trigger/[path]/route.ts‎

Lines changed: 7 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ import { tasks } from '@trigger.dev/sdk'
22
import { and, eq } from 'drizzle-orm'
33
import { type NextRequest, NextResponse } from 'next/server'
44
import { checkServerSideUsageLimits } from '@/lib/billing'
5+
import { getHighestPrioritySubscription } from '@/lib/billing/core/subscription'
56
import { env, isTruthy } from '@/lib/env'
67
import { createLogger } from '@/lib/logs/console/logger'
78
import {
@@ -11,7 +12,7 @@ import {
1112
} from '@/lib/webhooks/utils'
1213
import { executeWebhookJob } from '@/background/webhook-execution'
1314
import { db } from '@/db'
14-
import { subscription, webhook, workflow } from '@/db/schema'
15+
import { webhook, workflow } from '@/db/schema'
1516
import { RateLimiter } from '@/services/queue'
1617
import type { SubscriptionPlan } from '@/services/queue/types'
1718

@@ -249,21 +250,17 @@ export async function POST(
249250
// --- PHASE 3: Rate limiting for webhook execution ---
250251
let isEnterprise = false
251252
try {
252-
// Get user subscription for rate limiting
253-
const [subscriptionRecord] = await db
254-
.select({ plan: subscription.plan })
255-
.from(subscription)
256-
.where(eq(subscription.referenceId, foundWorkflow.userId))
257-
.limit(1)
253+
// Get user subscription for rate limiting (checks both personal and org subscriptions)
254+
const userSubscription = await getHighestPrioritySubscription(foundWorkflow.userId)
258255

259-
const subscriptionPlan = (subscriptionRecord?.plan || 'free') as SubscriptionPlan
256+
const subscriptionPlan = (userSubscription?.plan || 'free') as SubscriptionPlan
260257
isEnterprise = subscriptionPlan === 'enterprise'
261258

262259
// Check async rate limits (webhooks are processed asynchronously)
263260
const rateLimiter = new RateLimiter()
264-
const rateLimitCheck = await rateLimiter.checkRateLimit(
261+
const rateLimitCheck = await rateLimiter.checkRateLimitWithSubscription(
265262
foundWorkflow.userId,
266-
subscriptionPlan,
263+
userSubscription,
267264
'webhook',
268265
true // isAsync = true for webhook execution
269266
)

‎apps/sim/app/api/workflows/[id]/execute/route.test.ts‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,11 @@ describe('Workflow Execution API Route', () => {
4646
remaining: 10,
4747
resetAt: new Date(),
4848
}),
49+
checkRateLimitWithSubscription: vi.fn().mockResolvedValue({
50+
allowed: true,
51+
remaining: 10,
52+
resetAt: new Date(),
53+
}),
4954
})),
5055
RateLimitError: class RateLimitError extends Error {
5156
constructor(
@@ -66,6 +71,13 @@ describe('Workflow Execution API Route', () => {
6671
}),
6772
}))
6873

74+
vi.doMock('@/lib/billing/core/subscription', () => ({
75+
getHighestPrioritySubscription: vi.fn().mockResolvedValue({
76+
plan: 'free',
77+
referenceId: 'user-id',
78+
}),
79+
}))
80+
6981
vi.doMock('@/db/schema', () => ({
7082
subscription: {
7183
plan: 'plan',

‎apps/sim/app/api/workflows/[id]/execute/route.ts‎

Lines changed: 14 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import { v4 as uuidv4 } from 'uuid'
55
import { z } from 'zod'
66
import { getSession } from '@/lib/auth'
77
import { checkServerSideUsageLimits } from '@/lib/billing'
8+
import { getHighestPrioritySubscription } from '@/lib/billing/core/subscription'
89
import { getPersonalAndWorkspaceEnv } from '@/lib/environment/utils'
910
import { createLogger } from '@/lib/logs/console/logger'
1011
import { LoggingSession } from '@/lib/logs/execution/logging-session'
@@ -19,7 +20,7 @@ import {
1920
import { validateWorkflowAccess } from '@/app/api/workflows/middleware'
2021
import { createErrorResponse, createSuccessResponse } from '@/app/api/workflows/utils'
2122
import { db } from '@/db'
22-
import { subscription, userStats } from '@/db/schema'
23+
import { userStats } from '@/db/schema'
2324
import { Executor } from '@/executor'
2425
import { Serializer } from '@/serializer'
2526
import {
@@ -374,19 +375,15 @@ export async function GET(request: NextRequest, { params }: { params: Promise<{
374375
try {
375376
// Check rate limits BEFORE entering queue for GET requests
376377
if (triggerType === 'api') {
377-
// Get user subscription
378-
const [subscriptionRecord] = await db
379-
.select({ plan: subscription.plan })
380-
.from(subscription)
381-
.where(eq(subscription.referenceId, validation.workflow.userId))
382-
.limit(1)
378+
// Get user subscription (checks both personal and org subscriptions)
379+
const userSubscription = await getHighestPrioritySubscription(validation.workflow.userId)
383380

384-
const subscriptionPlan = (subscriptionRecord?.plan || 'free') as SubscriptionPlan
381+
const subscriptionPlan = (userSubscription?.plan || 'free') as SubscriptionPlan
385382

386383
const rateLimiter = new RateLimiter()
387-
const rateLimitCheck = await rateLimiter.checkRateLimit(
384+
const rateLimitCheck = await rateLimiter.checkRateLimitWithSubscription(
388385
validation.workflow.userId,
389-
subscriptionPlan,
386+
userSubscription,
390387
triggerType,
391388
false // isAsync = false for sync calls
392389
)
@@ -505,20 +502,17 @@ export async function POST(
505502
return createErrorResponse('Authentication required', 401)
506503
}
507504

508-
const [subscriptionRecord] = await db
509-
.select({ plan: subscription.plan })
510-
.from(subscription)
511-
.where(eq(subscription.referenceId, authenticatedUserId))
512-
.limit(1)
505+
// Get user subscription (checks both personal and org subscriptions)
506+
const userSubscription = await getHighestPrioritySubscription(authenticatedUserId)
513507

514-
const subscriptionPlan = (subscriptionRecord?.plan || 'free') as SubscriptionPlan
508+
const subscriptionPlan = (userSubscription?.plan || 'free') as SubscriptionPlan
515509

516510
if (isAsync) {
517511
try {
518512
const rateLimiter = new RateLimiter()
519-
const rateLimitCheck = await rateLimiter.checkRateLimit(
513+
const rateLimitCheck = await rateLimiter.checkRateLimitWithSubscription(
520514
authenticatedUserId,
521-
subscriptionPlan,
515+
userSubscription,
522516
'api',
523517
true // isAsync = true
524518
)
@@ -580,9 +574,9 @@ export async function POST(
580574

581575
try {
582576
const rateLimiter = new RateLimiter()
583-
const rateLimitCheck = await rateLimiter.checkRateLimit(
577+
const rateLimitCheck = await rateLimiter.checkRateLimitWithSubscription(
584578
authenticatedUserId,
585-
subscriptionPlan,
579+
userSubscription,
586580
triggerType,
587581
false // isAsync = false for sync calls
588582
)

‎apps/sim/db/migrations/0084_even_lockheed.sql‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
ALTER TABLE "user_rate_limits" RENAME COLUMN "user_id" TO "reference_id";--> statement-breakpoint
2+
ALTER TABLE "user_rate_limits" DROP CONSTRAINT "user_rate_limits_user_id_user_id_fk";

0 commit comments

Comments
 (0)