diff --git a/apps/sim/app/api/cron/cleanup-soft-deletes/route.ts b/apps/sim/app/api/cron/cleanup-soft-deletes/route.ts index 1df6df035e3..520db937568 100644 --- a/apps/sim/app/api/cron/cleanup-soft-deletes/route.ts +++ b/apps/sim/app/api/cron/cleanup-soft-deletes/route.ts @@ -1,18 +1,36 @@ import { createLogger } from '@sim/logger' import { type NextRequest, NextResponse } from 'next/server' +import { softDeletesCleanupContract } from '@/lib/api/contracts/cleanup' +import { parseRequest } from '@/lib/api/server/validation' import { verifyCronAuth } from '@/lib/auth/internal' -import { dispatchCleanupJobs } from '@/lib/billing/cleanup-dispatcher' +import { dispatchBoundedCleanup, dispatchCleanupJobs } from '@/lib/billing/cleanup-dispatcher' import { withRouteHandler } from '@/lib/core/utils/with-route-handler' export const dynamic = 'force-dynamic' const logger = createLogger('SoftDeleteCleanupAPI') +/** Cron-secret maintenance protocol is global; workspace principal authorization does not apply. */ export const GET = withRouteHandler(async (request: NextRequest) => { try { const authError = verifyCronAuth(request, 'soft-delete cleanup') if (authError) return authError + const parsed = await parseRequest( + softDeletesCleanupContract, + request, + {}, + { + rejectDuplicateQueryValues: true, + rejectBlankQueryValues: true, + } + ) + if (!parsed.success) return parsed.response + if (parsed.data.query) { + const result = await dispatchBoundedCleanup('cleanup-soft-deletes', parsed.data.query) + return NextResponse.json(result, { status: 202 }) + } + const result = await dispatchCleanupJobs('cleanup-soft-deletes') logger.info('Soft-delete cleanup jobs dispatched', result) diff --git a/apps/sim/app/api/logs/cleanup/route.test.ts b/apps/sim/app/api/logs/cleanup/route.test.ts new file mode 100644 index 00000000000..2841acf3a0f --- /dev/null +++ b/apps/sim/app/api/logs/cleanup/route.test.ts @@ -0,0 +1,72 @@ +import { createMockRequest } from '@sim/testing' +import { beforeEach, describe, expect, it, vi } from 'vitest' + +const { auth, bounded, scheduled } = vi.hoisted(() => ({ + auth: vi.fn(), + bounded: vi.fn(), + scheduled: vi.fn(), +})) +vi.mock('@/lib/auth/internal', () => ({ verifyCronAuth: auth })) +vi.mock('@/lib/billing/cleanup-dispatcher', () => ({ + dispatchBoundedCleanup: bounded, + dispatchCleanupJobs: scheduled, +})) + +import { GET as softDeletes } from '@/app/api/cron/cleanup-soft-deletes/route' +import { GET as logs } from '@/app/api/logs/cleanup/route' + +for (const [path, GET, type, limit] of [ + ['/api/logs/cleanup', logs, 'cleanup-logs', 'workflowLogs'], + ['/api/cron/cleanup-soft-deletes', softDeletes, 'cleanup-soft-deletes', 'workflows'], +] as const) { + describe(path, () => { + beforeEach(() => { + vi.clearAllMocks() + auth.mockReturnValue(null) + bounded.mockResolvedValue({ triggered: true, runId: 'run-one', limits: { [limit]: 2 } }) + scheduled.mockResolvedValue({ + jobIds: ['batch-one'], + jobCount: 1, + chunkCount: 2, + workspaceCount: 3, + }) + }) + const request = (query = '') => + createMockRequest('GET', undefined, {}, `http://localhost:3000${path}${query}`) + it('authenticates before parsing invalid limits', async () => { + auth.mockReturnValue(new Response(null, { status: 401 })) + expect((await GET(request('?unknown=1'))).status).toBe(401) + expect(bounded).not.toHaveBeenCalled() + expect(scheduled).not.toHaveBeenCalled() + }) + it('keeps no-parameter scheduled dispatch unchanged', async () => { + const response = await GET(request()) + expect(response.status).toBe(200) + expect(scheduled).toHaveBeenCalledWith(type) + expect(bounded).not.toHaveBeenCalled() + }) + it('accepts one bounded run', async () => { + const response = await GET(request(`?${limit}=2`)) + expect(response.status).toBe(202) + expect(bounded).toHaveBeenCalledWith(type, { [limit]: 2 }) + expect(await response.json()).toEqual({ + triggered: true, + runId: 'run-one', + limits: { [limit]: 2 }, + }) + expect(scheduled).not.toHaveBeenCalled() + }) + it.each(['?dryRun=true', '?unknown=1', '?batchSize=3', `?${limit}=2&${limit}=3`])( + 'rejects invalid query %s', + async (query) => { + expect((await GET(request(query))).status).toBe(400) + expect(bounded).not.toHaveBeenCalled() + expect(scheduled).not.toHaveBeenCalled() + } + ) + it('reports a dispatch failure', async () => { + bounded.mockRejectedValue(new Error('Trigger unavailable')) + expect((await GET(request(`?${limit}=2`))).status).toBe(500) + }) + }) +} diff --git a/apps/sim/app/api/logs/cleanup/route.ts b/apps/sim/app/api/logs/cleanup/route.ts index 7891a763bc6..7b8d4259cb4 100644 --- a/apps/sim/app/api/logs/cleanup/route.ts +++ b/apps/sim/app/api/logs/cleanup/route.ts @@ -1,18 +1,36 @@ import { createLogger } from '@sim/logger' import { type NextRequest, NextResponse } from 'next/server' +import { logsCleanupContract } from '@/lib/api/contracts/cleanup' +import { parseRequest } from '@/lib/api/server/validation' import { verifyCronAuth } from '@/lib/auth/internal' -import { dispatchCleanupJobs } from '@/lib/billing/cleanup-dispatcher' +import { dispatchBoundedCleanup, dispatchCleanupJobs } from '@/lib/billing/cleanup-dispatcher' import { withRouteHandler } from '@/lib/core/utils/with-route-handler' export const dynamic = 'force-dynamic' const logger = createLogger('LogsCleanupAPI') +/** Cron-secret maintenance protocol is global; workspace principal authorization does not apply. */ export const GET = withRouteHandler(async (request: NextRequest) => { try { const authError = verifyCronAuth(request, 'logs cleanup') if (authError) return authError + const parsed = await parseRequest( + logsCleanupContract, + request, + {}, + { + rejectDuplicateQueryValues: true, + rejectBlankQueryValues: true, + } + ) + if (!parsed.success) return parsed.response + if (parsed.data.query) { + const result = await dispatchBoundedCleanup('cleanup-logs', parsed.data.query) + return NextResponse.json(result, { status: 202 }) + } + const result = await dispatchCleanupJobs('cleanup-logs') logger.info('Log cleanup jobs dispatched', result) diff --git a/apps/sim/background/cleanup-logs.test.ts b/apps/sim/background/cleanup-logs.test.ts index 2b1adb3cdc5..feb631202fa 100644 --- a/apps/sim/background/cleanup-logs.test.ts +++ b/apps/sim/background/cleanup-logs.test.ts @@ -45,9 +45,12 @@ const { mockTask: vi.fn((config: unknown) => config), })) -vi.mock('@trigger.dev/sdk', () => ({ task: mockTask })) +vi.mock('@trigger.dev/sdk', () => ({ task: mockTask, queue: vi.fn((config) => config) })) + +vi.mock('@/lib/billing/cleanup-dispatcher', () => ({ runCleanupWithLimits: vi.fn() })) vi.mock('@/lib/cleanup/batch-delete', () => ({ + consumeRowBudget: vi.fn(), batchDeleteByWorkspaceAndTimestamp: mockBatchDeleteByWorkspaceAndTimestamp, chunkedBatchDelete: mockChunkedBatchDelete, })) @@ -199,7 +202,7 @@ describe('cleanup logs worker', () => { it('caps Trigger.dev concurrency for log cleanup tasks', () => { expect(cleanupLogsTask).toMatchObject({ - queue: { concurrencyLimit: 2 }, + queue: { name: 'retention-cleanup', concurrencyLimit: 1 }, }) }) }) diff --git a/apps/sim/background/cleanup-logs.ts b/apps/sim/background/cleanup-logs.ts index 2d560ef74c0..e46127a8e6a 100644 --- a/apps/sim/background/cleanup-logs.ts +++ b/apps/sim/background/cleanup-logs.ts @@ -12,12 +12,16 @@ import { createLogger } from '@sim/logger' import { chunkArray } from '@sim/utils/helpers' import { task } from '@trigger.dev/sdk' import { and, asc, eq, inArray, isNull, lt, notInArray, or, sql } from 'drizzle-orm' -import type { CleanupJobPayload } from '@/lib/billing/cleanup-dispatcher' +import { type CleanupJobPayload, runCleanupWithLimits } from '@/lib/billing/cleanup-dispatcher' import { batchDeleteByWorkspaceAndTimestamp, chunkedBatchDelete, + consumeRowBudget, + type RowBudget, type TableCleanupResult, } from '@/lib/cleanup/batch-delete' +import type { CleanupBudgets, LimitedCleanupPayload } from '@/lib/cleanup/limits' +import { retentionCleanupQueue } from '@/lib/cleanup/queue' import { LIVE_PAUSED_REFERENCE_STATUSES, markLargeValuesDeleted, @@ -43,7 +47,6 @@ const WORKFLOW_LOG_CLEANUP_BATCH_SIZE = 500 const WORKFLOW_LOG_CLEANUP_MAX_BATCHES = 50 const WORKFLOW_LOG_CLEANUP_ROW_LIMIT = WORKFLOW_LOG_CLEANUP_BATCH_SIZE * WORKFLOW_LOG_CLEANUP_MAX_BATCHES -const LOG_CLEANUP_CONCURRENCY_LIMIT = 2 const LARGE_VALUE_CLEANUP_BATCH_SIZE = 500 const LARGE_VALUE_CLEANUP_TOTAL_KEY_LIMIT = 5_000 const LARGE_VALUE_CLEANUP_GRACE_HOURS = 7 * 24 @@ -135,7 +138,8 @@ async function deleteLargeValueKeys(keys: string[]): Promise<{ deleted: number; async function cleanupLargeExecutionValues( workspaceIds: string[], retentionDate: Date, - label: string + label: string, + budget?: RowBudget ): Promise { const stats: LargeValueCleanupStats = { largeValuesTotal: 0, @@ -151,10 +155,11 @@ async function cleanupLargeExecutionValues( let attempted = 0 for (const chunkIds of workspaceChunks) { - while (attempted < LARGE_VALUE_CLEANUP_TOTAL_KEY_LIMIT) { + while (attempted < LARGE_VALUE_CLEANUP_TOTAL_KEY_LIMIT && budget?.remaining !== 0) { const limit = Math.min( LARGE_VALUE_CLEANUP_BATCH_SIZE, - LARGE_VALUE_CLEANUP_TOTAL_KEY_LIMIT - attempted + LARGE_VALUE_CLEANUP_TOTAL_KEY_LIMIT - attempted, + budget?.remaining ?? LARGE_VALUE_CLEANUP_BATCH_SIZE ) const rows = await cleanupDb .select({ key: executionLargeValues.key }) @@ -176,19 +181,21 @@ async function cleanupLargeExecutionValues( if (rows.length === 0) break + consumeRowBudget(budget, rows.length) const keys = rows.map((row) => row.key) stats.largeValuesTotal += keys.length attempted += keys.length const result = await deleteLargeValueKeys(keys) stats.largeValuesDeleted += result.deleted stats.largeValuesDeleteFailed += result.failed + if (budget && result.failed) throw new Error('Large value cleanup failed') if (result.deleted === 0) { break } } - if (attempted >= LARGE_VALUE_CLEANUP_TOTAL_KEY_LIMIT) break + if (attempted >= LARGE_VALUE_CLEANUP_TOTAL_KEY_LIMIT || budget?.remaining === 0) break } logger.info( @@ -201,7 +208,8 @@ async function cleanupLargeExecutionValues( async function cleanupLegacyLargeExecutionValues( workspaceIds: string[], retentionDate: Date, - label: string + label: string, + budget?: RowBudget ): Promise { const stats: LargeValueCleanupStats = { largeValuesTotal: 0, @@ -217,10 +225,11 @@ async function cleanupLegacyLargeExecutionValues( let attempted = 0 for (const chunkIds of workspaceChunks) { - while (attempted < LARGE_VALUE_CLEANUP_TOTAL_KEY_LIMIT) { + while (attempted < LARGE_VALUE_CLEANUP_TOTAL_KEY_LIMIT && budget?.remaining !== 0) { const limit = Math.min( LARGE_VALUE_CLEANUP_BATCH_SIZE, - LARGE_VALUE_CLEANUP_TOTAL_KEY_LIMIT - attempted + LARGE_VALUE_CLEANUP_TOTAL_KEY_LIMIT - attempted, + budget?.remaining ?? LARGE_VALUE_CLEANUP_BATCH_SIZE ) const rows = await cleanupDb .select({ key: workspaceFiles.key }) @@ -329,19 +338,21 @@ async function cleanupLegacyLargeExecutionValues( if (rows.length === 0) break + consumeRowBudget(budget, rows.length) const keys = rows.map((row) => row.key) stats.largeValuesTotal += keys.length attempted += keys.length const result = await deleteLargeValueKeys(keys) stats.largeValuesDeleted += result.deleted stats.largeValuesDeleteFailed += result.failed + if (budget && result.failed) throw new Error('Large value cleanup failed') if (result.deleted === 0) { break } } - if (attempted >= LARGE_VALUE_CLEANUP_TOTAL_KEY_LIMIT) break + if (attempted >= LARGE_VALUE_CLEANUP_TOTAL_KEY_LIMIT || budget?.remaining === 0) break } logger.info( @@ -351,7 +362,11 @@ async function cleanupLegacyLargeExecutionValues( return stats } -async function cleanupLargeValueMetadata(workspaceIds: string[], label: string): Promise { +async function cleanupLargeValueMetadata( + workspaceIds: string[], + label: string, + budgets?: CleanupBudgets +): Promise { try { const tombstonesDeletedBefore = new Date( Date.now() - LARGE_VALUE_TOMBSTONE_RETENTION_HOURS * 60 * 60 * 1000 @@ -359,12 +374,14 @@ async function cleanupLargeValueMetadata(workspaceIds: string[], label: string): const result = await pruneLargeValueMetadata({ workspaceIds, tombstonesDeletedBefore, + budgets, dbClient: cleanupDb, }) logger.info( `[${label}/execution_large_value_metadata] Pruned ${result.referencesDeleted} stale references, ${result.dependenciesDeleted} dependencies, ${result.tombstonesDeleted} tombstones` ) } catch (error) { + if (budgets) throw error logger.error(`[${label}/execution_large_value_metadata] Failed to prune metadata`, { error }) } } @@ -372,7 +389,8 @@ async function cleanupLargeValueMetadata(workspaceIds: string[], label: string): async function cleanupWorkflowExecutionLogs( workspaceIds: string[], retentionDate: Date, - label: string + label: string, + budget?: RowBudget ): Promise { const fileStats: FileDeleteStats = { filesTotal: 0, @@ -381,6 +399,7 @@ async function cleanupWorkflowExecutionLogs( } const dbStats = await chunkedBatchDelete({ + budget, tableDef: workflowExecutionLogs, workspaceIds, tableName: `${label}/workflow_execution_logs`, @@ -420,17 +439,27 @@ async function cleanupWorkflowExecutionLogs( return { ...dbStats, ...fileStats } } -async function cleanupFreePlanOrphanedSnapshots(retentionHours: number): Promise { +async function cleanupFreePlanOrphanedSnapshots( + retentionHours: number, + budget?: RowBudget +): Promise { try { const retentionDays = Math.floor(retentionHours / 24) - const snapshotsCleaned = await snapshotService.cleanupOrphanedSnapshots(retentionDays + 1) + const snapshotsCleaned = await snapshotService.cleanupOrphanedSnapshots( + retentionDays + 1, + budget + ) logger.info(`Cleaned up ${snapshotsCleaned} orphaned snapshots`) } catch (snapshotError) { + if (budget) throw snapshotError logger.error('Error cleaning up orphaned snapshots:', { snapshotError }) } } -export async function runCleanupLogs(payload: CleanupJobPayload): Promise { +export async function runCleanupLogs( + payload: CleanupJobPayload, + budgets?: CleanupBudgets +): Promise { const startTime = Date.now() const { workspaceIds, retentionHours, label, plan, runGlobalHousekeeping } = payload @@ -439,7 +468,7 @@ export async function runCleanupLogs(payload: CleanupJobPayload): Promise if (workspaceIds.length === 0) { logger.info(`[${label}] No workspaces to process`) if (runGlobalHousekeeping && plan === 'free') { - await cleanupFreePlanOrphanedSnapshots(retentionHours) + await cleanupFreePlanOrphanedSnapshots(retentionHours, budgets?.orphanSnapshots) } return } @@ -448,25 +477,38 @@ export async function runCleanupLogs(payload: CleanupJobPayload): Promise `[${label}] Cleaning ${workspaceIds.length} workspaces, cutoff: ${retentionDate.toISOString()}` ) - const workflowResults = await cleanupWorkflowExecutionLogs(workspaceIds, retentionDate, label) + const workflowResults = await cleanupWorkflowExecutionLogs( + workspaceIds, + retentionDate, + label, + budgets?.workflowLogs + ) logger.info( `[${label}] workflow_execution_logs files: ${workflowResults.filesDeleted}/${workflowResults.filesTotal} deleted, ${workflowResults.filesDeleteFailed} failed` ) - const largeValueResults = await cleanupLargeExecutionValues(workspaceIds, retentionDate, label) + if (budgets && workflowResults.filesDeleteFailed) throw new Error('Log file cleanup failed') + const largeValueResults = await cleanupLargeExecutionValues( + workspaceIds, + retentionDate, + label, + budgets?.largeValues + ) logger.info( `[${label}] execution_large_values: ${largeValueResults.largeValuesDeleted}/${largeValueResults.largeValuesTotal} deleted, ${largeValueResults.largeValuesDeleteFailed} failed` ) const legacyLargeValueResults = await cleanupLegacyLargeExecutionValues( workspaceIds, retentionDate, - label + label, + budgets?.legacyLargeValues ) logger.info( `[${label}] legacy_execution_large_values: ${legacyLargeValueResults.largeValuesDeleted}/${legacyLargeValueResults.largeValuesTotal} deleted, ${legacyLargeValueResults.largeValuesDeleteFailed} failed` ) - await cleanupLargeValueMetadata(workspaceIds, label) + await cleanupLargeValueMetadata(workspaceIds, label, budgets) await batchDeleteByWorkspaceAndTimestamp({ + budget: budgets?.jobLogs, tableDef: jobExecutionLogs, workspaceIdCol: jobExecutionLogs.workspaceId, timestampCol: jobExecutionLogs.startedAt, @@ -477,7 +519,7 @@ export async function runCleanupLogs(payload: CleanupJobPayload): Promise }) if (runGlobalHousekeeping && plan === 'free') { - await cleanupFreePlanOrphanedSnapshots(retentionHours) + await cleanupFreePlanOrphanedSnapshots(retentionHours, budgets?.orphanSnapshots) } const timeElapsed = (Date.now() - startTime) / 1000 @@ -487,6 +529,10 @@ export async function runCleanupLogs(payload: CleanupJobPayload): Promise export const cleanupLogsTask = task({ id: 'cleanup-logs', machine: 'large-1x', - queue: { concurrencyLimit: LOG_CLEANUP_CONCURRENCY_LIMIT }, - run: runCleanupLogs, + queue: retentionCleanupQueue, + retry: { maxAttempts: 1 }, + run: (payload: CleanupJobPayload | LimitedCleanupPayload) => + 'limits' in payload + ? runCleanupWithLimits('cleanup-logs', payload.limits, runCleanupLogs) + : runCleanupLogs(payload), }) diff --git a/apps/sim/background/cleanup-soft-deletes.test.ts b/apps/sim/background/cleanup-soft-deletes.test.ts index 2d6816d525a..cb9d262cdf4 100644 --- a/apps/sim/background/cleanup-soft-deletes.test.ts +++ b/apps/sim/background/cleanup-soft-deletes.test.ts @@ -48,7 +48,10 @@ const { mockSelectRowsByIdChunks: vi.fn(async () => [] as unknown[]), })) +vi.mock('@/lib/billing/cleanup-dispatcher', () => ({ runCleanupWithLimits: vi.fn() })) + vi.mock('@/lib/cleanup/batch-delete', () => ({ + consumeRowBudget: vi.fn(), batchDeleteByWorkspaceAndTimestamp: mockBatchDeleteByWorkspaceAndTimestamp, chunkedBatchDelete: mockChunkedBatchDelete, chunkedBatchDeleteByScope: mockScopedChunkedBatchDelete, @@ -57,6 +60,10 @@ vi.mock('@/lib/cleanup/batch-delete', () => ({ selectRowsByIdChunks: mockSelectRowsByIdChunks, })) +vi.mock('@/lib/cleanup/queue', () => ({ + retentionCleanupQueue: { name: 'retention-cleanup', concurrencyLimit: 1 }, +})) + vi.mock('@/lib/cleanup/chat-cleanup', () => ({ prepareChatCleanup: mockPrepareChatCleanup })) vi.mock('@/lib/billing/storage', () => ({ diff --git a/apps/sim/background/cleanup-soft-deletes.ts b/apps/sim/background/cleanup-soft-deletes.ts index 2f48f1d51d7..5497948fd26 100644 --- a/apps/sim/background/cleanup-soft-deletes.ts +++ b/apps/sim/background/cleanup-soft-deletes.ts @@ -16,7 +16,7 @@ import { createLogger } from '@sim/logger' import { chunkArray } from '@sim/utils/helpers' import { task } from '@trigger.dev/sdk' import { and, asc, eq, inArray, isNotNull, isNull, lt, sql } from 'drizzle-orm' -import type { CleanupJobPayload } from '@/lib/billing/cleanup-dispatcher' +import { type CleanupJobPayload, runCleanupWithLimits } from '@/lib/billing/cleanup-dispatcher' import { decrementStorageUsageForBillingContextInTx, resolveStorageBillingContext, @@ -26,10 +26,14 @@ import { batchDeleteByWorkspaceAndTimestamp, chunkedBatchDelete, chunkedBatchDeleteByScope, + consumeRowBudget, DEFAULT_DELETE_CHUNK_SIZE, + type RowBudget, selectRowsByIdChunks, } from '@/lib/cleanup/batch-delete' import { prepareChatCleanup } from '@/lib/cleanup/chat-cleanup' +import type { CleanupBudgets, LimitedCleanupPayload } from '@/lib/cleanup/limits' +import { retentionCleanupQueue } from '@/lib/cleanup/queue' import { type CleanupOwnerScope, cleanupOwnerCondition, @@ -93,47 +97,54 @@ interface WorkspaceFileStorageCleanupResult { */ async function selectExpiredWorkspaceFiles( scope: CleanupOwnerScope, - retentionDate: Date + retentionDate: Date, + budgets?: CleanupBudgets ): Promise { const [legacyRows, multiContextRows] = await Promise.all([ - selectRowsByIdChunks(scope.kind === 'workspace' ? scope.ids : [], (chunkIds, chunkLimit) => - cleanupDb - .select({ - id: workspaceFile.id, - key: workspaceFile.key, - workspaceId: workspaceFile.workspaceId, - }) - .from(workspaceFile) - .where( - and( - inArray(workspaceFile.workspaceId, chunkIds), - isNotNull(workspaceFile.deletedAt), - lt(workspaceFile.deletedAt, retentionDate) + selectRowsByIdChunks( + scope.kind === 'workspace' ? scope.ids : [], + (chunkIds, chunkLimit) => + cleanupDb + .select({ + id: workspaceFile.id, + key: workspaceFile.key, + workspaceId: workspaceFile.workspaceId, + }) + .from(workspaceFile) + .where( + and( + inArray(workspaceFile.workspaceId, chunkIds), + isNotNull(workspaceFile.deletedAt), + lt(workspaceFile.deletedAt, retentionDate) + ) ) - ) - .limit(chunkLimit) + .limit(chunkLimit), + { budget: budgets?.legacyFiles } ), - selectRowsByIdChunks(scope.ids, (chunkIds, chunkLimit) => - cleanupDb - .select({ - id: workspaceFiles.id, - key: workspaceFiles.key, - workspaceId: workspaceFiles.workspaceId, - context: workspaceFiles.context, - sizeBytes: workspaceFiles.sizeBytes, - }) - .from(workspaceFiles) - .where( - and( - cleanupOwnerCondition(workspaceFiles, scope, chunkIds), - scope.kind === 'organization' - ? eq(workspaceFiles.context, 'knowledge-base') - : undefined, - isNotNull(workspaceFiles.deletedAt), - lt(workspaceFiles.deletedAt, retentionDate) + selectRowsByIdChunks( + scope.ids, + (chunkIds, chunkLimit) => + cleanupDb + .select({ + id: workspaceFiles.id, + key: workspaceFiles.key, + workspaceId: workspaceFiles.workspaceId, + context: workspaceFiles.context, + sizeBytes: workspaceFiles.sizeBytes, + }) + .from(workspaceFiles) + .where( + and( + cleanupOwnerCondition(workspaceFiles, scope, chunkIds), + scope.kind === 'organization' + ? eq(workspaceFiles.context, 'knowledge-base') + : undefined, + isNotNull(workspaceFiles.deletedAt), + lt(workspaceFiles.deletedAt, retentionDate) + ) ) - ) - .limit(chunkLimit) + .limit(chunkLimit), + { budget: budgets?.files } ), ]) @@ -393,9 +404,11 @@ async function hardDeleteKnowledgeBaseDocuments( async function cleanupExpiredKnowledgeBases( scope: CleanupOwnerScope, retentionDate: Date, - label: string + label: string, + budget?: RowBudget ) { const options = { + budget, tableDef: knowledgeBase, tableName: `${label}/knowledgeBase`, batchSize: KB_RETENTION_BATCH_SIZE, @@ -687,25 +700,35 @@ const CLEANUP_TARGETS = [ ctx.retentionDate, ctx.label ), + budgetKey: 'folders', name: 'folder', }, { table: userTableDefinitions, softDeleteCol: userTableDefinitions.archivedAt, wsCol: userTableDefinitions.workspaceId, + budgetKey: 'userTables', name: 'userTableDefinitions', }, - { table: memory, softDeleteCol: memory.deletedAt, wsCol: memory.workspaceId, name: 'memory' }, + { + table: memory, + softDeleteCol: memory.deletedAt, + wsCol: memory.workspaceId, + budgetKey: 'memories', + name: 'memory', + }, { table: mcpServers, softDeleteCol: mcpServers.deletedAt, wsCol: mcpServers.workspaceId, + budgetKey: 'mcpServers', name: 'mcpServers', }, { table: workflowMcpServer, softDeleteCol: workflowMcpServer.deletedAt, wsCol: workflowMcpServer.workspaceId, + budgetKey: 'workflowMcpServers', name: 'workflowMcpServer', }, ] as const @@ -721,7 +744,8 @@ const CLEANUP_TARGETS = [ */ async function cleanupOrphanedKnowledgeBaseBindings( scope: CleanupOwnerScope, - label: string + label: string, + budget?: RowBudget ): Promise<{ total: number; deleted: number; failed: number }> { const stats = { total: 0, deleted: 0, failed: 0 } if (scope.ids.length === 0) return stats @@ -730,10 +754,11 @@ async function cleanupOrphanedKnowledgeBaseBindings( for (const chunkIds of chunkArray(scope.ids, KB_ORPHAN_BINDING_OWNER_CHUNK_SIZE)) { let attempted = 0 - while (attempted < KB_ORPHAN_BINDING_TOTAL_LIMIT) { + while (attempted < KB_ORPHAN_BINDING_TOTAL_LIMIT && budget?.remaining !== 0) { const limit = Math.min( KB_ORPHAN_BINDING_BATCH_SIZE, - KB_ORPHAN_BINDING_TOTAL_LIMIT - attempted + KB_ORPHAN_BINDING_TOTAL_LIMIT - attempted, + budget?.remaining ?? KB_ORPHAN_BINDING_BATCH_SIZE ) const rows = await cleanupDb .select({ key: workspaceFiles.key }) @@ -755,6 +780,7 @@ async function cleanupOrphanedKnowledgeBaseBindings( if (rows.length === 0) break + consumeRowBudget(budget, rows.length) const keys = rows.map((row) => row.key) stats.total += keys.length attempted += keys.length @@ -781,6 +807,7 @@ async function cleanupOrphanedKnowledgeBaseBindings( } } stats.deleted += deletedThisBatch + if (budget && stats.failed) throw new Error('Orphan binding cleanup failed') // No progress (every delete failed) — stop rather than reselect the same rows. if (deletedThisBatch === 0) break @@ -793,7 +820,10 @@ async function cleanupOrphanedKnowledgeBaseBindings( return stats } -export async function runCleanupSoftDeletes(payload: CleanupJobPayload): Promise { +export async function runCleanupSoftDeletes( + payload: CleanupJobPayload, + budgets?: CleanupBudgets +): Promise { const startTime = Date.now() const { workspaceIds, retentionHours, label } = payload const scope = resolveCleanupOwnerScope(payload) @@ -813,32 +843,38 @@ export async function runCleanupSoftDeletes(payload: CleanupJobPayload): Promise // could return different subsets above the LIMIT cap and orphan or // prematurely purge data. const [doomedWorkflows, fileScope, expiredSoftDeletedChats] = await Promise.all([ - selectRowsByIdChunks(workspaceIds, (chunkIds, chunkLimit) => - cleanupDb - .select({ id: workflow.id }) - .from(workflow) - .where( - and( - inArray(workflow.workspaceId, chunkIds), - isNotNull(workflow.archivedAt), - lt(workflow.archivedAt, retentionDate) + selectRowsByIdChunks( + workspaceIds, + (chunkIds, chunkLimit) => + cleanupDb + .select({ id: workflow.id }) + .from(workflow) + .where( + and( + inArray(workflow.workspaceId, chunkIds), + isNotNull(workflow.archivedAt), + lt(workflow.archivedAt, retentionDate) + ) ) - ) - .limit(chunkLimit) + .limit(chunkLimit), + { budget: budgets?.workflows } ), - selectExpiredWorkspaceFiles(scope, retentionDate), - selectRowsByIdChunks(scope.ids, (chunkIds, chunkLimit) => - cleanupDb - .select({ id: copilotChats.id }) - .from(copilotChats) - .where( - and( - cleanupOwnerCondition(copilotChats, scope, chunkIds), - isNotNull(copilotChats.deletedAt), - lt(copilotChats.deletedAt, retentionDate) + selectExpiredWorkspaceFiles(scope, retentionDate, budgets), + selectRowsByIdChunks( + scope.ids, + (chunkIds, chunkLimit) => + cleanupDb + .select({ id: copilotChats.id }) + .from(copilotChats) + .where( + and( + cleanupOwnerCondition(copilotChats, scope, chunkIds), + isNotNull(copilotChats.deletedAt), + lt(copilotChats.deletedAt, retentionDate) + ) ) - ) - .limit(chunkLimit) + .limit(chunkLimit), + { budget: budgets?.chats } ), ]) @@ -864,6 +900,7 @@ export async function runCleanupSoftDeletes(payload: CleanupJobPayload): Promise } const fileCleanup = await cleanupWorkspaceFileStorage(fileScope) + if (budgets && fileCleanup.filesFailed) throw new Error('File storage cleanup failed') let totalDeleted = 0 @@ -886,6 +923,7 @@ export async function runCleanupSoftDeletes(payload: CleanupJobPayload): Promise .returning({ id: workflow.id }) totalDeleted += deleted.length } catch (error) { + if (budgets) throw error logger.error(`[${label}/workflow] Archived workflow delete failed`, { error }) } } @@ -910,6 +948,7 @@ export async function runCleanupSoftDeletes(payload: CleanupJobPayload): Promise .returning({ id: copilotChats.id }) totalDeleted += deleted.length } catch (error) { + if (budgets) throw error logger.error(`[${label}/copilotChats] Soft-deleted chat delete failed`, { error }) } } @@ -934,12 +973,23 @@ export async function runCleanupSoftDeletes(payload: CleanupJobPayload): Promise label ) totalDeleted += unbilledFileResult.deleted + if ( + budgets && + (legacyFileResult.failed || billableFileResult.failed || unbilledFileResult.failed) + ) + throw new Error('File row cleanup failed') - const knowledgeBaseResult = await cleanupExpiredKnowledgeBases(scope, retentionDate, label) + const knowledgeBaseResult = await cleanupExpiredKnowledgeBases( + scope, + retentionDate, + label, + budgets?.knowledgeBases + ) totalDeleted += knowledgeBaseResult.deleted for (const target of scope.kind === 'workspace' ? CLEANUP_TARGETS : []) { const result = await batchDeleteByWorkspaceAndTimestamp({ + budget: budgets?.[target.budgetKey], tableDef: target.table, workspaceIdCol: target.wsCol, timestampCol: target.softDeleteCol, @@ -957,7 +1007,11 @@ export async function runCleanupSoftDeletes(payload: CleanupJobPayload): Promise totalDeleted += result.deleted } - const orphanBindingStats = await cleanupOrphanedKnowledgeBaseBindings(scope, label) + const orphanBindingStats = await cleanupOrphanedKnowledgeBaseBindings( + scope, + label, + budgets?.orphanKnowledgeBaseBindings + ) logger.info( `[${label}] Complete: ${totalDeleted} rows deleted, ${fileCleanup.filesDeleted} files cleaned, ${orphanBindingStats.deleted} orphan KB bindings cleaned` @@ -975,6 +1029,10 @@ export async function runCleanupSoftDeletes(payload: CleanupJobPayload): Promise export const cleanupSoftDeletesTask = task({ id: 'cleanup-soft-deletes', machine: 'large-1x', - queue: { concurrencyLimit: 5 }, - run: runCleanupSoftDeletes, + queue: retentionCleanupQueue, + retry: { maxAttempts: 1 }, + run: (payload: CleanupJobPayload | LimitedCleanupPayload) => + 'limits' in payload + ? runCleanupWithLimits('cleanup-soft-deletes', payload.limits, runCleanupSoftDeletes) + : runCleanupSoftDeletes(payload), }) diff --git a/apps/sim/lib/api/contracts/cleanup.test.ts b/apps/sim/lib/api/contracts/cleanup.test.ts new file mode 100644 index 00000000000..4a7ae0b9b6d --- /dev/null +++ b/apps/sim/lib/api/contracts/cleanup.test.ts @@ -0,0 +1,22 @@ +import { describe, expect, it } from 'vitest' +import { logsCleanupQuerySchema, softDeletesCleanupQuerySchema } from '@/lib/api/contracts/cleanup' + +describe('cleanup query limits', () => { + it('preserves scheduled calls without parameters', () => { + expect(logsCleanupQuerySchema.parse({})).toBeUndefined() + }) + it('parses per-type counts', () => { + expect(logsCleanupQuerySchema.parse({ workflowLogs: '25', jobLogs: '0' })).toEqual({ + workflowLogs: 25, + jobLogs: 0, + }) + expect(softDeletesCleanupQuerySchema.parse({ files: '1' })).toEqual({ files: 1 }) + }) + it.each(['', '-1', '1.5', '5001', 'abc', '0'])('rejects invalid count %s', (value) => { + expect(logsCleanupQuerySchema.safeParse({ workflowLogs: value }).success).toBe(false) + }) + it('rejects unknown or wrong-endpoint types', () => { + expect(logsCleanupQuerySchema.safeParse({ files: '1' }).success).toBe(false) + expect(softDeletesCleanupQuerySchema.safeParse({ workflowLogs: '1' }).success).toBe(false) + }) +}) diff --git a/apps/sim/lib/api/contracts/cleanup.ts b/apps/sim/lib/api/contracts/cleanup.ts new file mode 100644 index 00000000000..2a7f9c81111 --- /dev/null +++ b/apps/sim/lib/api/contracts/cleanup.ts @@ -0,0 +1,68 @@ +import { z } from 'zod' +import { defineRouteContract } from '@/lib/api/contracts/types' +import { + type CleanupLimits, + type CleanupType, + LOG_CLEANUP_TYPES, + SOFT_DELETE_CLEANUP_TYPES, +} from '@/lib/cleanup/limits' + +const limitSchema = z + .union([z.number(), z.string().regex(/^\d+$/).transform(Number)]) + .pipe(z.number().int().min(0).max(5000)) +function cleanupQuerySchema(types: readonly CleanupType[]) { + return z + .object(Object.fromEntries(types.map((type) => [type, limitSchema.optional()]))) + .strict() + .transform((query, ctx): CleanupLimits | undefined => { + if (Object.keys(query).length === 0) return undefined + if (!Object.values(query).some((limit) => limit !== undefined && limit > 0)) { + ctx.addIssue({ code: 'custom', message: 'At least one positive cleanup limit is required' }) + return z.NEVER + } + return query + }) +} +export const logsCleanupQuerySchema = cleanupQuerySchema(LOG_CLEANUP_TYPES) +export const softDeletesCleanupQuerySchema = cleanupQuerySchema(SOFT_DELETE_CLEANUP_TYPES) +const responseSchema = z.union([ + z.object({ + triggered: z.literal(true), + jobIds: z.array(z.string()), + jobCount: z.number(), + chunkCount: z.number(), + workspaceCount: z.number(), + }), + z.object({ + triggered: z.literal(true), + runId: z.string(), + limits: z.partialRecord( + z.enum([...LOG_CLEANUP_TYPES, ...SOFT_DELETE_CLEANUP_TYPES]), + z.number().int().min(0).max(5000) + ), + }), +]) +export const logsCleanupContract = defineRouteContract({ + method: 'GET', + path: '/api/logs/cleanup', + query: logsCleanupQuerySchema, + response: { mode: 'json', schema: responseSchema, status: [200, 202] }, +}) +export const softDeletesCleanupContract = defineRouteContract({ + method: 'GET', + path: '/api/cron/cleanup-soft-deletes', + query: softDeletesCleanupQuerySchema, + response: { mode: 'json', schema: responseSchema, status: [200, 202] }, +}) + +/** Apply the same bounds to direct task submissions as HTTP requests. */ +export function validateCleanupLimits( + jobType: 'cleanup-logs' | 'cleanup-soft-deletes', + limits: CleanupLimits +): CleanupLimits { + const parsed = ( + jobType === 'cleanup-logs' ? logsCleanupQuerySchema : softDeletesCleanupQuerySchema + ).parse(limits) + if (!parsed) throw new Error('Cleanup limits are required') + return parsed +} diff --git a/apps/sim/lib/billing/cleanup-dispatcher.test.ts b/apps/sim/lib/billing/cleanup-dispatcher.test.ts index fbf6a593652..f6df7c6f5a3 100644 --- a/apps/sim/lib/billing/cleanup-dispatcher.test.ts +++ b/apps/sim/lib/billing/cleanup-dispatcher.test.ts @@ -36,7 +36,14 @@ vi.mock('@/lib/workspaces/policy', () => ({ isOrganizationWorkspace: vi.fn(), })) -import { dispatchCleanupJobs } from '@/lib/billing/cleanup-dispatcher' +import { tasks } from '@trigger.dev/sdk' +import { + dispatchBoundedCleanup, + dispatchCleanupJobs, + runCleanupWithLimits, +} from '@/lib/billing/cleanup-dispatcher' +import { getHighestPriorityPersonalSubscription } from '@/lib/billing/core/subscription' +import { isOrganizationWorkspace } from '@/lib/workspaces/policy' afterAll(resetEnvFlagsMock) @@ -222,3 +229,103 @@ describe('organization-owned Search retention dispatch', () => { expect(mockEnqueue).not.toHaveBeenCalled() }) }) + +describe('cleanup limits', () => { + beforeEach(() => { + vi.clearAllMocks() + resetDbChainMock() + setEnvFlags({ isBillingEnabled: false, isDataRetentionEnabled: true }) + mockIsTriggerAvailable.mockReturnValue(true) + vi.mocked(getHighestPriorityPersonalSubscription).mockReset() + mockGetOrganizationSubscription.mockReset() + vi.mocked(isOrganizationWorkspace).mockReset() + }) + + it('enqueues one job without querying owners or dispatching child jobs', async () => { + vi.mocked(tasks.trigger).mockResolvedValueOnce({ id: 'run-limited' } as never) + expect(await dispatchBoundedCleanup('cleanup-logs', { workflowLogs: 3 })).toEqual({ + triggered: true, + runId: 'run-limited', + limits: { workflowLogs: 3 }, + }) + expect(tasks.trigger).toHaveBeenCalledWith( + 'cleanup-logs', + { limits: { workflowLogs: 3 } }, + expect.objectContaining({ maxAttempts: 1 }) + ) + expect(dbChainMockFns.select).not.toHaveBeenCalled() + expect(tasks.batchTrigger).not.toHaveBeenCalled() + }) + + it('shares budgets across owners and stops before the next page', async () => { + queueTableRows( + schemaMock.workspace, + ['a', 'b', 'c'].map((id) => ({ + id, + billedAccountUserId: 'user', + organizationId: null, + workspaceMode: 'personal', + organizationSettings: { logRetentionHours: 24 }, + })) + ) + const seen: number[] = [] + await runCleanupWithLimits('cleanup-logs', { workflowLogs: 2 }, async (_scope, budgets) => { + seen.push(budgets.workflowLogs.remaining) + expect(budgets.jobLogs.remaining).toBe(0) + budgets.workflowLogs.remaining-- + }) + expect(seen).toEqual([2, 1]) + expect(dbChainMockFns.select).toHaveBeenCalledOnce() + }) + + it('rejects invalid direct task input before querying', async () => { + await expect( + runCleanupWithLimits('cleanup-logs', { workflowLogs: -1 }, vi.fn()) + ).rejects.toThrow() + expect(dbChainMockFns.select).not.toHaveBeenCalled() + }) + + it.each(['personal', 'organization-workspace', 'organization'] as const)( + 'fails a manual job when %s subscription lookup fails', + async (kind) => { + setEnvFlags({ isBillingEnabled: true }) + const error = new Error('subscription lookup unavailable') + vi.mocked(getHighestPriorityPersonalSubscription).mockRejectedValueOnce(error) + mockGetOrganizationSubscription.mockRejectedValueOnce(error) + vi.mocked(isOrganizationWorkspace).mockReturnValue(true) + queueTableRows( + schemaMock.workspace, + kind === 'organization' + ? [] + : [ + { + id: 'workspace', + billedAccountUserId: 'user', + organizationId: 'organization', + workspaceMode: kind === 'personal' ? 'personal' : 'organization', + organizationSettings: null, + }, + ] + ) + if (kind === 'organization') + queueTableRows(schemaMock.organization, [{ id: 'organization', settings: null }]) + const runScope = vi.fn() + await expect( + runCleanupWithLimits('cleanup-soft-deletes', { files: 1 }, runScope) + ).rejects.toBe(error) + expect(runScope).not.toHaveBeenCalled() + } + ) + + it('requires queued execution and respects the retention switch', async () => { + mockIsTriggerAvailable.mockReturnValue(false) + await expect(dispatchBoundedCleanup('cleanup-logs', { workflowLogs: 1 })).rejects.toThrow( + 'requires Trigger.dev' + ) + setEnvFlags({ isDataRetentionEnabled: false }) + await expect( + runCleanupWithLimits('cleanup-logs', { workflowLogs: 1 }, vi.fn()) + ).rejects.toThrow('retention is disabled') + expect(dbChainMockFns.select).not.toHaveBeenCalled() + }) +}) diff --git a/apps/sim/lib/billing/cleanup-dispatcher.ts b/apps/sim/lib/billing/cleanup-dispatcher.ts index d70bc439369..ea9551c3209 100644 --- a/apps/sim/lib/billing/cleanup-dispatcher.ts +++ b/apps/sim/lib/billing/cleanup-dispatcher.ts @@ -5,10 +5,12 @@ import { createLogger } from '@sim/logger' import { chunkArray } from '@sim/utils/helpers' import { tasks } from '@trigger.dev/sdk' import { and, asc, eq, gt, isNull } from 'drizzle-orm' +import { validateCleanupLimits } from '@/lib/api/contracts/cleanup' import { getOrganizationSubscription } from '@/lib/billing/core/billing' import { getHighestPriorityPersonalSubscription } from '@/lib/billing/core/subscription' import { getPlanType, type PlanCategory } from '@/lib/billing/plan-helpers' import { type RetentionHoursKey, resolveEffectiveRetentionHours } from '@/lib/billing/retention' +import { type CleanupBudgets, type CleanupLimits, createCleanupBudgets } from '@/lib/cleanup/limits' import { getJobQueue } from '@/lib/core/async-jobs' import { shouldExecuteInline } from '@/lib/core/async-jobs/config' import { resolveTriggerRegion } from '@/lib/core/async-jobs/region' @@ -60,8 +62,8 @@ const DAY = 24 type PlanResolutionEntry = readonly [string, PlanCategory] -function getCleanupConcurrencyKey(jobType: CleanupJobType): string { - return `cleanup:${jobType}` +function getCleanupConcurrencyKey(jobType: CleanupJobType): string | undefined { + return jobType === 'cleanup-tasks' ? `cleanup:${jobType}` : undefined } /** @@ -86,7 +88,8 @@ export const CLEANUP_CONFIG = { } as const satisfies Record async function listActiveWorkspaceCleanupScopeRowsPage( - afterId: string | null + afterId: string | null, + pageSize: number ): Promise { const rows = await db .select({ @@ -104,7 +107,7 @@ async function listActiveWorkspaceCleanupScopeRowsPage( : isNull(workspace.archivedAt) ) .orderBy(asc(workspace.id)) - .limit(WORKSPACE_SCOPE_PAGE_SIZE) + .limit(pageSize) return rows.map((row) => ({ ...row, @@ -113,7 +116,8 @@ async function listActiveWorkspaceCleanupScopeRowsPage( } async function resolvePersonalPlanTypesByBilledUserId( - rows: WorkspaceCleanupScopeRow[] + rows: WorkspaceCleanupScopeRow[], + failOnLookupError: boolean ): Promise> { const billedUserIds = Array.from(new Set(rows.map((row) => row.billedAccountUserId))) const entries = await Promise.all( @@ -124,6 +128,7 @@ async function resolvePersonalPlanTypesByBilledUserId( }) return [userId, getPlanType(subscription?.plan)] as const } catch (error) { + if (failOnLookupError) throw error logger.error('Skipping cleanup for billed user after plan lookup failed', { userId, error, @@ -137,7 +142,8 @@ async function resolvePersonalPlanTypesByBilledUserId( } async function resolvePlanTypesByWorkspaceId( - rows: WorkspaceCleanupScopeRow[] + rows: WorkspaceCleanupScopeRow[], + failOnLookupError: boolean ): Promise> { /** * Without billing there are no subscription rows to read, and the per-plan @@ -156,12 +162,16 @@ async function resolvePlanTypesByWorkspaceId( } const userScopedRows = rows.filter((row) => row.workspaceMode !== WORKSPACE_MODE.ORGANIZATION) - const userPlanByBilledUserId = await resolvePersonalPlanTypesByBilledUserId(userScopedRows) + const userPlanByBilledUserId = await resolvePersonalPlanTypesByBilledUserId( + userScopedRows, + failOnLookupError + ) const entries = await Promise.all( rows.map(async (row) => { if (row.workspaceMode === WORKSPACE_MODE.ORGANIZATION) { const organizationId = isOrganizationWorkspace(row) ? row.organizationId : null if (!organizationId) { + if (failOnLookupError) throw new Error('Malformed organization workspace') logger.error('Skipping cleanup for malformed organization workspace', { workspaceId: row.id, organizationId: row.organizationId, @@ -183,6 +193,7 @@ async function resolvePlanTypesByWorkspaceId( return [row.id, getPlanType(subscription?.plan)] as const } catch (error) { + if (failOnLookupError) throw error logger.error('Skipping cleanup for organization workspace after plan lookup failed', { workspaceId: row.id, organizationId, @@ -225,7 +236,16 @@ const GLOBAL_HOUSEKEEPING_PLAN: Partial> = async function forEachCleanupChunk( jobType: CleanupJobType, - onChunk: (payload: CleanupJobPayload) => Promise + onChunk: (payload: CleanupJobPayload) => Promise, + { + shouldStop = () => false, + pageSize = WORKSPACE_SCOPE_PAGE_SIZE, + failOnLookupError = false, + }: { + shouldStop?: () => boolean + pageSize?: number + failOnLookupError?: boolean + } = {} ): Promise<{ chunkCount: number; workspaceCount: number }> { const config = CLEANUP_CONFIG[jobType] const chunkCountByPlan: Partial> = {} @@ -236,6 +256,7 @@ async function forEachCleanupChunk( let afterId: string | null = null const emitChunk = async (payload: CleanupJobPayload) => { + if (shouldStop()) return if (payload.plan === housekeepingPlan && !housekeepingAssigned) { payload.runGlobalHousekeeping = true housekeepingAssigned = true @@ -244,12 +265,12 @@ async function forEachCleanupChunk( await onChunk(payload) } - while (true) { - const rows = await listActiveWorkspaceCleanupScopeRowsPage(afterId) + while (!shouldStop()) { + const rows = await listActiveWorkspaceCleanupScopeRowsPage(afterId, pageSize) if (rows.length === 0) break afterId = rows[rows.length - 1].id - const planByWorkspaceId = await resolvePlanTypesByWorkspaceId(rows) + const planByWorkspaceId = await resolvePlanTypesByWorkspaceId(rows, failOnLookupError) for (const plan of NON_ENTERPRISE_PLANS) { const retentionHours = config.defaults[plan] @@ -294,16 +315,17 @@ async function forEachCleanupChunk( if (jobType === 'cleanup-soft-deletes' || jobType === 'cleanup-tasks') { let afterOrganizationId: string | null = null - while (true) { + while (!shouldStop()) { const organizations = await db .select({ id: organization.id, settings: organization.dataRetentionSettings }) .from(organization) .where(afterOrganizationId ? gt(organization.id, afterOrganizationId) : undefined) .orderBy(asc(organization.id)) - .limit(WORKSPACE_SCOPE_PAGE_SIZE) + .limit(pageSize) if (organizations.length === 0) break afterOrganizationId = organizations[organizations.length - 1].id for (const row of organizations) { + if (shouldStop()) break let plan: PlanCategory = 'enterprise' if (isBillingEnabled) { try { @@ -311,6 +333,7 @@ async function forEachCleanupChunk( if (!subscription) continue plan = getPlanType(subscription.plan) } catch (error) { + if (failOnLookupError) throw error logger.error('Skipping organization cleanup after plan lookup failed', { organizationId: row.id, error, @@ -463,3 +486,38 @@ export async function dispatchCleanupJobs(jobType: CleanupJobType): Promise<{ return { jobIds, jobCount: jobIds.length, chunkCount, workspaceCount } } + +/** Enqueue one job; owner discovery and all deletion happen in the worker. */ +export async function dispatchBoundedCleanup( + jobType: 'cleanup-logs' | 'cleanup-soft-deletes', + input: CleanupLimits +) { + const limits = validateCleanupLimits(jobType, input) + if (!isBillingEnabled && !isDataRetentionEnabled) throw new Error('Data retention is disabled') + if (!isTriggerAvailable()) throw new Error('Queued cleanup requires Trigger.dev') + const run = await tasks.trigger( + jobType, + { limits }, + { + maxAttempts: 1, + region: await resolveTriggerRegion(), + } + ) + return { triggered: true as const, runId: run.id, limits } +} + +/** Reuse existing cleanup functions with one budget across all workspace/organization chunks. */ +export async function runCleanupWithLimits( + jobType: 'cleanup-logs' | 'cleanup-soft-deletes', + input: CleanupLimits, + runScope: (payload: CleanupJobPayload, budgets: CleanupBudgets) => Promise +): Promise { + const limits = validateCleanupLimits(jobType, input) + if (!isBillingEnabled && !isDataRetentionEnabled) throw new Error('Data retention is disabled') + const budgets = createCleanupBudgets(limits) + await forEachCleanupChunk(jobType, (scope) => runScope(scope, budgets), { + shouldStop: () => Object.values(budgets).every((budget) => budget.remaining === 0), + pageSize: 25, + failOnLookupError: true, + }) +} diff --git a/apps/sim/lib/cleanup/batch-delete.test.ts b/apps/sim/lib/cleanup/batch-delete.test.ts index 5e25ce1307f..724e95ce39d 100644 --- a/apps/sim/lib/cleanup/batch-delete.test.ts +++ b/apps/sim/lib/cleanup/batch-delete.test.ts @@ -4,7 +4,11 @@ import { schemaMock } from '@sim/testing' import { describe, expect, it, vi } from 'vitest' -import { batchDeleteByWorkspaceAndTimestamp, chunkedBatchDelete } from '@/lib/cleanup/batch-delete' +import { + batchDeleteByWorkspaceAndTimestamp, + chunkedBatchDelete, + selectRowsByIdChunks, +} from '@/lib/cleanup/batch-delete' /** * Minimal stand-in for the drizzle client `chunkedBatchDelete` calls. Only the DELETE path is @@ -78,3 +82,72 @@ describe('chunkedBatchDelete onBatch contract', () => { expect(order[0]).toBe('onBatch') }) }) + +describe('shared cleanup row budgets', () => { + it('caps selection across ID chunks and subsequent owner scopes', async () => { + const budget = { remaining: 3 } + const select = vi.fn(async (_ids: string[], limit: number) => + [{ id: 'one' }, { id: 'two' }].slice(0, limit) + ) + expect(await selectRowsByIdChunks(['a', 'b'], select, { chunkSize: 1, budget })).toHaveLength(3) + expect(select.mock.calls.map(([, limit]) => limit)).toEqual([3, 1]) + expect(await selectRowsByIdChunks(['c'], select, { budget })).toEqual([]) + expect(select).toHaveBeenCalledTimes(2) + }) + + it('charges restored rows as attempts and uses the remaining limit for each delete batch', async () => { + const budget = { remaining: 3 } + const select = vi.fn(async (_ids: string[], limit: number) => + [{ id: 'one' }, { id: 'two' }].slice(0, limit) + ) + const options = { + tableDef: schemaMock.folder as never, + workspaceIds: ['a'], + tableName: 'folder', + dbClient: createDbClient(() => {}), + selectChunk: select, + budget, + batchSize: 2, + } + const result = await chunkedBatchDelete(options) + expect(result).toMatchObject({ deleted: 2, failed: 1 }) + expect(select.mock.calls.map(([, limit]) => limit)).toEqual([2, 1]) + await chunkedBatchDelete({ ...options, workspaceIds: ['b'] }) + expect(select).toHaveBeenCalledTimes(2) + }) + + it('stops on an error after charging selected rows', async () => { + const budget = { remaining: 2 } + const onDelete = vi.fn() + await expect( + chunkedBatchDelete({ + tableDef: schemaMock.folder as never, + workspaceIds: ['a', 'b'], + tableName: 'folder', + budget, + dbClient: createDbClient(onDelete), + selectChunk: async () => [{ id: 'one' }], + onBatch: async () => { + throw new Error('storage failed') + }, + }) + ).rejects.toThrow('storage failed') + expect(budget.remaining).toBe(1) + expect(onDelete).not.toHaveBeenCalled() + }) + + it('passes budgets through the timestamp helper', async () => { + const onDelete = vi.fn() + await batchDeleteByWorkspaceAndTimestamp({ + tableDef: schemaMock.folder as never, + workspaceIdCol: schemaMock.folder.workspaceId as never, + timestampCol: schemaMock.folder.deletedAt as never, + workspaceIds: ['a'], + retentionDate: new Date(0), + tableName: 'folder', + budget: { remaining: 0 }, + dbClient: createDbClient(onDelete, [{ id: 'one' }]), + }) + expect(onDelete).not.toHaveBeenCalled() + }) +}) diff --git a/apps/sim/lib/cleanup/batch-delete.ts b/apps/sim/lib/cleanup/batch-delete.ts index 468a1f2619f..cd6506d174c 100644 --- a/apps/sim/lib/cleanup/batch-delete.ts +++ b/apps/sim/lib/cleanup/batch-delete.ts @@ -26,7 +26,19 @@ export const DEFAULT_WORKSPACE_CHUNK_SIZE = 50 /** Bounds FK cascade trigger queue (per-statement in-memory) and bind-parameter count. */ export const DEFAULT_DELETE_CHUNK_SIZE = 1000 +export interface RowBudget { + remaining: number +} + +/** Charge selected roots before side effects, including rows that fail or are restored. */ +export function consumeRowBudget(budget: RowBudget | undefined, count: number): void { + if (!budget) return + if (count > budget.remaining) throw new Error('Cleanup selection exceeded its row budget') + budget.remaining -= count +} + export interface SelectByIdChunksOptions { + budget?: RowBudget /** Cap on rows returned across all chunks. Defaults to a full per-table cleanup budget. */ overallLimit?: number chunkSize?: number @@ -48,6 +60,7 @@ export async function selectRowsByIdChunks( { overallLimit = DEFAULT_BATCH_SIZE * DEFAULT_MAX_BATCHES_PER_TABLE, chunkSize = DEFAULT_WORKSPACE_CHUNK_SIZE, + budget, }: SelectByIdChunksOptions = {} ): Promise { if (ids.length === 0) return [] @@ -55,8 +68,10 @@ export async function selectRowsByIdChunks( const rows: T[] = [] for (const chunkIds of chunkArray(ids, chunkSize)) { if (rows.length >= overallLimit) break - const remaining = overallLimit - rows.length + const remaining = Math.min(overallLimit - rows.length, budget?.remaining ?? overallLimit) + if (remaining === 0) break const chunkRows = await query(chunkIds, remaining) + consumeRowBudget(budget, chunkRows.length) rows.push(...chunkRows) } return rows @@ -69,6 +84,7 @@ export interface TableCleanupResult { } export interface ChunkedBatchDeleteOptions { + budget?: RowBudget tableDef: PgTable workspaceIds: string[] tableName: string @@ -132,6 +148,7 @@ export interface ScopedChunkedBatchDeleteOptions /** Shares bounded deletion and side effects across explicit workspace and organization owners. */ export async function chunkedBatchDeleteByScope({ tableDef, + budget, scopeIds, tableName, selectChunk, @@ -155,7 +172,7 @@ export async function chunkedBatchDeleteByScope({ let attempted = 0 for (const [chunkIdx, chunkIds] of chunks.entries()) { - if (attempted >= totalRowLimit) { + if (attempted >= totalRowLimit || budget?.remaining === 0) { stoppedEarly = true break } @@ -167,7 +184,11 @@ export async function chunkedBatchDeleteByScope({ let rows: TRow[] = [] try { const remainingLimit = totalRowLimit - attempted - const effectiveBatchSize = Math.min(batchSize, remainingLimit) + const effectiveBatchSize = Math.min( + batchSize, + remainingLimit, + budget?.remaining ?? batchSize + ) if (effectiveBatchSize <= 0) { hasMore = false break @@ -180,6 +201,7 @@ export async function chunkedBatchDeleteByScope({ break } + consumeRowBudget(budget, rows.length) attempted += rows.length if (onBatch) await onBatch(rows) @@ -194,6 +216,7 @@ export async function chunkedBatchDeleteByScope({ hasMore = rows.length === effectiveBatchSize && attempted < totalRowLimit batchesProcessed++ } catch (error) { + if (budget) throw error // Count rows we tried to delete; SELECT-stage errors leave rows=[]. result.failed += rows.length logger.error( @@ -213,6 +236,7 @@ export async function chunkedBatchDeleteByScope({ } export interface BatchDeleteOptions { + budget?: RowBudget tableDef: PgTable workspaceIdCol: PgColumn timestampCol: PgColumn diff --git a/apps/sim/lib/cleanup/bounded-cleanup.md b/apps/sim/lib/cleanup/bounded-cleanup.md new file mode 100644 index 00000000000..c03ea1735c9 --- /dev/null +++ b/apps/sim/lib/cleanup/bounded-cleanup.md @@ -0,0 +1,18 @@ +# Manual cleanup limits + +The existing cron Lambda can call these endpoints with query parameters: + +- `/api/logs/cleanup?workflowLogs=25&jobLogs=25` +- `/api/cron/cleanup-soft-deletes?files=10&legacyFiles=10` + +Use the existing cron authentication. Each request returns HTTP 202 with `runId` and queues one job. The worker uses the existing retention rules and cleanup functions. Limits are integers from 0 to 5000; omitted types are zero. At least one positive limit is required. Calls without parameters retain scheduled dispatch. + +Log types: `workflowLogs`, `jobLogs`, `largeValues`, `legacyLargeValues`, `orphanSnapshots`, `staleReferences`, `staleDependencies`, `largeValueTombstones`. + +Soft-delete types: `workflows`, `chats`, `legacyFiles`, `files`, `knowledgeBases`, `folders`, `userTables`, `memories`, `mcpServers`, `workflowMcpServers`, `orphanKnowledgeBaseBindings`. + +Each type has one budget across all workspace and organization chunks in that job. Limits count selected root rows, including rows restored or unsuccessfully deleted after selection. Existing child cascades and attached-file cleanup still follow the selected parents; the limit is not a cap on every physical row affected by a cascade. + +Log and soft-delete tasks share a queue with concurrency one. Jobs have one attempt so automatic retries cannot reset a spent row budget. Each new API call creates a new job; inspect the returned run before repeating a call whose response was lost. + +Deploy the worker before the API. Keep schedules disabled while draining the backlog. Start with small limits for one type, inspect the job logs and database load, then repeat and increase counts gradually. Storage, billing, and concurrent-restore behavior follow the existing cleanup implementation. diff --git a/apps/sim/lib/cleanup/limits.ts b/apps/sim/lib/cleanup/limits.ts new file mode 100644 index 00000000000..f6194569ffb --- /dev/null +++ b/apps/sim/lib/cleanup/limits.ts @@ -0,0 +1,41 @@ +import type { RowBudget } from '@/lib/cleanup/batch-delete' + +export const LOG_CLEANUP_TYPES = [ + 'workflowLogs', + 'jobLogs', + 'largeValues', + 'legacyLargeValues', + 'orphanSnapshots', + 'staleReferences', + 'staleDependencies', + 'largeValueTombstones', +] as const +export const SOFT_DELETE_CLEANUP_TYPES = [ + 'workflows', + 'chats', + 'legacyFiles', + 'files', + 'knowledgeBases', + 'folders', + 'userTables', + 'memories', + 'mcpServers', + 'workflowMcpServers', + 'orphanKnowledgeBaseBindings', +] as const +export type CleanupType = + | (typeof LOG_CLEANUP_TYPES)[number] + | (typeof SOFT_DELETE_CLEANUP_TYPES)[number] +export type CleanupLimits = Partial> +export type CleanupBudgets = Record +export type LimitedCleanupPayload = { limits: CleanupLimits } + +/** One mutable budget per type, shared across every owner scope in the queued job. */ +export function createCleanupBudgets(limits: CleanupLimits): CleanupBudgets { + return Object.fromEntries( + [...LOG_CLEANUP_TYPES, ...SOFT_DELETE_CLEANUP_TYPES].map((type) => [ + type, + { remaining: limits[type] ?? 0 }, + ]) + ) as CleanupBudgets +} diff --git a/apps/sim/lib/cleanup/queue.ts b/apps/sim/lib/cleanup/queue.ts new file mode 100644 index 00000000000..06474dcf1f5 --- /dev/null +++ b/apps/sim/lib/cleanup/queue.ts @@ -0,0 +1,2 @@ +import { queue } from '@trigger.dev/sdk' +export const retentionCleanupQueue = queue({ name: 'retention-cleanup', concurrencyLimit: 1 }) diff --git a/apps/sim/lib/execution/payloads/large-value-metadata.ts b/apps/sim/lib/execution/payloads/large-value-metadata.ts index dcbc03f4ff6..efe5ab76482 100644 --- a/apps/sim/lib/execution/payloads/large-value-metadata.ts +++ b/apps/sim/lib/execution/payloads/large-value-metadata.ts @@ -9,6 +9,8 @@ import { import { createLogger } from '@sim/logger' import { chunkArray } from '@sim/utils/helpers' import { and, eq, inArray, notInArray, sql } from 'drizzle-orm' +import { consumeRowBudget } from '@/lib/cleanup/batch-delete' +import type { CleanupBudgets } from '@/lib/cleanup/limits' import { collectLargeValueKeys } from '@/lib/execution/payloads/large-execution-value' const logger = createLogger('LargeValueMetadata') @@ -50,6 +52,7 @@ export interface LargeValueMetadataPruneResult { } interface PruneLargeValueMetadataOptions { + budgets?: CleanupBudgets workspaceIds: string[] tombstonesDeletedBefore: Date batchSize?: number @@ -473,6 +476,7 @@ async function pruneDeletedLargeValueTombstones( export async function pruneLargeValueMetadata({ workspaceIds, tombstonesDeletedBefore, + budgets, batchSize = LARGE_VALUE_METADATA_PRUNE_BATCH_SIZE, maxRowsPerTable = LARGE_VALUE_METADATA_PRUNE_MAX_ROWS_PER_TABLE, dbClient = db, @@ -488,32 +492,47 @@ export async function pruneLargeValueMetadata({ workspaceIds, LARGE_VALUE_METADATA_WORKSPACE_CHUNK_SIZE )) { - const referencesRemaining = maxRowsPerTable - result.referencesDeleted + const referencesRemaining = Math.min( + maxRowsPerTable - result.referencesDeleted, + budgets?.staleReferences.remaining ?? maxRowsPerTable + ) if (referencesRemaining > 0) { - result.referencesDeleted += await pruneStaleReferences( + const deleted = await pruneStaleReferences( workspaceChunk, Math.min(batchSize, referencesRemaining), dbClient ) + consumeRowBudget(budgets?.staleReferences, deleted) + result.referencesDeleted += deleted } - const dependenciesRemaining = maxRowsPerTable - result.dependenciesDeleted + const dependenciesRemaining = Math.min( + maxRowsPerTable - result.dependenciesDeleted, + budgets?.staleDependencies.remaining ?? maxRowsPerTable + ) if (dependenciesRemaining > 0) { - result.dependenciesDeleted += await pruneDeletedParentDependencies( + const deleted = await pruneDeletedParentDependencies( workspaceChunk, Math.min(batchSize, dependenciesRemaining), dbClient ) + consumeRowBudget(budgets?.staleDependencies, deleted) + result.dependenciesDeleted += deleted } - const tombstonesRemaining = maxRowsPerTable - result.tombstonesDeleted + const tombstonesRemaining = Math.min( + maxRowsPerTable - result.tombstonesDeleted, + budgets?.largeValueTombstones.remaining ?? maxRowsPerTable + ) if (tombstonesRemaining > 0) { - result.tombstonesDeleted += await pruneDeletedLargeValueTombstones( + const deleted = await pruneDeletedLargeValueTombstones( workspaceChunk, tombstonesDeletedBefore, Math.min(batchSize, tombstonesRemaining), dbClient ) + consumeRowBudget(budgets?.largeValueTombstones, deleted) + result.tombstonesDeleted += deleted } if ( diff --git a/apps/sim/lib/logs/execution/snapshot/service.ts b/apps/sim/lib/logs/execution/snapshot/service.ts index c6cc2d22ddd..33db22dba97 100644 --- a/apps/sim/lib/logs/execution/snapshot/service.ts +++ b/apps/sim/lib/logs/execution/snapshot/service.ts @@ -4,6 +4,7 @@ import { createLogger } from '@sim/logger' import { sha256Hex } from '@sim/security/hash' import { generateId } from '@sim/utils/id' import { and, eq, inArray, lt, notExists, sql } from 'drizzle-orm' +import { consumeRowBudget, type RowBudget } from '@/lib/cleanup/batch-delete' import type { SnapshotService as ISnapshotService, SnapshotCreationResult, @@ -103,7 +104,7 @@ export class SnapshotService implements ISnapshotService { } /** Only invoked from the cleanup-logs background job, so it runs on the cleanup pool. */ - async cleanupOrphanedSnapshots(olderThanDays: number): Promise { + async cleanupOrphanedSnapshots(olderThanDays: number, budget?: RowBudget): Promise { const cleanupDb = dbFor('cleanup') const cutoffDate = new Date() cutoffDate.setDate(cutoffDate.getDate() - olderThanDays) @@ -115,6 +116,7 @@ export class SnapshotService implements ISnapshotService { let stoppedEarly = false for (let batch = 0; batch < MAX_BATCHES; batch++) { + if (budget?.remaining === 0) break const candidates = await cleanupDb .select({ id: workflowExecutionSnapshots.id }) .from(workflowExecutionSnapshots) @@ -129,10 +131,11 @@ export class SnapshotService implements ISnapshotService { ) ) ) - .limit(BATCH_SIZE) + .limit(Math.min(BATCH_SIZE, budget?.remaining ?? BATCH_SIZE)) if (candidates.length === 0) break + consumeRowBudget(budget, candidates.length) const ids = candidates.map((c) => c.id) const deleted = await cleanupDb .delete(workflowExecutionSnapshots)