Skip to content

Commit 15f96d1

Browse files
committed
fix(billing): report purchased credit usage to Stripe
1 parent 4b7bc16 commit 15f96d1

File tree

11 files changed

+199
-17
lines changed

11 files changed

+199
-17
lines changed

common/src/types/contracts/database.ts

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,9 +6,17 @@ type User = {
66
email: string
77
discord_id: string | null
88
referral_code: string | null
9+
stripe_customer_id: string | null
910
banned: boolean
1011
}
11-
export const userColumns = ['id', 'email', 'discord_id', 'referral_code', 'banned'] as const
12+
export const userColumns = [
13+
'id',
14+
'email',
15+
'discord_id',
16+
'referral_code',
17+
'stripe_customer_id',
18+
'banned',
19+
] as const
1220
export type UserColumn = keyof User
1321
export type GetUserInfoFromApiKeyInput<T extends UserColumn> = {
1422
apiKey: string

evals/impl/agent-runtime.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ export const EVALS_AGENT_RUNTIME_IMPL = Object.freeze<AgentRuntimeDeps>({
3737
email: 'test-email',
3838
discord_id: 'test-discord-id',
3939
referral_code: 'ref-test-code',
40+
stripe_customer_id: null,
4041
banned: false,
4142
}),
4243
fetchAgentFromDatabase: async () => null,

packages/billing/src/balance-calculator.ts

Lines changed: 34 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@ import * as schema from '@codebuff/internal/db/schema'
66
import { withSerializableTransaction } from '@codebuff/internal/db/transaction'
77
import { and, asc, gt, isNull, or, eq, sql } from 'drizzle-orm'
88

9+
import { reportPurchasedCreditsToStripe } from './stripe-metering'
10+
911
import type { Logger } from '@codebuff/common/types/contracts/logger'
1012
import type {
1113
ParamsExcluding,
@@ -330,12 +332,13 @@ export async function calculateUsageAndBalance(
330332
*/
331333
export async function consumeCredits(params: {
332334
userId: string
335+
stripeCustomerId?: string | null
333336
creditsToConsume: number
334337
logger: Logger
335338
}): Promise<CreditConsumptionResult> {
336339
const { userId, creditsToConsume, logger } = params
337340

338-
return await withSerializableTransaction({
341+
const result = await withSerializableTransaction({
339342
callback: async (tx) => {
340343
const now = new Date()
341344
const activeGrants = await getOrderedActiveGrants({
@@ -364,11 +367,24 @@ export async function consumeCredits(params: {
364367
context: { userId, creditsToConsume },
365368
logger,
366369
})
370+
371+
await reportPurchasedCreditsToStripe({
372+
userId,
373+
stripeCustomerId: params.stripeCustomerId,
374+
purchasedCredits: result.fromPurchased,
375+
logger,
376+
extraPayload: {
377+
source: 'consumeCredits',
378+
},
379+
})
380+
381+
return result
367382
}
368383

369384
export async function consumeCreditsAndAddAgentStep(params: {
370385
messageId: string
371386
userId: string
387+
stripeCustomerId?: string | null
372388
agentId: string
373389
clientId: string | null
374390
clientRequestId: string | null
@@ -421,8 +437,7 @@ export async function consumeCreditsAndAddAgentStep(params: {
421437
const latencyMs = finishedAt.getTime() - startTime.getTime()
422438

423439
try {
424-
return success(
425-
await withSerializableTransaction({
440+
const result = await withSerializableTransaction({
426441
callback: async (tx) => {
427442
const now = new Date()
428443

@@ -493,8 +508,22 @@ export async function consumeCreditsAndAddAgentStep(params: {
493508
},
494509
context: { userId, credits },
495510
logger,
496-
}),
497-
)
511+
})
512+
513+
await reportPurchasedCreditsToStripe({
514+
userId,
515+
stripeCustomerId: params.stripeCustomerId,
516+
purchasedCredits: result.fromPurchased,
517+
logger,
518+
eventId: messageId,
519+
timestamp: finishedAt,
520+
extraPayload: {
521+
source: 'consumeCreditsAndAddAgentStep',
522+
message_id: messageId,
523+
},
524+
})
525+
526+
return success(result)
498527
} catch (error) {
499528
logger.error(
500529
{ error: getErrorObject(error) },
Lines changed: 120 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,120 @@
1+
import { TEST_USER_ID } from '@codebuff/common/old-constants'
2+
import { withRetry, withTimeout } from '@codebuff/common/util/promise'
3+
import db from '@codebuff/internal/db'
4+
import * as schema from '@codebuff/internal/db/schema'
5+
import { stripeServer } from '@codebuff/internal/util/stripe'
6+
import { eq } from 'drizzle-orm'
7+
8+
import type { Logger } from '@codebuff/common/types/contracts/logger'
9+
10+
const STRIPE_METER_EVENT_NAME = 'credits'
11+
const STRIPE_METER_REQUEST_TIMEOUT_MS = 10_000
12+
13+
function shouldAttemptStripeMetering(): boolean {
14+
// Avoid sending Stripe metering events in CI/tests, and when Stripe isn't configured.
15+
if (process.env.CI === 'true' || process.env.CI === '1') return false
16+
if (process.env.NODE_ENV === 'test') return false
17+
return Boolean(process.env.STRIPE_SECRET_KEY)
18+
}
19+
20+
export async function reportPurchasedCreditsToStripe(params: {
21+
userId: string
22+
stripeCustomerId?: string | null
23+
purchasedCredits: number
24+
logger: Logger
25+
/**
26+
* Optional unique identifier used for Stripe idempotency + debugging.
27+
* For message-based usage, pass the message ID.
28+
*/
29+
eventId?: string
30+
/**
31+
* Optional timestamp for the usage event.
32+
* Defaults to "now".
33+
*/
34+
timestamp?: Date
35+
/**
36+
* Optional additional payload fields (must be strings).
37+
*/
38+
extraPayload?: Record<string, string>
39+
}): Promise<void> {
40+
const {
41+
userId,
42+
stripeCustomerId: providedStripeCustomerId,
43+
purchasedCredits,
44+
logger,
45+
eventId,
46+
timestamp = new Date(),
47+
extraPayload,
48+
} = params
49+
50+
if (purchasedCredits <= 0) return
51+
if (userId === TEST_USER_ID) return
52+
if (!shouldAttemptStripeMetering()) return
53+
54+
const logContext = { userId, purchasedCredits, eventId }
55+
56+
let stripeCustomerId = providedStripeCustomerId
57+
if (stripeCustomerId === undefined) {
58+
let user: { stripe_customer_id: string | null } | undefined
59+
try {
60+
user = await db.query.user.findFirst({
61+
where: eq(schema.user.id, userId),
62+
columns: { stripe_customer_id: true },
63+
})
64+
} catch (error) {
65+
logger.error(
66+
{ ...logContext, error },
67+
'Failed to fetch user for Stripe metering',
68+
)
69+
return
70+
}
71+
72+
stripeCustomerId = user?.stripe_customer_id ?? null
73+
}
74+
if (!stripeCustomerId) {
75+
logger.warn(logContext, 'Skipping Stripe metering (missing stripe_customer_id)')
76+
return
77+
}
78+
79+
const stripeTimestamp = Math.floor(timestamp.getTime() / 1000)
80+
const idempotencyKey = eventId ? `meter-${eventId}` : undefined
81+
82+
try {
83+
await withTimeout(
84+
withRetry(
85+
() =>
86+
stripeServer.billing.meterEvents.create(
87+
{
88+
event_name: STRIPE_METER_EVENT_NAME,
89+
timestamp: stripeTimestamp,
90+
payload: {
91+
stripe_customer_id: stripeCustomerId,
92+
value: purchasedCredits.toString(),
93+
...(eventId ? { event_id: eventId } : {}),
94+
...(extraPayload ?? {}),
95+
},
96+
},
97+
idempotencyKey ? { idempotencyKey } : undefined,
98+
),
99+
{
100+
maxRetries: 3,
101+
retryIf: (error: any) =>
102+
error?.type === 'StripeConnectionError' ||
103+
error?.type === 'StripeAPIError' ||
104+
error?.type === 'StripeRateLimitError',
105+
onRetry: (error: any, attempt: number) => {
106+
logger.warn(
107+
{ ...logContext, attempt, error },
108+
'Retrying Stripe metering call',
109+
)
110+
},
111+
retryDelayMs: 500,
112+
},
113+
),
114+
STRIPE_METER_REQUEST_TIMEOUT_MS,
115+
`Stripe metering timed out after ${STRIPE_METER_REQUEST_TIMEOUT_MS}ms`,
116+
)
117+
} catch (error) {
118+
logger.error({ ...logContext, error }, 'Failed to report purchased credits to Stripe')
119+
}
120+
}

sdk/src/__tests__/run-handle-event.test.ts

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,10 @@
11
import { afterEach, describe, expect, it, mock, spyOn } from 'bun:test'
22

3+
import * as mainPromptModule from '@codebuff/agent-runtime/main-prompt'
34
import { getInitialSessionState } from '@codebuff/common/types/session-state'
45
import { getStubProjectFileContext } from '@codebuff/common/util/file'
6+
import { CodebuffClient } from '../client'
7+
import * as databaseModule from '../impl/database'
58
import type { PrintModeEvent } from '@codebuff/common/types/print-mode'
69
import type { CodebuffClientOptions } from '../run'
710

@@ -11,14 +14,12 @@ describe('CodebuffClient handleEvent / handleStreamChunk', () => {
1114
})
1215

1316
it('streams subagent start/finish once and forwards subagent chunks to handleStreamChunk', async () => {
14-
const databaseModule = await import('../impl/database')
15-
const mainPromptModule = await import('@codebuff/agent-runtime/main-prompt')
16-
1717
spyOn(databaseModule, 'getUserInfoFromApiKey').mockResolvedValue({
1818
id: 'user-123',
1919
email: 'test@example.com',
2020
discord_id: null,
2121
referral_code: null,
22+
stripe_customer_id: null,
2223
banned: false,
2324
})
2425
spyOn(databaseModule, 'fetchAgentFromDatabase').mockResolvedValue(null)
@@ -108,8 +109,6 @@ describe('CodebuffClient handleEvent / handleStreamChunk', () => {
108109
const events: PrintModeEvent[] = []
109110
const streamChunks: StreamChunk[] = []
110111

111-
const { CodebuffClient } = await import('../client')
112-
113112
const client = new CodebuffClient({
114113
apiKey: 'test-key',
115114
})

web/src/app/api/v1/chat/completions/_post.ts

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -128,7 +128,7 @@ export async function postChatCompletions(params: {
128128
// Get user info
129129
const userInfo = await getUserInfoFromApiKey({
130130
apiKey,
131-
fields: ['id', 'email', 'discord_id', 'banned'],
131+
fields: ['id', 'email', 'discord_id', 'stripe_customer_id', 'banned'],
132132
logger,
133133
})
134134
if (!userInfo) {
@@ -148,6 +148,7 @@ export async function postChatCompletions(params: {
148148
logger = loggerWithContext({ userInfo })
149149

150150
const userId = userInfo.id
151+
const stripeCustomerId = userInfo.stripe_customer_id ?? null
151152

152153
// Check if user is banned.
153154
// We use a clear, helpful message rather than a cryptic error because:
@@ -269,6 +270,7 @@ export async function postChatCompletions(params: {
269270
const stream = await handleOpenRouterStream({
270271
body,
271272
userId,
273+
stripeCustomerId,
272274
agentId,
273275
openrouterApiKey,
274276
fetch,
@@ -312,6 +314,7 @@ export async function postChatCompletions(params: {
312314
? handleOpenAINonStream({
313315
body,
314316
userId,
317+
stripeCustomerId,
315318
agentId,
316319
fetch,
317320
logger,
@@ -320,6 +323,7 @@ export async function postChatCompletions(params: {
320323
: handleOpenRouterNonStream({
321324
body,
322325
userId,
326+
stripeCustomerId,
323327
agentId,
324328
openrouterApiKey,
325329
fetch,

web/src/app/api/v1/me/__tests__/me.test.ts

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -23,13 +23,15 @@ describe('/api/v1/me route', () => {
2323
email: 'test@example.com',
2424
discord_id: 'discord-123',
2525
referral_code: 'ref-user-123',
26+
stripe_customer_id: 'cus_test_123',
2627
banned: false,
2728
},
2829
'test-api-key-456': {
2930
id: 'user-456',
3031
email: 'test2@example.com',
3132
discord_id: null,
3233
referral_code: 'ref-user-456',
34+
stripe_customer_id: null,
3335
banned: false,
3436
},
3537
}
@@ -44,7 +46,7 @@ describe('/api/v1/me route', () => {
4446
return null
4547
}
4648
return Object.fromEntries(
47-
fields.map((field) => [field, userData[field]]),
49+
fields.map((field) => [field, (userData as any)[field]]),
4850
) as any
4951
},
5052
}
@@ -212,7 +214,7 @@ describe('/api/v1/me route', () => {
212214
const body = await response.json()
213215
expect(body.error).toContain('Invalid fields: invalid_field')
214216
expect(body.error).toContain(
215-
'Valid fields are: id, email, discord_id, referral_code, referral_link',
217+
'Valid fields are: id, email, discord_id, referral_code, stripe_customer_id, banned, referral_link',
216218
)
217219
})
218220

web/src/db/user.ts

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -5,19 +5,19 @@ import { eq } from 'drizzle-orm'
55
import type {
66
GetUserInfoFromApiKeyInput,
77
GetUserInfoFromApiKeyOutput,
8+
UserColumn,
89
} from '@codebuff/common/types/contracts/database'
910

1011
export const VALID_USER_INFO_FIELDS = [
1112
'id',
1213
'email',
1314
'discord_id',
1415
'referral_code',
16+
'stripe_customer_id',
1517
'banned',
1618
] as const
1719

18-
export async function getUserInfoFromApiKey<
19-
T extends (typeof VALID_USER_INFO_FIELDS)[number],
20-
>({
20+
export async function getUserInfoFromApiKey<T extends UserColumn>({
2121
apiKey,
2222
fields,
2323
}: GetUserInfoFromApiKeyInput<T>): GetUserInfoFromApiKeyOutput<T> {

web/src/llm-api/helpers.ts

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,7 @@ export async function insertMessageToBigQuery(params: {
9090
export async function consumeCreditsForMessage(params: {
9191
messageId: string
9292
userId: string
93+
stripeCustomerId?: string | null
9394
agentId: string
9495
clientId: string | null
9596
clientRequestId: string | null
@@ -104,6 +105,7 @@ export async function consumeCreditsForMessage(params: {
104105
const {
105106
messageId,
106107
userId,
108+
stripeCustomerId,
107109
agentId,
108110
clientId,
109111
clientRequestId,
@@ -119,6 +121,7 @@ export async function consumeCreditsForMessage(params: {
119121
await consumeCreditsAndAddAgentStep({
120122
messageId,
121123
userId,
124+
stripeCustomerId,
122125
agentId,
123126
clientId,
124127
clientRequestId,

0 commit comments

Comments
 (0)