summaryrefslogtreecommitdiffhomepage
path: root/packages/console/core/src/black.ts
blob: 753d25808f31c5a1c7344d73281b8b64f5b5ad42 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
import { z } from "zod"
import { fn } from "./util/fn"
import { Resource } from "@opencode-ai/console-resource"
import { centsToMicroCents } from "./util/price"
import { getWeekBounds } from "./util/date"

export namespace BlackData {
  const Schema = z.object({
    fixedLimit: z.number().int(),
    rollingLimit: z.number().int(),
    rollingWindow: z.number().int(),
  })

  export const validate = fn(Schema, (input) => {
    return input
  })

  export const get = fn(z.void(), () => {
    const json = JSON.parse(Resource.ZEN_BLACK.value)
    return Schema.parse(json)
  })
}

export namespace Black {
  export const analyzeRollingUsage = fn(
    z.object({
      usage: z.number().int(),
      timeUpdated: z.date(),
    }),
    ({ usage, timeUpdated }) => {
      const now = new Date()
      const black = BlackData.get()
      const rollingWindowMs = black.rollingWindow * 3600 * 1000
      const rollingLimitInMicroCents = centsToMicroCents(black.rollingLimit * 100)
      const windowStart = new Date(now.getTime() - rollingWindowMs)
      if (timeUpdated < windowStart) {
        return {
          status: "ok" as const,
          resetInSec: black.rollingWindow * 3600,
          usagePercent: 0,
        }
      }

      const windowEnd = new Date(timeUpdated.getTime() + rollingWindowMs)
      if (usage < rollingLimitInMicroCents) {
        return {
          status: "ok" as const,
          resetInSec: Math.ceil((windowEnd.getTime() - now.getTime()) / 1000),
          usagePercent: Math.ceil(Math.min(100, (usage / rollingLimitInMicroCents) * 100)),
        }
      }
      return {
        status: "rate-limited" as const,
        resetInSec: Math.ceil((windowEnd.getTime() - now.getTime()) / 1000),
        usagePercent: 100,
      }
    },
  )

  export const analyzeWeeklyUsage = fn(
    z.object({
      usage: z.number().int(),
      timeUpdated: z.date(),
    }),
    ({ usage, timeUpdated }) => {
      const black = BlackData.get()
      const now = new Date()
      const week = getWeekBounds(now)
      const fixedLimitInMicroCents = centsToMicroCents(black.fixedLimit * 100)
      if (timeUpdated < week.start) {
        return {
          status: "ok" as const,
          resetInSec: Math.ceil((week.end.getTime() - now.getTime()) / 1000),
          usagePercent: 0,
        }
      }
      if (usage < fixedLimitInMicroCents) {
        return {
          status: "ok" as const,
          resetInSec: Math.ceil((week.end.getTime() - now.getTime()) / 1000),
          usagePercent: Math.ceil(Math.min(100, (usage / fixedLimitInMicroCents) * 100)),
        }
      }

      return {
        status: "rate-limited" as const,
        resetInSec: Math.ceil((week.end.getTime() - now.getTime()) / 1000),
        usagePercent: 100,
      }
    },
  )
}