Skip to content

Commit 6afa672

Browse files
steipetePeter Steinbergerhaoli
authored
feat: enrich Kimi monthly usage from desktop (#2622)
Co-authored-by: Peter Steinberger <steipete@mac-studio-sf2.local> Co-authored-by: haoli <haoli@local.dev>
1 parent 7002b57 commit 6afa672

8 files changed

Lines changed: 477 additions & 6 deletions

File tree

CHANGELOG.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
## 0.47.1 — Unreleased
44

55
### Added
6+
- Kimi: enrich Code API and CLI usage with the monthly membership pool from a signed-in Kimi Desktop session, using WAL-safe read-only cookie access (#2351). Thanks @Leehow!
67
- Kimi/GLM: distinguish Kimi Code from the regional Open Platform, bind China and international keys to their issuing hosts, and show GLM Coding Plan's 5-hour window as primary with MCP separate (#2351). Thanks @Leehow!
78

89
### Changed

Sources/CodexBarCore/Providers/Kimi/KimiCookieImporter.swift

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,10 @@ import Foundation
44
import SweetCookieKit
55

66
public enum KimiCookieImporter {
7+
public static func desktopAuthToken() -> String? {
8+
KimiDesktopAuthToken.load()
9+
}
10+
711
private static let log = CodexBarLog.logger(LogCategories.kimiCookie)
812
private static let cookieClient = BrowserCookieClient()
913
private static let cookieDomains = ["www.kimi.com", "kimi.com"]
Lines changed: 127 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,127 @@
1+
import Foundation
2+
3+
#if canImport(SQLite3)
4+
import SQLite3
5+
#elseif canImport(CSQLite3)
6+
import CSQLite3
7+
#endif
8+
9+
#if canImport(SQLite3) || canImport(CSQLite3)
10+
/// Read-only access to the official Kimi Desktop Chromium cookie store.
11+
public enum KimiDesktopAuthToken: Sendable {
12+
private static let log = CodexBarLog.logger(LogCategories.kimiCookie)
13+
14+
public static func cookiesDatabaseURL(
15+
homeDirectory: URL = FileManager.default.homeDirectoryForCurrentUser) -> URL
16+
{
17+
homeDirectory
18+
.appendingPathComponent("Library", isDirectory: true)
19+
.appendingPathComponent("Application Support", isDirectory: true)
20+
.appendingPathComponent("kimi-desktop", isDirectory: true)
21+
.appendingPathComponent("Cookies", isDirectory: false)
22+
}
23+
24+
public static func load(
25+
homeDirectory: URL = FileManager.default.homeDirectoryForCurrentUser) -> String?
26+
{
27+
self.load(databaseURL: self.cookiesDatabaseURL(homeDirectory: homeDirectory))
28+
}
29+
30+
static func load(databaseURL: URL) -> String? {
31+
guard FileManager.default.isReadableFile(atPath: databaseURL.path) else { return nil }
32+
do {
33+
return try self.read(databaseURL: databaseURL, immutable: false)
34+
} catch let failure as SQLiteReadFailure {
35+
// Chromium can leave the main database in WAL mode after a clean shutdown removes both sidecars.
36+
// Immutable mode reads that idle file without recreating sidecars; active WAL databases stay on the
37+
// normal read-only path so committed WAL records remain visible.
38+
guard failure.code == SQLITE_CANTOPEN, self.walSidecarsAreMissing(databaseURL: databaseURL) else {
39+
Self.log.debug("Kimi Desktop Cookies read failed: \(failure.message)")
40+
return nil
41+
}
42+
do {
43+
return try self.read(databaseURL: databaseURL, immutable: true)
44+
} catch let fallbackFailure as SQLiteReadFailure {
45+
Self.log.debug("Kimi Desktop Cookies immutable read failed: \(fallbackFailure.message)")
46+
return nil
47+
} catch {
48+
return nil
49+
}
50+
} catch {
51+
return nil
52+
}
53+
}
54+
55+
private static func read(databaseURL: URL, immutable: Bool) throws -> String? {
56+
var db: OpaquePointer?
57+
let filename = immutable ? "\(databaseURL.absoluteURL.absoluteString)?immutable=1" : databaseURL.path
58+
let flags = immutable ? SQLITE_OPEN_READONLY | SQLITE_OPEN_URI : SQLITE_OPEN_READONLY
59+
let openResult = sqlite3_open_v2(filename, &db, flags, nil)
60+
guard openResult == SQLITE_OK else {
61+
let failure = self.sqliteFailure(db: db, resultCode: openResult)
62+
sqlite3_close(db)
63+
throw failure
64+
}
65+
defer { sqlite3_close(db) }
66+
sqlite3_busy_timeout(db, 250)
67+
68+
let sql = """
69+
SELECT value
70+
FROM cookies
71+
WHERE name = 'kimi-auth'
72+
AND host_key IN ('www.kimi.com', '.www.kimi.com', '.kimi.com', 'kimi.com')
73+
ORDER BY last_access_utc DESC
74+
LIMIT 1
75+
"""
76+
var statement: OpaquePointer?
77+
let prepareResult = sqlite3_prepare_v2(db, sql, -1, &statement, nil)
78+
guard prepareResult == SQLITE_OK else {
79+
throw self.sqliteFailure(db: db, resultCode: prepareResult)
80+
}
81+
defer { sqlite3_finalize(statement) }
82+
83+
let step = sqlite3_step(statement)
84+
if step == SQLITE_DONE {
85+
return nil
86+
}
87+
guard step == SQLITE_ROW else {
88+
throw self.sqliteFailure(db: db, resultCode: step)
89+
}
90+
guard let text = sqlite3_column_text(statement, 0) else { return nil }
91+
let token = String(cString: text).trimmingCharacters(in: .whitespacesAndNewlines)
92+
return token.isEmpty ? nil : token
93+
}
94+
95+
private static func walSidecarsAreMissing(databaseURL: URL) -> Bool {
96+
!FileManager.default.fileExists(atPath: databaseURL.path + "-wal") &&
97+
!FileManager.default.fileExists(atPath: databaseURL.path + "-shm")
98+
}
99+
100+
private static func sqliteFailure(db: OpaquePointer?, resultCode: Int32) -> SQLiteReadFailure {
101+
SQLiteReadFailure(
102+
code: db.map { sqlite3_errcode($0) } ?? resultCode,
103+
message: db.map { String(cString: sqlite3_errmsg($0)) } ?? "unknown error")
104+
}
105+
106+
private struct SQLiteReadFailure: Error {
107+
let code: Int32
108+
let message: String
109+
}
110+
}
111+
#else
112+
public enum KimiDesktopAuthToken: Sendable {
113+
public static func cookiesDatabaseURL(
114+
homeDirectory: URL = FileManager.default.homeDirectoryForCurrentUser) -> URL
115+
{
116+
homeDirectory
117+
.appendingPathComponent("Library", isDirectory: true)
118+
.appendingPathComponent("Application Support", isDirectory: true)
119+
.appendingPathComponent("kimi-desktop", isDirectory: true)
120+
.appendingPathComponent("Cookies", isDirectory: false)
121+
}
122+
123+
public static func load(homeDirectory _: URL = FileManager.default.homeDirectoryForCurrentUser) -> String? {
124+
nil
125+
}
126+
}
127+
#endif

Sources/CodexBarCore/Providers/Kimi/KimiProviderDescriptor.swift

Lines changed: 50 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -66,9 +66,15 @@ struct KimiAPIFetchStrategy: ProviderFetchStrategy {
6666
let id: String = "kimi.api"
6767
let kind: ProviderFetchKind = .apiToken
6868
private let transport: any ProviderHTTPTransport
69+
private let resolveWebAuthToken: @Sendable (ProviderFetchContext) -> String?
6970

70-
init(transport: any ProviderHTTPTransport = ProviderHTTPClient.shared) {
71+
init(
72+
transport: any ProviderHTTPTransport = ProviderHTTPClient.shared,
73+
resolveWebAuthToken: @escaping @Sendable (ProviderFetchContext) -> String? =
74+
KimiWebEnrichmentTokenResolver.resolve)
75+
{
7176
self.transport = transport
77+
self.resolveWebAuthToken = resolveWebAuthToken
7278
}
7379

7480
func isAvailable(_ context: ProviderFetchContext) async -> Bool {
@@ -83,6 +89,7 @@ struct KimiAPIFetchStrategy: ProviderFetchStrategy {
8389
let snapshot = try await KimiUsageFetcher.fetchCodeAPIUsage(
8490
apiKey: apiKey,
8591
baseURL: baseURL,
92+
webAuthToken: self.enrichmentToken(context),
8693
transport: self.transport)
8794
return self.makeResult(
8895
usage: snapshot.toUsageSnapshot(),
@@ -92,15 +99,26 @@ struct KimiAPIFetchStrategy: ProviderFetchStrategy {
9299
func shouldFallback(on error: Error, context: ProviderFetchContext) -> Bool {
93100
KimiCodeAPIFallbackPolicy.shouldFallback(on: error, context: context)
94101
}
102+
103+
private func enrichmentToken(_ context: ProviderFetchContext) -> String? {
104+
guard let settings = context.settings?.kimi, settings.cookieSource != .off else { return nil }
105+
return self.resolveWebAuthToken(context)
106+
}
95107
}
96108

97109
struct KimiCLICredentialFetchStrategy: ProviderFetchStrategy {
98110
let id: String = "kimi.cli"
99111
let kind: ProviderFetchKind = .oauth
100112
private let transport: any ProviderHTTPTransport
113+
private let resolveWebAuthToken: @Sendable (ProviderFetchContext) -> String?
101114

102-
init(transport: any ProviderHTTPTransport = ProviderHTTPClient.shared) {
115+
init(
116+
transport: any ProviderHTTPTransport = ProviderHTTPClient.shared,
117+
resolveWebAuthToken: @escaping @Sendable (ProviderFetchContext) -> String? =
118+
KimiWebEnrichmentTokenResolver.resolve)
119+
{
103120
self.transport = transport
121+
self.resolveWebAuthToken = resolveWebAuthToken
104122
}
105123

106124
func isAvailable(_ context: ProviderFetchContext) async -> Bool {
@@ -120,6 +138,7 @@ struct KimiCLICredentialFetchStrategy: ProviderFetchStrategy {
120138
apiKey: token,
121139
baseURL: baseURL,
122140
identityHeaders: identityHeaders,
141+
webAuthToken: self.enrichmentToken(context),
123142
transport: self.transport)
124143
} catch {
125144
throw Self.normalizedCodeAPIError(error)
@@ -137,6 +156,29 @@ struct KimiCLICredentialFetchStrategy: ProviderFetchStrategy {
137156
guard case KimiAPIError.invalidAPIKey = error else { return error }
138157
return KimiAPIError.invalidCodeCredential
139158
}
159+
160+
private func enrichmentToken(_ context: ProviderFetchContext) -> String? {
161+
guard let settings = context.settings?.kimi, settings.cookieSource != .off else { return nil }
162+
return self.resolveWebAuthToken(context)
163+
}
164+
}
165+
166+
enum KimiWebEnrichmentTokenResolver {
167+
static func resolve(_ context: ProviderFetchContext) -> String? {
168+
guard let settings = context.settings?.kimi, settings.cookieSource != .off else { return nil }
169+
if let override = KimiCookieHeader.resolveCookieOverride(context: context) {
170+
return override.token
171+
}
172+
#if os(macOS)
173+
if let token = KimiCookieImporter.desktopAuthToken() {
174+
return token
175+
}
176+
if let token = try? KimiCookieImporter.importSession().authToken {
177+
return token
178+
}
179+
#endif
180+
return nil
181+
}
140182
}
141183

142184
private enum KimiCodeAPIFallbackPolicy {
@@ -183,6 +225,9 @@ struct KimiWebFetchStrategy: ProviderFetchStrategy {
183225

184226
#if os(macOS)
185227
if context.settings?.kimi?.cookieSource != .off {
228+
if KimiCookieImporter.desktopAuthToken() != nil {
229+
return true
230+
}
186231
return KimiCookieImporter.hasSession()
187232
}
188233
#endif
@@ -220,6 +265,9 @@ struct KimiWebFetchStrategy: ProviderFetchStrategy {
220265
// Try browser cookie import when auto mode is enabled
221266
#if os(macOS)
222267
if context.settings?.kimi?.cookieSource != .off {
268+
if let token = KimiCookieImporter.desktopAuthToken() {
269+
return token
270+
}
223271
do {
224272
let session = try KimiCookieImporter.importSession()
225273
if let token = session.authToken {

Sources/CodexBarCore/Providers/Kimi/KimiUsageFetcher.swift

Lines changed: 37 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ public struct KimiUsageFetcher: Sendable {
1616
apiKey: String,
1717
baseURL: URL = KimiSettingsReader.defaultCodeAPIBaseURL,
1818
identityHeaders: [String: String] = [:],
19+
webAuthToken: String? = nil,
1920
now: Date = Date(),
2021
transport: any ProviderHTTPTransport = ProviderHTTPClient.shared) async throws -> KimiUsageSnapshot
2122
{
@@ -44,7 +45,13 @@ public struct KimiUsageFetcher: Sendable {
4445
throw self.codeAPIError(statusCode: response.statusCode)
4546
}
4647

47-
return try self.parseCodeAPIUsage(from: data, now: now)
48+
let snapshot = try self.parseCodeAPIUsage(from: data, now: now)
49+
guard let webAuthToken else { return snapshot }
50+
return try await self.enrichCodeAPIUsage(
51+
snapshot,
52+
webAuthToken: webAuthToken,
53+
now: now,
54+
transport: transport)
4855
}
4956

5057
static func _parseCodeAPIUsageForTesting(_ data: Data, now: Date = Date()) throws -> KimiUsageSnapshot {
@@ -179,6 +186,35 @@ public struct KimiUsageFetcher: Sendable {
179186
return codingUsage
180187
}
181188

189+
private static func enrichCodeAPIUsage(
190+
_ snapshot: KimiUsageSnapshot,
191+
webAuthToken: String,
192+
now: Date,
193+
transport: any ProviderHTTPTransport) async throws -> KimiUsageSnapshot
194+
{
195+
let sessionInfo = self.decodeSessionInfo(from: webAuthToken)
196+
let subscriptionStats: KimiSubscriptionStatsResponse?
197+
do {
198+
subscriptionStats = try await self.fetchSubscriptionStats(
199+
authToken: webAuthToken,
200+
sessionInfo: sessionInfo,
201+
transport: transport)
202+
} catch is CancellationError {
203+
throw CancellationError()
204+
} catch {
205+
Self.log.warning("Kimi Code monthly enrichment unavailable: \(error.localizedDescription)")
206+
return snapshot
207+
}
208+
guard let subscriptionStats else { return snapshot }
209+
return KimiUsageSnapshot(
210+
weekly: snapshot.weekly,
211+
rateLimit: snapshot.rateLimit,
212+
rateLimitWindow: snapshot.rateLimitWindow,
213+
subscriptionBalance: subscriptionStats.subscriptionBalance,
214+
subscriptionCodeWeeklyLimit: subscriptionStats.ratelimitCode7d,
215+
updatedAt: now)
216+
}
217+
182218
private static func parseCodeAPIUsage(from data: Data, now: Date) throws -> KimiUsageSnapshot {
183219
let response = try JSONDecoder().decode(KimiCodeAPIUsageResponse.self, from: data)
184220
let rateLimit = response.limits?.first

0 commit comments

Comments
 (0)