From ef4f50c17a357b0fb8cbccfd272e148e17b6a0b5 Mon Sep 17 00:00:00 2001 From: yishu-ziyu Date: Tue, 13 Oct 2026 17:58:51 +0800 Subject: [PATCH 1/3] Menubar: read the Cursor agent CLI session from the login keychain The quota adapter only consulted the Cursor IDE's state.vscdb, so users who only run the cursor-agent CLI (no IDE installed) always showed Not connected. The CLI stores the same WorkOS session JWT in the login keychain as cursor-access-token/cursor-user; read it there as a fallback, preferring /usr/bin/security because the agent-created item's ACL lets Apple's signed binary read it silently, while SecItemCopyMatching from the menu-bar app raises an interactive prompt. --- .../Data/CursorSubscriptionService.swift | 65 ++++++++++++++++++- 1 file changed, 64 insertions(+), 1 deletion(-) diff --git a/mac/Sources/CodeBurnMenubar/Data/CursorSubscriptionService.swift b/mac/Sources/CodeBurnMenubar/Data/CursorSubscriptionService.swift index 46bec8343..4bfcd8729 100644 --- a/mac/Sources/CodeBurnMenubar/Data/CursorSubscriptionService.swift +++ b/mac/Sources/CodeBurnMenubar/Data/CursorSubscriptionService.swift @@ -1,4 +1,5 @@ import Foundation +import Security import SQLite3 /// CodeBurn-owned Cursor quota adapter. It reads Cursor.app's existing local @@ -307,12 +308,70 @@ enum CursorSubscriptionService { /// Narrow read-only adapter for Cursor's own VS Code state database. struct CursorAppSessionStore: Sendable { private static let accessTokenKey = "cursorAuth/accessToken" + /// cursor-agent CLI keeps the same WorkOS session JWT in the login + /// keychain; no Cursor IDE install produces no state.vscdb at all. + private static let cliTokenService = "cursor-access-token" + private static let cliTokenAccount = "cursor-user" let databaseURL: URL init(databaseURL: URL = Self.defaultDatabaseURL()) { self.databaseURL = databaseURL } + /// Read-only lookup of the cursor-agent CLI's session token. Prefers + /// `/usr/bin/security`: agent-created items let Apple's signed `security` + /// binary read them silently, while SecItemCopyMatching from this app + /// raises an interactive prompt that a menu-bar agent cannot anchor. + static func loadCLIAccessToken() -> String? { + loadCLIAccessTokenViaSecurityTool() ?? loadCLIAccessTokenViaKeychain() + } + + private static func loadCLIAccessTokenViaSecurityTool() -> String? { + let process = Process() + let output = Pipe() + process.executableURL = URL(fileURLWithPath: "/usr/bin/security") + process.arguments = [ + "find-generic-password", "-s", cliTokenService, "-a", cliTokenAccount, "-w", + ] + process.standardOutput = output + process.standardError = FileHandle.nullDevice + do { + try process.run() + } catch { + return nil + } + // A missing or ACL-blocked item exits fast; anything else gets 5s + // before we kill it rather than block the refresh loop on a prompt. + let deadline = Date().addingTimeInterval(5) + while process.isRunning && Date() < deadline { + Thread.sleep(forTimeInterval: 0.02) + } + if process.isRunning { + process.terminate() + process.waitUntilExit() + return nil + } + let token = String(data: output.fileHandleForReading.readDataToEndOfFile(), encoding: .utf8)? + .trimmingCharacters(in: .whitespacesAndNewlines) + return (token?.isEmpty == false) ? token : nil + } + + private static func loadCLIAccessTokenViaKeychain() -> String? { + let query: [String: Any] = [ + kSecClass as String: kSecClassGenericPassword, + kSecAttrService as String: cliTokenService, + kSecAttrAccount as String: cliTokenAccount, + kSecReturnData as String: true, + kSecMatchLimit as String: kSecMatchLimitOne, + ] + var item: CFTypeRef? + guard SecItemCopyMatching(query as CFDictionary, &item) == errSecSuccess, + let data = item as? Data else { return nil } + let token = String(data: data, encoding: .utf8)? + .trimmingCharacters(in: .whitespacesAndNewlines) + return (token?.isEmpty == false) ? token : nil + } + static func defaultDatabaseURL(home: URL = FileManager.default.homeDirectoryForCurrentUser) -> URL { home .appendingPathComponent("Library/Application Support/Cursor/User/globalStorage", isDirectory: true) @@ -320,12 +379,16 @@ struct CursorAppSessionStore: Sendable { } func loadAccessToken() throws -> String? { - guard FileManager.default.fileExists(atPath: databaseURL.path) else { return nil } + guard FileManager.default.fileExists(atPath: databaseURL.path) else { + return Self.loadCLIAccessToken() + } do { return try value(for: Self.accessTokenKey, immutable: false) + ?? Self.loadCLIAccessToken() } catch let error as SQLiteFailure where error.code == SQLITE_CANTOPEN && walSidecarsAreMissing { return try value(for: Self.accessTokenKey, immutable: true) + ?? Self.loadCLIAccessToken() } catch { throw CursorAppSessionStoreError.unreadable } From 0af23dd373de73d73d28333058b7aa1a971eacf7 Mon Sep 17 00:00:00 2001 From: yishu-ziyu Date: Tue, 13 Oct 2026 17:58:51 +0800 Subject: [PATCH 2/3] Menubar: renew expired Kimi Code tokens through the CLI's OAuth grant MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Kimi access tokens live ~15 minutes and the adapter only read the file's access_token, so the dock showed a permanent login-expired row unless the user happened to run the Kimi CLI within the window. The credential file already carries a refresh_token; exchange it at the same endpoint the CLI uses (auth.kimi.com/api/oauth/token) and keep the renewed token in memory only — the CLI's credential file is never written to. --- .../Data/KimiSubscriptionService.swift | 100 ++++++++++++++++-- 1 file changed, 91 insertions(+), 9 deletions(-) diff --git a/mac/Sources/CodeBurnMenubar/Data/KimiSubscriptionService.swift b/mac/Sources/CodeBurnMenubar/Data/KimiSubscriptionService.swift index 4b6aa9c59..4cce864ca 100644 --- a/mac/Sources/CodeBurnMenubar/Data/KimiSubscriptionService.swift +++ b/mac/Sources/CodeBurnMenubar/Data/KimiSubscriptionService.swift @@ -33,11 +33,14 @@ struct KimiUsage: Sendable, Equatable { /// 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. +/// no keychain bootstrap. Access tokens are short-lived (~15 min), so an +/// expired one is renewed with the file's refresh_token against the same +/// endpoint the Kimi CLI uses. The rotated credentials are kept in memory +/// only; the CLI's file is never written back to. enum KimiSubscriptionService { private static let usageURL = URL(string: "https://api.kimi.com/coding/v1/usages")! + private static let tokenURL = URL(string: "https://auth.kimi.com/api/oauth/token")! + private static let oauthClientID = "17e5f671-d194-4dfb-9706-5516cb48c098" private static let usageBlockedUntilKey = "codeburn.kimi.usage.blockedUntil" enum FetchError: Error, LocalizedError { @@ -81,16 +84,19 @@ enum KimiSubscriptionService { private struct CredentialFile: Decodable { let accessToken: String + let refreshToken: String? let expiresAt: Double enum CodingKeys: String, CodingKey { case accessToken = "access_token" + case refreshToken = "refresh_token" case expiresAt = "expires_at" } init(from decoder: Decoder) throws { let c = try decoder.container(keyedBy: CodingKeys.self) accessToken = try c.decode(String.self, forKey: .accessToken) + refreshToken = try? c.decode(String.self, forKey: .refreshToken) // Tolerate Int / Double / String epoch-seconds. if let n = try? c.decode(Double.self, forKey: .expiresAt) { expiresAt = n @@ -112,18 +118,94 @@ enum KimiSubscriptionService { FileManager.default.fileExists(atPath: credentialsURL.path) } - /// Returns the access token only when it is still fresh (60s skew). - /// Throws noCredentials / tokenExpired otherwise. - private static func freshToken() throws -> String { + /// Fresh access token, renewing through the CLI's OAuth refresh grant + /// when the file's access token is stale. A missing refresh grant is + /// still terminal — the UI then asks the user to run the Kimi CLI once. + private static func accessToken() async throws -> String { guard let data = FileManager.default.contents(atPath: credentialsURL.path), let cred = try? JSONDecoder().decode(CredentialFile.self, from: data), !cred.accessToken.isEmpty else { throw FetchError.noCredentials } - guard cred.expiresAt > Date().timeIntervalSince1970 + 60 else { + if cred.expiresAt > Date().timeIntervalSince1970 + 60 { + return cred.accessToken + } + if let cached = freshCachedToken(), cached.expiresAt > Date() + 60 { + return cached.token + } + guard let refreshToken = cred.refreshToken, !refreshToken.isEmpty else { + throw FetchError.tokenExpired + } + return try await renewAccessToken(refreshToken) + } + + /// The CLI's rotated access token lives only for this process — the + /// credential file keeps the CLI's own copy untouched. Locking stays in + /// synchronous helpers: NSLock is unavailable from async contexts. + private static let refreshedTokenLock = NSLock() + private nonisolated(unsafe) static var cachedRefreshedToken: (token: String, expiresAt: Date)? + + private static func freshCachedToken() -> (token: String, expiresAt: Date)? { + refreshedTokenLock.lock() + defer { refreshedTokenLock.unlock() } + return cachedRefreshedToken + } + + private static func storeCachedToken(_ token: String, expiresAt: Date) { + refreshedTokenLock.lock() + cachedRefreshedToken = (token, expiresAt) + refreshedTokenLock.unlock() + } + + private static func renewAccessToken(_ refreshToken: String) async throws -> String { + var request = URLRequest(url: tokenURL) + request.httpMethod = "POST" + request.timeoutInterval = 15 + request.setValue("application/x-www-form-urlencoded", forHTTPHeaderField: "Content-Type") + request.setValue("CodeBurn", forHTTPHeaderField: "User-Agent") + var body = URLComponents() + body.queryItems = [ + URLQueryItem(name: "grant_type", value: "refresh_token"), + URLQueryItem(name: "refresh_token", value: refreshToken), + URLQueryItem(name: "client_id", value: oauthClientID), + ] + request.httpBody = body.percentEncodedQuery?.data(using: .utf8) + + let data: Data + let response: URLResponse + do { + (data, response) = try await URLSession.shared.data(for: request) + } catch { + throw FetchError.network(error) + } + guard let http = response as? HTTPURLResponse, http.statusCode == 200, + let tokens = try? JSONDecoder().decode(RenewedTokens.self, from: data), + !tokens.accessToken.isEmpty else { throw FetchError.tokenExpired } - return cred.accessToken + let ttl = max(60, tokens.expiresIn ?? 900) + storeCachedToken(tokens.accessToken, expiresAt: Date().addingTimeInterval(ttl)) + return tokens.accessToken + } + + private struct RenewedTokens: Decodable { + let accessToken: String + let expiresIn: Double? + enum CodingKeys: String, CodingKey { + case accessToken = "access_token" + case expiresIn = "expires_in" + } + init(from decoder: Decoder) throws { + let c = try decoder.container(keyedBy: CodingKeys.self) + accessToken = try c.decode(String.self, forKey: .accessToken) + if let n = try? c.decode(Double.self, forKey: .expiresIn) { + expiresIn = n + } else if let s = try? c.decode(String.self, forKey: .expiresIn) { + expiresIn = Double(s) + } else { + expiresIn = nil + } + } } private static func deviceId() -> String? { @@ -141,7 +223,7 @@ enum KimiSubscriptionService { if let until = usageBlockedUntil(), until > Date() { throw FetchError.rateLimited(retryAt: until) } - let token = try freshToken() + let token = try await accessToken() var request = URLRequest(url: usageURL) request.httpMethod = "GET" From 3422429d1f2b6ddfbeeac3c23e7655a22c7032bf Mon Sep 17 00:00:00 2001 From: yishu-ziyu Date: Tue, 13 Oct 2026 17:58:51 +0800 Subject: [PATCH 3/3] Menubar: fall back to agy -p /usage when no Antigravity server runs The probe only found the quota endpoints on a live language server, so users who only drive Antigravity through the agy CLI read Not connected between sessions. When no candidate answers, run the CLI's own /usage --output-format json path (it boots a throwaway server internally) and parse the same group/bucket shape; one spaced retry covers the throwaway server's 503 cold start when other pollers run concurrently. --- .../Data/AntigravitySubscriptionService.swift | 90 ++++++++++++++++++- 1 file changed, 89 insertions(+), 1 deletion(-) diff --git a/mac/Sources/CodeBurnMenubar/Data/AntigravitySubscriptionService.swift b/mac/Sources/CodeBurnMenubar/Data/AntigravitySubscriptionService.swift index 381e6c0ec..c46f1ed76 100644 --- a/mac/Sources/CodeBurnMenubar/Data/AntigravitySubscriptionService.swift +++ b/mac/Sources/CodeBurnMenubar/Data/AntigravitySubscriptionService.swift @@ -71,6 +71,12 @@ enum AntigravitySubscriptionService { /// POSTs the Connect-RPC probe to 127.0.0.1:. Nil means /// timeout / transport error — the probe simply moves on. var request: @Sendable (_ port: Int, _ tls: Bool, _ path: String, _ body: String, _ csrf: String?) async -> (status: Int, text: String)? + /// Runs `agy -p /usage --output-format json` for the CLI fallback. + /// Long timeout: the CLI boots a throwaway language server (~20-60s cold). + /// Defaulted so existing fixture-built Deps keep compiling. + var execAgy: @Sendable (_ executable: String) async throws -> String = { _ in + throw FetchError.disconnected + } var now: @Sendable () -> Date /// Log sink for sanitized diagnostics; tests capture it. var log: @Sendable (String) -> Void @@ -106,6 +112,17 @@ enum AntigravitySubscriptionService { return nil } }, + execAgy: { executable in + let process = Process() + process.executableURL = URL(fileURLWithPath: executable) + process.arguments = ["-p", "/usage", "--output-format", "json", "--print-timeout", "60s"] + let result = try await DataClient.runProcess( + process, + timeoutSeconds: 75, + label: "antigravity agy" + ) + return String(decoding: result.stdout, as: UTF8.self) + }, now: { Date() }, log: { NSLog("%@", $0) } ) @@ -250,6 +267,34 @@ enum AntigravitySubscriptionService { return windows } + /// `agy -p /usage --output-format json` payload: command.data.groups carry + /// buckets with snake_case fields (name / remaining_fraction / reset_time). + static func decodeAgyOutput(_ text: String) -> [AntigravityUsage.Window] { + guard let root = parseJson(text) as? [String: Any], + let command = root["command"] as? [String: Any], + let data = command["data"] as? [String: Any] else { return [] } + var windows: [AntigravityUsage.Window] = [] + for group in data["groups"] as? [Any] ?? [] { + guard let group = group as? [String: Any] else { continue } + let groupName = group["name"] as? String ?? "" + for bucket in group["buckets"] as? [Any] ?? [] { + guard let bucket = bucket as? [String: Any] else { continue } + let name = [groupName, bucket["name"] as? String ?? ""] + .filter { !$0.isEmpty } + .joined(separator: " · ") + if name.isEmpty { continue } + if let window = makeWindow( + label: name, + remainingFraction: bucket["remaining_fraction"], + resetTime: bucket["reset_time"]) + { + windows.append(window) + } + } + } + return windows + } + /// Legacy payload: flat per-model quota rows under GetUserStatus. static func decodeStatus(_ body: Any?) -> [AntigravityUsage.Window] { guard let data = body as? [String: Any], @@ -377,7 +422,6 @@ enum AntigravitySubscriptionService { static func refresh(deps: Deps = .live) async throws -> AntigravityUsage { do { let candidates = try await discoverCandidates(deps: deps) - if candidates.isEmpty { throw FetchError.disconnected } for candidate in candidates { var ports = await listeningPorts(deps: deps, pid: candidate.pid) if let extPort = candidate.extPort, !ports.contains(extPort) { @@ -390,6 +434,10 @@ enum AntigravitySubscriptionService { } } } + // No language server answered — the `agy` CLI prints the same + // quota groups itself (it boots a throwaway server internally), + // so a signed-in CLI keeps the dock live between app sessions. + if let usage = try await refreshViaAgyCLI(deps: deps) { return usage } throw FetchError.disconnected } catch let error as FetchError { throw error @@ -399,6 +447,46 @@ enum AntigravitySubscriptionService { } } + private static func agyBinaryPath() -> String? { + let home = FileManager.default.homeDirectoryForCurrentUser.path + for path in [ + "\(home)/.local/bin/agy", + "/opt/homebrew/bin/agy", + "/usr/local/bin/agy", + ] where FileManager.default.isExecutableFile(atPath: path) { + return path + } + return nil + } + + private static func refreshViaAgyCLI(deps: Deps) async throws -> AntigravityUsage? { + guard let agy = agyBinaryPath() else { return nil } + // Other pollers (e.g. Vibe Island) hit `agy -p /usage` too, and + // concurrent runs make the CLI's throwaway server answer 503 before it + // is ready — the error JSON carries no `command.data.groups`. One + // spaced retry covers that cold-start window. + for attempt in 1 ... 2 { + let output: String + do { + output = try await deps.execAgy(agy) + } catch { + deps.log("Antigravity agy quota fallback failed: \(sanitizeForLog(error.localizedDescription))") + return nil + } + let windows = decodeAgyOutput(output).sorted { $0.usedPercent > $1.usedPercent } + if !windows.isEmpty { + return AntigravityUsage(details: windows, plan: nil, fetchedAt: deps.now()) + } + if attempt == 1 { + deps.log("Antigravity agy returned no quota buckets (transient CLI error); retrying once") + try? await Task.sleep(nanoseconds: 8_000_000_000) + } else { + deps.log("Antigravity agy quota fallback produced no windows after retry: \(sanitizeForLog(String(output.prefix(160))))") + } + } + return nil + } + /// Strip control characters and any token-shaped substrings from error /// strings before they land in NSLog — a `ps` failure message can quote /// arbitrary command lines. Mirrors the Electron side's sanitizeError.