diff --git a/apps/sim/background/async-preprocessing-correlation.test.ts b/apps/sim/background/async-preprocessing-correlation.test.ts index 26a03fb7314..15820a1d6b7 100644 --- a/apps/sim/background/async-preprocessing-correlation.test.ts +++ b/apps/sim/background/async-preprocessing-correlation.test.ts @@ -682,36 +682,79 @@ describe('async preprocessing correlation threading', () => { } it.each([ - { statusCode: 404, message: 'Workflow not found' }, - { statusCode: 403, message: 'Workflow is not deployed' }, - { statusCode: 402, message: 'Usage limit exceeded. Please upgrade your plan to continue.' }, - { statusCode: 429, message: 'Rate limit exceeded. Please try again later.' }, - ])('completes the job for a non-retryable $statusCode', async ({ statusCode, message }) => { + { statusCode: 404, message: 'Workflow not found', workflowState: 'not_found' }, + { + statusCode: 403, + message: 'Workflow is not deployed', + code: 'WORKFLOW_NOT_DEPLOYED', + workflowState: 'not_deployed', + }, + { + statusCode: 403, + message: + 'This workflow is not attached to a workspace. Personal workflows are deprecated and cannot execute.', + workflowState: 'not_attached', + }, + ])('completes the job for a $statusCode $message that never started a log', async (error) => { + loggingSessionMockFns.mockHasCompleted.mockReturnValueOnce(false) + + await expect(runRejected(error)).resolves.toMatchObject({ + success: false, + workflowId: 'workflow-1', + executionId: 'execution-rejected', + error: error.message, + }) + }) + + const loggedRejections = [ + { + statusCode: 402, + message: 'Usage limit exceeded. Please upgrade your plan to continue.', + code: 'USAGE_LIMIT_EXCEEDED', + }, + { + statusCode: 402, + message: 'Usage headroom exhausted', + code: 'EXECUTION_PAYER_HEADROOM_EXHAUSTED', + retryable: false, + }, + { + statusCode: 429, + message: 'Rate limit exceeded. Please try again later.', + code: 'RATE_LIMIT_EXCEEDED', + }, + { statusCode: 403, message: 'Account suspended', code: 'ACCOUNT_SUSPENDED' }, + ] + + it.each(loggedRejections)('completes the job for a logged $code rejection', async (error) => { loggingSessionMockFns.mockHasCompleted.mockReturnValueOnce(true) - await expect(runRejected({ statusCode, message })).resolves.toMatchObject({ + + await expect(runRejected(error)).resolves.toMatchObject({ success: false, workflowId: 'workflow-1', executionId: 'execution-rejected', - error: message, + error: error.message, }) expect(loggingSessionMockFns.mockSafeCompleteWithError).toHaveBeenCalledWith( - expect.objectContaining({ error: expect.objectContaining({ message }) }) + expect.objectContaining({ error: expect.objectContaining({ message: error.message }) }) ) }) - it('fails the job when the rejection could not be logged', async () => { - loggingSessionMockFns.mockHasCompleted.mockReturnValueOnce(false) + it.each(loggedRejections)( + 'fails the job when a $code rejection could not be logged', + async (error) => { + loggingSessionMockFns.mockHasCompleted.mockReturnValueOnce(false) - await expect( - runRejected({ statusCode: 403, message: 'Workflow is not deployed' }) - ).rejects.toThrow('Workflow is not deployed') - }) + await expect(runRejected(error)).rejects.toThrow(error.message) + } + ) it.each([ { statusCode: 429, message: 'Too many concurrent executions', retryable: true }, { statusCode: 402, message: 'Usage is temporarily unavailable', retryable: true }, { statusCode: 503, message: 'Usage admission is temporarily unavailable', retryable: true }, { statusCode: 500, message: 'Internal error while fetching workflow' }, + { statusCode: 404, message: 'Workflow not found' }, ])('fails the job for $statusCode $message', async (error) => { await expect(runRejected(error)).rejects.toThrow(error.message) }) diff --git a/apps/sim/background/workflow-execution.ts b/apps/sim/background/workflow-execution.ts index aaf022ecf63..391d73b168f 100644 --- a/apps/sim/background/workflow-execution.ts +++ b/apps/sim/background/workflow-execution.ts @@ -16,6 +16,8 @@ import { assertBillingAttributionSnapshot, type BillingAttributionSnapshot, } from '@/lib/billing/core/attribution' +import { getDeterministicAdmissionRejectionCode } from '@/lib/core/admission/rejection' +import { isReservationHeadroomDenial } from '@/lib/core/admission/transient-failure' import { buildRecordedJobFailureResult } from '@/lib/core/async-jobs/recorded-failure' import type { AsyncExecutionCorrelation } from '@/lib/core/async-jobs/types' import type { ExecutionScope } from '@/lib/core/errors/execution-scope-deleted' @@ -30,7 +32,12 @@ import { RESERVATION_TTL_BUFFER_MS, } from '@/lib/core/execution-limits' import type { RequestAttribution } from '@/lib/core/utils/request-attribution' -import { type PreprocessExecutionError, preprocessExecution } from '@/lib/execution/preprocessing' +import { + isWorkflowStateRejection, + type PreprocessExecutionError, + preprocessExecution, + RATE_LIMIT_EXCEEDED_CODE, +} from '@/lib/execution/preprocessing' import { LoggingSession } from '@/lib/logs/execution/logging-session' import { buildTraceSpans } from '@/lib/logs/execution/trace-spans/trace-spans' import { cleanupExecutionBase64Cache } from '@/lib/uploads/utils/user-file-base64.server' @@ -109,18 +116,23 @@ type LegacyWorkflowExecutionPayload = Omit = new Set([402, 403, 404, 429]) - -function isNonFaultPreprocessingRejection(rejection: PreprocessExecutionError): boolean { - return ( - NON_FAULT_PREPROCESSING_STATUS_CODES.has(rejection.statusCode) && rejection.retryable !== true - ) +function isNonFaultPreprocessingRejection( + rejection: PreprocessExecutionError, + rejectionLogged: boolean +): boolean { + if (rejection.retryable === true) return false + if (isWorkflowStateRejection(rejection)) return true + const quotaRejection = + getDeterministicAdmissionRejectionCode(rejection) !== undefined || + isReservationHeadroomDenial(rejection) || + rejection.code === RATE_LIMIT_EXCEEDED_CODE + return quotaRejection && rejectionLogged } /** Completes the job with a failure already in the execution log, keeping it on the run's span. */ @@ -378,11 +390,9 @@ export async function executeWorkflowJob( const scopeDeletedOnCompletion = loggingSession.getDeletedScope() if (scopeDeletedOnCompletion) return completeWithDeletedScope(scopeDeletedOnCompletion) - // Complete only when the rejection is durably logged; an unsaved one must still alert. if ( preprocessingRejection && - isNonFaultPreprocessingRejection(preprocessingRejection) && - loggingSession.hasCompleted() + isNonFaultPreprocessingRejection(preprocessingRejection, loggingSession.hasCompleted()) ) { return completeWithRecordedFailure(workflowId, executionId, error) } diff --git a/apps/sim/lib/billing/core/usage.ts b/apps/sim/lib/billing/core/usage.ts index 22d9f5b1475..f9ba5b85d60 100644 --- a/apps/sim/lib/billing/core/usage.ts +++ b/apps/sim/lib/billing/core/usage.ts @@ -262,7 +262,7 @@ export async function handleNewSignInAccount( ): Promise { if (!isUserCreatedInAuthRequest(userId, authRequest)) return if (!(await isDeletedAccountIdentity({ oauthSubjects: [{ providerId, accountId }] }))) return - await forfeitFreeTierAllowance(userId) + await forfeitFreeTierAllowance(db, userId) } /** @@ -283,29 +283,45 @@ export async function ensureUserStatsExists( userId: string, authRequest?: object | null ): Promise { - const created = await db + await createUserStatsIfMissing(db, userId, { + includeSignIns: isUserCreatedInAuthRequest(userId, authRequest), + }) +} + +/** + * Creates a user's `user_stats` row with the free allowance when none exists, on `executor` + * (the caller's transaction, or `db`). Only a row this call creates gets the deleted-identity + * check signup makes: the email always, and the linked sign-ins only when `includeSignIns` is + * set. An existing row costs only the conflicting insert. A failed check leaves the allowance in + * place and never fails the caller or aborts its transaction. + * + * @returns whether this call created the row + */ +export async function createUserStatsIfMissing( + executor: DbOrTx, + userId: string, + { includeSignIns = false }: { includeSignIns?: boolean } = {} +): Promise { + const created = await executor .insert(userStats) .values({ id: generateId(), - userId: userId, + userId, currentUsageLimit: getFreeTierLimit().toString(), usageLimitUpdatedAt: new Date(), }) .onConflictDoNothing({ target: userStats.userId }) .returning({ id: userStats.id }) - if (created.length === 0) return + if (created.length === 0) return false try { - const includeSignIns = isUserCreatedInAuthRequest(userId, authRequest) - if (await userHasDeletedAccountIdentity(userId, { includeSignIns })) { - await forfeitFreeTierAllowance(userId) + if (await userHasDeletedAccountIdentity(userId, { includeSignIns }, executor)) { + await forfeitFreeTierAllowance(executor, userId) } } catch (error) { - logger.error('Failed to withhold the free allowance from a fallback stats row', { - userId, - error, - }) + logger.error('Failed to withhold the free allowance from a new stats row', { userId, error }) } + return true } interface ResolvedUserUsageData { @@ -690,12 +706,12 @@ export async function checkUsageStatus(userId: string): Promise<{ /** * Withholds the free allowance from a new account found to reuse the sign-in * identity of a recently deleted one after its stats row was created. The flag - * and the limit it implies commit together, so a failed re-derivation can never - * leave the account flagged but still holding the free limit; repeating the - * call is harmless. + * and the limit it implies commit together (a transaction on `db`, a savepoint + * inside a caller's transaction), so a failed re-derivation can never leave the + * account flagged but still holding the free limit; repeating the call is harmless. */ -async function forfeitFreeTierAllowance(userId: string): Promise { - await db.transaction(async (tx) => { +async function forfeitFreeTierAllowance(executor: DbOrTx, userId: string): Promise { + await executor.transaction(async (tx) => { await tx.update(userStats).set({ freeTierEligible: false }).where(eq(userStats.userId, userId)) await syncUsageLimitsFromSubscription(userId, tx) }) diff --git a/apps/sim/lib/billing/organizations/membership.ts b/apps/sim/lib/billing/organizations/membership.ts index 6af43569466..17e7925920a 100644 --- a/apps/sim/lib/billing/organizations/membership.ts +++ b/apps/sim/lib/billing/organizations/membership.ts @@ -29,7 +29,7 @@ import { invalidateSecurityPolicyVersionCache, } from '@/lib/auth/security-policy' import { applySessionPolicyToNewMember } from '@/lib/auth/session-policy' -import { syncUsageLimitsFromSubscription } from '@/lib/billing/core/usage' +import { createUserStatsIfMissing, syncUsageLimitsFromSubscription } from '@/lib/billing/core/usage' import { assertNoUnresolvedEnterpriseIssuance, resolveEnterpriseMetadataIntent, @@ -1935,14 +1935,7 @@ export async function transferOrganizationOwnership( .limit(1) if (oldStats) { - await tx - .insert(userStats) - .values({ - id: generateId(), - userId: newOwnerUserId, - usageLimitUpdatedAt: new Date(), - }) - .onConflictDoNothing({ target: userStats.userId }) + await createUserStatsIfMissing(tx, newOwnerUserId) const overage = oldStats.billedOverageThisPeriod || '0' const overageNum = toNumber(toDecimal(overage)) diff --git a/apps/sim/lib/billing/storage/tracking.ts b/apps/sim/lib/billing/storage/tracking.ts index 2f4875d6ff7..0a5858621d1 100644 --- a/apps/sim/lib/billing/storage/tracking.ts +++ b/apps/sim/lib/billing/storage/tracking.ts @@ -19,11 +19,11 @@ import { organization, project, userStats, workspace } from '@sim/db/schema' import { createLogger } from '@sim/logger' -import { generateId } from '@sim/utils/id' import { isRecordLike } from '@sim/utils/object' import { and, asc, eq, inArray, sql } from 'drizzle-orm' import { maybeNotifyLimit } from '@/lib/billing/core/limit-notifications' import type { HighestPrioritySubscription } from '@/lib/billing/core/plan' +import { createUserStatsIfMissing } from '@/lib/billing/core/usage' import type { BillingEntity } from '@/lib/billing/core/usage-log' import type { ProjectStorageBillingContext, @@ -39,7 +39,7 @@ import { isStorageEnforcementEnabled, StorageLimitExceededError, } from '@/lib/billing/storage/limits' -import { getFreeTierLimit, isOrgScopedSubscription } from '@/lib/billing/subscriptions/utils' +import { isOrgScopedSubscription } from '@/lib/billing/subscriptions/utils' import type { DbOrTx, DbTransaction } from '@/lib/db/types' import { lockProject } from '@/lib/projects/membership' @@ -764,14 +764,12 @@ export async function incrementAdmittedStorageUsageForBillingContextInTx( * window between the read and the increment. This helper does not send * notifications because it runs mid-transaction. * - * For a personal (non-org-scoped) `userId`, this first upserts the - * `userStats` row on `tx` — a documented possibility for OAuth account - * linking (see `ensureUserStatsExists` in `lib/billing/core/usage.ts`, whose - * insert values this mirrors) — because the conditional `UPDATE` below - * matches 0 rows, and therefore reads as "quota exceeded", if that row - * doesn't exist yet. `ensureUserStatsExists` itself isn't reused here since - * it writes through the standalone `db` client, which would open a second - * pooled connection while this transaction's is held. + * For a personal (non-org-scoped) `userId`, this first creates a missing + * `userStats` row on `tx` with `createUserStatsIfMissing` — a documented + * possibility for OAuth account linking — because the conditional `UPDATE` + * below matches 0 rows, and therefore reads as "quota exceeded", if that row + * doesn't exist yet. It runs on `tx` rather than the standalone `db` client, + * which would open a second pooled connection while this transaction's is held. */ export async function checkAndIncrementStorageUsageInTx( tx: DbOrTx, @@ -791,17 +789,7 @@ export async function checkAndIncrementStorageUsageInTx( const orgScoped = isOrgScopedSubscription(sub, userId) && sub - if (!orgScoped) { - await tx - .insert(userStats) - .values({ - id: generateId(), - userId, - currentUsageLimit: getFreeTierLimit().toString(), - usageLimitUpdatedAt: new Date(), - }) - .onConflictDoNothing({ target: userStats.userId }) - } + if (!orgScoped) await createUserStatsIfMissing(tx, userId) const [updated] = orgScoped ? await tx diff --git a/apps/sim/lib/core/admission/transient-failure.ts b/apps/sim/lib/core/admission/transient-failure.ts index c91f269cc88..1e25bf3188c 100644 --- a/apps/sim/lib/core/admission/transient-failure.ts +++ b/apps/sim/lib/core/admission/transient-failure.ts @@ -83,6 +83,23 @@ function hasAdmissionCode(failure: AdmissionFailureLike, code: string): boolean return failure.code === code || failure.cause?.code === code } +const RESERVATION_HEADROOM_DESCRIPTORS = [ + ADMISSION_ERROR_DESCRIPTOR.RESERVATION_PAYER_HEADROOM, + ADMISSION_ERROR_DESCRIPTOR.RESERVATION_MEMBER_HEADROOM, +] as const + +/** + * Whether a reservation was denied because the payer's or member's usage headroom is spent. The + * run is refused for good, though a later one can be admitted as in-flight runs settle, which is + * why deterministic admission rejections leave these out. + */ +export function isReservationHeadroomDenial(failure: AdmissionFailureLike): boolean { + return RESERVATION_HEADROOM_DESCRIPTORS.some( + (descriptor) => + failure.statusCode === descriptor.statusCode && hasAdmissionCode(failure, descriptor.code) + ) +} + /** * Identifies only admission failures that unattended callers can safely retry. */ diff --git a/apps/sim/lib/core/async-jobs/types.ts b/apps/sim/lib/core/async-jobs/types.ts index 8efb0e9893a..49583eaa6de 100644 --- a/apps/sim/lib/core/async-jobs/types.ts +++ b/apps/sim/lib/core/async-jobs/types.ts @@ -338,9 +338,12 @@ export interface JobQueueBackend { cancelByKey(cancelKey: string): boolean /** - * Cancel every queued, delayed or running workflow job (`WORKSPACE_EXECUTION_JOB_TYPES`) - * enqueued for a workspace, for when the workspace itself is being removed. Throws when any - * listing or cancellation failed, after attempting the rest, so the caller can retry. + * Cancel every queued, delayed or running persisted workflow job + * (`WORKSPACE_EXECUTION_JOB_TYPES`) enqueued for a workspace, for when the workspace itself is + * being removed. Runs the database backend executes through `batchEnqueueAndWait` keep no job + * row, so this never reaches them; they stop through their execution's cancellation signal. + * Throws when any listing or cancellation failed, after attempting the rest, so the caller can + * retry. */ cancelByWorkspace( workspaceId: string, diff --git a/apps/sim/lib/execution/preprocessing.test.ts b/apps/sim/lib/execution/preprocessing.test.ts index 376972d1649..d51a554f8dc 100644 --- a/apps/sim/lib/execution/preprocessing.test.ts +++ b/apps/sim/lib/execution/preprocessing.test.ts @@ -107,6 +107,69 @@ beforeEach(() => { }) describe('preprocessExecution deployment checks', () => { + type PrefetchedWorkflow = NonNullable[0]['workflowRecord']> + const prefetched = { + id: 'workflow-1', + userId: 'user-1', + workspaceId: 'workspace-1', + isDeployed: true, + archivedAt: null, + } as PrefetchedWorkflow + + it.each([ + { name: 'a missing workflow', active: null, workflowState: 'not_found' }, + { + name: 'a prefetched archived workflow', + workflowRecord: { ...prefetched, archivedAt: new Date() }, + workflowState: 'not_found', + }, + { + name: 'a workflow archived before the run started', + active: null, + workflowRecord: prefetched, + workflowState: 'not_found', + }, + { + name: 'a workflow attached to no workspace', + active: { ...prefetched, workspaceId: null }, + workflowState: 'not_attached', + }, + ])( + 'marks $name as a workflow-state rejection', + async ({ active, workflowRecord, workflowState }) => { + if (active !== undefined) + workflowAuthzMockFns.mockGetActiveWorkflowRecord.mockResolvedValueOnce(active) + + const result = await preprocessExecution({ + workflowId: 'workflow-1', + userId: 'user-1', + triggerType: 'api', + executionId: 'execution-1', + requestId: 'request-1', + checkDeployment: true, + ...(workflowRecord ? { workflowRecord } : {}), + }) + + expect(result).toMatchObject({ success: false, error: { workflowState } }) + } + ) + + it('never marks a failed workflow read as a workflow-state rejection', async () => { + workflowAuthzMockFns.mockGetActiveWorkflowRecord.mockRejectedValue(new Error('read failed')) + + const result = await preprocessExecution({ + workflowId: 'workflow-1', + userId: 'user-1', + triggerType: 'api', + executionId: 'execution-1', + requestId: 'request-1', + checkDeployment: true, + }) + + expect(result).toMatchObject({ success: false, error: { statusCode: 500 } }) + expect(result.success === false && result.error.workflowState).toBeUndefined() + }) + it('returns a structured code when a required deployment is missing', async () => { workflowAuthzMockFns.mockGetActiveWorkflowRecord.mockResolvedValueOnce({ id: 'workflow-1', @@ -129,6 +192,7 @@ describe('preprocessExecution deployment checks', () => { message: 'Workflow is not deployed', statusCode: 403, code: WORKFLOW_NOT_DEPLOYED_CODE, + workflowState: 'not_deployed', }, }) }) diff --git a/apps/sim/lib/execution/preprocessing.ts b/apps/sim/lib/execution/preprocessing.ts index 3ca53f852b6..0d4de3132a3 100644 --- a/apps/sim/lib/execution/preprocessing.ts +++ b/apps/sim/lib/execution/preprocessing.ts @@ -155,6 +155,11 @@ export interface PreprocessExecutionError { */ retryAfterMs?: number cause?: Record + /** + * Set when the workflow's own state refused the run; see {@link isWorkflowStateRejection}. + * Internal: unlike `code`, no surface copies it onto a response. + */ + workflowState?: 'not_found' | 'not_attached' | 'not_deployed' } /** @@ -174,6 +179,20 @@ function retryAfterMsFrom(retryAfterSeconds: number | undefined): { retryAfterMs export const WORKFLOW_NOT_DEPLOYED_CODE = 'WORKFLOW_NOT_DEPLOYED' +/** A rate-limit quota denial, distinct from the retryable concurrency-slot 429. */ +export const RATE_LIMIT_EXCEEDED_CODE = 'RATE_LIMIT_EXCEEDED' + +/** + * Whether the workflow's own state refused the run: deleted or archived, not attached to a + * workspace, or not deployed. Produced only after a successful read (a failed read is a 5xx), and + * these rejections write no execution log by design, so callers need not wait for one. + */ +export function isWorkflowStateRejection( + error: Pick +): boolean { + return error.workflowState !== undefined +} + export interface PreprocessExecutionSuccess { success: true actorUserId: string @@ -455,6 +474,7 @@ export async function preprocessExecution( error: { message: 'Workflow not found', statusCode: 404, + workflowState: 'not_found', }, } } @@ -490,6 +510,7 @@ export async function preprocessExecution( error: { message: 'Workflow not found', statusCode: 404, + workflowState: 'not_found', }, } } else { @@ -503,6 +524,7 @@ export async function preprocessExecution( error: { message: 'Workflow not found', statusCode: 404, + workflowState: 'not_found', }, } } @@ -519,6 +541,7 @@ export async function preprocessExecution( message: 'This workflow is not attached to a workspace. Personal workflows are deprecated and cannot execute.', statusCode: 403, + workflowState: 'not_attached', }, } } @@ -533,6 +556,7 @@ export async function preprocessExecution( message: 'Workflow is not deployed', statusCode: 403, code: WORKFLOW_NOT_DEPLOYED_CODE, + workflowState: 'not_deployed', }, } } @@ -869,7 +893,7 @@ export async function preprocessExecution( * (`EXECUTION_CONCURRENCY_LIMIT`, retryable in seconds) — the two * need different caller behavior. */ - code: 'RATE_LIMIT_EXCEEDED', + code: RATE_LIMIT_EXCEEDED_CODE, retryAfterMs: info.retryAfterMs ?? Math.max(0, info.resetAt.getTime() - Date.now()), }, }, diff --git a/apps/sim/lib/projects/files/prefix-cleanup.test.ts b/apps/sim/lib/projects/files/prefix-cleanup.test.ts index 0305976c60a..9993c3e22b3 100644 --- a/apps/sim/lib/projects/files/prefix-cleanup.test.ts +++ b/apps/sim/lib/projects/files/prefix-cleanup.test.ts @@ -10,10 +10,15 @@ import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' vi.mock('@/lib/uploads/core/setup.server', () => uploadsSetupMock) vi.mock('@/lib/uploads/config', () => uploadsConfigMock) vi.mock('@/lib/uploads/server/metadata', () => uploadsMetadataMock) +vi.mock('@/lib/uploads/core/prefix-sweep', async (importOriginal) => { + const actual = await importOriginal() + return { ...actual, listStorageKeyPage: vi.fn(actual.listStorageKeyPage) } +}) import type { OutboxEventContext } from '@/lib/core/outbox/service' import { PROJECT_FILE_PREFIX_CLEANUP_EVENT } from '@/lib/projects/files/outbox-events' import { projectFilePrefixCleanupOutboxHandlers } from '@/lib/projects/files/prefix-cleanup' +import { listStorageKeyPage } from '@/lib/uploads/core/prefix-sweep' const cleanup = projectFilePrefixCleanupOutboxHandlers[PROJECT_FILE_PREFIX_CLEANUP_EVENT] @@ -68,12 +73,56 @@ describe('Project prefix cleanup on local storage', () => { expect(await exists('project/p10/keep.txt')).toBe(true) }) - it('spends an attempt when a full page holds only objects it cannot remove', async () => { + it('reaches objects behind a full page it cannot remove', async () => { for (let index = 0; index < 520; index++) await put(`project/p1/legacy file ${index}.txt`) + await put('project/p1/zz.txt') + + await expect( + cleanup({ projectId: 'p1' }, outboxContext(Date.now() + 60_000)) + ).resolves.toBeUndefined() + + expect(await exists('project/p1/zz.txt')).toBe(false) + }) + + it('keeps sweeping after a short page that is not the last', async () => { + await put('project/p1/a.txt') + await put('project/p1/b.txt') + vi.mocked(listStorageKeyPage).mockResolvedValueOnce({ + keys: ['project/p1/a.txt'], + done: false, + nextPageToken: 'next', + }) + + const result = await cleanup({ projectId: 'p1' }, outboxContext()) + + expect(result).toMatchObject({ outcome: 'deferred', consumeAttempt: false }) + expect(await exists('project/p1/a.txt')).toBe(false) + expect(await exists('project/p1/b.txt')).toBe(true) + }) + + it('follows an empty page that is not the last to the objects after it', async () => { + await put('project/p1/a.txt') + vi.mocked(listStorageKeyPage) + .mockResolvedValueOnce({ keys: [], done: false, nextPageToken: 'next' }) + .mockResolvedValueOnce({ keys: ['project/p1/a.txt'], done: true }) + + await expect(cleanup({ projectId: 'p1' }, outboxContext())).resolves.toBeUndefined() + + expect(vi.mocked(listStorageKeyPage).mock.calls[1]?.[3]).toBe('next') + expect(await exists('project/p1/a.txt')).toBe(false) + }) + + it('spends an attempt when its listing stops advancing without removing anything', async () => { + vi.mocked(listStorageKeyPage).mockResolvedValue({ + keys: [], + done: false, + nextPageToken: 'stuck', + }) - const result = await cleanup({ projectId: 'p1' }, outboxContext(Date.now() + 60_000)) + const result = await cleanup({ projectId: 'p1' }, outboxContext()) expect(result).toMatchObject({ outcome: 'deferred' }) expect(result).not.toMatchObject({ consumeAttempt: false }) + expect(listStorageKeyPage).toHaveBeenCalledTimes(2) }) }) diff --git a/apps/sim/lib/projects/files/prefix-cleanup.ts b/apps/sim/lib/projects/files/prefix-cleanup.ts index 8f096a050a0..2b39b2eac8d 100644 --- a/apps/sim/lib/projects/files/prefix-cleanup.ts +++ b/apps/sim/lib/projects/files/prefix-cleanup.ts @@ -21,7 +21,6 @@ import { deleteStorageKeysUnderPrefix, listStorageKeyPage, outboxDeadlineNear, - STORAGE_KEY_PAGE_SIZE, storageCleanupAvailableAt, } from '@/lib/uploads/core/prefix-sweep' import { abortProviderUpload } from '@/lib/uploads/upload-session/provider' @@ -202,21 +201,33 @@ export async function queueRetiredProjectUploadCleanup( }) } -/** `stalled` when a page that must be re-listed removed nothing, so continuing cannot progress. */ +/** + * Deletes one listing page's objects per run, re-listing from the prefix start next time. A + * provider (Azure Blob, GCS, S3-compatibles) may return a short or empty page that is not the last, + * so only `page.done` ends the sweep; a non-final page that removes nothing is followed to the next + * one within the run, so objects behind it are still reached. `stalled` when no page removed + * anything before the listing stopped advancing or the deadline neared. + */ async function sweepPrefix( projectId: string, context: OutboxEventContext ): Promise<'done' | 'continue' | 'stalled'> { const prefix = prefixFor(projectId) - const page = await listStorageKeyPage(prefix, 'project', context.signal) - const { deleted, complete } = await deleteStorageKeysUnderPrefix( - prefix, - 'project', - page.keys, - context - ) - if (complete && page.keys.length < STORAGE_KEY_PAGE_SIZE) return 'done' - return deleted > 0 ? 'continue' : 'stalled' + let pageToken: string | undefined + for (;;) { + const page = await listStorageKeyPage(prefix, 'project', context.signal, pageToken) + const { deleted, complete } = await deleteStorageKeysUnderPrefix( + prefix, + 'project', + page.keys, + context + ) + if (complete && page.done) return 'done' + if (deleted > 0) return 'continue' + const advanced = page.nextPageToken !== undefined && page.nextPageToken !== pageToken + if (!complete || !advanced || outboxDeadlineNear(context)) return 'stalled' + pageToken = page.nextPageToken + } } const cleanupPrefix: OutboxHandler = async (raw, context) => { @@ -259,7 +270,7 @@ const cleanupPrefix: OutboxHandler = async (raw, context) => { const sweep = await sweepPrefix(payload.projectId, context) if (sweep === 'continue') return continueOutboxHandler('Continuing Project storage cleanup') if (sweep === 'stalled') - return deferOutboxHandler('Project storage cleanup removed nothing from a full listing page') + return deferOutboxHandler('Project storage cleanup removed nothing from its listing') } /** diff --git a/apps/sim/lib/uploads/core/prefix-sweep.ts b/apps/sim/lib/uploads/core/prefix-sweep.ts index 7e7719f422c..6786bcd8d55 100644 --- a/apps/sim/lib/uploads/core/prefix-sweep.ts +++ b/apps/sim/lib/uploads/core/prefix-sweep.ts @@ -16,7 +16,7 @@ import { deleteFile } from '@/lib/uploads/core/storage-service' import { UPLOAD_URL_TTL_MS } from '@/lib/uploads/upload-session/provider' /** Most keys one {@link listStorageKeyPage} call returns. */ -export const STORAGE_KEY_PAGE_SIZE = 500 +const STORAGE_KEY_PAGE_SIZE = 500 const CLOCK_SKEW_MS = 60_000 const DEADLINE_MARGIN_MS = 5000 diff --git a/apps/sim/lib/users/deleted-identity.integration.ts b/apps/sim/lib/users/deleted-identity.integration.ts index e98597a168b..54071e8ea10 100644 --- a/apps/sim/lib/users/deleted-identity.integration.ts +++ b/apps/sim/lib/users/deleted-identity.integration.ts @@ -4,19 +4,30 @@ * and plan-sync paths. */ import { db } from '@sim/db' -import { account, deletedAccountIdentity, user, userStats } from '@sim/db/schema' +import { + account, + deletedAccountIdentity, + member, + organization, + user, + userStats, +} from '@sim/db/schema' +import { createDeferred } from '@sim/testing' import { authMock } from '@sim/testing/mocks/auth.mock' import { envFlagsMock } from '@sim/testing/mocks/env-flags.mock' import { generateId } from '@sim/utils/id' import { generateRandomHex } from '@sim/utils/random' -import { eq, inArray } from 'drizzle-orm' +import { eq, inArray, sql } from 'drizzle-orm' import { afterAll, describe, expect, it, vi } from 'vitest' import { + createUserStatsIfMissing, ensureUserStatsExists, handleNewSignInAccount, handleNewUser, syncUsageLimitsFromSubscription, } from '@/lib/billing/core/usage' +import { transferOrganizationOwnership } from '@/lib/billing/organizations/membership' +import { checkAndIncrementStorageUsageInTx } from '@/lib/billing/storage/tracking' import { getFreeTierLimit } from '@/lib/billing/subscriptions/utils' import { deleteUserAccount } from '@/lib/users/account-deletion' import { purgeExpiredDeletedAccountIdentities } from '@/lib/users/deleted-identity' @@ -29,6 +40,7 @@ vi.mock('@/lib/core/config/env-flags', () => ({ vi.mock('@/lib/auth', () => authMock) const userIds: string[] = [] +const organizationIds: string[] = [] const recordedHashes: string[] = [] interface SignIn { @@ -111,6 +123,8 @@ afterAll(async () => { .delete(deletedAccountIdentity) .where(inArray(deletedAccountIdentity.identityHash, recordedHashes)) } + if (organizationIds.length) + await db.delete(organization).where(inArray(organization.id, organizationIds)) if (userIds.length) await db.delete(user).where(inArray(user.id, userIds)) }) @@ -250,6 +264,108 @@ describe('deleted account identities and the free allowance', () => { expect(await statsOf(stranger)).toEqual({ freeTierEligible: true, limit: getFreeTierLimit() }) }) + it('withholds the allowance from a stats row a storage write creates', async () => { + const email = `storage.${generateId()}@deleted-identity.test` + await deleteOwnAccount(await signUp(email)) + const returning = await insertUser(email) + const stranger = await insertUser(`storage-stranger.${generateId()}@deleted-identity.test`) + const store = (userId: string) => + db.transaction((tx) => checkAndIncrementStorageUsageInTx(tx, null, userId, 1024)) + + const first = await store(returning) + const second = await store(returning) + await store(stranger) + + expect([first.allowed, second.allowed]).toEqual([true, true]) + expect(await statsOf(returning)).toEqual({ freeTierEligible: false, limit: 0 }) + expect(await statsOf(stranger)).toEqual({ freeTierEligible: true, limit: getFreeTierLimit() }) + const [stored] = await db + .select({ bytes: userStats.storageUsedBytes }) + .from(userStats) + .where(eq(userStats.userId, returning)) + expect(stored.bytes).toBe(2048) + }) + + it('withholds the allowance from a stats row an ownership transfer creates for the new owner', async () => { + const email = `new-owner.${generateId()}@deleted-identity.test` + await deleteOwnAccount(await signUp(email)) + const owner = await signUp(`owner.${generateId()}@deleted-identity.test`) + const newOwner = await insertUser(email) + const organizationId = generateId() + organizationIds.push(organizationId) + const now = new Date() + await db.insert(organization).values({ + id: organizationId, + name: 'Transfer fixture', + slug: `transfer-${organizationId}`, + createdAt: now, + updatedAt: now, + }) + await db.insert(member).values([ + { id: generateId(), organizationId, userId: owner, role: 'owner', createdAt: now }, + { id: generateId(), organizationId, userId: newOwner, role: 'member', createdAt: now }, + ]) + + const result = await transferOrganizationOwnership({ + organizationId, + currentOwnerUserId: owner, + newOwnerUserId: newOwner, + }) + + expect(result).toMatchObject({ success: true }) + expect(await statsOf(newOwner)).toEqual({ freeTierEligible: false, limit: 0 }) + }) + + it('leaves an existing stats row untouched on a returning identity', async () => { + const email = `existing.${generateId()}@deleted-identity.test` + await deleteOwnAccount(await signUp(email)) + const returning = await insertUser(email) + await db.insert(userStats).values({ id: generateId(), userId: returning }) + + const created = await db.transaction((tx) => createUserStatsIfMissing(tx, returning)) + + expect(created).toBe(false) + expect(await statsOf(returning)).toMatchObject({ freeTierEligible: true }) + }) + + it('fails open without aborting the caller’s transaction when the identity lookup fails', async () => { + const email = `lookup-fails.${generateId()}@deleted-identity.test` + await deleteOwnAccount(await signUp(email)) + const returning = await insertUser(email) + const locked = createDeferred() + const release = createDeferred() + const holder = db.transaction(async (tx) => { + await tx.execute(sql`LOCK TABLE deleted_account_identity IN ACCESS EXCLUSIVE MODE`) + locked.resolve() + await release.promise + }) + await locked.promise + + let created: boolean + try { + created = await db.transaction(async (tx) => { + await tx.execute(sql`SET LOCAL lock_timeout = '200ms'`) + const result = await createUserStatsIfMissing(tx, returning) + await tx + .update(userStats) + .set({ storageUsedBytes: 7 }) + .where(eq(userStats.userId, returning)) + return result + }) + } finally { + release.resolve() + await holder + } + + expect(created).toBe(true) + expect(await statsOf(returning)).toEqual({ freeTierEligible: true, limit: getFreeTierLimit() }) + const [stored] = await db + .select({ bytes: userStats.storageUsedBytes }) + .from(userStats) + .where(eq(userStats.userId, returning)) + expect(stored.bytes).toBe(7) + }) + it('never withholds from a stats-less user’s fallback row for a sign-in linked in a later request', async () => { const microsoft = { providerId: 'microsoft', accountId: `ms-${generateId()}` } await deleteOwnAccount(await signUp(`ms-gone.${generateId()}@deleted-identity.test`, microsoft)) diff --git a/apps/sim/lib/users/deleted-identity.ts b/apps/sim/lib/users/deleted-identity.ts index c84aa096f78..e6973630529 100644 --- a/apps/sim/lib/users/deleted-identity.ts +++ b/apps/sim/lib/users/deleted-identity.ts @@ -214,6 +214,25 @@ export async function recordDeletedAccountIdentities(tx: DbOrTx, userId: string) } } +async function findDeletedIdentityMatch( + executor: DbOrTx, + identity: SignInIdentity +): Promise { + const hashes = identityHashes(identity) + if (hashes.length === 0) return false + const [match] = await executor + .select({ identityHash: deletedAccountIdentity.identityHash }) + .from(deletedAccountIdentity) + .where( + and( + inArray(deletedAccountIdentity.identityHash, hashes), + gt(deletedAccountIdentity.expiresAt, sql`now()`) + ) + ) + .limit(1) + return Boolean(match) +} + /** * Whether any of these identities belonged to an account deleted within the * retention window. Fails open: a lookup error returns `false` and is logged, @@ -222,19 +241,7 @@ export async function recordDeletedAccountIdentities(tx: DbOrTx, userId: string) export async function isDeletedAccountIdentity(identity: SignInIdentity): Promise { if (!isBillingEnabled) return false try { - const hashes = identityHashes(identity) - if (hashes.length === 0) return false - const [match] = await db - .select({ identityHash: deletedAccountIdentity.identityHash }) - .from(deletedAccountIdentity) - .where( - and( - inArray(deletedAccountIdentity.identityHash, hashes), - gt(deletedAccountIdentity.expiresAt, sql`now()`) - ) - ) - .limit(1) - return Boolean(match) + return await findDeletedIdentityMatch(db, identity) } catch (error) { logger.error('Deleted-identity lookup failed; granting the free allowance', { error }) return false @@ -246,16 +253,22 @@ export async function isDeletedAccountIdentity(identity: SignInIdentity): Promis * the retention window: always its email, and its linked social sign-ins only * when `includeSignIns` is set, which callers do only for a user created in the * current signup request (a sign-in linked later never withholds anything). - * Fails open like {@link isDeletedAccountIdentity}. + * Pass the caller's transaction as `executor` to read on it; the read then runs + * in a savepoint, so a failed lookup never aborts that transaction. Fails open + * like {@link isDeletedAccountIdentity}. */ export async function userHasDeletedAccountIdentity( userId: string, - { includeSignIns }: { includeSignIns: boolean } + { includeSignIns }: { includeSignIns: boolean }, + executor: DbOrTx = db ): Promise { if (!isBillingEnabled) return false + const read = async (reader: DbOrTx) => { + const identity = await loadSignInIdentity(reader, userId) + return findDeletedIdentityMatch(reader, includeSignIns ? identity : { email: identity.email }) + } try { - const identity = await loadSignInIdentity(db, userId) - return await isDeletedAccountIdentity(includeSignIns ? identity : { email: identity.email }) + return await (executor === db ? read(db) : executor.transaction(read)) } catch (error) { logger.error('Deleted-identity lookup failed; granting the free allowance', { userId, error }) return false diff --git a/apps/sim/lib/users/storage-reference-guard.ts b/apps/sim/lib/users/storage-reference-guard.ts index bf55b46af8f..d95991fb6dc 100644 --- a/apps/sim/lib/users/storage-reference-guard.ts +++ b/apps/sim/lib/users/storage-reference-guard.ts @@ -40,6 +40,7 @@ const SERVE_PATH = '/api/files/serve/' /** * The storage keys a stored URL can name: a `/api/files/serve/...` path (relative or absolute, * with or without a provider segment), or a direct bucket URL in virtual-host or path style. + * A relative path is also read without its `#fragment`, which a key may hold literally. * Over-matching only keeps an object. Never throws: a value whose percent-encoding is malformed * yields its raw path and that path with each valid escape decoded (a legacy key can hold a * literal `%`), so one bad stored string cannot disable every purge. @@ -72,10 +73,13 @@ function keysNamedByUrl(value: string): string[] { const serveAt = raw.indexOf(SERVE_PATH) path = serveAt >= 0 ? raw.slice(serveAt + SERVE_PATH.length) : raw } - path = path.replace(/^\/+/, '') - add(path) - const slash = path.indexOf('/') - if (slash > 0) add(path.slice(slash + 1)) + const fragmentAt = path.indexOf('#') + for (const candidate of fragmentAt >= 0 ? [path, path.slice(0, fragmentAt)] : [path]) { + const key = candidate.replace(/^\/+/, '') + add(key) + const slash = key.indexOf('/') + if (slash > 0) add(key.slice(slash + 1)) + } return [...keys] } diff --git a/apps/sim/lib/users/storage-sweep.test.ts b/apps/sim/lib/users/storage-sweep.test.ts index 6d702a43c01..7d25cc24c89 100644 --- a/apps/sim/lib/users/storage-sweep.test.ts +++ b/apps/sim/lib/users/storage-sweep.test.ts @@ -524,6 +524,26 @@ describe('account storage object cleanup', () => { expect(await exists('workspace-logos/wsforeign9/icon.png')).toBe(true) }) + it.each([ + { + key: 'workspace-logos/wsforeign9/icon.png', + url: '/api/files/serve/workspace-logos%2Fwsforeign9%2Ficon.png#v2', + }, + { + key: 'workspace-logos/wsforeign9/brand logo%.png', + url: '/api/files/serve/workspace-logos/wsforeign9/brand%20logo%.png#v2', + }, + ])('keeps a public key a relative URL with a fragment names: $url', async ({ key, url }) => { + const storedAs = key.replace(/[ %]/g, '_') + await put(storedAs) + queueTableRows(schemaMock.user, []) + queueTableRows(schemaMock.user, [{ value: url }]) + + await cleanup(payload([{ key, context: 'workspace-logos' }]), outboxContext()) + + expect(await exists(storedAs)).toBe(true) + }) + it('keeps a public key without an owner id to narrow on, reading no URL column', async () => { await put('og-images/logo.png') diff --git a/packages/testing/src/mocks/execution-preprocessing.mock.ts b/packages/testing/src/mocks/execution-preprocessing.mock.ts index cd926d63bdc..a5faaaa2ab8 100644 --- a/packages/testing/src/mocks/execution-preprocessing.mock.ts +++ b/packages/testing/src/mocks/execution-preprocessing.mock.ts @@ -29,4 +29,7 @@ export const executionPreprocessingMockFns = { export const executionPreprocessingMock = { preprocessExecution: executionPreprocessingMockFns.mockPreprocessExecution, WORKFLOW_NOT_DEPLOYED_CODE: 'WORKFLOW_NOT_DEPLOYED', + RATE_LIMIT_EXCEEDED_CODE: 'RATE_LIMIT_EXCEEDED', + isWorkflowStateRejection: (error: { workflowState?: string }) => + error.workflowState !== undefined, }