From 5de96c575d5b5ca05281cfdd2c87287859ed5a96 Mon Sep 17 00:00:00 2001 From: Miguel Machado Date: Mon, 31 Aug 2026 17:13:35 -0300 Subject: [PATCH 1/2] fix(mac): accept Kimi Code Plan API keys Kimi's special settings path bypassed provider-scoped API credentials and required a short-lived CLI token. Save the Code Plan key in CodeBurn's Keychain item and prefer it for the existing usage endpoint, with the CLI token retained as fallback. --- mac/Sources/CodeBurnMenubar/AppStore.swift | 33 ++++-- .../Data/KimiSubscriptionService.swift | 104 +++++++++++------- .../CodeBurnMenubar/Views/SettingsView.swift | 85 +++++++++++++- .../KimiUsageParsingTests.swift | 30 +++++ 4 files changed, 199 insertions(+), 53 deletions(-) diff --git a/mac/Sources/CodeBurnMenubar/AppStore.swift b/mac/Sources/CodeBurnMenubar/AppStore.swift index fd9bd63d3..dafdb2123 100644 --- a/mac/Sources/CodeBurnMenubar/AppStore.swift +++ b/mac/Sources/CodeBurnMenubar/AppStore.swift @@ -146,10 +146,12 @@ final class AppStore { var kimiUsage: KimiUsage? var kimiError: String? - // No keychain dance for Kimi — "connected" just means the CLI's - // credential file exists, so we start dormant and auto-activate on the - // first refresh tick. - var kimiLoadState: SubscriptionLoadState = KimiSubscriptionService.hasCredential ? .dormant : .notBootstrapped + // A provider-scoped Keychain key or the CLI credential makes Kimi eligible + // for prompt-free background activation. + var kimiLoadState: SubscriptionLoadState = ( + KimiSubscriptionService.hasCredential + || CapacityDockProviderCredentialPresence.contains(KimiSubscriptionService.providerID) + ) ? .dormant : .notBootstrapped var geminiUsage: GeminiUsage? var geminiError: String? @@ -1228,16 +1230,26 @@ final class AppStore { // MARK: - Kimi Code - /// Unlike Claude/Codex there is no keychain bootstrap: reading the CLI's - /// credential file is prompt-free, so the first refresh tick activates - /// the dormant state automatically. + /// A saved Code Plan API key wins; the CLI's short-lived token remains the + /// fallback for existing users who have not configured one. + private func fetchKimiUsage() async throws -> KimiUsage { + let apiKey: String? + if CapacityDockProviderCredentialPresence.contains(KimiSubscriptionService.providerID) { + apiKey = try await capacityDockCredentialLoader(KimiSubscriptionService.providerID) + .sanitizedOverride.apiKey + } else { + apiKey = nil + } + return try await KimiSubscriptionService.refresh(apiKey: apiKey) + } + func bootstrapKimi() async { // Capture the generation before the await so a disconnect that lands // mid-fetch cannot be resurrected into .loaded when the fetch returns. let gen = kimiRefreshGen kimiLoadState = .bootstrapping do { - let usage = try await KimiSubscriptionService.refresh() + let usage = try await fetchKimiUsage() guard gen == kimiRefreshGen else { return } kimiUsage = usage kimiError = nil @@ -1262,14 +1274,15 @@ final class AppStore { await bootstrapKimi() return kimiLoadState == .loaded } - guard KimiSubscriptionService.hasCredential else { + guard KimiSubscriptionService.hasCredential + || CapacityDockProviderCredentialPresence.contains(KimiSubscriptionService.providerID) else { if kimiLoadState != .notBootstrapped { kimiLoadState = .notBootstrapped } return false } let gen = kimiRefreshGen if kimiUsage == nil { kimiLoadState = .loading } do { - let usage = try await KimiSubscriptionService.refresh() + let usage = try await fetchKimiUsage() guard gen == kimiRefreshGen else { return false } kimiUsage = usage kimiError = nil diff --git a/mac/Sources/CodeBurnMenubar/Data/KimiSubscriptionService.swift b/mac/Sources/CodeBurnMenubar/Data/KimiSubscriptionService.swift index 4b6aa9c59..2d0b19922 100644 --- a/mac/Sources/CodeBurnMenubar/Data/KimiSubscriptionService.swift +++ b/mac/Sources/CodeBurnMenubar/Data/KimiSubscriptionService.swift @@ -31,18 +31,44 @@ struct KimiUsage: Sendable, Equatable { let fetchedAt: Date } -/// Mirror of CodexSubscriptionService for Kimi Code. Reads the CLI's -/// credential file directly (~/.kimi-code/credentials/kimi-code.json) — -/// no keychain bootstrap, no OAuth refresh. Tokens are short-lived -/// (~15 min) and only the Kimi CLI refreshes them, so an expired token is -/// a terminal state: the UI tells the user to run the CLI once. +/// Kimi Code quota client. A user-supplied Code Plan API key takes priority; +/// otherwise it reads the CLI credential file directly. CLI tokens are +/// short-lived and only the Kimi CLI refreshes them. enum KimiSubscriptionService { private static let usageURL = URL(string: "https://api.kimi.com/coding/v1/usages")! private static let usageBlockedUntilKey = "codeburn.kimi.usage.blockedUntil" + static let providerID = "kimi" + + struct Dependencies: Sendable { + var fetch: @Sendable (URLRequest) async throws -> (Data, HTTPURLResponse) + var readFile: @Sendable (URL) -> Data? + var credentialsURL: URL + var deviceIDURL: URL + var now: @Sendable () -> Date + + static let live: Dependencies = { + let home = ProcessInfo.processInfo.environment["KIMI_CODE_HOME"] + ?? NSHomeDirectory() + "/.kimi-code" + return Dependencies( + fetch: { request in + let (data, response) = try await URLSession.shared.data(for: request) + guard let http = response as? HTTPURLResponse else { + throw FetchError.usageHTTPError(-1, nil) + } + return (data, http) + }, + readFile: { FileManager.default.contents(atPath: $0.path) }, + credentialsURL: URL(fileURLWithPath: home + "/credentials/kimi-code.json"), + deviceIDURL: URL(fileURLWithPath: home + "/device_id"), + now: { Date() } + ) + }() + } enum FetchError: Error, LocalizedError { case noCredentials case tokenExpired + case apiKeyRejected case rateLimited(retryAt: Date) case usageHTTPError(Int, String?) case usageDecodeFailed @@ -51,9 +77,11 @@ enum KimiSubscriptionService { var errorDescription: String? { switch self { case .noCredentials: - return "No Kimi Code credentials found. Sign in with the Kimi CLI first." + return "No Kimi Code credentials found. Save a Code Plan API key or sign in with the Kimi CLI." case .tokenExpired: return "Kimi Code login expired. Run the Kimi CLI once to refresh, then try again." + case .apiKeyRejected: + return "Kimi Code rejected this API key. Check the Code Plan key saved in Settings." case let .rateLimited(retryAt): let f = RelativeDateTimeFormatter() f.unitsStyle = .short @@ -66,9 +94,10 @@ enum KimiSubscriptionService { } var isTerminal: Bool { - if case .tokenExpired = self { return true } - if case .noCredentials = self { return true } - return false + switch self { + case .noCredentials, .tokenExpired, .apiKeyRejected: return true + case .rateLimited, .usageHTTPError, .usageDecodeFailed, .network: return false + } } var rateLimitRetryAt: Date? { @@ -102,34 +131,26 @@ enum KimiSubscriptionService { } } - private static var credentialsURL: URL { - let home = ProcessInfo.processInfo.environment["KIMI_CODE_HOME"] - ?? NSHomeDirectory() + "/.kimi-code" - return URL(fileURLWithPath: home + "/credentials/kimi-code.json") - } - static var hasCredential: Bool { - FileManager.default.fileExists(atPath: credentialsURL.path) + FileManager.default.fileExists(atPath: Dependencies.live.credentialsURL.path) } /// Returns the access token only when it is still fresh (60s skew). /// Throws noCredentials / tokenExpired otherwise. - private static func freshToken() throws -> String { - guard let data = FileManager.default.contents(atPath: credentialsURL.path), + private static func freshToken(deps: Dependencies) throws -> String { + guard let data = deps.readFile(deps.credentialsURL), let cred = try? JSONDecoder().decode(CredentialFile.self, from: data), !cred.accessToken.isEmpty else { throw FetchError.noCredentials } - guard cred.expiresAt > Date().timeIntervalSince1970 + 60 else { + guard cred.expiresAt > deps.now().timeIntervalSince1970 + 60 else { throw FetchError.tokenExpired } return cred.accessToken } - private static func deviceId() -> String? { - let home = ProcessInfo.processInfo.environment["KIMI_CODE_HOME"] - ?? NSHomeDirectory() + "/.kimi-code" - guard let data = FileManager.default.contents(atPath: home + "/device_id"), + private static func deviceId(deps: Dependencies) -> String? { + guard let data = deps.readFile(deps.deviceIDURL), let id = String(data: data, encoding: .utf8)?.trimmingCharacters(in: .whitespacesAndNewlines), !id.isEmpty else { return nil } return id @@ -137,11 +158,18 @@ enum KimiSubscriptionService { // MARK: - Fetch - static func refresh() async throws -> KimiUsage { - if let until = usageBlockedUntil(), until > Date() { + static func refresh(apiKey: String? = nil, deps: Dependencies = .live) async throws -> KimiUsage { + if let until = usageBlockedUntil(), until > deps.now() { throw FetchError.rateLimited(retryAt: until) } - let token = try freshToken() + let configuredKey = apiKey?.trimmingCharacters(in: .whitespacesAndNewlines) + let usesAPIKey = configuredKey?.isEmpty == false + let token: String + if let configuredKey, !configuredKey.isEmpty { + token = configuredKey + } else { + token = try freshToken(deps: deps) + } var request = URLRequest(url: usageURL) request.httpMethod = "GET" @@ -149,22 +177,24 @@ enum KimiSubscriptionService { request.setValue("Bearer \(token)", forHTTPHeaderField: "Authorization") request.setValue("application/json", forHTTPHeaderField: "Accept") request.setValue("CodeBurn", forHTTPHeaderField: "User-Agent") - // Kimi server expects these platform headers. - request.setValue("kimi_code_cli", forHTTPHeaderField: "X-Msh-Platform") - if let deviceId = deviceId() { - request.setValue(deviceId, forHTTPHeaderField: "X-Msh-Device-Id") + // These identify the CLI credential path. A Code Plan API key needs + // only bearer authentication, matching Kimi's API contract. + if !usesAPIKey { + request.setValue("kimi_code_cli", forHTTPHeaderField: "X-Msh-Platform") + if let deviceId = deviceId(deps: deps) { + request.setValue(deviceId, forHTTPHeaderField: "X-Msh-Device-Id") + } } let data: Data - let response: URLResponse + let http: HTTPURLResponse do { - (data, response) = try await URLSession.shared.data(for: request) + (data, http) = try await deps.fetch(request) + } catch let error as FetchError { + throw error } catch { throw FetchError.network(error) } - guard let http = response as? HTTPURLResponse else { - throw FetchError.usageHTTPError(-1, nil) - } switch http.statusCode { case 200: @@ -177,9 +207,7 @@ enum KimiSubscriptionService { throw FetchError.usageDecodeFailed } case 401, 403: - // We don't self-refresh; surface as terminal so the UI prompts - // the user to run the CLI. - throw FetchError.tokenExpired + throw usesAPIKey ? FetchError.apiKeyRejected : FetchError.tokenExpired case 429: let retryAfter = parseRetryAfterHeader(http.value(forHTTPHeaderField: "Retry-After")) let until = recordUsageRateLimit(retryAfterSeconds: retryAfter) diff --git a/mac/Sources/CodeBurnMenubar/Views/SettingsView.swift b/mac/Sources/CodeBurnMenubar/Views/SettingsView.swift index 0e15161cc..171bb8be5 100644 --- a/mac/Sources/CodeBurnMenubar/Views/SettingsView.swift +++ b/mac/Sources/CodeBurnMenubar/Views/SettingsView.swift @@ -1115,8 +1115,9 @@ private struct KimiSettingsTab: View { Section("Connection") { KimiConnectionRow() } + KimiAPIKeySection() Section { - Text("Kimi Code live-quota tracking reads `~/.kimi-code/credentials/kimi-code.json` directly. Nothing is copied or stored. Access tokens are short-lived (~15 minutes) and only the Kimi CLI refreshes them, so if the connection shows as expired, run the Kimi CLI once and click Reconnect.") + Text("CodeBurn uses your saved Kimi Code Plan API key when present, then falls back to the Kimi CLI credential in `~/.kimi-code/credentials/kimi-code.json`. The API key stays in your login Keychain. The CLI file is read-only and never copied.") .font(.system(size: 11)) .foregroundStyle(.secondary) } header: { @@ -1128,6 +1129,80 @@ private struct KimiSettingsTab: View { } } +private struct KimiAPIKeySection: View { + @Environment(AppStore.self) private var store + @State private var apiKey = "" + @State private var isSaving = false + @State private var errorText: String? + + var body: some View { + Section { + SecureField("Kimi Code API key", text: $apiKey) + HStack { + Button("Save & Connect", action: save) + .buttonStyle(.borderedProminent) + .disabled(isSaving || apiKey.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty) + Button("Clear Key", action: clear) + .disabled(isSaving || !CapacityDockProviderCredentialPresence.contains(KimiSubscriptionService.providerID)) + if isSaving { + ProgressView().controlSize(.small) + } + } + Link("Create a Kimi Code API key", destination: URL(string: "https://www.kimi.com/code/console")!) + if let errorText { + Text(errorText) + .font(.system(size: 11)) + .foregroundStyle(.red) + } + } header: { + Text("Kimi Code Plan API key") + } footer: { + Text("Optional. A saved key avoids the Kimi CLI's short-lived login token and is used only to read quota from api.kimi.com.") + .font(.system(size: 11)) + } + } + + private func save() { + let raw = apiKey + isSaving = true + errorText = nil + Task { + defer { isSaving = false } + do { + try await store.saveCapacityDockCredential( + CapacityDockProviderCredential(sourceMode: "api", apiKey: raw), + for: .kimiCode + ) + apiKey = "" + KimiSubscriptionService.disconnect() + await store.bootstrapKimi() + } catch { + errorText = error.localizedDescription + } + } + } + + private func clear() { + isSaving = true + errorText = nil + Task { + defer { isSaving = false } + do { + try await CapacityDockProviderCredentialStore.removeAsync( + for: KimiSubscriptionService.providerID + ) + apiKey = "" + store.disconnectKimi() + if KimiSubscriptionService.hasCredential { + await store.bootstrapKimi() + } + } catch { + errorText = error.localizedDescription + } + } + } +} + private struct KimiConnectionRow: View { @Environment(AppStore.self) private var store @State private var showDisconnectConfirm = false @@ -1190,13 +1265,13 @@ private struct KimiConnectionRow: View { case .loaded: return "Live quota tracked from api.kimi.com." case .terminalFailure: - return "Run the Kimi CLI once to refresh your login, then click Reconnect." + return store.kimiError ?? "Check your saved API key or refresh the Kimi CLI login, then click Reconnect." case .transientFailure: return store.kimiError ?? "Kimi rate-limited; auto-retrying." - case .bootstrapping: return "Reading ~/.kimi-code credentials." + case .bootstrapping: return "Checking the saved API key, then Kimi CLI credentials." case .loading: return "Background refresh in progress." case .dormant: return "Tap Load Quota to fetch live usage from api.kimi.com." case .notBootstrapped, .noCredentials: - return "Sign in with the Kimi CLI first, then click Connect." + return "Paste a Kimi Code Plan API key below, or sign in with the Kimi CLI, then click Connect." case .failed: return store.kimiError ?? "" } } @@ -1215,7 +1290,7 @@ private struct KimiConnectionRow: View { } Button("Cancel", role: .cancel) {} } message: { - Text("CodeBurn will stop tracking Kimi Code quota. Your ~/.kimi-code credentials are untouched. The Kimi CLI keeps working.") + Text("CodeBurn will stop tracking Kimi Code quota. Your saved API key and ~/.kimi-code credentials are untouched.") } case .terminalFailure, .noCredentials, .failed: Button("Reconnect") { Task { await store.bootstrapKimi() } } diff --git a/mac/Tests/CodeBurnMenubarTests/KimiUsageParsingTests.swift b/mac/Tests/CodeBurnMenubarTests/KimiUsageParsingTests.swift index 2d78c43f0..62366bdde 100644 --- a/mac/Tests/CodeBurnMenubarTests/KimiUsageParsingTests.swift +++ b/mac/Tests/CodeBurnMenubarTests/KimiUsageParsingTests.swift @@ -111,4 +111,34 @@ final class KimiUsageParsingTests: XCTestCase { """.data(using: .utf8)! XCTAssertThrowsError(try KimiSubscriptionService.parseUsage(data: json)) } + + func testCodePlanAPIKeyFetchesWithoutCLICredentials() async throws { + KimiSubscriptionService.disconnect() + defer { KimiSubscriptionService.disconnect() } + let json = """ + {"usage": {"limit": 100, "used": 25, "remaining": 75}} + """.data(using: .utf8)! + let deps = KimiSubscriptionService.Dependencies( + fetch: { request in + guard request.value(forHTTPHeaderField: "Authorization") == "Bearer synthetic-code-plan-key" else { + throw URLError(.userAuthenticationRequired) + } + return ( + json, + HTTPURLResponse(url: request.url!, statusCode: 200, httpVersion: nil, headerFields: nil)! + ) + }, + readFile: { _ in nil }, + credentialsURL: URL(fileURLWithPath: "/missing/kimi-code.json"), + deviceIDURL: URL(fileURLWithPath: "/missing/device-id"), + now: { Date(timeIntervalSince1970: 1_700_000_000) } + ) + + let usage = try await KimiSubscriptionService.refresh( + apiKey: " synthetic-code-plan-key ", + deps: deps + ) + + XCTAssertEqual(usage.primary?.usedPercent ?? -1, 25, accuracy: 0.001) + } } From 61ca91fa71b907d655ed87297645fdcc6ca2549c Mon Sep 17 00:00:00 2001 From: Miguel Machado Date: Mon, 31 Aug 2026 17:20:09 -0300 Subject: [PATCH 2/2] fix(mac): show Kimi short window first --- mac/Sources/CodeBurnMenubar/AppStore.swift | 3 +++ 1 file changed, 3 insertions(+) diff --git a/mac/Sources/CodeBurnMenubar/AppStore.swift b/mac/Sources/CodeBurnMenubar/AppStore.swift index dafdb2123..8b51d2b6b 100644 --- a/mac/Sources/CodeBurnMenubar/AppStore.swift +++ b/mac/Sources/CodeBurnMenubar/AppStore.swift @@ -2027,6 +2027,9 @@ final class AppStore { if primary == nil { primary = row } details.append(row) } + if usage.primary != nil, details.count > 1 { + details.append(details.removeFirst()) + } } return QuotaSummary(providerFilter: filter, connection: connection, primary: primary, details: details, planLabel: kimiUsage?.plan ?? "Kimi Code", footerLines: []) }