Something went wrong. Try again.
zig library for building requests to query inference providers usage limits
Something went wrong. Try again.
7.9 kB · 186 lines
Zig
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187//! OpenAI Codex usage feed implementation.//!//! Mirrors prime-agent:packages/ai/src/utils/usage/providers/codex.ts field for field.const std = @import("std");const usage = @import("../usage.zig");
pub const CODEX_USAGE_URL = "https://chatgpt.com/backend-api/wham/usage";
fn parseWindow(alloc: std.mem.Allocator, id: []const u8, window_val: ?std.json.Value, now_ms: i64, scope: ?[]const u8) !?usage.UsageLimit { const value = usage.asObject(window_val) orelse return null; const used_percent = usage.asNumber(value.get("used_percent")); const reset_at = usage.asNumber(value.get("reset_at")); const reset_after = usage.asNumber(value.get("reset_after_seconds")); const limit_window = usage.asNumber(value.get("limit_window_seconds"));
const reset_ms: ?i64 = if (reset_at) |ra| (if (ra > 1_000_000_000_000.0) usage.toInt(i64, ra) else usage.toInt(i64, ra * 1000.0)) else if (reset_after) |after| blk: { const ms = usage.toInt(i64, after * 1000.0) orelse break :blk null; break :blk now_ms + ms; } else null;
const window_lbl = try usage.formatWindowLabel(alloc, limit_window); const window_seconds: ?i64 = if (limit_window) |lw| (if (lw > 0) usage.roundToInt(i64, lw) else null) else null;
return usage.UsageLimit{ .id = id, .label = "Usage", .usedPercent = used_percent, .windowLabel = window_lbl, .windowSeconds = window_seconds, .resetsAt = reset_ms, .scope = scope, .status = usage.statusFor(used_percent), };}
pub const ParseResult = struct { planType: ?[]const u8 = null, limits: []const usage.UsageLimit = &[_]usage.UsageLimit{},};
pub fn parseCodexUsage(alloc: std.mem.Allocator, payload: ?std.json.Value, now_ms: i64) !ParseResult { const root = usage.asObject(payload); const rate_limit = if (root) |r| usage.asObject(r.get("rate_limit")) else null; var limits = std.ArrayList(usage.UsageLimit).empty;
const primary_val = if (rate_limit) |rl| rl.get("primary_window") else null; if (try parseWindow(alloc, "primary", primary_val, now_ms, null)) |pw| { var w = pw; w.label = "Primary"; try limits.append(alloc, w); }
const secondary_val = if (rate_limit) |rl| rl.get("secondary_window") else null; if (try parseWindow(alloc, "secondary", secondary_val, now_ms, null)) |sw| { var w = sw; w.label = "Secondary"; try limits.append(alloc, w); }
const additional = if (root) |r| usage.asArray(r.get("additional_rate_limits")) else null; if (additional) |arr| { for (arr.items) |entry| { const item = usage.asObject(entry) orelse continue; const extra = usage.asObject(item.get("rate_limit")); const limit_name = usage.asString(item.get("limit_name")); const feature = usage.asString(item.get("metered_feature")); const name = limit_name orelse feature orelse "Additional"; // The metered feature is the machine key; the limit name is what it restricts. const key = feature orelse limit_name orelse "additional"; const scope = limit_name orelse feature;
const extra_primary = if (extra) |e| e.get("primary_window") else null; if (try parseWindow(alloc, try std.fmt.allocPrint(alloc, "{s}.primary", .{key}), extra_primary, now_ms, scope)) |ep| { var w = ep; w.label = try std.fmt.allocPrint(alloc, "{s} · Primary", .{name}); try limits.append(alloc, w); }
const extra_secondary = if (extra) |e| e.get("secondary_window") else null; if (try parseWindow(alloc, try std.fmt.allocPrint(alloc, "{s}.secondary", .{key}), extra_secondary, now_ms, scope)) |es| { var w = es; w.label = try std.fmt.allocPrint(alloc, "{s} · Secondary", .{name}); try limits.append(alloc, w); } } }
const plan_type = if (root) |r| usage.asString(r.get("plan_type")) else null; return ParseResult{ .planType = plan_type, .limits = limits.items, };}
fn jwtClaim(alloc: std.mem.Allocator, token: []const u8, namespace: []const u8, key: []const u8) ?[]const u8 { const payload = usage.jwtPayload(alloc, token) catch return null; const parsed = std.json.parseFromSlice(std.json.Value, alloc, payload, .{}) catch return null; const root = usage.asObject(parsed.value) orelse return null; const ns = usage.asObject(root.get(namespace)) orelse return null; return usage.asString(ns.get(key));}
fn jwtClaimEmail(alloc: std.mem.Allocator, token: []const u8) ?[]const u8 { const email_raw = jwtClaim(alloc, token, "https://api.openai.com/profile", "email") orelse return null; const trimmed = std.mem.trim(u8, email_raw, " \t\r\n"); if (trimmed.len == 0) return null; const lowered = alloc.alloc(u8, trimmed.len) catch return null; return std.ascii.lowerString(lowered, trimmed);}
pub fn advance( alloc: std.mem.Allocator, credential_json: []const u8, responses_json: []const u8, now_ms: i64,) !usage.Outcome { const cred_parsed = std.json.parseFromSlice(std.json.Value, alloc, credential_json, .{}) catch { return .{ .report = .{ .error_message = "invalid credential JSON" } }; }; const cred_root = usage.asObject(cred_parsed.value) orelse { return .{ .report = .{ .error_message = "invalid credential JSON" } }; }; const kind = usage.asString(cred_root.get("kind")) orelse ""; if (!std.mem.eql(u8, kind, "oauth")) { return .{ .report = .{ .error_message = "openai-codex needs an OAuth account" } }; }
const access = usage.asString(cred_root.get("access")) orelse ""; const cred_account_id = usage.asString(cred_root.get("accountId")); const account_id = if (cred_account_id != null and cred_account_id.?.len > 0) cred_account_id else jwtClaim(alloc, access, "https://api.openai.com/auth", "chatgpt_account_id");
const cred_email = usage.asString(cred_root.get("email")); const email = if (cred_email != null and cred_email.?.len > 0) cred_email else jwtClaimEmail(alloc, access);
const resp_parsed = std.json.parseFromSlice(std.json.Value, alloc, responses_json, .{}) catch { return .{ .report = .{ .error_message = "invalid responses JSON" } }; }; const resp_array = usage.asArray(resp_parsed.value) orelse { return .{ .report = .{ .error_message = "invalid responses JSON" } }; };
if (resp_array.items.len == 0) { var headers = std.ArrayList([2][]const u8).empty; const auth_hdr = try std.fmt.allocPrint(alloc, "Bearer {s}", .{access}); try headers.append(alloc, .{ "Authorization", auth_hdr }); try headers.append(alloc, .{ "User-Agent", "prime-agent" }); if (account_id) |aid| { if (aid.len > 0) { try headers.append(alloc, .{ "ChatGPT-Account-Id", aid }); } } const requests = try alloc.alloc(usage.UsageRequest, 1); requests[0] = .{ .method = "GET", .url = CODEX_USAGE_URL, .headers = headers.items, }; return .{ .requests = requests }; }
const first = switch (try usage.parseFirstResponse(alloc, resp_array, "usage body is not JSON")) { .ok => |f| f, .err => |msg| return .{ .report = .{ .error_message = msg } }, };
const parsed_usage = try parseCodexUsage(alloc, first.parsed.value, now_ms);
return .{ .report = .{ .planType = parsed_usage.planType, .limits = parsed_usage.limits, .accountId = account_id, .email = email, .error_message = null, }, };}