Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
44 changes: 43 additions & 1 deletion apps/sim/app/api/billing/update-cost/route.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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')
Expand All @@ -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(),
})

/**
Expand All @@ -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`)
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -113,6 +138,7 @@ export async function POST(req: NextRequest) {
],
additionalStats,
})
usageCommitted = true

logger.info(`[${requestId}] Recorded usage`, {
userId,
Expand Down Expand Up @@ -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,
Expand Down
9 changes: 9 additions & 0 deletions apps/sim/lib/core/idempotency/service.ts
Original file line number Diff line number Diff line change
Expand Up @@ -343,6 +343,10 @@ export class IdempotencyService {
logger.debug(`Stored idempotency result in database: ${normalizedKey}`)
}

async release(normalizedKey: string, storageMethod: 'redis' | 'database'): Promise<void> {
return this.deleteKey(normalizedKey, storageMethod)
}

private async deleteKey(
normalizedKey: string,
storageMethod: 'redis' | 'database'
Expand Down Expand Up @@ -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
})
Loading