billing.ts 2.1 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071
  1. import { Resource } from "sst"
  2. import { Stripe } from "stripe"
  3. import { Database, eq, sql } from "./drizzle"
  4. import { BillingTable, UsageTable } from "./schema/billing.sql"
  5. import { Actor } from "./actor"
  6. import { fn } from "./util/fn"
  7. import { z } from "zod"
  8. import { Identifier } from "./identifier"
  9. import { centsToMicroCents } from "./util/price"
  10. export namespace Billing {
  11. export const stripe = () =>
  12. new Stripe(Resource.STRIPE_SECRET_KEY.value, {
  13. apiVersion: "2025-03-31.basil",
  14. })
  15. export const get = async () => {
  16. return Database.use(async (tx) =>
  17. tx
  18. .select({
  19. customerID: BillingTable.customerID,
  20. paymentMethodID: BillingTable.paymentMethodID,
  21. balance: BillingTable.balance,
  22. reload: BillingTable.reload,
  23. })
  24. .from(BillingTable)
  25. .where(eq(BillingTable.workspaceID, Actor.workspace()))
  26. .then((r) => r[0]),
  27. )
  28. }
  29. export const consume = fn(
  30. z.object({
  31. requestID: z.string().optional(),
  32. model: z.string(),
  33. inputTokens: z.number(),
  34. outputTokens: z.number(),
  35. reasoningTokens: z.number().optional(),
  36. cacheReadTokens: z.number().optional(),
  37. cacheWriteTokens: z.number().optional(),
  38. costInCents: z.number(),
  39. }),
  40. async (input) => {
  41. const workspaceID = Actor.workspace()
  42. const cost = centsToMicroCents(input.costInCents)
  43. return await Database.transaction(async (tx) => {
  44. await tx.insert(UsageTable).values({
  45. workspaceID,
  46. id: Identifier.create("usage"),
  47. requestID: input.requestID,
  48. model: input.model,
  49. inputTokens: input.inputTokens,
  50. outputTokens: input.outputTokens,
  51. reasoningTokens: input.reasoningTokens,
  52. cacheReadTokens: input.cacheReadTokens,
  53. cacheWriteTokens: input.cacheWriteTokens,
  54. cost,
  55. })
  56. const [updated] = await tx
  57. .update(BillingTable)
  58. .set({
  59. balance: sql`${BillingTable.balance} - ${cost}`,
  60. })
  61. .where(eq(BillingTable.workspaceID, workspaceID))
  62. .returning()
  63. return updated.balance
  64. })
  65. },
  66. )
  67. }