|
| 1 | +import Foundation |
| 2 | +import Domain |
| 3 | + |
| 4 | +/// Infrastructure adapter that probes the Kimi CLI to fetch usage quotas. |
| 5 | +/// Starts the interactive `kimi` CLI, sends `/usage`, and parses the output. |
| 6 | +/// |
| 7 | +/// Sample CLI output: |
| 8 | +/// ``` |
| 9 | +/// ╭─────────────────────────────── API Usage ───────────────────────────────╮ |
| 10 | +/// │ Weekly limit ━━━━━━━━━━━━━━━━━━━━ 100% left (resets in 6d 23h 22m) │ |
| 11 | +/// │ 5h limit ━━━━━━━━━━━━━━━━━━━━ 100% left (resets in 4h 22m) │ |
| 12 | +/// ╰─────────────────────────────────────────────────────────────────────────╯ |
| 13 | +/// ``` |
| 14 | +public struct KimiCLIUsageProbe: UsageProbe { |
| 15 | + private let kimiBinary: String |
| 16 | + private let timeout: TimeInterval |
| 17 | + private let cliExecutor: CLIExecutor |
| 18 | + |
| 19 | + public init( |
| 20 | + kimiBinary: String = "kimi", |
| 21 | + timeout: TimeInterval = 15.0, |
| 22 | + cliExecutor: CLIExecutor? = nil |
| 23 | + ) { |
| 24 | + self.kimiBinary = kimiBinary |
| 25 | + self.timeout = timeout |
| 26 | + self.cliExecutor = cliExecutor ?? DefaultCLIExecutor() |
| 27 | + } |
| 28 | + |
| 29 | + public func isAvailable() async -> Bool { |
| 30 | + if cliExecutor.locate(kimiBinary) != nil { |
| 31 | + return true |
| 32 | + } |
| 33 | + AppLog.probes.error("Kimi binary '\(kimiBinary)' not found in PATH") |
| 34 | + return false |
| 35 | + } |
| 36 | + |
| 37 | + public func probe() async throws -> UsageSnapshot { |
| 38 | + guard cliExecutor.locate(kimiBinary) != nil else { |
| 39 | + throw ProbeError.cliNotFound(kimiBinary) |
| 40 | + } |
| 41 | + |
| 42 | + AppLog.probes.info("Starting Kimi CLI probe with /usage command...") |
| 43 | + |
| 44 | + let result: CLIResult |
| 45 | + do { |
| 46 | + result = try cliExecutor.execute( |
| 47 | + binary: kimiBinary, |
| 48 | + args: [], |
| 49 | + input: nil, |
| 50 | + timeout: timeout, |
| 51 | + workingDirectory: nil, |
| 52 | + autoResponses: [ |
| 53 | + "💫": "/usage\r", |
| 54 | + ] |
| 55 | + ) |
| 56 | + } catch { |
| 57 | + AppLog.probes.error("Kimi CLI probe failed: \(error.localizedDescription)") |
| 58 | + throw ProbeError.executionFailed(error.localizedDescription) |
| 59 | + } |
| 60 | + |
| 61 | + AppLog.probes.info("Kimi CLI /usage output:\n\(result.output)") |
| 62 | + |
| 63 | + let snapshot = try Self.parse(result.output) |
| 64 | + |
| 65 | + AppLog.probes.info("Kimi CLI probe success: \(snapshot.quotas.count) quotas found") |
| 66 | + for quota in snapshot.quotas { |
| 67 | + AppLog.probes.info(" - \(quota.quotaType.displayName): \(Int(quota.percentRemaining))% remaining") |
| 68 | + } |
| 69 | + |
| 70 | + return snapshot |
| 71 | + } |
| 72 | + |
| 73 | + // MARK: - Static Parsing (for testability) |
| 74 | + |
| 75 | + /// Parses the Kimi CLI `/usage` output into a UsageSnapshot. |
| 76 | + /// |
| 77 | + /// Expected format per quota line: |
| 78 | + /// ``` |
| 79 | + /// Weekly limit ━━━━━━━━━━━━━━━━━━━━ 100% left (resets in 6d 23h 22m) |
| 80 | + /// 5h limit ━━━━━━━━━━━━━━━━━━━━ 75% left (resets in 4h 22m) |
| 81 | + /// ``` |
| 82 | + public static func parse(_ text: String) throws -> UsageSnapshot { |
| 83 | + // Pattern: label + optional progress bar + percent + "left" + reset info |
| 84 | + // Real CLI output may omit progress bar chars (just whitespace between label and %) |
| 85 | + let pattern = #"([\w\s]+?)\s{2,}[━░─]*\s*(\d+)%\s+left\s+\(resets\s+in\s+(.+?)\)"# |
| 86 | + guard let regex = try? NSRegularExpression(pattern: pattern, options: [.caseInsensitive]) else { |
| 87 | + throw ProbeError.parseFailed("Invalid regex pattern") |
| 88 | + } |
| 89 | + |
| 90 | + let range = NSRange(text.startIndex..<text.endIndex, in: text) |
| 91 | + let matches = regex.matches(in: text, options: [], range: range) |
| 92 | + |
| 93 | + guard !matches.isEmpty else { |
| 94 | + throw ProbeError.parseFailed("No quota data found in Kimi CLI output") |
| 95 | + } |
| 96 | + |
| 97 | + var quotas: [UsageQuota] = [] |
| 98 | + |
| 99 | + for match in matches { |
| 100 | + guard match.numberOfRanges >= 4, |
| 101 | + let labelRange = Range(match.range(at: 1), in: text), |
| 102 | + let percentRange = Range(match.range(at: 2), in: text), |
| 103 | + let resetRange = Range(match.range(at: 3), in: text) else { |
| 104 | + continue |
| 105 | + } |
| 106 | + |
| 107 | + let label = String(text[labelRange]).trimmingCharacters(in: .whitespaces).lowercased() |
| 108 | + let percent = Double(text[percentRange]) ?? 0.0 |
| 109 | + let resetText = String(text[resetRange]).trimmingCharacters(in: .whitespaces) |
| 110 | + |
| 111 | + let quotaType: QuotaType |
| 112 | + if label.contains("weekly") { |
| 113 | + quotaType = .weekly |
| 114 | + } else if label.contains("5h") || label.contains("hour") { |
| 115 | + quotaType = .session |
| 116 | + } else { |
| 117 | + quotaType = .timeLimit(label) |
| 118 | + } |
| 119 | + |
| 120 | + let resetsAt = parseResetDuration(resetText) |
| 121 | + |
| 122 | + quotas.append(UsageQuota( |
| 123 | + percentRemaining: percent, |
| 124 | + quotaType: quotaType, |
| 125 | + providerId: "kimi", |
| 126 | + resetsAt: resetsAt, |
| 127 | + resetText: "Resets in \(resetText)" |
| 128 | + )) |
| 129 | + } |
| 130 | + |
| 131 | + return UsageSnapshot( |
| 132 | + providerId: "kimi", |
| 133 | + quotas: quotas, |
| 134 | + capturedAt: Date() |
| 135 | + ) |
| 136 | + } |
| 137 | + |
| 138 | + // MARK: - Private Helpers |
| 139 | + |
| 140 | + /// Parses a relative duration string like "6d 23h 22m" or "4h 22m" into a future Date. |
| 141 | + static func parseResetDuration(_ text: String) -> Date? { |
| 142 | + var totalSeconds: TimeInterval = 0 |
| 143 | + |
| 144 | + // Extract days |
| 145 | + if let dayMatch = text.range(of: #"(\d+)\s*d"#, options: .regularExpression) { |
| 146 | + let dayStr = String(text[dayMatch]) |
| 147 | + if let days = Int(dayStr.filter { $0.isNumber }) { |
| 148 | + totalSeconds += Double(days) * 24 * 3600 |
| 149 | + } |
| 150 | + } |
| 151 | + |
| 152 | + // Extract hours |
| 153 | + if let hourMatch = text.range(of: #"(\d+)\s*h"#, options: .regularExpression) { |
| 154 | + let hourStr = String(text[hourMatch]) |
| 155 | + if let hours = Int(hourStr.filter { $0.isNumber }) { |
| 156 | + totalSeconds += Double(hours) * 3600 |
| 157 | + } |
| 158 | + } |
| 159 | + |
| 160 | + // Extract minutes |
| 161 | + if let minMatch = text.range(of: #"(\d+)\s*m"#, options: .regularExpression) { |
| 162 | + let minStr = String(text[minMatch]) |
| 163 | + if let minutes = Int(minStr.filter { $0.isNumber }) { |
| 164 | + totalSeconds += Double(minutes) * 60 |
| 165 | + } |
| 166 | + } |
| 167 | + |
| 168 | + guard totalSeconds > 0 else { return nil } |
| 169 | + return Date().addingTimeInterval(totalSeconds) |
| 170 | + } |
| 171 | +} |
0 commit comments