Skip to content

Commit dccc855

Browse files
committed
feat(kimi): add CLI usage probe and parsing tests
- Implement KimiCLIUsageProbe for CLI-based quota retrieval - Add KimiProbeMode to switch between CLI and API modes - Provide unit tests for CLI usage probe and output parsing
1 parent da51ae8 commit dccc855

File tree

4 files changed

+554
-0
lines changed

4 files changed

+554
-0
lines changed
Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
import Foundation
2+
3+
/// The mode used by KimiProvider to fetch usage data.
4+
/// Users can switch between CLI (default) and API modes in Settings.
5+
public enum KimiProbeMode: String, Sendable, Equatable, CaseIterable {
6+
/// Use the Kimi CLI (`kimi` with `/usage` command) to fetch usage data.
7+
/// This is the default mode and works via interactive CLI prompt.
8+
case cli
9+
10+
/// Use the Kimi HTTP API to fetch usage data directly.
11+
/// Requires valid browser cookie authentication (kimi-auth).
12+
case api
13+
14+
/// Human-readable display name for the mode
15+
public var displayName: String {
16+
switch self {
17+
case .cli:
18+
return "CLI"
19+
case .api:
20+
return "API"
21+
}
22+
}
23+
24+
/// Description of what this mode does
25+
public var description: String {
26+
switch self {
27+
case .cli:
28+
return "Uses kimi /usage command"
29+
case .api:
30+
return "Calls Kimi API directly"
31+
}
32+
}
33+
}
Lines changed: 171 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,171 @@
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

Comments
 (0)