Skip to content
Source
Sign inSign up

usageLimits.ts

204 lines7.8 KB

1import type { ServerProviderUsageWindow } from "@t3tools/contracts";
2import * as DateTime from "effect/DateTime";
3import * as Effect from "effect/Effect";
4import * as FileSystem from "effect/FileSystem";
5import * as Option from "effect/Option";
6import * as Path from "effect/Path";
7import * as Schema from "effect/Schema";
8import * as HttpClient from "effect/http/HttpClient";
9import * as HttpClientRequest from "effect/http/HttpClientRequest";
10import * as HttpClientResponse from "effect/http/HttpClientResponse";
11import {
12 clampPercent,
13 makeUnavailableUsageLimits,
14 makeUsageLimits,
15} from "@t3tools/provider-core/server/usageLimits";
16import * as HostProcess from "@t3tools/shared/HostProcess";
17​
18const GrokCredentials = Schema.Record(
19 Schema.String,
20 Schema.Struct({
21 key: Schema.optional(Schema.String),
22 auth_mode: Schema.optional(Schema.String),
23 email: Schema.optional(Schema.String),
24 }),
25);
26const decodeCredentials = Schema.decodeEffect(Schema.fromJsonString(GrokCredentials));
27export const GROK_USAGE_URL = "https://cli-chat-proxy.grok.com/v1/billing?format=credits";
28export const GrokUsageResponse = Schema.Struct({
29 config: Schema.optional(
30 Schema.Struct({
31 creditUsagePercent: Schema.optional(Schema.Number),
32 currentPeriod: Schema.optional(
33 Schema.Struct({
34 type: Schema.optional(Schema.String),
35 start: Schema.optional(Schema.Unknown),
36 end: Schema.optional(Schema.String),
37 }),
38 ),
39 }),
40 ),
41});
42const GrokErrorResponse = Schema.Struct({ error: Schema.String });
43​
44function parsePeriodBoundary(value: string | undefined) {
45 if (
46 !value ||
47 value !== value.trim() ||
48 !/^(?:\d{4}|[+-]\d{6})-(?:0[1-9]|1[0-2])-(?:0[1-9]|[12]\d|3[01])T(?:(?:[01]\d|2[0-3]):[0-5]\d(?::[0-5]\d(?:\.\d+)?)?|24:00(?::00(?:\.0+)?)?)(?:Z|[+-](?:[01]\d|2[0-3]):[0-5]\d)$/.test(
49 value,
50 )
51 ) {
52 return Number.NaN;
53 }
54 // Date parsing can normalize impossible calendar dates instead of rejecting them.
55 const datePart = value.slice(0, value.indexOf("T"));
56 const date = DateTime.make(`${datePart}T00:00:00.000Z`);
57 if (Option.isNone(date)) return Number.NaN;
58 const parts = DateTime.toPartsUtc(date.value);
59 if (parts.month !== Number(datePart.slice(-5, -3)) || parts.day !== Number(datePart.slice(-2))) {
60 return Number.NaN;
61 }
62 return Date.parse(value);
64​
65export function grokUsageResponseToLimits(
66 response: typeof GrokUsageResponse.Type,
67 checkedAt: string,
68) {
69 const period = response.config?.currentPeriod;
70 const periodType = period?.type?.replace(/^USAGE_PERIOD_TYPE_/, "");
71 const kind = periodType === "WEEKLY" ? "weekly" : periodType === "MONTHLY" ? "monthly" : "other";
72 const start = typeof period?.start === "string" ? parsePeriodBoundary(period.start) : NaN;
73 const end = parsePeriodBoundary(period?.end);
74 const windowDurationMins = Math.round((end - start) / 60_000);
75 let usedPercent = response.config?.creditUsagePercent;
76 if (usedPercent === undefined && (kind === "weekly" || kind === "monthly") && start < end) {
77 // Grok omits a zero percentage while retaining the complete subscription period.
78 usedPercent = 0;
79 }
80 if (usedPercent === undefined || !Number.isFinite(usedPercent)) {
81 return makeUsageLimits({ checkedAt, windows: [] });
82 }
83 const reset = period?.end ? DateTime.make(period.end) : Option.none();
84 const window: ServerProviderUsageWindow = {
85 id: "subscription",
86 kind,
87 label: kind === "weekly" ? "Weekly" : kind === "monthly" ? "Monthly" : "Subscription",
88 usedPercent: clampPercent(usedPercent),
89 ...(Option.isSome(reset) ? { resetsAt: DateTime.formatIso(reset.value) } : {}),
90 ...(windowDurationMins > 0 ? { windowDurationMins } : {}),
91 };
92 return makeUsageLimits({ checkedAt, windows: [window] });
94​
95const readGrokCredential = Effect.fn("readGrokCredential")(function* (
96 environment: NodeJS.ProcessEnv,
97) {
98 // T3's ACP adapter explicitly selects API-key auth when this variable is set.
99 if (environment.XAI_API_KEY?.trim()) return undefined;
100 // Alternate auth deployments can select another scope or account from the same file.
101 if (
102 [
103 "GROK_OIDC_ISSUER",
104 "GROK_OIDC_CLIENT_ID",
105 "GROK_OAUTH2_ISSUER",
106 "GROK_OAUTH2_CLIENT_ID",
107 "GROK_OAUTH2_PRINCIPAL_TYPE",
108 "GROK_OAUTH2_PRINCIPAL_ID",
109 "GROK_AUTH_PROVIDER_COMMAND",
110 "GROK_LOCAL_AUTH",
111 "GROK_CLI_CHAT_PROXY_BASE_URL",
112 "GROK_MODELS_BASE_URL",
113 "GROK_CONFIG",
114 "GROK_CONFIG_PATH",
115 ].some((name) => environment[name]?.trim())
116 ) {
117 return undefined;
118 }
119 const fs = yield* FileSystem.FileSystem;
120 const path = yield* Path.Path;
121 const home =
122 environment.GROK_HOME?.trim() ||
123 path.join(
124 environment.HOME || environment.USERPROFILE || (yield* HostProcess.HomeDirectory),
125 ".grok",
126 );
127 for (const configPath of [
128 path.join(home, "config.toml"),
129 path.join(home, "managed_config.toml"),
130 path.join(home, "requirements.toml"),
131 "/etc/grok/managed_config.toml",
132 "/etc/grok/requirements.toml",
133 ]) {
134 const config = yield* fs.readFileString(configPath).pipe(
135 Effect.catchTags({
136 PlatformError: (error) =>
137 error.reason._tag === "NotFound" ? Effect.succeed("") : Effect.fail(error),
138 }),
139 );
140 // These sections can change the selected account or endpoint. Leave custom deployments to the CLI.
141 if (/^\s*(?:\[\[?\s*)?["']?(?:auth|grok_com_config|endpoints)["']?\s*[.\]=]/m.test(config)) {
142 return undefined;
143 }
144 }
145 const contents =
146 environment.GROK_AUTH?.trim() ||
147 (yield* fs.readFileString(path.join(home, "auth.json")).pipe(
148 Effect.catchTags({
149 PlatformError: (error) =>
150 error.reason._tag === "NotFound" ? Effect.succeed("{}") : Effect.fail(error),
151 }),
152 ));
153 const credentials = yield* decodeCredentials(contents);
154 // Never pick an arbitrary account from other deployments stored in the same file.
155 const credential =
156 credentials["https://auth.x.ai::b1a00492-073a-47ea-816f-4c329264a828"] ??
157 credentials["https://accounts.x.ai/sign-in"];
158 return credential?.auth_mode === "api_key" ? undefined : credential;
159});
160​
161export const readGrokAccount = Effect.fn("readGrokAccount")(function* (
162 environment: NodeJS.ProcessEnv = process.env,
163) {
164 const checkedAt = DateTime.formatIso(yield* DateTime.now);
165 const probeFailed = makeUnavailableUsageLimits({
166 checkedAt,
167 reason: "probeFailed",
168 message: "Grok could not read usage limits.",
169 });
170 const credential = yield* Effect.option(
171 readGrokCredential(environment).pipe(Effect.timeout("10 seconds")),
172 );
173 if (Option.isNone(credential)) return { email: undefined, usageLimits: probeFailed };
174 const email = credential.value?.email?.trim() || undefined;
175 const token = credential.value?.key?.trim();
176 if (!token) {
177 return { email, usageLimits: makeUnavailableUsageLimits({ checkedAt, reason: "unsupported" }) };
178 }
179 const usageLimits = yield* Effect.gen(function* () {
180 const client = yield* HttpClient.HttpClient;
181 const response = yield* client.execute(
182 HttpClientRequest.get(GROK_USAGE_URL).pipe(HttpClientRequest.bearerToken(token)),
183 );
184 if (response.status === 403) {
185 const { error } = yield* HttpClientResponse.schemaBodyJson(GrokErrorResponse)(response);
186 // xAI refuses Grok Build while the account's privacy setting is a legacy mode.
187 if (error.includes("WKE=privacy-mode:")) {
188 return makeUnavailableUsageLimits({
189 checkedAt,
190 reason: "probeFailed",
191 message: "Switch your xAI account to Privacy Mode to show Grok usage limits.",
192 });
193 }
194 }
195 const body = yield* HttpClientResponse.schemaBodyJson(GrokUsageResponse)(
196 yield* HttpClientResponse.filterStatusOk(response),
197 );
198 return grokUsageResponseToLimits(body, checkedAt);
199 }).pipe(
200 Effect.timeout("10 seconds"),
201 Effect.orElseSucceed(() => probeFailed),
202 );
203 return { email, usageLimits };
204});