Add WAL-safe Kimi Desktop monthly enrichment

CHANGELOG.md

@@ -3,6 +3,7 @@
 ## 0.47.1 — Unreleased
 
 ### Added
+- 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!
 - 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!
 
 ### Changed

Sources/CodexBarCore/Providers/Kimi/KimiCookieImporter.swift

@@ -4,6 +4,10 @@ import Foundation
 import SweetCookieKit
 
 public enum KimiCookieImporter {
+    public static func desktopAuthToken() -> String? {
+        KimiDesktopAuthToken.load()
+    }
+
     private static let log = CodexBarLog.logger(LogCategories.kimiCookie)
     private static let cookieClient = BrowserCookieClient()
     private static let cookieDomains = ["www.kimi.com", "kimi.com"]

Sources/CodexBarCore/Providers/Kimi/KimiDesktopAuthToken.swift

@@ -0,0 +1,127 @@
+import Foundation
+
+#if canImport(SQLite3)
+import SQLite3
+#elseif canImport(CSQLite3)
+import CSQLite3
+#endif
+
+#if canImport(SQLite3) || canImport(CSQLite3)
+/// Read-only access to the official Kimi Desktop Chromium cookie store.
+public enum KimiDesktopAuthToken: Sendable {
+    private static let log = CodexBarLog.logger(LogCategories.kimiCookie)
+
+    public static func cookiesDatabaseURL(
+        homeDirectory: URL = FileManager.default.homeDirectoryForCurrentUser) -> URL
+    {
+        homeDirectory
+            .appendingPathComponent("Library", isDirectory: true)
+            .appendingPathComponent("Application Support", isDirectory: true)
+            .appendingPathComponent("kimi-desktop", isDirectory: true)
+            .appendingPathComponent("Cookies", isDirectory: false)
+    }
+
+    public static func load(
+        homeDirectory: URL = FileManager.default.homeDirectoryForCurrentUser) -> String?
+    {
+        self.load(databaseURL: self.cookiesDatabaseURL(homeDirectory: homeDirectory))
+    }
+
+    static func load(databaseURL: URL) -> String? {
+        guard FileManager.default.isReadableFile(atPath: databaseURL.path) else { return nil }
+        do {
+            return try self.read(databaseURL: databaseURL, immutable: false)
+        } catch let failure as SQLiteReadFailure {
+            // Chromium can leave the main database in WAL mode after a clean shutdown removes both sidecars.
+            // Immutable mode reads that idle file without recreating sidecars; active WAL databases stay on the
+            // normal read-only path so committed WAL records remain visible.
+            guard failure.code == SQLITE_CANTOPEN, self.walSidecarsAreMissing(databaseURL: databaseURL) else {
+                Self.log.debug("Kimi Desktop Cookies read failed: \(failure.message)")
+                return nil
+            }
+            do {
+                return try self.read(databaseURL: databaseURL, immutable: true)
+            } catch let fallbackFailure as SQLiteReadFailure {
+                Self.log.debug("Kimi Desktop Cookies immutable read failed: \(fallbackFailure.message)")
+                return nil
+            } catch {
+                return nil
+            }
+        } catch {
+            return nil
+        }
+    }
+
+    private static func read(databaseURL: URL, immutable: Bool) throws -> String? {
+        var db: OpaquePointer?
+        let filename = immutable ? "\(databaseURL.absoluteURL.absoluteString)?immutable=1" : databaseURL.path
+        let flags = immutable ? SQLITE_OPEN_READONLY | SQLITE_OPEN_URI : SQLITE_OPEN_READONLY
+        let openResult = sqlite3_open_v2(filename, &db, flags, nil)
+        guard openResult == SQLITE_OK else {
+            let failure = self.sqliteFailure(db: db, resultCode: openResult)
+            sqlite3_close(db)
+            throw failure
+        }
+        defer { sqlite3_close(db) }
+        sqlite3_busy_timeout(db, 250)
+
+        let sql = """
+        SELECT value
+        FROM cookies
+        WHERE name = 'kimi-auth'
+          AND host_key IN ('www.kimi.com', '.www.kimi.com', '.kimi.com', 'kimi.com')
+        ORDER BY last_access_utc DESC
+        LIMIT 1
+        """
+        var statement: OpaquePointer?
+        let prepareResult = sqlite3_prepare_v2(db, sql, -1, &statement, nil)
+        guard prepareResult == SQLITE_OK else {
+            throw self.sqliteFailure(db: db, resultCode: prepareResult)
+        }
+        defer { sqlite3_finalize(statement) }
+
+        let step = sqlite3_step(statement)
+        if step == SQLITE_DONE {
+            return nil
+        }
+        guard step == SQLITE_ROW else {
+            throw self.sqliteFailure(db: db, resultCode: step)
+        }
+        guard let text = sqlite3_column_text(statement, 0) else { return nil }
+        let token = String(cString: text).trimmingCharacters(in: .whitespacesAndNewlines)
+        return token.isEmpty ? nil : token
+    }
+
+    private static func walSidecarsAreMissing(databaseURL: URL) -> Bool {
+        !FileManager.default.fileExists(atPath: databaseURL.path + "-wal") &&
+            !FileManager.default.fileExists(atPath: databaseURL.path + "-shm")
+    }
+
+    private static func sqliteFailure(db: OpaquePointer?, resultCode: Int32) -> SQLiteReadFailure {
+        SQLiteReadFailure(
+            code: db.map { sqlite3_errcode($0) } ?? resultCode,
+            message: db.map { String(cString: sqlite3_errmsg($0)) } ?? "unknown error")
+    }
+
+    private struct SQLiteReadFailure: Error {
+        let code: Int32
+        let message: String
+    }
+}
+#else
+public enum KimiDesktopAuthToken: Sendable {
+    public static func cookiesDatabaseURL(
+        homeDirectory: URL = FileManager.default.homeDirectoryForCurrentUser) -> URL
+    {
+        homeDirectory
+            .appendingPathComponent("Library", isDirectory: true)
+            .appendingPathComponent("Application Support", isDirectory: true)
+            .appendingPathComponent("kimi-desktop", isDirectory: true)
+            .appendingPathComponent("Cookies", isDirectory: false)
+    }
+
+    public static func load(homeDirectory _: URL = FileManager.default.homeDirectoryForCurrentUser) -> String? {
+        nil
+    }
+}
+#endif

Sources/CodexBarCore/Providers/Kimi/KimiProviderDescriptor.swift

@@ -66,9 +66,15 @@ struct KimiAPIFetchStrategy: ProviderFetchStrategy {
     let id: String = "kimi.api"
     let kind: ProviderFetchKind = .apiToken
     private let transport: any ProviderHTTPTransport
+    private let resolveWebAuthToken: @Sendable (ProviderFetchContext) -> String?
 
-    init(transport: any ProviderHTTPTransport = ProviderHTTPClient.shared) {
+    init(
+        transport: any ProviderHTTPTransport = ProviderHTTPClient.shared,
+        resolveWebAuthToken: @escaping @Sendable (ProviderFetchContext) -> String? =
+            KimiWebEnrichmentTokenResolver.resolve)
+    {
         self.transport = transport
+        self.resolveWebAuthToken = resolveWebAuthToken
     }
 
     func isAvailable(_ context: ProviderFetchContext) async -> Bool {
@@ -83,6 +89,7 @@ struct KimiAPIFetchStrategy: ProviderFetchStrategy {
         let snapshot = try await KimiUsageFetcher.fetchCodeAPIUsage(
             apiKey: apiKey,
             baseURL: baseURL,
+            webAuthToken: self.enrichmentToken(context),
             transport: self.transport)
         return self.makeResult(
             usage: snapshot.toUsageSnapshot(),
@@ -92,15 +99,26 @@ struct KimiAPIFetchStrategy: ProviderFetchStrategy {
     func shouldFallback(on error: Error, context: ProviderFetchContext) -> Bool {
         KimiCodeAPIFallbackPolicy.shouldFallback(on: error, context: context)
     }
+
+    private func enrichmentToken(_ context: ProviderFetchContext) -> String? {
+        guard let settings = context.settings?.kimi, settings.cookieSource != .off else { return nil }
+        return self.resolveWebAuthToken(context)
+    }
 }
 
 struct KimiCLICredentialFetchStrategy: ProviderFetchStrategy {
     let id: String = "kimi.cli"
     let kind: ProviderFetchKind = .oauth
     private let transport: any ProviderHTTPTransport
+    private let resolveWebAuthToken: @Sendable (ProviderFetchContext) -> String?
 
-    init(transport: any ProviderHTTPTransport = ProviderHTTPClient.shared) {
+    init(
+        transport: any ProviderHTTPTransport = ProviderHTTPClient.shared,
+        resolveWebAuthToken: @escaping @Sendable (ProviderFetchContext) -> String? =
+            KimiWebEnrichmentTokenResolver.resolve)
+    {
         self.transport = transport
+        self.resolveWebAuthToken = resolveWebAuthToken
     }
 
     func isAvailable(_ context: ProviderFetchContext) async -> Bool {
@@ -120,6 +138,7 @@ struct KimiCLICredentialFetchStrategy: ProviderFetchStrategy {
                 apiKey: token,
                 baseURL: baseURL,
                 identityHeaders: identityHeaders,
+                webAuthToken: self.enrichmentToken(context),
                 transport: self.transport)
         } catch {
             throw Self.normalizedCodeAPIError(error)
@@ -137,6 +156,29 @@ struct KimiCLICredentialFetchStrategy: ProviderFetchStrategy {
         guard case KimiAPIError.invalidAPIKey = error else { return error }
         return KimiAPIError.invalidCodeCredential
     }
+
+    private func enrichmentToken(_ context: ProviderFetchContext) -> String? {
+        guard let settings = context.settings?.kimi, settings.cookieSource != .off else { return nil }
+        return self.resolveWebAuthToken(context)
+    }
+}
+
+enum KimiWebEnrichmentTokenResolver {
+    static func resolve(_ context: ProviderFetchContext) -> String? {
+        guard let settings = context.settings?.kimi, settings.cookieSource != .off else { return nil }
+        if let override = KimiCookieHeader.resolveCookieOverride(context: context) {
+            return override.token
+        }
+        #if os(macOS)
+        if let token = KimiCookieImporter.desktopAuthToken() {
+            return token
+        }
+        if let token = try? KimiCookieImporter.importSession().authToken {
+            return token
+        }
+        #endif
+        return nil
+    }
 }
 
 private enum KimiCodeAPIFallbackPolicy {
@@ -183,6 +225,9 @@ struct KimiWebFetchStrategy: ProviderFetchStrategy {
 
         #if os(macOS)
         if context.settings?.kimi?.cookieSource != .off {
+            if KimiCookieImporter.desktopAuthToken() != nil {
+                return true
+            }
             return KimiCookieImporter.hasSession()
         }
         #endif
@@ -220,6 +265,9 @@ struct KimiWebFetchStrategy: ProviderFetchStrategy {
         // Try browser cookie import when auto mode is enabled
         #if os(macOS)
         if context.settings?.kimi?.cookieSource != .off {
+            if let token = KimiCookieImporter.desktopAuthToken() {
+                return token
+            }
             do {
                 let session = try KimiCookieImporter.importSession()
                 if let token = session.authToken {

Sources/CodexBarCore/Providers/Kimi/KimiUsageFetcher.swift

@@ -16,6 +16,7 @@ public struct KimiUsageFetcher: Sendable {
         apiKey: String,
         baseURL: URL = KimiSettingsReader.defaultCodeAPIBaseURL,
         identityHeaders: [String: String] = [:],
+        webAuthToken: String? = nil,
         now: Date = Date(),
         transport: any ProviderHTTPTransport = ProviderHTTPClient.shared) async throws -> KimiUsageSnapshot
     {
@@ -44,7 +45,13 @@ public struct KimiUsageFetcher: Sendable {
             throw self.codeAPIError(statusCode: response.statusCode)
         }
 
-        return try self.parseCodeAPIUsage(from: data, now: now)
+        let snapshot = try self.parseCodeAPIUsage(from: data, now: now)
+        guard let webAuthToken else { return snapshot }
+        return try await self.enrichCodeAPIUsage(
+            snapshot,
+            webAuthToken: webAuthToken,
+            now: now,
+            transport: transport)
     }
 
     static func _parseCodeAPIUsageForTesting(_ data: Data, now: Date = Date()) throws -> KimiUsageSnapshot {
@@ -179,6 +186,35 @@ public struct KimiUsageFetcher: Sendable {
         return codingUsage
     }
 
+    private static func enrichCodeAPIUsage(
+        _ snapshot: KimiUsageSnapshot,
+        webAuthToken: String,
+        now: Date,
+        transport: any ProviderHTTPTransport) async throws -> KimiUsageSnapshot
+    {
+        let sessionInfo = self.decodeSessionInfo(from: webAuthToken)
+        let subscriptionStats: KimiSubscriptionStatsResponse?
+        do {
+            subscriptionStats = try await self.fetchSubscriptionStats(
+                authToken: webAuthToken,
+                sessionInfo: sessionInfo,
+                transport: transport)
+        } catch is CancellationError {
+            throw CancellationError()
+        } catch {
+            Self.log.warning("Kimi Code monthly enrichment unavailable: \(error.localizedDescription)")
+            return snapshot
+        }
+        guard let subscriptionStats else { return snapshot }
+        return KimiUsageSnapshot(
+            weekly: snapshot.weekly,
+            rateLimit: snapshot.rateLimit,
+            rateLimitWindow: snapshot.rateLimitWindow,
+            subscriptionBalance: subscriptionStats.subscriptionBalance,
+            subscriptionCodeWeeklyLimit: subscriptionStats.ratelimitCode7d,
+            updatedAt: now)
+    }
+
     private static func parseCodeAPIUsage(from data: Data, now: Date) throws -> KimiUsageSnapshot {
         let response = try JSONDecoder().decode(KimiCodeAPIUsageResponse.self, from: data)
         let rateLimit = response.limits?.first