@@ -0,0 +1,388 @@
1+import Foundation
2+
3+// MARK: - Usage
4+
5+/// Usage data from z.ai's coding plan API — two endpoints that together give the
6+/// plan level, the 5-hour token quota, and the weekly token quota.
7+///
8+/// Every field in the report is optional. The quota endpoint may return fewer limits
9+/// than expected, and the subscription endpoint may return an empty list, so an
10+/// unrecognized shape degrades to "nothing to show" rather than failing the decode.
11+struct ZaiUsageReport: Equatable {
12+ /// Percent of the 5-hour token window consumed, 0–100.
13+ let fiveHourPercent: Int?
14+ /// Percent of the weekly token window consumed, 0–100.
15+ let weeklyPercent: Int?
16+ /// When the 5-hour window resets, or `nil` when not provided.
17+ let fiveHourResetsAt: Date?
18+ /// When the weekly window resets, or `nil` when not provided.
19+ let weeklyResetsAt: Date?
20+ /// The plan tier from the quota response: `"lite"`, `"pro"`, `"max"`.
21+ let planLevel: String?
22+
23+ /// Whether there is anything at all to render. A report with all-nil fields reads
24+ /// as a section header over nothing, which looks like a bug.
25+ var isEmpty: Bool {
26+ fiveHourPercent == nil && weeklyPercent == nil && planLevel == nil
27+ }
28+}
29+
30+// MARK: - API Response Models
31+
32+/// Response from `GET https://api.z.ai/api/monitor/usage/quota/limit`.
33+struct ZaiQuotaResponse: Decodable, Equatable {
34+ let code: Int
35+ let msg: String?
36+ let data: ZaiQuotaData?
37+ let success: Bool
38+
39+ private enum CodingKeys: String, CodingKey {
40+ case code, msg, data, success
41+ }
42+
43+ init(from decoder: Decoder) throws {
44+ let container = try decoder.container(keyedBy: CodingKeys.self)
45+ code = try container.decodeIfPresent(Int.self, forKey: .code) ?? 0
46+ msg = try container.decodeIfPresent(String.self, forKey: .msg)
47+ data = try container.decodeIfPresent(ZaiQuotaData.self, forKey: .data)
48+ success = try container.decodeIfPresent(Bool.self, forKey: .success) ?? false
49+ }
50+}
51+
52+struct ZaiQuotaData: Decodable, Equatable {
53+ let limits: [ZaiLimit]?
54+ let level: String?
55+
56+ private enum CodingKeys: String, CodingKey {
57+ case limits, level
58+ }
59+
60+ init(from decoder: Decoder) throws {
61+ let container = try decoder.container(keyedBy: CodingKeys.self)
62+ limits = try container.decodeIfPresent([ZaiLimit].self, forKey: .limits)
63+ level = try container.decodeIfPresent(String.self, forKey: .level)
64+ }
65+}
66+
67+/// One limit entry from the quota response.
68+///
69+/// `unit` identifies the window: 3 = 5-hour, 6 = weekly, 5 = monthly tools.
70+/// `type` is `"TOKENS_LIMIT"` for token quotas or `"TIME_LIMIT"` for tool-call quotas.
71+struct ZaiLimit: Decodable, Equatable {
72+ let type: String
73+ let unit: Int
74+ let percentage: Int
75+ /// Epoch milliseconds, or `nil` when nothing has been used.
76+ let nextResetTime: Int64?
77+
78+ private enum CodingKeys: String, CodingKey {
79+ case type, unit, percentage, nextResetTime
80+ }
81+
82+ init(from decoder: Decoder) throws {
83+ let container = try decoder.container(keyedBy: CodingKeys.self)
84+ type = try container.decodeIfPresent(String.self, forKey: .type) ?? ""
85+ unit = try container.decodeIfPresent(Int.self, forKey: .unit) ?? 0
86+ // Decoded as Double and rounded: the field is documented as an integer, but
87+ // a fractional value would throw and take the whole payload down with it.
88+ let rawPercent = try container.decodeIfPresent(Double.self, forKey: .percentage) ?? 0
89+ percentage = Int(rawPercent.rounded())
90+ nextResetTime = try container.decodeIfPresent(Int64.self, forKey: .nextResetTime)
91+ }
92+
93+ init(type: String, unit: Int, percentage: Int, nextResetTime: Int64? = nil) {
94+ self.type = type
95+ self.unit = unit
96+ self.percentage = percentage
97+ self.nextResetTime = nextResetTime
98+ }
99+}
100+
101+/// Response from `GET https://api.z.ai/api/biz/subscription/list`.
102+struct ZaiSubscriptionResponse: Decodable, Equatable {
103+ let code: Int
104+ let data: [ZaiSubscription]?
105+ let success: Bool
106+
107+ private enum CodingKeys: String, CodingKey {
108+ case code, data, success
109+ }
110+
111+ init(from decoder: Decoder) throws {
112+ let container = try decoder.container(keyedBy: CodingKeys.self)
113+ code = try container.decodeIfPresent(Int.self, forKey: .code) ?? 0
114+ data = try container.decodeIfPresent([ZaiSubscription].self, forKey: .data)
115+ success = try container.decodeIfPresent(Bool.self, forKey: .success) ?? false
116+ }
117+}
118+
119+struct ZaiSubscription: Decodable, Equatable {
120+ let productName: String?
121+ let status: String?
122+
123+ private enum CodingKeys: String, CodingKey {
124+ case productName, status
125+ }
126+
127+ init(from decoder: Decoder) throws {
128+ let container = try decoder.container(keyedBy: CodingKeys.self)
129+ productName = try container.decodeIfPresent(String.self, forKey: .productName)
130+ status = try container.decodeIfPresent(String.self, forKey: .status)
131+ }
132+
133+ init(productName: String? = nil, status: String? = nil) {
134+ self.productName = productName
135+ self.status = status
136+ }
137+}
138+
139+// MARK: - Errors
140+
141+/// Errors that can occur while fetching z.ai usage.
142+enum ZaiUsageError: LocalizedError, Equatable {
143+ /// No API key has been entered. Not an error the user needs to see — the popover
144+ /// just leaves the section out, same as missing Claude credentials.
145+ case noAPIKey
146+
147+ /// The API key was rejected by z.ai.
148+ case invalidAPIKey
149+ case requestFailed(String)
150+ case decodingFailed(String)
151+
152+ var errorDescription: String? {
153+ switch self {
154+ case .noAPIKey:
155+ return nil
156+ case .invalidAPIKey:
157+ return "z.ai API key is invalid. Check your key at z.ai/manage-apikey/coding-plan/personal/my-plan"
158+ case .requestFailed(let message):
159+ return "z.ai request failed: \(message)"
160+ case .decodingFailed(let message):
161+ return "Failed to decode z.ai usage: \(message)"
162+ }
163+ }
164+}
165+
166+// MARK: - Protocol
167+
168+/// Protocol abstraction for fetching z.ai usage, enabling test injection.
169+protocol ZaiUsageChecking {
170+ func fetchUsage(apiKey: String) async throws -> ZaiUsageReport
171+}
172+
173+// MARK: - Client
174+
175+/// Client for z.ai's coding plan usage API, authenticating with an API key the user
176+/// enters in settings (stored in Keychain, separate from the Hyper key).
177+///
178+/// Fetches the quota/limit and subscription endpoints concurrently to get both usage
179+/// percentages and the plan tier in one round-trip. Retries transient failures with
180+/// exponential backoff, same pattern as `CreditsChecker` and `ClaudeUsageClient`.
181+actor ZaiUsageClient: ZaiUsageChecking {
182+ static let baseURL = URL(string: "https://api.z.ai/api")!
183+ static let quotaEndpoint = URL(string: "https://api.z.ai/api/monitor/usage/quota/limit")!
184+ static let subscriptionEndpoint = URL(string: "https://api.z.ai/api/biz/subscription/list")!
185+
186+ /// Retries *after* the initial attempt, so at most 3 requests per call.
187+ static let maxRetries = 2
188+
189+ /// URL error codes worth retrying: the request never reached a server that had an
190+ /// opinion about it, so the same request may well succeed a moment later.
191+ private static let retryableURLErrorCodes: Set<URLError.Code> = [
192+ .timedOut,
193+ .cannotConnectToHost,
194+ .cannotFindHost,
195+ .dnsLookupFailed,
196+ .networkConnectionLost,
197+ .notConnectedToInternet,
198+ .resourceUnavailable,
199+ .badServerResponse,
200+ ]
201+
202+ /// Client-error status codes that are nonetheless worth retrying: both are the
203+ /// server telling us to come back later, not that the request itself is wrong.
204+ private static let retryableStatusCodes: Set<Int> = [
205+ 408, // Request Timeout
206+ 429, // Too Many Requests
207+ ]
208+
209+ private let session: URLSession
210+ private let retryBaseDelay: Duration
211+
212+ /// - Parameters:
213+ /// - retryBaseDelay: Delay before the first retry; doubles on each subsequent
214+ /// one (2s, then 4s). Injectable so tests need not wait in real time.
215+ init(
216+ session: URLSession? = nil,
217+ retryBaseDelay: Duration = .seconds(2)
218+ ) {
219+ self.retryBaseDelay = retryBaseDelay
220+ if let session {
221+ self.session = session
222+ } else {
223+ let config = URLSessionConfiguration.default
224+ config.timeoutIntervalForRequest = 10
225+ config.timeoutIntervalForResource = 15
226+ self.session = URLSession(configuration: config)
227+ }
228+ }
229+
230+ /// Fetches current usage with the given API key.
231+ ///
232+ /// The quota and subscription endpoints are fetched concurrently; neither waits on
233+ /// the other. The subscription endpoint is best-effort — if it fails, the plan
234+/// level from the quota response (`data.level`) is used instead.
235+ ///
236+ /// - Throws: `ZaiUsageError`. `.noAPIKey` means no key was provided, which callers
237+ /// should treat as "nothing to show" rather than as a failure.
238+ func fetchUsage(apiKey: String) async throws -> ZaiUsageReport {
239+ guard !apiKey.isEmpty else { throw ZaiUsageError.noAPIKey }
240+
241+ async let quotaOutcome = fetchQuota(apiKey: apiKey)
242+ async let subscriptionOutcome = fetchSubscription(apiKey: apiKey)
243+
244+ let quota = try await quotaOutcome
245+ let subscription = try? await subscriptionOutcome
246+
247+ // Extract percentages from the limits array.
248+ var fiveHourPercent: Int?
249+ var weeklyPercent: Int?
250+ var fiveHourResetsAt: Date?
251+ var weeklyResetsAt: Date?
252+
253+ for limit in quota?.limits ?? [] {
254+ switch limit.unit {
255+ case 3: // 5-hour window
256+ fiveHourPercent = limit.percentage
257+ fiveHourResetsAt = limit.nextResetTime.map { Date(timeIntervalSince1970: TimeInterval($0) / 1000) }
258+ case 6: // weekly window
259+ weeklyPercent = limit.percentage
260+ weeklyResetsAt = limit.nextResetTime.map { Date(timeIntervalSince1970: TimeInterval($0) / 1000) }
261+ default:
262+ break
263+ }
264+ }
265+
266+ // Plan level: prefer the product name from the subscription, fall back to the
267+ // quota response's `level` field.
268+ var planLevel: String? = quota?.level
269+ if let sub = subscription?.data?.first(where: { $0.status == "VALID" }),
270+ let name = sub.productName {
271+ planLevel = name
272+ }
273+
274+ return ZaiUsageReport(
275+ fiveHourPercent: fiveHourPercent,
276+ weeklyPercent: weeklyPercent,
277+ fiveHourResetsAt: fiveHourResetsAt,
278+ weeklyResetsAt: weeklyResetsAt,
279+ planLevel: planLevel
280+ )
281+ }
282+
283+ // MARK: - Quota Request
284+
285+ private func fetchQuota(apiKey: String) async throws -> ZaiQuotaData? {
286+ var request = URLRequest(url: Self.quotaEndpoint, cachePolicy: .reloadIgnoringLocalCacheData)
287+ request.httpMethod = "GET"
288+ request.setValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization")
289+ request.setValue("application/json", forHTTPHeaderField: "Accept")
290+
291+ let data = try await send(request)
292+
293+ do {
294+ let response = try JSONDecoder().decode(ZaiQuotaResponse.self, from: data)
295+ guard response.success else {
296+ throw ZaiUsageError.requestFailed(response.msg ?? "Unknown error")
297+ }
298+ return response.data
299+ } catch let error as ZaiUsageError {
300+ throw error
301+ } catch {
302+ throw ZaiUsageError.decodingFailed(error.localizedDescription)
303+ }
304+ }
305+
306+ // MARK: - Subscription Request
307+
308+ private func fetchSubscription(apiKey: String) async throws -> ZaiSubscriptionResponse {
309+ var request = URLRequest(url: Self.subscriptionEndpoint, cachePolicy: .reloadIgnoringLocalCacheData)
310+ request.httpMethod = "GET"
311+ request.setValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization")
312+ request.setValue("application/json", forHTTPHeaderField: "Accept")
313+
314+ let data = try await send(request)
315+
316+ do {
317+ return try JSONDecoder().decode(ZaiSubscriptionResponse.self, from: data)
318+ } catch {
319+ throw ZaiUsageError.decodingFailed(error.localizedDescription)
320+ }
321+ }
322+
323+ // MARK: - Transport
324+
325+ /// Sends a request, retrying transient failures with exponential backoff.
326+ private func send(_ request: URLRequest) async throws -> Data {
327+ var attempt = 0
328+
329+ while true {
330+ do {
331+ return try await performAttempt(request)
332+ } catch let failure as AttemptFailure {
333+ guard failure.isRetryable, attempt < Self.maxRetries else {
334+ throw failure.underlying
335+ }
336+ try await Task.sleep(for: retryBaseDelay * (1 << attempt))
337+ attempt += 1
338+ }
339+ }
340+ }
341+
342+ /// Performs a single request, tagging the failure with whether a retry could
343+ /// plausibly help.
344+ private func performAttempt(_ request: URLRequest) async throws -> Data {
345+ let data: Data
346+ let response: URLResponse
347+ do {
348+ (data, response) = try await session.data(for: request)
349+ } catch let urlError as URLError {
350+ throw AttemptFailure(
351+ underlying: .requestFailed(urlError.localizedDescription),
352+ isRetryable: Self.retryableURLErrorCodes.contains(urlError.code)
353+ )
354+ } catch {
355+ throw AttemptFailure(
356+ underlying: .requestFailed(error.localizedDescription),
357+ isRetryable: false
358+ )
359+ }
360+
361+ guard let httpResponse = response as? HTTPURLResponse else {
362+ throw AttemptFailure(underlying: .requestFailed("Invalid response type"), isRetryable: false)
363+ }
364+
365+ switch httpResponse.statusCode {
366+ case 200...299:
367+ return data
368+ case 401:
369+ throw AttemptFailure(underlying: .invalidAPIKey, isRetryable: false)
370+ case 500...599:
371+ throw AttemptFailure(
372+ underlying: .requestFailed("HTTP \(httpResponse.statusCode)"),
373+ isRetryable: true
374+ )
375+ default:
376+ throw AttemptFailure(
377+ underlying: .requestFailed("HTTP \(httpResponse.statusCode)"),
378+ isRetryable: Self.retryableStatusCodes.contains(httpResponse.statusCode)
379+ )
380+ }
381+ }
382+
383+ /// A failed attempt, carrying the error to surface plus whether retrying is worthwhile.
384+ private struct AttemptFailure: Error {
385+ let underlying: ZaiUsageError
386+ let isRetryable: Bool
387+ }
388+}