//! Z.ai usage feed implementation.
//!
//! Mirrors prime-agent:packages/ai/src/utils/usage/providers/zai.ts field for field.
const std = @import("std");
const usage = @import("../usage.zig");
pub const ZAI_USAGE_URL = "https://api.z.ai/api/monitor/usage/quota/limit";
pub const ParseZaiResult = struct {
planType: ?[]const u8 = null,
limits: []const usage.UsageLimit = &[_]usage.UsageLimit{},
};
pub fn parseZaiUsage(alloc: std.mem.Allocator, payload: ?std.json.Value) !ParseZaiResult {
const root = usage.asObject(payload);
const data = if (root) |r| usage.asObject(r.get("data")) else null;
if (data == null) return ParseZaiResult{};
var limits = std.ArrayList(usage.UsageLimit).empty;
const raw_limits = usage.asArray(data.?.get("limits"));
if (raw_limits) |arr| {
for (arr.items) |entry| {
const item = usage.asObject(entry) orelse continue;
const type_str = usage.asString(item.get("type")) orelse "";
const unit = usage.asNumber(item.get("unit"));
const raw_pct = usage.asNumber(item.get("percentage"));
const cur = usage.asNumber(item.get("currentValue"));
const total = usage.asNumber(item.get("usage")) orelse usage.asNumber(item.get("number"));
var used_percent = raw_pct;
if (used_percent == null and cur != null and total != null and total.? > 0.0) {
used_percent = (cur.? / total.?) * 100.0;
}
const raw_reset = usage.asNumber(item.get("nextResetTime"));
const resets_at: ?i64 = if (raw_reset) |rr|
(if (rr > 100_000_000_000.0) usage.toInt(i64, rr) else usage.toInt(i64, rr * 1000.0))
else
null;
var label: []const u8 = undefined;
var window_label: ?[]const u8 = null;
var detail: ?[]const u8 = null;
// One type can repeat with a different period, so the period is part of the key.
const number = usage.asNumber(item.get("number"));
const unit_int = usage.toInt(i64, unit);
const number_int = usage.toInt(i64, number);
const base_id = if (type_str.len > 0) try std.ascii.allocLowerString(alloc, type_str) else "limit";
const id = if (unit_int != null and number_int != null)
try std.fmt.allocPrint(alloc, "{s}:{d}:{d}", .{ base_id, unit_int.?, number_int.? })
else
base_id;
// Unit 3 counts hours; the other units are not pinned down, so they carry no length.
const window_seconds: ?i64 = if (unit_int == 3 and number_int != null) number_int.? * 3600 else null;
if (std.mem.eql(u8, type_str, "TIME_LIMIT") or (unit != null and unit.? == 5.0)) {
label = "Requests";
if (cur != null and total != null) {
if (usage.toInt(i64, cur.?) != null and usage.toInt(i64, total.?) != null) {
detail = try std.fmt.allocPrint(alloc, "{d} / {d} reqs", .{
usage.toInt(i64, cur.?).?,
usage.toInt(i64, total.?).?,
});
}
}
} else if (std.mem.eql(u8, type_str, "TOKENS_LIMIT") or (unit != null and (unit.? == 3.0 or unit.? == 6.0))) {
label = "Tokens";
window_label = "5h window";
} else {
label = if (type_str.len > 0) try usage.titleCase(alloc, type_str) else "Limit";
}
const status: []const u8 = if (used_percent != null and used_percent.? >= 100.0)
"exhausted"
else
usage.statusFor(used_percent);
try limits.append(alloc, .{
.id = id,
.label = label,
.usedPercent = used_percent,
.windowLabel = window_label,
.windowSeconds = window_seconds,
.resetsAt = resets_at,
.detail = detail,
.status = status,
});
}
}
const level = usage.asString(data.?.get("level"));
const plan_type: ?[]const u8 = if (level) |l| blk: {
const trimmed = std.mem.trim(u8, l, " \t\r\n");
if (trimmed.len == 0) break :blk null;
break :blk try usage.titleCase(alloc, trimmed);
} else null;
return ParseZaiResult{
.planType = plan_type,
.limits = limits.items,
};
}
pub fn advance(
alloc: std.mem.Allocator,
credential_json: []const u8,
responses_json: []const u8,
_now_ms: i64,
) !usage.Outcome {
_ = _now_ms;
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, "apiKey")) {
return .{ .report = .{ .error_message = "zai needs an apiKey credential" } };
}
const api_key = usage.asString(cred_root.get("apiKey")) orelse "";
const cred_account_id = usage.asString(cred_root.get("accountId"));
const cred_email = usage.asString(cred_root.get("email"));
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;
// zai requires authorization without Bearer prefix
try headers.append(alloc, .{ "Authorization", api_key });
try headers.append(alloc, .{ "Accept", "application/json" });
try headers.append(alloc, .{ "User-Agent", "OpenCode-Status-Plugin/1.0" });
const requests = try alloc.alloc(usage.UsageRequest, 1);
requests[0] = .{
.method = "GET",
.url = ZAI_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 body_parsed = first.parsed;
const root = usage.asObject(body_parsed.value);
const success_val = if (root) |r| r.get("success") else null;
const is_success = if (success_val) |sv| switch (sv) {
.bool => |b| b,
else => true,
} else true;
const code_num = if (root) |r| usage.asNumber(r.get("code")) else null;
if (!is_success or (code_num != null and code_num.? != 200.0)) {
const msg = if (root) |r| usage.asString(r.get("msg")) else null;
if (msg) |m| {
return .{ .report = .{ .error_message = m } };
}
const err_msg = try std.fmt.allocPrint(alloc, "Z.ai returned code {d}", .{usage.toInt(i64, code_num) orelse 0});
return .{ .report = .{ .error_message = err_msg } };
}
const parsed = try parseZaiUsage(alloc, body_parsed.value);
return .{
.report = .{
.planType = parsed.planType,
.accountId = cred_account_id,
.email = cred_email,
.limits = parsed.limits,
.error_message = null,
},
};
}