diff --git a/Sources/CodexBar/MenuCardView.swift b/Sources/CodexBar/MenuCardView.swift index 14703fae9..d211a9d96 100644 --- a/Sources/CodexBar/MenuCardView.swift +++ b/Sources/CodexBar/MenuCardView.swift @@ -943,6 +943,7 @@ extension UsageMenuCardView.Model { let zaiUsage = input.provider == .zai ? snapshot.zaiUsage : nil let zaiTokenDetail = Self.zaiLimitDetailText(limit: zaiUsage?.tokenLimit) let zaiTimeDetail = Self.zaiLimitDetailText(limit: zaiUsage?.timeLimit) + let zaiSessionDetail = Self.zaiLimitDetailText(limit: zaiUsage?.sessionTokenLimit) let openRouterQuotaDetail = Self.openRouterQuotaDetail(provider: input.provider, snapshot: snapshot) if input.provider == .codex, let codexProjection = input.codexProjection { metrics.append(contentsOf: Self.codexRateMetrics( @@ -984,6 +985,9 @@ extension UsageMenuCardView.Model { { tertiaryDetailText = detail } + if input.provider == .zai, let detail = zaiSessionDetail { + tertiaryDetailText = detail + } // Perplexity purchased credits don't reset; show balance without "Resets" prefix. let opusResetText: String? = input.provider == .perplexity ? opus.resetDescription?.trimmingCharacters(in: .whitespacesAndNewlines) diff --git a/Sources/CodexBarCore/Providers/Zai/ZaiProviderDescriptor.swift b/Sources/CodexBarCore/Providers/Zai/ZaiProviderDescriptor.swift index 430066a10..d644451c6 100644 --- a/Sources/CodexBarCore/Providers/Zai/ZaiProviderDescriptor.swift +++ b/Sources/CodexBarCore/Providers/Zai/ZaiProviderDescriptor.swift @@ -12,8 +12,8 @@ public enum ZaiProviderDescriptor { displayName: "z.ai", sessionLabel: "Tokens", weeklyLabel: "MCP", - opusLabel: nil, - supportsOpus: false, + opusLabel: "5-hour", + supportsOpus: true, supportsCredits: false, creditsHint: "", toggleTitle: "Show z.ai usage", diff --git a/Sources/CodexBarCore/Providers/Zai/ZaiUsageStats.swift b/Sources/CodexBarCore/Providers/Zai/ZaiUsageStats.swift index 1592a6181..1936c9f7e 100644 --- a/Sources/CodexBarCore/Providers/Zai/ZaiUsageStats.swift +++ b/Sources/CodexBarCore/Providers/Zai/ZaiUsageStats.swift @@ -15,6 +15,7 @@ public enum ZaiLimitUnit: Int, Sendable { case days = 1 case hours = 3 case minutes = 5 + case weeks = 6 } /// A single limit entry from the z.ai API @@ -69,6 +70,8 @@ extension ZaiLimitEntry { return self.number * 60 case .days: return self.number * 24 * 60 + case .weeks: + return self.number * 7 * 24 * 60 case .unknown: return nil } @@ -80,6 +83,7 @@ extension ZaiLimitEntry { case .minutes: "minute" case .hours: "hour" case .days: "day" + case .weeks: "week" case .unknown: nil } guard let unitLabel else { return nil } @@ -129,12 +133,21 @@ public struct ZaiUsageDetail: Sendable, Codable { /// Complete z.ai usage response public struct ZaiUsageSnapshot: Sendable { public let tokenLimit: ZaiLimitEntry? + /// Shorter-window TOKENS_LIMIT (e.g. 5-hour), present only when the API returns two TOKENS_LIMIT entries. + public let sessionTokenLimit: ZaiLimitEntry? public let timeLimit: ZaiLimitEntry? public let planName: String? public let updatedAt: Date - public init(tokenLimit: ZaiLimitEntry?, timeLimit: ZaiLimitEntry?, planName: String?, updatedAt: Date) { + public init( + tokenLimit: ZaiLimitEntry?, + sessionTokenLimit: ZaiLimitEntry? = nil, + timeLimit: ZaiLimitEntry?, + planName: String?, + updatedAt: Date) + { self.tokenLimit = tokenLimit + self.sessionTokenLimit = sessionTokenLimit self.timeLimit = timeLimit self.planName = planName self.updatedAt = updatedAt @@ -150,13 +163,13 @@ extension ZaiUsageSnapshot { public func toUsageSnapshot() -> UsageSnapshot { let primaryLimit = self.tokenLimit ?? self.timeLimit let secondaryLimit = (self.tokenLimit != nil && self.timeLimit != nil) ? self.timeLimit : nil - let primary = primaryLimit.map { Self.rateWindow(for: $0) } ?? RateWindow( usedPercent: 0, windowMinutes: nil, resetsAt: nil, resetDescription: nil) let secondary = secondaryLimit.map { Self.rateWindow(for: $0) } + let tertiary = self.sessionTokenLimit.map { Self.rateWindow(for: $0) } let planName = self.planName?.trimmingCharacters(in: .whitespacesAndNewlines) let loginMethod = (planName?.isEmpty ?? true) ? nil : planName @@ -168,7 +181,7 @@ extension ZaiUsageSnapshot { return UsageSnapshot( primary: primary, secondary: secondary, - tertiary: nil, + tertiary: tertiary, providerCost: nil, zaiUsage: self, updatedAt: self.updatedAt, @@ -364,22 +377,38 @@ public struct ZaiUsageFetcher: Sendable { throw ZaiUsageError.parseFailed("Missing data") } - var tokenLimit: ZaiLimitEntry? + var tokenLimits: [ZaiLimitEntry] = [] var timeLimit: ZaiLimitEntry? for limit in responseData.limits { if let entry = limit.toLimitEntry() { switch entry.type { case .tokensLimit: - tokenLimit = entry + tokenLimits.append(entry) case .timeLimit: timeLimit = entry } } } + // Multiple TOKENS_LIMIT entries: shortest window → sessionTokenLimit (tertiary), + // longest → tokenLimit (primary). + let tokenLimit: ZaiLimitEntry? + let sessionTokenLimit: ZaiLimitEntry? + if tokenLimits.count >= 2 { + let sorted = tokenLimits.sorted { + ($0.windowMinutes ?? Int.max) < ($1.windowMinutes ?? Int.max) + } + sessionTokenLimit = sorted.first + tokenLimit = sorted.last + } else { + tokenLimit = tokenLimits.first + sessionTokenLimit = nil + } + return ZaiUsageSnapshot( tokenLimit: tokenLimit, + sessionTokenLimit: sessionTokenLimit, timeLimit: timeLimit, planName: responseData.planName, updatedAt: Date()) diff --git a/Tests/CodexBarTests/CLISnapshotTests.swift b/Tests/CodexBarTests/CLISnapshotTests.swift index d5218d1bc..b3ce7ae52 100644 --- a/Tests/CodexBarTests/CLISnapshotTests.swift +++ b/Tests/CodexBarTests/CLISnapshotTests.swift @@ -425,4 +425,27 @@ struct CLISnapshotTests { #expect(!output.contains("\u{001B}[")) #expect(output.contains("Status: Operational – Operational")) } + + @Test + func `renders 5-hour tertiary row for zai`() { + let snap = UsageSnapshot( + primary: .init(usedPercent: 9, windowMinutes: 10080, resetsAt: nil, resetDescription: nil), + secondary: .init(usedPercent: 50, windowMinutes: nil, resetsAt: nil, resetDescription: nil), + tertiary: .init(usedPercent: 25, windowMinutes: 300, resetsAt: nil, resetDescription: nil), + updatedAt: Date(timeIntervalSince1970: 0)) + + let output = CLIRenderer.renderText( + provider: .zai, + snapshot: snap, + credits: nil, + context: RenderContext( + header: "z.ai 0.0.0 (zai)", + status: nil, + useColor: false, + resetStyle: .absolute)) + + #expect(output.contains("5-hour:")) + #expect(output.contains("Tokens:")) + #expect(output.contains("MCP:")) + } } diff --git a/Tests/CodexBarTests/CodexPresentationCharacterizationTests.swift b/Tests/CodexBarTests/CodexPresentationCharacterizationTests.swift index 81fd77603..047a15246 100644 --- a/Tests/CodexBarTests/CodexPresentationCharacterizationTests.swift +++ b/Tests/CodexBarTests/CodexPresentationCharacterizationTests.swift @@ -346,6 +346,56 @@ struct CodexPresentationCharacterizationTests { #expect(store.codexCookieCacheScopeForOpenAIWeb() == nil) } + @Test + func `zai menu descriptor includes Tokens MCP and 5-hour rows`() { + let settings = self.makeSettingsStore(suite: "CodexPresentationCharacterizationTests-zai-three-quota") + settings.statusChecksEnabled = false + + let fetcher = UsageFetcher(environment: [:]) + let store = UsageStore( + fetcher: fetcher, + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing) + store._setSnapshotForTesting( + UsageSnapshot( + primary: RateWindow( + usedPercent: 9, + windowMinutes: 10080, + resetsAt: nil, + resetDescription: nil), + secondary: RateWindow( + usedPercent: 50, + windowMinutes: nil, + resetsAt: nil, + resetDescription: nil), + tertiary: RateWindow( + usedPercent: 25, + windowMinutes: 300, + resetsAt: nil, + resetDescription: nil), + updatedAt: Date(), + identity: ProviderIdentitySnapshot( + providerID: .zai, + accountEmail: nil, + accountOrganization: nil, + loginMethod: "pro")), + provider: .zai) + + let descriptor = MenuDescriptor.build( + provider: .zai, + store: store, + settings: settings, + account: fetcher.loadAccountInfo(), + updateReady: false, + includeContextualActions: false) + + let lines = self.textLines(from: descriptor) + #expect(lines.contains(where: { $0.hasPrefix("Tokens:") })) + #expect(lines.contains(where: { $0.hasPrefix("MCP:") })) + #expect(lines.contains(where: { $0.hasPrefix("5-hour:") })) + } + private func makeSettingsStore(suite: String) -> SettingsStore { let defaults = UserDefaults(suiteName: suite)! defaults.removePersistentDomain(forName: suite) diff --git a/Tests/CodexBarTests/ZaiMenuCardTests.swift b/Tests/CodexBarTests/ZaiMenuCardTests.swift new file mode 100644 index 000000000..4433fdf9f --- /dev/null +++ b/Tests/CodexBarTests/ZaiMenuCardTests.swift @@ -0,0 +1,70 @@ +import CodexBarCore +import Foundation +import Testing +@testable import CodexBar + +struct ZaiMenuCardTests { + @Test + func `zai metrics titles are Tokens MCP and 5-hour when session token limit present`() throws { + let now = Date() + let zai = ZaiUsageSnapshot( + tokenLimit: ZaiLimitEntry( + type: .tokensLimit, + unit: .weeks, + number: 1, + usage: nil, + currentValue: nil, + remaining: nil, + percentage: 9, + usageDetails: [], + nextResetTime: nil), + sessionTokenLimit: ZaiLimitEntry( + type: .tokensLimit, + unit: .hours, + number: 5, + usage: 1000, + currentValue: 750, + remaining: 250, + percentage: 25, + usageDetails: [], + nextResetTime: nil), + timeLimit: ZaiLimitEntry( + type: .timeLimit, + unit: .minutes, + number: 1, + usage: 100, + currentValue: 50, + remaining: 50, + percentage: 50, + usageDetails: [], + nextResetTime: nil), + planName: "pro", + updatedAt: now) + let snapshot = zai.toUsageSnapshot() + let metadata = try #require(ProviderDefaults.metadata[.zai]) + + let model = UsageMenuCardView.Model.make(.init( + provider: .zai, + metadata: metadata, + snapshot: snapshot, + credits: nil, + creditsError: nil, + dashboard: nil, + dashboardError: nil, + tokenSnapshot: nil, + tokenError: nil, + account: AccountInfo(email: nil, plan: nil), + isRefreshing: false, + lastError: nil, + usageBarsShowUsed: false, + resetTimeDisplayStyle: .countdown, + tokenCostUsageEnabled: false, + showOptionalCreditsAndExtraUsage: true, + hidePersonalInfo: false, + now: now)) + + #expect(model.metrics.map(\.title) == ["Tokens", "MCP", "5-hour"]) + let tertiary = try #require(model.metrics.first(where: { $0.title == "5-hour" })) + #expect(tertiary.detailText == "750 / 1K (250 remaining)") + } +} diff --git a/Tests/CodexBarTests/ZaiProviderTests.swift b/Tests/CodexBarTests/ZaiProviderTests.swift index a112e281e..ce23d2351 100644 --- a/Tests/CodexBarTests/ZaiProviderTests.swift +++ b/Tests/CodexBarTests/ZaiProviderTests.swift @@ -67,7 +67,9 @@ struct ZaiUsageSnapshotTests { #expect(usage.primary?.resetDescription == "5 hours window") #expect(usage.secondary?.usedPercent == 20) #expect(usage.secondary?.resetDescription == "30 days window") + #expect(usage.tertiary == nil) #expect(usage.zaiUsage?.tokenLimit?.usage == 100) + #expect(usage.zaiUsage?.sessionTokenLimit == nil) } @Test @@ -320,6 +322,138 @@ struct ZaiUsageParsingTests { } } +struct ZaiThreeLimitTests { + @Test + func `parses three limit entries into session weekly and mcp slots`() throws { + let json = """ + { + "code": 200, + "msg": "操作成功", + "data": { + "limits": [ + { + "type": "TOKENS_LIMIT", + "unit": 3, + "number": 5, + "percentage": 25, + "nextResetTime": 1775020168897 + }, + { + "type": "TOKENS_LIMIT", + "unit": 6, + "number": 1, + "percentage": 9, + "nextResetTime": 1775588029998 + }, + { + "type": "TIME_LIMIT", + "unit": 5, + "number": 1, + "usage": 1000, + "currentValue": 224, + "remaining": 776, + "percentage": 22, + "nextResetTime": 1777575229998, + "usageDetails": [ + { "modelCode": "search-prime", "usage": 210 }, + { "modelCode": "web-reader", "usage": 14 } + ] + } + ], + "level": "pro" + }, + "success": true + } + """ + + let snapshot = try ZaiUsageFetcher.parseUsageSnapshot(from: Data(json.utf8)) + + // Weekly token limit (unit:6=weeks, longer window) → tokenLimit (primary) + #expect(snapshot.tokenLimit?.unit == .weeks) + #expect(snapshot.tokenLimit?.number == 1) + #expect(snapshot.tokenLimit?.percentage == 9.0) + #expect(snapshot.tokenLimit?.windowMinutes == 10080) + + // 5-hour token limit (unit:3=hours, number:5 → 300 min) → sessionTokenLimit (tertiary) + #expect(snapshot.sessionTokenLimit?.unit == .hours) + #expect(snapshot.sessionTokenLimit?.number == 5) + #expect(snapshot.sessionTokenLimit?.percentage == 25.0) + #expect(snapshot.sessionTokenLimit?.windowMinutes == 300) + + // MCP time limit → timeLimit (secondary) + #expect(snapshot.timeLimit?.usage == 1000) + #expect(snapshot.timeLimit?.usageDetails.first?.modelCode == "search-prime") + + // UsageSnapshot slot mapping + let usage = snapshot.toUsageSnapshot() + #expect(usage.primary?.usedPercent == 9.0) + #expect(usage.primary?.windowMinutes == 10080) + #expect(usage.secondary != nil) // MCP + #expect(usage.tertiary?.usedPercent == 25.0) + #expect(usage.tertiary?.windowMinutes == 300) + } + + @Test + func `unit 6 maps to weeks with correct window minutes`() { + let entry = ZaiLimitEntry( + type: .tokensLimit, + unit: .weeks, + number: 1, + usage: nil, + currentValue: nil, + remaining: nil, + percentage: 9, + usageDetails: [], + nextResetTime: nil) + #expect(entry.windowMinutes == 10080) + #expect(entry.windowDescription == "1 week") + #expect(entry.windowLabel == "1 week window") + } + + @Test + func `two limit entries remain backward compatible`() throws { + let json = """ + { + "code": 200, + "msg": "Operation successful", + "data": { + "limits": [ + { + "type": "TIME_LIMIT", + "unit": 5, + "number": 1, + "usage": 100, + "currentValue": 50, + "remaining": 50, + "percentage": 50, + "usageDetails": [] + }, + { + "type": "TOKENS_LIMIT", + "unit": 3, + "number": 5, + "percentage": 34, + "nextResetTime": 1768507567547 + } + ] + }, + "success": true + } + """ + + let snapshot = try ZaiUsageFetcher.parseUsageSnapshot(from: Data(json.utf8)) + + #expect(snapshot.tokenLimit != nil) + #expect(snapshot.sessionTokenLimit == nil) + #expect(snapshot.timeLimit != nil) + + let usage = snapshot.toUsageSnapshot() + #expect(usage.primary != nil) + #expect(usage.secondary != nil) + #expect(usage.tertiary == nil) + } +} + struct ZaiAPIRegionTests { @Test func `defaults to global endpoint`() {