diff --git a/apps/sim/app/api/billing/update-cost/route.ts b/apps/sim/app/api/billing/update-cost/route.ts index 743ddb17011..f01ec13f939 100644 --- a/apps/sim/app/api/billing/update-cost/route.ts +++ b/apps/sim/app/api/billing/update-cost/route.ts @@ -6,6 +6,7 @@ import { recordUsage } from '@/lib/billing/core/usage-log' import { checkAndBillOverageThreshold } from '@/lib/billing/threshold-billing' import { checkInternalApiKey } from '@/lib/copilot/request/http' import { isBillingEnabled } from '@/lib/core/config/feature-flags' +import { type AtomicClaimResult, billingIdempotency } from '@/lib/core/idempotency/service' import { generateRequestId } from '@/lib/core/utils/request' const logger = createLogger('BillingUpdateCostAPI') @@ -19,6 +20,7 @@ const UpdateCostSchema = z.object({ source: z .enum(['copilot', 'workspace-chat', 'mcp_copilot', 'mothership_block']) .default('copilot'), + idempotencyKey: z.string().min(1).optional(), }) /** @@ -28,6 +30,8 @@ const UpdateCostSchema = z.object({ export async function POST(req: NextRequest) { const requestId = generateRequestId() const startTime = Date.now() + let claim: AtomicClaimResult | null = null + let usageCommitted = false try { logger.info(`[${requestId}] Update cost request started`) @@ -75,9 +79,30 @@ export async function POST(req: NextRequest) { ) } - const { userId, cost, model, inputTokens, outputTokens, source } = validation.data + const { userId, cost, model, inputTokens, outputTokens, source, idempotencyKey } = + validation.data const isMcp = source === 'mcp_copilot' + claim = idempotencyKey + ? await billingIdempotency.atomicallyClaim('update-cost', idempotencyKey) + : null + + if (claim && !claim.claimed) { + logger.warn(`[${requestId}] Duplicate billing update rejected`, { + idempotencyKey, + userId, + source, + }) + return NextResponse.json( + { + success: false, + error: 'Duplicate request: idempotency key already processed', + requestId, + }, + { status: 409 } + ) + } + logger.info(`[${requestId}] Processing cost update`, { userId, cost, @@ -113,6 +138,7 @@ export async function POST(req: NextRequest) { ], additionalStats, }) + usageCommitted = true logger.info(`[${requestId}] Recorded usage`, { userId, @@ -149,6 +175,22 @@ export async function POST(req: NextRequest) { duration, }) + if (claim?.claimed && !usageCommitted) { + await billingIdempotency + .release(claim.normalizedKey, claim.storageMethod) + .catch((releaseErr) => { + logger.warn(`[${requestId}] Failed to release idempotency claim`, { + error: releaseErr instanceof Error ? releaseErr.message : String(releaseErr), + normalizedKey: claim?.normalizedKey, + }) + }) + } else if (claim?.claimed && usageCommitted) { + logger.warn( + `[${requestId}] Error occurred after usage committed; retaining idempotency claim to prevent double-billing`, + { normalizedKey: claim.normalizedKey } + ) + } + return NextResponse.json( { success: false, diff --git a/apps/sim/lib/core/idempotency/service.ts b/apps/sim/lib/core/idempotency/service.ts index a96627bba34..b9920f49e81 100644 --- a/apps/sim/lib/core/idempotency/service.ts +++ b/apps/sim/lib/core/idempotency/service.ts @@ -343,6 +343,10 @@ export class IdempotencyService { logger.debug(`Stored idempotency result in database: ${normalizedKey}`) } + async release(normalizedKey: string, storageMethod: 'redis' | 'database'): Promise { + return this.deleteKey(normalizedKey, storageMethod) + } + private async deleteKey( normalizedKey: string, storageMethod: 'redis' | 'database' @@ -482,3 +486,8 @@ export const pollingIdempotency = new IdempotencyService({ ttlSeconds: 60 * 60 * 24 * 3, // 3 days retryFailures: true, }) + +export const billingIdempotency = new IdempotencyService({ + namespace: 'billing', + ttlSeconds: 60 * 60, // 1 hour +})