chore: import upstream snapshot with attribution

This commit is contained in:
wehub-resource-sync
2026-07-13 12:22:33 +08:00
commit fc4fcbab58
1848 changed files with 472303 additions and 0 deletions
@@ -0,0 +1,72 @@
import Foundation
public struct APITokenFetchStrategy: ProviderFetchStrategy {
public typealias TokenResolver = @Sendable ([String: String]) -> String?
public typealias MissingCredentialsError = @Sendable () -> (any Error & Sendable)
public typealias UsageLoader = @Sendable (String, ProviderFetchContext) async throws -> UsageSnapshot
public let id: String
public let kind: ProviderFetchKind = .apiToken
private let sourceLabel: String
private let reportsMissingCredentials: Bool
private let resolveToken: TokenResolver
private let missingCredentialsError: MissingCredentialsError
private let loadUsage: UsageLoader
public init(
id: String,
sourceLabel: String = "api",
reportsMissingCredentials: Bool = false,
resolveToken: @escaping TokenResolver,
missingCredentialsError: @escaping MissingCredentialsError,
loadUsage: @escaping UsageLoader)
{
self.id = id
self.sourceLabel = sourceLabel
self.reportsMissingCredentials = reportsMissingCredentials
self.resolveToken = resolveToken
self.missingCredentialsError = missingCredentialsError
self.loadUsage = loadUsage
}
public func isAvailable(_ context: ProviderFetchContext) async -> Bool {
self.reportsMissingCredentials || self.resolveToken(context.env) != nil
}
public func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult {
guard let token = self.resolveToken(context.env) else {
throw self.missingCredentialsError()
}
return try await self.makeResult(
usage: self.loadUsage(token, context),
sourceLabel: self.sourceLabel)
}
public func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool {
false
}
}
extension ProviderFetchPlan {
public static func apiToken(
strategyID: String,
sourceLabel: String = "api",
reportsMissingCredentials: Bool = false,
resolveToken: @escaping APITokenFetchStrategy.TokenResolver,
missingCredentialsError: @escaping APITokenFetchStrategy.MissingCredentialsError,
loadUsage: @escaping APITokenFetchStrategy.UsageLoader) -> ProviderFetchPlan
{
ProviderFetchPlan(
sourceModes: [.auto, .api],
pipeline: ProviderFetchPipeline(resolveStrategies: { _ in
[APITokenFetchStrategy(
id: strategyID,
sourceLabel: sourceLabel,
reportsMissingCredentials: reportsMissingCredentials,
resolveToken: resolveToken,
missingCredentialsError: missingCredentialsError,
loadUsage: loadUsage)]
}))
}
}
@@ -0,0 +1,134 @@
import Foundation
#if os(macOS)
import SweetCookieKit
// MARK: - Abacus Cookie Importer
public enum AbacusCookieImporter {
private static let log = CodexBarLog.logger(LogCategories.abacusCookie)
private static let cookieClient = BrowserCookieClient()
private static let cookieDomains = ["abacus.ai", "apps.abacus.ai"]
private static let cookieImportOrder: BrowserCookieImportOrder =
ProviderDefaults.metadata[.abacus]?.browserCookieOrder ?? Browser.defaultImportOrder
/// Exact cookie names known to carry Abacus session state.
/// CSRF tokens are deliberately excluded they are present in anonymous
/// jars and do not indicate an authenticated session.
private static let knownSessionCookieNames: Set<String> = [
"sessionid", "session_id", "session_token",
"auth_token", "access_token",
]
/// Substrings that indicate a session or auth cookie (applied only when
/// no exact-name match is found). Deliberately excludes overly broad
/// patterns like "id" and "token" that match analytics/CSRF cookies.
private static let sessionCookieSubstrings = ["session", "auth", "sid", "jwt"]
/// Cookie name prefixes that indicate a non-session cookie even when a
/// substring match would otherwise accept it (e.g. "csrftoken").
private static let excludedCookiePrefixes = ["csrf", "_ga", "_gid", "tracking", "analytics"]
public struct SessionInfo: Sendable {
public let cookies: [HTTPCookie]
public let sourceLabel: String
public init(cookies: [HTTPCookie], sourceLabel: String) {
self.cookies = cookies
self.sourceLabel = sourceLabel
}
public var cookieHeader: String {
self.cookies.map { "\($0.name)=\($0.value)" }.joined(separator: "; ")
}
}
/// Returns all candidate sessions across browsers/profiles, ordered by
/// import priority. Callers should try each in turn so that a stale
/// session in the first source doesn't block a valid one further down.
///
/// Defaults to Chrome-only per AGENTS.md guideline. Pass an empty
/// `preferredBrowsers` list to fall back to the full descriptor-defined
/// import order (Safari, Firefox, etc.) when Chrome has no cookies.
public static func importSessions(
browserDetection: BrowserDetection = BrowserDetection(),
preferredBrowsers: [Browser] = [.chrome],
logger: ((String) -> Void)? = nil) throws -> [SessionInfo]
{
var candidates: [SessionInfo] = []
let installedBrowsers = preferredBrowsers.isEmpty
? self.cookieImportOrder.cookieImportCandidates(using: browserDetection)
: preferredBrowsers.cookieImportCandidates(using: browserDetection)
for browserSource in installedBrowsers {
do {
let query = BrowserCookieQuery(domains: self.cookieDomains)
let sources = try Self.cookieClient.codexBarRecords(
matching: query,
in: browserSource,
logger: { msg in self.emit(msg, logger: logger) })
for source in sources where !source.records.isEmpty {
let httpCookies = BrowserCookieClient.makeHTTPCookies(source.records, origin: query.origin)
guard !httpCookies.isEmpty else { continue }
guard Self.containsSessionCookie(httpCookies) else {
self.emit(
"Skipping \(source.label): no session cookie found",
logger: logger)
continue
}
self.emit(
"Found \(httpCookies.count) session cookies in \(source.label)",
logger: logger)
candidates.append(SessionInfo(cookies: httpCookies, sourceLabel: source.label))
}
} catch {
BrowserCookieAccessGate.recordIfNeeded(error)
self.emit(
"\(browserSource.displayName) cookie import failed: \(error.localizedDescription)",
logger: logger)
}
}
guard !candidates.isEmpty else {
throw AbacusUsageError.noSessionCookie
}
return candidates
}
/// Cheap check for whether any browser has an Abacus session cookie,
/// used by the fetch strategy's `isAvailable()`.
public static func hasSession(
browserDetection: BrowserDetection = BrowserDetection(),
preferredBrowsers: [Browser] = [.chrome],
logger: ((String) -> Void)? = nil) -> Bool
{
do {
return try !self.importSessions(
browserDetection: browserDetection,
preferredBrowsers: preferredBrowsers,
logger: logger).isEmpty
} catch {
return false
}
}
/// Returns `true` if the cookie set contains at least one cookie whose name
/// indicates session or authentication state. Checks exact known names
/// first, then falls back to conservative substring matching.
private static func containsSessionCookie(_ cookies: [HTTPCookie]) -> Bool {
cookies.contains { cookie in
let lower = cookie.name.lowercased()
if self.knownSessionCookieNames.contains(lower) { return true }
if self.excludedCookiePrefixes.contains(where: { lower.hasPrefix($0) }) { return false }
return self.sessionCookieSubstrings.contains { lower.contains($0) }
}
}
private static func emit(_ message: String, logger: ((String) -> Void)?) {
logger?("[abacus-cookie] \(message)")
self.log.debug(message)
}
}
#endif
@@ -0,0 +1,91 @@
import Foundation
#if os(macOS)
import SweetCookieKit
#endif
public enum AbacusProviderDescriptor {
public static let descriptor: ProviderDescriptor = Self.makeDescriptor()
static func makeDescriptor() -> ProviderDescriptor {
ProviderDescriptor(
id: .abacus,
metadata: ProviderMetadata(
id: .abacus,
displayName: "Abacus AI",
sessionLabel: "Credits",
weeklyLabel: "Weekly",
opusLabel: nil,
supportsOpus: false,
supportsCredits: false,
creditsHint: "",
toggleTitle: "Show Abacus AI usage",
cliName: "abacusai",
defaultEnabled: false,
isPrimaryProvider: false,
usesAccountFallback: false,
browserCookieOrder: ProviderBrowserCookieDefaults.defaultImportOrder,
dashboardURL: "https://apps.abacus.ai/chatllm/admin/compute-points-usage",
statusPageURL: nil,
statusLinkURL: nil),
branding: ProviderBranding(
iconStyle: .abacus,
iconResourceName: "ProviderIcon-abacus",
color: ProviderColor(red: 56 / 255, green: 189 / 255, blue: 248 / 255)),
tokenCost: ProviderTokenCostConfig(
supportsTokenCost: false,
noDataMessage: { "Abacus AI cost summary is not supported." }),
fetchPlan: ProviderFetchPlan(
sourceModes: [.auto, .web],
pipeline: ProviderFetchPipeline(resolveStrategies: { _ in
[AbacusWebFetchStrategy()]
})),
cli: ProviderCLIConfig(
name: "abacusai",
aliases: ["abacus-ai"],
versionDetector: nil))
}
}
// MARK: - Fetch Strategy
struct AbacusWebFetchStrategy: ProviderFetchStrategy {
let id: String = "abacus.web"
let kind: ProviderFetchKind = .web
func isAvailable(_ context: ProviderFetchContext) async -> Bool {
context.settings?.abacus?.cookieSource != .off
}
func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult {
let manual: String?
if context.settings?.abacus?.cookieSource == .manual {
guard let header = Self.manualCookieHeader(from: context) else {
throw AbacusUsageError.noSessionCookie
}
manual = header
} else {
manual = nil
}
let logger: ((String) -> Void)? = context.verbose
? { msg in CodexBarLog.logger(LogCategories.abacusUsage).verbose(msg) }
: nil
let snap = try await AbacusUsageFetcher.fetchUsage(
cookieHeaderOverride: manual,
browserDetection: context.browserDetection,
timeout: context.webTimeout,
logger: logger)
return self.makeResult(
usage: snap.toUsageSnapshot(),
sourceLabel: "web")
}
func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool {
false
}
private static func manualCookieHeader(from context: ProviderFetchContext) -> String? {
guard context.settings?.abacus?.cookieSource == .manual else { return nil }
return CookieHeaderNormalizer.normalize(context.settings?.abacus?.manualCookieHeader)
}
}
@@ -0,0 +1,67 @@
import Foundation
// MARK: - Abacus Usage Error
public enum AbacusUsageError: LocalizedError, Sendable, Equatable {
case noSessionCookie
case sessionExpired
case networkError(String)
case parseFailed(String)
case unauthorized
/// Whether this error indicates an authentication/session problem that
/// should trigger cache eviction.
public var isRecoverable: Bool {
switch self {
case .unauthorized, .sessionExpired: true
default: false
}
}
public var isAuthRelated: Bool {
self.isRecoverable
}
/// Whether browser-import scanning should continue to later sessions after
/// this failure. Imported sessions can differ by profile/browser, so we keep
/// scanning on per-session fetch failures and surface the first one only if
/// every candidate is exhausted.
var shouldTryNextImportedSession: Bool {
switch self {
case .unauthorized, .sessionExpired, .networkError, .parseFailed: true
case .noSessionCookie: false
}
}
/// Whether a cached cookie header should be evicted before falling back to
/// a fresh browser import. Parse/auth failures usually indicate that the
/// cached session is stale or no longer accepted.
var shouldClearCachedCookie: Bool {
switch self {
case .unauthorized, .sessionExpired, .parseFailed: true
case .networkError, .noSessionCookie: false
}
}
public var errorDescription: String? {
switch self {
case .noSessionCookie:
"No Abacus AI session found. Please log in to apps.abacus.ai in your browser "
+ "or paste a Cookie header in manual mode."
case .sessionExpired:
"Abacus AI session expired. Please log in again."
case let .networkError(msg):
"Abacus AI API error: \(msg)"
case let .parseFailed(msg):
"Could not parse Abacus AI usage: \(msg)"
case .unauthorized:
"Unauthorized. Please log in to Abacus AI."
}
}
}
#if !os(macOS)
extension AbacusUsageError {
public static let notSupported = AbacusUsageError.networkError("Abacus AI is only supported on macOS.")
}
#endif
@@ -0,0 +1,336 @@
import Foundation
#if canImport(FoundationNetworking)
import FoundationNetworking
#endif
#if os(macOS)
import SweetCookieKit
// MARK: - Abacus Usage Fetcher
public enum AbacusUsageFetcher {
private struct BrowserFetchRequest {
let browserDetection: BrowserDetection
let preferredBrowsers: [Browser]
let label: String
let timeout: TimeInterval
let logger: ((String) -> Void)?
}
/// Parsed JSON dictionaries are treated as immutable snapshots here and are
/// only moved between sibling fetch tasks before being consumed locally.
private struct JSONDictionaryBox: @unchecked Sendable {
let value: [String: Any]
}
private static let log = CodexBarLog.logger(LogCategories.abacusUsage)
private static let computePointsURL =
URL(string: "https://apps.abacus.ai/api/_getOrganizationComputePoints")!
private static let billingInfoURL =
URL(string: "https://apps.abacus.ai/api/_getBillingInfo")!
public static func fetchUsage(
cookieHeaderOverride: String? = nil,
browserDetection: BrowserDetection = BrowserDetection(),
timeout: TimeInterval = 15.0,
logger: ((String) -> Void)? = nil) async throws -> AbacusUsageSnapshot
{
// Manual cookie header no fallback, errors propagate directly
if let override = CookieHeaderNormalizer.normalize(cookieHeaderOverride) {
self.emit("Using manual cookie header", logger: logger)
return try await self.fetchWithCookieHeader(override, timeout: timeout, logger: logger)
}
// Cached cookie header fall back to a fresh browser import when the
// cached session is rejected or looks stale.
if let cached = CookieHeaderCache.load(provider: .abacus),
!cached.cookieHeader.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty
{
self.emit("Using cached cookie header from \(cached.sourceLabel)", logger: logger)
do {
return try await self.fetchWithCookieHeader(
cached.cookieHeader, timeout: timeout, logger: logger)
} catch let error as AbacusUsageError where error.shouldTryNextImportedSession {
if error.shouldClearCachedCookie {
CookieHeaderCache.clear(provider: .abacus)
self.emit(
"Cached cookie failed (\(error.localizedDescription)); cleared, trying fresh import",
logger: logger)
} else {
self.emit(
"Cached cookie failed (\(error.localizedDescription)); trying fresh import",
logger: logger)
}
}
}
// Fresh browser import try Chrome first (AGENTS.md default), then broaden
// to all browsers if Chrome has no sessions OR if every imported Chrome
// session is exhausted without a successful fetch.
var lastError: AbacusUsageError = .noSessionCookie
if let snapshot = try await self.tryFetchFromBrowsers(
BrowserFetchRequest(
browserDetection: browserDetection,
preferredBrowsers: [.chrome],
label: "Chrome",
timeout: timeout,
logger: logger),
lastError: &lastError)
{
return snapshot
}
self.emit("Chrome sessions exhausted; falling back to all browsers", logger: logger)
if let snapshot = try await self.tryFetchFromBrowsers(
BrowserFetchRequest(
browserDetection: browserDetection,
preferredBrowsers: [],
label: "all browsers",
timeout: timeout,
logger: logger),
lastError: &lastError)
{
return snapshot
}
throw lastError
}
/// Tries to import sessions from `preferredBrowsers` and fetch usage. Returns
/// the snapshot on success, nil if no sessions were available or every
/// imported session was exhausted without success.
private static func tryFetchFromBrowsers(
_ request: BrowserFetchRequest,
lastError: inout AbacusUsageError) async throws -> AbacusUsageSnapshot?
{
let sessions: [AbacusCookieImporter.SessionInfo]
do {
sessions = try AbacusCookieImporter.importSessions(
browserDetection: request.browserDetection,
preferredBrowsers: request.preferredBrowsers,
logger: request.logger)
} catch {
BrowserCookieAccessGate.recordIfNeeded(error)
self.emit(
"\(request.label) cookie import failed: \(error.localizedDescription)",
logger: request.logger)
return nil
}
for session in sessions {
self.emit("Trying cookies from \(session.sourceLabel)", logger: request.logger)
do {
let snapshot = try await self.fetchWithCookieHeader(
session.cookieHeader,
timeout: request.timeout,
logger: request.logger)
CookieHeaderCache.store(
provider: .abacus,
cookieHeader: session.cookieHeader,
sourceLabel: session.sourceLabel)
return snapshot
} catch let error as AbacusUsageError where error.shouldTryNextImportedSession {
self.emit(
"\(session.sourceLabel): \(error.localizedDescription), trying next source",
logger: request.logger)
lastError = error
continue
} catch {
self.emit(
"\(session.sourceLabel): \(error.localizedDescription), trying next source",
logger: request.logger)
lastError = .networkError(error.localizedDescription)
continue
}
}
return nil
}
// MARK: - API Requests
private static func fetchWithCookieHeader(
_ cookieHeader: String,
timeout: TimeInterval,
logger: ((String) -> Void)? = nil) async throws -> AbacusUsageSnapshot
{
enum FetchPart: Sendable {
case computePoints(JSONDictionaryBox)
case billingInfoSuccess(JSONDictionaryBox)
case billingInfoFailure(String)
}
// Fetch compute points (required, full timeout) and billing info
// (optional, shorter budget) concurrently. Billing is bounded so a
// slow/flaky billing endpoint can't delay credit rendering.
let billingBudget = min(timeout, 5.0)
var computePointsResult: [String: Any]?
var billingInfoResult: [String: Any] = [:]
try await withThrowingTaskGroup(of: FetchPart.self) { group in
group.addTask {
let result = try await self.fetchJSON(
url: self.computePointsURL,
method: "GET",
cookieHeader: cookieHeader,
timeout: timeout)
return .computePoints(JSONDictionaryBox(value: result))
}
group.addTask {
do {
let result = try await self.fetchJSON(
url: self.billingInfoURL,
method: "POST",
cookieHeader: cookieHeader,
timeout: billingBudget)
return .billingInfoSuccess(JSONDictionaryBox(value: result))
} catch {
return .billingInfoFailure(error.localizedDescription)
}
}
while let result = try await group.next() {
switch result {
case let .computePoints(value):
computePointsResult = value.value
case let .billingInfoSuccess(value):
billingInfoResult = value.value
case let .billingInfoFailure(message):
self.emit(
"Billing info fetch failed: \(message); credits shown without plan/reset",
logger: logger)
}
}
}
guard let computePointsResult else {
throw AbacusUsageError.networkError("Abacus compute points fetch did not complete")
}
return try self.parseResults(computePoints: computePointsResult, billingInfo: billingInfoResult)
}
private static func fetchJSON(
url: URL, method: String, cookieHeader: String, timeout: TimeInterval) async throws -> [String: Any]
{
var request = URLRequest(url: url)
request.httpMethod = method
request.timeoutInterval = timeout
request.setValue("application/json", forHTTPHeaderField: "Content-Type")
request.setValue("application/json", forHTTPHeaderField: "Accept")
request.setValue(cookieHeader, forHTTPHeaderField: "Cookie")
if method == "POST" {
request.httpBody = Data("{}".utf8)
}
let response = try await ProviderHTTPClient.shared.response(for: request)
let data = response.data
let statusCode = response.statusCode
if statusCode == 401 || statusCode == 403 {
throw AbacusUsageError.unauthorized
}
guard statusCode == 200 else {
let body = String(data: data.prefix(200), encoding: .utf8) ?? ""
throw AbacusUsageError.networkError("HTTP \(statusCode): \(body)")
}
let parsed: Any
do {
parsed = try JSONSerialization.jsonObject(with: data)
} catch {
let preview = String(data: data.prefix(200), encoding: .utf8) ?? "<non-UTF8>"
throw AbacusUsageError.parseFailed(
"\(url.lastPathComponent): \(error.localizedDescription) — preview: \(preview)")
}
guard let root = parsed as? [String: Any] else {
throw AbacusUsageError.parseFailed("\(url.lastPathComponent): top-level JSON is not a dictionary")
}
guard root["success"] as? Bool == true,
let result = root["result"] as? [String: Any]
else {
let errorMsg = (root["error"] as? String ?? "Unknown error").lowercased()
if errorMsg.contains("expired") || errorMsg.contains("session")
|| errorMsg.contains("login") || errorMsg.contains("authenticate")
|| errorMsg.contains("unauthorized") || errorMsg.contains("unauthenticated")
|| errorMsg.contains("forbidden")
{
throw AbacusUsageError.unauthorized
}
throw AbacusUsageError.parseFailed("\(url.lastPathComponent): \(errorMsg)")
}
return result
}
// MARK: - Parsing
private static func parseResults(
computePoints: [String: Any], billingInfo: [String: Any]) throws -> AbacusUsageSnapshot
{
let totalCredits = self.double(from: computePoints["totalComputePoints"])
let creditsLeft = self.double(from: computePoints["computePointsLeft"])
guard let totalCredits, let creditsLeft else {
let keys = computePoints.keys.sorted().joined(separator: ", ")
throw AbacusUsageError.parseFailed(
"Missing credit fields in compute points response. Keys: [\(keys)]")
}
let creditsUsed = totalCredits - creditsLeft
let nextBillingDate = billingInfo["nextBillingDate"] as? String
let currentTier = billingInfo["currentTier"] as? String
let resetsAt = self.parseDate(nextBillingDate)
return AbacusUsageSnapshot(
creditsUsed: creditsUsed,
creditsTotal: totalCredits,
resetsAt: resetsAt,
planName: currentTier)
}
private static func double(from value: Any?) -> Double? {
if let d = value as? Double { return d }
if let i = value as? Int { return Double(i) }
if let n = value as? NSNumber { return n.doubleValue }
return nil
}
private static func parseDate(_ isoString: String?) -> Date? {
guard let isoString else { return nil }
let formatter = ISO8601DateFormatter()
formatter.formatOptions = [.withInternetDateTime, .withFractionalSeconds]
if let date = formatter.date(from: isoString) { return date }
formatter.formatOptions = [.withInternetDateTime]
return formatter.date(from: isoString)
}
// MARK: - Logging
private static func emit(_ message: String, logger: ((String) -> Void)?) {
logger?("[abacus] \(message)")
self.log.debug(message)
}
}
#else
// MARK: - Abacus (Unsupported)
public enum AbacusUsageFetcher {
public static func fetchUsage(
cookieHeaderOverride _: String? = nil,
browserDetection _: BrowserDetection = BrowserDetection(),
timeout _: TimeInterval = 15.0,
logger _: ((String) -> Void)? = nil) async throws -> AbacusUsageSnapshot
{
throw AbacusUsageError.notSupported
}
}
#endif
@@ -0,0 +1,77 @@
import Foundation
// MARK: - Abacus Usage Snapshot
public struct AbacusUsageSnapshot: Sendable {
public let creditsUsed: Double?
public let creditsTotal: Double?
public let resetsAt: Date?
public let planName: String?
public init(
creditsUsed: Double? = nil,
creditsTotal: Double? = nil,
resetsAt: Date? = nil,
planName: String? = nil)
{
self.creditsUsed = creditsUsed
self.creditsTotal = creditsTotal
self.resetsAt = resetsAt
self.planName = planName
}
public func toUsageSnapshot() -> UsageSnapshot {
let percentUsed: Double = if let used = self.creditsUsed, let total = self.creditsTotal, total > 0 {
(used / total) * 100.0
} else {
0
}
let resetDesc: String? = if let used = self.creditsUsed, let total = self.creditsTotal {
"\(Self.formatCredits(used)) / \(Self.formatCredits(total)) credits"
} else {
nil
}
// Derive window from actual billing cycle when possible.
// Assume the cycle started one calendar month before resetsAt.
let windowMinutes: Int = if let resetDate = self.resetsAt,
let cycleStart = Calendar.current.date(byAdding: .month, value: -1, to: resetDate)
{
max(1, Int(resetDate.timeIntervalSince(cycleStart) / 60))
} else {
30 * 24 * 60
}
let primary = RateWindow(
usedPercent: percentUsed,
windowMinutes: windowMinutes,
resetsAt: self.resetsAt,
resetDescription: resetDesc)
let identity = ProviderIdentitySnapshot(
providerID: .abacus,
accountEmail: nil,
accountOrganization: nil,
loginMethod: self.planName)
return UsageSnapshot(
primary: primary,
secondary: nil,
tertiary: nil,
providerCost: nil,
updatedAt: Date(),
identity: identity)
}
// MARK: - Formatting
/// Thread-safe credit formatting allocates per call to avoid shared mutable state.
private static func formatCredits(_ value: Double) -> String {
let formatter = NumberFormatter()
formatter.numberStyle = .decimal
formatter.locale = Locale(identifier: "en_US")
formatter.maximumFractionDigits = value >= 1000 ? 0 : 1
return formatter.string(from: NSNumber(value: value)) ?? String(format: "%.0f", value)
}
}
@@ -0,0 +1,133 @@
import Foundation
public enum AlibabaCodingPlanAPIRegion: String, CaseIterable, Sendable {
case international = "intl"
case chinaMainland = "cn"
private static let endpointPath = "data/api.json"
public var displayName: String {
switch self {
case .international:
"International (modelstudio.console.alibabacloud.com)"
case .chinaMainland:
"China mainland (bailian.console.aliyun.com)"
}
}
public var gatewayBaseURLString: String {
switch self {
case .international:
"https://modelstudio.console.alibabacloud.com"
case .chinaMainland:
"https://bailian.console.aliyun.com"
}
}
public var dashboardURL: URL {
switch self {
case .international:
URL(
string: "https://modelstudio.console.alibabacloud.com/ap-southeast-1/?tab=coding-plan#/efm/coding_plan")!
case .chinaMainland:
URL(string: "https://bailian.console.aliyun.com/cn-beijing/?tab=model#/efm/coding_plan")!
}
}
public var consoleDomain: String {
switch self {
case .international:
"modelstudio.console.alibabacloud.com"
case .chinaMainland:
"bailian.console.aliyun.com"
}
}
public var consoleSite: String {
switch self {
case .international:
"MODELSTUDIO_ALIBABACLOUD"
case .chinaMainland:
"BAILIAN_ALIYUN"
}
}
public var consoleRefererURL: URL {
switch self {
case .international:
URL(string: "https://modelstudio.console.alibabacloud.com/ap-southeast-1/?tab=coding-plan")!
case .chinaMainland:
URL(string: "https://bailian.console.aliyun.com/cn-beijing/?tab=model")!
}
}
public var quotaURL: URL {
var components = URLComponents(string: self.gatewayBaseURLString)!
components.path = "/" + Self.endpointPath
components.queryItems = [
URLQueryItem(
name: "action",
value: "zeldaEasy.broadscope-bailian.codingPlan.queryCodingPlanInstanceInfoV2"),
URLQueryItem(name: "product", value: "broadscope-bailian"),
URLQueryItem(name: "api", value: "queryCodingPlanInstanceInfoV2"),
URLQueryItem(name: "currentRegionId", value: self.currentRegionID),
]
return components.url!
}
public var consoleRPCBaseURLString: String {
switch self {
case .international:
"https://bailian-singapore-cs.alibabacloud.com"
case .chinaMainland:
"https://bailian-cs.console.aliyun.com"
}
}
public var consoleRPCURL: URL {
var components = URLComponents(string: self.consoleRPCBaseURLString)!
components.path = "/" + Self.endpointPath
components.queryItems = [
URLQueryItem(name: "action", value: self.consoleRPCAction),
URLQueryItem(name: "product", value: self.consoleRPCProduct),
URLQueryItem(name: "api", value: self.consoleQuotaAPIName),
URLQueryItem(name: "_v", value: "undefined"),
]
return components.url!
}
public var consoleRPCAction: String {
switch self {
case .international:
"IntlBroadScopeAspnGateway"
case .chinaMainland:
"BroadScopeAspnGateway"
}
}
public var consoleRPCProduct: String {
"sfm_bailian"
}
public var consoleQuotaAPIName: String {
"zeldaEasy.broadscope-bailian.codingPlan.queryCodingPlanInstanceInfoV2"
}
public var commodityCode: String {
switch self {
case .international:
"sfm_codingplan_public_intl"
case .chinaMainland:
"sfm_codingplan_public_cn"
}
}
public var currentRegionID: String {
switch self {
case .international:
"ap-southeast-1"
case .chinaMainland:
"cn-beijing"
}
}
}
@@ -0,0 +1,510 @@
import Foundation
#if os(macOS)
import CommonCrypto
import Security
import SQLite3
import SweetCookieKit
private let alibabaCookieImportOrder: BrowserCookieImportOrder =
ProviderDefaults.metadata[.alibaba]?.browserCookieOrder ?? Browser.defaultImportOrder
public enum AlibabaCodingPlanCookieImporter {
private static let cookieClient = BrowserCookieClient()
private static let cookieDomains = [
"bailian-singapore-cs.alibabacloud.com",
"bailian-cs.console.aliyun.com",
"bailian-beijing-cs.aliyuncs.com",
"modelstudio.console.alibabacloud.com",
"bailian.console.aliyun.com",
"free.aliyun.com",
"account.aliyun.com",
"signin.aliyun.com",
"passport.alibabacloud.com",
"console.alibabacloud.com",
"console.aliyun.com",
"alibabacloud.com",
"aliyun.com",
]
public struct SessionInfo: Sendable {
public let cookies: [HTTPCookie]
public let sourceLabel: String
public init(cookies: [HTTPCookie], sourceLabel: String) {
self.cookies = cookies
self.sourceLabel = sourceLabel
}
public var cookieHeader: String {
var byName: [String: HTTPCookie] = [:]
byName.reserveCapacity(self.cookies.count)
for cookie in self.cookies {
if let expiry = cookie.expiresDate, expiry < Date() {
continue
}
guard !cookie.value.isEmpty else { continue }
if let existing = byName[cookie.name] {
let existingExpiry = existing.expiresDate ?? .distantPast
let candidateExpiry = cookie.expiresDate ?? .distantPast
if candidateExpiry >= existingExpiry {
byName[cookie.name] = cookie
}
} else {
byName[cookie.name] = cookie
}
}
return byName.keys.sorted().compactMap { name in
guard let cookie = byName[name] else { return nil }
return "\(cookie.name)=\(cookie.value)"
}.joined(separator: "; ")
}
}
nonisolated(unsafe) static var importSessionOverrideForTesting:
((BrowserDetection, ((String) -> Void)?) throws -> SessionInfo)?
public static func importSession(
browserDetection: BrowserDetection,
logger: ((String) -> Void)? = nil) throws -> SessionInfo
{
if let override = self.importSessionOverrideForTesting {
return try override(browserDetection, logger)
}
let log: (String) -> Void = { msg in logger?("[alibaba-cookie] \(msg)") }
var accessDeniedHints: [String] = []
var failureDetails: [String] = []
let installedBrowsers = self.cookieImportCandidates(browserDetection: browserDetection)
log("Cookie import candidates: \(installedBrowsers.map(\.displayName).joined(separator: ", "))")
for browserSource in installedBrowsers {
do {
log("Checking \(browserSource.displayName)")
let query = BrowserCookieQuery(domains: self.cookieDomains)
let sources = try Self.cookieClient.codexBarRecords(
matching: query,
in: browserSource,
logger: log)
if sources.isEmpty {
log("No matching cookie records in \(browserSource.displayName)")
if let fallbackSession = try Self.importChromiumFallbackSession(
browser: browserSource,
logger: log)
{
return fallbackSession
}
}
for source in sources where !source.records.isEmpty {
let httpCookies = BrowserCookieClient.makeHTTPCookies(source.records, origin: query.origin)
if self.isAuthenticatedSession(cookies: httpCookies) {
log("Found \(httpCookies.count) Alibaba cookies in \(source.label)")
return SessionInfo(cookies: httpCookies, sourceLabel: source.label)
}
let cookieNames = Set(httpCookies.map(\.name))
let hasTicket = cookieNames.contains("login_aliyunid_ticket")
let hasAccount =
cookieNames.contains("login_aliyunid_pk") ||
cookieNames.contains("login_current_pk") ||
cookieNames.contains("login_aliyunid")
log("Skipping \(source.label): missing auth cookies (ticket=\(hasTicket), account=\(hasAccount))")
}
if let fallbackSession = try Self.importChromiumFallbackSession(browser: browserSource, logger: log) {
return fallbackSession
}
} catch let error as BrowserCookieError {
BrowserCookieAccessGate.recordIfNeeded(error)
if let hint = error.accessDeniedHint {
accessDeniedHints.append(hint)
}
failureDetails.append("\(browserSource.displayName): \(error.localizedDescription)")
log("\(browserSource.displayName) cookie import failed: \(error.localizedDescription)")
} catch {
failureDetails.append("\(browserSource.displayName): \(error.localizedDescription)")
log("\(browserSource.displayName) cookie import failed: \(error.localizedDescription)")
}
}
let details = (Array(Set(accessDeniedHints)).sorted() + Array(Set(failureDetails)).sorted())
.joined(separator: " ")
throw AlibabaCodingPlanSettingsError.missingCookie(details: details.isEmpty ? nil : details)
}
private static func isAuthenticatedSession(cookies: [HTTPCookie]) -> Bool {
guard !cookies.isEmpty else { return false }
let names = Set(cookies.map(\.name))
let hasTicket = names.contains("login_aliyunid_ticket")
let hasAccount =
names.contains("login_aliyunid_pk") ||
names.contains("login_current_pk") ||
names.contains("login_aliyunid")
return hasTicket && hasAccount
}
public static func hasSession(
browserDetection: BrowserDetection,
logger: ((String) -> Void)? = nil) -> Bool
{
do {
_ = try self.importSession(browserDetection: browserDetection, logger: logger)
return true
} catch {
return false
}
}
private static func importChromiumFallbackSession(
browser: Browser,
logger: ((String) -> Void)? = nil) throws -> SessionInfo?
{
guard browser.usesChromiumProfileStore else { return nil }
return try AlibabaChromiumCookieFallbackImporter.importSession(
browser: browser,
domains: self.cookieDomains,
logger: logger)
}
static func cookieImportCandidates(
browserDetection: BrowserDetection,
importOrder: BrowserCookieImportOrder = alibabaCookieImportOrder) -> [Browser]
{
importOrder.cookieImportCandidates(using: browserDetection)
}
static func matchesCookieDomain(_ domain: String, patterns: [String] = Self.cookieDomains) -> Bool {
let normalized = self.normalizeCookieDomain(domain)
return patterns.contains { pattern in
let normalizedPattern = self.normalizeCookieDomain(pattern)
return normalized == normalizedPattern || normalized.hasSuffix(".\(normalizedPattern)")
}
}
static func normalizeCookieDomain(_ domain: String) -> String {
let trimmed = domain.trimmingCharacters(in: .whitespacesAndNewlines)
let normalized = trimmed.hasPrefix(".") ? String(trimmed.dropFirst()) : trimmed
return normalized.lowercased()
}
}
enum AlibabaChromiumCookieFallbackImporter {
private struct ChromiumCookieRecord {
let domain: String
let name: String
let path: String
let value: String
let expires: Date?
let isSecure: Bool
}
enum ImportError: LocalizedError {
case keyUnavailable(browser: Browser)
case keychainDenied(browser: Browser)
case sqliteFailed(label: String, details: String)
var errorDescription: String? {
switch self {
case let .keyUnavailable(browser):
"\(browser.displayName) Safe Storage key not found."
case let .keychainDenied(browser):
"macOS Keychain denied access to \(browser.displayName) Safe Storage."
case let .sqliteFailed(label, details):
"\(label) cookie fallback failed: \(details)"
}
}
}
static func importSession(
browser: Browser,
domains: [String],
cookieClient: BrowserCookieClient = BrowserCookieClient(),
logger: ((String) -> Void)? = nil) throws -> AlibabaCodingPlanCookieImporter.SessionInfo?
{
let stores = try cookieClient.codexBarStores(for: browser).filter { $0.databaseURL != nil }
guard !stores.isEmpty else { return nil }
logger?("[alibaba-cookie] Trying \(browser.displayName) Chromium fallback")
let keys = try self.derivedKeys(for: browser)
for store in stores {
let cookies = try self.loadCookies(from: store, domains: domains, keys: keys)
guard !cookies.isEmpty else { continue }
if self.isAuthenticatedSession(cookies) {
logger?("[alibaba-cookie] Found \(cookies.count) Alibaba cookies via \(store.label) fallback")
return AlibabaCodingPlanCookieImporter.SessionInfo(cookies: cookies, sourceLabel: store.label)
}
}
return nil
}
private static func isAuthenticatedSession(_ cookies: [HTTPCookie]) -> Bool {
let names = Set(cookies.map(\.name))
let hasTicket = names.contains("login_aliyunid_ticket")
let hasAccount =
names.contains("login_aliyunid_pk") ||
names.contains("login_current_pk") ||
names.contains("login_aliyunid")
return hasTicket && hasAccount
}
private static func loadCookies(
from store: BrowserCookieStore,
domains: [String],
keys: [Data]) throws -> [HTTPCookie]
{
guard let sourceDB = store.databaseURL else { return [] }
let records = try self.readCookiesFromLockedDB(
sourceDB: sourceDB,
domains: domains,
keys: keys,
label: store.label)
return records.compactMap(self.makeCookie)
}
private static func readCookiesFromLockedDB(
sourceDB: URL,
domains: [String],
keys: [Data],
label: String) throws -> [ChromiumCookieRecord]
{
let tempDir = FileManager.default.temporaryDirectory
.appendingPathComponent("alibaba-chromium-cookies-\(UUID().uuidString)", isDirectory: true)
try FileManager.default.createDirectory(at: tempDir, withIntermediateDirectories: true)
let copiedDB = tempDir.appendingPathComponent("Cookies")
try FileManager.default.copyItem(at: sourceDB, to: copiedDB)
for suffix in ["-wal", "-shm"] {
let src = URL(fileURLWithPath: sourceDB.path + suffix)
if FileManager.default.fileExists(atPath: src.path) {
let dst = URL(fileURLWithPath: copiedDB.path + suffix)
try? FileManager.default.copyItem(at: src, to: dst)
}
}
defer { try? FileManager.default.removeItem(at: tempDir) }
return try self.readCookies(fromDB: copiedDB.path, domains: domains, keys: keys, label: label)
}
private static func readCookies(
fromDB path: String,
domains: [String],
keys: [Data],
label: String) throws -> [ChromiumCookieRecord]
{
var db: OpaquePointer?
guard sqlite3_open_v2(path, &db, SQLITE_OPEN_READONLY, nil) == SQLITE_OK else {
throw ImportError.sqliteFailed(label: label, details: String(cString: sqlite3_errmsg(db)))
}
defer { sqlite3_close(db) }
let sql = "SELECT host_key, name, path, expires_utc, is_secure, value, encrypted_value FROM cookies"
var stmt: OpaquePointer?
guard sqlite3_prepare_v2(db, sql, -1, &stmt, nil) == SQLITE_OK else {
throw ImportError.sqliteFailed(label: label, details: String(cString: sqlite3_errmsg(db)))
}
defer { sqlite3_finalize(stmt) }
var records: [ChromiumCookieRecord] = []
while sqlite3_step(stmt) == SQLITE_ROW {
guard let hostKey = self.readText(stmt, index: 0), self.matches(domain: hostKey, patterns: domains) else {
continue
}
guard let name = self.readText(stmt, index: 1), let path = self.readText(stmt, index: 2) else {
continue
}
let value: String? = if let plain = self.readText(stmt, index: 5), !plain.isEmpty {
plain
} else if let encrypted = self.readBlob(stmt, index: 6) {
self.decrypt(encrypted, usingAnyOf: keys)
} else {
nil
}
guard let value, !value.isEmpty else { continue }
records.append(ChromiumCookieRecord(
domain: AlibabaCodingPlanCookieImporter.normalizeCookieDomain(hostKey),
name: name,
path: path,
value: value,
expires: self.chromiumExpiry(sqlite3_column_int64(stmt, 3)),
isSecure: sqlite3_column_int(stmt, 4) != 0))
}
return records.filter { record in
guard let expires = record.expires else { return true }
return expires >= Date()
}
}
private static func derivedKeys(for browser: Browser) throws -> [Data] {
var keys: [Data] = []
var sawDenied = false
for label in browser.safeStorageLabels {
switch KeychainAccessPreflight.checkGenericPassword(service: label.service, account: label.account) {
case .interactionRequired:
sawDenied = true
continue
case .allowed, .notFound, .failure:
break
}
if let password = self.safeStoragePassword(service: label.service, account: label.account) {
keys.append(self.deriveKey(from: password))
}
}
if !keys.isEmpty {
return keys
}
if sawDenied {
throw ImportError.keychainDenied(browser: browser)
}
throw ImportError.keyUnavailable(browser: browser)
}
private static func safeStoragePassword(service: String, account: String) -> String? {
// The preflight classifies prompt-requiring items as .interactionRequired, but its
// .notFound (gate disabled) and .failure outcomes still reach this read. Honor the
// access gate and keep the read strictly non-interactive so it can never prompt.
guard !KeychainAccessGate.isDisabled else { return nil }
var query: [String: Any] = [
kSecClass as String: kSecClassGenericPassword,
kSecAttrService as String: service,
kSecAttrAccount as String: account,
kSecMatchLimit as String: kSecMatchLimitOne,
kSecReturnData as String: true,
]
KeychainNoUIQuery.apply(to: &query)
var result: AnyObject?
let status = KeychainSecurity.copyMatching(query as CFDictionary, &result)
guard status == errSecSuccess, let data = result as? Data else { return nil }
return String(data: data, encoding: .utf8)
}
private static func deriveKey(from password: String) -> Data {
let salt = Data("saltysalt".utf8)
var key = Data(count: kCCKeySizeAES128)
let keyLength = key.count
_ = key.withUnsafeMutableBytes { keyBytes in
password.utf8CString.withUnsafeBytes { passBytes in
salt.withUnsafeBytes { saltBytes in
CCKeyDerivationPBKDF(
CCPBKDFAlgorithm(kCCPBKDF2),
passBytes.bindMemory(to: Int8.self).baseAddress,
passBytes.count - 1,
saltBytes.bindMemory(to: UInt8.self).baseAddress,
salt.count,
CCPseudoRandomAlgorithm(kCCPRFHmacAlgSHA1),
1003,
keyBytes.bindMemory(to: UInt8.self).baseAddress,
keyLength)
}
}
}
return key
}
private static func decrypt(_ encryptedValue: Data, usingAnyOf keys: [Data]) -> String? {
for key in keys {
if let value = self.decrypt(encryptedValue, key: key) {
return value
}
}
return nil
}
private static func decrypt(_ encryptedValue: Data, key: Data) -> String? {
guard encryptedValue.count > 3 else { return nil }
let prefix = String(data: encryptedValue.prefix(3), encoding: .utf8)
guard prefix == "v10" else { return nil }
let payload = Data(encryptedValue.dropFirst(3))
let iv = Data(repeating: 0x20, count: kCCBlockSizeAES128)
var outLength = 0
var out = Data(count: payload.count + kCCBlockSizeAES128)
let outCapacity = out.count
let status = out.withUnsafeMutableBytes { outBytes in
payload.withUnsafeBytes { payloadBytes in
key.withUnsafeBytes { keyBytes in
iv.withUnsafeBytes { ivBytes in
CCCrypt(
CCOperation(kCCDecrypt),
CCAlgorithm(kCCAlgorithmAES),
CCOptions(kCCOptionPKCS7Padding),
keyBytes.baseAddress,
key.count,
ivBytes.baseAddress,
payloadBytes.baseAddress,
payload.count,
outBytes.baseAddress,
outCapacity,
&outLength)
}
}
}
}
guard status == kCCSuccess else { return nil }
out.count = outLength
if let value = String(data: out, encoding: .utf8), !value.isEmpty {
return value
}
if out.count > 32 {
let trimmed = out.dropFirst(32)
if let value = String(data: trimmed, encoding: .utf8), !value.isEmpty {
return value
}
}
return nil
}
private static func makeCookie(from record: ChromiumCookieRecord) -> HTTPCookie? {
var properties: [HTTPCookiePropertyKey: Any] = [
.domain: record.domain,
.path: record.path,
.name: record.name,
.value: record.value,
]
if record.isSecure {
properties[.secure] = true
}
if let expires = record.expires {
properties[.expires] = expires
}
return HTTPCookie(properties: properties)
}
private static func readText(_ stmt: OpaquePointer?, index: Int32) -> String? {
guard sqlite3_column_type(stmt, index) != SQLITE_NULL,
let value = sqlite3_column_text(stmt, index)
else {
return nil
}
return String(cString: value)
}
private static func readBlob(_ stmt: OpaquePointer?, index: Int32) -> Data? {
guard sqlite3_column_type(stmt, index) != SQLITE_NULL,
let bytes = sqlite3_column_blob(stmt, index)
else {
return nil
}
return Data(bytes: bytes, count: Int(sqlite3_column_bytes(stmt, index)))
}
private static func matches(domain: String, patterns: [String]) -> Bool {
AlibabaCodingPlanCookieImporter.matchesCookieDomain(domain, patterns: patterns)
}
private static func chromiumExpiry(_ expiresUTC: Int64) -> Date? {
guard expiresUTC > 0 else { return nil }
let seconds = (Double(expiresUTC) / 1_000_000.0) - 11_644_473_600.0
guard seconds > 0 else { return nil }
return Date(timeIntervalSince1970: seconds)
}
}
#endif
@@ -0,0 +1,266 @@
import Foundation
#if os(macOS)
import SweetCookieKit
#endif
public enum AlibabaCodingPlanProviderDescriptor {
public static let descriptor: ProviderDescriptor = Self.makeDescriptor()
static func makeDescriptor() -> ProviderDescriptor {
#if os(macOS)
let browserOrder: BrowserCookieImportOrder = [
.chrome,
.chromeBeta,
.brave,
.edge,
.arc,
.firefox,
.safari,
]
#else
let browserOrder: BrowserCookieImportOrder? = nil
#endif
return ProviderDescriptor(
id: .alibaba,
metadata: ProviderMetadata(
id: .alibaba,
displayName: "Alibaba",
sessionLabel: "5-hour",
weeklyLabel: "Weekly",
opusLabel: "Monthly",
supportsOpus: true,
supportsCredits: false,
creditsHint: "",
toggleTitle: "Show Alibaba usage",
cliName: "alibaba-coding-plan",
defaultEnabled: false,
isPrimaryProvider: false,
usesAccountFallback: false,
browserCookieOrder: browserOrder,
dashboardURL: AlibabaCodingPlanAPIRegion.international.dashboardURL.absoluteString,
statusPageURL: nil,
statusLinkURL: "https://status.aliyun.com"),
branding: ProviderBranding(
iconStyle: .alibaba,
iconResourceName: "ProviderIcon-alibaba",
color: ProviderColor(red: 1.0, green: 106 / 255, blue: 0)),
tokenCost: ProviderTokenCostConfig(
supportsTokenCost: false,
noDataMessage: { "Alibaba Coding Plan cost summary is not supported." }),
fetchPlan: ProviderFetchPlan(
sourceModes: [.auto, .web, .api],
pipeline: ProviderFetchPipeline(resolveStrategies: self.resolveStrategies)),
cli: ProviderCLIConfig(
name: "alibaba-coding-plan",
aliases: ["alibaba", "bailian"],
versionDetector: nil))
}
private static func resolveStrategies(context: ProviderFetchContext) async -> [any ProviderFetchStrategy] {
switch context.sourceMode {
case .web:
return [AlibabaCodingPlanWebFetchStrategy()]
case .api:
return [AlibabaCodingPlanAPIFetchStrategy()]
case .cli, .oauth:
return []
case .auto:
break
}
if context.settings?.alibaba?.cookieSource == .off {
return [AlibabaCodingPlanAPIFetchStrategy()]
}
return [AlibabaCodingPlanWebFetchStrategy(), AlibabaCodingPlanAPIFetchStrategy()]
}
}
struct AlibabaCodingPlanWebFetchStrategy: ProviderFetchStrategy {
let id: String = "alibaba-coding-plan.web"
let kind: ProviderFetchKind = .web
func isAvailable(_ context: ProviderFetchContext) async -> Bool {
guard context.settings?.alibaba?.cookieSource != .off else { return false }
if AlibabaCodingPlanSettingsReader.cookieHeader(environment: context.env) != nil {
return true
}
if let settings = context.settings?.alibaba,
settings.cookieSource == .manual
{
return CookieHeaderNormalizer.normalize(settings.manualCookieHeader) != nil
}
#if os(macOS)
if let cached = CookieHeaderCache.load(provider: .alibaba),
!cached.cookieHeader.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty
{
return true
}
if AlibabaCodingPlanCookieImporter.hasSession(browserDetection: context.browserDetection) {
return true
}
return false
#else
return false
#endif
}
func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult {
let cookieSource = context.settings?.alibaba?.cookieSource ?? .auto
let cookieHeader = try Self.resolveCookieHeader(context: context, allowCached: true)
do {
let region = context.settings?.alibaba?.apiRegion ?? .international
let usage = try await AlibabaCodingPlanUsageFetcher.fetchUsage(
cookieHeader: cookieHeader,
region: region,
environment: context.env)
return self.makeResult(usage: usage.toUsageSnapshot(), sourceLabel: "web")
} catch let error as AlibabaCodingPlanUsageError
where (error == .invalidCredentials || error == .loginRequired) && cookieSource != .manual
{
#if os(macOS)
CookieHeaderCache.clear(provider: .alibaba)
let refreshedHeader = try Self.resolveCookieHeader(context: context, allowCached: false)
let region = context.settings?.alibaba?.apiRegion ?? .international
let usage = try await AlibabaCodingPlanUsageFetcher.fetchUsage(
cookieHeader: refreshedHeader,
region: region,
environment: context.env)
return self.makeResult(usage: usage.toUsageSnapshot(), sourceLabel: "web")
#else
throw AlibabaCodingPlanUsageError.invalidCredentials
#endif
}
}
func shouldFallback(on error: Error, context: ProviderFetchContext) -> Bool {
guard context.sourceMode == .auto else { return false }
if let urlError = error as? URLError {
switch urlError.code {
case .timedOut,
.cannotFindHost,
.cannotConnectToHost,
.networkConnectionLost,
.dnsLookupFailed,
.secureConnectionFailed,
.serverCertificateHasBadDate,
.serverCertificateUntrusted,
.serverCertificateHasUnknownRoot,
.serverCertificateNotYetValid,
.clientCertificateRejected,
.clientCertificateRequired,
.cannotLoadFromNetwork,
.internationalRoamingOff,
.callIsActive,
.dataNotAllowed,
.requestBodyStreamExhausted,
.resourceUnavailable,
.notConnectedToInternet:
return true
default:
break
}
}
if let settingsError = error as? AlibabaCodingPlanSettingsError {
switch settingsError {
case .missingCookie, .invalidCookie:
return true
case .missingToken:
return false
}
}
guard let alibabaError = error as? AlibabaCodingPlanUsageError else { return false }
switch alibabaError {
case .loginRequired:
return true
case .invalidCredentials:
return true
case .apiKeyUnavailableInRegion:
return false
case let .apiError(message):
return message.contains("HTTP 404") || message.contains("HTTP 403")
case .networkError:
return true
case .parseFailed:
return false
}
}
static func resolveCookieHeader(context: ProviderFetchContext, allowCached: Bool) throws -> String {
if let settings = context.settings?.alibaba,
settings.cookieSource == .manual
{
guard let header = CookieHeaderNormalizer.normalize(settings.manualCookieHeader) else {
throw AlibabaCodingPlanSettingsError.invalidCookie
}
return header
}
if let envCookie = AlibabaCodingPlanSettingsReader.cookieHeader(environment: context.env),
let normalized = CookieHeaderNormalizer.normalize(envCookie)
{
return normalized
}
#if os(macOS)
if allowCached,
let cached = CookieHeaderCache.load(provider: .alibaba),
!cached.cookieHeader.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty
{
return cached.cookieHeader
}
do {
let session = try AlibabaCodingPlanCookieImporter.importSession(browserDetection: context.browserDetection)
CookieHeaderCache.store(
provider: .alibaba,
cookieHeader: session.cookieHeader,
sourceLabel: session.sourceLabel)
return session.cookieHeader
} catch {
throw error
}
#else
throw AlibabaCodingPlanSettingsError.missingCookie()
#endif
}
}
struct AlibabaCodingPlanAPIFetchStrategy: ProviderFetchStrategy {
let id: String = "alibaba-coding-plan.api"
let kind: ProviderFetchKind = .apiToken
func isAvailable(_ context: ProviderFetchContext) async -> Bool {
Self.resolveToken(environment: context.env) != nil
}
func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult {
guard let apiKey = Self.resolveToken(environment: context.env) else {
throw AlibabaCodingPlanSettingsError.missingToken
}
let region = context.settings?.alibaba?.apiRegion ?? .international
let usage = try await AlibabaCodingPlanUsageFetcher.fetchUsage(
apiKey: apiKey,
region: region,
environment: context.env)
return self.makeResult(
usage: usage.toUsageSnapshot(),
sourceLabel: "api")
}
func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool {
false
}
private static func resolveToken(environment: [String: String]) -> String? {
ProviderTokenResolver.alibabaToken(environment: environment)
}
}
@@ -0,0 +1,118 @@
import Foundation
public struct AlibabaCodingPlanSettingsReader: Sendable {
private static let endpointValidator = ProviderEndpointOverrideValidator(
allowedHosts: [
"modelstudio.console.alibabacloud.com",
"bailian-singapore-cs.alibabacloud.com",
"bailian.console.aliyun.com",
"bailian-cs.console.aliyun.com",
"bailian-beijing-cs.aliyuncs.com",
])
public static let apiTokenKey = "ALIBABA_CODING_PLAN_API_KEY"
public static let qwenAPITokenKey = "ALIBABA_QWEN_API_KEY"
public static let dashScopeAPITokenKey = "DASHSCOPE_API_KEY"
public static let apiTokenEnvironmentKeys = [
Self.apiTokenKey,
Self.qwenAPITokenKey,
Self.dashScopeAPITokenKey,
]
public static let cookieHeaderKey = "ALIBABA_CODING_PLAN_COOKIE"
public static let hostKey = "ALIBABA_CODING_PLAN_HOST"
public static let quotaURLKey = "ALIBABA_CODING_PLAN_QUOTA_URL"
public static let requireProviderEndpointOverridesKey = "ALIBABA_CODING_PLAN_REQUIRE_PROVIDER_ENDPOINT_OVERRIDES"
private static let endpointOverrideKeys = [
Self.hostKey,
Self.quotaURLKey,
]
public static func apiToken(
environment: [String: String] = ProcessInfo.processInfo.environment) -> String?
{
for key in self.apiTokenEnvironmentKeys {
if let token = self.cleaned(environment[key]) { return token }
}
return nil
}
public static func hostOverride(
environment: [String: String] = ProcessInfo.processInfo.environment) -> String?
{
self.endpointValidator.validatedHost(
self.cleaned(environment[self.hostKey]),
policy: self.endpointOverrideHostPolicy(environment: environment))
}
public static func rejectedEndpointOverrideKey(
environment: [String: String] = ProcessInfo.processInfo.environment) -> String?
{
let policy = self.endpointOverrideHostPolicy(environment: environment)
return self.endpointOverrideKeys.first { key in
guard let value = self.cleaned(environment[key]) else { return false }
if key == Self.hostKey {
return self.endpointValidator.validatedHost(value, policy: policy) == nil
}
return self.endpointValidator.validatedURL(value, policy: policy) == nil
}
}
public static func cookieHeader(
environment: [String: String] = ProcessInfo.processInfo.environment) -> String?
{
self.cleaned(environment[self.cookieHeaderKey])
}
public static func quotaURL(
environment: [String: String] = ProcessInfo.processInfo.environment) -> URL?
{
self.endpointValidator.validatedURL(
self.cleaned(environment[self.quotaURLKey]),
policy: self.endpointOverrideHostPolicy(environment: environment))
}
static func endpointOverrideHostPolicy(environment: [String: String]) -> ProviderEndpointOverrideValidator
.HostPolicy {
guard let value = self.cleaned(environment[self.requireProviderEndpointOverridesKey])?.lowercased(),
["1", "true", "yes", "on"].contains(value)
else { return .allowAnyHTTPSHost }
return .providerOwnedOnly
}
static func cleaned(_ raw: String?) -> String? {
guard var value = raw?.trimmingCharacters(in: .whitespacesAndNewlines), !value.isEmpty else {
return nil
}
if (value.hasPrefix("\"") && value.hasSuffix("\"")) ||
(value.hasPrefix("'") && value.hasSuffix("'"))
{
value = String(value.dropFirst().dropLast())
}
value = value.trimmingCharacters(in: .whitespacesAndNewlines)
return value.isEmpty ? nil : value
}
}
public enum AlibabaCodingPlanSettingsError: LocalizedError, Sendable {
case missingToken
case missingCookie(details: String? = nil)
case invalidCookie
public var errorDescription: String? {
switch self {
case .missingToken:
return "Alibaba Coding Plan API key not found. " +
"Set apiKey in ~/.codexbar/config.json, ALIBABA_CODING_PLAN_API_KEY, " +
"ALIBABA_QWEN_API_KEY, or DASHSCOPE_API_KEY."
case let .missingCookie(details):
let base = "No Alibaba Coding Plan session cookies found in browsers. " +
"If you use Safari, enable Full Disk Access for CodexBar/Terminal or paste a manual Cookie header."
guard let details, !details.isEmpty else { return base }
return "\(base) \(details)"
case .invalidCookie:
return "Alibaba Coding Plan cookie header is invalid."
}
}
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,126 @@
import Foundation
public struct AlibabaCodingPlanUsageSnapshot: Sendable {
public let planName: String?
public let fiveHourUsedQuota: Int?
public let fiveHourTotalQuota: Int?
public let fiveHourNextRefreshTime: Date?
public let weeklyUsedQuota: Int?
public let weeklyTotalQuota: Int?
public let weeklyNextRefreshTime: Date?
public let monthlyUsedQuota: Int?
public let monthlyTotalQuota: Int?
public let monthlyNextRefreshTime: Date?
public let updatedAt: Date
public init(
planName: String?,
fiveHourUsedQuota: Int?,
fiveHourTotalQuota: Int?,
fiveHourNextRefreshTime: Date?,
weeklyUsedQuota: Int?,
weeklyTotalQuota: Int?,
weeklyNextRefreshTime: Date?,
monthlyUsedQuota: Int?,
monthlyTotalQuota: Int?,
monthlyNextRefreshTime: Date?,
updatedAt: Date)
{
self.planName = planName
self.fiveHourUsedQuota = fiveHourUsedQuota
self.fiveHourTotalQuota = fiveHourTotalQuota
self.fiveHourNextRefreshTime = fiveHourNextRefreshTime
self.weeklyUsedQuota = weeklyUsedQuota
self.weeklyTotalQuota = weeklyTotalQuota
self.weeklyNextRefreshTime = weeklyNextRefreshTime
self.monthlyUsedQuota = monthlyUsedQuota
self.monthlyTotalQuota = monthlyTotalQuota
self.monthlyNextRefreshTime = monthlyNextRefreshTime
self.updatedAt = updatedAt
}
}
extension AlibabaCodingPlanUsageSnapshot {
public func toUsageSnapshot() -> UsageSnapshot {
let primaryPercent = Self.usedPercent(used: self.fiveHourUsedQuota, total: self.fiveHourTotalQuota)
let secondaryPercent = Self.usedPercent(used: self.weeklyUsedQuota, total: self.weeklyTotalQuota)
let tertiaryPercent = Self.usedPercent(used: self.monthlyUsedQuota, total: self.monthlyTotalQuota)
let primary: RateWindow? = if let primaryPercent {
RateWindow(
usedPercent: primaryPercent,
windowMinutes: 5 * 60,
resetsAt: Self.normalizedFiveHourReset(
self.fiveHourNextRefreshTime,
updatedAt: self.updatedAt),
resetDescription: Self.usageDetail(used: self.fiveHourUsedQuota, total: self.fiveHourTotalQuota))
} else {
nil
}
let secondary: RateWindow? = if let secondaryPercent {
RateWindow(
usedPercent: secondaryPercent,
windowMinutes: 7 * 24 * 60,
resetsAt: self.weeklyNextRefreshTime,
resetDescription: Self.usageDetail(used: self.weeklyUsedQuota, total: self.weeklyTotalQuota))
} else {
nil
}
let tertiary: RateWindow? = if let tertiaryPercent {
RateWindow(
usedPercent: tertiaryPercent,
windowMinutes: 30 * 24 * 60,
resetsAt: self.monthlyNextRefreshTime,
resetDescription: Self.usageDetail(used: self.monthlyUsedQuota, total: self.monthlyTotalQuota))
} else {
nil
}
let loginMethod = self.planName?.trimmingCharacters(in: .whitespacesAndNewlines)
let identity = ProviderIdentitySnapshot(
providerID: .alibaba,
accountEmail: nil,
accountOrganization: nil,
loginMethod: loginMethod)
return UsageSnapshot(
primary: primary,
secondary: secondary,
tertiary: tertiary,
providerCost: nil,
updatedAt: self.updatedAt,
identity: identity)
}
private static func usedPercent(used: Int?, total: Int?) -> Double? {
guard let used, let total, total > 0 else { return nil }
let normalizedUsed = max(0, min(used, total))
return Double(normalizedUsed) / Double(total) * 100
}
private static func limitDescription(total: Int?, label: String) -> String? {
guard let total, total > 0 else { return nil }
return "\(total) requests / \(label)"
}
private static func usageDetail(used: Int?, total: Int?) -> String? {
guard let used, let total, total > 0 else { return nil }
return "\(used) / \(total) used"
}
private static func normalizedFiveHourReset(_ raw: Date?, updatedAt: Date) -> Date? {
guard let raw else { return nil }
if raw.timeIntervalSince(updatedAt) >= 60 {
return raw
}
let shifted = raw.addingTimeInterval(TimeInterval(5 * 60 * 60))
if shifted.timeIntervalSince(updatedAt) >= 60 {
return shifted
}
return updatedAt.addingTimeInterval(TimeInterval(5 * 60 * 60))
}
}
@@ -0,0 +1,60 @@
import Foundation
public enum AlibabaTokenPlanAPIRegion: String, CaseIterable, Sendable {
case international = "intl"
case chinaMainland = "cn"
public var displayName: String {
switch self {
case .international:
"International (modelstudio.console.alibabacloud.com)"
case .chinaMainland:
"China mainland (bailian.console.aliyun.com)"
}
}
public var gatewayBaseURLString: String {
switch self {
case .international:
"https://modelstudio.console.alibabacloud.com"
case .chinaMainland:
"https://bailian.console.aliyun.com"
}
}
public var dashboardOriginURLString: String {
self.gatewayBaseURLString
}
public var dashboardURL: URL {
switch self {
case .international:
URL(
string: "https://modelstudio.console.alibabacloud.com/ap-southeast-1/?tab=plan#/efm/subscription/token-plan")!
case .chinaMainland:
URL(string: "https://bailian.console.aliyun.com/cn-beijing?tab=plan#/efm/subscription/token-plan")!
}
}
public var currentRegionID: String {
switch self {
case .international:
"ap-southeast-1"
case .chinaMainland:
"cn-beijing"
}
}
public var tokenPlanProductCode: String {
switch self {
case .international:
"sfm_tokenplanteams_dp_intl"
case .chinaMainland:
"sfm_tokenplanteams_dp_cn"
}
}
public var cookieCacheScope: CookieHeaderCache.Scope {
.providerVariant("alibaba-token-plan.\(self.rawValue)")
}
}
@@ -0,0 +1,163 @@
import Foundation
#if canImport(FoundationNetworking)
import FoundationNetworking
#endif
struct AlibabaTokenPlanCookieHeaders {
private static let cachedAPIHeaderName = "__codexbar_alibaba_token_plan_api"
private static let cachedDashboardHeaderName = "__codexbar_alibaba_token_plan_dashboard"
let apiCookieHeader: String
let dashboardCookieHeader: String
init(apiCookieHeader: String, dashboardCookieHeader: String) {
self.apiCookieHeader = apiCookieHeader
self.dashboardCookieHeader = dashboardCookieHeader
}
init?(singleHeader raw: String?) {
guard let normalized = CookieHeaderNormalizer.normalize(raw) else { return nil }
self.apiCookieHeader = normalized
self.dashboardCookieHeader = normalized
}
init?(cachedHeader raw: String?) {
var valuesByName: [String: String] = [:]
for pair in CookieHeaderNormalizer.pairs(from: raw ?? "") {
valuesByName[pair.name] = pair.value
}
if let encodedAPI = valuesByName[Self.cachedAPIHeaderName],
let encodedDashboard = valuesByName[Self.cachedDashboardHeaderName],
let apiHeader = Self.decodeCachedHeader(encodedAPI),
let dashboardHeader = Self.decodeCachedHeader(encodedDashboard),
let normalizedAPI = CookieHeaderNormalizer.normalize(apiHeader),
let normalizedDashboard = CookieHeaderNormalizer.normalize(dashboardHeader)
{
self.init(apiCookieHeader: normalizedAPI, dashboardCookieHeader: normalizedDashboard)
return
}
self.init(singleHeader: raw)
}
var cacheCookieHeader: String {
[
"\(Self.cachedAPIHeaderName)=\(Self.encodeCachedHeader(self.apiCookieHeader))",
"\(Self.cachedDashboardHeaderName)=\(Self.encodeCachedHeader(self.dashboardCookieHeader))",
].joined(separator: "; ")
}
var apiCookieNames: [String] {
Self.cookieNames(from: self.apiCookieHeader)
}
var dashboardCookieNames: [String] {
Self.cookieNames(from: self.dashboardCookieHeader)
}
func hasCookie(named name: String) -> Bool {
Self.cookieNames(from: self.apiCookieHeader).contains(name) ||
Self.cookieNames(from: self.dashboardCookieHeader).contains(name)
}
private static func cookieNames(from header: String) -> [String] {
CookieHeaderNormalizer.pairs(from: header)
.map(\.name)
.filter { !$0.isEmpty }
.uniquedSorted()
}
private static func encodeCachedHeader(_ header: String) -> String {
Data(header.utf8).base64EncodedString()
}
private static func decodeCachedHeader(_ encoded: String) -> String? {
guard let data = Data(base64Encoded: encoded) else { return nil }
return String(data: data, encoding: .utf8)
}
}
enum AlibabaTokenPlanCookieHeader {
static func headers(
from cookies: [HTTPCookie],
region: AlibabaTokenPlanAPIRegion = .international,
environment: [String: String] = ProcessInfo.processInfo.environment) -> AlibabaTokenPlanCookieHeaders?
{
guard let apiHeader = self.header(
from: cookies,
targetURL: AlibabaTokenPlanUsageFetcher.resolveQuotaURL(region: region, environment: environment)),
let dashboardHeader = self.header(
from: cookies,
targetURL: AlibabaTokenPlanUsageFetcher.dashboardURL(region: region, environment: environment))
else {
return nil
}
return AlibabaTokenPlanCookieHeaders(apiCookieHeader: apiHeader, dashboardCookieHeader: dashboardHeader)
}
static func header(from cookies: [HTTPCookie], targetURL: URL) -> String? {
var byName: [String: HTTPCookie] = [:]
for cookie in cookies {
guard !cookie.name.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty else { continue }
guard !cookie.value.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty else { continue }
if let expiry = cookie.expiresDate, expiry < Date() { continue }
guard self.matchesRequestURL(cookie: cookie, url: targetURL) else { continue }
if let existing = byName[cookie.name] {
if self.cookieSortKey(for: cookie) >= self.cookieSortKey(for: existing) {
byName[cookie.name] = cookie
}
} else {
byName[cookie.name] = cookie
}
}
guard !byName.isEmpty else { return nil }
return byName.keys.sorted().compactMap { name in
guard let cookie = byName[name] else { return nil }
return "\(cookie.name)=\(cookie.value)"
}.joined(separator: "; ")
}
private static func matchesRequestURL(cookie: HTTPCookie, url: URL) -> Bool {
guard let host = url.host?.lowercased() else { return false }
let normalizedDomain = cookie.domain.lowercased().trimmingCharacters(in: CharacterSet(charactersIn: "."))
guard !normalizedDomain.isEmpty else { return false }
guard host == normalizedDomain || host.hasSuffix(".\(normalizedDomain)") else { return false }
let cookiePath = cookie.path.isEmpty ? "/" : cookie.path
let requestPath = url.path.isEmpty ? "/" : url.path
if requestPath == cookiePath {
return true
}
guard requestPath.hasPrefix(cookiePath) else { return false }
guard cookiePath != "/" else { return true }
if cookiePath.hasSuffix("/") {
return true
}
guard
let boundaryIndex = requestPath.index(
requestPath.startIndex,
offsetBy: cookiePath.count,
limitedBy: requestPath.endIndex),
boundaryIndex < requestPath.endIndex
else {
return true
}
return requestPath[boundaryIndex] == "/"
}
private static func cookieSortKey(for cookie: HTTPCookie) -> (Int, Int, Date) {
let pathLength = cookie.path.count
let normalizedDomain = cookie.domain.lowercased().trimmingCharacters(in: CharacterSet(charactersIn: "."))
let domainLength = normalizedDomain.count
let expiry = cookie.expiresDate ?? .distantPast
return (pathLength, domainLength, expiry)
}
}
extension [String] {
fileprivate func uniquedSorted() -> [String] {
Array(Set(self)).sorted()
}
}
@@ -0,0 +1,289 @@
import Foundation
#if os(macOS)
import SweetCookieKit
#endif
public enum AlibabaTokenPlanProviderDescriptor {
public static let descriptor: ProviderDescriptor = Self.makeDescriptor()
static func makeDescriptor() -> ProviderDescriptor {
#if os(macOS)
let browserOrder: BrowserCookieImportOrder = [
.chrome,
.chromeBeta,
.brave,
.edge,
.arc,
.firefox,
.safari,
]
#else
let browserOrder: BrowserCookieImportOrder? = nil
#endif
return ProviderDescriptor(
id: .alibabatokenplan,
metadata: ProviderMetadata(
id: .alibabatokenplan,
displayName: "Alibaba Token Plan",
sessionLabel: "Credits",
weeklyLabel: "Usage",
opusLabel: nil,
supportsOpus: false,
supportsCredits: false,
creditsHint: "",
toggleTitle: "Show Alibaba Token Plan usage",
cliName: "alibaba-token-plan",
defaultEnabled: false,
isPrimaryProvider: false,
usesAccountFallback: false,
browserCookieOrder: browserOrder,
dashboardURL: AlibabaTokenPlanUsageFetcher.dashboardURL.absoluteString,
statusPageURL: nil,
statusLinkURL: "https://status.aliyun.com"),
branding: ProviderBranding(
iconStyle: .alibaba,
iconResourceName: "ProviderIcon-alibaba",
color: ProviderColor(red: 1.0, green: 106 / 255, blue: 0)),
tokenCost: ProviderTokenCostConfig(
supportsTokenCost: false,
noDataMessage: { "Alibaba Token Plan cost summary is not supported." }),
fetchPlan: ProviderFetchPlan(
sourceModes: [.auto, .web],
pipeline: ProviderFetchPipeline(resolveStrategies: self.resolveStrategies)),
cli: ProviderCLIConfig(
name: "alibaba-token-plan",
aliases: ["alibaba-token", "bailian-token-plan"],
versionDetector: nil))
}
private static func resolveStrategies(context: ProviderFetchContext) async -> [any ProviderFetchStrategy] {
guard context.settings?.alibabaTokenPlan?.cookieSource != .off else { return [] }
switch context.sourceMode {
case .auto, .web:
return [AlibabaTokenPlanWebFetchStrategy()]
case .api, .cli, .oauth:
return []
}
}
}
struct AlibabaTokenPlanWebFetchStrategy: ProviderFetchStrategy {
private static let log = CodexBarLog.logger("alibaba-token-plan")
let id: String = "alibaba-token-plan.web"
let kind: ProviderFetchKind = .web
func isAvailable(_ context: ProviderFetchContext) async -> Bool {
guard context.settings?.alibabaTokenPlan?.cookieSource != .off else { return false }
let region = context.settings?.alibabaTokenPlan?.apiRegion ?? .international
if AlibabaTokenPlanSettingsReader.cookieHeader(environment: context.env) != nil {
return true
}
if let settings = context.settings?.alibabaTokenPlan,
settings.cookieSource == .manual
{
return CookieHeaderNormalizer.normalize(settings.manualCookieHeader) != nil
}
#if os(macOS)
if let cached = Self.cachedCookieEntry(region: region),
!cached.cookieHeader.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty
{
return true
}
return true
#else
return false
#endif
}
func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult {
let cookieSource = context.settings?.alibabaTokenPlan?.cookieSource ?? .auto
let region = context.settings?.alibabaTokenPlan?.apiRegion ?? .international
let cookieHeaders = try Self.resolveCookieHeaders(context: context, allowCached: true, region: region)
do {
let usage = try await AlibabaTokenPlanUsageFetcher.fetchUsage(
apiCookieHeader: cookieHeaders.apiCookieHeader,
dashboardCookieHeader: cookieHeaders.dashboardCookieHeader,
region: region,
environment: context.env)
return self.makeResult(usage: usage.toUsageSnapshot(), sourceLabel: "web")
} catch let error as AlibabaTokenPlanUsageError
where error.isCredentialFailure && cookieSource != .manual
{
#if os(macOS)
CookieHeaderCache.clear(provider: .alibabatokenplan, scope: region.cookieCacheScope)
let refreshedHeaders = try Self.resolveCookieHeaders(context: context, allowCached: false, region: region)
let usage = try await AlibabaTokenPlanUsageFetcher.fetchUsage(
apiCookieHeader: refreshedHeaders.apiCookieHeader,
dashboardCookieHeader: refreshedHeaders.dashboardCookieHeader,
region: region,
environment: context.env)
return self.makeResult(usage: usage.toUsageSnapshot(), sourceLabel: "web")
#else
throw error
#endif
}
}
func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool {
false
}
static func resolveCookieHeader(context: ProviderFetchContext, allowCached: Bool) throws -> String {
try self.resolveCookieHeaders(context: context, allowCached: allowCached, region: .international)
.apiCookieHeader
}
static func resolveCookieHeaders(
context: ProviderFetchContext,
allowCached: Bool,
region: AlibabaTokenPlanAPIRegion = .international) throws -> AlibabaTokenPlanCookieHeaders
{
if let settings = context.settings?.alibabaTokenPlan,
settings.cookieSource == .manual
{
guard let headers = AlibabaTokenPlanCookieHeaders(singleHeader: settings.manualCookieHeader) else {
self.log.warning("Alibaba Token Plan manual cookie header is invalid")
throw AlibabaTokenPlanSettingsError.invalidCookie
}
Self.log.info(
"Alibaba Token Plan using manual cookie header",
metadata: [
"apiCookieNames": headers.apiCookieNames.joined(separator: ","),
"dashboardCookieNames": headers.dashboardCookieNames.joined(separator: ","),
"hasSecToken": headers.hasCookie(named: "sec_token") ? "1" : "0",
])
return headers
}
if let envCookie = AlibabaTokenPlanSettingsReader.cookieHeader(environment: context.env),
let headers = AlibabaTokenPlanCookieHeaders(singleHeader: envCookie)
{
Self.log.info(
"Alibaba Token Plan using environment cookie header",
metadata: [
"apiCookieNames": headers.apiCookieNames.joined(separator: ","),
"dashboardCookieNames": headers.dashboardCookieNames.joined(separator: ","),
"hasSecToken": headers.hasCookie(named: "sec_token") ? "1" : "0",
])
return headers
}
#if os(macOS)
if allowCached,
let cached = Self.cachedCookieEntry(region: region),
let headers = AlibabaTokenPlanCookieHeaders(cachedHeader: cached.cookieHeader)
{
Self.log.info(
"Alibaba Token Plan using cached browser cookie header",
metadata: [
"source": cached.sourceLabel,
"apiCookieNames": headers.apiCookieNames.joined(separator: ","),
"dashboardCookieNames": headers.dashboardCookieNames.joined(separator: ","),
"hasSecToken": headers.hasCookie(named: "sec_token") ? "1" : "0",
])
return headers
}
do {
var importLog: [String] = []
let session = try AlibabaCodingPlanCookieImporter.importSession(
browserDetection: context.browserDetection,
logger: { importLog.append($0) })
let rawCookieNames = session.cookies.map(\.name).filter { !$0.isEmpty }.uniquedSorted()
guard let headers = AlibabaTokenPlanCookieHeader.headers(
from: session.cookies,
region: region,
environment: context.env)
else {
Self.log.warning(
"Alibaba Token Plan browser cookie header was empty",
metadata: [
"source": session.sourceLabel,
"rawCookieNames": rawCookieNames.joined(separator: ","),
])
throw AlibabaTokenPlanSettingsError.missingCookie(
details: "No Alibaba Token Plan browser cookies were available after import.")
}
CookieHeaderCache.store(
provider: .alibabatokenplan,
scope: region.cookieCacheScope,
cookieHeader: headers.cacheCookieHeader,
sourceLabel: session.sourceLabel)
Self.log.info(
"Alibaba Token Plan imported browser cookies",
metadata: [
"source": session.sourceLabel,
"rawCookieNames": rawCookieNames.joined(separator: ","),
"apiCookieNames": headers.apiCookieNames.joined(separator: ","),
"dashboardCookieNames": headers.dashboardCookieNames.joined(separator: ","),
"hasSecToken": headers.hasCookie(named: "sec_token") ? "1" : "0",
"importLogLines": "\(importLog.count)",
])
return headers
} catch {
Self.log.warning(
"Alibaba Token Plan cookie resolution failed",
metadata: ["error": error.localizedDescription])
throw AlibabaTokenPlanSettingsError.missingCookie(details: Self.missingCookieDetails(from: error))
}
#else
throw AlibabaTokenPlanSettingsError.missingCookie()
#endif
}
#if os(macOS)
/// The former unscoped cache only ever represented the China gateway. Never expose it to
/// International requests; migrate it into the China scope after a successful scoped write.
private static func cachedCookieEntry(region: AlibabaTokenPlanAPIRegion) -> CookieHeaderCache.Entry? {
if let scoped = CookieHeaderCache.load(provider: .alibabatokenplan, scope: region.cookieCacheScope) {
return scoped
}
guard region == .chinaMainland,
let legacy = CookieHeaderCache.load(provider: .alibabatokenplan)
else { return nil }
CookieHeaderCache.store(
provider: .alibabatokenplan,
scope: region.cookieCacheScope,
cookieHeader: legacy.cookieHeader,
sourceLabel: legacy.sourceLabel,
now: legacy.storedAt)
if let migrated = CookieHeaderCache.load(provider: .alibabatokenplan, scope: region.cookieCacheScope) {
CookieHeaderCache.clear(provider: .alibabatokenplan)
return migrated
}
return legacy
}
#endif
private static func missingCookieDetails(from error: Error) -> String? {
if case let AlibabaCodingPlanSettingsError.missingCookie(details) = error {
return details
}
let message = error.localizedDescription.trimmingCharacters(in: .whitespacesAndNewlines)
return message.isEmpty ? nil : message
}
}
extension [String] {
fileprivate func uniquedSorted() -> [String] {
Array(Set(self)).sorted()
}
}
extension AlibabaTokenPlanUsageError {
fileprivate var isCredentialFailure: Bool {
switch self {
case .loginRequired, .invalidCredentials:
true
case .apiError, .networkError, .parseFailed:
false
}
}
}
@@ -0,0 +1,67 @@
import Foundation
public struct AlibabaTokenPlanSettingsReader: Sendable {
public static let cookieHeaderKey = "ALIBABA_TOKEN_PLAN_COOKIE"
public static let hostKey = "ALIBABA_TOKEN_PLAN_HOST"
public static let quotaURLKey = "ALIBABA_TOKEN_PLAN_QUOTA_URL"
private static let endpointValidator = ProviderEndpointOverrideValidator(
allowedHosts: [
"modelstudio.console.alibabacloud.com",
"bailian.console.aliyun.com",
])
public static func cookieHeader(
environment: [String: String] = ProcessInfo.processInfo.environment) -> String?
{
self.cleaned(environment[self.cookieHeaderKey])
}
public static func hostOverride(
environment: [String: String] = ProcessInfo.processInfo.environment) -> String?
{
self.endpointValidator.validatedHost(
self.cleaned(environment[self.hostKey]),
policy: .allowAnyHTTPSHost)
}
public static func quotaURL(
environment: [String: String] = ProcessInfo.processInfo.environment) -> URL?
{
self.endpointValidator.validatedURL(
self.cleaned(environment[self.quotaURLKey]),
policy: .allowAnyHTTPSHost)
}
static func cleaned(_ raw: String?) -> String? {
guard var value = raw?.trimmingCharacters(in: .whitespacesAndNewlines), !value.isEmpty else {
return nil
}
if (value.hasPrefix("\"") && value.hasSuffix("\"")) ||
(value.hasPrefix("'") && value.hasSuffix("'"))
{
value = String(value.dropFirst().dropLast())
}
value = value.trimmingCharacters(in: .whitespacesAndNewlines)
return value.isEmpty ? nil : value
}
}
public enum AlibabaTokenPlanSettingsError: LocalizedError, Sendable {
case missingCookie(details: String? = nil)
case invalidCookie
public var errorDescription: String? {
switch self {
case let .missingCookie(details):
let base = "No Alibaba Token Plan session cookies found in browsers. " +
"Sign in to Model Studio/Bailian in Chrome, " +
"allow CodexBar to access Chrome Safe Storage in Keychain Access, " +
"or paste a manual Cookie header."
guard let details, !details.isEmpty else { return base }
return "\(base) \(details)"
case .invalidCookie:
return "Alibaba Token Plan cookie header is invalid."
}
}
}
@@ -0,0 +1,997 @@
import Foundation
#if canImport(FoundationNetworking)
import FoundationNetworking
#endif
public enum AlibabaTokenPlanUsageError: LocalizedError, Sendable, Equatable {
case loginRequired
case invalidCredentials
case apiError(String)
case networkError(String)
case parseFailed(String)
public var errorDescription: String? {
switch self {
case .loginRequired:
"Alibaba Token Plan login required."
case .invalidCredentials:
"Alibaba Token Plan credentials are invalid."
case let .apiError(message):
"Alibaba Token Plan API error: \(message)"
case let .networkError(message):
"Alibaba Token Plan network error: \(message)"
case let .parseFailed(message):
"Could not parse Alibaba Token Plan usage: \(message)"
}
}
}
// swiftlint:disable:next type_body_length
public struct AlibabaTokenPlanUsageFetcher: Sendable {
private static let log = CodexBarLog.logger("alibaba-token-plan")
private static let bssServiceCode = "BssOpenAPI-V3"
private static let subscriptionSummaryAction = "GetSubscriptionSummary"
private static let browserLikeUserAgent =
"Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) " +
"AppleWebKit/537.36 (KHTML, like Gecko) Chrome/143.0.0.0 Safari/537.36"
private static let safariLikeUserAgent =
"Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) " +
"AppleWebKit/605.1.15 (KHTML, like Gecko) Version/26.3 Safari/605.1.15"
public static var dashboardURL: URL {
Self.dashboardURL(region: .international, environment: ProcessInfo.processInfo.environment)
}
public static func dashboardURL(
region: AlibabaTokenPlanAPIRegion,
environment: [String: String] = ProcessInfo.processInfo.environment) -> URL
{
if let host = AlibabaTokenPlanSettingsReader.hostOverride(environment: environment),
let base = ProviderEndpointOverrideValidator.normalizedHTTPSURL(from: host),
var components = URLComponents(url: base, resolvingAgainstBaseURL: false),
let dashboardComponents = URLComponents(url: region.dashboardURL, resolvingAgainstBaseURL: false)
{
components.path = dashboardComponents.path
components.percentEncodedQuery = dashboardComponents.percentEncodedQuery
components.fragment = dashboardComponents.fragment
return components.url ?? region.dashboardURL
}
return region.dashboardURL
}
public static func fetchUsage(
cookieHeader: String,
region: AlibabaTokenPlanAPIRegion = .international,
environment: [String: String] = ProcessInfo.processInfo.environment,
now: Date = Date()) async throws -> AlibabaTokenPlanUsageSnapshot
{
guard let headers = AlibabaTokenPlanCookieHeaders(singleHeader: cookieHeader) else {
throw AlibabaTokenPlanSettingsError.invalidCookie
}
return try await self.fetchUsage(
apiCookieHeader: headers.apiCookieHeader,
dashboardCookieHeader: headers.dashboardCookieHeader,
region: region,
environment: environment,
now: now)
}
static func fetchUsage(
apiCookieHeader: String,
dashboardCookieHeader: String,
region: AlibabaTokenPlanAPIRegion = .international,
environment: [String: String] = ProcessInfo.processInfo.environment,
now: Date = Date(),
session overrideSession: URLSession? = nil) async throws -> AlibabaTokenPlanUsageSnapshot
{
guard let normalizedAPIHeader = CookieHeaderNormalizer.normalize(apiCookieHeader),
let normalizedDashboardHeader = CookieHeaderNormalizer.normalize(dashboardCookieHeader)
else {
throw AlibabaTokenPlanSettingsError.invalidCookie
}
let url = self.resolveQuotaURL(region: region, environment: environment)
let apiRedirectDiagnostics = RedirectDiagnostics(cookieHeader: normalizedAPIHeader)
let dashboardRedirectDiagnostics: RedirectDiagnostics?
let apiSession: URLSession
let dashboardSession: URLSession
if let overrideSession {
apiSession = overrideSession
dashboardSession = overrideSession
dashboardRedirectDiagnostics = nil
} else {
let dashboardDiagnostics = RedirectDiagnostics(cookieHeader: normalizedDashboardHeader)
apiSession = URLSession(
configuration: .default,
delegate: apiRedirectDiagnostics,
delegateQueue: nil)
dashboardSession = URLSession(
configuration: .default,
delegate: dashboardDiagnostics,
delegateQueue: nil)
dashboardRedirectDiagnostics = dashboardDiagnostics
}
defer {
if overrideSession == nil {
apiSession.invalidateAndCancel()
dashboardSession.invalidateAndCancel()
}
}
let secToken = await self.resolveSECToken(
dashboardCookieHeader: normalizedDashboardHeader,
apiCookieHeader: normalizedAPIHeader,
region: region,
environment: environment,
session: dashboardSession)
Self.log.info(
"Fetching Alibaba Token Plan usage",
metadata: [
"apiHost": url.host ?? "unknown",
"region": region.rawValue,
"apiCookieNames": self.cookieNamesDescription(from: normalizedAPIHeader),
"dashboardCookieNames": self.cookieNamesDescription(from: normalizedDashboardHeader),
"hasCSRF": self.hasCSRF(in: normalizedAPIHeader) ? "1" : "0",
"secTokenSource": secToken == nil ? "missing" : "resolved",
])
var request = URLRequest(url: url)
request.httpMethod = "POST"
request.timeoutInterval = 20
request.httpBody = self.subscriptionSummaryRequestBody(region: region, secToken: secToken)
request.setValue("application/x-www-form-urlencoded", forHTTPHeaderField: "Content-Type")
request.setValue("*/*", forHTTPHeaderField: "Accept")
request.setValue(normalizedAPIHeader, forHTTPHeaderField: "Cookie")
if let csrf = self.extractCookieValue(name: "login_aliyunid_csrf", from: normalizedAPIHeader) ??
self.extractCookieValue(name: "csrf", from: normalizedAPIHeader)
{
request.setValue(csrf, forHTTPHeaderField: "x-xsrf-token")
request.setValue(csrf, forHTTPHeaderField: "x-csrf-token")
}
request.setValue("XMLHttpRequest", forHTTPHeaderField: "X-Requested-With")
request.setValue(Self.browserLikeUserAgent, forHTTPHeaderField: "User-Agent")
request.setValue(region.dashboardOriginURLString, forHTTPHeaderField: "Origin")
request.setValue(
Self.dashboardURL(region: region, environment: environment).absoluteString,
forHTTPHeaderField: "Referer")
let data: Data
let response: URLResponse
do {
(data, response) = try await apiSession.data(for: request)
} catch {
Self.log.error(
"Alibaba Token Plan request failed",
metadata: [
"apiHost": url.host ?? "unknown",
"error": error.localizedDescription,
])
throw AlibabaTokenPlanUsageError.networkError(error.localizedDescription)
}
if let dashboardRedirectDiagnostics, !dashboardRedirectDiagnostics.redirects.isEmpty {
Self.log.info(
"Alibaba Token Plan dashboard redirects",
metadata: [
"count": "\(dashboardRedirectDiagnostics.redirects.count)",
"items": dashboardRedirectDiagnostics.redirects.joined(separator: " | "),
])
}
if !apiRedirectDiagnostics.redirects.isEmpty {
Self.log.info(
"Alibaba Token Plan redirects",
metadata: [
"count": "\(apiRedirectDiagnostics.redirects.count)",
"items": apiRedirectDiagnostics.redirects.joined(separator: " | "),
])
}
guard let httpResponse = response as? HTTPURLResponse else {
Self.log.error("Alibaba Token Plan response was not HTTP")
throw AlibabaTokenPlanUsageError.networkError("Invalid response")
}
Self.log.info(
"Alibaba Token Plan HTTP response",
metadata: [
"status": "\(httpResponse.statusCode)",
"bodyBytes": "\(data.count)",
"contentType": httpResponse.value(forHTTPHeaderField: "Content-Type") ?? "none",
])
guard httpResponse.statusCode == 200 else {
if httpResponse.statusCode == 401 || httpResponse.statusCode == 403 {
throw AlibabaTokenPlanUsageError.loginRequired
}
Self.log.error("Alibaba Token Plan returned HTTP \(httpResponse.statusCode)")
throw AlibabaTokenPlanUsageError.apiError("HTTP \(httpResponse.statusCode)")
}
return try self.parseUsageSnapshot(from: data, now: now)
}
static func resolveQuotaURL(
region: AlibabaTokenPlanAPIRegion,
environment: [String: String]) -> URL
{
if let override = AlibabaTokenPlanSettingsReader.quotaURL(environment: environment) {
return override
}
if let host = AlibabaTokenPlanSettingsReader.hostOverride(environment: environment),
let hostURL = self.quotaURL(from: host)
{
return hostURL
}
return self.defaultQuotaURL(region: region)
}
static var defaultQuotaURL: URL {
self.defaultQuotaURL(region: .international)
}
static func defaultQuotaURL(region: AlibabaTokenPlanAPIRegion) -> URL {
var components = URLComponents(string: region.gatewayBaseURLString)!
components.path = "/data/api.json"
components.queryItems = [
URLQueryItem(name: "action", value: Self.subscriptionSummaryAction),
URLQueryItem(name: "product", value: Self.bssServiceCode),
URLQueryItem(name: "_tag", value: ""),
]
return components.url!
}
static func parseUsageSnapshot(from data: Data, now: Date = Date()) throws -> AlibabaTokenPlanUsageSnapshot {
guard !data.isEmpty else {
throw AlibabaTokenPlanUsageError.parseFailed("Empty response body")
}
let object: Any
do {
object = try JSONSerialization.jsonObject(with: data, options: [])
} catch {
if self.isLikelyLoginHTML(data) {
throw AlibabaTokenPlanUsageError.loginRequired
}
throw AlibabaTokenPlanUsageError.parseFailed("Invalid JSON response")
}
let expanded = self.expandedJSON(object)
guard let dictionary = expanded as? [String: Any] else {
throw AlibabaTokenPlanUsageError.parseFailed("Unexpected payload")
}
try self.throwIfErrorPayload(dictionary)
let summary = self.findSubscriptionSummary(in: dictionary) ?? dictionary
let total = self.anyDouble(for: Self.totalQuotaKeys, in: summary)
let remaining = self.anyDouble(for: Self.remainingQuotaKeys, in: summary)
let used = self.anyDouble(for: Self.usedQuotaKeys, in: summary) ??
total.flatMap { total in remaining.map { max(0, total - $0) } }
let resetsAt = self.findResetDate(in: summary) ?? self.findResetDate(in: dictionary)
let totalCount = self.anyDouble(for: Self.subscriptionCountKeys, in: summary)
let planName = self.findPlanName(in: summary) ?? ((totalCount ?? 0) > 0 || total != nil ? "TOKEN PLAN" : nil)
if planName == nil, total == nil, used == nil, remaining == nil, totalCount == nil {
let diagnostics = self.payloadDiagnostics(payload: dictionary)
Self.log.error("Alibaba Token Plan payload missing expected fields: \(diagnostics)")
throw AlibabaTokenPlanUsageError.parseFailed("Missing token plan data (\(diagnostics))")
}
return AlibabaTokenPlanUsageSnapshot(
planName: planName,
usedQuota: used,
totalQuota: total,
remainingQuota: remaining,
resetsAt: resetsAt,
updatedAt: now)
}
private static func subscriptionSummaryRequestBody(region: AlibabaTokenPlanAPIRegion, secToken: String?) -> Data {
let paramsObject = ["ProductCode": region.tokenPlanProductCode]
guard let paramsData = try? JSONSerialization.data(withJSONObject: paramsObject, options: []),
let paramsString = String(data: paramsData, encoding: .utf8)
else {
return Data()
}
var components = URLComponents()
var queryItems = [
URLQueryItem(name: "product", value: Self.bssServiceCode),
URLQueryItem(name: "action", value: Self.subscriptionSummaryAction),
URLQueryItem(name: "params", value: paramsString),
URLQueryItem(name: "region", value: region.currentRegionID),
]
if let secToken, !secToken.isEmpty {
queryItems.append(URLQueryItem(name: "sec_token", value: secToken))
}
components.queryItems = queryItems
return Data((components.percentEncodedQuery ?? "").utf8)
}
private static func resolveSECToken(
dashboardCookieHeader: String,
apiCookieHeader: String,
region: AlibabaTokenPlanAPIRegion,
environment: [String: String],
session: URLSession) async -> String?
{
let cookieSECToken = self.extractCookieValue(name: "sec_token", from: dashboardCookieHeader) ??
self.extractCookieValue(name: "sec_token", from: apiCookieHeader)
var request = URLRequest(url: self.dashboardURL(region: region, environment: environment))
request.httpMethod = "GET"
request.timeoutInterval = 10
request.setValue(dashboardCookieHeader, forHTTPHeaderField: "Cookie")
request.setValue(Self.safariLikeUserAgent, forHTTPHeaderField: "User-Agent")
request.setValue(
"text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8",
forHTTPHeaderField: "Accept")
if let (data, response) = try? await session.data(for: request),
let httpResponse = response as? HTTPURLResponse,
httpResponse.statusCode == 200,
let html = String(data: data, encoding: .utf8),
let token = self.extractSECToken(from: html)
{
Self.log.info(
"Resolved Alibaba Token Plan sec_token from dashboard HTML",
metadata: [
"dashboardHost": request.url?.host ?? "unknown",
"htmlBytes": "\(data.count)",
])
return token
}
if let token = await self.fetchSECTokenFromUserInfo(
cookieHeader: dashboardCookieHeader,
region: region,
environment: environment,
session: session)
{
return token
}
if let cookieSECToken, !cookieSECToken.isEmpty {
Self.log.info("Resolved Alibaba Token Plan sec_token from cookies")
return cookieSECToken
}
Self.log.info(
"Alibaba Token Plan sec_token missing; continuing with cookie-only request",
metadata: [
"dashboardCookieNames": self.cookieNamesDescription(from: dashboardCookieHeader),
"apiCookieNames": self.cookieNamesDescription(from: apiCookieHeader),
])
return nil
}
private static func fetchSECTokenFromUserInfo(
cookieHeader: String,
region: AlibabaTokenPlanAPIRegion,
environment: [String: String],
session: URLSession) async -> String?
{
let baseURL = self.consoleBaseURL(region: region, environment: environment)
let userInfoURL = baseURL.appendingPathComponent("tool/user/info.json")
var request = URLRequest(url: userInfoURL)
request.httpMethod = "GET"
request.timeoutInterval = 10
request.setValue(cookieHeader, forHTTPHeaderField: "Cookie")
request.setValue(Self.safariLikeUserAgent, forHTTPHeaderField: "User-Agent")
request.setValue("application/json, text/plain, */*", forHTTPHeaderField: "Accept")
let referer = baseURL.absoluteString.hasSuffix("/") ? baseURL.absoluteString : "\(baseURL.absoluteString)/"
request.setValue(referer, forHTTPHeaderField: "Referer")
guard let (data, response) = try? await session.data(for: request),
let httpResponse = response as? HTTPURLResponse,
httpResponse.statusCode == 200,
let object = try? JSONSerialization.jsonObject(with: data, options: [])
else {
return nil
}
let expanded = self.expandedJSON(object)
guard let token = self.findFirstString(forKeys: ["secToken", "sec_token"], in: expanded),
!token.isEmpty
else {
return nil
}
Self.log.info(
"Resolved Alibaba Token Plan sec_token from user info",
metadata: [
"userInfoHost": userInfoURL.host ?? "unknown",
"bodyBytes": "\(data.count)",
])
return token
}
private static func consoleBaseURL(
region: AlibabaTokenPlanAPIRegion,
environment: [String: String]) -> URL
{
let dashboard = self.dashboardURL(region: region, environment: environment)
var components = URLComponents()
components.scheme = dashboard.scheme
components.host = dashboard.host
components.port = dashboard.port
return components.url ?? URL(string: region.dashboardOriginURLString)!
}
private static func quotaURL(from rawHost: String) -> URL? {
let cleaned = AlibabaTokenPlanSettingsReader.cleaned(rawHost)
guard let cleaned else { return nil }
guard let base = ProviderEndpointOverrideValidator.normalizedHTTPSURL(from: cleaned) else { return nil }
var components = URLComponents(url: base, resolvingAgainstBaseURL: false)
let defaultComponents = URLComponents(
url: Self.defaultQuotaURL(region: .international),
resolvingAgainstBaseURL: false)
components?.path = "/data/api.json"
components?.queryItems = defaultComponents?.queryItems
return components?.url
}
private final class RedirectDiagnostics: NSObject, URLSessionTaskDelegate, @unchecked Sendable {
private let cookieHeader: String
var redirects: [String] = []
init(cookieHeader: String) {
self.cookieHeader = cookieHeader
}
func urlSession(
_: URLSession,
task _: URLSessionTask,
willPerformHTTPRedirection response: HTTPURLResponse,
newRequest request: URLRequest,
completionHandler: @escaping (URLRequest?) -> Void)
{
let from = AlibabaTokenPlanUsageFetcher.redactedURLDescription(response.url)
let to = AlibabaTokenPlanUsageFetcher.redactedURLDescription(request.url)
self.redirects.append("\(response.statusCode) \(from) -> \(to)")
completionHandler(AlibabaTokenPlanUsageFetcher.redirectedRequest(
response: response,
request: request,
cookieHeader: self.cookieHeader))
}
}
private static func redactedURLDescription(_ url: URL?) -> String {
guard let url else { return "unknown" }
var components = URLComponents(url: url, resolvingAgainstBaseURL: false)
components?.query = nil
components?.fragment = nil
return components?.string ?? "\(url.scheme ?? "unknown")://\(url.host ?? "unknown")"
}
static func redirectedRequest(
response: HTTPURLResponse,
request: URLRequest,
cookieHeader: String) -> URLRequest?
{
guard request.url?.scheme?.lowercased() == "https" else {
return nil
}
var updated = request
if self.shouldForwardRedirectCookies(from: response.url, to: request.url) {
updated.setValue(cookieHeader, forHTTPHeaderField: "Cookie")
} else {
updated.setValue(nil, forHTTPHeaderField: "Cookie")
}
return updated
}
private static func shouldForwardRedirectCookies(from sourceURL: URL?, to targetURL: URL?) -> Bool {
guard let sourceHost = sourceURL?.host?.lowercased(),
let targetHost = targetURL?.host?.lowercased()
else {
return false
}
return sourceHost == targetHost
}
private static func throwIfErrorPayload(_ dictionary: [String: Any]) throws {
if self.parseBool(dictionary["successResponse"]) == false {
if let statusCode = self.findFirstInt(forKeys: ["statusCode", "status_code", "code"], in: dictionary),
statusCode == 401 || statusCode == 403
{
throw AlibabaTokenPlanUsageError.invalidCredentials
}
let code = self.findFirstString(forKeys: ["code", "status", "statusCode"], in: dictionary)
let message = self.findFirstString(forKeys: ["message", "msg", "statusMessage"], in: dictionary) ??
code ??
"request was not successful"
if self.isLoginOrTokenError(code: code, message: message) {
throw AlibabaTokenPlanUsageError.loginRequired
}
throw AlibabaTokenPlanUsageError.apiError(message)
}
if self.findBoolValues(forKeys: ["Success", "success"], in: dictionary).contains(false) {
let code = self.findFirstString(forKeys: ["Code", "code"], in: dictionary)
let message = self.findFirstString(forKeys: ["Message", "message", "msg", "Code", "code"], in: dictionary)
?? "request was not successful"
if self.isLoginOrTokenError(code: code, message: message) {
throw AlibabaTokenPlanUsageError.loginRequired
}
throw AlibabaTokenPlanUsageError.apiError(message)
}
if let statusCode = self.findFirstInt(forKeys: ["statusCode", "status_code", "code"], in: dictionary),
statusCode != 0,
statusCode != 200
{
let message = self.findFirstString(
forKeys: ["statusMessage", "status_msg", "message", "msg"],
in: dictionary)
?? "status code \(statusCode)"
if statusCode == 401 || statusCode == 403 {
throw AlibabaTokenPlanUsageError.invalidCredentials
}
throw AlibabaTokenPlanUsageError.apiError(message)
}
let codeText = self.findFirstString(forKeys: ["code", "status", "statusCode"], in: dictionary)?.lowercased()
let messageText = self.findFirstString(forKeys: ["message", "msg", "statusMessage"], in: dictionary)?
.lowercased()
if self.isLoginOrTokenError(code: codeText, message: messageText) {
throw AlibabaTokenPlanUsageError.loginRequired
}
}
private static func isLoginOrTokenError(code: String?, message: String?) -> Bool {
let combined = [code, message]
.compactMap { $0?.lowercased() }
.joined(separator: " ")
return combined.contains("needlogin") ||
combined.contains("login") ||
combined.contains("postonlyortokenerror") ||
combined.contains("tokenerror") ||
combined.contains("request has expired") ||
combined.contains("refresh page") ||
combined.contains("请求已经过期")
}
private static let planNameKeys = [
"planName",
"plan_name",
"packageName",
"package_name",
"commodityName",
"commodity_name",
"instanceName",
"instance_name",
"displayName",
"display_name",
"ProductName",
"productName",
"name",
"title",
"planType",
"plan_type",
]
private static let usedQuotaKeys = [
"usedQuota",
"used_quota",
"usedCredits",
"usedCredit",
"consumedCredits",
"usage",
"used",
"usedAmount",
"consumeAmount",
"usedValue",
"UsedValue",
"consumedValue",
"ConsumedValue",
]
private static let totalQuotaKeys = [
"totalQuota",
"total_quota",
"totalCredits",
"totalCredit",
"quota",
"creditLimit",
"creditsTotal",
"monthlyTotalQuota",
"amount",
"totalValue",
"TotalValue",
]
private static let remainingQuotaKeys = [
"remainingQuota",
"remainQuota",
"remainingCredits",
"remainingCredit",
"availableCredits",
"balance",
"remaining",
"availableAmount",
"remainAmount",
"totalSurplusValue",
"TotalSurplusValue",
"surplusValue",
"SurplusValue",
]
private static let subscriptionCountKeys = [
"totalCount",
"TotalCount",
"subscriptionTotalNumber",
"SubscriptionTotalNumber",
]
private static let resetDateKeys = [
"nextRefreshTime",
"resetTime",
"periodEndTime",
"billingCycleEnd",
"billCycleEndTime",
"expireTime",
"expirationTime",
"endTime",
"validEndTime",
"instanceEndTime",
"nearestExpireDate",
"NearestExpireDate",
]
private static func findSubscriptionSummary(in payload: [String: Any]) -> [String: Any]? {
if let data = self.findFirstDictionary(
forKeys: ["Data", "data", "successResponse", "success_response"],
in: payload),
self.containsSubscriptionSummaryFields(data)
{
return data
}
return self.findFirstDictionary(
matchingAnyKey: Self.usedQuotaKeys + Self.totalQuotaKeys + Self.remainingQuotaKeys +
Self.subscriptionCountKeys,
in: payload)
}
private static func containsSubscriptionSummaryFields(_ payload: [String: Any]) -> Bool {
let keys = self.usedQuotaKeys + self.totalQuotaKeys + self.remainingQuotaKeys + self.subscriptionCountKeys
return keys.contains { payload[$0] != nil }
}
private static func findPlanName(in payload: [String: Any]) -> String? {
self.anyString(for: self.planNameKeys, in: payload) ??
self.findFirstString(forKeys: self.planNameKeys, in: payload)
}
private static func findResetDate(in payload: [String: Any]) -> Date? {
self.anyDate(for: self.resetDateKeys, in: payload) ??
self.findFirstDate(forKeys: self.resetDateKeys, in: payload)
}
private static func payloadDiagnostics(payload: [String: Any]) -> String {
let topKeys = payload.keys.sorted()
let dataDict = self.findFirstDictionary(
forKeys: ["Data", "data", "successResponse", "success_response"],
in: payload)
let dataKeys = dataDict?.keys.sorted() ?? []
return "topKeys=\(topKeys.joined(separator: ",")) dataKeys=\(dataKeys.joined(separator: ","))"
}
private static func isLikelyLoginHTML(_ data: Data) -> Bool {
guard let text = String(data: data, encoding: .utf8)?.lowercased() else { return false }
return text.contains("<html") &&
(text.contains("login") || text.contains("sign in") || text.contains("signin"))
}
private static func findFirstDictionary(forKeys keys: [String], in value: Any) -> [String: Any]? {
if let dict = value as? [String: Any] {
for key in keys {
if let nested = dict[key] as? [String: Any] {
return nested
}
}
for nestedValue in dict.values {
if let nested = self.findFirstDictionary(forKeys: keys, in: nestedValue) {
return nested
}
}
return nil
}
if let array = value as? [Any] {
for item in array {
if let nested = self.findFirstDictionary(forKeys: keys, in: item) {
return nested
}
}
}
return nil
}
private static func findFirstDictionary(matchingAnyKey keys: [String], in value: Any) -> [String: Any]? {
if let dict = value as? [String: Any] {
if keys.contains(where: { dict[$0] != nil }) {
return dict
}
for nestedValue in dict.values {
if let nested = self.findFirstDictionary(matchingAnyKey: keys, in: nestedValue) {
return nested
}
}
return nil
}
if let array = value as? [Any] {
for item in array {
if let nested = self.findFirstDictionary(matchingAnyKey: keys, in: item) {
return nested
}
}
}
return nil
}
private static func findFirstString(forKeys keys: [String], in value: Any) -> String? {
if let dict = value as? [String: Any] {
for key in keys {
if let parsed = self.parseString(dict[key]) {
return parsed
}
}
for nestedValue in dict.values {
if let parsed = self.findFirstString(forKeys: keys, in: nestedValue) {
return parsed
}
}
return nil
}
if let array = value as? [Any] {
for item in array {
if let parsed = self.findFirstString(forKeys: keys, in: item) {
return parsed
}
}
}
return nil
}
private static func findBoolValues(forKeys keys: [String], in value: Any) -> [Bool] {
if let dict = value as? [String: Any] {
let directValues = keys.compactMap { self.parseBool(dict[$0]) }
let nestedValues = dict.values.flatMap { self.findBoolValues(forKeys: keys, in: $0) }
return directValues + nestedValues
}
if let array = value as? [Any] {
return array.flatMap { self.findBoolValues(forKeys: keys, in: $0) }
}
return []
}
private static func findFirstInt(forKeys keys: [String], in value: Any) -> Int? {
if let dict = value as? [String: Any] {
for key in keys {
if let parsed = self.parseInt(dict[key]) {
return parsed
}
}
for nestedValue in dict.values {
if let parsed = self.findFirstInt(forKeys: keys, in: nestedValue) {
return parsed
}
}
return nil
}
if let array = value as? [Any] {
for item in array {
if let parsed = self.findFirstInt(forKeys: keys, in: item) {
return parsed
}
}
}
return nil
}
private static func findFirstDate(forKeys keys: [String], in value: Any) -> Date? {
if let dict = value as? [String: Any] {
for key in keys {
if let parsed = self.parseDate(dict[key]) {
return parsed
}
}
for nestedValue in dict.values {
if let parsed = self.findFirstDate(forKeys: keys, in: nestedValue) {
return parsed
}
}
return nil
}
if let array = value as? [Any] {
for item in array {
if let parsed = self.findFirstDate(forKeys: keys, in: item) {
return parsed
}
}
}
return nil
}
private static func expandedJSON(_ value: Any) -> Any {
if let dict = value as? [String: Any] {
var expanded: [String: Any] = [:]
expanded.reserveCapacity(dict.count)
for (key, nested) in dict {
expanded[key] = self.expandedJSON(nested)
}
return expanded
}
if let array = value as? [Any] {
return array.map { self.expandedJSON($0) }
}
if let string = value as? String,
let data = string.data(using: .utf8),
let nested = try? JSONSerialization.jsonObject(with: data, options: []),
nested is [String: Any] || nested is [Any]
{
return self.expandedJSON(nested)
}
return value
}
private static func anyString(for keys: [String], in dict: [String: Any]) -> String? {
for key in keys {
if let value = self.parseString(dict[key]) {
return value
}
}
return nil
}
private static func anyDouble(for keys: [String], in dict: [String: Any]) -> Double? {
for key in keys {
if let value = self.parseDouble(dict[key]) {
return value
}
}
return nil
}
private static func anyDate(for keys: [String], in dict: [String: Any]) -> Date? {
for key in keys {
if let value = self.parseDate(dict[key]) {
return value
}
}
return nil
}
private static func anyBool(for keys: [String], in dict: [String: Any]) -> Bool? {
for key in keys {
if let value = self.parseBool(dict[key]) {
return value
}
}
return nil
}
private static func parseInt(_ raw: Any?) -> Int? {
if let value = raw as? Int { return value }
if let value = raw as? Int64 { return Int(value) }
if let value = raw as? Double { return Int(value) }
if let value = raw as? NSNumber { return value.intValue }
if let value = self.parseString(raw) { return Int(value) }
return nil
}
private static func parseDouble(_ raw: Any?) -> Double? {
if let value = raw as? Double { return value }
if let value = raw as? Int { return Double(value) }
if let value = raw as? Int64 { return Double(value) }
if let value = raw as? NSNumber { return value.doubleValue }
if let value = self.parseString(raw) {
let cleaned = value.replacingOccurrences(of: ",", with: "")
return Double(cleaned)
}
return nil
}
private static func parseString(_ raw: Any?) -> String? {
guard let value = raw as? String else { return nil }
let trimmed = value.trimmingCharacters(in: .whitespacesAndNewlines)
return trimmed.isEmpty ? nil : trimmed
}
private static func parseDate(_ raw: Any?) -> Date? {
if let intValue = self.parseInt(raw) {
if intValue > 1_000_000_000_000 {
return Date(timeIntervalSince1970: TimeInterval(intValue) / 1000)
}
if intValue > 1_000_000_000 {
return Date(timeIntervalSince1970: TimeInterval(intValue))
}
}
if let string = self.parseString(raw) {
let formatter = ISO8601DateFormatter()
if let date = formatter.date(from: string) {
return date
}
let dateFormatter = DateFormatter()
dateFormatter.locale = Locale(identifier: "en_US_POSIX")
for format in ["yyyy-MM-dd", "yyyy-MM-dd HH:mm", "yyyy-MM-dd HH:mm:ss"] {
dateFormatter.dateFormat = format
if let date = dateFormatter.date(from: string) {
return date
}
}
}
return nil
}
private static func parseBool(_ raw: Any?) -> Bool? {
if let value = raw as? Bool { return value }
if let number = raw as? NSNumber { return number.boolValue }
guard let string = self.parseString(raw)?.lowercased() else { return nil }
switch string {
case "true", "1", "yes", "active", "valid", "normal":
return true
case "false", "0", "no", "inactive", "invalid", "expired":
return false
default:
return nil
}
}
private static func extractCookieValue(name: String, from cookieHeader: String) -> String? {
cookieHeader
.split(separator: ";")
.compactMap { part -> (String, String)? in
let pieces = part.split(separator: "=", maxSplits: 1).map {
$0.trimmingCharacters(in: .whitespacesAndNewlines)
}
guard pieces.count == 2 else { return nil }
return (pieces[0], pieces[1])
}
.first { $0.0 == name }?
.1
}
private static func hasCSRF(in cookieHeader: String) -> Bool {
self.extractCookieValue(name: "login_aliyunid_csrf", from: cookieHeader) != nil ||
self.extractCookieValue(name: "csrf", from: cookieHeader) != nil
}
static func cookieNames(from cookieHeader: String) -> [String] {
CookieHeaderNormalizer.pairs(from: cookieHeader)
.map(\.name)
.filter { !$0.isEmpty }
.uniquedSorted()
}
static func cookieNamesDescription(from cookieHeader: String) -> String {
let names = self.cookieNames(from: cookieHeader)
return names.isEmpty ? "none" : names.joined(separator: ",")
}
private static func extractSECToken(from html: String) -> String? {
let patterns = [
#""secToken"\s*:\s*"([^"]+)""#,
#""sec_token"\s*:\s*"([^"]+)""#,
#"secToken['"]?\s*[:=]\s*['"]([^'"]+)['"]"#,
#"sec_token['"]?\s*[:=]\s*['"]([^'"]+)['"]"#,
]
for pattern in patterns {
if let token = self.matchFirstGroup(pattern: pattern, in: html), !token.isEmpty {
return token
}
}
return nil
}
private static func matchFirstGroup(pattern: String, in text: String) -> String? {
guard let regex = try? NSRegularExpression(pattern: pattern, options: []) else {
return nil
}
let range = NSRange(text.startIndex..<text.endIndex, in: text)
guard let match = regex.firstMatch(in: text, options: [], range: range),
match.numberOfRanges > 1,
let valueRange = Range(match.range(at: 1), in: text)
else {
return nil
}
let value = text[valueRange].trimmingCharacters(in: .whitespacesAndNewlines)
return value.isEmpty ? nil : String(value)
}
}
extension [String] {
fileprivate func uniquedSorted() -> [String] {
Array(Set(self)).sorted()
}
}
@@ -0,0 +1,96 @@
import Foundation
public struct AlibabaTokenPlanUsageSnapshot: Sendable {
public let planName: String?
public let usedQuota: Double?
public let totalQuota: Double?
public let remainingQuota: Double?
public let resetsAt: Date?
public let updatedAt: Date
public init(
planName: String?,
usedQuota: Double?,
totalQuota: Double?,
remainingQuota: Double?,
resetsAt: Date?,
updatedAt: Date)
{
self.planName = planName
self.usedQuota = usedQuota
self.totalQuota = totalQuota
self.remainingQuota = remainingQuota
self.resetsAt = resetsAt
self.updatedAt = updatedAt
}
}
extension AlibabaTokenPlanUsageSnapshot {
public func toUsageSnapshot() -> UsageSnapshot {
let primary: RateWindow? = Self.usedPercent(
used: self.usedQuota,
total: self.totalQuota,
remaining: self.remainingQuota).map {
RateWindow(
usedPercent: $0,
windowMinutes: 30 * 24 * 60,
resetsAt: self.resetsAt,
resetDescription: Self.quotaDetail(
used: self.usedQuota,
total: self.totalQuota,
remaining: self.remainingQuota))
}
let planName = self.planName?.trimmingCharacters(in: .whitespacesAndNewlines)
let loginMethod = (planName?.isEmpty ?? true) ? nil : planName
let identity = ProviderIdentitySnapshot(
providerID: .alibabatokenplan,
accountEmail: nil,
accountOrganization: nil,
loginMethod: loginMethod)
return UsageSnapshot(
primary: primary,
secondary: nil,
tertiary: nil,
providerCost: nil,
updatedAt: self.updatedAt,
identity: identity)
}
private static func usedPercent(used: Double?, total: Double?, remaining: Double?) -> Double? {
guard let total, total > 0 else { return nil }
let usedValue: Double? = if let used {
used
} else if let remaining {
total - remaining
} else {
nil
}
guard let usedValue else { return nil }
let normalizedUsed = max(0, min(usedValue, total))
return normalizedUsed / total * 100
}
private static func quotaDetail(used: Double?, total: Double?, remaining: Double?) -> String? {
if let used, let total, total > 0 {
return "\(self.format(used)) / \(self.format(total)) credits used"
}
if let remaining, let total, total > 0 {
return "\(Self.format(remaining)) / \(Self.format(total)) credits left"
}
if let remaining {
return "\(Self.format(remaining)) credits left"
}
return nil
}
private static func format(_ value: Double) -> String {
let formatter = NumberFormatter()
formatter.numberStyle = .decimal
formatter.usesGroupingSeparator = true
formatter.maximumFractionDigits = value.rounded() == value ? 0 : 2
formatter.locale = Locale(identifier: "en_US_POSIX")
return formatter.string(from: NSNumber(value: value)) ?? String(format: "%.2f", value)
}
}
@@ -0,0 +1,46 @@
import Foundation
public struct AmpCLIProbe: Sendable {
private static let commandTimeout: TimeInterval = 15
private let arguments: [String]
public init() {
self.arguments = ["usage"]
}
init(arguments: [String]) {
self.arguments = arguments
}
public func fetch(
environment: [String: String] = ProcessInfo.processInfo.environment,
now: Date = Date()) async throws -> AmpUsageSnapshot
{
let loginPATH = LoginShellPathCache.shared.current
guard let executable = BinaryLocator.resolveAmpBinary(env: environment, loginPATH: loginPATH) else {
throw SubprocessRunnerError.binaryNotFound("amp")
}
var commandEnvironment = environment
commandEnvironment["NO_COLOR"] = "1"
commandEnvironment["PATH"] = PathBuilder.effectivePATH(
purposes: [.tty, .nodeTooling],
env: environment,
loginPATH: loginPATH)
let result = try await SubprocessRunner.run(
binary: executable,
arguments: self.arguments,
environment: commandEnvironment,
timeout: Self.commandTimeout,
standardInput: FileHandle.nullDevice,
label: "amp-usage")
let output = result.stdout.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty
? result.stderr
: result.stdout
guard !output.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty else {
throw AmpUsageError.parseFailed("The Amp CLI returned no usage data.")
}
return try AmpUsageParser.parse(displayText: output, now: now)
}
}
@@ -0,0 +1,169 @@
import Foundation
public enum AmpProviderDescriptor {
public static let descriptor: ProviderDescriptor = Self.makeDescriptor()
static func makeDescriptor() -> ProviderDescriptor {
ProviderDescriptor(
id: .amp,
metadata: ProviderMetadata(
id: .amp,
displayName: "Amp",
sessionLabel: "Amp Free",
weeklyLabel: "Balance",
opusLabel: nil,
supportsOpus: false,
supportsCredits: true,
creditsHint: "Individual and workspace credit balances from Amp.",
toggleTitle: "Show Amp usage",
cliName: "amp",
defaultEnabled: false,
isPrimaryProvider: false,
usesAccountFallback: false,
browserCookieOrder: ProviderBrowserCookieDefaults.defaultImportOrder,
dashboardURL: "https://ampcode.com/settings/usage",
statusPageURL: nil),
branding: ProviderBranding(
iconStyle: .amp,
iconResourceName: "ProviderIcon-amp",
color: ProviderColor(red: 220 / 255, green: 38 / 255, blue: 38 / 255)),
tokenCost: ProviderTokenCostConfig(
supportsTokenCost: false,
noDataMessage: { "Amp cost summary is not supported." }),
fetchPlan: ProviderFetchPlan(
sourceModes: [.auto, .api, .web, .cli],
pipeline: ProviderFetchPipeline(resolveStrategies: self.resolveStrategies)),
cli: ProviderCLIConfig(
name: "amp",
versionDetector: nil))
}
private static func resolveStrategies(context: ProviderFetchContext) async -> [any ProviderFetchStrategy] {
switch context.sourceMode {
case .auto:
[AmpCLIFetchStrategy(), AmpAPIFetchStrategy(), AmpStatusFetchStrategy()]
case .cli:
[AmpCLIFetchStrategy()]
case .api:
[AmpAPIFetchStrategy()]
case .web:
[AmpStatusFetchStrategy()]
case .oauth:
[]
}
}
}
struct AmpCLIFetchStrategy: ProviderFetchStrategy {
let id: String = "amp.cli"
let kind: ProviderFetchKind = .cli
func isAvailable(_ context: ProviderFetchContext) async -> Bool {
BinaryLocator.resolveAmpBinary(
env: context.env,
loginPATH: LoginShellPathCache.shared.current) != nil
}
func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult {
let snapshot = try await AmpCLIProbe().fetch(environment: context.env)
return self.makeResult(
usage: snapshot.toUsageSnapshot(now: snapshot.updatedAt),
sourceLabel: "cli")
}
func shouldFallback(on error: Error, context: ProviderFetchContext) -> Bool {
if error is CancellationError || (error as? URLError)?.code == .cancelled {
return false
}
return context.sourceMode == .auto
}
}
struct AmpAPIFetchStrategy: ProviderFetchStrategy {
let id: String = "amp.api"
let kind: ProviderFetchKind = .apiToken
func isAvailable(_ context: ProviderFetchContext) async -> Bool {
_ = context
return true
}
func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult {
guard let token = ProviderTokenResolver.ampToken(environment: context.env) else {
throw AmpUsageError.missingAPIToken
}
let logger: ((String) -> Void)? = context.verbose
? { msg in CodexBarLog.logger(LogCategories.amp).verbose(msg) }
: nil
let snapshot = try await AmpUsageFetcher(browserDetection: context.browserDetection)
.fetch(apiToken: token, logger: logger)
return self.makeResult(
usage: snapshot.toUsageSnapshot(now: snapshot.updatedAt),
sourceLabel: "api")
}
func shouldFallback(on error: Error, context: ProviderFetchContext) -> Bool {
guard context.sourceMode == .auto else { return false }
return !(error is CancellationError) && (error as? URLError)?.code != .cancelled
}
}
struct AmpStatusFetchStrategy: ProviderFetchStrategy {
let id: String = "amp.web"
let kind: ProviderFetchKind = .web
func isAvailable(_ context: ProviderFetchContext) async -> Bool {
guard context.sourceMode.usesWeb else { return false }
#if os(macOS)
let canImportBrowserCookies = true
#else
let canImportBrowserCookies = false
#endif
return Self.canUseWebFallback(
settings: context.settings?.amp,
canImportBrowserCookies: canImportBrowserCookies)
}
func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult {
let fetcher = AmpUsageFetcher(browserDetection: context.browserDetection)
let manual = Self.manualCookieHeader(from: context)
let logger: ((String) -> Void)? = context.verbose
? { msg in CodexBarLog.logger(LogCategories.amp).verbose(msg) }
: nil
let snap = try await fetcher.fetch(cookieHeaderOverride: manual, logger: logger)
return self.makeResult(
usage: snap.toUsageSnapshot(now: snap.updatedAt),
sourceLabel: "web")
}
func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool {
false
}
static func canUseWebFallback(
settings: ProviderSettingsSnapshot.AmpProviderSettings?,
canImportBrowserCookies: Bool) -> Bool
{
guard let settings else { return canImportBrowserCookies }
switch settings.cookieSource {
case .auto:
return canImportBrowserCookies
case .manual:
return Self.manualCookieHeader(from: settings.manualCookieHeader) != nil
case .off:
return false
}
}
private static func manualCookieHeader(from context: ProviderFetchContext) -> String? {
guard let settings = context.settings?.amp, settings.cookieSource == .manual else { return nil }
return Self.manualCookieHeader(from: settings.manualCookieHeader)
}
private static func manualCookieHeader(from rawHeader: String?) -> String? {
guard let header = CookieHeaderNormalizer.normalize(rawHeader),
CookieHeaderNormalizer.pairs(from: header).contains(where: { $0.name == "session" })
else { return nil }
return header
}
}
@@ -0,0 +1,22 @@
import Foundation
public enum AmpSettingsReader {
public static let apiTokenKey = "AMP_API_KEY"
public static func apiToken(environment: [String: String] = ProcessInfo.processInfo.environment) -> String? {
self.cleaned(environment[self.apiTokenKey])
}
static func cleaned(_ raw: String?) -> String? {
guard var value = raw?.trimmingCharacters(in: .whitespacesAndNewlines), !value.isEmpty else {
return nil
}
if (value.hasPrefix("\"") && value.hasSuffix("\"")) ||
(value.hasPrefix("'") && value.hasSuffix("'"))
{
value = String(value.dropFirst().dropLast())
}
value = value.trimmingCharacters(in: .whitespacesAndNewlines)
return value.isEmpty ? nil : value
}
}
@@ -0,0 +1,491 @@
import Foundation
#if canImport(FoundationNetworking)
import FoundationNetworking
#endif
#if os(macOS)
import SweetCookieKit
#endif
public enum AmpUsageError: LocalizedError, Sendable {
case notLoggedIn
case invalidCredentials
case missingAPIToken
case invalidAPIToken
case parseFailed(String)
case networkError(String)
case noSessionCookie
public var errorDescription: String? {
switch self {
case .notLoggedIn:
"Not logged in to Amp. Please log in via ampcode.com."
case .invalidCredentials:
"Amp session cookie expired. Please log in again."
case .missingAPIToken:
"Amp access token not configured. Set AMP_API_KEY or add it in Settings."
case .invalidAPIToken:
"Amp access token is invalid or expired."
case let .parseFailed(message):
"Could not parse Amp usage: \(message)"
case let .networkError(message):
"Amp request failed: \(message)"
case .noSessionCookie:
"No Amp session cookie found. Please log in to ampcode.com in your browser."
}
}
}
#if os(macOS)
private let ampCookieImportOrder: BrowserCookieImportOrder =
ProviderDefaults.metadata[.amp]?.browserCookieOrder ?? Browser.defaultImportOrder
public enum AmpCookieImporter {
private static let cookieClient = BrowserCookieClient()
private static let cookieDomains = ["ampcode.com", "www.ampcode.com"]
private static let sessionCookieNames: Set<String> = [
"session",
]
public struct SessionInfo: Sendable {
public let cookies: [HTTPCookie]
public let sourceLabel: String
public init(cookies: [HTTPCookie], sourceLabel: String) {
self.cookies = cookies
self.sourceLabel = sourceLabel
}
public var cookieHeader: String {
self.cookies.map { "\($0.name)=\($0.value)" }.joined(separator: "; ")
}
}
public static func importSession(
browserDetection: BrowserDetection,
logger: ((String) -> Void)? = nil) throws -> SessionInfo
{
let log: (String) -> Void = { msg in logger?("[amp-cookie] \(msg)") }
let installed = ampCookieImportOrder.cookieImportCandidates(using: browserDetection)
for browserSource in installed {
do {
let query = BrowserCookieQuery(domains: self.cookieDomains)
let sources = try Self.cookieClient.codexBarRecords(
matching: query,
in: browserSource,
logger: log)
for source in sources where !source.records.isEmpty {
let cookies = BrowserCookieClient.makeHTTPCookies(source.records, origin: query.origin)
guard !cookies.isEmpty else { continue }
let names = cookies.map(\.name).joined(separator: ", ")
log("\(source.label) cookies: \(names)")
let sessionCookies = cookies.filter { Self.sessionCookieNames.contains($0.name) }
if !sessionCookies.isEmpty {
log("Found Amp session cookie in \(source.label)")
return SessionInfo(cookies: sessionCookies, sourceLabel: source.label)
}
log("\(source.label) cookies found, but no Amp session cookie present")
log("Expected one of: \(Self.sessionCookieNames.joined(separator: ", "))")
}
} catch {
BrowserCookieAccessGate.recordIfNeeded(error)
log("\(browserSource.displayName) cookie import failed: \(error.localizedDescription)")
}
}
throw AmpUsageError.noSessionCookie
}
}
#endif
public struct AmpUsageFetcher: Sendable {
private static let settingsURL = URL(string: "https://ampcode.com/settings")!
static let usageURL = URL(string: "https://ampcode.com/api/internal?userDisplayBalanceInfo")!
@MainActor private static var recentDumps: [String] = []
public let browserDetection: BrowserDetection
public init(browserDetection: BrowserDetection) {
self.browserDetection = browserDetection
}
public func fetch(
cookieHeaderOverride: String? = nil,
logger: ((String) -> Void)? = nil,
now: Date = Date()) async throws -> AmpUsageSnapshot
{
let log: (String) -> Void = { msg in logger?("[amp] \(msg)") }
let cookieHeader = try await self.resolveCookieHeader(override: cookieHeaderOverride, logger: log)
if let logger {
let names = self.cookieNames(from: cookieHeader)
if !names.isEmpty {
logger("[amp] Cookie names: \(names.joined(separator: ", "))")
}
let diagnostics = RedirectDiagnostics(cookieHeader: cookieHeader, logger: logger)
do {
let (html, responseInfo) = try await self.fetchLegacyHTMLWithDiagnostics(
cookieHeader: cookieHeader,
diagnostics: diagnostics)
self.logDiagnostics(responseInfo: responseInfo, diagnostics: diagnostics, logger: logger)
return try AmpUsageParser.parse(html: html, now: now)
} catch {
self.logDiagnostics(responseInfo: nil, diagnostics: diagnostics, logger: logger)
logger("[amp] Fetch failed: \(error.localizedDescription)")
throw error
}
}
let diagnostics = RedirectDiagnostics(cookieHeader: cookieHeader, logger: nil)
let (html, _) = try await self.fetchLegacyHTMLWithDiagnostics(
cookieHeader: cookieHeader,
diagnostics: diagnostics)
return try AmpUsageParser.parse(html: html, now: now)
}
public func fetch(
apiToken: String,
logger: ((String) -> Void)? = nil,
now: Date = Date()) async throws -> AmpUsageSnapshot
{
guard let token = AmpSettingsReader.cleaned(apiToken) else {
throw AmpUsageError.missingAPIToken
}
let request = try Self.makeUsageAPIRequest(apiToken: token)
let diagnostics = APIRedirectDiagnostics(logger: logger)
let session = URLSession(configuration: .ephemeral, delegate: diagnostics, delegateQueue: nil)
let httpResponse = try await session.response(for: request)
logger?("[amp] API response: \(httpResponse.statusCode) " +
"\(httpResponse.response.url?.absoluteString ?? "unknown")")
try Self.validateAPIResponse(httpResponse)
return try Self.parseUsageAPIResponse(httpResponse.data, now: now)
}
public func debugRawProbe(cookieHeaderOverride: String? = nil) async -> String {
let stamp = ISO8601DateFormatter().string(from: Date())
var lines: [String] = []
lines.append("=== Amp Debug Probe @ \(stamp) ===")
lines.append("")
do {
let cookieHeader = try await self.resolveCookieHeader(
override: cookieHeaderOverride,
logger: { msg in lines.append("[cookie] \(msg)") })
let diagnostics = RedirectDiagnostics(cookieHeader: cookieHeader, logger: nil)
let cookieNames = CookieHeaderNormalizer.pairs(from: cookieHeader).map(\.name)
lines.append("Cookie names: \(cookieNames.joined(separator: ", "))")
let (html, responseInfo) = try await self.fetchLegacyHTMLWithDiagnostics(
cookieHeader: cookieHeader,
diagnostics: diagnostics)
let snapshot = try AmpUsageParser.parse(html: html)
lines.append("")
lines.append("Fetch Success")
lines.append("Status: \(responseInfo.statusCode) \(responseInfo.url)")
if !diagnostics.redirects.isEmpty {
lines.append("")
lines.append("Redirects:")
for entry in diagnostics.redirects {
lines.append(" \(entry)")
}
}
lines.append("")
lines.append("Amp Free:")
lines.append(" quota=\(snapshot.freeQuota?.description ?? "nil")")
lines.append(" used=\(snapshot.freeUsed?.description ?? "nil")")
lines.append(" hourlyReplenishment=\(snapshot.hourlyReplenishment?.description ?? "nil")")
lines.append(" windowHours=\(snapshot.windowHours?.description ?? "nil")")
lines.append(" individualCredits=\(snapshot.individualCredits?.description ?? "nil")")
for workspace in snapshot.workspaceBalances {
lines.append(" workspace[\(workspace.name)]=\(workspace.remaining)")
}
let output = lines.joined(separator: "\n")
Task { @MainActor in Self.recordDump(output) }
return output
} catch {
lines.append("")
lines.append("Probe Failed: \(error.localizedDescription)")
let output = lines.joined(separator: "\n")
Task { @MainActor in Self.recordDump(output) }
return output
}
}
public static func latestDumps() async -> String {
await MainActor.run {
let result = Self.recentDumps.joined(separator: "\n\n---\n\n")
return result.isEmpty ? "No Amp probe dumps captured yet." : result
}
}
private func resolveCookieHeader(
override: String?,
logger: ((String) -> Void)?) async throws -> String
{
if let override = CookieHeaderNormalizer.normalize(override) {
if let sessionHeader = self.sessionCookieHeader(from: override) {
logger?("[amp] Using manual session cookie")
return sessionHeader
}
throw AmpUsageError.noSessionCookie
}
#if os(macOS)
let session = try AmpCookieImporter.importSession(browserDetection: self.browserDetection, logger: logger)
logger?("[amp] Using cookies from \(session.sourceLabel)")
return session.cookieHeader
#else
throw AmpUsageError.noSessionCookie
#endif
}
static func makeUsageAPIRequest(apiToken: String) throws -> URLRequest {
var request = URLRequest(url: Self.usageURL)
request.httpMethod = "POST"
request.httpBody = try JSONSerialization.data(withJSONObject: [
"method": "userDisplayBalanceInfo",
"params": [:],
])
request.setValue("Bearer \(apiToken)", forHTTPHeaderField: "Authorization")
request.setValue("application/json", forHTTPHeaderField: "accept")
request.setValue("application/json", forHTTPHeaderField: "content-type")
return request
}
private func fetchLegacyHTMLWithDiagnostics(
cookieHeader: String,
diagnostics: RedirectDiagnostics) async throws -> (String, ResponseInfo)
{
var request = URLRequest(url: Self.settingsURL)
request.httpMethod = "GET"
request.setValue(cookieHeader, forHTTPHeaderField: "Cookie")
request.setValue(
"text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8",
forHTTPHeaderField: "accept")
Self.applyBrowserHeaders(to: &request)
let session = URLSession(configuration: .ephemeral, delegate: diagnostics, delegateQueue: nil)
let httpResponse = try await session.response(for: request)
let responseInfo = ResponseInfo(
statusCode: httpResponse.statusCode,
url: httpResponse.response.url?.absoluteString ?? "unknown")
try Self.validateBrowserResponse(response: httpResponse, diagnostics: diagnostics)
let html = String(data: httpResponse.data, encoding: .utf8) ?? ""
return (html, responseInfo)
}
static func parseUsageAPIResponse(_ data: Data, now: Date = Date()) throws -> AmpUsageSnapshot {
let response: UsageAPIResponse
do {
response = try JSONDecoder().decode(UsageAPIResponse.self, from: data)
} catch {
throw AmpUsageError.parseFailed("Invalid Amp usage API response.")
}
guard response.ok else {
if response.error?.code == "auth-required" {
throw AmpUsageError.invalidAPIToken
}
throw AmpUsageError.networkError(response.error?.message ?? "Amp usage API returned an error.")
}
guard let displayText = response.result?.displayText, !displayText.isEmpty else {
throw AmpUsageError.parseFailed("Missing Amp usage display text.")
}
return try AmpUsageParser.parse(displayText: displayText, now: now)
}
private static func applyBrowserHeaders(to request: inout URLRequest) {
request.setValue(
"Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) " +
"AppleWebKit/537.36 (KHTML, like Gecko) Chrome/143.0.0.0 Safari/537.36",
forHTTPHeaderField: "user-agent")
request.setValue("en-US,en;q=0.9", forHTTPHeaderField: "accept-language")
request.setValue("https://ampcode.com", forHTTPHeaderField: "origin")
request.setValue(self.settingsURL.absoluteString, forHTTPHeaderField: "referer")
}
private static func validateBrowserResponse(
response: ProviderHTTPResponse,
diagnostics: RedirectDiagnostics) throws
{
guard response.statusCode == 200 else {
if response.statusCode == 401 || response.statusCode == 403 || diagnostics.detectedLoginRedirect {
throw AmpUsageError.invalidCredentials
}
throw AmpUsageError.networkError("HTTP \(response.statusCode)")
}
}
private static func validateAPIResponse(_ response: ProviderHTTPResponse) throws {
guard response.statusCode == 200 else {
if response.statusCode == 401 || response.statusCode == 403 {
throw AmpUsageError.invalidAPIToken
}
throw AmpUsageError.networkError("HTTP \(response.statusCode)")
}
}
@MainActor private static func recordDump(_ text: String) {
if self.recentDumps.count >= 5 { self.recentDumps.removeFirst() }
self.recentDumps.append(text)
}
private final class RedirectDiagnostics: NSObject, URLSessionTaskDelegate, @unchecked Sendable {
private let cookieHeader: String
private let logger: ((String) -> Void)?
var redirects: [String] = []
private(set) var detectedLoginRedirect = false
init(cookieHeader: String, logger: ((String) -> Void)?) {
self.cookieHeader = cookieHeader
self.logger = logger
}
func urlSession(
_: URLSession,
task: URLSessionTask,
willPerformHTTPRedirection response: HTTPURLResponse,
newRequest request: URLRequest,
completionHandler: @escaping (URLRequest?) -> Void)
{
let from = response.url?.absoluteString ?? "unknown"
let to = request.url?.absoluteString ?? "unknown"
self.redirects.append("\(response.statusCode) \(from) -> \(to)")
if let toURL = request.url, AmpUsageFetcher.isLoginRedirect(toURL) {
if let logger {
logger("[amp] Detected login redirect, aborting (invalid session)")
}
self.detectedLoginRedirect = true
completionHandler(nil)
return
}
var updated = request
if AmpUsageFetcher.shouldAttachCookie(to: request.url), !self.cookieHeader.isEmpty {
updated.setValue(self.cookieHeader, forHTTPHeaderField: "Cookie")
} else {
updated.setValue(nil, forHTTPHeaderField: "Cookie")
}
if let referer = response.url?.absoluteString {
updated.setValue(referer, forHTTPHeaderField: "referer")
}
if let logger {
logger("[amp] Redirect \(response.statusCode) \(from) -> \(to)")
}
completionHandler(updated)
}
}
/// Amp's balance RPC should not redirect. Refusing redirects guarantees the bearer token cannot cross hosts.
private final class APIRedirectDiagnostics: NSObject, URLSessionTaskDelegate, @unchecked Sendable {
private let logger: ((String) -> Void)?
init(logger: ((String) -> Void)?) {
self.logger = logger
}
func urlSession(
_: URLSession,
task _: URLSessionTask,
willPerformHTTPRedirection response: HTTPURLResponse,
newRequest request: URLRequest,
completionHandler: @escaping (URLRequest?) -> Void)
{
let from = response.url?.absoluteString ?? "unknown"
let to = request.url?.absoluteString ?? "unknown"
self.logger?("[amp] API redirect blocked: \(response.statusCode) \(from) -> \(to)")
completionHandler(nil)
}
}
private struct ResponseInfo {
let statusCode: Int
let url: String
}
private struct UsageAPIResponse: Decodable {
let ok: Bool
let result: Result?
let error: APIError?
struct Result: Decodable {
let displayText: String
}
struct APIError: Decodable {
let code: String?
let message: String?
}
}
private func logDiagnostics(
responseInfo: ResponseInfo?,
diagnostics: RedirectDiagnostics,
logger: (String) -> Void)
{
if let responseInfo {
logger("[amp] Response: \(responseInfo.statusCode) \(responseInfo.url)")
}
if !diagnostics.redirects.isEmpty {
logger("[amp] Redirects:")
for entry in diagnostics.redirects {
logger("[amp] \(entry)")
}
}
}
private func cookieNames(from header: String) -> [String] {
header.split(separator: ";").compactMap { part in
let trimmed = part.trimmingCharacters(in: .whitespacesAndNewlines)
guard let idx = trimmed.firstIndex(of: "=") else { return nil }
let name = trimmed[..<idx].trimmingCharacters(in: .whitespacesAndNewlines)
return name.isEmpty ? nil : String(name)
}
}
private func sessionCookieHeader(from header: String) -> String? {
let pairs = CookieHeaderNormalizer.pairs(from: header)
let sessionPairs = pairs.filter { $0.name == "session" }
guard !sessionPairs.isEmpty else { return nil }
return sessionPairs.map { "\($0.name)=\($0.value)" }.joined(separator: "; ")
}
static func shouldAttachCookie(to url: URL?) -> Bool {
guard url?.scheme?.lowercased() == "https" else { return false }
return self.isAmpHost(url)
}
private static func isAmpHost(_ url: URL?) -> Bool {
guard let host = url?.host?.lowercased() else { return false }
if host == "ampcode.com" || host == "www.ampcode.com" { return true }
return host.hasSuffix(".ampcode.com")
}
static func isLoginRedirect(_ url: URL) -> Bool {
guard self.isAmpHost(url) else { return false }
if url.host?.lowercased() == "auth.ampcode.com" { return true }
let path = url.path.lowercased()
let components = path.split(separator: "/").map(String.init)
if components.contains("login") { return true }
if components.contains("signin") { return true }
if components.contains("sign-in") { return true }
// Amp currently redirects to /auth/sign-in?returnTo=... when session is invalid. Keep this slightly broader
// than one exact path so we keep working if Amp changes auth routes.
if components.contains("auth") {
let query = url.query?.lowercased() ?? ""
if query.contains("returnto=") { return true }
if query.contains("redirect=") { return true }
if query.contains("redirectto=") { return true }
}
return false
}
}
@@ -0,0 +1,221 @@
import Foundation
enum AmpUsageParser {
static func parse(html: String, now: Date = Date()) throws -> AmpUsageSnapshot {
guard let usage = self.parseFreeTierUsage(html) else {
if self.looksSignedOut(html) {
throw AmpUsageError.notLoggedIn
}
throw AmpUsageError.parseFailed("Missing Amp Free usage data.")
}
return AmpUsageSnapshot(
freeQuota: usage.quota,
freeUsed: usage.used,
hourlyReplenishment: usage.hourlyReplenishment,
windowHours: usage.windowHours,
updatedAt: now)
}
static func parse(displayText: String, now: Date = Date()) throws -> AmpUsageSnapshot {
let text = TextParsing.stripANSICodes(displayText)
let identityPattern = #"(?im)^\s*Signed in as\s+([^\s(]+)(?:\s+\(([^\r\n)]+)\))?\s*$"#
let identity = self.captures(in: text, pattern: identityPattern)
if identity == nil, self.looksSignedOut(text) {
throw AmpUsageError.notLoggedIn
}
let amountPattern = #"([0-9][0-9,]*(?:\.[0-9]+)?)"#
let freePattern = #"(?im)^\s*Amp Free:\s*\$?"# + amountPattern +
#"\s*/\s*\$?"# + amountPattern +
#"\s+remaining(?:\s*\(replenishes\s*\+\$?"# + amountPattern + #"\s*/\s*hour\))?"#
let freePercentPattern = #"(?im)^\s*Amp Free:\s*"# + amountPattern +
#"\s*%\s+remaining(?:\s+today)?(?:\s*\(resets\s+daily\))?"#
let creditsPattern = #"(?im)^\s*Individual credits:\s*\$?"# + amountPattern + #"\s+remaining"#
let individualCredits = self.captures(in: text, pattern: creditsPattern)?.first
.flatMap(self.number(from:))
let workspacePattern = #"(?im)^\s*Workspace\s+(.+?):\s*\$?"# + amountPattern + #"\s+remaining"#
let workspaceBalances: [AmpWorkspaceBalance] = self.allCaptures(
in: text,
pattern: workspacePattern).compactMap { captures -> AmpWorkspaceBalance? in
guard captures.count == 2,
let name = self.nonEmpty(captures[0]),
let remaining = self.number(from: captures[1])
else { return nil }
return AmpWorkspaceBalance(name: name, remaining: remaining)
}
let freeUsage: FreeTierUsage? = {
guard let free = self.captures(in: text, pattern: freePattern),
let remaining = self.number(from: free[0]),
let quota = self.number(from: free[1])
else { return nil }
let hourlyReplenishment = self.number(from: free[2]) ?? 0
let windowHours = hourlyReplenishment > 0
? max(1, (quota / hourlyReplenishment).rounded())
: nil
return FreeTierUsage(
quota: quota,
used: max(0, quota - remaining),
hourlyReplenishment: hourlyReplenishment,
windowHours: windowHours,
resetDescription: nil)
}()
let freePercentUsage: FreeTierUsage? = {
guard let remainingText = self.captures(in: text, pattern: freePercentPattern)?.first,
let remaining = self.number(from: remainingText)
else { return nil }
let clampedRemaining = min(100, max(0, remaining))
return FreeTierUsage(
quota: 100,
used: 100 - clampedRemaining,
hourlyReplenishment: 0,
windowHours: 24,
resetDescription: "resets daily")
}()
let resolvedFreeUsage = freeUsage ?? freePercentUsage
guard resolvedFreeUsage != nil || individualCredits != nil || !workspaceBalances.isEmpty else {
throw AmpUsageError.parseFailed("Missing Amp usage data.")
}
return AmpUsageSnapshot(
freeQuota: resolvedFreeUsage?.quota,
freeUsed: resolvedFreeUsage?.used,
hourlyReplenishment: resolvedFreeUsage?.hourlyReplenishment,
windowHours: resolvedFreeUsage?.windowHours,
individualCredits: individualCredits,
workspaceBalances: workspaceBalances,
accountEmail: self.nonEmpty(identity?[0]),
accountOrganization: self.nonEmpty(identity?[1]),
updatedAt: now,
freeResetDescription: resolvedFreeUsage?.resetDescription)
}
private struct FreeTierUsage {
let quota: Double
let used: Double
let hourlyReplenishment: Double
let windowHours: Double?
let resetDescription: String?
}
private static func parseFreeTierUsage(_ html: String) -> FreeTierUsage? {
let tokens = ["freeTierUsage", "getFreeTierUsage"]
for token in tokens {
if let object = self.extractObject(named: token, in: html),
let usage = self.parseFreeTierUsageObject(object)
{
return usage
}
}
return nil
}
private static func parseFreeTierUsageObject(_ object: String) -> FreeTierUsage? {
guard let quota = self.number(for: "quota", in: object),
let used = self.number(for: "used", in: object),
let hourly = self.number(for: "hourlyReplenishment", in: object)
else { return nil }
let windowHours = self.number(for: "windowHours", in: object)
return FreeTierUsage(
quota: quota,
used: used,
hourlyReplenishment: hourly,
windowHours: windowHours,
resetDescription: nil)
}
private static func extractObject(named token: String, in text: String) -> String? {
guard let tokenRange = text.range(of: token) else { return nil }
guard let braceIndex = text[tokenRange.upperBound...].firstIndex(of: "{") else { return nil }
var depth = 0
var inString = false
var isEscaped = false
var index = braceIndex
while index < text.endIndex {
let char = text[index]
if inString {
if isEscaped {
isEscaped = false
} else if char == "\\" {
isEscaped = true
} else if char == "\"" {
inString = false
}
} else {
if char == "\"" {
inString = true
} else if char == "{" {
depth += 1
} else if char == "}" {
depth -= 1
if depth == 0 {
return String(text[braceIndex...index])
}
}
}
index = text.index(after: index)
}
return nil
}
private static func number(for key: String, in text: String) -> Double? {
let pattern = "\\b\(NSRegularExpression.escapedPattern(for: key))\\b\\s*:\\s*([0-9]+(?:\\.[0-9]+)?)"
guard let regex = try? NSRegularExpression(pattern: pattern) else { return nil }
let range = NSRange(text.startIndex..<text.endIndex, in: text)
guard let match = regex.firstMatch(in: text, options: [], range: range),
match.numberOfRanges > 1,
let valueRange = Range(match.range(at: 1), in: text)
else { return nil }
return Double(text[valueRange])
}
private static func captures(in text: String, pattern: String) -> [String]? {
guard let regex = try? NSRegularExpression(pattern: pattern) else { return nil }
let range = NSRange(text.startIndex..<text.endIndex, in: text)
guard let match = regex.firstMatch(in: text, options: [], range: range) else { return nil }
return self.captures(in: text, match: match)
}
private static func allCaptures(in text: String, pattern: String) -> [[String]] {
guard let regex = try? NSRegularExpression(pattern: pattern) else { return [] }
let range = NSRange(text.startIndex..<text.endIndex, in: text)
return regex.matches(in: text, options: [], range: range).map { match in
self.captures(in: text, match: match)
}
}
private static func captures(in text: String, match: NSTextCheckingResult) -> [String] {
(1..<match.numberOfRanges).map { index in
let captureRange = match.range(at: index)
guard captureRange.location != NSNotFound,
let range = Range(captureRange, in: text)
else { return "" }
return String(text[range]).trimmingCharacters(in: .whitespacesAndNewlines)
}
}
private static func number(from text: String) -> Double? {
Double(text.replacingOccurrences(of: ",", with: ""))
}
private static func nonEmpty(_ text: String?) -> String? {
guard let text, !text.isEmpty else { return nil }
return text
}
private static func looksSignedOut(_ html: String) -> Bool {
let lower = html.lowercased()
if lower.contains("sign in") || lower.contains("log in") || lower.contains("login") {
return true
}
if lower.contains("/login") || lower.contains("ampcode.com/login") {
return true
}
return false
}
}
@@ -0,0 +1,109 @@
import Foundation
public struct AmpWorkspaceBalance: Codable, Equatable, Sendable {
public let name: String
public let remaining: Double
public init(name: String, remaining: Double) {
self.name = name
self.remaining = remaining
}
}
public struct AmpUsageDetails: Codable, Equatable, Sendable {
public let individualCredits: Double?
public let workspaceBalances: [AmpWorkspaceBalance]
public init(individualCredits: Double?, workspaceBalances: [AmpWorkspaceBalance]) {
self.individualCredits = individualCredits
self.workspaceBalances = workspaceBalances
}
}
public struct AmpUsageSnapshot: Sendable {
public let freeQuota: Double?
public let freeUsed: Double?
public let hourlyReplenishment: Double?
public let windowHours: Double?
public let individualCredits: Double?
public let workspaceBalances: [AmpWorkspaceBalance]
public let accountEmail: String?
public let accountOrganization: String?
public let updatedAt: Date
public let freeResetDescription: String?
public init(
freeQuota: Double?,
freeUsed: Double?,
hourlyReplenishment: Double?,
windowHours: Double?,
individualCredits: Double? = nil,
workspaceBalances: [AmpWorkspaceBalance] = [],
accountEmail: String? = nil,
accountOrganization: String? = nil,
updatedAt: Date,
freeResetDescription: String? = nil)
{
self.freeQuota = freeQuota
self.freeUsed = freeUsed
self.hourlyReplenishment = hourlyReplenishment
self.windowHours = windowHours
self.individualCredits = individualCredits
self.workspaceBalances = workspaceBalances
self.accountEmail = accountEmail
self.accountOrganization = accountOrganization
self.updatedAt = updatedAt
self.freeResetDescription = freeResetDescription
}
}
extension AmpUsageSnapshot {
public func toUsageSnapshot(now: Date = Date()) -> UsageSnapshot {
let primary: RateWindow? = if let freeQuota, let freeUsed {
{
let quota = max(0, freeQuota)
let used = max(0, freeUsed)
let percent = quota > 0 ? min(100, (used / quota) * 100) : 0
let windowMinutes: Int? = if let hours = self.windowHours, hours > 0 {
Int((hours * 60).rounded())
} else {
nil
}
let resetsAt: Date? = {
guard quota > 0, let hourlyReplenishment, hourlyReplenishment > 0 else { return nil }
return now.addingTimeInterval(max(0, used / hourlyReplenishment * 3600))
}()
return RateWindow(
usedPercent: percent,
windowMinutes: windowMinutes,
resetsAt: resetsAt,
resetDescription: self.freeResetDescription)
}()
} else {
nil
}
let identity = ProviderIdentitySnapshot(
providerID: .amp,
accountEmail: self.accountEmail,
accountOrganization: self.accountOrganization,
loginMethod: primary == nil ? "Amp" : "Amp Free")
let ampUsage: AmpUsageDetails? = if self.individualCredits != nil || !self.workspaceBalances.isEmpty {
AmpUsageDetails(
individualCredits: self.individualCredits,
workspaceBalances: self.workspaceBalances)
} else {
nil
}
return UsageSnapshot(
primary: primary,
secondary: nil,
tertiary: nil,
ampUsage: ampUsage,
providerCost: nil,
updatedAt: self.updatedAt,
identity: identity)
}
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,37 @@
import Foundation
extension AntigravityStatusProbe {
static func preferredLocalSnapshot(
_ snapshots: [AntigravityStatusSnapshot],
matchingAccountEmail expectedAccountEmail: String?) -> AntigravityStatusSnapshot?
{
var candidates = snapshots
if let expected = self.normalizedAccountEmail(expectedAccountEmail) {
let matches = snapshots.filter {
guard let found = self.normalizedAccountEmail($0.accountEmail) else { return false }
return found.caseInsensitiveCompare(expected) == .orderedSame
}
if !matches.isEmpty {
candidates = matches
}
}
var bestSnapshot: AntigravityStatusSnapshot?
var bestScore = Int.min
for snapshot in candidates {
let score = self.localSnapshotScore(snapshot)
if score > bestScore {
bestSnapshot = snapshot
bestScore = score
}
}
return bestSnapshot
}
private static func normalizedAccountEmail(_ email: String?) -> String? {
guard let trimmed = email?.trimmingCharacters(in: .whitespacesAndNewlines), !trimmed.isEmpty else {
return nil
}
return trimmed
}
}
@@ -0,0 +1,109 @@
import Foundation
#if canImport(FoundationNetworking)
import FoundationNetworking
#endif
enum LocalhostTrustPolicy {
static func shouldAcceptServerTrust(
host: String,
authenticationMethod: String,
hasServerTrust: Bool) -> Bool
{
#if !os(Linux)
guard authenticationMethod == NSURLAuthenticationMethodServerTrust else { return false }
#endif
let normalizedHost = host.lowercased()
guard normalizedHost == "127.0.0.1" || normalizedHost == "localhost" else { return false }
return hasServerTrust
}
}
final class LocalhostSessionDelegate: NSObject {
func data(for request: URLRequest, session: URLSession) async throws -> (Data, URLResponse) {
let state = LocalhostSessionTaskState()
return try await withTaskCancellationHandler {
try await withCheckedThrowingContinuation { continuation in
let task = session.dataTask(with: request) { data, response, error in
if let error {
continuation.resume(throwing: error)
return
}
guard let data, let response else {
continuation.resume(throwing: AntigravityStatusProbeError.apiError("Invalid response"))
return
}
continuation.resume(returning: (data, response))
}
state.setTask(task)
task.resume()
}
} onCancel: {
state.cancel()
}
}
private func challengeResult(_ challenge: URLAuthenticationChallenge) -> (
disposition: URLSession.AuthChallengeDisposition,
credential: URLCredential?)
{
#if os(Linux)
return (.performDefaultHandling, nil)
#else
let protectionSpace = challenge.protectionSpace
let trust = protectionSpace.serverTrust
guard LocalhostTrustPolicy.shouldAcceptServerTrust(
host: protectionSpace.host,
authenticationMethod: protectionSpace.authenticationMethod,
hasServerTrust: trust != nil),
let trust
else {
return (.performDefaultHandling, nil)
}
return (.useCredential, URLCredential(trust: trust))
#endif
}
}
extension LocalhostSessionDelegate: URLSessionDelegate {
func urlSession(
_ session: URLSession,
didReceive challenge: URLAuthenticationChallenge) async -> (URLSession.AuthChallengeDisposition, URLCredential?)
{
self.challengeResult(challenge)
}
}
extension LocalhostSessionDelegate: URLSessionTaskDelegate {
func urlSession(
_ session: URLSession,
task: URLSessionTask,
didReceive challenge: URLAuthenticationChallenge) async -> (URLSession.AuthChallengeDisposition, URLCredential?)
{
self.challengeResult(challenge)
}
}
private final class LocalhostSessionTaskState: @unchecked Sendable {
private let lock = NSLock()
private var task: URLSessionDataTask?
private var isCancelled = false
func setTask(_ task: URLSessionDataTask) {
self.lock.lock()
self.task = task
let shouldCancel = self.isCancelled
self.lock.unlock()
if shouldCancel {
task.cancel()
}
}
func cancel() {
self.lock.lock()
self.isCancelled = true
let task = self.task
self.lock.unlock()
task?.cancel()
}
}
@@ -0,0 +1,516 @@
import Foundation
public struct AntigravityOAuthCredentials: Codable, Sendable, Equatable {
public var accessToken: String?
public var refreshToken: String?
public var expiryDateMilliseconds: Double?
public var idToken: String?
public var email: String?
public var projectID: String?
public var clientID: String?
public var clientSecret: String?
public init(
accessToken: String?,
refreshToken: String?,
expiryDate: Date?,
idToken: String? = nil,
email: String? = nil,
projectID: String? = nil,
clientID: String? = nil,
clientSecret: String? = nil)
{
self.accessToken = accessToken
self.refreshToken = refreshToken
self.expiryDateMilliseconds = expiryDate.map { $0.timeIntervalSince1970 * 1000 }
self.idToken = idToken
self.email = email
self.projectID = projectID
self.clientID = clientID
self.clientSecret = clientSecret
}
public var expiryDate: Date? {
guard let expiryDateMilliseconds else { return nil }
return Date(timeIntervalSince1970: expiryDateMilliseconds / 1000)
}
/// Email of the Google account these credentials authenticate, preferring the
/// signed `id_token` claim (what the remote OAuth fetcher reports) and falling
/// back to the stored `email` field. Used to verify that an ambient local/CLI
/// Antigravity snapshot belongs to the account the user explicitly selected.
public var resolvedAccountEmail: String? {
Self.email(fromIDToken: self.idToken) ?? self.email?.trimmedNonEmptyEmail
}
static func email(fromIDToken idToken: String?) -> String? {
guard let idToken else { return nil }
let parts = idToken.components(separatedBy: ".")
guard parts.count >= 2 else { return nil }
var payload = parts[1]
.replacingOccurrences(of: "-", with: "+")
.replacingOccurrences(of: "_", with: "/")
let remainder = payload.count % 4
if remainder > 0 {
payload += String(repeating: "=", count: 4 - remainder)
}
guard let data = Data(base64Encoded: payload, options: .ignoreUnknownCharacters),
let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any]
else {
return nil
}
return (json["email"] as? String)?.trimmedNonEmptyEmail
}
public init(from decoder: Decoder) throws {
let container = try decoder.container(keyedBy: CodingKeys.self)
self.accessToken =
try container.decodeIfPresent(String.self, forKey: .accessTokenSnake)
?? container.decodeIfPresent(String.self, forKey: .accessTokenCamel)
self.refreshToken =
try container.decodeIfPresent(String.self, forKey: .refreshTokenSnake)
?? container.decodeIfPresent(String.self, forKey: .refreshTokenCamel)
self.idToken =
try container.decodeIfPresent(String.self, forKey: .idTokenSnake)
?? container.decodeIfPresent(String.self, forKey: .idTokenCamel)
self.email = try container.decodeIfPresent(String.self, forKey: .email)
self.projectID =
try container.decodeIfPresent(String.self, forKey: .projectIDSnake)
?? container.decodeIfPresent(String.self, forKey: .projectIDCamel)
self.clientID =
try container.decodeIfPresent(String.self, forKey: .clientIDSnake)
?? container.decodeIfPresent(String.self, forKey: .clientIDCamel)
self.clientSecret =
try container.decodeIfPresent(String.self, forKey: .clientSecretSnake)
?? container.decodeIfPresent(String.self, forKey: .clientSecretCamel)
if let expiryDateMilliseconds = try container.decodeIfPresent(Double.self, forKey: .expiryDateSnake)
?? container.decodeIfPresent(Double.self, forKey: .expiresAtCamel)
{
self.expiryDateMilliseconds = expiryDateMilliseconds
} else if let expiryDateMilliseconds = try container.decodeIfPresent(Int.self, forKey: .expiryDateSnake)
?? container.decodeIfPresent(Int.self, forKey: .expiresAtCamel)
{
self.expiryDateMilliseconds = Double(expiryDateMilliseconds)
} else {
self.expiryDateMilliseconds = nil
}
}
public func encode(to encoder: Encoder) throws {
var container = encoder.container(keyedBy: CodingKeys.self)
try container.encodeIfPresent(self.accessToken, forKey: .accessTokenSnake)
try container.encodeIfPresent(self.refreshToken, forKey: .refreshTokenSnake)
try container.encodeIfPresent(self.expiryDateMilliseconds, forKey: .expiryDateSnake)
try container.encodeIfPresent(self.idToken, forKey: .idTokenSnake)
try container.encodeIfPresent(self.email, forKey: .email)
try container.encodeIfPresent(self.projectID, forKey: .projectIDSnake)
try container.encodeIfPresent(self.clientID, forKey: .clientIDSnake)
try container.encodeIfPresent(self.clientSecret, forKey: .clientSecretSnake)
}
enum CodingKeys: String, CodingKey {
case accessTokenSnake = "access_token"
case accessTokenCamel = "accessToken"
case refreshTokenSnake = "refresh_token"
case refreshTokenCamel = "refreshToken"
case expiryDateSnake = "expiry_date"
case expiresAtCamel = "expiresAt"
case idTokenSnake = "id_token"
case idTokenCamel = "idToken"
case email
case projectIDSnake = "project_id"
case projectIDCamel = "projectId"
case clientIDSnake = "client_id"
case clientIDCamel = "clientId"
case clientSecretSnake = "client_secret"
case clientSecretCamel = "clientSecret"
}
}
public struct AntigravityOAuthClient: Sendable, Equatable {
public let clientID: String
public let clientSecret: String
public init(clientID: String, clientSecret: String) {
self.clientID = clientID
self.clientSecret = clientSecret
}
}
public enum AntigravityOAuthConfig {
public static var configuredClientID: String? {
let value = ProcessInfo.processInfo.environment["ANTIGRAVITY_OAUTH_CLIENT_ID"]
return value?.trimmingCharacters(in: .whitespacesAndNewlines).nilIfEmpty
}
public static var configuredClientSecret: String? {
let value = ProcessInfo.processInfo.environment["ANTIGRAVITY_OAUTH_CLIENT_SECRET"]
return value?.trimmingCharacters(in: .whitespacesAndNewlines).nilIfEmpty
}
public static let authURL = URL(string: "https://accounts.google.com/o/oauth2/v2/auth")!
public static let tokenURL = URL(string: "https://oauth2.googleapis.com/token")!
public static let userInfoURL = URL(string: "https://www.googleapis.com/oauth2/v2/userinfo")!
public static let scopes = [
"https://www.googleapis.com/auth/cloud-platform",
"https://www.googleapis.com/auth/userinfo.email",
]
public static let missingCredentialsMessage =
"""
Antigravity OAuth client is not configured. Install Antigravity.app or set \
ANTIGRAVITY_OAUTH_CLIENT_ID and ANTIGRAVITY_OAUTH_CLIENT_SECRET before logging in.
"""
public static func resolvedClient() -> AntigravityOAuthClient? {
if let client = environmentClient() {
return client
}
return Self.discoverClientFromInstalledApp()
}
private static func environmentClient() -> AntigravityOAuthClient? {
guard let clientID = configuredClientID,
let clientSecret = configuredClientSecret
else {
return nil
}
return AntigravityOAuthClient(clientID: clientID, clientSecret: clientSecret)
}
static func discoverClientFromInstalledApp(
applicationRoots: [URL]? = nil,
fileManager: FileManager = .default) -> AntigravityOAuthClient?
{
for url in self.candidateOAuthClientArtifactURLs(
applicationRoots: applicationRoots,
fileManager: fileManager)
where fileManager.fileExists(atPath: url.path)
{
guard let data = try? Data(contentsOf: url),
let client = Self.parseClient(fromInstalledArtifactData: data)
else {
continue
}
return client
}
return nil
}
static func candidateOAuthClientArtifactURLs(
applicationRoots: [URL]? = nil,
fileManager: FileManager = .default) -> [URL]
{
let roots = [
URL(fileURLWithPath: "/Applications", isDirectory: true),
fileManager.homeDirectoryForCurrentUser.appendingPathComponent("Applications", isDirectory: true),
]
let applicationRoots = applicationRoots ?? roots
let appBundleURLs = self.candidateAntigravityAppBundleURLs(
applicationRoots: applicationRoots,
fileManager: fileManager)
let relativePaths = [
"Contents/Resources/app/extensions/antigravity/bin/language_server_macos_arm",
"Contents/Resources/app/extensions/antigravity/bin/language_server_macos_x64",
"Contents/Resources/app/extensions/antigravity/bin/language_server_macos",
"Contents/Resources/app/out/main.js",
"Contents/Resources/bin/language_server",
"Contents/Resources/bin/language_server_macos",
]
return appBundleURLs.flatMap { bundleURL in
relativePaths.map { bundleURL.appendingPathComponent($0) }
}
}
private static func candidateAntigravityAppBundleURLs(
applicationRoots: [URL],
fileManager: FileManager) -> [URL]
{
var urls: [URL] = []
for root in applicationRoots {
urls.append(root.appendingPathComponent("Antigravity.app", isDirectory: true))
let appURLs = (try? fileManager.contentsOfDirectory(
at: root,
includingPropertiesForKeys: [.isDirectoryKey],
options: [.skipsHiddenFiles])) ?? []
for appURL in appURLs where appURL.pathExtension == "app" {
guard self.isAntigravityAppBundle(appURL) else { continue }
urls.append(appURL)
}
}
var seen = Set<String>()
return urls.filter { url in
let key = url.standardizedFileURL.path
guard !seen.contains(key) else { return false }
seen.insert(key)
return true
}
}
private static func isAntigravityAppBundle(_ url: URL) -> Bool {
switch Bundle(url: url)?.bundleIdentifier {
case "com.google.antigravity", "com.google.antigravity-ide":
true
default:
false
}
}
static func parseClient(fromInstalledArtifactData data: Data) -> AntigravityOAuthClient? {
if let content = String(data: data, encoding: .utf8),
let client = parseClient(fromInstalledArtifactText: content)
{
return client
}
let clientIDs = Self.clientIDs(in: data)
let clientSecrets = Self.clientSecrets(in: data)
guard let client = Self.preferredBinaryClient(
clientIDs: clientIDs,
clientSecrets: clientSecrets)
else {
return nil
}
return client
}
static func parseClient(fromInstalledArtifactText content: String) -> AntigravityOAuthClient? {
let marker = "vs/platform/cloudCode/common/oauthClient.js"
let searchStart = content.range(of: marker)?.lowerBound ?? content.startIndex
let searchEnd = content.index(searchStart, offsetBy: 4000, limitedBy: content.endIndex) ?? content.endIndex
let haystack = String(content[searchStart..<searchEnd])
guard let clientID = Self.firstMatch(
pattern: #"[0-9]+-[A-Za-z0-9_-]+\.apps\.googleusercontent\.com"#,
in: haystack),
let clientSecret = Self.firstMatch(
pattern: #"GOCSPX-[A-Za-z0-9_-]{28}"#,
in: haystack)
else {
return nil
}
return AntigravityOAuthClient(clientID: clientID, clientSecret: clientSecret)
}
private static func clientIDs(in data: Data) -> [String] {
let suffix = Data(".apps.googleusercontent.com".utf8)
var searchRange = data.startIndex..<data.endIndex
var values: [String] = []
while let range = data.range(of: suffix, options: [], in: searchRange) {
var start = range.lowerBound
while start > data.startIndex {
let previous = data.index(before: start)
guard Self.isOAuthClientIDPrefixByte(data[previous]) else { break }
start = previous
}
let candidateData = Data(data[start..<range.upperBound])
if let candidate = String(data: candidateData, encoding: .ascii),
let clientID = Self.firstMatch(
pattern: #"[0-9]+-[A-Za-z0-9_-]+\.apps\.googleusercontent\.com"#,
in: candidate)
{
values.append(clientID)
}
searchRange = range.upperBound..<data.endIndex
}
return self.unique(values)
}
private static func clientSecrets(in data: Data) -> [String] {
let prefix = Data("GOCSPX-".utf8)
let secretLength = 35
var searchRange = data.startIndex..<data.endIndex
var values: [String] = []
while let range = data.range(of: prefix, options: [], in: searchRange) {
let end = range.lowerBound + secretLength
if end <= data.endIndex {
let candidateData = Data(data[range.lowerBound..<end])
if candidateData.dropFirst(prefix.count).allSatisfy(Self.isOAuthClientSecretByte),
let candidate = String(data: candidateData, encoding: .ascii)
{
values.append(candidate)
}
}
searchRange = range.upperBound..<data.endIndex
}
return self.unique(values)
}
private static func preferredBinaryClient(
clientIDs: [String],
clientSecrets: [String]) -> AntigravityOAuthClient?
{
guard !clientIDs.isEmpty,
!clientSecrets.isEmpty
else {
return nil
}
if clientSecrets.count == 1, clientIDs.count > 1 {
return AntigravityOAuthClient(clientID: clientIDs[clientIDs.count - 1], clientSecret: clientSecrets[0])
}
let clientSecret: String = if clientSecrets.count == clientIDs.count, clientSecrets.count > 1 {
// Antigravity 2's language_server binary stores the secret table before the client id table.
clientSecrets[clientSecrets.count - 1]
} else {
clientSecrets[0]
}
return AntigravityOAuthClient(clientID: clientIDs[0], clientSecret: clientSecret)
}
private static func isOAuthClientIDPrefixByte(_ byte: UInt8) -> Bool {
(byte >= 48 && byte <= 57)
|| (byte >= 65 && byte <= 90)
|| (byte >= 97 && byte <= 122)
|| byte == 45
|| byte == 95
}
private static func isOAuthClientSecretByte(_ byte: UInt8) -> Bool {
self.isOAuthClientIDPrefixByte(byte)
}
private static func firstMatch(pattern: String, in text: String) -> String? {
guard let regex = try? NSRegularExpression(pattern: pattern) else { return nil }
let range = NSRange(text.startIndex..<text.endIndex, in: text)
guard let match = regex.firstMatch(in: text, range: range),
let swiftRange = Range(match.range, in: text)
else {
return nil
}
return String(text[swiftRange])
}
private static func unique(_ values: [String]) -> [String] {
var seen = Set<String>()
return values.filter { value in
guard !seen.contains(value) else { return false }
seen.insert(value)
return true
}
}
}
public struct AntigravityOAuthCredentialsStore: @unchecked Sendable {
public static let environmentCredentialsKey = "ANTIGRAVITY_OAUTH_CREDENTIALS_JSON"
private static let fileLock = NSLock()
public let fileURL: URL
private let fileManager: FileManager
public init(fileURL: URL = Self.defaultURL(), fileManager: FileManager = .default) {
self.fileURL = fileURL
self.fileManager = fileManager
}
public func load() throws -> AntigravityOAuthCredentials? {
try Self.fileLock.withLock {
try self.loadUnlocked()
}
}
private func loadUnlocked() throws -> AntigravityOAuthCredentials? {
guard self.fileManager.fileExists(atPath: self.fileURL.path) else { return nil }
let data = try Data(contentsOf: self.fileURL)
return try JSONDecoder().decode(AntigravityOAuthCredentials.self, from: data)
}
public func save(_ credentials: AntigravityOAuthCredentials) throws {
try Self.fileLock.withLock {
let data = try JSONEncoder.antigravityCredentials.encode(credentials)
let directory = self.fileURL.deletingLastPathComponent()
if !self.fileManager.fileExists(atPath: directory.path) {
try self.fileManager.createDirectory(at: directory, withIntermediateDirectories: true)
}
try data.write(to: self.fileURL, options: [.atomic])
try self.applySecurePermissionsIfNeeded()
}
}
public func deleteIfPresent() throws {
try Self.fileLock.withLock {
try self.deleteIfPresentUnlocked()
}
}
public func deleteIfPresent(
matching predicate: (AntigravityOAuthCredentials) -> Bool) throws
{
try Self.fileLock.withLock {
guard let credentials = try self.loadUnlocked(), predicate(credentials) else { return }
try self.deleteIfPresentUnlocked()
}
}
private func deleteIfPresentUnlocked() throws {
guard self.fileManager.fileExists(atPath: self.fileURL.path) else { return }
try self.fileManager.removeItem(at: self.fileURL)
}
public static func defaultDirectoryURL(home: URL = FileManager.default.homeDirectoryForCurrentUser) -> URL {
home
.appendingPathComponent(".codexbar", isDirectory: true)
.appendingPathComponent("antigravity", isDirectory: true)
}
public static func defaultURL(home: URL = FileManager.default.homeDirectoryForCurrentUser) -> URL {
self.defaultDirectoryURL(home: home)
.appendingPathComponent("oauth_creds.json")
}
public static func tokenAccountValue(for credentials: AntigravityOAuthCredentials) throws -> String {
let data = try JSONEncoder.antigravityCredentials.encode(credentials)
guard let value = String(data: data, encoding: .utf8) else {
throw CocoaError(.coderInvalidValue)
}
return value
}
public static func credentials(fromTokenAccountValue value: String) -> AntigravityOAuthCredentials? {
let trimmed = value.trimmingCharacters(in: .whitespacesAndNewlines)
guard !trimmed.isEmpty, let data = trimmed.data(using: .utf8) else { return nil }
return try? JSONDecoder().decode(AntigravityOAuthCredentials.self, from: data)
}
private func applySecurePermissionsIfNeeded() throws {
#if os(macOS) || os(Linux)
try self.fileManager.setAttributes([
.posixPermissions: NSNumber(value: Int16(0o600)),
], ofItemAtPath: self.fileURL.path)
#endif
}
}
extension JSONEncoder {
fileprivate static let antigravityCredentials: JSONEncoder = {
let encoder = JSONEncoder()
encoder.outputFormatting = [.prettyPrinted, .sortedKeys]
return encoder
}()
}
extension String {
fileprivate var nilIfEmpty: String? {
self.isEmpty ? nil : self
}
fileprivate var trimmedNonEmptyEmail: String? {
let trimmed = self.trimmingCharacters(in: .whitespacesAndNewlines)
return trimmed.isEmpty ? nil : trimmed
}
}
@@ -0,0 +1,673 @@
import Foundation
public enum AntigravityProviderDescriptor {
public static let descriptor: ProviderDescriptor = Self.makeDescriptor()
static func makeDescriptor() -> ProviderDescriptor {
ProviderDescriptor(
id: .antigravity,
metadata: ProviderMetadata(
id: .antigravity,
displayName: "Antigravity",
sessionLabel: "Gemini Models",
weeklyLabel: "Claude and GPT",
opusLabel: nil,
supportsOpus: false,
supportsCredits: false,
creditsHint: "",
toggleTitle: "Show Antigravity usage (experimental)",
cliName: "antigravity",
defaultEnabled: false,
isPrimaryProvider: false,
usesAccountFallback: false,
dashboardURL: nil,
statusPageURL: nil,
statusLinkURL: "https://www.google.com/appsstatus/dashboard/products/npdyhgECDJ6tB66MxXyo/history",
statusWorkspaceProductID: "npdyhgECDJ6tB66MxXyo"),
branding: ProviderBranding(
iconStyle: .antigravity,
iconResourceName: "ProviderIcon-antigravity",
color: ProviderColor(red: 96 / 255, green: 186 / 255, blue: 126 / 255)),
tokenCost: ProviderTokenCostConfig(
supportsTokenCost: false,
noDataMessage: { "Antigravity cost summary is not supported." }),
fetchPlan: ProviderFetchPlan(
sourceModes: [.auto, .cli, .oauth],
pipeline: ProviderFetchPipeline(resolveStrategies: self.resolveStrategies)),
cli: ProviderCLIConfig(
name: "antigravity",
versionDetector: nil))
}
private static func resolveStrategies(context: ProviderFetchContext) async -> [any ProviderFetchStrategy] {
let app = AntigravityStatusFetchStrategy(source: .app)
let cli = AntigravityCLIHTTPSFetchStrategy()
let ide = AntigravityStatusFetchStrategy(source: .ide)
let oauth = AntigravityOAuthFetchStrategy()
switch context.sourceMode {
case .cli:
return [app, cli, ide]
case .oauth:
return [oauth]
case .auto:
if context.selectedTokenAccountID != nil ||
context.env[AntigravityOAuthCredentialsStore.environmentCredentialsKey] != nil ||
self.hasSharedOAuthCredentials(context: context)
{
return [app, cli, ide, oauth]
}
return [app, cli, ide]
case .web, .api:
return []
}
}
private static func hasSharedOAuthCredentials(context: ProviderFetchContext) -> Bool {
let homeURL = context.env["HOME"]
.flatMap { $0.isEmpty ? nil : URL(fileURLWithPath: $0, isDirectory: true) }
?? FileManager.default.homeDirectoryForCurrentUser
let fileURL = AntigravityOAuthCredentialsStore.defaultURL(home: homeURL)
return FileManager.default.fileExists(atPath: fileURL.path)
}
}
struct AntigravityStatusFetchStrategy: ProviderFetchStrategy {
enum Source {
case app
case ide
var id: String {
switch self {
case .app: "antigravity.app-local"
case .ide: "antigravity.ide-local"
}
}
var processScope: AntigravityStatusProbe.ProcessScope {
switch self {
case .app: .appOnly
case .ide: .ideOnly
}
}
var sourceLabel: String {
switch self {
case .app: "app"
case .ide: "ide"
}
}
}
let source: Source
var id: String {
self.source.id
}
let kind: ProviderFetchKind = .localProbe
init(source: Source = .app) {
self.source = source
}
func isAvailable(_: ProviderFetchContext) async -> Bool {
true
}
func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult {
let probe = AntigravityStatusProbe(processScope: self.source.processScope)
let selectedAccountEmail: String? = if context.sourceMode == .auto, context.selectedTokenAccountID != nil {
AntigravitySelectedAccountGuard.selectedAccountEmail(context: context)
} else {
nil
}
let snap = try await probe.fetch(matchingAccountEmail: selectedAccountEmail)
let usage = try snap.toUsageSnapshot()
try AntigravitySelectedAccountGuard.validate(usage, context: context)
return self.makeResult(
usage: usage,
sourceLabel: self.source.sourceLabel)
}
func shouldFallback(on _: Error, context: ProviderFetchContext) -> Bool {
context.sourceMode == .auto || context.sourceMode == .cli
}
}
/// When the Antigravity 2.0 app is closed or unavailable, this strategy spawns
/// or reuses ``agy`` and talks to the localhost server embedded in that CLI
/// process. ``agy`` is an interactive REPL, not a query command, so CodexBar
/// never scrapes TUI output here; it only keeps the process alive long enough
/// for the server to answer quota endpoints.
struct AntigravityCLIHTTPSFetchStrategy: ProviderFetchStrategy {
static let sourceLabel = "cli"
let id: String = "antigravity.cli-https"
let kind: ProviderFetchKind = .cli
private static let log = CodexBarLog.logger(LogCategories.antigravity)
struct SnapshotWaitDependencies {
let pollIntervalNanoseconds: UInt64
let listeningPorts: @Sendable (Int, TimeInterval) async throws -> [Int]
let drainOutput: @Sendable () async -> Data
let fetchSnapshot: @Sendable ([Int]) async throws -> AntigravityStatusSnapshot
let now: @Sendable () -> Date
init(
pollIntervalNanoseconds: UInt64,
listeningPorts: @escaping @Sendable (Int, TimeInterval) async throws -> [Int],
drainOutput: @escaping @Sendable () async -> Data,
fetchSnapshot: @escaping @Sendable ([Int]) async throws -> AntigravityStatusSnapshot,
now: @escaping @Sendable () -> Date = Date.init)
{
self.pollIntervalNanoseconds = pollIntervalNanoseconds
self.listeningPorts = listeningPorts
self.drainOutput = drainOutput
self.fetchSnapshot = fetchSnapshot
self.now = now
}
}
/// Seams for discovering and reusing an already-running ``agy`` CLI language
/// server, so a fresh spawn (and its multi-second ``GetUserStatus`` warm-up)
/// can be skipped when a warm server is already present.
struct WarmAgyDependencies {
let processInfos: @Sendable (TimeInterval) async throws -> [AntigravityStatusProbe.ProcessInfoResult]
let listeningPorts: @Sendable (Int, TimeInterval) async throws -> [Int]
let fetchSnapshot: @Sendable ([Int], TimeInterval) async throws -> AntigravityStatusSnapshot
let processOwnerUserID: @Sendable (Int) -> UInt32?
let currentUserID: @Sendable () -> UInt32
/// The pid of an ``agy`` that CodexBar itself spawned and manages through
/// ``AntigravityCLISession`` (if any). Such a process must NOT be reused
/// through the warm path: doing so bypasses `beginProbe`/`finishProbe`, so
/// the idle timer is never cancelled/extended and `stopIfIdle` could tear
/// the managed session down mid-poll. Externally owned `agy` (an IDE, a
/// long-lived `agy`, or another CodexBar host) has no such accounting.
let ownedPID: @Sendable () async -> Int?
let now: @Sendable () -> Date
init(
processInfos: @escaping @Sendable (TimeInterval) async throws
-> [AntigravityStatusProbe.ProcessInfoResult],
listeningPorts: @escaping @Sendable (Int, TimeInterval) async throws -> [Int],
fetchSnapshot: @escaping @Sendable ([Int], TimeInterval) async throws -> AntigravityStatusSnapshot,
processOwnerUserID: @escaping @Sendable (Int) -> UInt32? = { _ in 0 },
currentUserID: @escaping @Sendable () -> UInt32 = { 0 },
ownedPID: @escaping @Sendable () async -> Int? = { nil },
now: @escaping @Sendable () -> Date = Date.init)
{
self.processInfos = processInfos
self.listeningPorts = listeningPorts
self.fetchSnapshot = fetchSnapshot
self.processOwnerUserID = processOwnerUserID
self.currentUserID = currentUserID
self.ownedPID = ownedPID
self.now = now
}
}
/// Discover an already-running, authenticated ``agy`` CLI language server and
/// reuse its listening ports instead of spawning a fresh process.
///
/// One-shot CLI invocations otherwise spawn a brand-new ``agy`` on every
/// call; a fresh server binds its port quickly but ``GetUserStatus`` returns
/// transient initialization failures for a few seconds, so the readiness
/// deadline is occasionally missed. When a warm CLI server is already up, we
/// can talk to it immediately it needs no CSRF token (``cliHTTPS``).
///
/// Returns the snapshot from the first warm server that answers with
/// parseable usage for the requested account, or `nil` when none is found or
/// none answers in which case the caller falls back to the existing spawn
/// path unchanged.
static func tryWarmAgyFetch(
timeout: TimeInterval,
expectedBinaryPath: String? = nil,
expectedAccountEmail: String? = nil,
dependencies: WarmAgyDependencies) async throws -> AntigravityStatusSnapshot?
{
try Task.checkCancellation()
let deadline = dependencies.now().addingTimeInterval(timeout)
guard let discoveryTimeout = Self.remainingWarmProbeTime(deadline: deadline, now: dependencies.now) else {
return nil
}
let processInfos: [AntigravityStatusProbe.ProcessInfoResult]
do {
processInfos = try await dependencies.processInfos(discoveryTimeout)
} catch let error as CancellationError {
throw error
} catch {
return nil
}
try Task.checkCancellation()
let ownedPID = await dependencies.ownedPID()
try Task.checkCancellation()
let currentUserID = dependencies.currentUserID()
// Only the CLI's language server needs no CSRF token; the IDE/app servers
// require one and must not be reused through this token-less fast path.
// Also exclude any `agy` CodexBar itself spawned and manages: reusing it
// here would bypass session lifecycle accounting (see `ownedPID`).
let cliProcesses = processInfos.filter { info in
info.pid != ownedPID &&
dependencies.processOwnerUserID(info.pid) == currentUserID &&
AntigravityStatusProbe.antigravityProcessKind(info.commandLine) == .cli
}
guard !cliProcesses.isEmpty else { return nil }
for info in cliProcesses {
if let expectedBinaryPath {
guard Self.commandLine(info.commandLine, matchesBinaryPath: expectedBinaryPath)
else {
continue
}
}
guard let portTimeout = Self.remainingWarmProbeTime(deadline: deadline, now: dependencies.now) else {
return nil
}
let ports: [Int]
do {
ports = try await dependencies.listeningPorts(info.pid, portTimeout)
} catch let error as CancellationError {
throw error
} catch {
continue
}
try Task.checkCancellation()
guard !ports.isEmpty else { continue }
guard let fetchTimeout = Self.remainingWarmProbeTime(deadline: deadline, now: dependencies.now) else {
return nil
}
let snapshot: AntigravityStatusSnapshot
do {
snapshot = try await dependencies.fetchSnapshot(ports, fetchTimeout)
} catch let error as CancellationError {
throw error
} catch {
continue
}
try Task.checkCancellation()
guard (try? snapshot.toUsageSnapshot()) != nil,
AntigravitySelectedAccountGuard.matches(
snapshotAccountEmail: snapshot.accountEmail,
expectedAccountEmail: expectedAccountEmail)
else {
continue
}
Self.log.debug("Antigravity CLI HTTPS reusing warm agy", metadata: [
"pid": "\(info.pid)",
"ports": ports.map(String.init).joined(separator: ","),
])
return snapshot
}
try Task.checkCancellation()
return nil
}
private static func remainingWarmProbeTime(
deadline: Date,
now: @Sendable () -> Date) -> TimeInterval?
{
let remaining = deadline.timeIntervalSince(now())
return remaining > 0 ? remaining : nil
}
private static func commandLine(_ commandLine: String, matchesBinaryPath binaryPath: String) -> Bool {
let candidates = [
URL(fileURLWithPath: binaryPath).standardizedFileURL.path,
URL(fileURLWithPath: binaryPath).resolvingSymlinksInPath().standardizedFileURL.path,
]
return candidates.contains { candidate in
commandLine == candidate || commandLine.hasPrefix("\(candidate) ")
}
}
/// Production wiring for ``tryWarmAgyFetch``: list processes via `ps`, find
/// listening ports via `lsof`, and probe the token-less CLI HTTPS endpoint.
static func liveWarmAgyDependencies() -> WarmAgyDependencies {
WarmAgyDependencies(
processInfos: { timeout in
// A missing-CSRF/notRunning throw means no reusable server; the
// caller maps any throw to "no warm agy" and spawns instead.
try await AntigravityStatusProbe.detectProcessInfos(
timeout: timeout,
scope: .ideAndCLI)
},
listeningPorts: { pid, timeout in
try await AntigravityStatusProbe.listeningPorts(pid: pid, timeout: timeout)
},
fetchSnapshot: { ports, timeout in
let deadline = Date().addingTimeInterval(timeout)
return try await AntigravityStatusProbe(timeout: timeout)
.fetchFromPorts(ports, deadline: deadline)
},
processOwnerUserID: { pid in
AntigravityProcessIdentityProvider().ownerUserID(for: pid_t(pid))
},
currentUserID: {
AntigravityProcessIdentityProvider.currentUserID
},
ownedPID: {
// The pid of the `agy` CodexBar manages through the shared
// session, so the warm scan never reuses our own process.
await AntigravityCLISession.shared.pid.map(Int.init)
})
}
func isAvailable(_ context: ProviderFetchContext) async -> Bool {
BinaryLocator.resolveAntigravityBinary(env: context.env) != nil
}
func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult {
guard let binary = BinaryLocator.resolveAntigravityBinary(env: context.env) else {
throw AntigravityStatusProbeError.notRunning
}
let expectedAccountEmail: String? = if context.sourceMode == .auto,
context.selectedTokenAccountID != nil
{
AntigravitySelectedAccountGuard.selectedAccountEmail(context: context)
} else {
nil
}
let result = try await self.fetchUsingWarmSession(
binary: binary,
idleWindow: context.persistentCLISessionIdleWindow,
resetAfterFetch: Self.shouldResetSessionAfterFetch(context),
expectedAccountEmail: expectedAccountEmail)
try AntigravitySelectedAccountGuard.validate(result.usage, context: context)
return result
}
private func fetchUsingWarmSession(
binary: String,
idleWindow: TimeInterval?,
resetAfterFetch: Bool,
expectedAccountEmail: String?) async throws -> ProviderFetchResult
{
try await self.fetchUsingWarmSession(
binary: binary,
idleWindow: idleWindow,
resetAfterFetch: resetAfterFetch,
expectedAccountEmail: expectedAccountEmail,
warmDependencies: Self.liveWarmAgyDependencies(),
spawnFetch: { binary, idleWindow, resetAfterFetch in
try await self.fetchBySpawning(
binary: binary,
idleWindow: idleWindow,
resetAfterFetch: resetAfterFetch)
})
}
/// Testable core of the CLI fetch: try the warm-reuse fast path first, then
/// fall back to spawning. The `spawnFetch` seam lets tests assert the spawn
/// path is skipped when a warm server is reused.
func fetchUsingWarmSession(
binary: String,
idleWindow: TimeInterval?,
resetAfterFetch: Bool,
expectedAccountEmail: String? = nil,
warmDependencies: WarmAgyDependencies,
spawnFetch: @Sendable (String, TimeInterval?, Bool) async throws -> ProviderFetchResult)
async throws -> ProviderFetchResult
{
// Fast path: reuse an already-running, authenticated `agy` CLI server if
// one is present, avoiding a fresh spawn and its multi-second warm-up.
// When none is found (or none answers), fall through to the spawn path.
//
// The warm path deliberately does NOT touch `AntigravityCLISession`
// (`beginProbe`/`finishProbe`): a discovered `agy` is owned by another
// process (an IDE, a long-lived `agy`, or another CodexBar host), so
// CodexBar must not manage its lifecycle, idle timeout, or
// `resetAfterFetch` teardown. Those apply only to processes CodexBar
// itself spawns on the fallback path below.
// Long-lived hosts already keep a managed session warm. Restrict external
// process reuse to one-shot CLI calls so app/server lifecycle accounting
// stays entirely inside AntigravityCLISession.
if resetAfterFetch, let warmSnapshot = try await Self.tryWarmAgyFetch(
timeout: 2.0,
expectedBinaryPath: binary,
expectedAccountEmail: expectedAccountEmail,
dependencies: warmDependencies)
{
// `tryWarmAgyFetch` only returns a snapshot whose `toUsageSnapshot()`
// already succeeded, so this conversion must not silently fail.
let warmUsage = try warmSnapshot.toUsageSnapshot()
return self.makeResult(
usage: warmUsage,
sourceLabel: Self.sourceLabel)
}
try Task.checkCancellation()
return try await spawnFetch(binary, idleWindow, resetAfterFetch)
}
/// Spawn (or reuse CodexBar's own warm) `agy` session and wait for the CLI
/// HTTPS endpoint to report ready. This is the original behavior, unchanged.
private func fetchBySpawning(
binary: String,
idleWindow: TimeInterval?,
resetAfterFetch: Bool) async throws -> ProviderFetchResult
{
let session = AntigravityCLISession.shared
let pid = try await session.beginProbe(binary: binary, idleWindow: idleWindow)
let deadline = Date().addingTimeInterval(5.0)
let snap: AntigravityStatusSnapshot
let usage: UsageSnapshot
do {
snap = try await Self.waitForSnapshot(
pid: pid,
deadline: deadline,
dependencies: SnapshotWaitDependencies(
pollIntervalNanoseconds: 200_000_000,
listeningPorts: { pid, timeout in
try await AntigravityStatusProbe.listeningPorts(pid: pid, timeout: timeout)
},
drainOutput: {
await session.drainOutput()
},
fetchSnapshot: { ports in
let timeout = min(2.0, max(0.2, deadline.timeIntervalSinceNow))
return try await AntigravityStatusProbe(timeout: timeout)
.fetchFromPorts(ports, deadline: deadline)
}))
usage = try snap.toUsageSnapshot()
await session.finishProbe(success: true, resetAfterFetch: resetAfterFetch)
} catch {
let authenticationRequired = (error as? AntigravityStatusProbeError) == .authenticationRequired
await session.finishProbe(
success: false,
resetAfterFetch: resetAfterFetch || authenticationRequired,
forceTerminate: authenticationRequired)
throw error
}
return self.makeResult(
usage: usage,
sourceLabel: Self.sourceLabel)
}
static func shouldResetSessionAfterFetch(_ context: ProviderFetchContext) -> Bool {
// Long-lived hosts (the app, `codexbar serve`) keep the warm `agy`
// session between fetches; only one-shot CLI invocations reset it.
context.runtime == .cli && !context.persistsCLISessions
}
/// Waits for real API readiness, not just socket readiness. Fresh ``agy``
/// processes bind ports quickly, but ``GetUserStatus`` can return transient
/// initialization failures for a few seconds after the port appears.
static func waitForSnapshot(
pid: pid_t,
deadline: Date,
dependencies: SnapshotWaitDependencies) async throws -> AntigravityStatusSnapshot
{
var lastFetchError: Error?
while dependencies.now() < deadline {
try await Self.checkAuthenticationPrompt(dependencies)
let remaining = deadline.timeIntervalSince(dependencies.now())
let portProbeTimeout = min(2.0, max(0.2, remaining))
let ports: [Int]
do {
ports = try await dependencies.listeningPorts(Int(pid), portProbeTimeout)
} catch {
guard Self.isNoListeningPortsError(error) else {
try await Self.checkAuthenticationPrompt(dependencies)
throw error
}
ports = []
}
if !ports.isEmpty {
var readySnapshot: AntigravityStatusSnapshot?
do {
let snapshot = try await dependencies.fetchSnapshot(ports)
_ = try snapshot.toUsageSnapshot()
readySnapshot = snapshot
} catch {
try await Self.checkAuthenticationPrompt(dependencies)
lastFetchError = error
Self.log.debug("Antigravity CLI HTTPS endpoint not ready", metadata: [
"pid": "\(pid)",
"ports": ports.map(String.init).joined(separator: ","),
"error": error.localizedDescription,
])
}
if let readySnapshot {
try await Self.checkAuthenticationPrompt(dependencies)
return readySnapshot
}
}
let remainingNanoseconds = UInt64(
max(0, deadline.timeIntervalSince(dependencies.now())) * 1_000_000_000)
guard remainingNanoseconds > 0 else { break }
let sleepNanoseconds = min(dependencies.pollIntervalNanoseconds, remainingNanoseconds)
if sleepNanoseconds > 0 {
try await Task.sleep(nanoseconds: sleepNanoseconds)
}
}
try await Self.checkAuthenticationPrompt(dependencies)
if let lastFetchError {
throw lastFetchError
}
Self.log.warning("Antigravity CLI HTTPS: no ports found for pid \(pid)")
throw AntigravityStatusProbeError.portDetectionFailed(
"Antigravity CLI started but no listening ports found")
}
static func containsAuthenticationPrompt(_ output: Data) -> Bool {
AntigravityCLIAuthenticationPrompt.contains(output)
}
private static func checkAuthenticationPrompt(_ dependencies: SnapshotWaitDependencies) async throws {
let terminalOutput = await dependencies.drainOutput()
if Self.containsAuthenticationPrompt(terminalOutput) {
throw AntigravityStatusProbeError.authenticationRequired
}
}
private static func isNoListeningPortsError(_ error: Error) -> Bool {
if case let AntigravityStatusProbeError.portDetectionFailed(message) = error {
return message == "no listening ports found"
}
if case let SubprocessRunnerError.nonZeroExit(code, stderr) = error {
return code == 1 && stderr.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty
}
return false
}
func shouldFallback(on _: Error, context: ProviderFetchContext) -> Bool {
context.sourceMode == .auto || context.sourceMode == .cli
}
}
struct AntigravityOAuthFetchStrategy: ProviderFetchStrategy {
let id: String = "antigravity.oauth"
let kind: ProviderFetchKind = .oauth
func isAvailable(_: ProviderFetchContext) async -> Bool {
true
}
static func usageSnapshot(
from snapshot: AntigravityStatusSnapshot,
updatedAt: Date = Date()) throws -> UsageSnapshot
{
if snapshot.modelQuotas.isEmpty {
return UsageSnapshot(
primary: nil,
secondary: nil,
tertiary: nil,
updatedAt: updatedAt,
identity: ProviderIdentitySnapshot(
providerID: .antigravity,
accountEmail: snapshot.accountEmail,
accountOrganization: nil,
loginMethod: snapshot.accountPlan))
}
return try snapshot.toUsageSnapshot()
}
func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult {
let fetcher = AntigravityRemoteUsageFetcher(
environment: context.env,
credentialsUpdateHandler: { credentials in
guard let accountID = context.selectedTokenAccountID,
let updater = context.tokenAccountTokenUpdater
else {
return
}
let token = try AntigravityOAuthCredentialsStore.tokenAccountValue(for: credentials)
await updater(.antigravity, accountID, token)
})
let snapshot = try await fetcher.fetch()
let usage = try Self.usageSnapshot(from: snapshot)
return self.makeResult(
usage: usage,
sourceLabel: "oauth")
}
func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool {
false
}
}
/// Guards ambient Antigravity snapshots against the explicitly selected account.
///
/// The local desktop probe and the ``agy`` CLI HTTPS server report whichever
/// Antigravity account is signed into the local session. When the user has
/// selected a specific saved Google account, an ambient probe can return a
/// *different* account's quota. Only the OAuth strategy is account-scoped (it
/// fetches with the selected account's injected credentials), so in ``auto``
/// mode we reject a snapshot whose identity does not match the selected account
/// and let the pipeline fall through to OAuth. Explicit ``cli``/``oauth`` source
/// modes stay authoritative and are never second-guessed here.
enum AntigravitySelectedAccountGuard {
static func matches(snapshotAccountEmail: String?, expectedAccountEmail: String?) -> Bool {
guard let expected = self.normalizedEmail(expectedAccountEmail) else { return true }
guard let found = self.normalizedEmail(snapshotAccountEmail) else { return false }
return found.caseInsensitiveCompare(expected) == .orderedSame
}
static func validate(_ usage: UsageSnapshot, context: ProviderFetchContext) throws {
guard context.sourceMode == .auto, context.selectedTokenAccountID != nil else { return }
let expected = self.selectedAccountEmail(context: context)
let found = self.normalizedEmail(usage.identity?.accountEmail)
guard let expected, let found, found.caseInsensitiveCompare(expected) == .orderedSame else {
throw AntigravityStatusProbeError.accountMismatch(expected: expected, found: found)
}
}
/// Email of the selected token account, read from the same injected
/// credentials the OAuth strategy would use (`ANTIGRAVITY_OAUTH_CREDENTIALS_JSON`).
static func selectedAccountEmail(context: ProviderFetchContext) -> String? {
guard let value = context.env[AntigravityOAuthCredentialsStore.environmentCredentialsKey],
let credentials = AntigravityOAuthCredentialsStore.credentials(fromTokenAccountValue: value)
else {
return nil
}
return credentials.resolvedAccountEmail
}
private static func normalizedEmail(_ email: String?) -> String? {
guard let trimmed = email?.trimmingCharacters(in: .whitespacesAndNewlines), !trimmed.isEmpty else {
return nil
}
return trimmed
}
}
@@ -0,0 +1,173 @@
import Foundation
// swiftformat:disable:next redundantSendable
struct AntigravityQuotaSummary: Sendable, Equatable {
let description: String?
let groups: [AntigravityQuotaSummaryGroup]
}
// swiftformat:disable:next redundantSendable
struct AntigravityQuotaSummaryGroup: Sendable, Equatable {
let displayName: String
let description: String?
let buckets: [AntigravityQuotaSummaryBucket]
}
// swiftformat:disable:next redundantSendable
struct AntigravityQuotaSummaryBucket: Sendable, Equatable {
let bucketId: String
let displayName: String
let remainingFraction: Double?
let resetTime: Date?
let resetDescription: String?
let disabled: Bool
init(
bucketId: String,
displayName: String,
remainingFraction: Double?,
resetTime: Date? = nil,
resetDescription: String?,
disabled: Bool)
{
self.bucketId = bucketId
self.displayName = displayName
self.remainingFraction = remainingFraction
self.resetTime = resetTime
self.resetDescription = resetDescription
self.disabled = disabled
}
}
extension AntigravityStatusProbe {
static func parseQuotaSummaryResponse(_ data: Data) throws -> AntigravityStatusSnapshot {
let decoder = JSONDecoder()
let response = try decoder.decode(QuotaSummaryResponse.self, from: data)
if let invalid = Self.invalidCode(response.code) {
throw AntigravityStatusProbeError.apiError(invalid)
}
let payload = response.response ?? response.summary ?? response.rootPayload
guard let payload else {
throw AntigravityStatusProbeError.parseFailed("Missing quota summary")
}
let groups = payload.groups.compactMap(self.quotaSummaryGroup(from:))
guard !groups.isEmpty else {
throw AntigravityStatusProbeError.parseFailed("Missing quota groups")
}
return AntigravityStatusSnapshot(
quotaSummary: AntigravityQuotaSummary(
description: payload.description,
groups: groups),
accountEmail: nil,
accountPlan: nil,
source: .local)
}
private static func quotaSummaryGroup(from payload: QuotaSummaryGroupPayload) -> AntigravityQuotaSummaryGroup? {
let displayName = payload.displayName?.trimmingCharacters(in: .whitespacesAndNewlines)
let buckets = (payload.buckets ?? []).compactMap(self.quotaSummaryBucket(from:))
guard !buckets.isEmpty else { return nil }
return AntigravityQuotaSummaryGroup(
displayName: self.nonEmpty(displayName) ?? "Quota",
description: payload.description,
buckets: buckets)
}
private static func quotaSummaryBucket(from payload: QuotaSummaryBucketPayload) -> AntigravityQuotaSummaryBucket? {
let bucketId = payload.bucketId?.trimmingCharacters(in: .whitespacesAndNewlines)
let displayName = payload.displayName?.trimmingCharacters(in: .whitespacesAndNewlines)
guard let resolvedBucketId = bucketId, !resolvedBucketId.isEmpty else { return nil }
let resetTime = payload.resetTime.flatMap { Self.parseDate($0) }
return AntigravityQuotaSummaryBucket(
bucketId: resolvedBucketId,
displayName: self.nonEmpty(displayName) ?? resolvedBucketId,
remainingFraction: payload.resolvedRemainingFraction,
resetTime: resetTime,
resetDescription: payload.description,
disabled: payload.disabled ?? false)
}
private static func nonEmpty(_ value: String?) -> String? {
guard let value, !value.isEmpty else { return nil }
return value
}
}
private struct QuotaSummaryResponse: Decodable {
let code: CodeValue?
let message: String?
let response: QuotaSummaryPayload?
let summary: QuotaSummaryPayload?
let description: String?
let groups: [QuotaSummaryGroupPayload]?
var rootPayload: QuotaSummaryPayload? {
guard let groups else { return nil }
return QuotaSummaryPayload(description: self.description, groups: groups)
}
}
private struct QuotaSummaryPayload: Decodable {
let description: String?
let groups: [QuotaSummaryGroupPayload]
init(description: String?, groups: [QuotaSummaryGroupPayload]) {
self.description = description
self.groups = groups
}
private enum CodingKeys: String, CodingKey {
case description
case groups
}
init(from decoder: Decoder) throws {
let container = try decoder.container(keyedBy: CodingKeys.self)
self.description = try container.decodeIfPresent(String.self, forKey: .description)
self.groups = try container.decodeIfPresent([QuotaSummaryGroupPayload].self, forKey: .groups) ?? []
}
}
private struct QuotaSummaryGroupPayload: Decodable {
let displayName: String?
let description: String?
let buckets: [QuotaSummaryBucketPayload]?
}
private struct QuotaSummaryBucketPayload: Decodable {
let bucketId: String?
let displayName: String?
let description: String?
let disabled: Bool?
let remainingFraction: Double?
let remaining: QuotaSummaryRemainingPayload?
let resetTime: String?
var resolvedRemainingFraction: Double? {
self.remainingFraction ?? self.remaining?.remainingFraction
}
}
private struct QuotaSummaryRemainingPayload: Decodable {
let remainingFraction: Double?
private enum CodingKeys: String, CodingKey {
case remainingFraction
case oneofCase = "case"
case value
}
init(from decoder: Decoder) throws {
let container = try decoder.container(keyedBy: CodingKeys.self)
if let remainingFraction = try container.decodeIfPresent(Double.self, forKey: .remainingFraction) {
self.remainingFraction = remainingFraction
return
}
let oneofCase = try container.decodeIfPresent(String.self, forKey: .oneofCase)
if oneofCase == "remainingFraction" {
self.remainingFraction = try container.decodeIfPresent(Double.self, forKey: .value)
} else {
self.remainingFraction = nil
}
}
}
@@ -0,0 +1,777 @@
import Foundation
#if canImport(FoundationNetworking)
import FoundationNetworking
#endif
public enum AntigravityRemoteFetchError: LocalizedError, Sendable, Equatable {
case notLoggedIn
case permissionDenied(String)
case apiError(String)
case parseFailed(String)
public var errorDescription: String? {
switch self {
case .notLoggedIn:
"Antigravity Google auth not found. Use Antigravity login to authenticate."
case let .permissionDenied(message):
"Antigravity remote API permission denied: \(message)"
case let .apiError(message):
"Antigravity remote API error: \(message)"
case let .parseFailed(message):
"Could not parse Antigravity remote usage: \(message)"
}
}
}
public struct AntigravityRemoteUsageFetcher: Sendable {
public var timeout: TimeInterval = 10.0
public var homeDirectory: String
public var environment: [String: String]
public var dataLoader: @Sendable (URLRequest) async throws -> (Data, URLResponse)
public var oauthClientResolver: @Sendable () -> AntigravityOAuthClient?
public var credentialsUpdateHandler: @Sendable (AntigravityOAuthCredentials) async throws -> Void
private static let log = CodexBarLog.logger(LogCategories.antigravity)
private static let userAgent = "antigravity"
private static let baseURL = "https://cloudcode-pa.googleapis.com"
private static let loadCodeAssistEndpoint = "\(baseURL)/v1internal:loadCodeAssist"
private static let onboardUserEndpoint = "\(baseURL)/v1internal:onboardUser"
private static let fetchAvailableModelsEndpoint = "\(baseURL)/v1internal:fetchAvailableModels"
private static let retrieveUserQuotaEndpoint = "\(baseURL)/v1internal:retrieveUserQuota"
private static let refreshSafetyWindow: TimeInterval = 60
private struct FetchContext {
let timeout: TimeInterval
let store: AntigravityOAuthCredentialsStore?
let dataLoader: @Sendable (URLRequest) async throws -> (Data, URLResponse)
let oauthClientResolver: @Sendable () -> AntigravityOAuthClient?
let credentialsUpdateHandler: @Sendable (AntigravityOAuthCredentials) async throws -> Void
func persistCredentials(_ credentials: AntigravityOAuthCredentials) async throws {
if let store {
try store.save(credentials)
} else {
try await self.credentialsUpdateHandler(credentials)
}
}
}
public init(
timeout: TimeInterval = 10.0,
homeDirectory: String = NSHomeDirectory(),
environment: [String: String] = ProcessInfo.processInfo.environment,
dataLoader: @escaping @Sendable (URLRequest) async throws -> (Data, URLResponse) = { request in
try await ProviderHTTPClient.shared.data(for: request)
},
oauthClientResolver: @escaping @Sendable () -> AntigravityOAuthClient? = {
AntigravityOAuthConfig.resolvedClient()
},
credentialsUpdateHandler: @escaping @Sendable (AntigravityOAuthCredentials) async throws -> Void = { _ in })
{
self.timeout = timeout
self.homeDirectory = homeDirectory
self.environment = environment
self.dataLoader = dataLoader
self.oauthClientResolver = oauthClientResolver
self.credentialsUpdateHandler = credentialsUpdateHandler
}
public func fetch() async throws -> AntigravityStatusSnapshot {
let source = try Self.resolveCredentialSource(homeDirectory: self.homeDirectory, environment: self.environment)
guard let credentials = source.credentials else {
throw AntigravityRemoteFetchError.notLoggedIn
}
return try await self.fetchSnapshot(
using: credentials,
store: source.store)
}
private func fetchSnapshot(
using initialCredentials: AntigravityOAuthCredentials,
store: AntigravityOAuthCredentialsStore?) async throws
-> AntigravityStatusSnapshot
{
guard let storedAccessToken = initialCredentials.accessToken?.trimmedNonEmpty else {
throw AntigravityRemoteFetchError.notLoggedIn
}
var credentials = initialCredentials
var accessToken = storedAccessToken
let context = FetchContext(
timeout: self.timeout,
store: store,
dataLoader: self.dataLoader,
oauthClientResolver: self.oauthClientResolver,
credentialsUpdateHandler: self.credentialsUpdateHandler)
if Self.shouldRefresh(expiryDate: credentials.expiryDate, now: Date()) {
guard let refreshToken = credentials.refreshToken?.trimmedNonEmpty else {
throw AntigravityRemoteFetchError.notLoggedIn
}
let refreshed = try await Self.refreshAccessToken(
credentials: credentials,
refreshToken: refreshToken,
context: context)
accessToken = refreshed.accessToken
credentials = refreshed.credentials
if let store {
credentials = try store.load() ?? credentials
}
credentials.accessToken = credentials.accessToken?.trimmedNonEmpty ?? accessToken
}
let claims = Self.extractClaims(from: credentials)
let codeAssist = try await Self.loadCodeAssist(
accessToken: accessToken,
timeout: self.timeout,
dataLoader: self.dataLoader)
let projectId = try await Self.resolveProjectID(
accessToken: accessToken,
storedProjectID: credentials.projectID?.trimmedNonEmpty,
initialResponse: codeAssist,
context: context)
if let projectId, credentials.projectID?.trimmedNonEmpty != projectId {
credentials.projectID = projectId
do {
try await context.persistCredentials(credentials)
} catch {
Self.log.warning("Could not persist Antigravity project ID: \(error.localizedDescription)")
}
}
let models = try await Self.fetchModelQuotas(
accessToken: accessToken,
projectId: projectId,
timeout: self.timeout,
dataLoader: self.dataLoader)
return AntigravityStatusSnapshot(
modelQuotas: models,
accountEmail: claims.email,
accountPlan: Self.resolvePlan(response: codeAssist, claims: claims),
source: .remote)
}
private static func shouldRefresh(expiryDate: Date?, now: Date) -> Bool {
guard let expiryDate else { return false }
return expiryDate.timeIntervalSince(now) <= Self.refreshSafetyWindow
}
private static func loadCodeAssist(
accessToken: String,
timeout: TimeInterval,
dataLoader: @escaping @Sendable (URLRequest) async throws -> (Data, URLResponse)) async throws
-> CodeAssistResponse
{
let body = [
"metadata": [
"ideType": "ANTIGRAVITY",
"platform": "PLATFORM_UNSPECIFIED",
"pluginType": "GEMINI",
],
]
return try await Self.sendRequest(
endpoint: Self.loadCodeAssistEndpoint,
accessToken: accessToken,
body: body,
timeout: timeout,
dataLoader: dataLoader)
}
private static func fetchAvailableModels(
accessToken: String,
projectId: String?,
timeout: TimeInterval,
dataLoader: @escaping @Sendable (URLRequest) async throws -> (Data, URLResponse)) async throws
-> FetchAvailableModelsResponse
{
let body: [String: Any] = if let projectId = projectId?.trimmedNonEmpty {
["project": projectId]
} else {
[:]
}
return try await Self.sendRequest(
endpoint: Self.fetchAvailableModelsEndpoint,
accessToken: accessToken,
body: body,
timeout: timeout,
dataLoader: dataLoader)
}
private static func fetchModelQuotas(
accessToken: String,
projectId: String?,
timeout: TimeInterval,
dataLoader: @escaping @Sendable (URLRequest) async throws -> (Data, URLResponse)) async throws
-> [AntigravityModelQuota]
{
do {
let response = try await Self.fetchAvailableModels(
accessToken: accessToken,
projectId: projectId,
timeout: timeout,
dataLoader: dataLoader)
let modelQuotas = try Self.parseModelQuotas(response)
if Self.shouldVerifyFullRemoteQuotas(modelQuotas) {
let quotaBuckets = try await Self.fetchQuotaBucketsIfPermitted(
accessToken: accessToken,
projectId: projectId,
timeout: timeout,
dataLoader: dataLoader)
guard let quotaBuckets, Self.hasQuotaFractionData(quotaBuckets) else {
return []
}
return Self.mergeVerifiedQuotas(modelQuotas: modelQuotas, verifiedQuotas: quotaBuckets)
}
return modelQuotas
} catch let error as AntigravityRemoteFetchError {
guard case .permissionDenied = error else {
throw error
}
Self.log.info("Falling back to retrieveUserQuota for Antigravity remote usage")
return try await Self.fetchQuotaBucketsIfPermitted(
accessToken: accessToken,
projectId: projectId,
timeout: timeout,
dataLoader: dataLoader) ?? []
}
}
private static func mergeVerifiedQuotas(
modelQuotas: [AntigravityModelQuota],
verifiedQuotas: [AntigravityModelQuota]) -> [AntigravityModelQuota]
{
var verifiedByModelID = Dictionary(
uniqueKeysWithValues: verifiedQuotas.map { (Self.quotaKey($0), $0) })
var merged = modelQuotas.compactMap { modelQuota -> AntigravityModelQuota? in
guard let verifiedQuota = verifiedByModelID.removeValue(forKey: Self.quotaKey(modelQuota)) else {
return nil
}
let resetTime = verifiedQuota.resetTime ?? modelQuota.resetTime
return AntigravityModelQuota(
label: modelQuota.label,
modelId: modelQuota.modelId,
remainingFraction: verifiedQuota.remainingFraction ?? modelQuota.remainingFraction,
resetTime: resetTime,
resetDescription: resetTime.map { UsageFormatter.resetDescription(from: $0) })
}
let unmatchedVerifiedQuotas = verifiedByModelID.values
.filter { $0.remainingFraction != nil }
.sorted { lhs, rhs in
lhs.modelId.localizedCaseInsensitiveCompare(rhs.modelId) == .orderedAscending
}
merged.append(contentsOf: unmatchedVerifiedQuotas)
return merged
}
private static func quotaKey(_ quota: AntigravityModelQuota) -> String {
quota.modelId.trimmingCharacters(in: .whitespacesAndNewlines).lowercased()
}
private static func shouldVerifyFullRemoteQuotas(_ quotas: [AntigravityModelQuota]) -> Bool {
guard !quotas.isEmpty else { return false }
return quotas.allSatisfy { quota in
guard let remaining = quota.remainingFraction else { return false }
return remaining >= 0.999
}
}
private static func hasQuotaFractionData(_ quotas: [AntigravityModelQuota]) -> Bool {
quotas.contains { quota in
quota.remainingFraction != nil
}
}
private static func fetchQuotaBucketsIfPermitted(
accessToken: String,
projectId: String?,
timeout: TimeInterval,
dataLoader: @escaping @Sendable (URLRequest) async throws -> (Data, URLResponse)) async throws
-> [AntigravityModelQuota]?
{
do {
let response = try await Self.retrieveUserQuota(
accessToken: accessToken,
projectId: projectId,
timeout: timeout,
dataLoader: dataLoader)
return try Self.parseQuotaBuckets(response)
} catch let quotaError as AntigravityRemoteFetchError {
guard case .permissionDenied = quotaError else {
throw quotaError
}
Self.log.info("Antigravity remote quota endpoint is not permitted for this account")
return nil
}
}
private static func retrieveUserQuota(
accessToken: String,
projectId: String?,
timeout: TimeInterval,
dataLoader: @escaping @Sendable (URLRequest) async throws -> (Data, URLResponse)) async throws
-> RetrieveUserQuotaResponse
{
let body: [String: Any] = if let projectId = projectId?.trimmedNonEmpty {
["project": projectId]
} else {
[:]
}
return try await Self.sendRequest(
endpoint: Self.retrieveUserQuotaEndpoint,
accessToken: accessToken,
body: body,
timeout: timeout,
dataLoader: dataLoader)
}
private static func resolveProjectID(
accessToken: String,
storedProjectID: String?,
initialResponse: CodeAssistResponse,
context: FetchContext) async throws
-> String?
{
if let storedProjectID {
return storedProjectID
}
if let projectID = initialResponse.projectID {
return projectID
}
guard let tierID = Self.pickOnboardTier(from: initialResponse) else {
return nil
}
let onboardBody: [String: Any] = [
"tierId": tierID,
"metadata": [
"ideType": "ANTIGRAVITY",
"platform": "PLATFORM_UNSPECIFIED",
"pluginType": "GEMINI",
],
]
do {
let onboardResponse: OnboardResponse = try await Self.sendRequest(
endpoint: Self.onboardUserEndpoint,
accessToken: accessToken,
body: onboardBody,
timeout: context.timeout,
dataLoader: context.dataLoader)
if let projectID = onboardResponse.projectID {
return projectID
}
} catch {
Self.log.warning("Antigravity onboarding request failed", metadata: [
"error": "\(error.localizedDescription)",
])
}
for _ in 0..<5 {
try? await Task.sleep(for: .milliseconds(2000))
let refreshed = try await Self.loadCodeAssist(
accessToken: accessToken,
timeout: context.timeout,
dataLoader: context.dataLoader)
if let projectID = refreshed.projectID {
return projectID
}
}
return nil
}
private static func sendRequest<Response: Decodable>(
endpoint: String,
accessToken: String,
body: [String: Any],
timeout: TimeInterval,
dataLoader: @escaping @Sendable (URLRequest) async throws -> (Data, URLResponse)) async throws
-> Response
{
guard let url = URL(string: endpoint) else {
throw AntigravityRemoteFetchError.apiError("Invalid endpoint URL")
}
var request = URLRequest(url: url)
request.httpMethod = "POST"
request.timeoutInterval = timeout
request.setValue("Bearer \(accessToken)", forHTTPHeaderField: "Authorization")
request.setValue("application/json", forHTTPHeaderField: "Content-Type")
request.setValue(Self.userAgent, forHTTPHeaderField: "User-Agent")
request.httpBody = try JSONSerialization.data(withJSONObject: body)
let httpResponse = try await ProviderHTTPTransportHandler(dataLoader).response(for: request)
switch httpResponse.statusCode {
case 200:
break
case 401:
throw AntigravityRemoteFetchError.notLoggedIn
case 403:
let message = String(data: httpResponse.data, encoding: .utf8)?.trimmedNonEmpty ?? "HTTP 403"
throw AntigravityRemoteFetchError.permissionDenied(message)
default:
let message = String(data: httpResponse.data, encoding: .utf8)?.trimmedNonEmpty
?? "HTTP \(httpResponse.statusCode)"
throw AntigravityRemoteFetchError.apiError("HTTP \(httpResponse.statusCode): \(message)")
}
do {
return try JSONDecoder().decode(Response.self, from: httpResponse.data)
} catch {
throw AntigravityRemoteFetchError.parseFailed(error.localizedDescription)
}
}
private static func parseModelQuotas(_ response: FetchAvailableModelsResponse) throws -> [AntigravityModelQuota] {
let models = response.models ?? [:]
return models.compactMap { modelID, model in
guard let quotaInfo = model.quotaInfo else { return nil }
let resetTime = quotaInfo.resetTime.flatMap(Self.parseResetTime(_:))
let label = model.displayName?.trimmedNonEmpty
?? model.label?.trimmedNonEmpty
?? modelID
return AntigravityModelQuota(
label: label,
modelId: modelID,
remainingFraction: quotaInfo.remainingFraction,
resetTime: resetTime,
resetDescription: resetTime.map { UsageFormatter.resetDescription(from: $0) })
}
}
private static func parseQuotaBuckets(_ response: RetrieveUserQuotaResponse) throws -> [AntigravityModelQuota] {
guard let buckets = response.buckets else {
throw AntigravityRemoteFetchError.parseFailed("No quota buckets in response")
}
guard !buckets.isEmpty else { return [] }
var modelQuotaMap: [String: (fraction: Double?, resetTime: String?)] = [:]
for bucket in buckets {
guard let modelID = bucket.modelId?.trimmedNonEmpty else { continue }
let next = (bucket.remainingFraction, bucket.resetTime)
if let existing = modelQuotaMap[modelID] {
let existingValue = existing.fraction ?? Double.greatestFiniteMagnitude
let nextValue = next.0 ?? Double.greatestFiniteMagnitude
if nextValue < existingValue {
modelQuotaMap[modelID] = next
}
} else {
modelQuotaMap[modelID] = next
}
}
return modelQuotaMap.keys.sorted().compactMap { modelID in
guard let info = modelQuotaMap[modelID] else { return nil }
let resetTime = info.resetTime.flatMap(Self.parseResetTime(_:))
return AntigravityModelQuota(
label: modelID,
modelId: modelID,
remainingFraction: info.fraction,
resetTime: resetTime,
resetDescription: resetTime.map { UsageFormatter.resetDescription(from: $0) })
}
}
private static func resolvePlan(response: CodeAssistResponse, claims: TokenClaims) -> String? {
if let planType = response.planInfo?.planType?.trimmedNonEmpty {
return planType
}
switch (response.currentTier?.id?.trimmedNonEmpty, claims.hostedDomain) {
case ("standard-tier", _):
return "Paid"
case ("free-tier", .some):
return "Workspace"
case ("free-tier", .none):
return "Free"
case ("legacy-tier", _):
return "Legacy"
default:
return response.currentTier?.name?.trimmedNonEmpty
}
}
private static func pickOnboardTier(from response: CodeAssistResponse) -> String? {
if let defaultTier = response.allowedTiers?
.first(where: { $0.isDefault == true && $0.id?.trimmedNonEmpty != nil })?.id?.trimmedNonEmpty
{
return defaultTier
}
if let firstTier = response.allowedTiers?
.first(where: { $0.id?.trimmedNonEmpty != nil })?.id?.trimmedNonEmpty
{
return firstTier
}
if let paidTier = response.paidTier?.id?.trimmedNonEmpty {
return paidTier
}
if let currentTier = response.currentTier?.id?.trimmedNonEmpty {
return currentTier
}
return nil
}
private static func parseResetTime(_ value: String) -> Date? {
let formatter = ISO8601DateFormatter()
formatter.formatOptions = [.withInternetDateTime, .withFractionalSeconds]
if let date = formatter.date(from: value) {
return date
}
formatter.formatOptions = [.withInternetDateTime]
return formatter.date(from: value)
}
private static func credentialsStore(homeDirectory: String) -> AntigravityOAuthCredentialsStore {
let homeURL = URL(fileURLWithPath: homeDirectory, isDirectory: true)
return AntigravityOAuthCredentialsStore(fileURL: AntigravityOAuthCredentialsStore.defaultURL(home: homeURL))
}
private static func resolveCredentialSource(
homeDirectory: String,
environment: [String: String]) throws -> (
credentials: AntigravityOAuthCredentials?,
store: AntigravityOAuthCredentialsStore?)
{
let primaryStore = Self.credentialsStore(homeDirectory: homeDirectory)
if let tokenValue = environment[AntigravityOAuthCredentialsStore.environmentCredentialsKey] {
guard let credentials = AntigravityOAuthCredentialsStore.credentials(fromTokenAccountValue: tokenValue)
else {
throw AntigravityRemoteFetchError.parseFailed("Could not decode selected account credentials.")
}
return (credentials, nil)
}
return try (primaryStore.load(), primaryStore)
}
private struct RefreshResult {
let accessToken: String
let credentials: AntigravityOAuthCredentials
}
private static func refreshAccessToken(
credentials: AntigravityOAuthCredentials,
refreshToken: String,
context: FetchContext) async throws
-> RefreshResult
{
let oauthClient = try Self.refreshOAuthClient(
from: credentials,
oauthClientResolver: context.oauthClientResolver)
var request = URLRequest(url: AntigravityOAuthConfig.tokenURL)
request.httpMethod = "POST"
request.timeoutInterval = context.timeout
request.setValue("application/x-www-form-urlencoded", forHTTPHeaderField: "Content-Type")
request.httpBody = Self.formBody([
"client_id": oauthClient.clientID,
"client_secret": oauthClient.clientSecret,
"refresh_token": refreshToken,
"grant_type": "refresh_token",
])
let httpResponse = try await ProviderHTTPTransportHandler(context.dataLoader).response(for: request)
guard httpResponse.statusCode == 200 else {
throw AntigravityRemoteFetchError.notLoggedIn
}
guard let json = try JSONSerialization.jsonObject(with: httpResponse.data) as? [String: Any],
let accessToken = json["access_token"] as? String
else {
throw AntigravityRemoteFetchError.parseFailed("Could not parse refresh response")
}
let updatedCredentials = Self.updatedCredentials(credentials, refreshResponse: json)
try await context.persistCredentials(updatedCredentials)
return RefreshResult(accessToken: accessToken, credentials: updatedCredentials)
}
private static func refreshOAuthClient(
from credentials: AntigravityOAuthCredentials,
oauthClientResolver: @escaping @Sendable () -> AntigravityOAuthClient?) throws
-> AntigravityOAuthClient
{
if let clientID = credentials.clientID?.trimmedNonEmpty,
let clientSecret = credentials.clientSecret?.trimmedNonEmpty
{
return AntigravityOAuthClient(clientID: clientID, clientSecret: clientSecret)
}
guard let client = oauthClientResolver() else {
throw AntigravityRemoteFetchError.apiError(AntigravityOAuthConfig.missingCredentialsMessage)
}
return client
}
private static func updatedCredentials(
_ credentials: AntigravityOAuthCredentials,
refreshResponse: [String: Any]) -> AntigravityOAuthCredentials
{
var credentials = credentials
if let accessToken = refreshResponse["access_token"] as? String {
credentials.accessToken = accessToken
}
if let expiresIn = refreshResponse["expires_in"] as? Double {
credentials.expiryDateMilliseconds = (Date().timeIntervalSince1970 + expiresIn) * 1000
}
if let expiresIn = refreshResponse["expires_in"] as? Int {
credentials.expiryDateMilliseconds = (Date().timeIntervalSince1970 + Double(expiresIn)) * 1000
}
if let idToken = refreshResponse["id_token"] as? String {
credentials.idToken = idToken
}
return credentials
}
private static func formBody(_ values: [String: String]) -> Data? {
var components = URLComponents()
components.queryItems = values.map { key, value in
URLQueryItem(name: key, value: value)
}
return components.query?.data(using: .utf8)
}
private struct TokenClaims {
let email: String?
let hostedDomain: String?
}
private static func extractClaims(from credentials: AntigravityOAuthCredentials) -> TokenClaims {
let tokenClaims = Self.extractClaimsFromToken(credentials.idToken)
return TokenClaims(
email: tokenClaims.email ?? credentials.email?.trimmedNonEmpty,
hostedDomain: tokenClaims.hostedDomain)
}
private static func extractClaimsFromToken(_ idToken: String?) -> TokenClaims {
guard let idToken else {
return TokenClaims(email: nil, hostedDomain: nil)
}
let parts = idToken.components(separatedBy: ".")
guard parts.count >= 2 else {
return TokenClaims(email: nil, hostedDomain: nil)
}
var payload = parts[1]
.replacingOccurrences(of: "-", with: "+")
.replacingOccurrences(of: "_", with: "/")
let remainder = payload.count % 4
if remainder > 0 {
payload += String(repeating: "=", count: 4 - remainder)
}
guard let data = Data(base64Encoded: payload, options: .ignoreUnknownCharacters),
let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any]
else {
return TokenClaims(email: nil, hostedDomain: nil)
}
return TokenClaims(
email: (json["email"] as? String)?.trimmedNonEmpty,
hostedDomain: (json["hd"] as? String)?.trimmedNonEmpty)
}
}
extension String {
fileprivate var trimmedNonEmpty: String? {
let trimmed = self.trimmingCharacters(in: .whitespacesAndNewlines)
return trimmed.isEmpty ? nil : trimmed
}
}
private struct ProjectReference: Decodable {
let value: String?
init(from decoder: Decoder) throws {
let single = try decoder.singleValueContainer()
if let stringValue = try? single.decode(String.self) {
self.value = stringValue
return
}
let keyed = try decoder.container(keyedBy: CodingKeys.self)
self.value = try keyed.decodeIfPresent(String.self, forKey: .id)
?? keyed.decodeIfPresent(String.self, forKey: .projectID)
}
private enum CodingKeys: String, CodingKey {
case id
case projectID = "projectId"
}
}
private struct CodeAssistResponse: Decodable {
let planInfo: CodeAssistPlanInfo?
let currentTier: TierInfo?
let paidTier: TierInfo?
let allowedTiers: [AllowedTier]?
let cloudaicompanionProject: ProjectReference?
var projectID: String? {
self.cloudaicompanionProject?.value?.trimmingCharacters(in: .whitespacesAndNewlines).nilIfEmpty
}
}
private struct CodeAssistPlanInfo: Decodable {
let planType: String?
}
private struct TierInfo: Decodable {
let id: String?
let name: String?
}
private struct AllowedTier: Decodable {
let id: String?
let isDefault: Bool?
}
private struct OnboardResponse: Decodable {
let response: OnboardInnerResponse?
var projectID: String? {
self.response?.cloudaicompanionProject?.value?.trimmingCharacters(in: .whitespacesAndNewlines).nilIfEmpty
}
}
private struct OnboardInnerResponse: Decodable {
let cloudaicompanionProject: ProjectReference?
}
private struct FetchAvailableModelsResponse: Decodable {
let models: [String: AntigravityRemoteModel]?
}
private struct RetrieveUserQuotaResponse: Decodable {
let buckets: [RetrieveUserQuotaBucket]?
}
private struct RetrieveUserQuotaBucket: Decodable {
let modelId: String?
let remainingFraction: Double?
let resetTime: String?
}
private struct AntigravityRemoteModel: Decodable {
let displayName: String?
let label: String?
let quotaInfo: AntigravityRemoteQuotaInfo?
}
private struct AntigravityRemoteQuotaInfo: Decodable {
let remainingFraction: Double?
let resetTime: String?
}
extension String? {
fileprivate var trimmedNonEmpty: String? {
self?.trimmingCharacters(in: .whitespacesAndNewlines).nilIfEmpty
}
}
extension String {
fileprivate var nilIfEmpty: String? {
self.isEmpty ? nil : self
}
}
@@ -0,0 +1,128 @@
import Foundation
extension AntigravityStatusProbe {
private static let processProbeLog = CodexBarLog.logger(LogCategories.antigravity)
private enum ProcessSnapshotFetchFailure {
case antigravity(AntigravityStatusProbeError)
case url(URLError)
case cancellation
case other(String)
init(_ error: Error) {
if let error = error as? AntigravityStatusProbeError {
self = .antigravity(error)
} else if let error = error as? URLError, error.code == .cancelled {
self = .cancellation
} else if let error = error as? URLError {
self = .url(error)
} else if error is CancellationError {
self = .cancellation
} else {
self = .other(error.localizedDescription)
}
}
var error: Error {
switch self {
case let .antigravity(error):
error
case let .url(error):
error
case .cancellation:
CancellationError()
case let .other(message):
AntigravityStatusProbeError.apiError(message)
}
}
}
private enum ProcessSnapshotFetchOutcome {
case success(index: Int, snapshot: AntigravityStatusSnapshot)
case failure(index: Int, pid: Int, failure: ProcessSnapshotFetchFailure)
}
static func fetchProcessSnapshots(
processInfos: [ProcessInfoResult],
fetch: @escaping @Sendable (ProcessInfoResult) async throws -> AntigravityStatusSnapshot)
async throws -> (snapshots: [AntigravityStatusSnapshot], lastError: Error?)
{
let outcomes = await withTaskGroup(of: ProcessSnapshotFetchOutcome.self) { group in
for (index, processInfo) in processInfos.enumerated() {
group.addTask {
do {
return try await .success(index: index, snapshot: fetch(processInfo))
} catch {
return .failure(index: index, pid: processInfo.pid, failure: .init(error))
}
}
}
var ordered = [ProcessSnapshotFetchOutcome?](repeating: nil, count: processInfos.count)
for await outcome in group {
switch outcome {
case let .success(index, _), let .failure(index, _, _):
ordered[index] = outcome
}
}
return ordered.compactMap(\.self)
}
var snapshots: [AntigravityStatusSnapshot] = []
var lastError: Error?
for outcome in outcomes {
switch outcome {
case let .success(_, snapshot):
snapshots.append(snapshot)
case let .failure(_, pid, failure):
if case .cancellation = failure {
throw CancellationError()
}
let error = failure.error
lastError = error
Self.processProbeLog.debug("Antigravity local process probe failed", metadata: [
"pid": "\(pid)",
"error": error.localizedDescription,
])
}
}
return (snapshots, lastError)
}
static func fetch(
processInfo: ProcessInfoResult,
timeout: TimeInterval,
deadline: Date) async throws -> AntigravityStatusSnapshot
{
guard let portTimeout = timeoutForNextAttempt(timeout: timeout, deadline: deadline) else {
throw AntigravityStatusProbeError.timedOut
}
let ports = try await Self.listeningPorts(pid: processInfo.pid, timeout: portTimeout)
let endpoint = try await Self.resolveWorkingEndpoint(
candidateEndpoints: Self.connectionCandidates(
listeningPorts: ports,
languageServerCSRFToken: processInfo.csrfToken,
extensionServerPort: processInfo.extensionPort,
extensionServerCSRFToken: processInfo.extensionServerCSRFToken),
timeout: timeout,
deadline: deadline)
let context = RequestContext(
endpoints: Self.requestEndpoints(
resolvedEndpoint: endpoint,
listeningPorts: ports,
languageServerCSRFToken: processInfo.csrfToken,
extensionServerPort: processInfo.extensionPort,
extensionServerCSRFToken: processInfo.extensionServerCSRFToken),
timeout: timeout,
deadline: deadline)
return try await Self.fetchSnapshot(context: context)
}
static func timeoutForNextAttempt(timeout: TimeInterval, deadline: Date?) -> TimeInterval? {
guard let deadline else { return timeout }
let remaining = deadline.timeIntervalSinceNow
guard remaining > 0 else { return nil }
return min(timeout, remaining)
}
}
@@ -0,0 +1,23 @@
extension AntigravityStatusProbe {
static func cliEndpoints(ports: [Int]) -> [AntigravityConnectionEndpoint] {
ports.flatMap { port in
self.localProbeSchemes.map { scheme in
AntigravityConnectionEndpoint(
scheme: scheme,
port: port,
csrfToken: "",
source: .cliHTTPS)
}
}
}
static var localProbeSchemes: [String] {
#if os(Linux)
// FoundationNetworking cannot trust Antigravity's self-signed TLS cert.
// Requests remain pinned to 127.0.0.1 in makeRequest.
["https", "http"]
#else
["https"]
#endif
}
}
@@ -0,0 +1,137 @@
import Foundation
/// Parses Linux `/proc/<pid>/net/tcp{,6}` output to recover the listening ports
/// owned by a process. The parsing is platform-independent for focused tests.
enum ProcNetTCPListeningPortParser {
/// The `st` column value for a socket in the LISTEN state.
private static let listenState = "0A"
/// Extracts the socket inode from a `/proc/<pid>/fd` symlink destination such
/// as `socket:[12345]`. Returns nil for non-socket descriptors.
static func socketInode(fromLink destination: String) -> String? {
let prefix = "socket:["
guard destination.hasPrefix(prefix), destination.hasSuffix("]") else { return nil }
let inode = destination.dropFirst(prefix.count).dropLast()
return inode.isEmpty ? nil : String(inode)
}
/// Returns the local ports of LISTEN sockets whose inode is in `socketInodes`.
///
/// `content` is the raw text of a process-scoped `tcp` or `tcp6` table. Each
/// row encodes the local endpoint as `ADDRESS:PORT` (for example,
/// `0100007F:1F90` uses port 8080) and the owning socket inode in column ten.
static func listeningPorts(_ content: String, socketInodes: Set<String>) -> Set<Int> {
var ports: Set<Int> = []
for line in content.split(separator: "\n") {
let columns = line.split(separator: " ", omittingEmptySubsequences: true)
// Columns: sl local_address rem_address st ... uid timeout inode
guard columns.count > 9,
columns[3] == self.listenState,
socketInodes.contains(String(columns[9]))
else { continue }
let localAddress = columns[1]
guard let separator = localAddress.lastIndex(of: ":"),
let port = Int(localAddress[localAddress.index(after: separator)...], radix: 16),
(0...Int(UInt16.max)).contains(port)
else { continue }
ports.insert(port)
}
return ports
}
}
extension AntigravityStatusProbe {
/// Resolves the TCP ports the process `pid` is listening on. Uses `lsof` when
/// present (the common denominator across macOS and Linux) and falls back to
/// the kernel's `/proc` interface on Linux hosts without `lsof`.
static func listeningPorts(pid: Int, timeout: TimeInterval) async throws -> [Int] {
let lsof = ["/usr/sbin/lsof", "/usr/bin/lsof"].first(where: {
FileManager.default.isExecutableFile(atPath: $0)
})
if let lsof {
return try await Self.lsofListeningPorts(lsof: lsof, pid: pid, timeout: timeout)
}
#if os(Linux)
// `lsof` is frequently absent on minimal Linux hosts. Fall back to the
// kernel's /proc interface, mirroring the /proc/<pid>/cwd fallback that
// LocalAgentSessionScanner.cwdByPID already uses when lsof is missing.
let ports = Self.procListeningPorts(pid: pid)
if ports.isEmpty {
throw AntigravityStatusProbeError.portDetectionFailed("no listening ports found")
}
return ports
#else
throw AntigravityStatusProbeError.portDetectionFailed("lsof not available")
#endif
}
private static func lsofListeningPorts(
lsof: String,
pid: Int,
timeout: TimeInterval) async throws -> [Int]
{
let env = ProcessInfo.processInfo.environment
let result: SubprocessResult
do {
result = try await SubprocessRunner.run(
binary: lsof,
arguments: ["-nP", "-iTCP", "-sTCP:LISTEN", "-a", "-p", String(pid)],
environment: env,
timeout: timeout,
label: "antigravity-lsof")
} catch let SubprocessRunnerError.nonZeroExit(code, stderr)
where code == 1 && stderr.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty
{
throw AntigravityStatusProbeError.portDetectionFailed("no listening ports found")
}
let ports = Self.parseListeningPorts(result.stdout)
if ports.isEmpty {
throw AntigravityStatusProbeError.portDetectionFailed("no listening ports found")
}
return ports
}
private static func parseListeningPorts(_ output: String) -> [Int] {
guard let regex = try? NSRegularExpression(pattern: #":(\d+)\s+\(LISTEN\)"#) else { return [] }
let range = NSRange(output.startIndex..<output.endIndex, in: output)
var ports: Set<Int> = []
regex.enumerateMatches(in: output, options: [], range: range) { match, _, _ in
guard let match,
let range = Range(match.range(at: 1), in: output),
let value = Int(output[range]) else { return }
ports.insert(value)
}
return ports.sorted()
}
/// Recovers the listening ports owned by `pid` by matching its open socket
/// inodes against the TCP tables from the same process/network namespace.
static func procListeningPorts(pid: Int, procRoot: String = "/proc") -> [Int] {
let processRoot = "\(procRoot)/\(pid)"
let inodes = Self.socketInodes(processRoot: processRoot)
guard !inodes.isEmpty else { return [] }
var ports: Set<Int> = []
for path in ["\(processRoot)/net/tcp", "\(processRoot)/net/tcp6"] {
guard let content = try? String(contentsOfFile: path, encoding: .utf8) else { continue }
ports.formUnion(ProcNetTCPListeningPortParser.listeningPorts(content, socketInodes: inodes))
}
return ports.sorted()
}
/// Collects the socket inodes referenced by the process's open descriptors.
private static func socketInodes(processRoot: String) -> Set<String> {
let fdDirectory = "\(processRoot)/fd"
guard let entries = try? FileManager.default.contentsOfDirectory(atPath: fdDirectory) else { return [] }
var inodes: Set<String> = []
for entry in entries {
guard let destination = try? FileManager.default.destinationOfSymbolicLink(
atPath: "\(fdDirectory)/\(entry)"),
let inode = ProcNetTCPListeningPortParser.socketInode(fromLink: destination)
else { continue }
inodes.insert(inode)
}
return inodes
}
}
@@ -0,0 +1,111 @@
import Foundation
struct UserStatusResponse: Decodable {
let code: CodeValue?
let message: String?
let userStatus: UserStatus?
}
struct CommandModelConfigResponse: Decodable {
let code: CodeValue?
let message: String?
let clientModelConfigs: [ModelConfig]?
}
struct UserStatus: Decodable {
let email: String?
let planStatus: PlanStatus?
let cascadeModelConfigData: ModelConfigData?
let userTier: UserTier?
}
struct UserTier: Decodable {
let id: String?
let name: String?
let description: String?
var preferredName: String? {
guard let value = self.name?.trimmingCharacters(in: .whitespacesAndNewlines) else { return nil }
return value.isEmpty ? nil : value
}
}
struct PlanStatus: Decodable {
let planInfo: PlanInfo?
}
struct PlanInfo: Decodable {
let planName: String?
let planDisplayName: String?
let displayName: String?
let productName: String?
let planShortName: String?
var preferredName: String? {
let candidates = [
self.planDisplayName,
self.displayName,
self.productName,
self.planName,
self.planShortName,
]
for candidate in candidates {
guard let value = candidate?.trimmingCharacters(in: .whitespacesAndNewlines) else { continue }
if !value.isEmpty { return value }
}
return nil
}
}
struct ModelConfigData: Decodable {
let clientModelConfigs: [ModelConfig]?
}
struct ModelConfig: Decodable {
let label: String
let modelOrAlias: ModelAlias
let quotaInfo: QuotaInfo?
}
struct ModelAlias: Decodable {
let model: String
}
struct QuotaInfo: Decodable {
let remainingFraction: Double?
let resetTime: String?
}
enum CodeValue: Decodable {
case int(Int)
case string(String)
var isOK: Bool {
switch self {
case let .int(value):
value == 0
case let .string(value):
value.lowercased() == "ok" || value.lowercased() == "success" || value == "0"
}
}
var rawValue: String {
switch self {
case let .int(value): "\(value)"
case let .string(value): value
}
}
init(from decoder: Decoder) throws {
let container = try decoder.singleValueContainer()
if let value = try? container.decode(Int.self) {
self = .int(value)
return
}
if let value = try? container.decode(String.self) {
self = .string(value)
return
}
throw DecodingError.dataCorruptedError(in: container, debugDescription: "Unsupported code type")
}
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,30 @@
import Foundation
public enum AntigravityUsageDataSource: String, CaseIterable, Identifiable, Sendable {
case auto
case oauth
case cli
public var id: String {
self.rawValue
}
public var displayName: String {
switch self {
case .auto: "Auto"
case .oauth: "Google OAuth"
case .cli: "Local API / agy CLI"
}
}
public var sourceLabel: String {
switch self {
case .auto:
"auto"
case .oauth:
"oauth"
case .cli:
"cli"
}
}
}
@@ -0,0 +1,191 @@
import Foundation
#if os(macOS)
/// Fetches Augment usage via `auggie account status` CLI command
public struct AuggieCLIProbe: Sendable {
private static let log = CodexBarLog.logger(LogCategories.auggieCLI)
public init() {}
public func fetch() async throws -> AugmentStatusSnapshot {
let output = try await self.runAuggieAccountStatus()
return try self.parse(output)
}
/// Timeout for the `auggie account status` command.
private static let commandTimeout: TimeInterval = 15
private func runAuggieAccountStatus() async throws -> String {
let env = ProcessInfo.processInfo.environment
let loginPATH = LoginShellPathCache.shared.current
let executable = BinaryLocator.resolveAuggieBinary(env: env, loginPATH: loginPATH) ?? "auggie"
var pathEnv = env
pathEnv["PATH"] = PathBuilder.effectivePATH(
purposes: [.tty, .nodeTooling],
env: env,
loginPATH: loginPATH)
let result = try await SubprocessRunner.run(
binary: executable,
arguments: ["account", "status"],
environment: pathEnv,
timeout: Self.commandTimeout,
label: "auggie-account-status")
let output = result.stdout
let errorOutput = result.stderr
guard !output.isEmpty else {
if !errorOutput.isEmpty {
Self.log.error("Auggie stderr: \(errorOutput)")
}
throw AuggieCLIError.noOutput
}
// Check for auth errors
if output.contains("Authentication failed") || output.contains("auggie login") {
throw AuggieCLIError.notAuthenticated
}
return output
}
func parse(_ output: String) throws -> AugmentStatusSnapshot {
// Legacy output:
// Max Plan 450,000 credits / month
// 11,657 remaining · 953,170 / 964,827 credits used
// 2 days remaining in this billing cycle (ends 1/8/2026)
//
// Current output (2026+):
// 319,054 credits remaining Max Plan
// 450,000 credits / month
// 9 days remaining in this billing cycle (ends 6/9/2026)
var maxCredits: Int?
var remaining: Int?
var used: Int?
var total: Int?
var billingCycleEnd: Date?
for line in output.split(separator: "\n") {
let trimmed = line.trimmingCharacters(in: .whitespaces)
if trimmed.contains("credits / month") {
if let match = trimmed.range(of: #"([\d,]+)\s+credits\s*/\s*month"#, options: .regularExpression) {
let numberStr = String(trimmed[match]).replacingOccurrences(of: ",", with: "")
.replacingOccurrences(of: " credits", with: "")
.replacingOccurrences(of: " / month", with: "")
maxCredits = Int(numberStr)
total = total ?? Int(numberStr)
}
} else if trimmed.contains("Max Plan"), trimmed.contains("credits"), !trimmed.contains("remaining") {
if let match = trimmed.range(of: #"([\d,]+)\s+credits"#, options: .regularExpression) {
let numberStr = String(trimmed[match]).replacingOccurrences(of: ",", with: "")
.replacingOccurrences(of: " credits", with: "")
maxCredits = Int(numberStr)
}
}
if trimmed.contains("credits remaining"), !trimmed.contains("billing cycle") {
if let match = trimmed.range(of: #"([\d,]+)\s+credits\s+remaining"#, options: .regularExpression) {
let numberStr = String(trimmed[match]).replacingOccurrences(of: ",", with: "")
.replacingOccurrences(of: " credits", with: "")
.replacingOccurrences(of: " remaining", with: "")
remaining = Int(numberStr)
}
}
// Parse "11,657 remaining · 953,170 / 964,827 credits used"
if trimmed.contains("remaining"), trimmed.contains("credits used") {
if let remMatch = trimmed.range(of: #"([\d,]+)\s+remaining"#, options: .regularExpression) {
let numStr = String(trimmed[remMatch])
.replacingOccurrences(of: ",", with: "")
.replacingOccurrences(of: " remaining", with: "")
remaining = Int(numStr)
}
if let usedMatch = trimmed.range(
of: #"([\d,]+)\s*/\s*([\d,]+)\s+credits used"#,
options: .regularExpression)
{
let parts = String(trimmed[usedMatch])
.replacingOccurrences(of: " credits used", with: "")
.split(separator: "/")
if parts.count == 2 {
used = Int(parts[0].replacingOccurrences(of: ",", with: "")
.trimmingCharacters(in: .whitespaces))
total = Int(parts[1].replacingOccurrences(of: ",", with: "")
.trimmingCharacters(in: .whitespaces))
}
}
}
if trimmed.contains("billing cycle"), trimmed.contains("ends") {
if let dateMatch = trimmed.range(of: #"ends\s+([\d/]+)"#, options: .regularExpression) {
let dateStr = String(trimmed[dateMatch])
.replacingOccurrences(of: "ends", with: "")
.trimmingCharacters(in: .whitespaces)
let formatter = DateFormatter()
formatter.dateFormat = "M/d/yyyy"
formatter.locale = Locale(identifier: "en_US_POSIX")
formatter.timeZone = TimeZone.current
billingCycleEnd = formatter.date(from: dateStr)
}
}
}
guard let finalRemaining = remaining else {
Self.log.error("Failed to parse auggie output: \(output)")
throw AuggieCLIError.parseError("Could not extract credits from output")
}
let finalTotal = total ?? maxCredits
guard let finalTotal else {
Self.log.error("Failed to parse auggie output: \(output)")
throw AuggieCLIError.parseError("Could not extract credits from output")
}
let finalUsed = used ?? max(0, finalTotal - finalRemaining)
return AugmentStatusSnapshot(
creditsRemaining: Double(finalRemaining),
creditsUsed: Double(finalUsed),
creditsLimit: Double(finalTotal),
billingCycleEnd: billingCycleEnd,
accountEmail: nil,
accountPlan: maxCredits.map { "\($0.formatted()) credits/month" },
rawJSON: nil)
}
}
#else
public struct AuggieCLIProbe: Sendable {
public init() {}
public func fetch() async throws -> AugmentStatusSnapshot {
throw AugmentStatusProbeError.notSupported
}
}
#endif
public enum AuggieCLIError: LocalizedError {
case noOutput
case notAuthenticated
case parseError(String)
public var errorDescription: String? {
switch self {
case .noOutput:
"Auggie CLI returned no output"
case .notAuthenticated:
"Not authenticated. Run 'auggie login' to authenticate."
case let .parseError(msg):
"Failed to parse auggie output: \(msg)"
}
}
}
@@ -0,0 +1,133 @@
import Foundation
#if os(macOS)
import SweetCookieKit
#endif
public enum AugmentProviderDescriptor {
public static let descriptor: ProviderDescriptor = Self.makeDescriptor()
static func makeDescriptor() -> ProviderDescriptor {
#if os(macOS)
// Custom browser order that includes Chrome Beta and other variants
// to support users running beta/canary versions
let browserOrder: BrowserCookieImportOrder = [
.safari,
.chrome,
.chromeBeta, // Added for Chrome Beta support
.chromeCanary, // Added for Chrome Canary support
.edge,
.edgeBeta,
.brave,
.arc,
.dia,
.arcBeta,
.firefox,
]
#else
let browserOrder: BrowserCookieImportOrder? = nil
#endif
return ProviderDescriptor(
id: .augment,
metadata: ProviderMetadata(
id: .augment,
displayName: "Augment",
sessionLabel: "Credits",
weeklyLabel: "Usage",
opusLabel: nil,
supportsOpus: false,
supportsCredits: true,
creditsHint: "Augment Code credits for AI-powered coding assistance.",
toggleTitle: "Show Augment usage",
cliName: "augment",
defaultEnabled: false,
isPrimaryProvider: false,
usesAccountFallback: false,
browserCookieOrder: browserOrder,
dashboardURL: "https://app.augmentcode.com/account/subscription",
statusPageURL: "https://status.augmentcode.com"),
branding: ProviderBranding(
iconStyle: .augment,
iconResourceName: "ProviderIcon-augment",
color: ProviderColor(red: 99 / 255, green: 102 / 255, blue: 241 / 255)),
tokenCost: ProviderTokenCostConfig(
supportsTokenCost: false,
noDataMessage: { "Augment cost summary is not supported." }),
fetchPlan: ProviderFetchPlan(
sourceModes: [.auto, .cli],
pipeline: ProviderFetchPipeline(resolveStrategies: { _ in
var strategies: [any ProviderFetchStrategy] = []
// Try CLI first (no browser prompts!)
strategies.append(AugmentCLIFetchStrategy())
// Fallback to web (browser cookies)
strategies.append(AugmentStatusFetchStrategy())
return strategies
})),
cli: ProviderCLIConfig(
name: "augment",
versionDetector: nil))
}
}
struct AugmentCLIFetchStrategy: ProviderFetchStrategy {
let id: String = "augment.cli"
let kind: ProviderFetchKind = .cli
func isAvailable(_ context: ProviderFetchContext) async -> Bool {
// Check if auggie CLI is installed
let env = ProcessInfo.processInfo.environment
let loginPATH = LoginShellPathCache.shared.current
return BinaryLocator.resolveAuggieBinary(env: env, loginPATH: loginPATH) != nil
}
func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult {
let probe = AuggieCLIProbe()
let snap = try await probe.fetch()
return self.makeResult(
usage: snap.toUsageSnapshot(),
sourceLabel: "cli")
}
func shouldFallback(on error: Error, context _: ProviderFetchContext) -> Bool {
// Fallback to web if CLI fails (not authenticated, etc.)
if let cliError = error as? AuggieCLIError {
switch cliError {
case .notAuthenticated, .noOutput, .parseError:
return true
}
}
return true
}
}
struct AugmentStatusFetchStrategy: ProviderFetchStrategy {
let id: String = "augment.web"
let kind: ProviderFetchKind = .web
func isAvailable(_ context: ProviderFetchContext) async -> Bool {
guard context.settings?.augment?.cookieSource != .off else { return false }
return true
}
func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult {
let probe = AugmentStatusProbe()
let manual = Self.manualCookieHeader(from: context)
let logger: ((String) -> Void)? = context.verbose
? { msg in CodexBarLog.logger(LogCategories.augment).verbose(msg) }
: nil
let snap = try await probe.fetch(cookieHeaderOverride: manual, logger: logger)
return self.makeResult(
usage: snap.toUsageSnapshot(),
sourceLabel: "web")
}
func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool {
false
}
private static func manualCookieHeader(from context: ProviderFetchContext) -> String? {
guard context.settings?.augment?.cookieSource == .manual else { return nil }
return CookieHeaderNormalizer.normalize(context.settings?.augment?.manualCookieHeader)
}
}
@@ -0,0 +1,491 @@
import Foundation
#if os(macOS)
import AppKit
import UserNotifications
/// Manages automatic session keepalive for Augment to prevent cookie expiration.
///
/// This actor monitors cookie expiration and proactively refreshes the session
/// before cookies expire, ensuring uninterrupted access to Augment APIs.
@MainActor
public final class AugmentSessionKeepalive {
// MARK: - Configuration
/// How often to check if session needs refresh (default: 1 minute for faster recovery)
private let checkInterval: TimeInterval = 60
/// Refresh session this many seconds before cookie expiration (default: 5 minutes)
private let refreshBufferSeconds: TimeInterval = 300
/// Minimum time between refresh attempts (default: 1 minute for faster recovery)
private let minRefreshInterval: TimeInterval = 60
/// Maximum time to wait for session refresh (default: 30 seconds)
private let refreshTimeout: TimeInterval = 30
// MARK: - State
private var timerTask: Task<Void, Never>?
private var lastRefreshAttempt: Date?
private var lastSuccessfulRefresh: Date?
private var isRefreshing = false
private let logger: ((String) -> Void)?
private var onSessionRecovered: (() async -> Void)?
/// Track consecutive failures to stop retrying after too many failures
private var consecutiveFailures = 0
private let maxConsecutiveFailures = 3 // Stop after 3 failures
private var hasGivenUp = false
// MARK: - Initialization
public init(logger: ((String) -> Void)? = nil, onSessionRecovered: (() async -> Void)? = nil) {
self.logger = logger
self.onSessionRecovered = onSessionRecovered
}
deinit {
self.timerTask?.cancel()
}
// MARK: - Public API
/// Start the automatic session keepalive timer
public func start() {
guard self.timerTask == nil else {
self.log("Keepalive already running")
return
}
self.log("🚀 Starting Augment session keepalive")
self.log(
" - Check interval: \(Int(self.checkInterval))s "
+ "(\(Self.durationDescription(seconds: self.checkInterval)))")
self.log(
" - Refresh buffer: \(Int(self.refreshBufferSeconds))s "
+ "(\(Self.durationDescription(seconds: self.refreshBufferSeconds)) before expiry)")
self.log(
" - Min refresh interval: \(Int(self.minRefreshInterval))s "
+ "(\(Self.durationDescription(seconds: self.minRefreshInterval)))")
self.timerTask = Task.detached(priority: .utility) { [weak self] in
while !Task.isCancelled {
try? await Task.sleep(for: .seconds(self?.checkInterval ?? 300))
await self?.checkAndRefreshIfNeeded()
}
}
self.log("✅ Keepalive timer started successfully")
}
/// Stop the automatic session keepalive timer
public func stop() {
self.log("Stopping Augment session keepalive")
self.timerTask?.cancel()
self.timerTask = nil
}
/// Manually trigger a session refresh (bypasses rate limiting)
public func forceRefresh() async {
self.log("Force refresh requested")
await self.performRefresh(forced: true)
}
// MARK: - Private Implementation
private func checkAndRefreshIfNeeded() async {
guard !self.isRefreshing else {
self.log("Refresh already in progress, skipping check")
return
}
// Stop trying if we've given up
if self.hasGivenUp {
self.log("⏸️ Keepalive has given up after \(self.maxConsecutiveFailures) consecutive failures")
self.log(" User must manually log in to Augment and click 'Refresh Session'")
return
}
// Rate limit: don't refresh too frequently
if let lastAttempt = self.lastRefreshAttempt {
let timeSinceLastAttempt = Date().timeIntervalSince(lastAttempt)
if timeSinceLastAttempt < self.minRefreshInterval {
self.log(
"Skipping refresh (last attempt \(Int(timeSinceLastAttempt))s ago, " +
"min interval: \(Int(self.minRefreshInterval))s)")
return
}
}
// Check if cookies are about to expire
let shouldRefresh = await self.shouldRefreshSession()
if shouldRefresh {
await self.performRefresh(forced: false)
}
}
private func shouldRefreshSession() async -> Bool {
do {
let session = try AugmentCookieImporter.importSession(logger: self.logger)
self.log("📊 Cookie Status Check:")
self.log(" Total cookies: \(session.cookies.count)")
self.log(" Source: \(session.sourceLabel)")
// Log each cookie's expiration status
for cookie in session.cookies {
if let expiry = cookie.expiresDate {
let timeUntil = expiry.timeIntervalSinceNow
let status = timeUntil > 0 ? "expires in \(Int(timeUntil))s" : "EXPIRED \(Int(-timeUntil))s ago"
self.log(" - \(cookie.name): \(status)")
} else {
self.log(" - \(cookie.name): session cookie (no expiry)")
}
}
// Find the earliest expiration date among session cookies
let expirationDates = session.cookies.compactMap(\.expiresDate)
guard !expirationDates.isEmpty else {
// Session cookies (no expiration) - refresh periodically
self.log(" All cookies are session cookies (no expiration dates)")
if let lastRefresh = self.lastSuccessfulRefresh {
let timeSinceRefresh = Date().timeIntervalSince(lastRefresh)
// Refresh every 30 minutes for session cookies
if timeSinceRefresh > 1800 {
self.log(" ⚠️ Need periodic refresh (\(Int(timeSinceRefresh))s since last refresh)")
return true
} else {
self.log(" ✅ Recently refreshed (\(Int(timeSinceRefresh))s ago)")
return false
}
} else {
// Never refreshed - do it now
self.log(" ⚠️ Never refreshed - doing initial refresh")
return true
}
}
let earliestExpiration = expirationDates.min()!
let timeUntilExpiration = earliestExpiration.timeIntervalSinceNow
let expiringCookie = session.cookies.first { $0.expiresDate == earliestExpiration }
if timeUntilExpiration < self.refreshBufferSeconds {
self.log(" ⚠️ REFRESH NEEDED:")
self.log(" Earliest expiring cookie: \(expiringCookie?.name ?? "unknown")")
self.log(" Time until expiration: \(Int(timeUntilExpiration))s")
self.log(" Refresh threshold: \(Int(self.refreshBufferSeconds))s")
return true
} else {
self.log(" ✅ Session healthy:")
self.log(" Earliest expiring cookie: \(expiringCookie?.name ?? "unknown")")
self.log(" Time until expiration: \(Int(timeUntilExpiration))s")
return false
}
} catch {
self.log("✗ Failed to check session: \(error.localizedDescription)")
return false
}
}
private func performRefresh(forced: Bool) async {
self.isRefreshing = true
self.lastRefreshAttempt = Date()
defer { self.isRefreshing = false }
self.log(forced ? "Performing forced session refresh..." : "Performing automatic session refresh...")
// If this is a forced refresh (user clicked "Refresh Session"), reset failure tracking
if forced {
self.consecutiveFailures = 0
self.hasGivenUp = false
self.log("🔄 Manual refresh - resetting failure tracking")
}
do {
// Step 1: Ping the session endpoint to trigger cookie refresh
let refreshed = try await self.pingSessionEndpoint()
if refreshed {
// Step 2: Re-import cookies from browser
try await Task.sleep(for: .seconds(1)) // Brief delay for browser to update cookies
let newSession = try AugmentCookieImporter.importSession(logger: self.logger)
await AugmentSessionStore.shared.setCookies(newSession.cookies)
CookieHeaderCache.store(
provider: .augment,
cookieHeader: newSession.cookieHeader,
sourceLabel: newSession.sourceLabel)
self.log(
"✅ Session refresh successful - imported \(newSession.cookies.count) cookies " +
"from \(newSession.sourceLabel)")
self.lastSuccessfulRefresh = Date()
// Reset failure tracking on success
self.consecutiveFailures = 0
self.hasGivenUp = false
if let callback = self.onSessionRecovered {
self.log("🔄 Triggering usage refresh after session refresh")
await callback()
}
} else {
self.log("⚠️ Session refresh returned no new cookies")
self.consecutiveFailures += 1
self.checkIfShouldGiveUp()
}
} catch AugmentSessionKeepaliveError.sessionExpired {
self.log("🔐 Session expired - attempting automatic recovery...")
self.consecutiveFailures += 1
if self.consecutiveFailures >= self.maxConsecutiveFailures {
self.log("❌ Too many consecutive failures (\(self.consecutiveFailures)) - giving up")
self.log(" User must manually log in to Augment and click 'Refresh Session'")
self.hasGivenUp = true
self.notifyUserLoginRequired()
} else {
await self.attemptSessionRecovery()
}
} catch {
self.log("✗ Session refresh failed: \(error.localizedDescription)")
self.consecutiveFailures += 1
self.checkIfShouldGiveUp()
}
}
private func checkIfShouldGiveUp() {
if self.consecutiveFailures >= self.maxConsecutiveFailures {
self.log("❌ Too many consecutive failures (\(self.consecutiveFailures)) - giving up")
self.log(" User must manually log in to Augment and click 'Refresh Session'")
self.hasGivenUp = true
self.notifyUserLoginRequired()
}
}
/// Attempt to recover from an expired session by triggering browser re-authentication
private func attemptSessionRecovery() async {
self.log("🔄 Attempting automatic session recovery...")
self.log(" Strategy: Open Augment dashboard to trigger browser re-auth")
#if os(macOS)
// Open the Augment dashboard in the default browser
// This will trigger the browser to re-authenticate if the user is still logged in
if let url = URL(string: "https://app.augmentcode.com") {
_ = await MainActor.run {
NSWorkspace.shared.open(url)
}
self.log(" ✅ Opened Augment dashboard in browser")
self.log(" ⏳ Waiting 5 seconds for browser to re-authenticate...")
// Wait for browser to potentially re-authenticate
try? await Task.sleep(for: .seconds(5))
// Try to import cookies again
do {
let newSession = try AugmentCookieImporter.importSession(logger: self.logger)
self.log(" ✅ Session recovery successful - imported \(newSession.cookies.count) cookies")
self.lastSuccessfulRefresh = Date()
// Verify the session is actually valid by pinging the API
let isValid = try await self.pingSessionEndpoint()
if isValid {
self.log(" ✅ Session verified - recovery complete!")
// Notify UsageStore to refresh Augment usage
if let callback = self.onSessionRecovered {
self.log(" 🔄 Triggering usage refresh after successful recovery")
await callback()
}
} else {
self.log(" ⚠️ Session imported but not yet valid - may need manual login")
self.notifyUserLoginRequired()
}
} catch {
self.log(" ✗ Session recovery failed: \(error.localizedDescription)")
self.log(" ️ User needs to manually log in to Augment")
self.notifyUserLoginRequired()
}
}
#else
self.log(" ✗ Automatic recovery not supported on this platform")
#endif
}
/// Notify the user that they need to log in to Augment
private func notifyUserLoginRequired() {
#if os(macOS)
self.log("📢 Sending notification: Augment session expired")
Task {
let center = UNUserNotificationCenter.current()
// Request authorization if needed
do {
let granted = try await center.requestAuthorization(options: [.alert, .sound])
guard granted else {
self.log("⚠️ Notification permission denied")
return
}
} catch {
self.log("✗ Failed to request notification permission: \(error)")
return
}
// Create notification content
let content = UNMutableNotificationContent()
content.title = "Augment Session Expired"
content.body = "Please log in to app.augmentcode.com to restore your session."
content.sound = .default
// Create trigger (deliver immediately)
let request = UNNotificationRequest(
identifier: "augment-session-expired-\(UUID().uuidString)",
content: content,
trigger: nil)
// Deliver notification
do {
try await center.add(request)
self.log("✅ Notification delivered successfully")
} catch {
self.log("✗ Failed to deliver notification: \(error)")
}
}
#endif
}
/// Ping Augment's session endpoint to trigger cookie refresh
private func pingSessionEndpoint() async throws -> Bool {
// Try to get current cookies first
let currentSession = try? AugmentCookieImporter.importSession(logger: self.logger)
guard let cookieHeader = currentSession?.cookieHeader else {
self.log("No cookies available for session ping")
return false
}
self.log("🔄 Attempting session refresh...")
self.log(" Cookies being sent: \(cookieHeader.prefix(100))...")
// Try multiple endpoints - Augment might use different auth patterns
let endpoints = [
"https://app.augmentcode.com/api/auth/session", // NextAuth pattern
"https://app.augmentcode.com/api/session", // Alternative
"https://app.augmentcode.com/api/user", // User endpoint
]
var receivedUnauthorized = false
for (index, urlString) in endpoints.enumerated() {
self.log(" Trying endpoint \(index + 1)/\(endpoints.count): \(urlString)")
guard let sessionURL = URL(string: urlString) else { continue }
var request = URLRequest(url: sessionURL)
request.timeoutInterval = self.refreshTimeout
request.setValue(cookieHeader, forHTTPHeaderField: "Cookie")
request.setValue("application/json", forHTTPHeaderField: "Accept")
request.setValue("https://app.augmentcode.com", forHTTPHeaderField: "Origin")
request.setValue("https://app.augmentcode.com", forHTTPHeaderField: "Referer")
do {
let (data, response) = try await ProviderHTTPClient.shared.data(for: request)
guard let httpResponse = response as? HTTPURLResponse else {
self.log(" ✗ Invalid response type")
continue
}
self.log(" Response: HTTP \(httpResponse.statusCode)")
// Log Set-Cookie headers if present
if let setCookies = httpResponse.allHeaderFields["Set-Cookie"] as? String {
self.log(" Set-Cookie headers received: \(setCookies.prefix(100))...")
}
if httpResponse.statusCode == 200 {
// Check if we got a valid session response
if let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] {
self.log(" JSON response keys: \(json.keys.joined(separator: ", "))")
if json["user"] != nil || json["email"] != nil || json["session"] != nil {
self.log(" ✅ Valid session data found!")
return true
} else {
self.log(" ⚠️ 200 OK but no session data in response")
// Try next endpoint
continue
}
} else {
self.log(" ⚠️ 200 OK but response is not JSON")
if let responseText = String(data: data, encoding: .utf8) {
self.log(" Response text: \(responseText.prefix(200))...")
}
continue
}
} else if httpResponse.statusCode == 401 {
self.log(" ✗ 401 Unauthorized - session expired")
receivedUnauthorized = true
// Don't throw immediately - try all endpoints first
continue
} else if httpResponse.statusCode == 404 {
self.log(" ✗ 404 Not Found - trying next endpoint")
continue
} else {
self.log(" ✗ HTTP \(httpResponse.statusCode) - trying next endpoint")
continue
}
} catch {
self.log(" ✗ Request failed: \(error.localizedDescription)")
continue
}
}
// If we got 401 from all endpoints, the session is definitely expired
if receivedUnauthorized {
self.log("⚠️ All endpoints returned 401 - session is expired")
throw AugmentSessionKeepaliveError.sessionExpired
}
self.log("⚠️ All session endpoints failed or returned no valid data")
return false
}
private static let log = CodexBarLog.logger(LogCategories.augmentKeepalive)
private static func durationDescription(seconds: TimeInterval) -> String {
let totalSeconds = max(0, Int(seconds.rounded()))
if totalSeconds >= 60, totalSeconds % 60 == 0 {
let minutes = totalSeconds / 60
return "\(minutes) minute\(minutes == 1 ? "" : "s")"
}
return "\(totalSeconds) second\(totalSeconds == 1 ? "" : "s")"
}
private func log(_ message: String) {
let timestamp = Date().formatted(date: .omitted, time: .standard)
let fullMessage = "[\(timestamp)] [AugmentKeepalive] \(message)"
self.logger?(fullMessage)
Self.log.debug(fullMessage)
}
}
// MARK: - Errors
public enum AugmentSessionKeepaliveError: LocalizedError, Sendable {
case invalidResponse
case sessionExpired
case networkError(String)
public var errorDescription: String? {
switch self {
case .invalidResponse:
"Invalid response from session endpoint"
case .sessionExpired:
"Session has expired"
case let .networkError(message):
"Network error: \(message)"
}
}
}
#endif
@@ -0,0 +1,698 @@
import Foundation
import SweetCookieKit
#if os(macOS)
private let augmentCookieImportOrder: BrowserCookieImportOrder =
ProviderDefaults.metadata[.augment]?.browserCookieOrder ?? Browser.defaultImportOrder
// MARK: - Augment Cookie Importer
/// Imports Augment session cookies from browser cookies.
public enum AugmentCookieImporter {
private static let cookieClient = BrowserCookieClient()
/// Auth0 session cookies used by Augment
/// NOTE: This list may not be exhaustive. If authentication fails with cookies present,
/// check debug logs for cookie names and report them.
private static let sessionCookieNames: Set<String> = [
"session", // Augment auth session (auth.augmentcode.com)
"_session", // Legacy session cookie (app.augmentcode.com)
"web_rpc_proxy_session", // Augment RPC proxy session
"auth0", // Auth0 session
"auth0.is.authenticated", // Auth0 authentication flag
"a0.spajs.txs", // Auth0 SPA transaction state
"__Secure-next-auth.session-token", // NextAuth secure session
"next-auth.session-token", // NextAuth session
"__Secure-authjs.session-token", // AuthJS secure session
"__Host-authjs.csrf-token", // AuthJS CSRF token
"authjs.session-token", // AuthJS session
]
public struct SessionInfo: Sendable {
public let cookies: [HTTPCookie]
public let sourceLabel: String
public init(cookies: [HTTPCookie], sourceLabel: String) {
self.cookies = cookies
self.sourceLabel = sourceLabel
}
public var cookieHeader: String {
self.cookies.map { "\($0.name)=\($0.value)" }.joined(separator: "; ")
}
/// Returns cookie header filtered for a specific target URL
public func cookieHeader(for url: URL) -> String {
guard let host = url.host else { return "" }
let matchingCookies = self.cookies.filter { cookie in
let domain = cookie.domain
// Handle wildcard domains (e.g., ".augmentcode.com")
if domain.hasPrefix(".") {
let baseDomain = String(domain.dropFirst())
return host == baseDomain || host.hasSuffix(".\(baseDomain)")
}
// Exact match or subdomain match
return host == domain || host.hasSuffix(".\(domain)")
}
return matchingCookies.map { "\($0.name)=\($0.value)" }.joined(separator: "; ")
}
}
/// Attempts to import Augment cookies using the standard browser import order.
public static func importSession(logger: ((String) -> Void)? = nil) throws -> SessionInfo {
let log: (String) -> Void = { msg in logger?("[augment-cookie] \(msg)") }
let cookieDomains = ["augmentcode.com", "app.augmentcode.com"]
for browserSource in augmentCookieImportOrder {
do {
let query = BrowserCookieQuery(domains: cookieDomains)
let sources = try Self.cookieClient.codexBarRecords(
matching: query,
in: browserSource,
logger: log)
for source in sources where !source.records.isEmpty {
let httpCookies = BrowserCookieClient.makeHTTPCookies(source.records, origin: query.origin)
// Log all cookies found for debugging
let cookieNames = httpCookies.map(\.name).joined(separator: ", ")
log("Found \(httpCookies.count) cookies in \(source.label): \(cookieNames)")
// Check if we have any session cookies
let matchingCookies = httpCookies.filter { Self.sessionCookieNames.contains($0.name) }
if !matchingCookies.isEmpty {
let matchingCookieNames = matchingCookies.map(\.name).joined(separator: ", ")
log("✓ Found Augment session cookies in \(source.label): \(matchingCookieNames)")
return SessionInfo(cookies: httpCookies, sourceLabel: source.label)
} else if !httpCookies.isEmpty {
// Log unrecognized cookies to help discover missing session cookie names
log(
"⚠️ \(source.label) has \(httpCookies.count) cookies but none match known session cookies")
log(" Cookie names found: \(cookieNames)")
log(" If you're logged into Augment, please report these cookie names")
log(" to help improve detection")
}
}
} catch {
BrowserCookieAccessGate.recordIfNeeded(error)
log("\(browserSource.displayName) cookie import failed: \(error.localizedDescription)")
}
}
throw AugmentStatusProbeError.noSessionCookie
}
/// Check if Augment session cookies are available
public static func hasSession(logger: ((String) -> Void)? = nil) -> Bool {
do {
let session = try self.importSession(logger: logger)
return !session.cookies.isEmpty
} catch {
return false
}
}
}
// MARK: - Augment API Models
public struct AugmentCreditsResponse: Codable, Sendable {
public let usageUnitsRemaining: Double?
public let usageUnitsConsumedThisBillingCycle: Double?
public let usageUnitsAvailable: Double?
public let usageBalanceStatus: String?
/// Computed properties for compatibility with existing code
public var credits: Double? {
self.usageUnitsRemaining
}
public var creditsUsed: Double? {
self.usageUnitsConsumedThisBillingCycle
}
public var creditsLimit: Double? {
if let available = self.usageUnitsAvailable, available > 0 {
return available
}
guard let remaining = self.usageUnitsRemaining,
let consumed = self.usageUnitsConsumedThisBillingCycle
else {
return nil
}
return remaining + consumed
}
private enum CodingKeys: String, CodingKey {
case usageUnitsRemaining
case usageUnitsConsumedThisBillingCycle
case usageUnitsAvailable
case usageBalanceStatus
}
}
public struct AugmentSubscriptionResponse: Codable, Sendable {
public let planName: String?
public let billingPeriodEnd: String?
public let email: String?
public let organization: String?
private enum CodingKeys: String, CodingKey {
case planName
case billingPeriodEnd
case email
case organization
}
}
// MARK: - Augment Status Snapshot
public struct AugmentStatusSnapshot: Sendable {
public let creditsRemaining: Double?
public let creditsUsed: Double?
public let creditsLimit: Double?
public let billingCycleEnd: Date?
public let accountEmail: String?
public let accountPlan: String?
public let rawJSON: String?
public init(
creditsRemaining: Double?,
creditsUsed: Double?,
creditsLimit: Double?,
billingCycleEnd: Date?,
accountEmail: String?,
accountPlan: String?,
rawJSON: String?)
{
self.creditsRemaining = creditsRemaining
self.creditsUsed = creditsUsed
self.creditsLimit = creditsLimit
self.billingCycleEnd = billingCycleEnd
self.accountEmail = accountEmail
self.accountPlan = accountPlan
self.rawJSON = rawJSON
}
public func toUsageSnapshot() -> UsageSnapshot {
let percentUsed: Double = if let used = self.creditsUsed, let limit = self.creditsLimit, limit > 0 {
(used / limit) * 100.0
} else if let remaining = self.creditsRemaining, let limit = self.creditsLimit, limit > 0 {
((limit - remaining) / limit) * 100.0
} else {
0
}
let primary = RateWindow(
usedPercent: percentUsed,
windowMinutes: nil,
resetsAt: self.billingCycleEnd,
resetDescription: self.billingCycleEnd.map { "Resets \(Self.formatResetDate($0))" })
let identity = ProviderIdentitySnapshot(
providerID: .augment,
accountEmail: self.accountEmail,
accountOrganization: nil,
loginMethod: self.accountPlan)
return UsageSnapshot(
primary: primary,
secondary: nil,
tertiary: nil,
providerCost: nil,
updatedAt: Date(),
identity: identity)
}
private static func formatResetDate(_ date: Date) -> String {
let formatter = RelativeDateTimeFormatter()
formatter.locale = Locale(identifier: "en_US")
formatter.unitsStyle = .full
return formatter.localizedString(for: date, relativeTo: Date())
}
}
// MARK: - Augment Status Probe Error
public enum AugmentStatusProbeError: LocalizedError, Sendable {
case notLoggedIn
case networkError(String)
case parseFailed(String)
case noSessionCookie
case sessionExpired
public var errorDescription: String? {
switch self {
case .notLoggedIn:
"Not logged in to Augment. Please log in via the CodexBar menu."
case let .networkError(msg):
"Augment API error: \(msg)"
case let .parseFailed(msg):
"Could not parse Augment usage: \(msg)"
case .noSessionCookie:
"No Augment session found. Please log in to app.augmentcode.com in \(augmentCookieImportOrder.loginHint)."
case .sessionExpired:
"Augment session expired. Please log in again."
}
}
}
// MARK: - Augment Session Store
public actor AugmentSessionStore {
public static let shared = AugmentSessionStore()
private var sessionCookies: [HTTPCookie] = []
private var hasLoadedFromDisk = false
private let fileURL: URL
private init() {
let fm = FileManager.default
let appSupport = fm.urls(for: .applicationSupportDirectory, in: .userDomainMask).first
?? fm.temporaryDirectory
let dir = appSupport.appendingPathComponent("CodexBar", isDirectory: true)
try? fm.createDirectory(at: dir, withIntermediateDirectories: true)
self.fileURL = dir.appendingPathComponent("augment-session.json")
// Load saved cookies on init
Task { await self.loadFromDiskIfNeeded() }
}
public func setCookies(_ cookies: [HTTPCookie]) {
self.hasLoadedFromDisk = true
self.sessionCookies = cookies
self.saveToDisk()
}
public func getCookies() -> [HTTPCookie] {
self.loadFromDiskIfNeeded()
return self.sessionCookies
}
public func clearCookies() {
self.hasLoadedFromDisk = true
self.sessionCookies = []
try? FileManager.default.removeItem(at: self.fileURL)
}
public func hasValidSession() -> Bool {
self.loadFromDiskIfNeeded()
return !self.sessionCookies.isEmpty
}
#if DEBUG
func resetForTesting(clearDisk: Bool = true) {
self.hasLoadedFromDisk = false
self.sessionCookies = []
if clearDisk {
try? FileManager.default.removeItem(at: self.fileURL)
}
}
#endif
private func loadFromDiskIfNeeded() {
guard !self.hasLoadedFromDisk else { return }
self.hasLoadedFromDisk = true
self.loadFromDisk()
}
private func saveToDisk() {
// Convert cookie properties to JSON-serializable format
// Date values must be converted to TimeInterval (Double)
let cookieData = self.sessionCookies.compactMap { cookie -> [String: Any]? in
guard let props = cookie.properties else { return nil }
var serializable: [String: Any] = [:]
for (key, value) in props {
let keyString = key.rawValue
if let date = value as? Date {
// Convert Date to TimeInterval for JSON compatibility
serializable[keyString] = date.timeIntervalSince1970
serializable[keyString + "_isDate"] = true
} else if let url = value as? URL {
serializable[keyString] = url.absoluteString
serializable[keyString + "_isURL"] = true
} else if JSONSerialization.isValidJSONObject([value]) ||
value is String ||
value is Bool ||
value is NSNumber
{
serializable[keyString] = value
}
}
return serializable
}
guard !cookieData.isEmpty,
let data = try? JSONSerialization.data(withJSONObject: cookieData, options: [.prettyPrinted])
else {
return
}
try? data.write(to: self.fileURL)
}
private func loadFromDisk() {
guard let data = try? Data(contentsOf: self.fileURL),
let cookieArray = try? JSONSerialization.jsonObject(with: data) as? [[String: Any]]
else { return }
self.sessionCookies = cookieArray.compactMap { props in
// Convert back to HTTPCookiePropertyKey dictionary
var cookieProps: [HTTPCookiePropertyKey: Any] = [:]
for (key, value) in props {
// Skip marker keys
if key.hasSuffix("_isDate") || key.hasSuffix("_isURL") { continue }
let propKey = HTTPCookiePropertyKey(key)
// Check if this was a Date
if props[key + "_isDate"] as? Bool == true, let interval = value as? TimeInterval {
cookieProps[propKey] = Date(timeIntervalSince1970: interval)
}
// Check if this was a URL
else if props[key + "_isURL"] as? Bool == true, let urlString = value as? String {
cookieProps[propKey] = URL(string: urlString)
} else {
cookieProps[propKey] = value
}
}
return HTTPCookie(properties: cookieProps)
}
}
}
// MARK: - Augment Status Probe
public struct AugmentStatusProbe: Sendable {
public let baseURL: URL
public var timeout: TimeInterval = 15.0
public init(baseURL: URL = URL(string: "https://app.augmentcode.com")!, timeout: TimeInterval = 15.0) {
self.baseURL = baseURL
self.timeout = timeout
}
/// Fetch Augment usage with manual cookie header (for debugging).
public func fetchWithManualCookies(_ cookieHeader: String) async throws -> AugmentStatusSnapshot {
try await self.fetchWithCookieHeader(cookieHeader)
}
/// Fetch Augment usage using browser cookies with fallback to stored session.
public func fetch(cookieHeaderOverride: String? = nil, logger: ((String) -> Void)? = nil)
async throws -> AugmentStatusSnapshot
{
let log: (String) -> Void = { msg in logger?("[augment] \(msg)") }
if let override = CookieHeaderNormalizer.normalize(cookieHeaderOverride) {
log("Using manual cookie header")
return try await self.fetchWithCookieHeader(override)
}
if let cached = CookieHeaderCache.load(provider: .augment),
!cached.cookieHeader.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty
{
log("Using cached cookie header from \(cached.sourceLabel)")
do {
return try await self.fetchWithCookieHeader(cached.cookieHeader)
} catch let error as AugmentStatusProbeError {
switch error {
case .notLoggedIn, .sessionExpired:
CookieHeaderCache.clear(provider: .augment)
default:
throw error
}
} catch {
throw error
}
}
// Try importing cookies from the configured browser order first.
do {
let session = try AugmentCookieImporter.importSession(logger: log)
log("Using cookies from \(session.sourceLabel)")
let snapshot = try await self.fetchWithCookieHeader(session.cookieHeader)
// SUCCESS: Save cookies to fallback store for future use
await AugmentSessionStore.shared.setCookies(session.cookies)
log("Saved session cookies to fallback store")
CookieHeaderCache.store(
provider: .augment,
cookieHeader: session.cookieHeader,
sourceLabel: session.sourceLabel)
return snapshot
} catch {
BrowserCookieAccessGate.recordIfNeeded(error)
log("Browser cookie import failed: \(error.localizedDescription)")
}
// Fall back to stored session cookies (from previous successful fetch or "Add Account" login flow)
let storedCookies = await AugmentSessionStore.shared.getCookies()
if !storedCookies.isEmpty {
log("Using stored session cookies")
let cookieHeader = storedCookies.map { "\($0.name)=\($0.value)" }.joined(separator: "; ")
do {
return try await self.fetchWithCookieHeader(cookieHeader)
} catch {
if case AugmentStatusProbeError.notLoggedIn = error {
// Clear only when auth is invalid; keep for transient failures.
await AugmentSessionStore.shared.clearCookies()
log("Stored session invalid, cleared")
} else if case AugmentStatusProbeError.sessionExpired = error {
await AugmentSessionStore.shared.clearCookies()
log("Stored session expired, cleared")
} else {
log("Stored session failed: \(error.localizedDescription)")
}
}
}
throw AugmentStatusProbeError.noSessionCookie
}
private func fetchWithCookieHeader(_ cookieHeader: String) async throws -> AugmentStatusSnapshot {
// Fetch credits (required)
let (creditsResponse, creditsJSON) = try await self.fetchCredits(cookieHeader: cookieHeader)
// Fetch subscription (optional - provides plan name and billing cycle)
let subscriptionResult: (AugmentSubscriptionResponse?, String?) = await {
do {
let (response, json) = try await self.fetchSubscription(cookieHeader: cookieHeader)
return (response, json)
} catch {
// Subscription API is optional - don't fail the whole fetch if it's unavailable
return (nil, nil)
}
}()
return self.parseResponse(
credits: creditsResponse,
subscription: subscriptionResult.0,
creditsJSON: creditsJSON,
subscriptionJSON: subscriptionResult.1)
}
private func fetchCredits(cookieHeader: String) async throws -> (AugmentCreditsResponse, String) {
let url = self.baseURL.appendingPathComponent("/api/credits")
var request = URLRequest(url: url)
request.timeoutInterval = self.timeout
request.setValue("application/json", forHTTPHeaderField: "Accept")
request.setValue(cookieHeader, forHTTPHeaderField: "Cookie")
let (data, response) = try await ProviderHTTPClient.shared.data(for: request)
guard let httpResponse = response as? HTTPURLResponse else {
throw AugmentStatusProbeError.networkError("Invalid response")
}
if httpResponse.statusCode == 401 {
// Session expired - trigger automatic recovery
throw AugmentStatusProbeError.sessionExpired
}
if httpResponse.statusCode == 403 {
let responseBody = String(data: data, encoding: .utf8) ?? ""
throw AugmentStatusProbeError.networkError("HTTP \(httpResponse.statusCode): \(responseBody)")
}
guard httpResponse.statusCode == 200 else {
let responseBody = String(data: data, encoding: .utf8) ?? ""
throw AugmentStatusProbeError.networkError("HTTP \(httpResponse.statusCode): \(responseBody)")
}
let rawJSON = String(data: data, encoding: .utf8) ?? ""
let decoder = JSONDecoder()
do {
let response = try decoder.decode(AugmentCreditsResponse.self, from: data)
return (response, rawJSON)
} catch {
throw AugmentStatusProbeError.parseFailed("Credits response: \(error.localizedDescription)")
}
}
private func fetchSubscription(cookieHeader: String) async throws -> (AugmentSubscriptionResponse, String) {
let url = self.baseURL.appendingPathComponent("/api/subscription")
var request = URLRequest(url: url)
request.timeoutInterval = self.timeout
request.setValue("application/json", forHTTPHeaderField: "Accept")
request.setValue(cookieHeader, forHTTPHeaderField: "Cookie")
let (data, response) = try await ProviderHTTPClient.shared.data(for: request)
guard let httpResponse = response as? HTTPURLResponse else {
throw AugmentStatusProbeError.networkError("Invalid response")
}
if httpResponse.statusCode == 401 {
// Session expired - trigger automatic recovery
throw AugmentStatusProbeError.sessionExpired
}
if httpResponse.statusCode == 403 {
throw AugmentStatusProbeError.notLoggedIn
}
guard httpResponse.statusCode == 200 else {
throw AugmentStatusProbeError.networkError("HTTP \(httpResponse.statusCode)")
}
let rawJSON = String(data: data, encoding: .utf8) ?? ""
let decoder = JSONDecoder()
do {
let response = try decoder.decode(AugmentSubscriptionResponse.self, from: data)
return (response, rawJSON)
} catch {
throw AugmentStatusProbeError.parseFailed("Subscription response: \(error.localizedDescription)")
}
}
private func parseResponse(
credits: AugmentCreditsResponse,
subscription: AugmentSubscriptionResponse?,
creditsJSON: String,
subscriptionJSON: String?) -> AugmentStatusSnapshot
{
// Combine both API responses for debugging
var combinedJSON = "Credits API:\n\(creditsJSON)"
if let subJSON = subscriptionJSON {
combinedJSON += "\n\nSubscription API:\n\(subJSON)"
}
// Parse billing period end date from ISO8601 string
let billingCycleEnd: Date? = {
guard let dateString = subscription?.billingPeriodEnd else { return nil }
let formatter = ISO8601DateFormatter()
return formatter.date(from: dateString)
}()
return AugmentStatusSnapshot(
creditsRemaining: credits.credits,
creditsUsed: credits.creditsUsed,
creditsLimit: credits.creditsLimit,
billingCycleEnd: billingCycleEnd,
accountEmail: subscription?.email,
accountPlan: subscription?.planName,
rawJSON: combinedJSON)
}
/// Debug probe that returns raw API responses
public func debugRawProbe(cookieHeaderOverride: String? = nil) async -> String {
let stamp = ISO8601DateFormatter().string(from: Date())
var lines: [String] = []
lines.append("=== Augment Debug Probe @ \(stamp) ===")
lines.append("")
do {
let snapshot = try await self.fetch(
cookieHeaderOverride: cookieHeaderOverride,
logger: { msg in lines.append("[log] \(msg)") })
lines.append("")
lines.append("Probe Success")
lines.append("")
lines.append("Credits Balance:")
lines.append(" Remaining: \(snapshot.creditsRemaining?.description ?? "nil")")
lines.append(" Used: \(snapshot.creditsUsed?.description ?? "nil")")
lines.append(" Limit: \(snapshot.creditsLimit?.description ?? "nil")")
lines.append("")
lines.append("Billing Cycle End: \(snapshot.billingCycleEnd?.description ?? "nil")")
lines.append("Account Email: \(snapshot.accountEmail ?? "nil")")
lines.append("Account Plan: \(snapshot.accountPlan ?? "nil")")
if let rawJSON = snapshot.rawJSON {
lines.append("")
lines.append("Raw API Response:")
lines.append(rawJSON)
}
let output = lines.joined(separator: "\n")
await MainActor.run { Self.recordDump(output) }
return output
} catch {
lines.append("")
lines.append("Probe Failed: \(error.localizedDescription)")
let output = lines.joined(separator: "\n")
await MainActor.run { Self.recordDump(output) }
return output
}
}
// MARK: - Dump storage (in-memory ring buffer)
@MainActor private static var recentDumps: [String] = []
@MainActor private static func recordDump(_ text: String) {
if self.recentDumps.count >= 5 { self.recentDumps.removeFirst() }
self.recentDumps.append(text)
}
public static func latestDumps() async -> String {
await MainActor.run {
let result = Self.recentDumps.joined(separator: "\n\n---\n\n")
return result.isEmpty ? "No Augment probe dumps captured yet." : result
}
}
}
#else
// MARK: - Augment (Unsupported)
public enum AugmentStatusProbeError: LocalizedError, Sendable {
case notSupported
public var errorDescription: String? {
"Augment is only supported on macOS."
}
}
public struct AugmentStatusSnapshot: Sendable {
public init() {}
public func toUsageSnapshot() -> UsageSnapshot {
UsageSnapshot(
primary: RateWindow(usedPercent: 0, windowMinutes: nil, resetsAt: nil, resetDescription: nil),
secondary: nil,
tertiary: nil,
providerCost: nil,
updatedAt: Date(),
identity: nil)
}
}
public struct AugmentStatusProbe: Sendable {
public init(baseURL: URL = URL(string: "https://app.augmentcode.com")!, timeout: TimeInterval = 15.0) {
_ = baseURL
_ = timeout
}
public func fetch(cookieHeaderOverride: String? = nil, logger: ((String) -> Void)? = nil)
async throws -> AugmentStatusSnapshot
{
_ = cookieHeaderOverride
_ = logger
throw AugmentStatusProbeError.notSupported
}
}
#endif
@@ -0,0 +1,92 @@
import Foundation
public enum AzureOpenAIProviderDescriptor {
public static let descriptor: ProviderDescriptor = Self.makeDescriptor()
static func makeDescriptor() -> ProviderDescriptor {
ProviderDescriptor(
id: .azureopenai,
metadata: ProviderMetadata(
id: .azureopenai,
displayName: "Azure OpenAI",
sessionLabel: "Status",
weeklyLabel: "Deployment",
opusLabel: nil,
supportsOpus: false,
supportsCredits: false,
creditsHint: "",
toggleTitle: "Show Azure OpenAI status",
cliName: "azure-openai",
defaultEnabled: false,
isPrimaryProvider: false,
usesAccountFallback: false,
browserCookieOrder: nil,
dashboardURL: "https://ai.azure.com",
statusPageURL: nil,
statusLinkURL: "https://azure.status.microsoft/en-us/status"),
branding: ProviderBranding(
iconStyle: .openai,
iconResourceName: "ProviderIcon-codex",
color: ProviderColor(red: 0, green: 120 / 255, blue: 212 / 255)),
tokenCost: ProviderTokenCostConfig(
supportsTokenCost: false,
noDataMessage: { "Azure OpenAI usage history is not exposed by the deployment validation probe." }),
fetchPlan: ProviderFetchPlan(
sourceModes: [.auto, .api],
pipeline: ProviderFetchPipeline(resolveStrategies: { _ in [AzureOpenAIAPIFetchStrategy()] })),
cli: ProviderCLIConfig(
name: "azure-openai",
aliases: ["azureopenai", "aoai"],
versionDetector: nil))
}
}
struct AzureOpenAIAPIFetchStrategy: ProviderFetchStrategy {
let id: String = "azureopenai.api"
let kind: ProviderFetchKind = .apiToken
func isAvailable(_ context: ProviderFetchContext) async -> Bool {
_ = context
// Keep the strategy available so missing partial configuration surfaces
// as a precise settings error instead of a generic no-strategy failure.
return true
}
func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult {
guard let apiKey = Self.resolveAPIKey(environment: context.env) else {
throw AzureOpenAIUsageError.missingAPIKey
}
try AzureOpenAISettingsReader.validateEndpointOverrides(environment: context.env)
guard let endpoint = Self.resolveEndpoint(environment: context.env) else {
throw AzureOpenAIUsageError.missingEndpoint
}
guard let deploymentName = Self.resolveDeploymentName(environment: context.env) else {
throw AzureOpenAIUsageError.missingDeploymentName
}
let usage = try await AzureOpenAIUsageFetcher.fetchUsage(
apiKey: apiKey,
endpoint: endpoint,
deploymentName: deploymentName,
apiVersion: AzureOpenAISettingsReader.apiVersion(environment: context.env))
return self.makeResult(
usage: usage.toUsageSnapshot(),
sourceLabel: "deployment")
}
func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool {
false
}
private static func resolveAPIKey(environment: [String: String]) -> String? {
ProviderTokenResolver.azureOpenAIToken(environment: environment)
}
private static func resolveEndpoint(environment: [String: String]) -> URL? {
AzureOpenAISettingsReader.endpoint(environment: environment)
}
private static func resolveDeploymentName(environment: [String: String]) -> String? {
AzureOpenAISettingsReader.deploymentName(environment: environment)
}
}
@@ -0,0 +1,82 @@
import Foundation
public enum AzureOpenAISettingsReader {
public static let apiKeyEnvironmentKey = "AZURE_OPENAI_API_KEY"
public static let endpointEnvironmentKey = "AZURE_OPENAI_ENDPOINT"
public static let deploymentNameEnvironmentKey = "AZURE_OPENAI_DEPLOYMENT_NAME"
public static let apiVersionEnvironmentKey = "AZURE_OPENAI_API_VERSION"
public static let defaultAPIVersion = "2024-10-21"
public static func apiKey(environment: [String: String] = ProcessInfo.processInfo.environment) -> String? {
self.cleaned(environment[self.apiKeyEnvironmentKey])
}
public static func endpoint(environment: [String: String] = ProcessInfo.processInfo.environment) -> URL? {
guard let rawEndpoint = self.rawEndpoint(environment: environment) else { return nil }
return self.endpointURL(from: rawEndpoint)
}
public static func rawEndpoint(environment: [String: String] = ProcessInfo.processInfo.environment) -> String? {
self.cleaned(environment[self.endpointEnvironmentKey])
}
public static func deploymentName(environment: [String: String] = ProcessInfo.processInfo.environment) -> String? {
self.cleaned(environment[self.deploymentNameEnvironmentKey])
}
public static func apiVersion(environment: [String: String] = ProcessInfo.processInfo.environment) -> String {
self.cleaned(environment[self.apiVersionEnvironmentKey]) ?? self.defaultAPIVersion
}
public static func endpointURL(from rawEndpoint: String) -> URL? {
let trimmed = rawEndpoint.trimmingCharacters(in: .whitespacesAndNewlines)
guard !trimmed.isEmpty else { return nil }
return ProviderEndpointOverrideValidator.normalizedHTTPSURL(from: trimmed)
}
public static func validateEndpointOverrides(
environment: [String: String] = ProcessInfo.processInfo.environment) throws
{
guard let rawEndpoint = self.rawEndpoint(environment: environment) else { return }
guard self.endpointURL(from: rawEndpoint) != nil else {
throw AzureOpenAISettingsError.invalidEndpointOverride(self.endpointEnvironmentKey)
}
}
static func cleaned(_ raw: String?) -> String? {
guard var value = raw?.trimmingCharacters(in: .whitespacesAndNewlines), !value.isEmpty else {
return nil
}
if (value.hasPrefix("\"") && value.hasSuffix("\"")) ||
(value.hasPrefix("'") && value.hasSuffix("'"))
{
value = String(value.dropFirst().dropLast())
}
value = value.trimmingCharacters(in: .whitespacesAndNewlines)
return value.isEmpty ? nil : value
}
}
public enum AzureOpenAISettingsError: LocalizedError, Sendable, Equatable {
case missingAPIKey
case missingEndpoint
case missingDeploymentName
case invalidEndpointOverride(String)
public var errorDescription: String? {
switch self {
case .missingAPIKey:
"Azure OpenAI API key not configured. Set AZURE_OPENAI_API_KEY or configure an API key in Settings."
case .missingEndpoint:
"Azure OpenAI endpoint not configured. Set AZURE_OPENAI_ENDPOINT or configure an endpoint in Settings."
case .missingDeploymentName:
"Azure OpenAI deployment not configured. Set AZURE_OPENAI_DEPLOYMENT_NAME or configure a deployment " +
"in Settings."
case let .invalidEndpointOverride(key):
"Azure OpenAI endpoint override \(key) is not allowed. " +
"Use an HTTPS endpoint without user info or encoded host tricks."
}
}
}
@@ -0,0 +1,244 @@
import Foundation
#if canImport(FoundationNetworking)
import FoundationNetworking
#endif
public struct AzureOpenAIUsageSnapshot: Codable, Sendable, Equatable {
public let endpointHost: String
public let deploymentName: String
public let model: String?
public let apiVersion: String
public let updatedAt: Date
public init(
endpointHost: String,
deploymentName: String,
model: String?,
apiVersion: String,
updatedAt: Date)
{
self.endpointHost = endpointHost
self.deploymentName = deploymentName
self.model = model
self.apiVersion = apiVersion
self.updatedAt = updatedAt
}
public func toUsageSnapshot() -> UsageSnapshot {
let detail = Self.detailText(deploymentName: self.deploymentName, model: self.model)
return UsageSnapshot(
primary: RateWindow(
usedPercent: 0,
windowMinutes: nil,
resetsAt: nil,
resetDescription: detail),
secondary: nil,
tertiary: nil,
updatedAt: self.updatedAt,
identity: ProviderIdentitySnapshot(
providerID: .azureopenai,
accountEmail: nil,
accountOrganization: self.endpointHost,
loginMethod: "Deployment: \(self.deploymentName)"))
}
private static func detailText(deploymentName: String, model: String?) -> String {
let cleanedModel = model?.trimmingCharacters(in: .whitespacesAndNewlines)
guard let cleanedModel, !cleanedModel.isEmpty else {
return "Deployment: \(deploymentName)"
}
return "Deployment: \(deploymentName) · Model: \(cleanedModel)"
}
}
public enum AzureOpenAIUsageError: LocalizedError, Sendable, Equatable {
case missingAPIKey
case missingEndpoint
case missingDeploymentName
case invalidEndpointOverride(String)
case invalidURL
case networkError(String)
case apiError(statusCode: Int, message: String)
case parseFailed(String)
public var errorDescription: String? {
switch self {
case .missingAPIKey:
AzureOpenAISettingsError.missingAPIKey.errorDescription
case .missingEndpoint:
AzureOpenAISettingsError.missingEndpoint.errorDescription
case .missingDeploymentName:
AzureOpenAISettingsError.missingDeploymentName.errorDescription
case let .invalidEndpointOverride(key):
AzureOpenAISettingsError.invalidEndpointOverride(key).errorDescription
case .invalidURL:
"Azure OpenAI validation URL is invalid."
case let .networkError(message):
"Azure OpenAI network error: \(message)"
case let .apiError(statusCode, message):
if message.isEmpty {
"Azure OpenAI API error: HTTP \(statusCode)"
} else {
"Azure OpenAI API error: HTTP \(statusCode): \(message)"
}
case let .parseFailed(message):
"Azure OpenAI response parse error: \(message)"
}
}
}
private struct AzureOpenAIChatCompletionResponse: Decodable {
let model: String?
}
public enum AzureOpenAIUsageFetcher {
private static let timeoutSeconds: TimeInterval = 20
private static let maxErrorBodyLength = 240
public static func fetchUsage(
apiKey: String,
endpoint: URL,
deploymentName: String,
apiVersion: String = AzureOpenAISettingsReader.defaultAPIVersion,
transport: any ProviderHTTPTransport = ProviderHTTPClient.shared,
updatedAt: Date = Date()) async throws -> AzureOpenAIUsageSnapshot
{
let apiKey = apiKey.trimmingCharacters(in: .whitespacesAndNewlines)
let deploymentName = deploymentName.trimmingCharacters(in: .whitespacesAndNewlines)
let apiVersion = apiVersion.trimmingCharacters(in: .whitespacesAndNewlines)
guard !apiKey.isEmpty else { throw AzureOpenAIUsageError.missingAPIKey }
guard !deploymentName.isEmpty else { throw AzureOpenAIUsageError.missingDeploymentName }
guard let endpoint = ProviderEndpointOverrideValidator.normalizedHTTPSURL(from: endpoint.absoluteString) else {
throw AzureOpenAIUsageError.invalidEndpointOverride(AzureOpenAISettingsReader.endpointEnvironmentKey)
}
let effectiveAPIVersion = apiVersion.isEmpty ? AzureOpenAISettingsReader.defaultAPIVersion : apiVersion
var request = try URLRequest(url: self.chatCompletionsURL(
endpoint: endpoint,
deploymentName: deploymentName,
apiVersion: effectiveAPIVersion))
request.httpMethod = "POST"
request.timeoutInterval = Self.timeoutSeconds
request.setValue(apiKey, forHTTPHeaderField: "api-key")
request.setValue("application/json", forHTTPHeaderField: "Accept")
request.setValue("application/json", forHTTPHeaderField: "Content-Type")
request.httpBody = try self.validationRequestBody(
deploymentName: deploymentName,
apiVersion: effectiveAPIVersion)
let response: ProviderHTTPResponse
do {
response = try await transport.response(for: request)
} catch {
throw AzureOpenAIUsageError.networkError(error.localizedDescription)
}
guard (200..<300).contains(response.statusCode) else {
throw AzureOpenAIUsageError.apiError(
statusCode: response.statusCode,
message: self.responseSummary(response.data))
}
let model: String?
do {
model = try JSONDecoder().decode(AzureOpenAIChatCompletionResponse.self, from: response.data).model
} catch {
throw AzureOpenAIUsageError.parseFailed(error.localizedDescription)
}
return AzureOpenAIUsageSnapshot(
endpointHost: endpoint.host ?? endpoint.absoluteString,
deploymentName: deploymentName,
model: model,
apiVersion: effectiveAPIVersion,
updatedAt: updatedAt)
}
public static func _chatCompletionsURLForTesting(
endpoint: URL,
deploymentName: String,
apiVersion: String) throws -> URL
{
try self.chatCompletionsURL(endpoint: endpoint, deploymentName: deploymentName, apiVersion: apiVersion)
}
private static func chatCompletionsURL(
endpoint: URL,
deploymentName: String,
apiVersion: String) throws -> URL
{
if self.usesV1API(apiVersion) {
let base = self.apiRoot(endpoint: endpoint, pathComponents: ["openai", "v1"])
.appendingPathComponent("chat")
.appendingPathComponent("completions")
guard let url = URLComponents(url: base, resolvingAgainstBaseURL: false)?.url else {
throw AzureOpenAIUsageError.invalidURL
}
return url
}
let base = self.apiRoot(endpoint: endpoint, pathComponents: ["openai"])
.appendingPathComponent("deployments")
.appendingPathComponent(deploymentName)
.appendingPathComponent("chat")
.appendingPathComponent("completions")
guard var components = URLComponents(url: base, resolvingAgainstBaseURL: false) else {
throw AzureOpenAIUsageError.invalidURL
}
components.queryItems = [URLQueryItem(name: "api-version", value: apiVersion)]
guard let url = components.url else { throw AzureOpenAIUsageError.invalidURL }
return url
}
private static func apiRoot(endpoint: URL, pathComponents expectedComponents: [String]) -> URL {
let existingComponents = endpoint.pathComponents
.filter { $0 != "/" }
.map { $0.lowercased() }
let expectedComponents = expectedComponents.map { $0.lowercased() }
let sharedCount = stride(
from: min(existingComponents.count, expectedComponents.count),
through: 0,
by: -1)
.first { count in
count == 0 || Array(existingComponents.suffix(count)) == Array(expectedComponents.prefix(count))
} ?? 0
return expectedComponents.dropFirst(sharedCount).reduce(endpoint) { url, component in
url.appendingPathComponent(component)
}
}
private static func validationRequestBody(
deploymentName: String,
apiVersion: String) throws -> Data
{
var payload: [String: Any] = [
"messages": [
["role": "user", "content": "ping"],
],
]
if self.usesV1API(apiVersion) {
payload["model"] = deploymentName
payload["max_completion_tokens"] = 1
} else {
payload["max_tokens"] = 1
}
return try JSONSerialization.data(withJSONObject: payload, options: [])
}
private static func usesV1API(_ apiVersion: String) -> Bool {
apiVersion.trimmingCharacters(in: .whitespacesAndNewlines).lowercased() == "v1"
}
private static func responseSummary(_ data: Data) -> String {
guard let body = String(data: data, encoding: .utf8)?
.replacingOccurrences(of: #"\s+"#, with: " ", options: .regularExpression)
.trimmingCharacters(in: .whitespacesAndNewlines),
!body.isEmpty
else {
return ""
}
guard body.count > Self.maxErrorBodyLength else { return body }
let index = body.index(body.startIndex, offsetBy: Self.maxErrorBodyLength)
return "\(body[..<index])… [truncated]"
}
}
@@ -0,0 +1,179 @@
#if canImport(CryptoKit)
import CryptoKit
#else
import Crypto
#endif
import Foundation
#if canImport(FoundationNetworking)
import FoundationNetworking
#endif
enum BedrockAWSSigner {
struct Credentials {
let accessKeyID: String
let secretAccessKey: String
let sessionToken: String?
}
static func sign(
request: inout URLRequest,
credentials: Credentials,
region: String,
service: String,
date: Date = Date())
{
let dateFormatter = Self.dateFormatter()
let dateStamp = Self.dateStamp(date: date)
let amzDate = dateFormatter.string(from: date)
request.setValue(amzDate, forHTTPHeaderField: "X-Amz-Date")
if let sessionToken = credentials.sessionToken {
request.setValue(sessionToken, forHTTPHeaderField: "X-Amz-Security-Token")
}
let host = request.url?.host ?? ""
request.setValue(host, forHTTPHeaderField: "Host")
let bodyHash = Self.sha256Hex(request.httpBody ?? Data())
request.setValue(bodyHash, forHTTPHeaderField: "x-amz-content-sha256")
let signedHeaders = Self.signedHeaders(request: request)
let canonicalRequest = Self.canonicalRequest(
request: request,
signedHeaders: signedHeaders,
bodyHash: bodyHash)
let credentialScope = "\(dateStamp)/\(region)/\(service)/aws4_request"
let stringToSign = [
"AWS4-HMAC-SHA256",
amzDate,
credentialScope,
Self.sha256Hex(Data(canonicalRequest.utf8)),
].joined(separator: "\n")
let signature = Self.calculateSignature(
secretKey: credentials.secretAccessKey,
dateStamp: dateStamp,
region: region,
service: service,
stringToSign: stringToSign)
let authorization = "AWS4-HMAC-SHA256 "
+ "Credential=\(credentials.accessKeyID)/\(credentialScope), "
+ "SignedHeaders=\(signedHeaders.keys), "
+ "Signature=\(signature)"
request.setValue(authorization, forHTTPHeaderField: "Authorization")
}
private struct SignedHeadersInfo {
let keys: String
let canonical: String
}
private static func signedHeaders(request: URLRequest) -> SignedHeadersInfo {
var headers: [(String, String)] = []
if let allHeaders = request.allHTTPHeaderFields {
for (key, value) in allHeaders {
headers.append((key.lowercased(), value.trimmingCharacters(in: .whitespaces)))
}
}
headers.sort { $0.0 < $1.0 }
let keys = headers.map(\.0).joined(separator: ";")
let canonical = headers.map { "\($0.0):\($0.1)" }.joined(separator: "\n")
return SignedHeadersInfo(keys: keys, canonical: canonical)
}
private static func canonicalRequest(
request: URLRequest,
signedHeaders: SignedHeadersInfo,
bodyHash: String) -> String
{
let method = request.httpMethod ?? "GET"
let url = request.url!
let path = url.path.isEmpty ? "/" : url.path
let query = Self.canonicalQueryString(url: url)
return [
method,
Self.uriEncodePath(path),
query,
signedHeaders.canonical + "\n",
signedHeaders.keys,
bodyHash,
].joined(separator: "\n")
}
private static func canonicalQueryString(url: URL) -> String {
guard let components = URLComponents(url: url, resolvingAgainstBaseURL: false),
let queryItems = components.queryItems,
!queryItems.isEmpty
else {
return ""
}
return queryItems
.map { item in
let key = Self.uriEncode(item.name)
let value = Self.uriEncode(item.value ?? "")
return "\(key)=\(value)"
}
.sorted()
.joined(separator: "&")
}
private static func calculateSignature(
secretKey: String,
dateStamp: String,
region: String,
service: String,
stringToSign: String) -> String
{
let kDate = Self.hmacSHA256(key: Data("AWS4\(secretKey)".utf8), data: Data(dateStamp.utf8))
let kRegion = Self.hmacSHA256(key: kDate, data: Data(region.utf8))
let kService = Self.hmacSHA256(key: kRegion, data: Data(service.utf8))
let kSigning = Self.hmacSHA256(key: kService, data: Data("aws4_request".utf8))
let signature = Self.hmacSHA256(key: kSigning, data: Data(stringToSign.utf8))
return signature.map { String(format: "%02x", $0) }.joined()
}
private static func hmacSHA256(key: Data, data: Data) -> Data {
let symmetricKey = SymmetricKey(data: key)
let mac = HMAC<SHA256>.authenticationCode(for: data, using: symmetricKey)
return Data(mac)
}
private static func sha256Hex(_ data: Data) -> String {
let digest = SHA256.hash(data: data)
return digest.map { String(format: "%02x", $0) }.joined()
}
private static func dateFormatter() -> DateFormatter {
let formatter = DateFormatter()
formatter.dateFormat = "yyyyMMdd'T'HHmmss'Z'"
formatter.timeZone = TimeZone(identifier: "UTC")
formatter.locale = Locale(identifier: "en_US_POSIX")
return formatter
}
private static func dateStamp(date: Date) -> String {
let formatter = DateFormatter()
formatter.dateFormat = "yyyyMMdd"
formatter.timeZone = TimeZone(identifier: "UTC")
formatter.locale = Locale(identifier: "en_US_POSIX")
return formatter.string(from: date)
}
private static func uriEncode(_ string: String) -> String {
var allowed = CharacterSet.alphanumerics
allowed.insert(charactersIn: "-._~")
return string.addingPercentEncoding(withAllowedCharacters: allowed) ?? string
}
private static func uriEncodePath(_ path: String) -> String {
path.split(separator: "/", omittingEmptySubsequences: false)
.map { self.uriEncode(String($0)) }
.joined(separator: "/")
}
}
@@ -0,0 +1,231 @@
import Foundation
#if canImport(FoundationNetworking)
import FoundationNetworking
#endif
struct BedrockClaudeActivity: Equatable {
let inputTokens: Int
let outputTokens: Int
let requestCount: Int
}
enum BedrockCloudWatchUsageFetcher {
static let lookbackDays = 14
private static let maxResponseBytes = 4 * 1024 * 1024
private static let maxPages = 20
private static let requestTimeoutSeconds: TimeInterval = 15
private enum Metric: String, CaseIterable {
case inputTokens
case outputTokens
case requests
var cloudWatchName: String {
switch self {
case .inputTokens: "InputTokenCount"
case .outputTokens: "OutputTokenCount"
case .requests: "Invocations"
}
}
}
private struct RequestContext {
let credentials: BedrockAWSSigner.Credentials
let region: String
let now: Date
let endpoint: URL
let transport: any ProviderHTTPTransport
}
static func fetch(
credentials: BedrockAWSSigner.Credentials,
region: String,
now: Date,
endpointOverride: String?,
transport: any ProviderHTTPTransport) async throws -> BedrockClaudeActivity
{
let totals = try await self.fetchTotals(
credentials: credentials,
region: region,
now: now,
endpointOverride: endpointOverride,
transport: transport)
return BedrockClaudeActivity(
inputTokens: totals[.inputTokens] ?? 0,
outputTokens: totals[.outputTokens] ?? 0,
requestCount: totals[.requests] ?? 0)
}
private static func fetchTotals(
credentials: BedrockAWSSigner.Credentials,
region: String,
now: Date,
endpointOverride: String?,
transport: any ProviderHTTPTransport) async throws -> [Metric: Int]
{
let context = try RequestContext(
credentials: credentials,
region: region,
now: now,
endpoint: self.endpoint(region: region, override: endpointOverride),
transport: transport)
var totals: [Metric: Double] = [:]
var nextToken: String?
var seenTokens: Set<String> = []
var pageCount = 0
repeat {
pageCount += 1
guard pageCount <= self.maxPages else {
throw BedrockUsageError.cloudWatchParseFailed("too many response pages")
}
let response = try await self.callPage(
context: context,
nextToken: nextToken)
for (metric, value) in response.totals {
totals[metric, default: 0] += value
}
nextToken = response.nextToken
if let nextToken, !seenTokens.insert(nextToken).inserted {
throw BedrockUsageError.cloudWatchParseFailed("repeated NextToken")
}
} while nextToken != nil
var converted: [Metric: Int] = [:]
for metric in Metric.allCases {
let total = totals[metric] ?? 0
guard total.isFinite, total >= 0, total <= Double(Int.max) else {
throw BedrockUsageError.cloudWatchParseFailed("invalid metric total")
}
converted[metric] = Int(total.rounded())
}
return converted
}
private static func callPage(
context: RequestContext,
nextToken: String?) async throws
-> (totals: [Metric: Double], nextToken: String?)
{
let start = context.now.addingTimeInterval(-Double(self.lookbackDays) * 24 * 60 * 60)
let queries: [[String: Any]] = Metric.allCases.map { metric in
let search = "SEARCH('{AWS/Bedrock,ModelId} "
+ "MetricName=\"\(metric.cloudWatchName)\" claude', 'Sum', 86400)"
return [
"Id": metric.rawValue,
"Expression": "SUM(\(search))",
"ReturnData": true,
]
}
var payload: [String: Any] = [
"StartTime": start.timeIntervalSince1970,
"EndTime": context.now.timeIntervalSince1970,
"ScanBy": "TimestampAscending",
"MetricDataQueries": queries,
]
if let nextToken {
payload["NextToken"] = nextToken
}
var request = URLRequest(url: context.endpoint)
request.httpMethod = "POST"
request.httpBody = try JSONSerialization.data(withJSONObject: payload)
request.timeoutInterval = self.requestTimeoutSeconds
request.setValue("application/x-amz-json-1.0", forHTTPHeaderField: "Content-Type")
request.setValue(
"GraniteServiceVersion20100801.GetMetricData",
forHTTPHeaderField: "X-Amz-Target")
BedrockAWSSigner.sign(
request: &request,
credentials: context.credentials,
region: context.region,
service: "monitoring")
let response = try await context.transport.response(for: request)
guard response.statusCode == 200 else {
throw BedrockUsageError.cloudWatchAPIError("HTTP \(response.statusCode)")
}
guard response.data.count <= self.maxResponseBytes else {
throw BedrockUsageError.cloudWatchParseFailed("response exceeds 4 MiB")
}
return try self.parsePage(response.data)
}
private static func endpoint(region: String, override: String?) throws -> URL {
if override != nil {
guard let override = BedrockSettingsReader.cleaned(override),
let url = ProviderEndpointOverrideValidator()
.validatedURLAllowingLoopbackHTTP(override)
else {
throw BedrockUsageError.cloudWatchParseFailed("invalid endpoint override")
}
return url
}
guard region.range(
of: #"^[a-z0-9]+(?:-[a-z0-9]+)+-[0-9]+$"#,
options: .regularExpression) != nil
else {
throw BedrockUsageError.cloudWatchParseFailed("invalid region endpoint")
}
// Match the CloudWatch partition suffixes published by the AWS SDK endpoint resolver.
let suffix = switch region {
case let region where region.hasPrefix("cn-"):
"amazonaws.com.cn"
case let region where region.hasPrefix("eusc-"):
"amazonaws.eu"
case let region where region.hasPrefix("us-iso-"):
"c2s.ic.gov"
case let region where region.hasPrefix("us-isob-"):
"sc2s.sgov.gov"
case let region where region.hasPrefix("eu-isoe-"):
"cloud.adc-e.uk"
case let region where region.hasPrefix("us-isof-"):
"csp.hci.ic.gov"
default:
"amazonaws.com"
}
guard let url = URL(string: "https://monitoring.\(region).\(suffix)") else {
throw BedrockUsageError.cloudWatchParseFailed("invalid region endpoint")
}
return url
}
private static func parsePage(_ data: Data) throws
-> (totals: [Metric: Double], nextToken: String?)
{
guard let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] else {
throw BedrockUsageError.cloudWatchParseFailed("invalid JSON response")
}
if let messages = json["Messages"] as? [[String: Any]], !messages.isEmpty {
throw BedrockUsageError.cloudWatchParseFailed("CloudWatch reported incomplete results")
}
let results = json["MetricDataResults"] as? [[String: Any]] ?? []
var totals: [Metric: Double] = [:]
for result in results {
guard let id = result["Id"] as? String, let metric = Metric(rawValue: id) else {
throw BedrockUsageError.cloudWatchParseFailed("metric result had an unknown ID")
}
guard result["StatusCode"] as? String == "Complete" else {
throw BedrockUsageError.cloudWatchParseFailed("metric result was incomplete")
}
guard let values = result["Values"] as? [Any] else { continue }
for value in values {
guard let number = value as? NSNumber else {
throw BedrockUsageError.cloudWatchParseFailed("metric value was not numeric")
}
let double = number.doubleValue
guard double.isFinite, double >= 0 else {
throw BedrockUsageError.cloudWatchParseFailed("metric value was invalid")
}
totals[metric, default: 0] += double
}
}
return (totals, BedrockSettingsReader.cleaned(json["NextToken"] as? String))
}
}
@@ -0,0 +1,78 @@
import Foundation
/// Resolves Bedrock signing credentials + region for the configured auth mode.
///
/// Shared by the main usage fetch (`BedrockAPIFetchStrategy`) and the daily
/// cost-history refresh (`CostUsageFetcher`) so both honor AWS-profile auth
/// identically instead of the history path silently requiring static keys.
enum BedrockCredentialResolver {
struct Resolved {
let credentials: BedrockAWSSigner.Credentials
let region: String
}
static func resolve(
environment: [String: String],
resolveAWSBinary: ([String: String]) -> String? = { BinaryLocator.resolveAWSBinary(env: $0) },
makeProvider: (String) -> BedrockProfileCredentialProvider = {
BedrockProfileCredentialProvider.live(awsBinaryPath: $0)
}) async throws -> Resolved
{
switch BedrockSettingsReader.authMode(environment: environment) {
case .keys:
guard let accessKeyID = BedrockSettingsReader.accessKeyID(environment: environment),
let secretAccessKey = BedrockSettingsReader.secretAccessKey(environment: environment)
else {
throw BedrockUsageError.missingCredentials
}
let credentials = BedrockAWSSigner.Credentials(
accessKeyID: accessKeyID,
secretAccessKey: secretAccessKey,
sessionToken: BedrockSettingsReader.sessionToken(environment: environment))
return Resolved(
credentials: credentials,
region: BedrockSettingsReader.region(environment: environment))
case .profile:
guard let profile = BedrockSettingsReader.profile(environment: environment) else {
throw BedrockUsageError.missingCredentials
}
guard let awsBinary = resolveAWSBinary(environment) else {
throw BedrockUsageError.awsCLINotFound
}
let cliEnvironment = Self.profileCLIEnvironment(environment)
let provider = makeProvider(awsBinary)
let credentials = try await provider.exportCredentials(profile: profile, environment: cliEnvironment)
let region = try await Self.resolveRegion(
provider: provider,
profile: profile,
environment: cliEnvironment)
return Resolved(credentials: credentials, region: region)
}
}
private static func profileCLIEnvironment(_ environment: [String: String]) -> [String: String] {
var cliEnvironment = environment
// `--profile` selects the requested profile. Leaving AWS_PROFILE set can
// break assume-role profiles that use `credential_source = Environment`;
// keep the source credential variables themselves available.
cliEnvironment[BedrockSettingsReader.profileKey] = nil
return cliEnvironment
}
private static func resolveRegion(
provider: BedrockProfileCredentialProvider,
profile: String,
environment: [String: String]) async throws -> String
{
if let explicit = BedrockSettingsReader.cleaned(environment[BedrockSettingsReader.regionKeys[0]])
?? BedrockSettingsReader.cleaned(environment[BedrockSettingsReader.regionKeys[1]])
{
return explicit
}
if let derived = try await provider.resolveRegion(profile: profile, environment: environment) {
return derived
}
return BedrockSettingsReader.defaultRegion
}
}
@@ -0,0 +1,87 @@
import Foundation
/// Resolves AWS credentials for a named profile by shelling out to the AWS CLI.
///
/// Uses `aws configure export-credentials`, which transparently resolves static
/// credentials, SSO sessions, assume-role chains, and `credential_process` profiles.
/// The runner is injected so tests never invoke a real `aws` binary.
struct BedrockProfileCredentialProvider {
typealias Runner = @Sendable (
_ arguments: [String],
_ environment: [String: String]) async throws -> SubprocessResult
let awsBinaryPath: String
let run: Runner
/// Production provider that drives the real `aws` binary via `SubprocessRunner`.
static func live(awsBinaryPath: String) -> BedrockProfileCredentialProvider {
BedrockProfileCredentialProvider(awsBinaryPath: awsBinaryPath) { arguments, environment in
try await SubprocessRunner.run(
binary: awsBinaryPath,
arguments: arguments,
environment: environment,
timeout: 20,
label: "aws-bedrock-credentials")
}
}
func exportCredentials(
profile: String,
environment: [String: String] = [:]) async throws -> BedrockAWSSigner.Credentials
{
let result: SubprocessResult
do {
result = try await self.run(
["configure", "export-credentials", "--profile", profile, "--format", "process"],
environment)
} catch let SubprocessRunnerError.nonZeroExit(_, stderr) {
throw Self.mapExportError(stderr: stderr, profile: profile)
}
guard let data = result.stdout.data(using: .utf8),
let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any],
let accessKeyID = Self.nonEmpty(json["AccessKeyId"] as? String),
let secretAccessKey = Self.nonEmpty(json["SecretAccessKey"] as? String)
else {
throw BedrockUsageError.parseFailed("Could not parse AWS CLI export-credentials output")
}
return BedrockAWSSigner.Credentials(
accessKeyID: accessKeyID,
secretAccessKey: secretAccessKey,
sessionToken: Self.nonEmpty(json["SessionToken"] as? String))
}
/// Returns the profile's configured region, or `nil` when unset.
/// `aws configure get region` exits non-zero when the value is not configured,
/// which is a normal case rather than an error.
func resolveRegion(
profile: String,
environment: [String: String] = [:]) async throws -> String?
{
do {
let result = try await self.run(
["configure", "get", "region", "--profile", profile],
environment)
return Self.nonEmpty(result.stdout)
} catch SubprocessRunnerError.nonZeroExit {
return nil
}
}
static func mapExportError(stderr: String, profile: String) -> BedrockUsageError {
let lower = stderr.lowercased()
if lower.contains("sso login") || lower.contains("expired") || lower.contains("token has expired") {
return .profileSessionExpired(profile)
}
let trimmed = stderr.trimmingCharacters(in: .whitespacesAndNewlines)
return .apiError(trimmed.isEmpty ? "AWS CLI failed to export credentials" : trimmed)
}
private static func nonEmpty(_ raw: String?) -> String? {
guard let value = raw?.trimmingCharacters(in: .whitespacesAndNewlines), !value.isEmpty else {
return nil
}
return value
}
}
@@ -0,0 +1,74 @@
import Foundation
public enum BedrockProviderDescriptor {
public static let descriptor: ProviderDescriptor = Self.makeDescriptor()
static func makeDescriptor() -> ProviderDescriptor {
ProviderDescriptor(
id: .bedrock,
metadata: ProviderMetadata(
id: .bedrock,
displayName: "AWS Bedrock",
sessionLabel: "Budget",
weeklyLabel: "Cost",
opusLabel: nil,
supportsOpus: false,
supportsCredits: false,
creditsHint: "",
toggleTitle: "Show AWS Bedrock usage",
cliName: "bedrock",
defaultEnabled: false,
isPrimaryProvider: false,
usesAccountFallback: false,
dashboardURL: "https://console.aws.amazon.com/bedrock",
statusPageURL: nil,
statusLinkURL: "https://health.aws.amazon.com/health/status"),
branding: ProviderBranding(
iconStyle: .bedrock,
iconResourceName: "ProviderIcon-bedrock",
color: ProviderColor(red: 1, green: 0.6, blue: 0)),
tokenCost: ProviderTokenCostConfig(
supportsTokenCost: true,
noDataMessage: { "No AWS Bedrock cost data available. Check your AWS access keys "
+ "or profile, and that the AWS CLI is installed for profile auth."
}),
fetchPlan: ProviderFetchPlan(
sourceModes: [.auto, .api],
pipeline: ProviderFetchPipeline(resolveStrategies: { _ in [BedrockAPIFetchStrategy()] })),
cli: ProviderCLIConfig(
name: "bedrock",
aliases: ["aws-bedrock"],
versionDetector: nil))
}
}
struct BedrockAPIFetchStrategy: ProviderFetchStrategy {
let id: String = "bedrock.api"
let kind: ProviderFetchKind = .apiToken
func isAvailable(_ context: ProviderFetchContext) async -> Bool {
switch BedrockSettingsReader.authMode(environment: context.env) {
case .keys:
BedrockSettingsReader.hasCredentials(environment: context.env)
case .profile:
BedrockSettingsReader.profile(environment: context.env) != nil
}
}
func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult {
let resolved = try await BedrockCredentialResolver.resolve(environment: context.env)
let budget = BedrockSettingsReader.budget(environment: context.env)
let usage = try await BedrockUsageFetcher.fetchUsage(
credentials: resolved.credentials,
region: resolved.region,
budget: budget,
environment: context.env)
return self.makeResult(
usage: usage.toUsageSnapshot(),
sourceLabel: "api")
}
func shouldFallback(on _: any Error, context _: ProviderFetchContext) -> Bool {
false
}
}
@@ -0,0 +1,107 @@
import Foundation
public enum BedrockAuthMode: String, Codable, Sendable, CaseIterable {
case keys
case profile
}
public enum BedrockSettingsReader {
public static let accessKeyIDKey = "AWS_ACCESS_KEY_ID"
public static let secretAccessKeyKey = "AWS_SECRET_ACCESS_KEY"
public static let sessionTokenKey = "AWS_SESSION_TOKEN"
public static let regionKeys = ["AWS_REGION", "AWS_DEFAULT_REGION"]
public static let budgetKey = "CODEXBAR_BEDROCK_BUDGET"
public static let apiURLKey = "CODEXBAR_BEDROCK_API_URL"
public static let cloudWatchAPIURLKey = "CODEXBAR_BEDROCK_CLOUDWATCH_API_URL"
public static let profileKey = "AWS_PROFILE"
public static let authModeKey = "CODEXBAR_BEDROCK_AUTH_MODE"
public static let defaultRegion = "us-east-1"
public static func accessKeyID(environment: [String: String] = ProcessInfo.processInfo.environment) -> String? {
self.cleaned(environment[self.accessKeyIDKey])
}
public static func secretAccessKey(
environment: [String: String] = ProcessInfo.processInfo.environment) -> String?
{
self.cleaned(environment[self.secretAccessKeyKey])
}
public static func sessionToken(
environment: [String: String] = ProcessInfo.processInfo.environment) -> String?
{
self.cleaned(environment[self.sessionTokenKey])
}
public static func region(environment: [String: String] = ProcessInfo.processInfo.environment) -> String {
for key in self.regionKeys {
if let value = self.cleaned(environment[key]) {
return value
}
}
return self.defaultRegion
}
public static func budget(environment: [String: String] = ProcessInfo.processInfo.environment) -> Double? {
guard let raw = self.cleaned(environment[self.budgetKey]),
let value = Double(raw),
value > 0
else {
return nil
}
return value
}
public static func profile(
environment: [String: String] = ProcessInfo.processInfo.environment) -> String?
{
self.cleaned(environment[self.profileKey])
}
public static func authMode(
environment: [String: String] = ProcessInfo.processInfo.environment) -> BedrockAuthMode
{
if let raw = self.cleaned(environment[self.authModeKey])?.lowercased(),
let mode = BedrockAuthMode(rawValue: raw)
{
return mode
}
if self.profile(environment: environment) != nil,
!self.hasStaticKeys(environment: environment)
{
return .profile
}
return .keys
}
static func hasStaticKeys(environment: [String: String]) -> Bool {
self.accessKeyID(environment: environment) != nil &&
self.secretAccessKey(environment: environment) != nil
}
public static func hasCredentials(
environment: [String: String] = ProcessInfo.processInfo.environment) -> Bool
{
switch self.authMode(environment: environment) {
case .keys:
self.hasStaticKeys(environment: environment)
case .profile:
self.profile(environment: environment) != nil
}
}
static func cleaned(_ raw: String?) -> String? {
guard var value = raw?.trimmingCharacters(in: .whitespacesAndNewlines), !value.isEmpty else {
return nil
}
if (value.hasPrefix("\"") && value.hasSuffix("\"")) ||
(value.hasPrefix("'") && value.hasSuffix("'"))
{
value = String(value.dropFirst().dropLast())
}
value = value.trimmingCharacters(in: .whitespacesAndNewlines)
return value.isEmpty ? nil : value
}
}
@@ -0,0 +1,503 @@
import Foundation
#if canImport(FoundationNetworking)
import FoundationNetworking
#endif
public struct BedrockUsageSnapshot: Codable, Sendable {
public let monthlySpend: Double
public let monthlyBudget: Double?
public let inputTokens: Int?
public let outputTokens: Int?
public let requestCount: Int?
public let region: String
public let updatedAt: Date
public init(
monthlySpend: Double,
monthlyBudget: Double?,
inputTokens: Int? = nil,
outputTokens: Int? = nil,
requestCount: Int? = nil,
region: String,
updatedAt: Date)
{
self.monthlySpend = monthlySpend
self.monthlyBudget = monthlyBudget
self.inputTokens = inputTokens
self.outputTokens = outputTokens
self.requestCount = requestCount
self.region = region
self.updatedAt = updatedAt
}
public var budgetUsedPercent: Double? {
guard let budget = self.monthlyBudget, budget > 0 else { return nil }
return min(100, max(0, (self.monthlySpend / budget) * 100))
}
public var totalTokens: Int? {
guard let input = self.inputTokens, let output = self.outputTokens else { return nil }
return input + output
}
}
extension BedrockUsageSnapshot {
public func toUsageSnapshot() -> UsageSnapshot {
let primary: RateWindow? = if let usedPercent = self.budgetUsedPercent {
RateWindow(
usedPercent: usedPercent,
windowMinutes: nil,
resetsAt: Self.endOfCurrentMonth(),
resetDescription: "Monthly budget")
} else {
nil
}
let cost = ProviderCostSnapshot(
used: self.monthlySpend,
limit: self.monthlyBudget ?? 0,
currencyCode: "USD",
period: "Monthly",
resetsAt: Self.endOfCurrentMonth(),
updatedAt: self.updatedAt)
var loginParts: [String] = []
loginParts.append(String(format: "Spend: $%.2f", self.monthlySpend))
if let budget = self.monthlyBudget {
loginParts.append(String(format: "Budget: $%.2f", budget))
}
if let total = self.totalTokens {
loginParts.append("Claude 14d: \(Self.formattedTokenCount(total)) tokens")
}
if let requestCount = self.requestCount {
loginParts.append("Requests: \(Self.formattedTokenCount(requestCount))")
}
let identity = ProviderIdentitySnapshot(
providerID: .bedrock,
accountEmail: nil,
accountOrganization: nil,
loginMethod: loginParts.joined(separator: " - "))
return UsageSnapshot(
primary: primary,
secondary: nil,
tertiary: nil,
providerCost: cost,
updatedAt: self.updatedAt,
identity: identity)
}
private static func endOfCurrentMonth() -> Date? {
let calendar = Calendar.current
guard let range = calendar.range(of: .day, in: .month, for: Date()) else { return nil }
let components = calendar.dateComponents([.year, .month], from: Date())
guard let startOfMonth = calendar.date(from: components) else { return nil }
return calendar.date(byAdding: .day, value: range.count, to: startOfMonth)
}
static func formattedTokenCount(_ count: Int) -> String {
if count >= 1_000_000 {
return String(format: "%.1fM", Double(count) / 1_000_000)
} else if count >= 1000 {
return String(format: "%.1fK", Double(count) / 1000)
}
return "\(count)"
}
}
enum BedrockUsageFetcher {
private static let log = CodexBarLog.logger(LogCategories.bedrockUsage)
private static let requestTimeoutSeconds: TimeInterval = 15
private struct CostExplorerQuery {
let startDate: String
let endDate: String
let granularity: String
let nextPageToken: String?
}
static func fetchUsage(
credentials: BedrockAWSSigner.Credentials,
region: String,
budget: Double?,
environment: [String: String] = ProcessInfo.processInfo.environment,
now: Date = Date(),
cloudWatchTransport: any ProviderHTTPTransport = ProviderHTTPClient.shared) async throws
-> BedrockUsageSnapshot
{
let spend = try await Self.fetchMonthlyCost(
credentials: credentials,
environment: environment)
var activity: BedrockClaudeActivity?
let cloudWatchOverride = environment[BedrockSettingsReader.cloudWatchAPIURLKey]
let shouldFetchCloudWatch = environment[BedrockSettingsReader.apiURLKey] == nil || cloudWatchOverride != nil
if shouldFetchCloudWatch {
do {
activity = try await BedrockCloudWatchUsageFetcher.fetch(
credentials: credentials,
region: region,
now: now,
endpointOverride: cloudWatchOverride,
transport: cloudWatchTransport)
} catch {
if Task.isCancelled { throw error }
Self.log.debug("AWS CloudWatch Claude activity unavailable; keeping Cost Explorer usage.")
}
}
return BedrockUsageSnapshot(
monthlySpend: spend,
monthlyBudget: budget,
inputTokens: activity?.inputTokens,
outputTokens: activity?.outputTokens,
requestCount: activity?.requestCount,
region: region,
updatedAt: now)
}
static func fetchDailyReport(
credentials: BedrockAWSSigner.Credentials,
since: Date,
until: Date,
environment: [String: String] = ProcessInfo.processInfo.environment) async throws
-> CostUsageDailyReport
{
let formatter = Self.dateFormatter()
let startDate = formatter.string(from: since)
let inclusiveEnd = Self.utcCalendar().date(byAdding: .day, value: 1, to: until) ?? until
let endDate = formatter.string(from: inclusiveEnd)
let pages = try await Self.callCostExplorerPages(
startDate: startDate,
endDate: endDate,
granularity: "DAILY",
credentials: credentials,
environment: environment)
let entries = try Self.parseDailyResponses(pages)
return CostUsageDailyReport(data: entries, summary: nil)
}
private static func fetchMonthlyCost(
credentials: BedrockAWSSigner.Credentials,
environment: [String: String]) async throws -> Double
{
let (startDate, endDate) = Self.currentMonthRange()
let pages = try await Self.callCostExplorerPages(
startDate: startDate,
endDate: endDate,
granularity: "MONTHLY",
credentials: credentials,
environment: environment)
return try Self.parseTotalCost(pages)
}
private static func callCostExplorerPages(
startDate: String,
endDate: String,
granularity: String,
credentials: BedrockAWSSigner.Credentials,
environment: [String: String]) async throws -> [Data]
{
var pages: [Data] = []
var nextPageToken: String?
var seenPageTokens: Set<String> = []
repeat {
let page = try await Self.callCostExplorerPage(
query: CostExplorerQuery(
startDate: startDate,
endDate: endDate,
granularity: granularity,
nextPageToken: nextPageToken),
credentials: credentials,
environment: environment)
pages.append(page)
nextPageToken = try Self.nextPageToken(from: page)
if let nextPageToken, !seenPageTokens.insert(nextPageToken).inserted {
throw BedrockUsageError.parseFailed("Cost Explorer returned repeated NextPageToken")
}
} while nextPageToken != nil
return pages
}
private static func callCostExplorerPage(
query: CostExplorerQuery,
credentials: BedrockAWSSigner.Credentials,
environment: [String: String]) async throws -> Data
{
let ceRegion = "us-east-1"
let baseURL: URL = if environment[BedrockSettingsReader.apiURLKey] != nil {
if let override = BedrockSettingsReader.cleaned(environment[BedrockSettingsReader.apiURLKey]),
let url = ProviderEndpointOverrideValidator().validatedURLAllowingLoopbackHTTP(override)
{
url
} else {
throw BedrockUsageError.parseFailed("invalid endpoint override")
}
} else {
URL(string: "https://ce.\(ceRegion).amazonaws.com")!
}
var requestBody: [String: Any] = [
"TimePeriod": [
"Start": query.startDate,
"End": query.endDate,
],
"Granularity": query.granularity,
"Metrics": ["UnblendedCost"],
"GroupBy": [
["Type": "DIMENSION", "Key": "SERVICE"],
],
]
if let nextPageToken = query.nextPageToken {
requestBody["NextPageToken"] = nextPageToken
}
let bodyData = try JSONSerialization.data(withJSONObject: requestBody)
var request = URLRequest(url: baseURL)
request.httpMethod = "POST"
request.httpBody = bodyData
request.setValue("application/x-amz-json-1.1", forHTTPHeaderField: "Content-Type")
request.setValue(
"AWSInsightsIndexService.GetCostAndUsage",
forHTTPHeaderField: "X-Amz-Target")
request.timeoutInterval = Self.requestTimeoutSeconds
BedrockAWSSigner.sign(
request: &request,
credentials: credentials,
region: ceRegion,
service: "ce")
let response = try await ProviderHTTPClient.shared.response(for: request)
guard response.statusCode == 200 else {
if Self.isDataUnavailableResponse(statusCode: response.statusCode, data: response.data) {
Self.log.info("AWS Cost Explorer data unavailable, assuming zero usage.")
return Data(#"{"ResultsByTime":[]}"#.utf8)
}
let summary = Self.sanitizedResponseBody(response.data)
Self.log.error("AWS Cost Explorer returned \(response.statusCode): \(summary)")
throw BedrockUsageError.apiError("HTTP \(response.statusCode)")
}
return response.data
}
private static func nextPageToken(from data: Data) throws -> String? {
guard let json = try JSONSerialization.jsonObject(with: data) as? [String: Any] else {
throw BedrockUsageError.parseFailed("Invalid Cost Explorer response")
}
return BedrockSettingsReader.cleaned(json["NextPageToken"] as? String)
}
private static func parseTotalCost(_ pages: [Data]) throws -> Double {
var total = 0.0
for page in pages {
total += try Self.parseTotalCost(page)
}
return total
}
private static func parseTotalCost(_ data: Data) throws -> Double {
var total = 0.0
for (_, cost, _) in try Self.parseGroupedResults(data) {
total += cost
}
return total
}
private static func parseDailyResponses(_ pages: [Data]) throws -> [CostUsageDailyReport.Entry] {
let reports = try pages.map { page in
try CostUsageDailyReport(data: Self.parseDailyResponse(page), summary: nil)
}
return CostUsageDailyReport.merged(reports).data
}
private static func parseDailyResponse(_ data: Data) throws -> [CostUsageDailyReport.Entry] {
guard let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any],
let results = json["ResultsByTime"] as? [[String: Any]]
else {
throw BedrockUsageError.parseFailed("Missing ResultsByTime in Cost Explorer response")
}
var entries: [CostUsageDailyReport.Entry] = []
for result in results {
guard let timePeriod = result["TimePeriod"] as? [String: String],
let dateStr = timePeriod["Start"]
else { continue }
var dayCost = 0.0
var breakdowns: [CostUsageDailyReport.ModelBreakdown] = []
if let groups = result["Groups"] as? [[String: Any]] {
for group in groups {
guard let keys = group["Keys"] as? [String],
let serviceName = keys.first,
serviceName.localizedCaseInsensitiveContains("Bedrock")
else { continue }
if let metrics = group["Metrics"] as? [String: Any],
let unblended = metrics["UnblendedCost"] as? [String: Any],
let amountStr = unblended["Amount"] as? String,
let amount = Double(amountStr),
amount > 0
{
dayCost += amount
breakdowns.append(CostUsageDailyReport.ModelBreakdown(
modelName: serviceName,
costUSD: amount))
}
}
}
guard dayCost > 0 else { continue }
entries.append(CostUsageDailyReport.Entry(
date: dateStr,
inputTokens: nil,
outputTokens: nil,
totalTokens: nil,
costUSD: dayCost,
modelsUsed: breakdowns.map(\.modelName),
modelBreakdowns: breakdowns.isEmpty ? nil : breakdowns))
}
return entries
}
private static func parseGroupedResults(_ data: Data) throws
-> [(service: String, cost: Double, date: String)]
{
guard let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any],
let results = json["ResultsByTime"] as? [[String: Any]]
else {
throw BedrockUsageError.parseFailed("Missing ResultsByTime in Cost Explorer response")
}
var items: [(service: String, cost: Double, date: String)] = []
for result in results {
let dateStr = (result["TimePeriod"] as? [String: String])?["Start"] ?? ""
guard let groups = result["Groups"] as? [[String: Any]] else { continue }
for group in groups {
guard let keys = group["Keys"] as? [String],
let serviceName = keys.first,
serviceName.localizedCaseInsensitiveContains("Bedrock")
else { continue }
if let metrics = group["Metrics"] as? [String: Any],
let unblended = metrics["UnblendedCost"] as? [String: Any],
let amountStr = unblended["Amount"] as? String,
let amount = Double(amountStr)
{
items.append((serviceName, amount, dateStr))
}
}
}
return items
}
private static func dateFormatter() -> DateFormatter {
let formatter = DateFormatter()
formatter.dateFormat = "yyyy-MM-dd"
formatter.timeZone = TimeZone(identifier: "UTC")
formatter.locale = Locale(identifier: "en_US_POSIX")
return formatter
}
private static func utcCalendar() -> Calendar {
var calendar = Calendar(identifier: .gregorian)
calendar.timeZone = TimeZone(secondsFromGMT: 0)!
return calendar
}
static func currentMonthRange(now: Date = Date()) -> (start: String, end: String) {
let calendar = Self.utcCalendar()
let components = calendar.dateComponents([.year, .month], from: now)
let startOfMonth = calendar.date(from: components)!
let formatter = Self.dateFormatter()
let startOfToday = calendar.startOfDay(for: now)
let tomorrow = calendar.date(byAdding: .day, value: 1, to: startOfToday)!
return (formatter.string(from: startOfMonth), formatter.string(from: tomorrow))
}
private static func isDataUnavailableResponse(statusCode: Int, data: Data) -> Bool {
guard statusCode == 400,
let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any]
else {
return false
}
let nestedError = json["Error"] as? [String: Any]
let candidates = [
json["__type"],
json["code"],
json["Code"],
nestedError?["Code"],
]
return candidates.compactMap { $0 as? String }.contains { rawCode in
rawCode.split(separator: "#").last == "DataUnavailableException"
}
}
private static func sanitizedResponseBody(_ data: Data) -> String {
guard !data.isEmpty,
let body = String(bytes: data, encoding: .utf8)
else {
return "empty body"
}
let trimmed = body.replacingOccurrences(
of: #"\s+"#,
with: " ",
options: .regularExpression)
.trimmingCharacters(in: .whitespacesAndNewlines)
if trimmed.count > 240 {
let index = trimmed.index(trimmed.startIndex, offsetBy: 240)
return "\(trimmed[..<index])... [truncated]"
}
return trimmed
}
}
public enum BedrockUsageError: LocalizedError, Sendable, Equatable {
case missingCredentials
case awsCLINotFound
case profileSessionExpired(String)
case networkError(String)
case apiError(String)
case parseFailed(String)
case cloudWatchAPIError(String)
case cloudWatchParseFailed(String)
public var errorDescription: String? {
switch self {
case .missingCredentials:
"AWS credentials not configured. Set AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY " +
"or configure Bedrock in Settings."
case .awsCLINotFound:
"AWS CLI not found. Install the AWS CLI (v2) or set AWS_CLI_PATH to its location."
case let .profileSessionExpired(profile):
"AWS profile session expired. Run `aws sso login --profile \(profile)` and try again."
case let .networkError(message):
"AWS Bedrock network error: \(message)"
case let .apiError(message):
"AWS Cost Explorer API error: \(message)"
case let .parseFailed(message):
"Failed to parse AWS Cost Explorer response: \(message)"
case let .cloudWatchAPIError(message):
"AWS CloudWatch API error: \(message)"
case let .cloudWatchParseFailed(message):
"Failed to parse AWS CloudWatch response: \(message)"
}
}
}
@@ -0,0 +1,9 @@
import Foundation
public enum CLIProbeSessionResetter {
public static func resetAll() async {
await ClaudeCLISession.shared.reset()
await CodexCLISession.shared.reset()
await AntigravityCLISession.shared.reset()
}
}
@@ -0,0 +1,68 @@
import Foundation
public enum ChutesProviderDescriptor {
public static let descriptor: ProviderDescriptor = Self.makeDescriptor()
static func makeDescriptor() -> ProviderDescriptor {
ProviderDescriptor(
id: .chutes,
metadata: ProviderMetadata(
id: .chutes,
displayName: "Chutes",
sessionLabel: "4-hour quota",
weeklyLabel: "Monthly quota",
opusLabel: nil,
supportsOpus: false,
supportsCredits: false,
creditsHint: "Subscription usage from the Chutes API.",
toggleTitle: "Show Chutes usage",
cliName: "chutes",
defaultEnabled: false,
isPrimaryProvider: false,
usesAccountFallback: false,
browserCookieOrder: nil,
dashboardURL: "https://chutes.ai",
statusPageURL: nil),
branding: ProviderBranding(
iconStyle: .chutes,
iconResourceName: "ProviderIcon-chutes",
color: ProviderColor(red: 49 / 255, green: 132 / 255, blue: 255 / 255)),
tokenCost: ProviderTokenCostConfig(
supportsTokenCost: false,
noDataMessage: { "Chutes cost history is not available from CodexBar." }),
fetchPlan: ProviderFetchPlan(
sourceModes: [.auto, .api],
pipeline: ProviderFetchPipeline(resolveStrategies: { _ in [ChutesAPIFetchStrategy()] })),
cli: ProviderCLIConfig(
name: "chutes",
aliases: ["chutes.ai"],
versionDetector: nil))
}
}
struct ChutesAPIFetchStrategy: ProviderFetchStrategy {
let id: String = "chutes.api"
let kind: ProviderFetchKind = .apiToken
func isAvailable(_ context: ProviderFetchContext) async -> Bool {
ChutesSettingsReader.apiKey(environment: context.env) != nil
}
func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult {
guard let apiKey = ChutesSettingsReader.apiKey(environment: context.env) else {
throw ChutesSettingsError.missingToken
}
let usage = try await ChutesUsageFetcher.fetchUsage(
apiKey: apiKey,
environment: context.env)
return self.makeResult(
usage: usage.toUsageSnapshot(),
sourceLabel: "api")
}
func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool {
false
}
}
@@ -0,0 +1,63 @@
import Foundation
public struct ChutesSettingsReader: Sendable {
public static let apiKeyEnvironmentKey = "CHUTES_API_KEY"
public static let apiURLEnvironmentKey = "CHUTES_API_URL"
public static func apiKey(
environment: [String: String] = ProcessInfo.processInfo.environment) -> String?
{
self.cleaned(environment[self.apiKeyEnvironmentKey])
}
public static func apiURL(
environment: [String: String] = ProcessInfo.processInfo.environment) -> URL
{
if let override = self.validAPIURL(environment: environment) {
return override
}
return URL(string: "https://api.chutes.ai")!
}
public static func validateEndpointOverrides(
environment: [String: String] = ProcessInfo.processInfo.environment) throws
{
guard let raw = self.cleaned(environment[self.apiURLEnvironmentKey]) else { return }
guard ProviderEndpointOverrideValidator.normalizedHTTPSURL(from: raw) == nil else { return }
throw ChutesSettingsError.invalidEndpointOverride(self.apiURLEnvironmentKey)
}
static func cleaned(_ raw: String?) -> String? {
guard var value = raw?.trimmingCharacters(in: .whitespacesAndNewlines), !value.isEmpty else {
return nil
}
if (value.hasPrefix("\"") && value.hasSuffix("\"")) ||
(value.hasPrefix("'") && value.hasSuffix("'"))
{
value = String(value.dropFirst().dropLast())
}
value = value.trimmingCharacters(in: .whitespacesAndNewlines)
return value.isEmpty ? nil : value
}
private static func validAPIURL(environment: [String: String]) -> URL? {
guard let raw = self.cleaned(environment[self.apiURLEnvironmentKey]) else { return nil }
return ProviderEndpointOverrideValidator.normalizedHTTPSURL(from: raw)
}
}
public enum ChutesSettingsError: LocalizedError, Sendable, Equatable {
case missingToken
case invalidEndpointOverride(String)
public var errorDescription: String? {
switch self {
case .missingToken:
"Chutes API key not found. Set apiKey in ~/.codexbar/config.json or CHUTES_API_KEY."
case let .invalidEndpointOverride(key):
"Chutes endpoint override \(key) must use HTTPS or a bare host."
}
}
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,43 @@
import Foundation
public enum ClaudeAdminAPISettingsReader {
public static let adminAPIKeyEnvironmentKey = "ANTHROPIC_ADMIN_KEY"
public static let alternateAdminAPIKeyEnvironmentKey = "ANTHROPIC_ADMIN_API_KEY"
public static let apiKeyEnvironmentKeys = [
Self.adminAPIKeyEnvironmentKey,
Self.alternateAdminAPIKeyEnvironmentKey,
]
public static func apiKey(environment: [String: String] = ProcessInfo.processInfo.environment) -> String? {
for key in self.apiKeyEnvironmentKeys {
if let token = self.cleaned(environment[key]) { return token }
}
return nil
}
public static func cleaned(_ raw: String?) -> String? {
guard var value = raw?.trimmingCharacters(in: .whitespacesAndNewlines), !value.isEmpty else {
return nil
}
if (value.hasPrefix("\"") && value.hasSuffix("\"")) ||
(value.hasPrefix("'") && value.hasSuffix("'"))
{
value = String(value.dropFirst().dropLast())
}
value = value.trimmingCharacters(in: .whitespacesAndNewlines)
return value.isEmpty ? nil : value
}
}
public enum ClaudeAdminAPISettingsError: LocalizedError, Sendable {
case missingToken
public var errorDescription: String? {
switch self {
case .missingToken:
"Claude API usage needs an Anthropic Admin API key."
}
}
}
@@ -0,0 +1,432 @@
import Foundation
#if canImport(FoundationNetworking)
import FoundationNetworking
#endif
public enum ClaudeAdminAPIUsageError: LocalizedError, Sendable, Equatable {
case missingCredentials
case networkError(String)
case apiError(endpoint: String, statusCode: Int)
case parseFailed(endpoint: String, message: String)
public var errorDescription: String? {
switch self {
case .missingCredentials:
"Missing Anthropic Admin API key."
case let .networkError(message):
"Claude API usage network error: \(message)"
case let .apiError(endpoint, statusCode):
"Claude API usage \(endpoint) error: HTTP \(statusCode)"
case let .parseFailed(endpoint, message):
"Failed to parse Claude API usage \(endpoint): \(message)"
}
}
}
public enum ClaudeAdminAPIUsageFetcher {
public static let costReportURL = URL(string: "https://api.anthropic.com/v1/organizations/cost_report")!
public static let messagesUsageURL =
URL(string: "https://api.anthropic.com/v1/organizations/usage_report/messages")!
private static let anthropicVersion = "2023-06-01"
private static let timeoutSeconds: TimeInterval = 20
private static let maxDailyBuckets = 31
public static func fetchUsage(
apiKey: String,
costURL: URL = Self.costReportURL,
messagesURL: URL = Self.messagesUsageURL,
session transport: any ProviderHTTPTransport = ProviderHTTPClient.shared,
now: Date = Date()) async throws -> ClaudeAdminAPIUsageSnapshot
{
let trimmed = apiKey.trimmingCharacters(in: .whitespacesAndNewlines)
guard !trimmed.isEmpty else {
throw ClaudeAdminAPIUsageError.missingCredentials
}
let calendar = Self.utcCalendar
let range = Self.dailyRange(now: now, calendar: calendar)
let costs = try await Self.fetchCostReport(
apiKey: trimmed,
baseURL: costURL,
range: range,
transport: transport)
let messages = try await Self.fetchMessagesUsage(
apiKey: trimmed,
baseURL: messagesURL,
range: range,
transport: transport)
return Self.makeSnapshot(costs: costs, messages: messages, now: now, calendar: calendar)
}
static func _parseSnapshotForTesting(
costs: Data,
messages: Data,
now: Date,
calendar: Calendar = Self.utcCalendar) throws -> ClaudeAdminAPIUsageSnapshot
{
let costs = try Self.decodeCosts(costs)
let messages = try Self.decodeMessages(messages)
return Self.makeSnapshot(costs: costs, messages: messages, now: now, calendar: calendar)
}
private static func fetchCostReport(
apiKey: String,
baseURL: URL,
range: DateRange,
transport: any ProviderHTTPTransport) async throws -> CostReportResponse
{
let url = Self.url(
baseURL: baseURL,
range: range,
queryItems: [
URLQueryItem(name: "group_by[]", value: "description"),
])
let data = try await Self.fetchData(url: url, apiKey: apiKey, endpoint: "cost_report", transport: transport)
return try Self.decodeCosts(data)
}
private static func fetchMessagesUsage(
apiKey: String,
baseURL: URL,
range: DateRange,
transport: any ProviderHTTPTransport) async throws -> MessagesUsageResponse
{
let url = Self.url(
baseURL: baseURL,
range: range,
queryItems: [
URLQueryItem(name: "group_by[]", value: "model"),
])
let data = try await Self.fetchData(url: url, apiKey: apiKey, endpoint: "messages", transport: transport)
return try Self.decodeMessages(data)
}
private static func fetchData(
url: URL,
apiKey: String,
endpoint: String,
transport: any ProviderHTTPTransport) async throws -> Data
{
var request = URLRequest(url: url)
request.httpMethod = "GET"
request.timeoutInterval = Self.timeoutSeconds
request.setValue(Self.anthropicVersion, forHTTPHeaderField: "anthropic-version")
request.setValue(apiKey, forHTTPHeaderField: "x-api-key")
request.setValue("application/json", forHTTPHeaderField: "Accept")
request.setValue("CodexBar/1.0", forHTTPHeaderField: "User-Agent")
let response: ProviderHTTPResponse
do {
response = try await transport.response(for: request)
} catch {
throw ClaudeAdminAPIUsageError.networkError(error.localizedDescription)
}
guard response.statusCode == 200 else {
throw ClaudeAdminAPIUsageError.apiError(endpoint: endpoint, statusCode: response.statusCode)
}
return response.data
}
private static func decodeCosts(_ data: Data) throws -> CostReportResponse {
do {
return try JSONDecoder().decode(CostReportResponse.self, from: data)
} catch {
throw ClaudeAdminAPIUsageError.parseFailed(endpoint: "cost_report", message: error.localizedDescription)
}
}
private static func decodeMessages(_ data: Data) throws -> MessagesUsageResponse {
do {
return try JSONDecoder().decode(MessagesUsageResponse.self, from: data)
} catch {
throw ClaudeAdminAPIUsageError.parseFailed(endpoint: "messages", message: error.localizedDescription)
}
}
private static func makeSnapshot(
costs: CostReportResponse,
messages: MessagesUsageResponse,
now: Date,
calendar: Calendar) -> ClaudeAdminAPIUsageSnapshot
{
var accumulators: [String: DailyAccumulator] = [:]
for bucket in costs.data {
var accumulator = accumulators[bucket.startingAt] ?? DailyAccumulator(
startingAt: bucket.startingAt,
endingAt: bucket.endingAt)
for result in bucket.results {
// Anthropic Usage & Cost API docs define `amount` as a decimal string in lowest USD units.
let value = Self.usdFromAnthropicLowestUnitAmount(result.amount)
accumulator.costUSD += value
let item = Self.displayName(result.description ?? result.costType, fallback: "Claude API")
accumulator.costItems[item, default: 0] += value
}
accumulators[bucket.startingAt] = accumulator
}
for bucket in messages.data {
var accumulator = accumulators[bucket.startingAt] ?? DailyAccumulator(
startingAt: bucket.startingAt,
endingAt: bucket.endingAt)
for result in bucket.results {
let input = result.uncachedInputTokens ?? 0
let cacheCreation = result.cacheCreation?.totalInputTokens ?? 0
let cacheRead = result.cacheReadInputTokens ?? 0
let output = result.outputTokens ?? 0
let total = input + cacheCreation + cacheRead + output
accumulator.inputTokens += input
accumulator.cacheCreationInputTokens += cacheCreation
accumulator.cacheReadInputTokens += cacheRead
accumulator.outputTokens += output
accumulator.totalTokens += total
let modelName = Self.displayName(result.model, fallback: "Claude API")
accumulator.models[modelName, default: ModelAccumulator()].add(
inputTokens: input,
cacheCreationInputTokens: cacheCreation,
cacheReadInputTokens: cacheRead,
outputTokens: output,
totalTokens: total)
}
accumulators[bucket.startingAt] = accumulator
}
let daily = accumulators.values
.compactMap { $0.makeBucket(calendar: calendar) }
.filter { $0.startTime <= now }
.sorted { $0.startTime < $1.startTime }
return ClaudeAdminAPIUsageSnapshot(daily: daily, updatedAt: now)
}
private static func displayName(_ raw: String?, fallback: String) -> String {
guard let trimmed = raw?.trimmingCharacters(in: .whitespacesAndNewlines), !trimmed.isEmpty else {
return fallback
}
return trimmed
}
private static func usdFromAnthropicLowestUnitAmount(_ raw: String) -> Double {
(Double(raw) ?? 0) / 100
}
private static func url(baseURL: URL, range: DateRange, queryItems extraItems: [URLQueryItem]) -> URL {
var components = URLComponents(url: baseURL, resolvingAgainstBaseURL: false)!
components.queryItems = [
URLQueryItem(name: "starting_at", value: Self.rfc3339String(from: range.start)),
URLQueryItem(name: "ending_at", value: Self.rfc3339String(from: range.end)),
URLQueryItem(name: "bucket_width", value: "1d"),
URLQueryItem(name: "limit", value: String(Self.maxDailyBuckets)),
] + extraItems
return components.url!
}
private static func dailyRange(now: Date, calendar: Calendar) -> DateRange {
let today = calendar.startOfDay(for: now)
let start = calendar.date(byAdding: .day, value: -(Self.maxDailyBuckets - 1), to: today) ?? today
let end = calendar.date(byAdding: .day, value: 1, to: today) ?? now
return DateRange(start: start, end: end)
}
private static var utcCalendar: Calendar {
var calendar = Calendar(identifier: .gregorian)
calendar.locale = Locale(identifier: "en_US_POSIX")
calendar.timeZone = TimeZone(secondsFromGMT: 0)!
return calendar
}
private static func rfc3339Formatter() -> ISO8601DateFormatter {
let formatter = ISO8601DateFormatter()
formatter.formatOptions = [.withInternetDateTime]
formatter.timeZone = TimeZone(secondsFromGMT: 0)
return formatter
}
private static func rfc3339String(from date: Date) -> String {
self.rfc3339Formatter().string(from: date)
}
fileprivate static func dayKey(from date: Date, calendar: Calendar) -> String {
let formatter = DateFormatter()
formatter.locale = Locale(identifier: "en_US_POSIX")
formatter.timeZone = calendar.timeZone
formatter.dateFormat = "yyyy-MM-dd"
return formatter.string(from: date)
}
fileprivate static func parseDate(_ raw: String) -> Date? {
self.rfc3339Formatter().date(from: raw)
}
}
private struct DateRange {
let start: Date
let end: Date
}
private struct DailyAccumulator {
let startingAt: String
let endingAt: String
var costUSD: Double = 0
var inputTokens: Int = 0
var cacheCreationInputTokens: Int = 0
var cacheReadInputTokens: Int = 0
var outputTokens: Int = 0
var totalTokens: Int = 0
var costItems: [String: Double] = [:]
var models: [String: ModelAccumulator] = [:]
func makeBucket(calendar: Calendar) -> ClaudeAdminAPIUsageSnapshot.DailyBucket? {
guard let start = ClaudeAdminAPIUsageFetcher.parseDate(self.startingAt),
let end = ClaudeAdminAPIUsageFetcher.parseDate(self.endingAt)
else { return nil }
return ClaudeAdminAPIUsageSnapshot.DailyBucket(
day: ClaudeAdminAPIUsageFetcher.dayKey(from: start, calendar: calendar),
startTime: start,
endTime: end,
costUSD: self.costUSD,
inputTokens: self.inputTokens,
cacheCreationInputTokens: self.cacheCreationInputTokens,
cacheReadInputTokens: self.cacheReadInputTokens,
outputTokens: self.outputTokens,
totalTokens: self.totalTokens,
costItems: self.costItems
.map { ClaudeAdminAPIUsageSnapshot.CostBreakdown(name: $0.key, costUSD: $0.value) }
.sorted {
if $0.costUSD == $1.costUSD { return $0.name < $1.name }
return $0.costUSD > $1.costUSD
},
models: self.models
.map { name, total in total.makeModel(name: name) }
.sorted {
if $0.totalTokens == $1.totalTokens { return $0.name < $1.name }
return $0.totalTokens > $1.totalTokens
})
}
}
private struct ModelAccumulator {
var inputTokens = 0
var cacheCreationInputTokens = 0
var cacheReadInputTokens = 0
var outputTokens = 0
var totalTokens = 0
mutating func add(
inputTokens: Int,
cacheCreationInputTokens: Int,
cacheReadInputTokens: Int,
outputTokens: Int,
totalTokens: Int)
{
self.inputTokens += inputTokens
self.cacheCreationInputTokens += cacheCreationInputTokens
self.cacheReadInputTokens += cacheReadInputTokens
self.outputTokens += outputTokens
self.totalTokens += totalTokens
}
func makeModel(name: String) -> ClaudeAdminAPIUsageSnapshot.ModelBreakdown {
ClaudeAdminAPIUsageSnapshot.ModelBreakdown(
name: name,
inputTokens: self.inputTokens,
cacheCreationInputTokens: self.cacheCreationInputTokens,
cacheReadInputTokens: self.cacheReadInputTokens,
outputTokens: self.outputTokens,
totalTokens: self.totalTokens)
}
}
private struct CostReportResponse: Decodable {
let data: [CostBucket]
let hasMore: Bool?
let nextPage: String?
private enum CodingKeys: String, CodingKey {
case data
case hasMore = "has_more"
case nextPage = "next_page"
}
}
private struct CostBucket: Decodable {
let startingAt: String
let endingAt: String
let results: [CostResult]
private enum CodingKeys: String, CodingKey {
case startingAt = "starting_at"
case endingAt = "ending_at"
case results
}
}
private struct CostResult: Decodable {
let currency: String?
let amount: String
let description: String?
let costType: String?
private enum CodingKeys: String, CodingKey {
case currency
case amount
case description
case costType = "cost_type"
}
}
private struct MessagesUsageResponse: Decodable {
let data: [MessagesBucket]
let hasMore: Bool?
let nextPage: String?
private enum CodingKeys: String, CodingKey {
case data
case hasMore = "has_more"
case nextPage = "next_page"
}
}
private struct MessagesBucket: Decodable {
let startingAt: String
let endingAt: String
let results: [MessagesResult]
private enum CodingKeys: String, CodingKey {
case startingAt = "starting_at"
case endingAt = "ending_at"
case results
}
}
private struct MessagesResult: Decodable {
let uncachedInputTokens: Int?
let cacheCreation: CacheCreation?
let cacheReadInputTokens: Int?
let outputTokens: Int?
let model: String?
private enum CodingKeys: String, CodingKey {
case uncachedInputTokens = "uncached_input_tokens"
case cacheCreation = "cache_creation"
case cacheReadInputTokens = "cache_read_input_tokens"
case outputTokens = "output_tokens"
case model
}
}
private struct CacheCreation: Decodable {
let ephemeral1HInputTokens: Int?
let ephemeral5MInputTokens: Int?
var totalInputTokens: Int {
(self.ephemeral1HInputTokens ?? 0) + (self.ephemeral5MInputTokens ?? 0)
}
private enum CodingKeys: String, CodingKey {
case ephemeral1HInputTokens = "ephemeral_1h_input_tokens"
case ephemeral5MInputTokens = "ephemeral_5m_input_tokens"
}
}
@@ -0,0 +1,242 @@
import Foundation
public struct ClaudeAdminAPIUsageSnapshot: Codable, Equatable, Sendable {
public struct DailyBucket: Codable, Equatable, Sendable, Identifiable {
public let day: String
public let startTime: Date
public let endTime: Date
public let costUSD: Double
public let inputTokens: Int
public let cacheCreationInputTokens: Int
public let cacheReadInputTokens: Int
public let outputTokens: Int
public let totalTokens: Int
public let costItems: [CostBreakdown]
public let models: [ModelBreakdown]
public var id: String {
self.day
}
public init(
day: String,
startTime: Date,
endTime: Date,
costUSD: Double,
inputTokens: Int,
cacheCreationInputTokens: Int,
cacheReadInputTokens: Int,
outputTokens: Int,
totalTokens: Int,
costItems: [CostBreakdown],
models: [ModelBreakdown])
{
self.day = day
self.startTime = startTime
self.endTime = endTime
self.costUSD = costUSD
self.inputTokens = inputTokens
self.cacheCreationInputTokens = cacheCreationInputTokens
self.cacheReadInputTokens = cacheReadInputTokens
self.outputTokens = outputTokens
self.totalTokens = totalTokens
self.costItems = costItems
self.models = models
}
}
public struct CostBreakdown: Codable, Equatable, Sendable, Identifiable {
public let name: String
public let costUSD: Double
public var id: String {
self.name
}
public init(name: String, costUSD: Double) {
self.name = name
self.costUSD = costUSD
}
}
public struct ModelBreakdown: Codable, Equatable, Sendable, Identifiable {
public let name: String
public let inputTokens: Int
public let cacheCreationInputTokens: Int
public let cacheReadInputTokens: Int
public let outputTokens: Int
public let totalTokens: Int
public var id: String {
self.name
}
public init(
name: String,
inputTokens: Int,
cacheCreationInputTokens: Int,
cacheReadInputTokens: Int,
outputTokens: Int,
totalTokens: Int)
{
self.name = name
self.inputTokens = inputTokens
self.cacheCreationInputTokens = cacheCreationInputTokens
self.cacheReadInputTokens = cacheReadInputTokens
self.outputTokens = outputTokens
self.totalTokens = totalTokens
}
}
public struct Summary: Equatable, Sendable {
public let costUSD: Double
public let inputTokens: Int
public let cacheCreationInputTokens: Int
public let cacheReadInputTokens: Int
public let outputTokens: Int
public let totalTokens: Int
public init(
costUSD: Double,
inputTokens: Int,
cacheCreationInputTokens: Int,
cacheReadInputTokens: Int,
outputTokens: Int,
totalTokens: Int)
{
self.costUSD = costUSD
self.inputTokens = inputTokens
self.cacheCreationInputTokens = cacheCreationInputTokens
self.cacheReadInputTokens = cacheReadInputTokens
self.outputTokens = outputTokens
self.totalTokens = totalTokens
}
}
public let daily: [DailyBucket]
public let updatedAt: Date
public init(daily: [DailyBucket], updatedAt: Date) {
self.daily = daily.sorted { $0.startTime < $1.startTime }
self.updatedAt = updatedAt
}
public var last30Days: Summary {
self.summary(days: 30)
}
public var last7Days: Summary {
self.summary(days: 7)
}
public var currentDay: Summary {
self.summary(forLocalDayContaining: self.updatedAt)
}
public var latestDay: Summary {
self.summary(days: 1)
}
public func summary(forLocalDayContaining date: Date, calendar _: Calendar = .current) -> Summary {
let selected = self.daily.filter { bucket in
CostUsageBucketInterval.contains(
date,
startTime: bucket.startTime,
endTime: bucket.endTime)
}
return Summary(
costUSD: selected.reduce(0) { $0 + $1.costUSD },
inputTokens: selected.reduce(0) { $0 + $1.inputTokens },
cacheCreationInputTokens: selected.reduce(0) { $0 + $1.cacheCreationInputTokens },
cacheReadInputTokens: selected.reduce(0) { $0 + $1.cacheReadInputTokens },
outputTokens: selected.reduce(0) { $0 + $1.outputTokens },
totalTokens: selected.reduce(0) { $0 + $1.totalTokens })
}
public func summary(days: Int) -> Summary {
let selected = self.daily.suffix(max(1, days))
return Summary(
costUSD: selected.reduce(0) { $0 + $1.costUSD },
inputTokens: selected.reduce(0) { $0 + $1.inputTokens },
cacheCreationInputTokens: selected.reduce(0) { $0 + $1.cacheCreationInputTokens },
cacheReadInputTokens: selected.reduce(0) { $0 + $1.cacheReadInputTokens },
outputTokens: selected.reduce(0) { $0 + $1.outputTokens },
totalTokens: selected.reduce(0) { $0 + $1.totalTokens })
}
public var topModels: [ModelBreakdown] {
var totals: [String: ModelAccumulator] = [:]
for day in self.daily {
for model in day.models {
totals[model.name, default: ModelAccumulator()].add(model)
}
}
return totals
.map { name, total in total.makeModel(name: name) }
.sorted {
if $0.totalTokens == $1.totalTokens { return $0.name < $1.name }
return $0.totalTokens > $1.totalTokens
}
}
public var topCostItems: [CostBreakdown] {
var totals: [String: Double] = [:]
for day in self.daily {
for item in day.costItems {
totals[item.name, default: 0] += item.costUSD
}
}
return totals
.map { CostBreakdown(name: $0.key, costUSD: $0.value) }
.sorted {
if $0.costUSD == $1.costUSD { return $0.name < $1.name }
return $0.costUSD > $1.costUSD
}
}
public func toUsageSnapshot() -> UsageSnapshot {
let total = self.last30Days
return UsageSnapshot(
primary: nil,
secondary: nil,
providerCost: ProviderCostSnapshot(
used: total.costUSD,
limit: 0,
currencyCode: "USD",
period: "Last 30 days",
updatedAt: self.updatedAt),
claudeAdminAPIUsage: self,
updatedAt: self.updatedAt,
identity: ProviderIdentitySnapshot(
providerID: .claude,
accountEmail: nil,
accountOrganization: nil,
loginMethod: "Admin API"))
}
private struct ModelAccumulator {
var inputTokens = 0
var cacheCreationInputTokens = 0
var cacheReadInputTokens = 0
var outputTokens = 0
var totalTokens = 0
mutating func add(_ model: ModelBreakdown) {
self.inputTokens += model.inputTokens
self.cacheCreationInputTokens += model.cacheCreationInputTokens
self.cacheReadInputTokens += model.cacheReadInputTokens
self.outputTokens += model.outputTokens
self.totalTokens += model.totalTokens
}
func makeModel(name: String) -> ModelBreakdown {
ModelBreakdown(
name: name,
inputTokens: self.inputTokens,
cacheCreationInputTokens: self.cacheCreationInputTokens,
cacheReadInputTokens: self.cacheReadInputTokens,
outputTokens: self.outputTokens,
totalTokens: self.totalTokens)
}
}
}
@@ -0,0 +1,35 @@
import Foundation
enum ClaudeCLIAuthStatusProbe {
private struct Response: Decodable {
let loggedIn: Bool
}
static func isLoggedIn(
binary: String,
environment: [String: String],
timeout: TimeInterval = 5) async -> Bool
{
do {
let result = try await SubprocessRunner.run(
binary: binary,
arguments: ["auth", "status", "--json"],
environment: ClaudeCLISession.launchEnvironment(baseEnv: environment),
timeout: timeout,
standardInput: FileHandle.nullDevice,
label: "claude-auth-status")
return self.parseLoggedIn(result.stdout)
} catch {
return false
}
}
static func parseLoggedIn(_ output: String) -> Bool {
guard let data = output.data(using: .utf8),
let response = try? JSONDecoder().decode(Response.self, from: data)
else {
return false
}
return response.loggedIn
}
}
@@ -0,0 +1,64 @@
import Foundation
enum ClaudeCLIRateLimitGate {
private static let blockedUntilKey = "claudeCLIUsageRateLimitBlockedUntilV1"
private static let defaultCooldown: TimeInterval = 60 * 5
static let message = "Claude CLI usage endpoint is rate limited right now. Please try again later."
static func blockedUntil(
interaction: ProviderInteraction = ProviderInteractionContext.current,
now: Date = Date()) -> Date?
{
guard interaction != .userInitiated else { return nil }
return self.currentBlockedUntil(now: now)
}
static func currentBlockedUntil(now: Date = Date()) -> Date? {
guard let raw = UserDefaults.standard.object(forKey: self.blockedUntilKey) as? Double else {
return nil
}
let blockedUntil = Date(timeIntervalSince1970: raw)
guard blockedUntil > now else {
UserDefaults.standard.removeObject(forKey: self.blockedUntilKey)
return nil
}
return blockedUntil
}
static func recordRateLimit(now: Date = Date()) {
UserDefaults.standard.set(
now.addingTimeInterval(self.defaultCooldown).timeIntervalSince1970,
forKey: self.blockedUntilKey)
}
static func recordSuccess() {
UserDefaults.standard.removeObject(forKey: self.blockedUntilKey)
}
static func isRateLimitError(_ error: Error) -> Bool {
if case let ClaudeStatusProbeError.parseFailed(message) = error {
return self.isRateLimitMessage(message, allowRawRateLimitToken: true)
}
if case let ClaudeUsageError.parseFailed(message) = error {
return self.isRateLimitMessage(message, allowRawRateLimitToken: true)
}
return self.isRateLimitMessage(error.localizedDescription, allowRawRateLimitToken: false)
}
private static func isRateLimitMessage(_ message: String, allowRawRateLimitToken: Bool) -> Bool {
let lower = message.lowercased()
return lower.contains(Self.message.lowercased()) ||
(allowRawRateLimitToken && lower.contains("rate_limit_error")) ||
(lower.contains("claude cli") &&
lower.contains("usage") &&
lower.contains("rate limited"))
}
#if DEBUG
static func resetForTesting() {
UserDefaults.standard.removeObject(forKey: self.blockedUntilKey)
}
#endif
}
@@ -0,0 +1,499 @@
#if canImport(Darwin)
import Darwin
#elseif canImport(Glibc)
import Glibc
#elseif canImport(Musl)
import Musl
#endif
import Foundation
actor ClaudeCLISession {
static let shared = ClaudeCLISession()
private static let log = CodexBarLog.logger(LogCategories.claudeCLI)
#if DEBUG
@TaskLocal private static var sessionOverrideForTesting: ClaudeCLISession?
static var current: ClaudeCLISession {
self.sessionOverrideForTesting ?? self.shared
}
static func withIsolatedSessionForTesting<T>(operation: () async throws -> T) async rethrows -> T {
let session = ClaudeCLISession()
defer { Task { await session.reset() } }
return try await self.$sessionOverrideForTesting.withValue(session) {
try await operation()
}
}
#else
static var current: ClaudeCLISession {
self.shared
}
#endif
enum SessionError: LocalizedError {
case launchFailed(String)
case ioFailed(String)
case timedOut
case processExited
var errorDescription: String? {
switch self {
case let .launchFailed(msg): "Failed to launch Claude CLI session: \(msg)"
case let .ioFailed(msg): "Claude CLI PTY I/O failed: \(msg)"
case .timedOut: "Claude CLI session timed out."
case .processExited: "Claude CLI session exited."
}
}
}
private var process: Process?
private var primaryFD: Int32 = -1
private var primaryHandle: FileHandle?
private var secondaryHandle: FileHandle?
private var processGroup: pid_t?
private var binaryPath: String?
private var startedAt: Date?
private let promptSends: [String: String] = [
"Do you trust the files in this folder?": "y\r",
"Quick safety check:": "\r",
"Yes, I trust this folder": "\r",
"Ready to code here?": "\r",
"Press Enter to continue": "\r",
]
private struct RollingBuffer {
private let maxNeedle: Int
private var tail = Data()
init(maxNeedle: Int) {
self.maxNeedle = max(0, maxNeedle)
}
mutating func append(_ data: Data) -> Data {
guard !data.isEmpty else { return Data() }
var combined = Data()
combined.reserveCapacity(self.tail.count + data.count)
combined.append(self.tail)
combined.append(data)
if self.maxNeedle > 1 {
if combined.count >= self.maxNeedle - 1 {
self.tail = combined.suffix(self.maxNeedle - 1)
} else {
self.tail = combined
}
} else {
self.tail.removeAll(keepingCapacity: true)
}
return combined
}
}
private static func normalizedNeedle(_ text: String) -> String {
String(text.lowercased().filter { !$0.isWhitespace })
}
private static func commandPaletteSends(for subcommand: String) -> [String: String] {
let normalized = subcommand.trimmingCharacters(in: .whitespacesAndNewlines).lowercased()
switch normalized {
case "/usage":
// Claude's command palette can render several "Show ..." actions together; only auto-confirm the
// usage-related actions here so we do not accidentally execute /status.
return [
"Show plan": "\r",
"Show plan usage limits": "\r",
]
case "/status":
return [
"Show Claude Code": "\r",
"Show Claude Code status": "\r",
]
default:
return [:]
}
}
func capture(
subcommand: String,
binary: String,
timeout: TimeInterval,
idleTimeout: TimeInterval? = 3.0,
stopOnSubstrings: [String] = [],
stopWhenNormalized: (@Sendable (String) -> Bool)? = nil,
settleAfterStop: TimeInterval = 0.25,
sendEnterEvery: TimeInterval? = nil) async throws -> String
{
try self.ensureStarted(binary: binary)
if let startedAt {
let sinceStart = Date().timeIntervalSince(startedAt)
// Claude's TUI can drop early keystrokes while it's still initializing. Wait a bit longer than the
// original 0.4s to ensure slash commands reliably open their panels.
if sinceStart < 2.0 {
let delay = UInt64((2.0 - sinceStart) * 1_000_000_000)
try await Task.sleep(nanoseconds: delay)
}
}
self.drainOutput()
let trimmed = subcommand.trimmingCharacters(in: .whitespacesAndNewlines)
if !trimmed.isEmpty {
try self.send(trimmed)
try self.send("\r")
}
let stopNeedles = stopOnSubstrings.map { Self.normalizedNeedle($0) }
var sendMap = self.promptSends
for (needle, keys) in Self.commandPaletteSends(for: trimmed) {
sendMap[needle] = keys
}
let sendNeedles = sendMap.map { (needle: Self.normalizedNeedle($0.key), keys: $0.value) }
let cursorQuery = Data([0x1B, 0x5B, 0x36, 0x6E])
let needleLengths =
stopOnSubstrings.map(\.utf8.count) +
sendMap.keys.map(\.utf8.count) +
[cursorQuery.count]
let maxNeedle = needleLengths.max() ?? cursorQuery.count
var scanBuffer = RollingBuffer(maxNeedle: maxNeedle)
var triggeredSends = Set<String>()
var buffer = Data()
var scanTailText = ""
var normalizedScan = ""
var utf8Carry = Data()
let deadline = Date().addingTimeInterval(timeout)
var lastOutputAt = Date()
var lastEnterAt = Date()
var stoppedEarly = false
// Only send periodic Enter when the caller explicitly asks for it (used for /usage rendering).
// For /status, periodic input can keep producing output and prevent idle-timeout short-circuiting.
let effectiveEnterEvery: TimeInterval? = sendEnterEvery
while Date() < deadline {
let newData = self.readChunk()
if !newData.isEmpty {
buffer.append(newData)
lastOutputAt = Date()
Self.appendScanText(newData: newData, scanTailText: &scanTailText, utf8Carry: &utf8Carry)
if scanTailText.count > 8192 {
scanTailText = String(scanTailText.suffix(8192))
}
normalizedScan = Self.normalizedNeedle(TextParsing.stripANSICodes(scanTailText))
let scanData = scanBuffer.append(newData)
if scanData.range(of: cursorQuery) != nil {
try? self.send("\u{1b}[1;1R")
}
for item in sendNeedles where !triggeredSends.contains(item.needle) {
if normalizedScan.contains(item.needle) {
try? self.send(item.keys)
triggeredSends.insert(item.needle)
}
}
if stopNeedles
.contains(where: normalizedScan.contains) || (stopWhenNormalized?(normalizedScan) == true)
{
stoppedEarly = true
break
}
}
if self.shouldStopForIdleTimeout(
idleTimeout: idleTimeout,
bufferIsEmpty: buffer.isEmpty,
lastOutputAt: lastOutputAt)
{
stoppedEarly = true
break
}
self.sendPeriodicEnterIfNeeded(every: effectiveEnterEvery, lastEnterAt: &lastEnterAt)
if let proc = self.process, !proc.isRunning {
throw SessionError.processExited
}
try await Task.sleep(nanoseconds: 60_000_000)
}
if stoppedEarly {
let settle = max(0, min(settleAfterStop, deadline.timeIntervalSinceNow))
if settle > 0 {
let settleDeadline = Date().addingTimeInterval(settle)
while Date() < settleDeadline {
let newData = self.readChunk()
if !newData.isEmpty {
buffer.append(newData)
}
try await Task.sleep(nanoseconds: 50_000_000)
}
}
}
guard !buffer.isEmpty, let text = String(data: buffer, encoding: .utf8) else {
throw SessionError.timedOut
}
return text
}
private static func appendScanText(newData: Data, scanTailText: inout String, utf8Carry: inout Data) {
// PTY reads can split multibyte UTF-8 sequences. Keep a small carry buffer so prompt/stop scanning doesn't
// drop chunks when the decode fails due to an incomplete trailing sequence.
var combined = Data()
combined.reserveCapacity(utf8Carry.count + newData.count)
combined.append(utf8Carry)
combined.append(newData)
if let chunk = String(data: combined, encoding: .utf8) {
scanTailText.append(chunk)
utf8Carry.removeAll(keepingCapacity: true)
return
}
for trimCount in 1...3 where combined.count > trimCount {
let prefix = combined.dropLast(trimCount)
if let chunk = String(data: prefix, encoding: .utf8) {
scanTailText.append(chunk)
utf8Carry = Data(combined.suffix(trimCount))
return
}
}
// If the data is still not UTF-8 decodable, keep only a small suffix to avoid unbounded growth.
utf8Carry = Data(combined.suffix(12))
}
func reset() {
self.cleanup()
}
private func ensureStarted(binary: String) throws {
if let proc = self.process, proc.isRunning, self.binaryPath == binary {
Self.log.debug("Claude CLI session reused")
return
}
self.cleanup()
var primaryFD: Int32 = -1
var secondaryFD: Int32 = -1
var win = winsize(ws_row: 50, ws_col: 160, ws_xpixel: 0, ws_ypixel: 0)
guard openpty(&primaryFD, &secondaryFD, nil, nil, &win) == 0 else {
Self.log.warning("Claude CLI PTY openpty failed")
throw SessionError.launchFailed("openpty failed")
}
_ = fcntl(primaryFD, F_SETFL, O_NONBLOCK)
let primaryHandle = FileHandle(fileDescriptor: primaryFD, closeOnDealloc: true)
let secondaryHandle = FileHandle(fileDescriptor: secondaryFD, closeOnDealloc: true)
let proc = Process()
let resolvedURL = URL(fileURLWithPath: binary)
let disableWatchdog = ProcessInfo.processInfo.environment["CODEXBAR_DISABLE_CLAUDE_WATCHDOG"] == "1"
if !disableWatchdog,
resolvedURL.lastPathComponent == "claude",
let watchdog = TTYCommandRunner.locateBundledHelper("CodexBarClaudeWatchdog")
{
proc.executableURL = URL(fileURLWithPath: watchdog)
proc.arguments = ["--", binary, "--allowed-tools", ""]
} else {
proc.executableURL = resolvedURL
proc.arguments = ["--allowed-tools", ""]
}
proc.standardInput = secondaryHandle
proc.standardOutput = secondaryHandle
proc.standardError = secondaryHandle
let workingDirectory = ClaudeStatusProbe.preparedProbeWorkingDirectoryURL()
proc.currentDirectoryURL = workingDirectory
var env = Self.launchEnvironment()
env["PWD"] = workingDirectory.path
proc.environment = env
guard TTYCommandRunner.beginActiveProcessLaunchForAppShutdown() else {
try? primaryHandle.close()
try? secondaryHandle.close()
throw SessionError.launchFailed("App shutdown in progress")
}
defer { TTYCommandRunner.endActiveProcessLaunchForAppShutdown() }
do {
try proc.run()
Self.log.debug(
"Claude CLI session started",
metadata: ["binary": URL(fileURLWithPath: binary).lastPathComponent])
} catch {
Self.log.warning("Claude CLI launch failed", metadata: ["error": error.localizedDescription])
try? primaryHandle.close()
try? secondaryHandle.close()
throw SessionError.launchFailed(error.localizedDescription)
}
let pid = proc.processIdentifier
guard TTYCommandRunner.registerActiveProcessForAppShutdown(
pid: pid,
binary: URL(fileURLWithPath: binary).lastPathComponent)
else {
proc.terminate()
kill(pid, SIGKILL)
try? primaryHandle.close()
try? secondaryHandle.close()
throw SessionError.launchFailed("App shutdown in progress")
}
var processGroup: pid_t?
if setpgid(pid, pid) == 0 {
processGroup = pid
TTYCommandRunner.updateActiveProcessGroupForAppShutdown(pid: pid, processGroup: processGroup)
}
self.process = proc
self.primaryFD = primaryFD
self.primaryHandle = primaryHandle
self.secondaryHandle = secondaryHandle
self.processGroup = processGroup
self.binaryPath = binary
self.startedAt = Date()
}
static func launchEnvironment(baseEnv: [String: String] = ProcessInfo.processInfo.environment) -> [String: String] {
var env = self.scrubbedClaudeEnvironment(from: TTYCommandRunner.enrichedEnvironment(baseEnv: baseEnv))
// Passive status and auth probes must not mutate or update the user's Claude CLI installation.
env["DISABLE_AUTOUPDATER"] = "1"
return env
}
private static func scrubbedClaudeEnvironment(from base: [String: String]) -> [String: String] {
var env = base
let explicitKeys: [String] = [
ClaudeOAuthCredentialsStore.environmentTokenKey,
ClaudeOAuthCredentialsStore.environmentScopesKey,
]
for key in explicitKeys {
env.removeValue(forKey: key)
}
for key in env.keys where key.hasPrefix("ANTHROPIC_") {
env.removeValue(forKey: key)
}
return env
}
private func cleanup() {
if self.process != nil {
Self.log.debug("Claude CLI session stopping")
}
if let proc = self.process, proc.isRunning {
try? self.writeAllToPrimary(Data("/exit\r".utf8))
}
try? self.primaryHandle?.close()
try? self.secondaryHandle?.close()
let descendants = self.process.map { TTYProcessTreeTerminator.descendantPIDs(of: $0.processIdentifier) } ?? []
if let proc = self.process, proc.isRunning {
proc.terminate()
}
if let proc = self.process {
TTYProcessTreeTerminator.terminateProcessTree(
rootPID: proc.processIdentifier,
processGroup: self.processGroup,
signal: SIGTERM,
knownDescendants: descendants)
}
let waitDeadline = Date().addingTimeInterval(1.0)
if let proc = self.process {
while proc.isRunning, Date() < waitDeadline {
usleep(100_000)
}
if proc.isRunning {
TTYProcessTreeTerminator.terminateProcessTree(
rootPID: proc.processIdentifier,
processGroup: self.processGroup,
signal: SIGKILL,
knownDescendants: descendants)
} else {
for pid in descendants where pid > 0 {
kill(pid, SIGKILL)
}
}
TTYCommandRunner.unregisterActiveProcessForAppShutdown(pid: proc.processIdentifier)
}
self.process = nil
self.primaryHandle = nil
self.secondaryHandle = nil
self.primaryFD = -1
self.processGroup = nil
self.startedAt = nil
}
private func readChunk() -> Data {
guard self.primaryFD >= 0 else { return Data() }
var appended = Data()
while true {
var tmp = [UInt8](repeating: 0, count: 8192)
let n = read(self.primaryFD, &tmp, tmp.count)
if n > 0 {
appended.append(contentsOf: tmp.prefix(n))
continue
}
break
}
return appended
}
private func drainOutput() {
_ = self.readChunk()
}
private func shouldStopForIdleTimeout(
idleTimeout: TimeInterval?,
bufferIsEmpty: Bool,
lastOutputAt: Date) -> Bool
{
guard let idleTimeout, !bufferIsEmpty else { return false }
return Date().timeIntervalSince(lastOutputAt) >= idleTimeout
}
private func sendPeriodicEnterIfNeeded(every: TimeInterval?, lastEnterAt: inout Date) {
guard let every, Date().timeIntervalSince(lastEnterAt) >= every else { return }
try? self.send("\r")
lastEnterAt = Date()
}
private func send(_ text: String) throws {
guard let data = text.data(using: .utf8) else { return }
guard self.primaryFD >= 0 else { throw SessionError.processExited }
try self.writeAllToPrimary(data)
}
private func writeAllToPrimary(_ data: Data) throws {
guard self.primaryFD >= 0 else { throw SessionError.processExited }
try data.withUnsafeBytes { rawBytes in
guard let baseAddress = rawBytes.baseAddress else { return }
var offset = 0
var retries = 0
while offset < rawBytes.count {
let written = write(self.primaryFD, baseAddress.advanced(by: offset), rawBytes.count - offset)
if written > 0 {
offset += written
retries = 0
continue
}
if written == 0 {
break
}
let err = errno
if err == EINTR || err == EAGAIN || err == EWOULDBLOCK {
retries += 1
if retries > 200 {
throw SessionError.ioFailed("write to PTY would block")
}
usleep(5000)
continue
}
throw SessionError.ioFailed("write to PTY failed: \(String(cString: strerror(err)))")
}
}
}
}
@@ -0,0 +1,102 @@
import Foundation
public enum ClaudeCredentialRouting: Sendable, Equatable {
case none
case oauth(accessToken: String)
case webCookie(header: String)
case adminAPIKey(String)
public static func resolve(tokenAccountToken: String?, manualCookieHeader: String?) -> Self {
if let tokenAccountToken,
let route = self.resolvePrimaryCredential(tokenAccountToken)
{
return route
}
guard let manualCookieHeader = self.normalizedWebCookie(manualCookieHeader) else {
return .none
}
return .webCookie(header: manualCookieHeader)
}
public var oauthAccessToken: String? {
guard case let .oauth(accessToken) = self else { return nil }
return accessToken
}
public var manualCookieHeader: String? {
guard case let .webCookie(header) = self else { return nil }
return header
}
public var isOAuth: Bool {
if case .oauth = self { return true }
return false
}
public var adminAPIKey: String? {
guard case let .adminAPIKey(key) = self else { return nil }
return key
}
private static func resolvePrimaryCredential(_ raw: String) -> Self? {
if let adminKey = self.normalizedAdminAPIKey(raw) {
return .adminAPIKey(adminKey)
}
if let accessToken = self.normalizedOAuthToken(raw) {
return .oauth(accessToken: accessToken)
}
if let cookieHeader = self.normalizedWebCookie(raw) {
return .webCookie(header: cookieHeader)
}
return nil
}
private static func normalizedOAuthToken(_ raw: String?) -> String? {
guard let trimmed = self.cleaned(raw) else { return nil }
let lower = trimmed.lowercased()
if lower.contains("cookie:") || trimmed.contains("=") {
return nil
}
if lower.hasPrefix("bearer ") {
let bearerTrimmed = trimmed.dropFirst("bearer ".count)
.trimmingCharacters(in: .whitespacesAndNewlines)
guard !bearerTrimmed.isEmpty else { return nil }
let lowerBearerTrimmed = bearerTrimmed.lowercased()
return lowerBearerTrimmed.hasPrefix("sk-ant-oat") ? bearerTrimmed : nil
}
return lower.hasPrefix("sk-ant-oat") ? trimmed : nil
}
private static func normalizedAdminAPIKey(_ raw: String?) -> String? {
guard let trimmed = self.cleaned(raw) else { return nil }
let lower = trimmed.lowercased()
if lower.contains("cookie:") || trimmed.contains("=") {
return nil
}
if lower.hasPrefix("bearer ") {
let bearerTrimmed = trimmed.dropFirst("bearer ".count)
.trimmingCharacters(in: .whitespacesAndNewlines)
guard !bearerTrimmed.isEmpty else { return nil }
return bearerTrimmed.lowercased().hasPrefix("sk-ant-admin") ? bearerTrimmed : nil
}
return lower.hasPrefix("sk-ant-admin") ? trimmed : nil
}
private static func normalizedWebCookie(_ raw: String?) -> String? {
guard let normalized = CookieHeaderNormalizer.normalize(raw) else { return nil }
if normalized.contains("=") {
return normalized
}
return "sessionKey=\(normalized)"
}
private static func cleaned(_ raw: String?) -> String? {
guard let trimmed = raw?.trimmingCharacters(in: .whitespacesAndNewlines),
!trimmed.isEmpty
else {
return nil
}
return trimmed
}
}
@@ -0,0 +1,89 @@
import Foundation
public enum ClaudeDesktopProjectsLocator {
private static let sessionDirectoryNames = [
"local-agent-mode-sessions",
"claude-code-sessions",
]
private static let skippedDirectoryNames = Set([
".build",
".git",
"build",
"DerivedData",
"node_modules",
"outputs",
"target",
])
public static func roots(
homeDirectory: URL = FileManager.default.homeDirectoryForCurrentUser,
fileManager: FileManager = .default) -> [URL]
{
let claudeApplicationSupport = homeDirectory
.appendingPathComponent("Library", isDirectory: true)
.appendingPathComponent("Application Support", isDirectory: true)
.appendingPathComponent("Claude", isDirectory: true)
let sessionRoots = self.sessionDirectoryNames.map {
claudeApplicationSupport.appendingPathComponent($0, isDirectory: true)
}
var roots: [URL] = []
var queue = sessionRoots.map { (url: $0, depth: 0) }
var visited = Set(sessionRoots.map(\.standardizedFileURL.path))
var nextIndex = 0
// Current Desktop entries under claude-code-sessions are metadata whose cliSessionId maps to the
// shared ~/.claude/projects logs. Seed that root anyway so embedded project stores are found when present.
// Covers observed Desktop layouts through account/workspace, session, agent, and local_agent
// without descending into arbitrary checked-out workspaces.
let maxDepth = 4
while nextIndex < queue.count {
let current = queue[nextIndex]
nextIndex += 1
if let projects = self.projectsRoot(under: current.url, fileManager: fileManager) {
roots.append(projects)
}
guard current.depth < maxDepth else { continue }
for child in self.childDirectories(at: current.url, fileManager: fileManager) {
let standardized = child.standardizedFileURL
guard visited.insert(standardized.path).inserted else { continue }
queue.append((standardized, current.depth + 1))
}
}
return roots
}
private static func projectsRoot(under base: URL, fileManager: FileManager) -> URL? {
let projects = base
.appendingPathComponent(".claude", isDirectory: true)
.appendingPathComponent("projects", isDirectory: true)
var isDirectory: ObjCBool = false
guard fileManager.fileExists(atPath: projects.path, isDirectory: &isDirectory),
isDirectory.boolValue
else { return nil }
return projects.standardizedFileURL
}
private static func childDirectories(at url: URL, fileManager: FileManager) -> [URL] {
guard let children = try? fileManager.contentsOfDirectory(
at: url,
includingPropertiesForKeys: [.isDirectoryKey, .isSymbolicLinkKey],
options: [.skipsHiddenFiles, .skipsPackageDescendants])
else {
return []
}
return children.compactMap { child in
guard !self.skippedDirectoryNames.contains(child.lastPathComponent) else { return nil }
guard let values = try? child.resourceValues(forKeys: [.isDirectoryKey, .isSymbolicLinkKey]),
values.isSymbolicLink != true,
values.isDirectory == true
else {
return nil
}
return child
}
}
}
@@ -0,0 +1,301 @@
import Foundation
#if canImport(CryptoKit)
import CryptoKit
#else
import Crypto
#endif
#if os(macOS)
import Security
#endif
public struct ClaudeOAuthCredentials: Sendable {
public let accessToken: String
public let refreshToken: String?
public let expiresAt: Date?
public let scopes: [String]
public let rateLimitTier: String?
public let subscriptionType: String?
public init(
accessToken: String,
refreshToken: String?,
expiresAt: Date?,
scopes: [String],
rateLimitTier: String?,
subscriptionType: String? = nil)
{
self.accessToken = accessToken
self.refreshToken = refreshToken
self.expiresAt = expiresAt
self.scopes = scopes
self.rateLimitTier = rateLimitTier
self.subscriptionType = subscriptionType
}
public var isExpired: Bool {
guard let expiresAt else { return true }
return Date() >= expiresAt
}
public var expiresIn: TimeInterval? {
guard let expiresAt else { return nil }
return expiresAt.timeIntervalSinceNow
}
/// A one-way discriminator for history owned by this credential.
///
/// Prefer the refresh token because access tokens routinely rotate for the same principal. If a provider
/// supplies only an access token, rotating that token intentionally starts a new history bucket rather than
/// risking that two identityless accounts share one. The source secret never leaves this computation.
var historyOwnerIdentifier: String? {
let normalizedRefreshToken = self.refreshToken?.trimmingCharacters(in: .whitespacesAndNewlines)
if let normalizedRefreshToken, !normalizedRefreshToken.isEmpty {
return Self.makeHistoryOwnerIdentifier(secretKind: "refresh", secret: normalizedRefreshToken)
}
let normalizedAccessToken = self.accessToken.trimmingCharacters(in: .whitespacesAndNewlines)
guard !normalizedAccessToken.isEmpty else { return nil }
return Self.makeHistoryOwnerIdentifier(secretKind: "access", secret: normalizedAccessToken)
}
static func historyOwnerIdentifier(forRefreshToken refreshToken: String) -> String? {
let normalized = refreshToken.trimmingCharacters(in: .whitespacesAndNewlines)
guard !normalized.isEmpty else { return nil }
return Self.makeHistoryOwnerIdentifier(secretKind: "refresh", secret: normalized)
}
private static func makeHistoryOwnerIdentifier(secretKind: String, secret: String) -> String? {
let material = Data("codexbar:claude-oauth-history-owner:v1\0\(secretKind)\0\(secret)".utf8)
return SHA256.hash(data: material).map { String(format: "%02x", $0) }.joined()
}
static func normalizedHistoryOwnerIdentifier(_ identifier: String?) -> String? {
guard let identifier else { return nil }
let normalized = identifier.trimmingCharacters(in: .whitespacesAndNewlines).lowercased()
guard normalized.count == 64,
normalized.unicodeScalars.allSatisfy({ scalar in
switch scalar.value {
case 48...57, 97...102:
true
default:
false
}
})
else { return nil }
return normalized
}
public static func isMcpOAuthOnlyPayload(data: Data) -> Bool {
guard let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] else {
return false
}
return json["claudeAiOauth"] == nil && json["mcpOAuth"] != nil
}
public static func parse(data: Data) throws -> ClaudeOAuthCredentials {
if ClaudeOAuthCredentials.isMcpOAuthOnlyPayload(data: data) {
throw ClaudeOAuthCredentialsError.mcpOAuthOnlyKeychain
}
let decoder = JSONDecoder()
guard let root = try? decoder.decode(Root.self, from: data) else {
throw ClaudeOAuthCredentialsError.decodeFailed
}
guard let oauth = root.claudeAiOauth else {
throw ClaudeOAuthCredentialsError.missingOAuth
}
let accessToken = oauth.accessToken?.trimmingCharacters(in: .whitespacesAndNewlines) ?? ""
guard !accessToken.isEmpty else {
throw ClaudeOAuthCredentialsError.missingAccessToken
}
let expiresAt = oauth.expiresAt.map { millis in
Date(timeIntervalSince1970: millis / 1000.0)
}
return ClaudeOAuthCredentials(
accessToken: accessToken,
refreshToken: oauth.refreshToken,
expiresAt: expiresAt,
scopes: oauth.scopes ?? [],
rateLimitTier: oauth.rateLimitTier,
subscriptionType: oauth.subscriptionType)
}
private struct Root: Decodable {
let claudeAiOauth: OAuth?
}
private struct OAuth: Decodable {
let accessToken: String?
let refreshToken: String?
let expiresAt: Double?
let scopes: [String]?
let rateLimitTier: String?
let subscriptionType: String?
enum CodingKeys: String, CodingKey {
case accessToken
case refreshToken
case expiresAt
case scopes
case rateLimitTier
case subscriptionType
}
}
}
extension ClaudeOAuthCredentials {
func diagnosticsMetadata(now: Date = Date()) -> [String: String] {
let hasRefreshToken = !(self.refreshToken?.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty ?? true)
let hasUserProfileScope = self.scopes.contains("user:profile")
var metadata: [String: String] = [
"hasRefreshToken": "\(hasRefreshToken)",
"scopesCount": "\(self.scopes.count)",
"hasUserProfileScope": "\(hasUserProfileScope)",
]
if let expiresAt = self.expiresAt {
let expiresAtMs = Int(expiresAt.timeIntervalSince1970 * 1000.0)
let expiresInSec = Int(expiresAt.timeIntervalSince(now).rounded())
metadata["expiresAtMs"] = "\(expiresAtMs)"
metadata["expiresInSec"] = "\(expiresInSec)"
metadata["isExpired"] = "\(now >= expiresAt)"
} else {
metadata["expiresAtMs"] = "nil"
metadata["expiresInSec"] = "nil"
metadata["isExpired"] = "true"
}
return metadata
}
}
public enum ClaudeOAuthCredentialOwner: String, Codable, Sendable {
case claudeCLI
case codexbar
case environment
}
public enum ClaudeOAuthCredentialSource: String, Sendable {
case environment
case memoryCache
case cacheKeychain
case credentialsFile
case claudeKeychain
}
enum ClaudeKeychainCredentialMatch: Equatable, Sendable {
case notApplicable
case absent
case unavailable
case mismatch
case matched(persistentRefHash: String)
var persistentRefHash: String? {
guard case let .matched(persistentRefHash) = self else { return nil }
return persistentRefHash
}
var isMismatch: Bool {
self == .mismatch
}
var isAbsent: Bool {
self == .absent
}
var isUnavailable: Bool {
self == .unavailable
}
}
public struct ClaudeOAuthCredentialRecord: Sendable {
public let credentials: ClaudeOAuthCredentials
public let owner: ClaudeOAuthCredentialOwner
public let source: ClaudeOAuthCredentialSource
private let inheritedHistoryOwnerIdentifier: String?
/// An opaque, one-way owner identifier that survives a refresh-token rotation proven by a successful refresh.
/// Records from unrelated credential sources do not inherit this value and derive a fresh identifier instead.
var historyOwnerIdentifier: String? {
self.inheritedHistoryOwnerIdentifier ?? self.credentials.historyOwnerIdentifier
}
public init(
credentials: ClaudeOAuthCredentials,
owner: ClaudeOAuthCredentialOwner,
source: ClaudeOAuthCredentialSource)
{
self.credentials = credentials
self.owner = owner
self.source = source
self.inheritedHistoryOwnerIdentifier = nil
}
init(
credentials: ClaudeOAuthCredentials,
owner: ClaudeOAuthCredentialOwner,
source: ClaudeOAuthCredentialSource,
historyOwnerIdentifier: String?)
{
self.credentials = credentials
self.owner = owner
self.source = source
self.inheritedHistoryOwnerIdentifier = ClaudeOAuthCredentials.normalizedHistoryOwnerIdentifier(
historyOwnerIdentifier)
}
}
public enum ClaudeOAuthCredentialsError: LocalizedError, Sendable {
case decodeFailed
case missingOAuth
case mcpOAuthOnlyKeychain
case missingAccessToken
case notFound
case keychainError(Int)
case readFailed(String)
case refreshFailed(String)
case noRefreshToken
case refreshDelegatedToClaudeCLI
public var errorDescription: String? {
switch self {
case .decodeFailed:
return "Claude OAuth credentials are invalid."
case .missingOAuth:
return "Claude OAuth credentials missing. Run `claude` to authenticate."
case .mcpOAuthOnlyKeychain:
return "Claude keychain contains MCP OAuth state only (no claudeAiOauth). "
+ "Claude Code may store subscription OAuth elsewhere now. "
+ "Open the CodexBar menu and click Refresh to re-authenticate, "
+ "or switch Claude Usage source to Web/CLI."
case .missingAccessToken:
return "Claude OAuth access token missing. Run `claude` to authenticate."
case .notFound:
return "Claude OAuth credentials not found. Run `claude` to authenticate."
case let .keychainError(status):
#if os(macOS)
if status == Int(errSecUserCanceled)
|| status == Int(errSecAuthFailed)
|| status == Int(errSecInteractionNotAllowed)
|| status == Int(errSecNoAccessForItem)
{
return "Claude Keychain access was denied. CodexBar will back off in the background until you retry "
+ "via a user action (menu open / manual refresh). "
+ "Switch Claude Usage source to Web/CLI, or allow access in Keychain Access."
}
#endif
return "Claude OAuth keychain error: \(status)"
case let .readFailed(message):
return "Claude OAuth credentials read failed: \(message)"
case let .refreshFailed(message):
return "Claude OAuth token refresh failed: \(message)"
case .noRefreshToken:
return "Claude OAuth refresh token missing. Run `claude` to authenticate."
case .refreshDelegatedToClaudeCLI:
return "Claude OAuth refresh is delegated to Claude CLI."
}
}
}
@@ -0,0 +1,18 @@
import Foundation
#if canImport(CryptoKit)
import CryptoKit
#endif
extension ClaudeOAuthCredentialsStore {
static func sha256Prefix(_ data: Data) -> String? {
#if canImport(CryptoKit)
let digest = SHA256.hash(data: data)
let hex = digest.compactMap { String(format: "%02x", $0) }.joined()
return String(hex.prefix(12))
#else
_ = data
return nil
#endif
}
}
@@ -0,0 +1,375 @@
import Dispatch
import Foundation
#if canImport(Darwin)
import Darwin
#elseif canImport(Glibc)
import Glibc
#elseif canImport(Musl)
import Musl
#endif
extension ClaudeOAuthCredentialsStore {
private static let securityBinaryPath = "/usr/bin/security"
private static let securityCLIReadTimeout: TimeInterval = 1.5
static let isolatedSecurityCLIKeychainEnvironmentKey = "CODEXBAR_CLAUDE_SECURITY_CLI_KEYCHAIN"
struct SecurityCLIReadRequest {
let account: String?
}
static func shouldPreferSecurityCLIKeychainRead(
readStrategy: ClaudeOAuthKeychainReadStrategy = ClaudeOAuthKeychainReadStrategyPreference.current())
-> Bool
{
readStrategy == .securityCLIExperimental
}
#if os(macOS)
private enum SecurityCLIReadError: Error {
case binaryUnavailable
case isolatedKeychainUnavailable
case launchFailed
case timedOut
case nonZeroExit(status: Int32, stderrLength: Int)
}
private struct SecurityCLIReadCommandResult {
let status: Int32
let stdout: Data
let stderrLength: Int
let durationMs: Double
}
/// Attempts a Claude keychain read via `/usr/bin/security` when the experimental reader is enabled.
/// - Important: `interaction` is diagnostics context only and does not gate CLI execution.
static func loadFromClaudeKeychainViaSecurityCLIIfEnabled(
interaction: ProviderInteraction,
readStrategy: ClaudeOAuthKeychainReadStrategy = ClaudeOAuthKeychainReadStrategyPreference.current())
-> Data?
{
guard let sanitized = self.readRawClaudeKeychainPayloadViaSecurityCLIIfEnabled(
interaction: interaction,
readStrategy: readStrategy)
else {
return nil
}
let interactionMetadata = interaction == .userInitiated ? "user" : "background"
let parsedCredentials: ClaudeOAuthCredentials
do {
parsedCredentials = try ClaudeOAuthCredentials.parse(data: sanitized)
} catch {
self.log.warning(
"Claude keychain security CLI output invalid; falling back",
metadata: [
"reader": "securityCLI",
"callerInteraction": interactionMetadata,
"payload_bytes": "\(sanitized.count)",
"parse_error_type": String(describing: type(of: error)),
])
return nil
}
var metadata: [String: String] = [
"reader": "securityCLI",
"callerInteraction": interactionMetadata,
"payload_bytes": "\(sanitized.count)",
]
for (key, value) in parsedCredentials.diagnosticsMetadata(now: Date()) {
metadata[key] = value
}
self.log.debug(
"Claude keychain security CLI read succeeded",
metadata: metadata)
return sanitized
}
static func readRawClaudeKeychainPayloadViaSecurityCLIIfEnabled(
interaction: ProviderInteraction,
readStrategy: ClaudeOAuthKeychainReadStrategy = ClaudeOAuthKeychainReadStrategyPreference.current(),
environment: [String: String] = ProcessInfo.processInfo.environment)
-> Data?
{
guard self.shouldPreferSecurityCLIKeychainRead(readStrategy: readStrategy) else { return nil }
let interactionMetadata = interaction == .userInitiated ? "user" : "background"
do {
let preferredAccount = self.preferredClaudeKeychainAccountForSecurityCLIRead(
interaction: interaction)
let output: Data
let status: Int32
let stderrLength: Int
let durationMs: Double
#if DEBUG
if let override = self.taskSecurityCLIReadOverride ?? self.securityCLIReadOverride {
switch override {
case let .data(data):
output = data ?? Data()
status = 0
stderrLength = 0
durationMs = 0
case .timedOut:
throw SecurityCLIReadError.timedOut
case .nonZeroExit:
throw SecurityCLIReadError.nonZeroExit(status: 1, stderrLength: 0)
case let .dynamic(read):
output = read(SecurityCLIReadRequest(account: preferredAccount)) ?? Data()
status = 0
stderrLength = 0
durationMs = 0
}
} else {
let result = try self.runClaudeSecurityCLIRead(
timeout: self.securityCLIReadTimeout,
account: preferredAccount,
environment: environment)
output = result.stdout
status = result.status
stderrLength = result.stderrLength
durationMs = result.durationMs
}
#else
let result = try self.runClaudeSecurityCLIRead(
timeout: self.securityCLIReadTimeout,
account: preferredAccount,
environment: environment)
output = result.stdout
status = result.status
stderrLength = result.stderrLength
durationMs = result.durationMs
#endif
let sanitized = self.sanitizeSecurityCLIOutput(output)
guard !sanitized.isEmpty else { return nil }
if ClaudeOAuthCredentials.isMcpOAuthOnlyPayload(data: sanitized) {
self.log.warning(
"Claude keychain security CLI output is MCP OAuth only; falling back",
metadata: [
"reader": "securityCLI",
"callerInteraction": interactionMetadata,
"status": "\(status)",
"duration_ms": String(format: "%.2f", durationMs),
"stderr_length": "\(stderrLength)",
"payload_bytes": "\(sanitized.count)",
])
} else {
self.log.debug(
"Claude keychain security CLI raw read succeeded",
metadata: [
"reader": "securityCLI",
"callerInteraction": interactionMetadata,
"status": "\(status)",
"duration_ms": String(format: "%.2f", durationMs),
"stderr_length": "\(stderrLength)",
"payload_bytes": "\(sanitized.count)",
])
}
return sanitized
} catch let error as SecurityCLIReadError {
var metadata: [String: String] = [
"reader": "securityCLI",
"callerInteraction": interactionMetadata,
"error_type": String(describing: type(of: error)),
]
switch error {
case .binaryUnavailable:
metadata["reason"] = "binaryUnavailable"
case .isolatedKeychainUnavailable:
metadata["reason"] = "isolatedKeychainUnavailable"
case .launchFailed:
metadata["reason"] = "launchFailed"
case .timedOut:
metadata["reason"] = "timedOut"
case let .nonZeroExit(status, stderrLength):
metadata["reason"] = "nonZeroExit"
metadata["status"] = "\(status)"
metadata["stderr_length"] = "\(stderrLength)"
}
self.log.warning("Claude keychain security CLI read failed; falling back", metadata: metadata)
return nil
} catch {
self.log.warning(
"Claude keychain security CLI read failed; falling back",
metadata: [
"reader": "securityCLI",
"callerInteraction": interactionMetadata,
"error_type": String(describing: type(of: error)),
])
return nil
}
}
private static func sanitizeSecurityCLIOutput(_ data: Data) -> Data {
var sanitized = data
while let last = sanitized.last, last == 0x0A || last == 0x0D {
sanitized.removeLast()
}
return sanitized
}
private static func runClaudeSecurityCLIRead(
timeout: TimeInterval,
account: String?,
environment: [String: String]) throws -> SecurityCLIReadCommandResult
{
guard FileManager.default.isExecutableFile(atPath: self.securityBinaryPath) else {
throw SecurityCLIReadError.binaryUnavailable
}
guard let arguments = self.securityCLIReadArguments(account: account, environment: environment) else {
throw SecurityCLIReadError.isolatedKeychainUnavailable
}
let process = Process()
process.executableURL = URL(fileURLWithPath: self.securityBinaryPath)
process.arguments = arguments
let stdoutPipe = Pipe()
let stderrPipe = Pipe()
process.standardOutput = stdoutPipe
process.standardError = stderrPipe
process.standardInput = nil
let startedAt = DispatchTime.now().uptimeNanoseconds
do {
try process.run()
} catch {
throw SecurityCLIReadError.launchFailed
}
var processGroup: pid_t?
let pid = process.processIdentifier
if setpgid(pid, pid) == 0 {
processGroup = pid
}
let deadline = Date().addingTimeInterval(timeout)
while process.isRunning, Date() < deadline {
Thread.sleep(forTimeInterval: 0.02)
}
if process.isRunning {
self.terminate(process: process, processGroup: processGroup)
throw SecurityCLIReadError.timedOut
}
let stdout = stdoutPipe.fileHandleForReading.readDataToEndOfFile()
let stderr = stderrPipe.fileHandleForReading.readDataToEndOfFile()
let status = process.terminationStatus
let durationMs = Double(DispatchTime.now().uptimeNanoseconds - startedAt) / 1_000_000.0
guard status == 0 else {
throw SecurityCLIReadError.nonZeroExit(status: status, stderrLength: stderr.count)
}
return SecurityCLIReadCommandResult(
status: status,
stdout: stdout,
stderrLength: stderr.count,
durationMs: durationMs)
}
static func securityCLIReadArguments(
account: String?,
environment: [String: String]) -> [String]?
{
let isolatedKeychainPath = self.isolatedSecurityCLIKeychainPath(environment: environment)
if KeychainTestSafety.shouldBlockRealKeychainAccess(environment: environment),
isolatedKeychainPath == nil
{
return nil
}
var arguments = [
"find-generic-password",
"-s",
self.claudeKeychainService,
]
if let account, !account.isEmpty {
arguments.append(contentsOf: ["-a", account])
}
arguments.append("-w")
let rawIsolatedPath = environment[self.isolatedSecurityCLIKeychainEnvironmentKey]?
.trimmingCharacters(in: .whitespacesAndNewlines)
if rawIsolatedPath != nil || KeychainAccessGate.isDisabledByEnvironment(environment) {
guard let isolatedKeychainPath else {
return nil
}
arguments.append(isolatedKeychainPath)
}
return arguments
}
private static func terminate(process: Process, processGroup: pid_t?) {
guard process.isRunning else { return }
process.terminate()
if let processGroup {
kill(-processGroup, SIGTERM)
}
let deadline = Date().addingTimeInterval(0.4)
while process.isRunning, Date() < deadline {
usleep(50000)
}
if process.isRunning {
if let processGroup {
kill(-processGroup, SIGKILL)
}
kill(process.processIdentifier, SIGKILL)
}
}
#else
static func loadFromClaudeKeychainViaSecurityCLIIfEnabled(
interaction _: ProviderInteraction,
readStrategy _: ClaudeOAuthKeychainReadStrategy = ClaudeOAuthKeychainReadStrategyPreference.current())
-> Data?
{
nil
}
static func readRawClaudeKeychainPayloadViaSecurityCLIIfEnabled(
interaction _: ProviderInteraction,
readStrategy _: ClaudeOAuthKeychainReadStrategy = ClaudeOAuthKeychainReadStrategyPreference.current(),
environment _: [String: String] = ProcessInfo.processInfo.environment)
-> Data?
{
nil
}
#endif
private static func isolatedSecurityCLIKeychainPath(environment: [String: String]) -> String? {
guard KeychainAccessGate.isDisabledByEnvironment(environment) else { return nil }
guard let rawPath = environment[self.isolatedSecurityCLIKeychainEnvironmentKey]?
.trimmingCharacters(in: .whitespacesAndNewlines),
!rawPath.isEmpty,
(rawPath as NSString).isAbsolutePath
else {
return nil
}
return URL(fileURLWithPath: rawPath).standardizedFileURL.path
}
static func isMcpOAuthOnlyClaudeKeychainPayloadPresent(
interaction: ProviderInteraction,
readStrategy: ClaudeOAuthKeychainReadStrategy = ClaudeOAuthKeychainReadStrategyPreference.current(),
keychainAccessDisabled: Bool = KeychainAccessGate.isDisabled,
environment: [String: String] = ProcessInfo.processInfo.environment) -> Bool
{
guard !keychainAccessDisabled || self.isolatedSecurityCLIKeychainPath(environment: environment) != nil else {
return false
}
guard ClaudeOAuthKeychainPromptPreference.effectiveMode(readStrategy: readStrategy) != .never else {
return false
}
let payload: Data? = switch readStrategy {
case .securityFramework:
self.readRawClaudeKeychainPayloadViaSecurityFrameworkWithoutPrompt()
case .securityCLIExperimental:
self.readRawClaudeKeychainPayloadViaSecurityCLIIfEnabled(
interaction: interaction,
readStrategy: readStrategy,
environment: environment)
}
guard let payload else { return false }
return ClaudeOAuthCredentials.isMcpOAuthOnlyPayload(data: payload)
}
}
@@ -0,0 +1,430 @@
import Foundation
#if DEBUG
extension ClaudeOAuthCredentialsStore {
nonisolated(unsafe) static var claudeKeychainDataOverride: Data?
nonisolated(unsafe) static var claudeKeychainFingerprintOverride: ClaudeKeychainFingerprint?
@TaskLocal static var taskClaudeKeychainDataOverride: Data?
@TaskLocal static var taskClaudeKeychainFingerprintOverride: ClaudeKeychainFingerprint?
@TaskLocal static var taskMemoryCacheStoreOverride: MemoryCacheStore?
@TaskLocal static var taskClaudeKeychainFingerprintStoreOverride: ClaudeKeychainFingerprintStore?
@TaskLocal static var taskPendingCacheClearStoreOverride: ClaudeOAuthPendingCacheClearStore?
typealias OAuthCacheOperation = KeychainCacheStore.Operation
typealias OAuthCacheOperationRecorder = KeychainCacheStore.OperationRecorder
final class PendingCacheClearMemoryStore: ClaudeOAuthPendingCacheClearStore, @unchecked Sendable {
private let lock = NSLock()
private var pending: Bool
init(isPending: Bool = false) {
self.pending = isPending
}
var isPending: Bool {
self.lock.lock()
defer { self.lock.unlock() }
return self.pending
}
func markPending() {
self.lock.lock()
self.pending = true
self.lock.unlock()
}
func withCacheTransaction(_ operation: (inout Bool) -> Void) {
self.lock.lock()
defer { self.lock.unlock() }
operation(&self.pending)
}
}
final class ClaudeKeychainFingerprintStore: @unchecked Sendable {
var fingerprint: ClaudeKeychainFingerprint?
init(fingerprint: ClaudeKeychainFingerprint? = nil) {
self.fingerprint = fingerprint
}
}
final class MemoryCacheStore: @unchecked Sendable {
var record: ClaudeOAuthCredentialRecord?
var timestamp: Date?
}
static func setClaudeKeychainDataOverrideForTesting(_ data: Data?) {
self.claudeKeychainDataOverride = data
}
static func setClaudeKeychainFingerprintOverrideForTesting(_ fingerprint: ClaudeKeychainFingerprint?) {
self.claudeKeychainFingerprintOverride = fingerprint
}
static func withClaudeKeychainOverridesForTesting<T>(
data: Data?,
fingerprint: ClaudeKeychainFingerprint?,
operation: () throws -> T) rethrows -> T
{
try self.$taskClaudeKeychainDataOverride.withValue(data) {
try self.$taskClaudeKeychainFingerprintOverride.withValue(fingerprint) {
try operation()
}
}
}
static func withClaudeKeychainOverridesForTesting<T>(
data: Data?,
fingerprint: ClaudeKeychainFingerprint?,
operation: () async throws -> T) async rethrows -> T
{
try await self.$taskClaudeKeychainDataOverride.withValue(data) {
try await self.$taskClaudeKeychainFingerprintOverride.withValue(fingerprint) {
try await operation()
}
}
}
static func withClaudeKeychainFingerprintStoreOverrideForTesting<T>(
_ store: ClaudeKeychainFingerprintStore?,
operation: () throws -> T) rethrows -> T
{
try self.$taskClaudeKeychainFingerprintStoreOverride.withValue(store) {
try operation()
}
}
static func withClaudeKeychainFingerprintStoreOverrideForTesting<T>(
_ store: ClaudeKeychainFingerprintStore?,
operation: () async throws -> T) async rethrows -> T
{
try await self.$taskClaudeKeychainFingerprintStoreOverride.withValue(store) {
try await operation()
}
}
static func withIsolatedMemoryCacheForTesting<T>(operation: () throws -> T) rethrows -> T {
let store = MemoryCacheStore()
let preAlertStore = ClaudeOAuthKeychainPreAlertGate.StateStore()
return try ClaudeOAuthKeychainPreAlertGate.withStateStoreOverrideForTesting(preAlertStore) {
try self.$taskMemoryCacheStoreOverride.withValue(store) {
try operation()
}
}
}
static func withIsolatedMemoryCacheForTesting<T>(operation: () async throws -> T) async rethrows -> T {
let store = MemoryCacheStore()
let preAlertStore = ClaudeOAuthKeychainPreAlertGate.StateStore()
return try await ClaudeOAuthKeychainPreAlertGate.withStateStoreOverrideForTesting(preAlertStore) {
try await self.$taskMemoryCacheStoreOverride.withValue(store) {
try await operation()
}
}
}
final class CredentialsFileFingerprintStore: @unchecked Sendable {
var fingerprint: CredentialsFileFingerprint?
init(fingerprint: CredentialsFileFingerprint? = nil) {
self.fingerprint = fingerprint
}
func load() -> CredentialsFileFingerprint? {
self.fingerprint
}
func save(_ fingerprint: CredentialsFileFingerprint?) {
self.fingerprint = fingerprint
}
}
enum SecurityCLIReadOverride {
case data(Data?)
case timedOut
case nonZeroExit
case dynamic(@Sendable (SecurityCLIReadRequest) -> Data?)
}
@TaskLocal static var taskKeychainAccessOverride: Bool?
@TaskLocal static var taskCredentialsFileFingerprintStoreOverride: CredentialsFileFingerprintStore?
@TaskLocal static var taskSecurityCLIReadOverride: SecurityCLIReadOverride?
@TaskLocal static var taskSecurityCLIReadAccountOverride: String?
nonisolated(unsafe) static var securityCLIReadOverride: SecurityCLIReadOverride?
public struct TestingOverridesSnapshot: Sendable {
let keychainOverrideStore: ClaudeKeychainOverrideStore?
let keychainData: Data?
let keychainFingerprint: ClaudeKeychainFingerprint?
let memoryCacheStore: MemoryCacheStore?
let fingerprintStore: ClaudeKeychainFingerprintStore?
let keychainAccessOverride: Bool?
let credentialsFileFingerprintStore: CredentialsFileFingerprintStore?
let securityCLIReadOverride: SecurityCLIReadOverride?
let securityCLIReadAccountOverride: String?
let pendingCacheClearStore: ClaudeOAuthPendingCacheClearStore?
let oauthCacheOperationRecorder: OAuthCacheOperationRecorder?
init(
keychainOverrideStore: ClaudeKeychainOverrideStore?,
keychainData: Data?,
keychainFingerprint: ClaudeKeychainFingerprint?,
memoryCacheStore: MemoryCacheStore?,
fingerprintStore: ClaudeKeychainFingerprintStore?,
keychainAccessOverride: Bool?,
credentialsFileFingerprintStore: CredentialsFileFingerprintStore?,
securityCLIReadOverride: SecurityCLIReadOverride?,
securityCLIReadAccountOverride: String?,
pendingCacheClearStore: ClaudeOAuthPendingCacheClearStore?,
oauthCacheOperationRecorder: OAuthCacheOperationRecorder?)
{
self.keychainOverrideStore = keychainOverrideStore
self.keychainData = keychainData
self.keychainFingerprint = keychainFingerprint
self.memoryCacheStore = memoryCacheStore
self.fingerprintStore = fingerprintStore
self.keychainAccessOverride = keychainAccessOverride
self.credentialsFileFingerprintStore = credentialsFileFingerprintStore
self.securityCLIReadOverride = securityCLIReadOverride
self.securityCLIReadAccountOverride = securityCLIReadAccountOverride
self.pendingCacheClearStore = pendingCacheClearStore
self.oauthCacheOperationRecorder = oauthCacheOperationRecorder
}
}
static func withKeychainAccessOverrideForTesting<T>(
_ disabled: Bool?,
operation: () throws -> T) rethrows -> T
{
try self.$taskKeychainAccessOverride.withValue(disabled) {
try operation()
}
}
static func withKeychainAccessOverrideForTesting<T>(
_ disabled: Bool?,
operation: () async throws -> T) async rethrows -> T
{
try await self.$taskKeychainAccessOverride.withValue(disabled) {
try await operation()
}
}
fileprivate static func withCredentialsFileFingerprintStoreOverrideForTesting<T>(
_ store: CredentialsFileFingerprintStore?,
operation: () throws -> T) rethrows -> T
{
try self.$taskCredentialsFileFingerprintStoreOverride.withValue(store) {
try operation()
}
}
fileprivate static func withCredentialsFileFingerprintStoreOverrideForTesting<T>(
_ store: CredentialsFileFingerprintStore?,
operation: () async throws -> T) async rethrows -> T
{
try await self.$taskCredentialsFileFingerprintStoreOverride.withValue(store) {
try await operation()
}
}
static func withIsolatedCredentialsFileTrackingForTesting<T>(
operation: () throws -> T) rethrows -> T
{
let store = CredentialsFileFingerprintStore()
return try self.$taskCredentialsFileFingerprintStoreOverride.withValue(store) {
try operation()
}
}
static func withPendingCacheClearStoreOverrideForTesting<T>(
_ store: ClaudeOAuthPendingCacheClearStore?,
operation: () throws -> T) rethrows -> T
{
try self.$taskPendingCacheClearStoreOverride.withValue(store) {
try operation()
}
}
static func withPendingCacheClearStoreOverrideForTesting<T>(
_ store: ClaudeOAuthPendingCacheClearStore?,
operation: () async throws -> T) async rethrows -> T
{
try await self.$taskPendingCacheClearStoreOverride.withValue(store) {
try await operation()
}
}
static func withOAuthCacheOperationRecorderForTesting<T>(
_ recorder: OAuthCacheOperationRecorder?,
operation: () throws -> T) rethrows -> T
{
try KeychainCacheStore.withOperationRecorderForTesting(recorder, operation: operation)
}
static func withOAuthCacheOperationRecorderForTesting<T>(
_ recorder: OAuthCacheOperationRecorder?,
operation: () async throws -> T) async rethrows -> T
{
try await KeychainCacheStore.withOperationRecorderForTesting(recorder, operation: operation)
}
static func withIsolatedCredentialsFileTrackingForTesting<T>(
operation: () async throws -> T) async rethrows -> T
{
let store = CredentialsFileFingerprintStore()
return try await self.$taskCredentialsFileFingerprintStoreOverride.withValue(store) {
try await operation()
}
}
static func withSecurityCLIReadOverrideForTesting<T>(
_ readOverride: SecurityCLIReadOverride?,
operation: () throws -> T) rethrows -> T
{
try self.$taskSecurityCLIReadOverride.withValue(readOverride) {
try operation()
}
}
static func withSecurityCLIReadOverrideForTesting<T>(
_ readOverride: SecurityCLIReadOverride?,
operation: () async throws -> T) async rethrows -> T
{
try await self.$taskSecurityCLIReadOverride.withValue(readOverride) {
try await operation()
}
}
static func currentSecurityCLIReadOverrideForTesting() -> SecurityCLIReadOverride? {
self.taskSecurityCLIReadOverride ?? self.securityCLIReadOverride
}
static func withSecurityCLIReadAccountOverrideForTesting<T>(
_ account: String?,
operation: () throws -> T) rethrows -> T
{
try self.$taskSecurityCLIReadAccountOverride.withValue(account) {
try operation()
}
}
static func withSecurityCLIReadAccountOverrideForTesting<T>(
_ account: String?,
operation: () async throws -> T) async rethrows -> T
{
try await self.$taskSecurityCLIReadAccountOverride.withValue(account) {
try await operation()
}
}
public static func withCurrentTestingOverridesForTask<T>(
operation: () async throws -> T) async rethrows -> T
{
try await self.withTestingOverridesSnapshotForTask(
self.currentTestingOverridesSnapshotForTask,
operation: operation)
}
public static func withCurrentTestingOverridesForTask<T>(
operation: () throws -> T) rethrows -> T
{
try self.withTestingOverridesSnapshotForTask(
self.currentTestingOverridesSnapshotForTask,
operation: operation)
}
public static var currentTestingOverridesSnapshotForTask: TestingOverridesSnapshot {
TestingOverridesSnapshot(
keychainOverrideStore: self.taskClaudeKeychainOverrideStore,
keychainData: self.taskClaudeKeychainDataOverride,
keychainFingerprint: self.taskClaudeKeychainFingerprintOverride,
memoryCacheStore: self.taskMemoryCacheStoreOverride,
fingerprintStore: self.taskClaudeKeychainFingerprintStoreOverride,
keychainAccessOverride: self.taskKeychainAccessOverride,
credentialsFileFingerprintStore: self.taskCredentialsFileFingerprintStoreOverride,
securityCLIReadOverride: self.taskSecurityCLIReadOverride,
securityCLIReadAccountOverride: self.taskSecurityCLIReadAccountOverride,
pendingCacheClearStore: self.taskPendingCacheClearStoreOverride,
oauthCacheOperationRecorder: KeychainCacheStore.currentOperationRecorderForTesting)
}
public static func withTestingOverridesSnapshotForTask<T>(
_ snapshot: TestingOverridesSnapshot,
operation: () async throws -> T) async rethrows -> T
{
try await self.$taskPendingCacheClearStoreOverride.withValue(snapshot.pendingCacheClearStore) {
try await KeychainCacheStore.withOperationRecorderForTesting(snapshot.oauthCacheOperationRecorder) {
try await self.$taskClaudeKeychainOverrideStore.withValue(snapshot.keychainOverrideStore) {
try await self.$taskClaudeKeychainDataOverride.withValue(snapshot.keychainData) {
try await self.$taskClaudeKeychainFingerprintOverride.withValue(snapshot.keychainFingerprint) {
try await self.$taskMemoryCacheStoreOverride.withValue(snapshot.memoryCacheStore) {
try await self.$taskClaudeKeychainFingerprintStoreOverride
.withValue(snapshot.fingerprintStore) {
try await self.$taskKeychainAccessOverride
.withValue(snapshot.keychainAccessOverride) {
try await self.$taskCredentialsFileFingerprintStoreOverride.withValue(
snapshot.credentialsFileFingerprintStore)
{
try await self.$taskSecurityCLIReadOverride.withValue(
snapshot.securityCLIReadOverride)
{
try await self.$taskSecurityCLIReadAccountOverride.withValue(
snapshot.securityCLIReadAccountOverride)
{
try await operation()
}
}
}
}
}
}
}
}
}
}
}
}
public static func withTestingOverridesSnapshotForTask<T>(
_ snapshot: TestingOverridesSnapshot,
operation: () throws -> T) rethrows -> T
{
try self.$taskPendingCacheClearStoreOverride.withValue(snapshot.pendingCacheClearStore) {
try KeychainCacheStore.withOperationRecorderForTesting(snapshot.oauthCacheOperationRecorder) {
try self.$taskClaudeKeychainOverrideStore.withValue(snapshot.keychainOverrideStore) {
try self.$taskClaudeKeychainDataOverride.withValue(snapshot.keychainData) {
try self.$taskClaudeKeychainFingerprintOverride.withValue(snapshot.keychainFingerprint) {
try self.$taskMemoryCacheStoreOverride.withValue(snapshot.memoryCacheStore) {
try self.$taskClaudeKeychainFingerprintStoreOverride
.withValue(snapshot.fingerprintStore) {
try self.$taskKeychainAccessOverride
.withValue(snapshot.keychainAccessOverride) {
try self.$taskCredentialsFileFingerprintStoreOverride.withValue(
snapshot.credentialsFileFingerprintStore)
{
try self.$taskSecurityCLIReadOverride.withValue(
snapshot.securityCLIReadOverride)
{
try self.$taskSecurityCLIReadAccountOverride.withValue(
snapshot.securityCLIReadAccountOverride)
{
try operation()
}
}
}
}
}
}
}
}
}
}
}
}
static func setSecurityCLIReadOverrideForTesting(_ readOverride: SecurityCLIReadOverride?) {
self.securityCLIReadOverride = readOverride
}
}
#endif
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,580 @@
import Foundation
public enum ClaudeOAuthDelegatedRefreshCoordinator {
private final class AttemptStateStorage: @unchecked Sendable {
let lock = NSLock()
let persistsCooldown: Bool
var hasLoadedState = false
var lastAttemptAt: Date?
var lastCooldownInterval: TimeInterval?
var inFlightAttemptID: UInt64?
var inFlightInteraction: ProviderInteraction?
var inFlightTask: Task<Outcome, Never>?
var nextAttemptID: UInt64 = 0
init(persistsCooldown: Bool) {
self.persistsCooldown = persistsCooldown
}
}
public enum Outcome: Sendable, Equatable {
case skippedByCooldown
case cliUnavailable
case attemptedSucceeded
case attemptedFailed(String)
}
private static let log = CodexBarLog.logger(LogCategories.claudeUsage)
private static let cooldownDefaultsKey = "claudeOAuthDelegatedRefreshLastAttemptAtV1"
private static let cooldownIntervalDefaultsKey = "claudeOAuthDelegatedRefreshCooldownIntervalSecondsV1"
private static let defaultCooldownInterval: TimeInterval = 60 * 5
private static let shortCooldownInterval: TimeInterval = 20
private static let sharedState = AttemptStateStorage(persistsCooldown: true)
public static func attempt(
now: Date = Date(),
timeout: TimeInterval = 8,
environment: [String: String] = ProcessInfo.processInfo.environment) async -> Outcome
{
if Task.isCancelled {
return .attemptedFailed("Cancelled.")
}
let decision = self.inFlightDecision(
now: now,
timeout: timeout,
environment: environment,
interaction: ProviderInteractionContext.current)
#if DEBUG
if case .joinThenRetry = decision {
self.userInitiatedBackgroundJoinObserverForTesting?()
}
#endif
switch decision {
case let .join(task):
return await task.value
case let .joinThenRetry(id, task, state):
let outcome = await task.value
self.clearInFlightTaskIfStillCurrent(id: id, state: state)
switch outcome {
case .attemptedFailed, .skippedByCooldown, .cliUnavailable:
return await self.attempt(now: now, timeout: timeout, environment: environment)
case .attemptedSucceeded:
return outcome
}
case let .start(id, task, state):
let outcome = await task.value
self.clearInFlightTaskIfStillCurrent(id: id, state: state)
return outcome
}
}
private enum InFlightDecision {
case join(Task<Outcome, Never>)
case joinThenRetry(UInt64, Task<Outcome, Never>, AttemptStateStorage)
case start(UInt64, Task<Outcome, Never>, AttemptStateStorage)
}
private struct AttemptConfiguration {
let environment: [String: String]
let interaction: ProviderInteraction
let readStrategy: ClaudeOAuthKeychainReadStrategy
let keychainAccessDisabled: Bool
#if DEBUG
let cliAvailableOverride: Bool?
let touchAuthPathOverride: (@Sendable (TimeInterval, [String: String]) async throws -> Void)?
let keychainFingerprintOverride: (@Sendable () -> ClaudeOAuthCredentialsStore.ClaudeKeychainFingerprint?)?
#endif
}
private static func inFlightDecision(
now: Date,
timeout: TimeInterval,
environment: [String: String],
interaction: ProviderInteraction) -> InFlightDecision
{
let state = self.currentStateStorage
state.lock.lock()
defer { state.lock.unlock() }
if let existing = state.inFlightTask {
if interaction == .userInitiated,
state.inFlightInteraction != .userInitiated,
let existingID = state.inFlightAttemptID
{
return .joinThenRetry(existingID, existing, state)
}
return .join(existing)
}
state.nextAttemptID += 1
let attemptID = state.nextAttemptID
// Detached to avoid inheriting the caller's executor context (e.g. MainActor) and cancellation state.
#if DEBUG
let configuration = AttemptConfiguration(
environment: environment,
interaction: interaction,
readStrategy: ClaudeOAuthKeychainReadStrategyPreference.current(),
keychainAccessDisabled: KeychainAccessGate.isDisabled,
cliAvailableOverride: self.cliAvailableOverrideForTesting,
touchAuthPathOverride: self.touchAuthPathOverrideForTesting,
keychainFingerprintOverride: self.keychainFingerprintOverrideForTesting)
let securityCLIReadOverride = ClaudeOAuthCredentialsStore.currentSecurityCLIReadOverrideForTesting()
#else
let configuration = AttemptConfiguration(
environment: environment,
interaction: interaction,
readStrategy: ClaudeOAuthKeychainReadStrategyPreference.current(),
keychainAccessDisabled: KeychainAccessGate.isDisabled)
#endif
let task = Task.detached(priority: .utility) {
#if DEBUG
return await ClaudeOAuthCredentialsStore.withSecurityCLIReadOverrideForTesting(securityCLIReadOverride) {
await self.performAttempt(
now: now,
timeout: timeout,
configuration: configuration,
state: state)
}
#else
await self.performAttempt(
now: now,
timeout: timeout,
configuration: configuration,
state: state)
#endif
}
state.inFlightAttemptID = attemptID
state.inFlightInteraction = interaction
state.inFlightTask = task
return .start(attemptID, task, state)
}
private static func performAttempt(
now: Date,
timeout: TimeInterval,
configuration: AttemptConfiguration,
state: AttemptStateStorage) async -> Outcome
{
guard self.isClaudeCLIAvailable(environment: configuration.environment, configuration: configuration) else {
self.log.info("Claude OAuth delegated refresh skipped: claude CLI unavailable")
return .cliUnavailable
}
// Atomically reserve an attempt under the lock so concurrent callers don't race past isInCooldown() and start
// multiple touches/poll loops.
guard self.reserveAttemptIfNotInCooldown(
now: now,
bypassCooldown: configuration.interaction == .userInitiated,
state: state)
else {
self.log.debug("Claude OAuth delegated refresh skipped by cooldown")
return .skippedByCooldown
}
if let mcpOAuthOnlyFailure = self.mcpOAuthOnlyKeychainFailureIfPresent(
interaction: configuration.interaction,
readStrategy: configuration.readStrategy,
keychainAccessDisabled: configuration.keychainAccessDisabled,
environment: configuration.environment)
{
self.recordAttempt(now: now, cooldown: self.defaultCooldownInterval, state: state)
self.log.warning(
"Claude OAuth delegated refresh skipped: Claude keychain has MCP OAuth state only",
metadata: ["readStrategy": configuration.readStrategy.rawValue])
return .attemptedFailed(mcpOAuthOnlyFailure)
}
let baseline = self.currentKeychainChangeObservationBaseline(
readStrategy: configuration.readStrategy,
keychainAccessDisabled: configuration.keychainAccessDisabled,
configuration: configuration)
var touchError: Error?
do {
try await self.touchOAuthAuthPath(
timeout: timeout,
environment: configuration.environment,
configuration: configuration)
} catch {
touchError = error
}
// "Touch succeeded" must mean we actually observed the Claude keychain entry change.
// Otherwise we end up in a long cooldown with still-expired credentials.
let changed = await self.waitForClaudeKeychainChange(
from: baseline,
readStrategy: configuration.readStrategy,
keychainAccessDisabled: configuration.keychainAccessDisabled,
configuration: configuration,
timeout: min(max(timeout, 1), 2))
if changed {
self.recordAttempt(now: now, cooldown: self.defaultCooldownInterval, state: state)
self.log.info("Claude OAuth delegated refresh touch succeeded")
return .attemptedSucceeded
}
self.recordAttempt(now: now, cooldown: self.shortCooldownInterval, state: state)
if let touchError {
let errorType = String(describing: type(of: touchError))
self.log.warning(
"Claude OAuth delegated refresh touch failed",
metadata: ["errorType": errorType])
self.log.debug("Claude OAuth delegated refresh touch error: \(touchError.localizedDescription)")
return .attemptedFailed(touchError.localizedDescription)
}
self.log.warning("Claude OAuth delegated refresh touch did not update Claude keychain")
return .attemptedFailed("Claude keychain did not update after Claude CLI touch.")
}
public static func isInCooldown(now: Date = Date()) -> Bool {
let state = self.currentStateStorage
state.lock.lock()
defer { state.lock.unlock() }
self.loadStateIfNeededLocked(state: state)
guard let lastAttemptAt = state.lastAttemptAt else { return false }
let cooldown = state.lastCooldownInterval ?? self.defaultCooldownInterval
return now.timeIntervalSince(lastAttemptAt) < cooldown
}
public static func cooldownRemainingSeconds(now: Date = Date()) -> Int? {
let state = self.currentStateStorage
state.lock.lock()
defer { state.lock.unlock() }
self.loadStateIfNeededLocked(state: state)
guard let lastAttemptAt = state.lastAttemptAt else { return nil }
let cooldown = state.lastCooldownInterval ?? self.defaultCooldownInterval
let remaining = cooldown - now.timeIntervalSince(lastAttemptAt)
guard remaining > 0 else { return nil }
return Int(remaining.rounded(.up))
}
public static func isClaudeCLIAvailable(
environment: [String: String] = ProcessInfo.processInfo.environment) -> Bool
{
self.isClaudeCLIAvailable(
environment: environment,
configuration: nil)
}
private static func isClaudeCLIAvailable(
environment: [String: String],
configuration: AttemptConfiguration?) -> Bool
{
#if DEBUG
if let override = configuration?.cliAvailableOverride ?? self.cliAvailableOverrideForTesting {
return override
}
#endif
return ClaudeCLIResolver.isAvailable(environment: environment)
}
private static func touchOAuthAuthPath(
timeout: TimeInterval,
environment: [String: String],
configuration: AttemptConfiguration?) async throws
{
#if DEBUG
if let override = configuration?.touchAuthPathOverride ?? self.touchAuthPathOverrideForTesting {
try await override(timeout, environment)
return
}
#endif
try await ClaudeStatusProbe.touchOAuthAuthPath(timeout: timeout, environment: environment)
}
private enum KeychainChangeObservationBaseline {
case securityFramework(fingerprint: ClaudeOAuthCredentialsStore.ClaudeKeychainFingerprint?)
case securityCLI(data: Data?)
}
private static func currentKeychainChangeObservationBaseline(
readStrategy: ClaudeOAuthKeychainReadStrategy,
keychainAccessDisabled: Bool,
configuration: AttemptConfiguration?) -> KeychainChangeObservationBaseline
{
if readStrategy == .securityCLIExperimental {
return .securityCLI(data: self.currentClaudeKeychainDataViaSecurityCLIForObservation(
readStrategy: readStrategy,
keychainAccessDisabled: keychainAccessDisabled,
interaction: .background))
}
return .securityFramework(fingerprint: self.currentClaudeKeychainFingerprint(configuration: configuration))
}
private static func waitForClaudeKeychainChange(
from baseline: KeychainChangeObservationBaseline,
readStrategy: ClaudeOAuthKeychainReadStrategy,
keychainAccessDisabled: Bool,
configuration: AttemptConfiguration?,
timeout: TimeInterval) async -> Bool
{
// Prefer correctness but bound the delay. Keychain writes can be slightly delayed after the CLI touch.
// Keep this short to avoid "prompt storms" on configurations where "no UI" queries can still surface UI.
let clampedTimeout = max(0, min(timeout, 2))
if clampedTimeout == 0 { return false }
let delays: [TimeInterval] = [0.2, 0.5, 0.8].filter { $0 <= clampedTimeout }
let deadline = Date().addingTimeInterval(clampedTimeout)
func isObservedChange() -> Bool {
switch baseline {
case let .securityFramework(fingerprintBefore):
// Treat "no fingerprint" as "not observed"; we only succeed if we can read a fingerprint and it
// differs.
guard let current = self.currentClaudeKeychainFingerprintForObservation(configuration: configuration)
else {
return false
}
return current != fingerprintBefore
case let .securityCLI(dataBefore):
// In experimental mode, avoid Security.framework observation entirely and detect change from
// /usr/bin/security output only.
// If baseline capture failed (nil), treat observation as inconclusive and do not infer a change from
// a later successful read.
guard let dataBefore else { return false }
guard let current = self.currentClaudeKeychainDataViaSecurityCLIForObservation(
readStrategy: readStrategy,
keychainAccessDisabled: keychainAccessDisabled,
interaction: .background)
else { return false }
return current != dataBefore
}
}
if isObservedChange() {
return true
}
for delay in delays {
if Date() >= deadline { break }
do {
try Task.checkCancellation()
try await Task.sleep(nanoseconds: UInt64(delay * 1_000_000_000))
} catch {
return false
}
if isObservedChange() {
return true
}
}
return false
}
private static func currentClaudeKeychainFingerprint(
configuration: AttemptConfiguration?) -> ClaudeOAuthCredentialsStore.ClaudeKeychainFingerprint?
{
#if DEBUG
if let override = configuration?.keychainFingerprintOverride ?? self.keychainFingerprintOverrideForTesting {
return override()
}
#endif
return ClaudeOAuthCredentialsStore.currentClaudeKeychainFingerprintWithoutPromptForAuthGate()
}
private static func currentClaudeKeychainFingerprintForObservation() -> ClaudeOAuthCredentialsStore
.ClaudeKeychainFingerprint?
{
self.currentClaudeKeychainFingerprintForObservation(configuration: nil)
}
private static func currentClaudeKeychainFingerprintForObservation(
configuration: AttemptConfiguration?) -> ClaudeOAuthCredentialsStore.ClaudeKeychainFingerprint?
{
#if DEBUG
if let override = configuration?.keychainFingerprintOverride ?? self.keychainFingerprintOverrideForTesting {
return override()
}
#endif
// Observation should not be blocked by the background cooldown gate; otherwise we can "false fail" even when
// the CLI refreshed successfully but we couldn't observe it due to a previous denied prompt/cooldown.
//
// This temporarily classifies the observation query as "user initiated" so it bypasses the gate that only
// applies to background probes. The query remains "no UI" and does not clear cooldown state itself.
return ProviderInteractionContext.$current.withValue(.userInitiated) {
ClaudeOAuthCredentialsStore.currentClaudeKeychainFingerprintWithoutPromptForAuthGate()
}
}
private static func currentClaudeKeychainDataViaSecurityCLIForObservation(
readStrategy: ClaudeOAuthKeychainReadStrategy,
keychainAccessDisabled: Bool,
interaction: ProviderInteraction) -> Data?
{
guard !keychainAccessDisabled else { return nil }
return ClaudeOAuthCredentialsStore.readRawClaudeKeychainPayloadViaSecurityCLIIfEnabled(
interaction: interaction,
readStrategy: readStrategy)
}
private static func mcpOAuthOnlyKeychainFailureIfPresent(
interaction: ProviderInteraction,
readStrategy: ClaudeOAuthKeychainReadStrategy,
keychainAccessDisabled: Bool,
environment: [String: String]) -> String?
{
guard interaction != .userInitiated else { return nil }
guard ClaudeOAuthCredentialsStore.isMcpOAuthOnlyClaudeKeychainPayloadPresent(
interaction: interaction,
readStrategy: readStrategy,
keychainAccessDisabled: keychainAccessDisabled,
environment: environment)
else {
return nil
}
return ClaudeOAuthCredentialsError.mcpOAuthOnlyKeychain.errorDescription
?? "Claude keychain contains MCP OAuth state only."
}
private static func clearInFlightTaskIfStillCurrent(id: UInt64, state: AttemptStateStorage) {
state.lock.lock()
if state.inFlightAttemptID == id {
state.inFlightAttemptID = nil
state.inFlightInteraction = nil
state.inFlightTask = nil
}
state.lock.unlock()
}
private static func recordAttempt(now: Date, cooldown: TimeInterval, state: AttemptStateStorage) {
state.lock.lock()
defer { state.lock.unlock() }
self.loadStateIfNeededLocked(state: state)
state.lastAttemptAt = now
state.lastCooldownInterval = cooldown
guard state.persistsCooldown else { return }
UserDefaults.standard.set(now.timeIntervalSince1970, forKey: self.cooldownDefaultsKey)
UserDefaults.standard.set(cooldown, forKey: self.cooldownIntervalDefaultsKey)
}
private static func reserveAttemptIfNotInCooldown(
now: Date,
bypassCooldown: Bool,
state: AttemptStateStorage) -> Bool
{
state.lock.lock()
defer { state.lock.unlock() }
self.loadStateIfNeededLocked(state: state)
let cooldown = state.lastCooldownInterval ?? self.defaultCooldownInterval
if !bypassCooldown,
let lastAttemptAt = state.lastAttemptAt,
now.timeIntervalSince(lastAttemptAt) < cooldown
{
return false
}
// Reserve with a short cooldown; the final outcome will extend or keep it short.
state.lastAttemptAt = now
state.lastCooldownInterval = self.shortCooldownInterval
guard state.persistsCooldown else { return true }
UserDefaults.standard.set(now.timeIntervalSince1970, forKey: self.cooldownDefaultsKey)
UserDefaults.standard.set(self.shortCooldownInterval, forKey: self.cooldownIntervalDefaultsKey)
return true
}
private static func loadStateIfNeededLocked(state: AttemptStateStorage) {
guard !state.hasLoadedState else { return }
state.hasLoadedState = true
guard state.persistsCooldown else {
state.lastAttemptAt = nil
state.lastCooldownInterval = nil
return
}
guard let raw = UserDefaults.standard.object(forKey: self.cooldownDefaultsKey) as? Double else {
state.lastAttemptAt = nil
state.lastCooldownInterval = nil
return
}
state.lastAttemptAt = Date(timeIntervalSince1970: raw)
if let interval = UserDefaults.standard.object(forKey: self.cooldownIntervalDefaultsKey) as? Double {
state.lastCooldownInterval = interval
} else {
state.lastCooldownInterval = nil
}
}
#if DEBUG
@TaskLocal private static var stateStorageForTesting: AttemptStateStorage?
@TaskLocal static var cliAvailableOverrideForTesting: Bool?
@TaskLocal static var touchAuthPathOverrideForTesting: (@Sendable (
TimeInterval,
[String: String]) async throws -> Void)?
@TaskLocal static var keychainFingerprintOverrideForTesting: (@Sendable () -> ClaudeOAuthCredentialsStore
.ClaudeKeychainFingerprint?)?
@TaskLocal static var userInitiatedBackgroundJoinObserverForTesting: (@Sendable () -> Void)?
static func withCLIAvailableOverrideForTesting<T>(
_ override: Bool?,
operation: () async throws -> T) async rethrows -> T
{
try await self.$cliAvailableOverrideForTesting.withValue(override) {
try await operation()
}
}
static func withTouchAuthPathOverrideForTesting<T>(
_ override: (@Sendable (TimeInterval, [String: String]) async throws -> Void)?,
operation: () async throws -> T) async rethrows -> T
{
try await self.$touchAuthPathOverrideForTesting.withValue(override) {
try await operation()
}
}
static func withKeychainFingerprintOverrideForTesting<T>(
_ override: (@Sendable () -> ClaudeOAuthCredentialsStore.ClaudeKeychainFingerprint?)?,
operation: () async throws -> T) async rethrows -> T
{
try await self.$keychainFingerprintOverrideForTesting.withValue(override) {
try await operation()
}
}
static func withUserInitiatedBackgroundJoinObserverForTesting<T>(
_ observer: (@Sendable () -> Void)?,
operation: () async throws -> T) async rethrows -> T
{
try await self.$userInitiatedBackgroundJoinObserverForTesting.withValue(observer) {
try await operation()
}
}
static func withIsolatedStateForTesting<T>(operation: () async throws -> T) async rethrows -> T {
let state = AttemptStateStorage(persistsCooldown: false)
return try await self.$stateStorageForTesting.withValue(state) {
try await operation()
}
}
static func resetForTesting() {
let state = self.currentStateStorage
state.lock.lock()
state.hasLoadedState = true
state.lastAttemptAt = nil
state.lastCooldownInterval = nil
state.inFlightAttemptID = nil
state.inFlightInteraction = nil
state.inFlightTask = nil
state.nextAttemptID = 0
state.lock.unlock()
guard state.persistsCooldown else { return }
UserDefaults.standard.removeObject(forKey: self.cooldownDefaultsKey)
UserDefaults.standard.removeObject(forKey: self.cooldownIntervalDefaultsKey)
}
#endif
private static var currentStateStorage: AttemptStateStorage {
#if DEBUG
self.stateStorageForTesting ?? self.sharedState
#else
self.sharedState
#endif
}
}
@@ -0,0 +1,184 @@
import Foundation
#if os(macOS)
import os.lock
public enum ClaudeOAuthKeychainAccessGate {
private struct State {
var loaded = false
var deniedUntil: Date?
}
private static let lock = OSAllocatedUnfairLock<State>(initialState: State())
private static let defaultsKey = "claudeOAuthKeychainDeniedUntil"
private static let cooldownInterval: TimeInterval = 60 * 60 * 6
@TaskLocal private static var taskOverrideShouldAllowPromptForTesting: Bool?
#if DEBUG
public final class DeniedUntilStore: @unchecked Sendable {
public var deniedUntil: Date?
public init() {}
}
@TaskLocal private static var taskDeniedUntilStoreOverrideForTesting: DeniedUntilStore?
#endif
public static func shouldAllowPrompt(now: Date = Date()) -> Bool {
guard !KeychainAccessGate.isDisabled else { return false }
if let override = self.taskOverrideShouldAllowPromptForTesting { return override }
#if DEBUG
if let store = self.taskDeniedUntilStoreOverrideForTesting {
if let deniedUntil = store.deniedUntil, deniedUntil > now {
return false
}
store.deniedUntil = nil
return true
}
#endif
return self.lock.withLock { state in
self.loadIfNeeded(&state)
if let deniedUntil = state.deniedUntil {
if deniedUntil > now {
return false
}
state.deniedUntil = nil
self.persist(state)
}
return true
}
}
public static func recordDenied(now: Date = Date()) {
let deniedUntil = now.addingTimeInterval(self.cooldownInterval)
#if DEBUG
if let store = self.taskDeniedUntilStoreOverrideForTesting {
store.deniedUntil = deniedUntil
return
}
#endif
self.lock.withLock { state in
self.loadIfNeeded(&state)
state.deniedUntil = deniedUntil
self.persist(state)
}
}
/// Clears the cooldown so the next attempt can proceed. Intended for user-initiated repairs.
/// - Returns: true if a cooldown was present and cleared.
public static func clearDenied(now: Date = Date()) -> Bool {
#if DEBUG
if let store = self.taskDeniedUntilStoreOverrideForTesting {
guard let deniedUntil = store.deniedUntil, deniedUntil > now else {
store.deniedUntil = nil
return false
}
store.deniedUntil = nil
return true
}
#endif
return self.lock.withLock { state in
self.loadIfNeeded(&state)
guard let deniedUntil = state.deniedUntil, deniedUntil > now else {
state.deniedUntil = nil
self.persist(state)
return false
}
state.deniedUntil = nil
self.persist(state)
return true
}
}
#if DEBUG
static func withShouldAllowPromptOverrideForTesting<T>(
_ value: Bool?,
operation: () throws -> T) rethrows -> T
{
try self.$taskOverrideShouldAllowPromptForTesting.withValue(value) {
try operation()
}
}
static func withShouldAllowPromptOverrideForTesting<T>(
_ value: Bool?,
operation: () async throws -> T) async rethrows -> T
{
try await self.$taskOverrideShouldAllowPromptForTesting.withValue(value) {
try await operation()
}
}
public static func withDeniedUntilStoreOverrideForTesting<T>(
_ store: DeniedUntilStore?,
operation: () throws -> T) rethrows -> T
{
try self.$taskDeniedUntilStoreOverrideForTesting.withValue(store) {
try operation()
}
}
public static func withDeniedUntilStoreOverrideForTesting<T>(
_ store: DeniedUntilStore?,
operation: () async throws -> T) async rethrows -> T
{
try await self.$taskDeniedUntilStoreOverrideForTesting.withValue(store) {
try await operation()
}
}
public static var currentDeniedUntilStoreOverrideForTesting: DeniedUntilStore? {
self.taskDeniedUntilStoreOverrideForTesting
}
public static func resetForTesting() {
self.lock.withLock { state in
// Keep deterministic during tests: avoid re-loading UserDefaults written by unrelated code paths.
state.loaded = true
state.deniedUntil = nil
UserDefaults.standard.removeObject(forKey: self.defaultsKey)
}
}
public static func resetInMemoryForTesting() {
self.lock.withLock { state in
state.loaded = false
state.deniedUntil = nil
}
}
#endif
private static func loadIfNeeded(_ state: inout State) {
guard !state.loaded else { return }
state.loaded = true
if let raw = UserDefaults.standard.object(forKey: self.defaultsKey) as? Double {
state.deniedUntil = Date(timeIntervalSince1970: raw)
}
}
private static func persist(_ state: State) {
if let deniedUntil = state.deniedUntil {
UserDefaults.standard.set(deniedUntil.timeIntervalSince1970, forKey: self.defaultsKey)
} else {
UserDefaults.standard.removeObject(forKey: self.defaultsKey)
}
}
}
#else
public enum ClaudeOAuthKeychainAccessGate {
public static func shouldAllowPrompt(now _: Date = Date()) -> Bool {
true
}
public static func recordDenied(now _: Date = Date()) {}
public static func clearDenied(now _: Date = Date()) -> Bool {
false
}
#if DEBUG
public static func resetForTesting() {}
public static func resetInMemoryForTesting() {}
#endif
}
#endif
@@ -0,0 +1,169 @@
import Foundation
#if os(macOS)
import os.lock
enum ClaudeOAuthKeychainPreAlertGate {
fileprivate struct State {
var loaded = false
var acknowledgedUntil: Date?
var presentationInFlight = false
}
private static let lock = OSAllocatedUnfairLock<State>(initialState: State())
private static let defaultsKey = "claudeOAuthKeychainPreAlertAcknowledgedUntilV1"
static let cooldownInterval: TimeInterval = 60 * 60 * 6
#if DEBUG
final class StateStore: @unchecked Sendable {
fileprivate let lock = OSAllocatedUnfairLock<State>(initialState: State(loaded: true))
}
@TaskLocal private static var taskStateStoreOverrideForTesting: StateStore?
#endif
/// Presents at most one explanatory alert and starts the cooldown only when it reaches a handler.
@discardableResult
static func presentIfNeeded(
now: Date = Date(),
completedAt: Date? = nil,
present: () -> Bool) -> Bool
{
guard self.beginPresentation(now: now) else { return false }
let wasPresented = present()
self.finishPresentation(wasPresented: wasPresented, now: completedAt ?? Date())
return wasPresented
}
private static func beginPresentation(now: Date) -> Bool {
#if DEBUG
if let store = self.taskStateStoreOverrideForTesting {
return store.lock.withLock { state in
self.reservePresentation(state: &state, now: now)
}
}
#endif
return self.lock.withLock { state in
self.loadIfNeeded(&state)
guard self.reservePresentation(state: &state, now: now) else { return false }
self.persist(state)
return true
}
}
private static func finishPresentation(wasPresented: Bool, now: Date) {
#if DEBUG
if let store = self.taskStateStoreOverrideForTesting {
store.lock.withLock { state in
self.completePresentation(state: &state, wasPresented: wasPresented, now: now)
}
return
}
#endif
self.lock.withLock { state in
self.loadIfNeeded(&state)
self.completePresentation(state: &state, wasPresented: wasPresented, now: now)
self.persist(state)
}
}
#if DEBUG
static func withStateStoreOverrideForTesting<T>(
_ store: StateStore?,
operation: () throws -> T) rethrows -> T
{
try self.$taskStateStoreOverrideForTesting.withValue(store) {
try operation()
}
}
static func withStateStoreOverrideForTesting<T>(
_ store: StateStore?,
operation: () async throws -> T) async rethrows -> T
{
try await self.$taskStateStoreOverrideForTesting.withValue(store) {
try await operation()
}
}
static func resetForTesting() {
self.lock.withLock { state in
state = State(loaded: true)
UserDefaults.standard.removeObject(forKey: self.defaultsKey)
}
}
static func resetInMemoryForTesting() {
self.lock.withLock { state in
state = State()
}
}
#endif
private static func loadIfNeeded(_ state: inout State) {
guard !state.loaded else { return }
state.loaded = true
if let raw = UserDefaults.standard.object(forKey: self.defaultsKey) as? Double {
state.acknowledgedUntil = Date(timeIntervalSince1970: raw)
}
}
private static func reservePresentation(state: inout State, now: Date) -> Bool {
guard !state.presentationInFlight else { return false }
if let acknowledgedUntil = state.acknowledgedUntil, acknowledgedUntil > now {
return false
}
state.acknowledgedUntil = nil
state.presentationInFlight = true
return true
}
private static func completePresentation(state: inout State, wasPresented: Bool, now: Date) {
state.presentationInFlight = false
if wasPresented {
state.acknowledgedUntil = now.addingTimeInterval(self.cooldownInterval)
}
}
private static func persist(_ state: State) {
if let acknowledgedUntil = state.acknowledgedUntil {
UserDefaults.standard.set(acknowledgedUntil.timeIntervalSince1970, forKey: self.defaultsKey)
} else {
UserDefaults.standard.removeObject(forKey: self.defaultsKey)
}
}
}
#else
enum ClaudeOAuthKeychainPreAlertGate {
static let cooldownInterval: TimeInterval = 60 * 60 * 6
#if DEBUG
final class StateStore: @unchecked Sendable {}
#endif
@discardableResult
static func presentIfNeeded(
now _: Date = Date(),
completedAt _: Date? = nil,
present _: () -> Bool) -> Bool
{
false
}
#if DEBUG
static func withStateStoreOverrideForTesting<T>(
_: StateStore?,
operation: () throws -> T) rethrows -> T
{
try operation()
}
static func withStateStoreOverrideForTesting<T>(
_: StateStore?,
operation: () async throws -> T) async rethrows -> T
{
try await operation()
}
#endif
}
#endif
@@ -0,0 +1,177 @@
import Foundation
public enum ClaudeOAuthKeychainPromptMode: String, Sendable, Codable, CaseIterable {
case never
case onlyOnUserAction
case always
}
public enum ClaudeOAuthKeychainPromptPreference {
static let releaseApplicationDefaultsDomain = "com.steipete.codexbar"
static let debugApplicationDefaultsDomain = "com.steipete.codexbar.debug"
private static let userDefaultsKey = "claudeOAuthKeychainPromptMode"
static var applicationDefaultsDomain: String {
self.resolveApplicationDefaultsDomain(
bundleIdentifier: Bundle.main.bundleIdentifier,
bundleURL: Bundle.main.bundleURL,
executableURL: Bundle.main.executableURL,
invocationURL: CommandLine.arguments.first.map(URL.init(fileURLWithPath:)))
}
#if DEBUG
private final class UserDefaultsBox: @unchecked Sendable {
let value: UserDefaults
init(_ value: UserDefaults) {
self.value = value
}
}
@TaskLocal private static var taskOverride: ClaudeOAuthKeychainPromptMode?
@TaskLocal private static var taskApplicationUserDefaultsOverride: UserDefaultsBox?
#endif
public static func current(userDefaults: UserDefaults? = nil) -> ClaudeOAuthKeychainPromptMode {
self.effectiveMode(userDefaults: userDefaults)
}
public static func storedMode(userDefaults: UserDefaults? = nil) -> ClaudeOAuthKeychainPromptMode {
#if DEBUG
if let taskOverride {
return taskOverride
}
#endif
let userDefaults = userDefaults ?? self.applicationUserDefaults
if let raw = userDefaults.string(forKey: self.userDefaultsKey),
let mode = ClaudeOAuthKeychainPromptMode(rawValue: raw)
{
return mode
}
return .onlyOnUserAction
}
public static func isApplicable(
readStrategy: ClaudeOAuthKeychainReadStrategy = ClaudeOAuthKeychainReadStrategyPreference.current()) -> Bool
{
readStrategy == .securityFramework
}
public static func effectiveMode(
userDefaults: UserDefaults? = nil,
readStrategy: ClaudeOAuthKeychainReadStrategy = ClaudeOAuthKeychainReadStrategyPreference.current())
-> ClaudeOAuthKeychainPromptMode
{
guard self.isApplicable(readStrategy: readStrategy) else {
return .always
}
return self.storedMode(userDefaults: userDefaults)
}
public static func securityFrameworkFallbackMode(
userDefaults: UserDefaults? = nil,
readStrategy: ClaudeOAuthKeychainReadStrategy = ClaudeOAuthKeychainReadStrategyPreference.current())
-> ClaudeOAuthKeychainPromptMode
{
if readStrategy == .securityCLIExperimental {
return self.storedMode(userDefaults: userDefaults)
}
return self.effectiveMode(userDefaults: userDefaults, readStrategy: readStrategy)
}
static var applicationUserDefaults: UserDefaults {
#if DEBUG
if let taskApplicationUserDefaultsOverride {
return taskApplicationUserDefaultsOverride.value
}
#endif
return UserDefaults(suiteName: self.applicationDefaultsDomain) ?? .standard
}
static func resolveApplicationDefaultsDomain(
bundleIdentifier: String?,
bundleURL: URL?,
executableURL: URL?,
invocationURL: URL?,
bundleIdentifierForApp: (URL) -> String? = { Bundle(url: $0)?.bundleIdentifier }) -> String
{
if let domain = self.defaultsDomain(forBundleIdentifier: bundleIdentifier) {
return domain
}
let candidates = [bundleURL, executableURL, invocationURL].compactMap(\.self)
var visitedPaths = Set<String>()
for candidate in candidates {
var current = candidate.standardizedFileURL.resolvingSymlinksInPath()
let ancestorCount = current.pathComponents.count
for _ in 0..<ancestorCount {
if current.pathExtension == "app",
visitedPaths.insert(current.path).inserted,
let domain = self.defaultsDomain(forBundleIdentifier: bundleIdentifierForApp(current))
{
return domain
}
current.deleteLastPathComponent()
}
}
return self.releaseApplicationDefaultsDomain
}
private static func defaultsDomain(forBundleIdentifier bundleIdentifier: String?) -> String? {
guard let bundleIdentifier else { return nil }
// Check debug first because its identifier is a child of the release identifier.
if bundleIdentifier == self.debugApplicationDefaultsDomain
|| bundleIdentifier.hasPrefix("\(self.debugApplicationDefaultsDomain).")
{
return self.debugApplicationDefaultsDomain
}
if bundleIdentifier == self.releaseApplicationDefaultsDomain
|| bundleIdentifier.hasPrefix("\(self.releaseApplicationDefaultsDomain).")
{
return self.releaseApplicationDefaultsDomain
}
return nil
}
#if DEBUG
public static func withTaskOverrideForTesting<T>(
_ mode: ClaudeOAuthKeychainPromptMode?,
operation: () throws -> T) rethrows -> T
{
try self.$taskOverride.withValue(mode) {
try operation()
}
}
public static func withTaskOverrideForTesting<T>(
_ mode: ClaudeOAuthKeychainPromptMode?,
operation: () async throws -> T) async rethrows -> T
{
try await self.$taskOverride.withValue(mode) {
try await operation()
}
}
public static var currentTaskOverrideForTesting: ClaudeOAuthKeychainPromptMode? {
self.taskOverride
}
static func withApplicationUserDefaultsOverrideForTesting<T>(
_ userDefaults: UserDefaults?,
operation: () throws -> T) rethrows -> T
{
try self.$taskApplicationUserDefaultsOverride.withValue(userDefaults.map(UserDefaultsBox.init)) {
try operation()
}
}
static func withApplicationUserDefaultsOverrideForTesting<T>(
_ userDefaults: UserDefaults?,
operation: () async throws -> T) async rethrows -> T
{
try await self.$taskApplicationUserDefaultsOverride.withValue(userDefaults.map(UserDefaultsBox.init)) {
try await operation()
}
}
#endif
}
@@ -0,0 +1,31 @@
import Dispatch
import Foundation
#if os(macOS)
import Security
enum ClaudeOAuthKeychainQueryTiming {
static func copyMatching(_ query: [String: Any]) -> (status: OSStatus, result: AnyObject?, durationMs: Double) {
var result: AnyObject?
let startedAtNs = DispatchTime.now().uptimeNanoseconds
let status = KeychainSecurity.copyMatching(query as CFDictionary, &result)
let durationMs = Double(DispatchTime.now().uptimeNanoseconds - startedAtNs) / 1_000_000.0
return (status, result, durationMs)
}
static func backoffIfSlowNoUIQuery(_ durationMs: Double, _ service: String, _ log: CodexBarLogger) -> Bool {
// Intentionally no longer treats "slow" no-UI Keychain queries as a denial. Some systems can have
// non-deterministic timing characteristics that would make this backoff too aggressive and surprising.
//
// Keep this hook so call sites can cheaply log slow queries during debugging without changing behavior.
guard ProviderInteractionContext.current == .background, durationMs > 1000 else { return false }
log.debug(
"Claude keychain no-UI query was slow",
metadata: [
"service": service,
"duration_ms": String(format: "%.2f", durationMs),
])
return false
}
}
#endif
@@ -0,0 +1,49 @@
import Foundation
public enum ClaudeOAuthKeychainReadStrategy: String, Sendable, Codable, CaseIterable {
case securityFramework
case securityCLIExperimental
}
public enum ClaudeOAuthKeychainReadStrategyPreference {
private static let userDefaultsKey = "claudeOAuthKeychainReadStrategy"
#if DEBUG
@TaskLocal private static var taskOverride: ClaudeOAuthKeychainReadStrategy?
#endif
public static func current(userDefaults: UserDefaults = .standard) -> ClaudeOAuthKeychainReadStrategy {
#if DEBUG
if let taskOverride { return taskOverride }
#endif
if let raw = userDefaults.string(forKey: self.userDefaultsKey) {
let strategy = ClaudeOAuthKeychainReadStrategy(rawValue: raw) ?? .securityFramework
return strategy == .securityCLIExperimental ? .securityFramework : strategy
}
return .securityFramework
}
#if DEBUG
public static func withTaskOverrideForTesting<T>(
_ strategy: ClaudeOAuthKeychainReadStrategy?,
operation: () throws -> T) rethrows -> T
{
try self.$taskOverride.withValue(strategy) {
try operation()
}
}
public static func withTaskOverrideForTesting<T>(
_ strategy: ClaudeOAuthKeychainReadStrategy?,
operation: () async throws -> T) async rethrows -> T
{
try await self.$taskOverride.withValue(strategy) {
try await operation()
}
}
public static var currentTaskOverrideForTesting: ClaudeOAuthKeychainReadStrategy? {
self.taskOverride
}
#endif
}
@@ -0,0 +1,35 @@
import Foundation
#if DEBUG
extension ClaudeOAuthCredentialsStore {
final class ClaudeKeychainOverrideStore: @unchecked Sendable {
var data: Data?
var fingerprint: ClaudeKeychainFingerprint?
init(data: Data? = nil, fingerprint: ClaudeKeychainFingerprint? = nil) {
self.data = data
self.fingerprint = fingerprint
}
}
@TaskLocal static var taskClaudeKeychainOverrideStore: ClaudeKeychainOverrideStore?
static func withMutableClaudeKeychainOverrideStoreForTesting<T>(
_ store: ClaudeKeychainOverrideStore?,
operation: () throws -> T) rethrows -> T
{
try self.$taskClaudeKeychainOverrideStore.withValue(store) {
try operation()
}
}
static func withMutableClaudeKeychainOverrideStoreForTesting<T>(
_ store: ClaudeKeychainOverrideStore?,
operation: () async throws -> T) async rethrows -> T
{
try await self.$taskClaudeKeychainOverrideStore.withValue(store) {
try await operation()
}
}
}
#endif
@@ -0,0 +1,142 @@
#if canImport(Darwin)
import Darwin
#elseif canImport(Glibc)
import Glibc
#elseif canImport(Musl)
import Musl
#endif
import Foundation
protocol ClaudeOAuthPendingCacheClearStore: Sendable {
var isPending: Bool { get }
func markPending()
func withCacheTransaction(_ operation: (inout Bool) -> Void)
}
final class ClaudeOAuthPendingCacheClearUserDefaultsStore: ClaudeOAuthPendingCacheClearStore, @unchecked Sendable {
private static let processLock = NSLock()
private static let log = CodexBarLog.logger(LogCategories.claudeUsage)
private let domain: String
private let key: String
private let lockURL: URL
init(
domain: String,
key: String,
lockURL: URL = URL(fileURLWithPath: NSHomeDirectory(), isDirectory: true)
.appendingPathComponent(".codexbar", isDirectory: true)
.appendingPathComponent("claude-oauth-cache.lock"))
{
self.domain = domain
self.key = key
self.lockURL = lockURL
}
var isPending: Bool {
do {
return try self.withInterprocessLock {
self.currentGeneration() != nil
}
} catch {
Self.log.error("Claude OAuth cache tombstone lock failed: \(error.localizedDescription)")
return true
}
}
func markPending() {
do {
try self.withInterprocessLock {
self.writeGeneration(UUID().uuidString)
}
} catch {
// A surviving tombstone is safer than allowing a stale cache read. If the lock itself is unavailable,
// write a fresh generation but never clear one through the unlocked fallback path.
Self.log.error("Claude OAuth cache tombstone lock failed: \(error.localizedDescription)")
self.writeGeneration(UUID().uuidString)
}
}
func withCacheTransaction(_ operation: (inout Bool) -> Void) {
do {
try self.withInterprocessLock {
let initialGeneration = self.currentGeneration()
var pending = initialGeneration != nil
operation(&pending)
self.persist(
pending: pending,
initialGeneration: initialGeneration)
}
} catch {
Self.log.error("Claude OAuth cache transaction lock failed: \(error.localizedDescription)")
// Fail closed: without the shared lock, do not touch the cache and leave a fresh invalidation marker.
self.writeGeneration(UUID().uuidString)
}
}
private func persist(
pending: Bool,
initialGeneration: String?)
{
let currentGeneration = self.currentGeneration()
if pending {
if currentGeneration == nil {
self.writeGeneration(UUID().uuidString)
}
return
}
// Compare the observed generation before removal. This is defensive against writers that do not yet honor
// the lock, while the lock serializes all current app and bundled-CLI cache mutations.
guard currentGeneration == initialGeneration else { return }
self.writeGeneration(nil)
}
private func currentGeneration() -> String? {
let userDefaults = UserDefaults(suiteName: self.domain) ?? .standard
userDefaults.synchronize()
if let generation = userDefaults.string(forKey: self.key), !generation.isEmpty {
return generation
}
// V1 stored a boolean. Preserve an outstanding invalidation across the generation-based upgrade.
if userDefaults.object(forKey: self.key) as? Bool == true {
return "legacy-boolean"
}
return nil
}
private func writeGeneration(_ generation: String?) {
let userDefaults = UserDefaults(suiteName: self.domain) ?? .standard
if let generation {
userDefaults.set(generation, forKey: self.key)
} else {
userDefaults.removeObject(forKey: self.key)
}
userDefaults.synchronize()
}
private func withInterprocessLock<T>(_ operation: () throws -> T) throws -> T {
Self.processLock.lock()
defer { Self.processLock.unlock() }
try FileManager.default.createDirectory(
at: self.lockURL.deletingLastPathComponent(),
withIntermediateDirectories: true)
let fd = open(self.lockURL.path, O_CREAT | O_RDWR | O_CLOEXEC, S_IRUSR | S_IWUSR)
guard fd >= 0 else {
throw POSIXError(POSIXErrorCode(rawValue: errno) ?? .EIO)
}
defer {
_ = flock(fd, LOCK_UN)
close(fd)
}
while flock(fd, LOCK_EX) != 0 {
guard errno == EINTR else {
throw POSIXError(POSIXErrorCode(rawValue: errno) ?? .EIO)
}
}
return try operation()
}
}
@@ -0,0 +1,371 @@
import Foundation
#if os(macOS)
import os.lock
public enum ClaudeOAuthRefreshFailureGate {
public enum BlockStatus: Equatable, Sendable {
case terminal(reason: String?, failures: Int)
case transient(until: Date, failures: Int)
}
struct AuthFingerprint: Codable, Equatable {
let keychain: ClaudeOAuthCredentialsStore.ClaudeKeychainFingerprint?
let credentialsFile: String?
}
private struct State {
var loaded = false
var terminalFailureCount = 0
var transientFailureCount = 0
var isTerminalBlocked = false
var transientBlockedUntil: Date?
var fingerprintAtFailure: AuthFingerprint?
var lastCredentialsRecheckAt: Date?
var terminalReason: String?
}
private static let lock = OSAllocatedUnfairLock<State>(initialState: State())
private static let blockedUntilKey = "claudeOAuthRefreshBackoffBlockedUntilV1" // legacy (migration)
private static let failureCountKey = "claudeOAuthRefreshBackoffFailureCountV1" // legacy + terminal count
private static let fingerprintKey = "claudeOAuthRefreshBackoffFingerprintV2"
private static let terminalBlockedKey = "claudeOAuthRefreshTerminalBlockedV1"
private static let terminalReasonKey = "claudeOAuthRefreshTerminalReasonV1"
private static let transientBlockedUntilKey = "claudeOAuthRefreshTransientBlockedUntilV1"
private static let transientFailureCountKey = "claudeOAuthRefreshTransientFailureCountV1"
private static let log = CodexBarLog.logger(LogCategories.claudeUsage)
private static let minimumCredentialsRecheckInterval: TimeInterval = 15
private static let unknownFingerprint = AuthFingerprint(keychain: nil, credentialsFile: nil)
private static let transientBaseInterval: TimeInterval = 60 * 5
private static let transientMaxInterval: TimeInterval = 60 * 60 * 6
#if DEBUG
@TaskLocal static var shouldAttemptOverride: Bool?
private nonisolated(unsafe) static var fingerprintProviderOverride: (() -> AuthFingerprint?)?
static func setFingerprintProviderOverrideForTesting(_ provider: (() -> AuthFingerprint?)?) {
self.fingerprintProviderOverride = provider
}
public static func resetInMemoryStateForTesting() {
self.lock.withLock { state in
state.loaded = false
state.terminalFailureCount = 0
state.transientFailureCount = 0
state.isTerminalBlocked = false
state.transientBlockedUntil = nil
state.fingerprintAtFailure = nil
state.lastCredentialsRecheckAt = nil
state.terminalReason = nil
}
}
public static func resetForTesting() {
self.lock.withLock { state in
state.loaded = false
state.terminalFailureCount = 0
state.transientFailureCount = 0
state.isTerminalBlocked = false
state.transientBlockedUntil = nil
state.fingerprintAtFailure = nil
state.lastCredentialsRecheckAt = nil
state.terminalReason = nil
UserDefaults.standard.removeObject(forKey: self.blockedUntilKey)
UserDefaults.standard.removeObject(forKey: self.failureCountKey)
UserDefaults.standard.removeObject(forKey: self.fingerprintKey)
UserDefaults.standard.removeObject(forKey: self.terminalBlockedKey)
UserDefaults.standard.removeObject(forKey: self.terminalReasonKey)
UserDefaults.standard.removeObject(forKey: self.transientBlockedUntilKey)
UserDefaults.standard.removeObject(forKey: self.transientFailureCountKey)
}
}
#endif
public static func shouldAttempt(now: Date = Date()) -> Bool {
#if DEBUG
if let override = self.shouldAttemptOverride { return override }
#endif
return self.lock.withLock { state in
let didMigrate = self.loadIfNeeded(&state, now: now)
if didMigrate {
self.persist(state)
}
if state.isTerminalBlocked {
guard self.shouldRecheckCredentials(now: now, state: state) else { return false }
state.lastCredentialsRecheckAt = now
if self.hasCredentialsChangedSinceFailure(state) {
self.resetState(&state)
self.persist(state)
return true
}
self.log.debug(
"Claude OAuth refresh blocked until auth changes",
metadata: [
"terminalFailures": "\(state.terminalFailureCount)",
"reason": state.terminalReason ?? "nil",
])
return false
}
if let blockedUntil = state.transientBlockedUntil {
if blockedUntil <= now {
self.clearTransientState(&state)
// Once transient backoff expires, forget its auth baseline so future failures capture fresh
// fingerprints and so we don't ratchet backoff across unrelated intermittent failures.
state.fingerprintAtFailure = nil
state.lastCredentialsRecheckAt = nil
self.persist(state)
return true
}
if self.shouldRecheckCredentials(now: now, state: state) {
state.lastCredentialsRecheckAt = now
if self.hasCredentialsChangedSinceFailure(state) {
self.resetState(&state)
self.persist(state)
return true
}
}
self.log.debug(
"Claude OAuth refresh transient backoff active",
metadata: [
"until": "\(blockedUntil.timeIntervalSince1970)",
"transientFailures": "\(state.transientFailureCount)",
])
return false
}
return true
}
}
public static func currentBlockStatus(now: Date = Date()) -> BlockStatus? {
self.lock.withLock { state in
_ = self.loadIfNeeded(&state, now: now)
if state.isTerminalBlocked {
return .terminal(reason: state.terminalReason, failures: state.terminalFailureCount)
}
if let blockedUntil = state.transientBlockedUntil, blockedUntil > now {
return .transient(until: blockedUntil, failures: state.transientFailureCount)
}
return nil
}
}
public static func recordTerminalAuthFailure(now: Date = Date()) {
self.lock.withLock { state in
_ = self.loadIfNeeded(&state, now: now)
state.terminalFailureCount += 1
state.isTerminalBlocked = true
state.terminalReason = "invalid_grant"
state.fingerprintAtFailure = self.currentFingerprint() ?? self.unknownFingerprint
state.lastCredentialsRecheckAt = now
self.clearTransientState(&state)
self.persist(state)
}
}
public static func recordTransientFailure(now: Date = Date()) {
self.lock.withLock { state in
_ = self.loadIfNeeded(&state, now: now)
// Keep terminal blocking monotonic: once we know auth is rejected (e.g. invalid_grant),
// do not downgrade it to time-based backoff unless auth changes (fingerprint) or we record success.
guard !state.isTerminalBlocked else { return }
self.clearTerminalState(&state)
state.transientFailureCount += 1
let interval = self.transientCooldownInterval(failures: state.transientFailureCount)
state.transientBlockedUntil = now.addingTimeInterval(interval)
state.fingerprintAtFailure = self.currentFingerprint() ?? self.unknownFingerprint
state.lastCredentialsRecheckAt = now
self.persist(state)
}
}
public static func recordAuthFailure(now: Date = Date()) {
// Legacy shim: treat as terminal auth failure.
self.recordTerminalAuthFailure(now: now)
}
public static func recordSuccess() {
self.lock.withLock { state in
_ = self.loadIfNeeded(&state, now: Date())
self.resetState(&state)
self.persist(state)
}
}
private static func shouldRecheckCredentials(now: Date, state: State) -> Bool {
guard let last = state.lastCredentialsRecheckAt else { return true }
return now.timeIntervalSince(last) >= self.minimumCredentialsRecheckInterval
}
private static func hasCredentialsChangedSinceFailure(_ state: State) -> Bool {
guard let current = self.currentFingerprint() else { return false }
guard let prior = state.fingerprintAtFailure else { return false }
return current != prior
}
private static func currentFingerprint() -> AuthFingerprint? {
#if DEBUG
if let override = self.fingerprintProviderOverride { return override() }
#endif
return AuthFingerprint(
keychain: ClaudeOAuthCredentialsStore.currentClaudeKeychainFingerprintWithoutPromptForAuthGate(),
credentialsFile: ClaudeOAuthCredentialsStore.currentCredentialsFileFingerprintWithoutPromptForAuthGate())
}
private static func loadIfNeeded(_ state: inout State, now: Date) -> Bool {
state.loaded = true
var didMutate = false
// Always refresh persisted fields from UserDefaults, even after first load.
//
// This avoids stale state when UserDefaults are modified while the app is running (or during tests),
// while still keeping ephemeral throttling state (like lastCredentialsRecheckAt) in memory.
state.terminalFailureCount = UserDefaults.standard.integer(forKey: self.failureCountKey)
state.transientFailureCount = UserDefaults.standard.integer(forKey: self.transientFailureCountKey)
if let raw = UserDefaults.standard.object(forKey: self.transientBlockedUntilKey) as? Double {
state.transientBlockedUntil = Date(timeIntervalSince1970: raw)
}
let legacyBlockedUntil = (UserDefaults.standard.object(forKey: self.blockedUntilKey) as? Double)
.map { Date(timeIntervalSince1970: $0) }
let legacyFailureCount = UserDefaults.standard.integer(forKey: self.failureCountKey)
if let data = UserDefaults.standard.data(forKey: self.fingerprintKey) {
state.fingerprintAtFailure = (try? JSONDecoder().decode(AuthFingerprint.self, from: data))
} else {
state.fingerprintAtFailure = nil
}
if UserDefaults.standard.object(forKey: self.terminalBlockedKey) != nil {
state.isTerminalBlocked = UserDefaults.standard.bool(forKey: self.terminalBlockedKey)
state.terminalReason = UserDefaults.standard.string(forKey: self.terminalReasonKey)
if legacyBlockedUntil != nil {
didMutate = true
}
} else {
// Migration: legacy keys represented a time-based backoff. Migrate to transient backoff (never terminal)
// unless we already have new transient keys persisted.
if UserDefaults.standard.object(forKey: self.transientFailureCountKey) == nil,
UserDefaults.standard.object(forKey: self.transientBlockedUntilKey) == nil,
legacyBlockedUntil != nil || legacyFailureCount > 0
{
state.isTerminalBlocked = false
state.terminalReason = nil
state.terminalFailureCount = 0
if let legacyBlockedUntil, legacyBlockedUntil > now {
state.transientFailureCount = max(legacyFailureCount, 0)
state.transientBlockedUntil = legacyBlockedUntil
} else {
state.transientFailureCount = 0
state.transientBlockedUntil = nil
}
didMutate = true
}
}
if state.isTerminalBlocked || state.transientBlockedUntil != nil, state.fingerprintAtFailure == nil {
state.fingerprintAtFailure = self.unknownFingerprint
didMutate = true
}
if legacyBlockedUntil != nil {
didMutate = true
}
return didMutate
}
private static func persist(_ state: State) {
UserDefaults.standard.set(state.terminalFailureCount, forKey: self.failureCountKey)
UserDefaults.standard.set(state.isTerminalBlocked, forKey: self.terminalBlockedKey)
if let reason = state.terminalReason {
UserDefaults.standard.set(reason, forKey: self.terminalReasonKey)
} else {
UserDefaults.standard.removeObject(forKey: self.terminalReasonKey)
}
UserDefaults.standard.set(state.transientFailureCount, forKey: self.transientFailureCountKey)
if let blockedUntil = state.transientBlockedUntil {
UserDefaults.standard.set(blockedUntil.timeIntervalSince1970, forKey: self.transientBlockedUntilKey)
} else {
UserDefaults.standard.removeObject(forKey: self.transientBlockedUntilKey)
}
UserDefaults.standard.removeObject(forKey: self.blockedUntilKey)
if let fingerprint = state.fingerprintAtFailure,
let data = try? JSONEncoder().encode(fingerprint)
{
UserDefaults.standard.set(data, forKey: self.fingerprintKey)
} else {
UserDefaults.standard.removeObject(forKey: self.fingerprintKey)
}
}
private static func transientCooldownInterval(failures: Int) -> TimeInterval {
guard failures > 0 else { return 0 }
let factor = pow(2.0, Double(failures - 1))
return min(self.transientBaseInterval * factor, self.transientMaxInterval)
}
private static func clearTerminalState(_ state: inout State) {
state.terminalFailureCount = 0
state.isTerminalBlocked = false
state.terminalReason = nil
}
private static func clearTransientState(_ state: inout State) {
state.transientFailureCount = 0
state.transientBlockedUntil = nil
}
private static func resetState(_ state: inout State) {
self.clearTerminalState(&state)
self.clearTransientState(&state)
state.fingerprintAtFailure = nil
state.lastCredentialsRecheckAt = nil
}
}
#else
public enum ClaudeOAuthRefreshFailureGate {
public enum BlockStatus: Equatable, Sendable {
case terminal(reason: String?, failures: Int)
case transient(until: Date, failures: Int)
}
public static func shouldAttempt(now _: Date = Date()) -> Bool {
true
}
public static func currentBlockStatus(now _: Date = Date()) -> BlockStatus? {
nil
}
public static func recordTerminalAuthFailure(now _: Date = Date()) {}
public static func recordTransientFailure(now _: Date = Date()) {}
public static func recordAuthFailure(now _: Date = Date()) {}
public static func recordSuccess() {}
#if DEBUG
static func setFingerprintProviderOverrideForTesting(_: (() -> Any?)?) {}
public static func resetInMemoryStateForTesting() {}
public static func resetForTesting() {}
#endif
}
#endif
@@ -0,0 +1,324 @@
import Foundation
#if canImport(FoundationNetworking)
import FoundationNetworking
#endif
public enum ClaudeOAuthFetchError: LocalizedError, Sendable {
case unauthorized
case rateLimited(retryAfter: Date?)
case invalidResponse
case serverError(Int, String?)
case networkError(Error)
public var errorDescription: String? {
switch self {
case .unauthorized:
return "Claude OAuth request unauthorized. Run `claude` to re-authenticate."
case .rateLimited:
return "Claude OAuth usage endpoint is rate limited by Anthropic right now. Wait a few minutes, "
+ "then click Refresh. If it keeps happening, run `claude logout && claude login`, then try again."
case .invalidResponse:
return "Claude OAuth response was invalid."
case let .serverError(code, body):
if let body, !body.isEmpty {
let cleaned = body
.replacingOccurrences(of: "\n", with: " ")
.trimmingCharacters(in: .whitespacesAndNewlines)
let shortened = cleaned.count > 400 ? String(cleaned.prefix(400)) + "" : cleaned
return "Claude OAuth error: HTTP \(code) \(shortened)"
}
return "Claude OAuth error: HTTP \(code)"
case let .networkError(error):
return "Claude OAuth network error: \(error.localizedDescription)"
}
}
}
enum ClaudeOAuthUsageFetcher {
private static let baseURL = "https://api.anthropic.com"
private static let usagePath = "/api/oauth/usage"
private static let betaHeader = "oauth-2025-04-20"
private static let fallbackClaudeCodeVersion = "2.1.0"
static func fetchUsage(
accessToken: String,
detectClaudeVersion: Bool = true) async throws -> OAuthUsageResponse
{
if let blockedUntil = ClaudeOAuthUsageRateLimitGate.blockedUntil() {
throw ClaudeOAuthFetchError.rateLimited(retryAfter: blockedUntil)
}
guard let url = URL(string: baseURL + usagePath) else {
throw ClaudeOAuthFetchError.invalidResponse
}
var request = URLRequest(url: url)
request.httpMethod = "GET"
request.timeoutInterval = 30
request.setValue("Bearer \(accessToken)", forHTTPHeaderField: "Authorization")
request.setValue("application/json", forHTTPHeaderField: "Accept")
request.setValue("application/json", forHTTPHeaderField: "Content-Type")
// OAuth usage endpoint currently requires the beta header.
request.setValue(Self.betaHeader, forHTTPHeaderField: "anthropic-beta")
request.setValue(
Self.claudeCodeUserAgent(detectClaudeVersion: detectClaudeVersion),
forHTTPHeaderField: "User-Agent")
do {
let response = try await ProviderHTTPClient.shared.response(for: request)
let data = response.data
switch response.statusCode {
case 200:
let usage = try Self.decodeUsageResponse(data)
ClaudeOAuthUsageRateLimitGate.recordSuccess()
return usage
case 401:
throw ClaudeOAuthFetchError.unauthorized
case 429:
let retryAfter = Self.retryAfterDate(from: response.response)
ClaudeOAuthUsageRateLimitGate.recordRateLimit(retryAfter: retryAfter)
throw ClaudeOAuthFetchError.rateLimited(
retryAfter: ClaudeOAuthUsageRateLimitGate.currentBlockedUntil() ?? retryAfter)
case 403:
let body = String(data: data, encoding: .utf8)
throw ClaudeOAuthFetchError.serverError(response.statusCode, body)
default:
let body = String(data: data, encoding: .utf8)
throw ClaudeOAuthFetchError.serverError(response.statusCode, body)
}
} catch let error as ClaudeOAuthFetchError {
throw error
} catch {
throw ClaudeOAuthFetchError.networkError(error)
}
}
static func decodeUsageResponse(_ data: Data) throws -> OAuthUsageResponse {
let decoder = JSONDecoder()
return try decoder.decode(OAuthUsageResponse.self, from: data)
}
static func parseISO8601Date(_ string: String?) -> Date? {
guard let string, !string.isEmpty else { return nil }
let formatter = ISO8601DateFormatter()
formatter.formatOptions = [.withInternetDateTime, .withFractionalSeconds]
if let date = formatter.date(from: string) { return date }
formatter.formatOptions = [.withInternetDateTime]
return formatter.date(from: string)
}
private static func retryAfterDate(from response: HTTPURLResponse, now: Date = Date()) -> Date? {
guard let raw = response.value(forHTTPHeaderField: "Retry-After")?
.trimmingCharacters(in: .whitespacesAndNewlines),
!raw.isEmpty
else { return nil }
if let seconds = TimeInterval(raw), seconds >= 0 {
return now.addingTimeInterval(seconds)
}
let formatter = DateFormatter()
formatter.locale = Locale(identifier: "en_US_POSIX")
formatter.timeZone = TimeZone(secondsFromGMT: 0)
formatter.dateFormat = "EEE',' dd MMM yyyy HH':'mm':'ss zzz"
return formatter.date(from: raw)
}
private static func claudeCodeUserAgent(
detectClaudeVersion: Bool,
versionDetector: () -> String? = { ProviderVersionDetector.claudeVersion() }) -> String
{
self.claudeCodeUserAgent(versionString: detectClaudeVersion ? versionDetector() : nil)
}
private static func claudeCodeUserAgent(versionString: String?) -> String {
let version = self.normalizedClaudeCodeVersion(versionString) ?? self.fallbackClaudeCodeVersion
return "claude-code/\(version)"
}
private static func normalizedClaudeCodeVersion(_ versionString: String?) -> String? {
guard let raw = versionString?.trimmingCharacters(in: .whitespacesAndNewlines), !raw.isEmpty else {
return nil
}
let token = raw.split(whereSeparator: \.isWhitespace).first.map(String.init) ?? raw
let trimmed = token.trimmingCharacters(in: .whitespacesAndNewlines)
return trimmed.isEmpty ? nil : trimmed
}
}
struct OAuthUsageResponse: Decodable {
let fiveHour: OAuthUsageWindow?
let sevenDay: OAuthUsageWindow?
let sevenDayOAuthApps: OAuthUsageWindow?
let sevenDayOpus: OAuthUsageWindow?
let sevenDaySonnet: OAuthUsageWindow?
let sevenDayRoutines: OAuthUsageWindow?
let sevenDayRoutinesSourceKey: String?
let iguanaNecktie: OAuthUsageWindow?
let extraUsage: OAuthExtraUsage?
/// Newer shape (superseding the flat `seven_day_*` fields above for scoped weekly
/// windows): a flat list of limit entries, each optionally naming the model it scopes
/// to via `scope.model.display_name` (e.g. "Fable" during a promotional access window).
let limits: [OAuthLimitEntry]?
init(from decoder: Decoder) throws {
let container = try decoder.container(keyedBy: DynamicCodingKey.self)
self.fiveHour = Self.decodeWindow(in: container, keys: ["five_hour"])
self.sevenDay = Self.decodeWindow(in: container, keys: ["seven_day"])
self.sevenDayOAuthApps = Self.decodeWindow(in: container, keys: ["seven_day_oauth_apps"])
self.sevenDayOpus = Self.decodeWindow(in: container, keys: ["seven_day_opus"])
self.sevenDaySonnet = Self.decodeWindow(in: container, keys: ["seven_day_sonnet"])
let routines = Self.decodeWindowWithSource(in: container, keys: [
"seven_day_routines",
"seven_day_claude_routines",
"claude_routines",
"routines",
"routine",
"seven_day_cowork",
"cowork",
])
self.sevenDayRoutines = routines.window
self.sevenDayRoutinesSourceKey = routines.sourceKey
self.iguanaNecktie = Self.decodeWindow(in: container, keys: ["iguana_necktie"])
self.extraUsage = Self.decodeValue(in: container, keys: ["extra_usage"])
self.limits = Self.decodeValue(in: container, keys: ["limits"])
}
private static func decodeWindow(
in container: KeyedDecodingContainer<DynamicCodingKey>,
keys: [String]) -> OAuthUsageWindow?
{
self.decodeValue(in: container, keys: keys)
}
private static func decodeWindowWithSource(
in container: KeyedDecodingContainer<DynamicCodingKey>,
keys: [String]) -> (window: OAuthUsageWindow?, sourceKey: String?)
{
var firstNullKey: String?
for keyName in keys {
guard let key = DynamicCodingKey(stringValue: keyName) else { continue }
guard container.contains(key) else { continue }
if let value = try? container.decodeIfPresent(OAuthUsageWindow.self, forKey: key) {
return (value, keyName)
}
if firstNullKey == nil {
firstNullKey = keyName
}
}
return (nil, firstNullKey)
}
private static func decodeValue<T: Decodable>(
in container: KeyedDecodingContainer<DynamicCodingKey>,
keys: [String]) -> T?
{
for keyName in keys {
guard let key = DynamicCodingKey(stringValue: keyName) else { continue }
if let value = try? container.decodeIfPresent(T.self, forKey: key) {
return value
}
}
return nil
}
}
private struct DynamicCodingKey: CodingKey {
let stringValue: String
let intValue: Int?
init?(stringValue: String) {
self.stringValue = stringValue
self.intValue = nil
}
init?(intValue: Int) {
nil
}
}
struct OAuthUsageWindow: Decodable {
let utilization: Double?
let resetsAt: String?
enum CodingKeys: String, CodingKey {
case utilization
case resetsAt = "resets_at"
}
}
/// A single entry from the `limits` array. `kind`/`group` classify the limit
/// (e.g. `kind: "weekly_scoped"`, `group: "weekly"`); `scope.model.display_name` names the
/// model it applies to when the limit is scoped to one, rather than the whole account.
struct OAuthLimitEntry: Decodable {
let kind: String?
let group: String?
let percent: Double?
let resetsAt: String?
let scope: OAuthLimitScope?
let isActive: Bool?
enum CodingKeys: String, CodingKey {
case kind
case group
case percent
case resetsAt = "resets_at"
case scope
case isActive = "is_active"
}
}
struct OAuthLimitScope: Decodable {
let model: OAuthLimitScopeModel?
}
struct OAuthLimitScopeModel: Decodable {
let id: String?
let displayName: String?
enum CodingKeys: String, CodingKey {
case id
case displayName = "display_name"
}
}
struct OAuthExtraUsage: Decodable {
let isEnabled: Bool?
let monthlyLimit: Double?
let usedCredits: Double?
let utilization: Double?
let currency: String?
enum CodingKeys: String, CodingKey {
case isEnabled = "is_enabled"
case monthlyLimit = "monthly_limit"
case usedCredits = "used_credits"
case utilization
case currency
}
}
#if DEBUG
extension ClaudeOAuthUsageFetcher {
static func _decodeUsageResponseForTesting(_ data: Data) throws -> OAuthUsageResponse {
try self.decodeUsageResponse(data)
}
static func _userAgentForTesting(versionString: String?) -> String {
self.claudeCodeUserAgent(versionString: versionString)
}
static func _userAgentForTesting(
detectClaudeVersion: Bool,
versionDetector: () -> String?) -> String
{
self.claudeCodeUserAgent(
detectClaudeVersion: detectClaudeVersion,
versionDetector: versionDetector)
}
static func _retryAfterDateForTesting(from response: HTTPURLResponse, now: Date) -> Date? {
self.retryAfterDate(from: response, now: now)
}
}
#endif
@@ -0,0 +1,46 @@
import Foundation
enum ClaudeOAuthUsageRateLimitGate {
private static let blockedUntilKey = "claudeOAuthUsageRateLimitBlockedUntilV1"
private static let defaultCooldown: TimeInterval = 60 * 5
static func blockedUntil(
interaction: ProviderInteraction = ProviderInteractionContext.current,
now: Date = Date()) -> Date?
{
guard interaction != .userInitiated else { return nil }
return self.currentBlockedUntil(now: now)
}
static func currentBlockedUntil(now: Date = Date()) -> Date? {
guard let raw = UserDefaults.standard.object(forKey: self.blockedUntilKey) as? Double else {
return nil
}
let blockedUntil = Date(timeIntervalSince1970: raw)
guard blockedUntil > now else {
UserDefaults.standard.removeObject(forKey: self.blockedUntilKey)
return nil
}
return blockedUntil
}
static func recordRateLimit(retryAfter: Date?, now: Date = Date()) {
let blockedUntil = if let retryAfter, retryAfter > now {
retryAfter
} else {
now.addingTimeInterval(self.defaultCooldown)
}
UserDefaults.standard.set(blockedUntil.timeIntervalSince1970, forKey: self.blockedUntilKey)
}
static func recordSuccess() {
UserDefaults.standard.removeObject(forKey: self.blockedUntilKey)
}
#if DEBUG
static func resetForTesting() {
UserDefaults.standard.removeObject(forKey: self.blockedUntilKey)
}
#endif
}
@@ -0,0 +1,179 @@
import Foundation
public enum ClaudePlan: String, CaseIterable, Sendable {
case max
case pro
case team
case enterprise
case ultra
public var brandedLoginMethod: String {
switch self {
case .max:
"Claude Max"
case .pro:
"Claude Pro"
case .team:
"Claude Team"
case .enterprise:
"Claude Enterprise"
case .ultra:
"Claude Ultra"
}
}
/// Branded label including the Max usage multiplier when the rate-limit tier carries one.
public func brandedLoginMethod(rateLimitTier: String?) -> String {
guard self == .max, let multiplier = Self.maxUsageMultiplier(rateLimitTier: rateLimitTier) else {
return self.brandedLoginMethod
}
return "\(self.brandedLoginMethod) \(multiplier)"
}
public var compactLoginMethod: String {
switch self {
case .max:
"Max"
case .pro:
"Pro"
case .team:
"Team"
case .enterprise:
"Enterprise"
case .ultra:
"Ultra"
}
}
public var countsAsSubscription: Bool {
switch self {
case .max, .pro, .team, .ultra:
true
case .enterprise:
false
}
}
public static func fromOAuthRateLimitTier(_ rateLimitTier: String?) -> Self? {
self.fromRateLimitTier(rateLimitTier)
}
public static func fromOAuthCredentials(subscriptionType: String?, rateLimitTier: String?) -> Self? {
self.fromCompatibilityLoginMethod(subscriptionType)
?? self.fromOAuthRateLimitTier(rateLimitTier)
}
public static func fromWebAccount(rateLimitTier: String?, billingType: String?) -> Self? {
if let plan = self.fromRateLimitTier(rateLimitTier) {
return plan
}
let tier = Self.normalized(rateLimitTier)
let billing = Self.normalized(billingType)
if billing.contains("stripe"), tier.contains("claude") {
return .pro
}
return nil
}
public static func fromCompatibilityLoginMethod(_ loginMethod: String?) -> Self? {
let words = Self.normalizedWords(loginMethod)
if words.isEmpty {
return nil
}
if words.contains("max") {
return .max
}
if words.contains("pro") {
return .pro
}
if words.contains("team") {
return .team
}
if words.contains("enterprise") {
return .enterprise
}
if words.contains("ultra") {
return .ultra
}
return nil
}
public static func oauthLoginMethod(rateLimitTier: String?) -> String? {
self.fromOAuthRateLimitTier(rateLimitTier)?.brandedLoginMethod(rateLimitTier: rateLimitTier)
}
public static func oauthLoginMethod(subscriptionType: String?, rateLimitTier: String?) -> String? {
self.fromOAuthCredentials(
subscriptionType: subscriptionType,
rateLimitTier: rateLimitTier)?.brandedLoginMethod(rateLimitTier: rateLimitTier)
}
public static func webLoginMethod(rateLimitTier: String?, billingType: String?) -> String? {
self.fromWebAccount(
rateLimitTier: rateLimitTier,
billingType: billingType)?.brandedLoginMethod(rateLimitTier: rateLimitTier)
}
public static func cliCompatibilityLoginMethod(_ loginMethod: String?) -> String? {
guard let loginMethod = loginMethod?.trimmingCharacters(in: .whitespacesAndNewlines),
!loginMethod.isEmpty
else {
return nil
}
if let plan = self.fromCompatibilityLoginMethod(loginMethod) {
return plan.compactLoginMethod
}
return loginMethod
}
public static func isSubscriptionLoginMethod(_ loginMethod: String?) -> Bool {
self.fromCompatibilityLoginMethod(loginMethod)?.countsAsSubscription ?? false
}
private static func fromRateLimitTier(_ rateLimitTier: String?) -> Self? {
let tier = Self.normalized(rateLimitTier)
if tier.contains("max") {
return .max
}
if tier.contains("pro") {
return .pro
}
if tier.contains("team") {
return .team
}
if tier.contains("enterprise") {
return .enterprise
}
return nil
}
private static func maxUsageMultiplier(rateLimitTier: String?) -> String? {
let words = Self.normalizedWords(rateLimitTier)
guard let maxIndex = words.firstIndex(of: Self.max.rawValue),
words.indices.contains(maxIndex + 1)
else {
return nil
}
let multiplier = words[maxIndex + 1]
guard multiplier.last == "x", Int(multiplier.dropLast()) != nil else {
return nil
}
return multiplier
}
private static func normalized(_ text: String?) -> String {
text?
.trimmingCharacters(in: .whitespacesAndNewlines)
.lowercased() ?? ""
}
private static func normalizedWords(_ text: String?) -> [String] {
self.normalized(text)
.split(whereSeparator: { !$0.isLetter && !$0.isNumber })
.map(String.init)
}
}
@@ -0,0 +1,105 @@
import Foundation
enum ClaudeProbeSessionArtifactCleaner {
private static let log = CodexBarLog.logger(LogCategories.claudeProbe)
private static let maxProjectDirectoryNameLength = 200
@discardableResult
static func cleanupProbeSessionArtifacts(
probeDirectory: URL = ClaudeStatusProbe.probeWorkingDirectoryURL(),
environment: [String: String] = ProcessInfo.processInfo.environment,
fileManager fm: FileManager = .default) -> [URL]
{
let projectDirectoryName = self.claudeProjectDirectoryName(for: probeDirectory)
var visitedDirectories = Set<String>()
var removedFiles: [URL] = []
for root in self.claudeConfigRoots(environment: environment, fileManager: fm) {
let projectsRoot = root.appendingPathComponent("projects", isDirectory: true)
let directories = [projectsRoot.appendingPathComponent(projectDirectoryName, isDirectory: true)]
for directory in directories where visitedDirectories.insert(directory.path).inserted {
guard let entries = try? fm.contentsOfDirectory(
at: directory,
includingPropertiesForKeys: [.isRegularFileKey],
options: [.skipsHiddenFiles])
else { continue }
for entry in entries where entry.pathExtension == "jsonl" {
let values = try? entry.resourceValues(forKeys: [.isRegularFileKey])
guard values?.isRegularFile == true else { continue }
do {
try fm.removeItem(at: entry)
removedFiles.append(entry)
} catch {
Self.log.debug(
"Claude probe session artifact cleanup skipped file",
metadata: ["error": error.localizedDescription])
}
}
if (try? fm.contentsOfDirectory(atPath: directory.path).isEmpty) == true {
try? fm.removeItem(at: directory)
}
}
}
return removedFiles
}
static func claudeProjectDirectoryName(for directory: URL) -> String {
let path = directory.path.precomposedStringWithCanonicalMapping
let sanitized = String(path.utf16.map { codeUnit in
switch codeUnit {
case 48...57, 65...90, 97...122:
Character(UnicodeScalar(codeUnit)!)
default:
"-"
}
})
guard sanitized.count > self.maxProjectDirectoryNameLength else { return sanitized }
return "\(sanitized.prefix(self.maxProjectDirectoryNameLength))-\(self.javaScriptHashBase36(path))"
}
private static func javaScriptHashBase36(_ string: String) -> String {
var hash: Int32 = 0
for codeUnit in string.utf16 {
hash = hash &* 31 &+ Int32(truncatingIfNeeded: codeUnit)
}
let magnitude = hash < 0 ? -Int64(hash) : Int64(hash)
return String(magnitude, radix: 36)
}
private static func claudeConfigRoots(
environment: [String: String],
fileManager fm: FileManager) -> [URL]
{
var roots: [URL] = []
var seen = Set<String>()
func append(_ url: URL) {
let standardized = url.standardizedFileURL
guard seen.insert(standardized.path).inserted else { return }
roots.append(standardized)
}
if let raw = environment["CLAUDE_CONFIG_DIR"] {
for part in raw.split(separator: ",") {
let path = part.trimmingCharacters(in: .whitespacesAndNewlines)
guard !path.isEmpty else { continue }
append(URL(fileURLWithPath: path))
}
}
let home = environment["HOME"].flatMap { $0.isEmpty ? nil : $0 } ?? NSHomeDirectory()
append(URL(fileURLWithPath: home).appendingPathComponent(".claude", isDirectory: true))
append(URL(fileURLWithPath: home).appendingPathComponent(".config/claude", isDirectory: true))
if roots.isEmpty {
append(fm.homeDirectoryForCurrentUser.appendingPathComponent(".claude", isDirectory: true))
}
return roots
}
}
@@ -0,0 +1,679 @@
import Foundation
public enum ClaudeProviderDescriptor {
public static let descriptor: ProviderDescriptor = Self.makeDescriptor()
static func makeDescriptor() -> ProviderDescriptor {
ProviderDescriptor(
id: .claude,
metadata: ProviderMetadata(
id: .claude,
displayName: "Claude",
sessionLabel: "Session",
weeklyLabel: "Weekly",
opusLabel: "Sonnet",
supportsOpus: true,
supportsCredits: false,
creditsHint: "",
toggleTitle: "Show Claude Code usage",
cliName: "claude",
defaultEnabled: false,
isPrimaryProvider: true,
usesAccountFallback: false,
browserCookieOrder: ProviderBrowserCookieDefaults.defaultImportOrder,
dashboardURL: "https://console.anthropic.com/settings/billing",
subscriptionDashboardURL: "https://claude.ai/settings/usage",
changelogURL: "https://github.com/anthropics/claude-code/releases",
statusPageURL: "https://status.claude.com/"),
branding: ProviderBranding(
iconStyle: .claude,
iconResourceName: "ProviderIcon-claude",
color: ProviderColor(red: 204 / 255, green: 124 / 255, blue: 94 / 255)),
tokenCost: ProviderTokenCostConfig(
supportsTokenCost: true,
noDataMessage: self.noDataMessage),
fetchPlan: ProviderFetchPlan(
sourceModes: [.auto, .api, .web, .cli, .oauth],
pipeline: ProviderFetchPipeline(resolveStrategies: self.resolveStrategies)),
cli: ProviderCLIConfig(
name: "claude",
versionDetector: { browserDetection in
ClaudeUsageFetcher(browserDetection: browserDetection).detectVersion()
}))
}
private static func resolveStrategies(context: ProviderFetchContext) async -> [any ProviderFetchStrategy] {
if context.sourceMode == .api || self.hasAutoAdminAPIKey(context: context) {
return [ClaudeAdminAPIFetchStrategy()]
}
if ClaudeAdminAPIFetchStrategy.isSelectedAdminAPIAccount(context: context) {
return [ClaudeAdminAPIFetchStrategy()]
}
let planningInput = Self.makePlanningInput(context: context)
let plan = ClaudeSourcePlanner.resolve(input: planningInput)
let manualCookieHeader = Self.manualCookieHeader(from: context)
return plan.orderedSteps.map { step in
let strategy: any ProviderFetchStrategy = switch step.dataSource {
case .api:
ClaudeAdminAPIFetchStrategy()
case .oauth:
ClaudeOAuthFetchStrategy()
case .web:
ClaudeWebFetchStrategy(browserDetection: context.browserDetection)
case .cli:
ClaudeCLIFetchStrategy(
useWebExtras: context.runtime == .app
&& planningInput.webExtrasEnabled,
manualCookieHeader: manualCookieHeader,
browserDetection: context.browserDetection,
hasWebFallback: planningInput.hasWebSession)
case .auto:
fatalError("Planner must not emit .auto as an executable step.")
}
return ClaudePlannedFetchStrategy(base: strategy, plannedStep: step)
}
}
private static func hasAutoAdminAPIKey(context: ProviderFetchContext) -> Bool {
context.sourceMode == .auto && ClaudeAdminAPISettingsReader.apiKey(environment: context.env) != nil
}
private static func makePlanningInput(context: ProviderFetchContext) -> ClaudeSourcePlanningInput {
let webExtrasEnabled = context.settings?.claude?.webExtrasEnabled ?? false
let needsOAuthAvailability = context.runtime == .app && context.sourceMode == .auto
let hasWebSession = Self.hasPlausibleWebSession(context: context)
return ClaudeSourcePlanningInput(
runtime: context.runtime,
selectedDataSource: Self.sourceDataSource(from: context.sourceMode),
webExtrasEnabled: webExtrasEnabled,
hasWebSession: hasWebSession,
hasCLI: ClaudeCLIResolver.isAvailable(environment: context.env),
hasOAuthCredentials: needsOAuthAvailability && ClaudeOAuthPlanningAvailability.isAvailable(
runtime: context.runtime,
sourceMode: context.sourceMode,
environment: context.env))
}
private static func hasPlausibleWebSession(context: ProviderFetchContext) -> Bool {
switch context.sourceMode {
case .api, .oauth, .cli:
return false
case .web:
// Explicit web performs its cookie/session work inside the bounded fetch.
return context.settings?.claude?.cookieSource != .off
case .auto:
break
}
guard ClaudeWebFetchStrategy.isSupportedOnCurrentPlatform else { return false }
switch context.settings?.claude?.cookieSource {
case .off?:
return false
case .manual?:
return ClaudeWebFetchStrategy.hasManualSessionKey(context: context)
case .auto?, nil:
// Browser/Keychain inspection can block. Keep planning synchronous and let the web step
// perform the bounded availability check if app-auto actually reaches its last fallback.
// CLI auto continues to perform its real session import inside the bounded web fetch.
return true
}
}
private static func manualCookieHeader(from context: ProviderFetchContext) -> String? {
guard context.settings?.claude?.cookieSource == .manual else { return nil }
return CookieHeaderNormalizer.normalize(context.settings?.claude?.manualCookieHeader)
}
private static func noDataMessage() -> String {
"No Claude usage logs found in ~/.config/claude/projects, ~/.claude/projects, " +
"or Claude Desktop sessions."
}
public static func resolveUsageStrategy(
selectedDataSource: ClaudeUsageDataSource,
webExtrasEnabled: Bool,
hasWebSession: Bool,
hasCLI: Bool,
hasOAuthCredentials: Bool) -> ClaudeUsageStrategy
{
let plan = ClaudeSourcePlanner.resolve(input: ClaudeSourcePlanningInput(
runtime: .app,
selectedDataSource: selectedDataSource,
webExtrasEnabled: webExtrasEnabled,
hasWebSession: hasWebSession,
hasCLI: hasCLI,
hasOAuthCredentials: hasOAuthCredentials))
return plan.compatibilityStrategy ?? ClaudeUsageStrategy(dataSource: selectedDataSource, useWebExtras: false)
}
private static func sourceDataSource(from mode: ProviderSourceMode) -> ClaudeUsageDataSource {
switch mode {
case .auto:
.auto
case .api:
.api
case .web:
.web
case .cli:
.cli
case .oauth:
.oauth
}
}
}
public struct ClaudeUsageStrategy: Equatable, Sendable {
public let dataSource: ClaudeUsageDataSource
public let useWebExtras: Bool
}
public enum ClaudeOAuthPlanningAvailability {
public static func isAvailable(
runtime: ProviderRuntime,
sourceMode: ProviderSourceMode,
environment: [String: String]) -> Bool
{
ClaudeOAuthFetchStrategy.isPlausiblyAvailable(
runtime: runtime,
sourceMode: sourceMode,
environment: environment)
}
}
private struct ClaudePlannedFetchStrategy: ProviderFetchStrategy {
let base: any ProviderFetchStrategy
let plannedStep: ClaudeFetchPlanStep
var id: String {
self.base.id
}
var kind: ProviderFetchKind {
self.base.kind
}
func isAvailable(_ context: ProviderFetchContext) async -> Bool {
guard context.sourceMode == .auto else {
return await self.base.isAvailable(context)
}
guard self.plannedStep.isPlausiblyAvailable else { return false }
if self.plannedStep.dataSource == .cli ||
(context.runtime == .app && self.plannedStep.dataSource == .web)
{
return await self.base.isAvailable(context)
}
return true
}
func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult {
try await self.base.fetch(context)
}
func shouldFallback(on error: Error, context: ProviderFetchContext) -> Bool {
self.base.shouldFallback(on: error, context: context)
}
}
struct ClaudeAdminAPIFetchStrategy: ProviderFetchStrategy {
let id: String = "claude.admin-api"
let kind: ProviderFetchKind = .apiToken
let usageFetcher: @Sendable (String) async throws -> ClaudeAdminAPIUsageSnapshot
init(
usageFetcher: @escaping @Sendable (String) async throws -> ClaudeAdminAPIUsageSnapshot = { apiKey in
try await ClaudeAdminAPIUsageFetcher.fetchUsage(apiKey: apiKey)
})
{
self.usageFetcher = usageFetcher
}
static func isSelectedAdminAPIAccount(context: ProviderFetchContext) -> Bool {
guard context.selectedTokenAccountID != nil else { return false }
return self.resolveToken(environment: context.env) != nil
}
func isAvailable(_ context: ProviderFetchContext) async -> Bool {
Self.resolveToken(environment: context.env) != nil
}
func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult {
guard let apiKey = Self.resolveToken(environment: context.env) else {
throw ClaudeAdminAPISettingsError.missingToken
}
let usage = try await self.usageFetcher(apiKey)
return self.makeResult(
usage: usage.toUsageSnapshot(),
sourceLabel: "admin-api")
}
func shouldFallback(on _: Error, context: ProviderFetchContext) -> Bool {
context.runtime == .app &&
context.sourceMode == .auto &&
!Self.isSelectedAdminAPIAccount(context: context)
}
private static func resolveToken(environment: [String: String]) -> String? {
ProviderTokenResolver.claudeAdminAPIToken(environment: environment)
}
}
struct ClaudeOAuthFetchStrategy: ProviderFetchStrategy {
let id: String = "claude.oauth"
let kind: ProviderFetchKind = .oauth
#if DEBUG
@TaskLocal static var nonInteractiveCredentialRecordOverride: ClaudeOAuthCredentialRecord?
@TaskLocal static var claudeCLIAvailableOverride: Bool?
#endif
private func loadNonInteractiveCredentialRecord(environment: [String: String]) -> ClaudeOAuthCredentialRecord? {
#if DEBUG
if let override = Self.nonInteractiveCredentialRecordOverride {
return override
}
#endif
return try? ClaudeOAuthCredentialsStore.loadRecord(
environment: environment,
allowKeychainPrompt: false,
respectKeychainPromptCooldown: true,
allowClaudeKeychainRepairWithoutPrompt: false)
}
private func isClaudeCLIAvailable(environment: [String: String]) -> Bool {
#if DEBUG
if let override = Self.claudeCLIAvailableOverride {
return override
}
#endif
return ClaudeCLIResolver.isAvailable(environment: environment)
}
static func isPlausiblyAvailable(
runtime: ProviderRuntime,
sourceMode: ProviderSourceMode,
environment: [String: String]) -> Bool
{
let hasEnvironmentOAuthToken = !(environment[ClaudeOAuthCredentialsStore.environmentTokenKey]?
.trimmingCharacters(in: .whitespacesAndNewlines)
.isEmpty ?? true)
if hasEnvironmentOAuthToken {
return true
}
let strategy = ClaudeOAuthFetchStrategy()
let nonInteractiveRecord = strategy.loadNonInteractiveCredentialRecord(environment: environment)
let nonInteractiveCredentials = nonInteractiveRecord?.credentials
let hasRequiredScopeWithoutPrompt = nonInteractiveCredentials?.scopes.contains("user:profile") == true
if hasRequiredScopeWithoutPrompt, nonInteractiveCredentials?.isExpired == false {
return true
}
let claudeCLIAvailable = strategy.isClaudeCLIAvailable(environment: environment)
if let nonInteractiveRecord, hasRequiredScopeWithoutPrompt, nonInteractiveRecord.credentials.isExpired {
switch nonInteractiveRecord.owner {
case .codexbar:
let refreshToken = nonInteractiveRecord.credentials.refreshToken?
.trimmingCharacters(in: .whitespacesAndNewlines) ?? ""
if sourceMode == .auto {
return !refreshToken.isEmpty
}
return true
case .claudeCLI:
guard sourceMode == .auto else { return true }
guard claudeCLIAvailable else { return false }
guard ProviderInteractionContext.current == .background else { return true }
return !Self.hasMcpOAuthOnlyClaudeKeychainPayload(environment: environment)
case .environment:
return sourceMode != .auto
}
}
guard sourceMode == .auto else { return true }
let fallbackPromptMode = ClaudeOAuthKeychainPromptPreference.securityFrameworkFallbackMode()
let promptPolicyApplicable = ClaudeOAuthKeychainPromptPreference.isApplicable()
if ProviderInteractionContext.current == .userInitiated {
_ = ClaudeOAuthKeychainAccessGate.clearDenied()
}
let shouldAllowStartupBootstrap = runtime == .app &&
ProviderRefreshContext.current == .startup &&
ProviderInteractionContext.current == .background &&
fallbackPromptMode == .onlyOnUserAction &&
!ClaudeOAuthCredentialsStore.hasCachedCredentials(environment: environment)
if shouldAllowStartupBootstrap {
return ClaudeOAuthKeychainAccessGate.shouldAllowPrompt()
}
if promptPolicyApplicable,
!ClaudeOAuthKeychainAccessGate.shouldAllowPrompt()
{
return false
}
return ClaudeOAuthCredentialsStore.hasClaudeKeychainCredentialsWithoutPrompt()
}
private static func hasMcpOAuthOnlyClaudeKeychainPayload(environment: [String: String]) -> Bool {
ClaudeOAuthCredentialsStore.isMcpOAuthOnlyClaudeKeychainPayloadPresent(
interaction: ProviderInteractionContext.current,
environment: environment)
}
func isAvailable(_ context: ProviderFetchContext) async -> Bool {
Self.isPlausiblyAvailable(
runtime: context.runtime,
sourceMode: context.sourceMode,
environment: context.env)
}
func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult {
let fetcher = ClaudeUsageFetcher(
browserDetection: context.browserDetection,
environment: context.env,
runtime: context.runtime,
dataSource: .oauth,
oauthKeychainPromptCooldownEnabled: context.sourceMode == .auto,
allowBackgroundDelegatedRefresh: false,
allowStartupBootstrapPrompt: context.runtime == .app &&
(context.sourceMode == .auto || context.sourceMode == .oauth),
useWebExtras: false)
let usage = try await fetcher.loadLatestUsage(model: "sonnet")
return ProviderFetchResult(
usage: Self.snapshot(from: usage),
credits: nil,
dashboard: nil,
sourceLabel: "oauth",
strategyID: self.id,
strategyKind: self.kind,
claudeOAuthKeychainPersistentRefHash: usage.oauthKeychainPersistentRefHash,
claudeOAuthHistoryOwnerIdentifier: usage.oauthHistoryOwnerIdentifier,
claudeOAuthKeychainCredentialMismatch: usage.oauthKeychainCredentialMismatch,
claudeOAuthKeychainCredentialAbsent: usage.oauthKeychainCredentialAbsent,
claudeOAuthKeychainCredentialUnavailable: usage.oauthKeychainCredentialUnavailable)
}
func shouldFallback(on _: Error, context: ProviderFetchContext) -> Bool {
// In Auto mode, fall back to the next strategy (cli/web) if OAuth fails (e.g. user cancels keychain prompt
// or auth breaks).
context.runtime == .app && context.sourceMode == .auto
}
fileprivate static func snapshot(from usage: ClaudeUsageSnapshot) -> UsageSnapshot {
let identity = ProviderIdentitySnapshot(
providerID: .claude,
accountEmail: usage.accountEmail,
accountOrganization: usage.accountOrganization,
loginMethod: usage.loginMethod)
let primary = usage.primaryWindowKind == .spendLimit ? nil : usage.primary
return UsageSnapshot(
primary: primary,
secondary: usage.secondary,
tertiary: usage.opus,
extraRateWindows: usage.extraRateWindows.isEmpty ? nil : usage.extraRateWindows,
providerCost: usage.providerCost,
updatedAt: usage.updatedAt,
identity: identity)
}
static func _snapshotForTesting(from usage: ClaudeUsageSnapshot) -> UsageSnapshot {
self.snapshot(from: usage)
}
}
private final class ClaudeWebFetchDeadlineState: @unchecked Sendable {
private let lock = NSLock()
private var deadline: ContinuousClock.Instant?
func remainingBudget(
fullBudget: Duration,
now: ContinuousClock.Instant) -> Duration
{
self.lock.lock()
defer { self.lock.unlock() }
let deadline: ContinuousClock.Instant
if let existingDeadline = self.deadline {
deadline = existingDeadline
} else {
deadline = now.advanced(by: fullBudget)
self.deadline = deadline
}
return max(.zero, now.duration(to: deadline))
}
}
struct ClaudeWebFetchStrategy: ProviderFetchStrategy {
typealias UsageLoader = @Sendable (ProviderFetchContext) async throws -> ClaudeUsageSnapshot
typealias DeadlineNow = @Sendable () -> ContinuousClock.Instant
#if DEBUG
@TaskLocal static var availabilityProbeOverrideForTesting:
(@Sendable (ProviderFetchContext, BrowserDetection) -> Bool)?
@TaskLocal static var usageLoaderOverrideForTesting: UsageLoader?
#endif
let id: String = "claude.web"
let kind: ProviderFetchKind = .web
let browserDetection: BrowserDetection
private let usageLoader: UsageLoader?
private let deadlineState: ClaudeWebFetchDeadlineState
private let deadlineNow: DeadlineNow
init(
browserDetection: BrowserDetection,
usageLoader: UsageLoader? = nil,
deadlineNow: @escaping DeadlineNow = { ContinuousClock.now })
{
self.browserDetection = browserDetection
self.usageLoader = usageLoader
self.deadlineState = ClaudeWebFetchDeadlineState()
self.deadlineNow = deadlineNow
}
func isAvailable(_ context: ProviderFetchContext) async -> Bool {
let appAutoRemainingBudget: Duration?
if context.runtime == .app, context.sourceMode == .auto {
guard let fullBudget = Self.timeoutDuration(context.webTimeout) else { return false }
let remainingBudget = self.deadlineState.remainingBudget(
fullBudget: fullBudget,
now: self.deadlineNow())
guard remainingBudget > .zero else { return false }
appAutoRemainingBudget = remainingBudget
} else {
appAutoRemainingBudget = nil
}
return switch context.settings?.claude?.cookieSource {
case .off?:
false
case .manual?:
Self.hasManualSessionKey(context: context)
case .auto?, nil:
if let appAutoRemainingBudget {
await Self.hasBrowserSessionKey(context: context, before: appAutoRemainingBudget)
} else {
// Explicit web and CLI auto perform the real browser import inside the bounded fetch.
true
}
}
}
func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult {
let usage = try await self.loadUsage(before: context.webTimeout, context: context)
return self.makeResult(
usage: ClaudeOAuthFetchStrategy.snapshot(from: usage),
sourceLabel: "web")
}
func shouldFallback(on error: Error, context: ProviderFetchContext) -> Bool {
guard context.sourceMode == .auto else { return false }
guard !Task.isCancelled,
!(error is CancellationError),
(error as? URLError)?.code != .cancelled
else {
return false
}
// In CLI runtime auto mode, web comes before CLI so fallback is required.
// In app runtime auto mode, web is terminal and should surface its concrete error.
return context.runtime == .cli
}
fileprivate static func hasManualSessionKey(context: ProviderFetchContext) -> Bool {
ClaudeWebAPIFetcher.hasSessionKey(cookieHeader: self.manualCookieHeader(from: context))
}
fileprivate static func hasBrowserSessionKey(
context: ProviderFetchContext,
before timeout: Duration) async -> Bool
{
let browserDetection = context.browserDetection
let sourceTask = Task<Bool, Error> {
#if DEBUG
if let override = Self.availabilityProbeOverrideForTesting {
return override(context, browserDetection)
}
#endif
return ClaudeWebAPIFetcher.hasSessionKey(browserDetection: browserDetection)
}
let race = BoundedTaskJoin(sourceTask: sourceTask)
switch await race.value(joinGrace: timeout) {
case let .value(isAvailable):
return isAvailable
case .failure, .timedOut:
return false
}
}
private static func manualCookieHeader(from context: ProviderFetchContext) -> String? {
guard context.settings?.claude?.cookieSource == .manual else { return nil }
return CookieHeaderNormalizer.normalize(context.settings?.claude?.manualCookieHeader)
}
private func loadUsage(
before timeout: TimeInterval,
context: ProviderFetchContext) async throws -> ClaudeUsageSnapshot
{
guard let timeoutDuration = Self.timeoutDuration(timeout) else {
throw ClaudeWebFetchStrategyError.invalidTimeout
}
let remainingBudget = self.deadlineState.remainingBudget(
fullBudget: timeoutDuration,
now: self.deadlineNow())
guard remainingBudget > .zero else {
try Task.checkCancellation()
throw ClaudeWebFetchStrategyError.timedOut(seconds: timeout)
}
let sourceTask = Task<ClaudeUsageSnapshot, Error> {
if let usageLoader = self.usageLoader {
return try await usageLoader(context)
}
#if DEBUG
if let usageLoader = Self.usageLoaderOverrideForTesting {
return try await usageLoader(context)
}
#endif
let fetcher = ClaudeUsageFetcher(
browserDetection: self.browserDetection,
dataSource: .web,
useWebExtras: false,
manualCookieHeader: Self.manualCookieHeader(from: context),
webOrganizationID: context.settings?.claude?.organizationID)
return try await fetcher.loadLatestUsage(model: "sonnet")
}
let race = BoundedTaskJoin(sourceTask: sourceTask)
switch await race.value(joinGrace: remainingBudget) {
case let .value(usage):
try Task.checkCancellation()
return usage
case let .failure(error):
throw error
case .timedOut:
try Task.checkCancellation()
throw ClaudeWebFetchStrategyError.timedOut(seconds: timeout)
}
}
private static func timeoutDuration(_ timeout: TimeInterval) -> Duration? {
guard timeout.isFinite,
timeout >= 0,
timeout <= TimeInterval(Int64.max)
else {
return nil
}
return .seconds(timeout)
}
fileprivate static var isSupportedOnCurrentPlatform: Bool {
#if os(macOS)
true
#else
false
#endif
}
}
public enum ClaudeWebFetchStrategyError: LocalizedError, Equatable, Sendable {
case invalidTimeout
case timedOut(seconds: TimeInterval)
public var errorDescription: String? {
switch self {
case .invalidTimeout:
"Claude web usage fetch timeout must be a finite, nonnegative value within the supported range."
case let .timedOut(seconds):
"Claude web usage fetch timed out after \(seconds.formatted()) seconds."
}
}
}
struct ClaudeCLIFetchStrategy: ProviderFetchStrategy {
let id: String = "claude.cli"
let kind: ProviderFetchKind = .cli
let useWebExtras: Bool
let manualCookieHeader: String?
let browserDetection: BrowserDetection
let hasWebFallback: Bool
func isAvailable(_ context: ProviderFetchContext) async -> Bool {
// The interactive Claude REPL can open browser OAuth when it starts logged out. CLI runtime and background
// app Auto refreshes must establish authentication through the non-interactive status command first.
let requiresAuthPreflight = context.runtime == .cli || (
context.runtime == .app &&
context.sourceMode == .auto &&
ProviderInteractionContext.current == .background)
guard requiresAuthPreflight else { return true }
guard let binary = ClaudeCLIResolver.resolvedBinaryPath(environment: context.env) else { return false }
return await ClaudeCLIAuthStatusProbe.isLoggedIn(binary: binary, environment: context.env)
}
func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult {
let keepAlive = context.settings?.debugKeepCLISessionsAlive ?? false
let fetcher = ClaudeUsageFetcher(
browserDetection: browserDetection,
environment: context.env,
dataSource: .cli,
useWebExtras: self.useWebExtras,
manualCookieHeader: self.manualCookieHeader,
webOrganizationID: context.settings?.claude?.organizationID,
keepCLISessionsAlive: keepAlive)
let usage = try await fetcher.loadLatestUsage(model: "sonnet")
return self.makeResult(
usage: ClaudeOAuthFetchStrategy.snapshot(from: usage),
sourceLabel: "claude")
}
func shouldFallback(on error: Error, context: ProviderFetchContext) -> Bool {
guard context.runtime == .app, context.sourceMode == .auto else { return false }
guard !ClaudeStatusProbe.isSubscriptionQuotaUnavailableDescription(error.localizedDescription) else {
return false
}
// Reuse the bounded planning result instead of repeating browser/Keychain work after CLI failure.
return self.hasWebFallback
}
}
@@ -0,0 +1,62 @@
import Foundation
/// Shared mapping for the model-scoped weekly limits returned by both Claude usage APIs.
enum ClaudeScopedWeeklyLimitMapper {
struct Limit {
let kind: String?
let group: String?
let percent: Double?
let resetsAt: Date?
let modelID: String?
let modelName: String?
}
static func extraRateWindows(
from limits: [Limit]?,
resetDescription: ((Date) -> String)? = nil) -> [NamedRateWindow]
{
guard let limits else { return [] }
var seenIDs: Set<String> = []
return limits.compactMap { limit in
guard limit.group == "weekly", limit.kind == "weekly_scoped" else { return nil }
guard let percent = limit.percent, percent.isFinite else { return nil }
guard let modelName = self.nonEmpty(limit.modelName) else { return nil }
let identity = self.nonEmpty(limit.modelID) ?? modelName
let slug = self.slug(identity)
guard !slug.isEmpty else { return nil }
let id = "claude-weekly-scoped-\(slug)"
guard seenIDs.insert(id).inserted else { return nil }
return NamedRateWindow(
id: id,
title: "\(modelName) only",
window: RateWindow(
usedPercent: percent,
windowMinutes: 7 * 24 * 60,
resetsAt: limit.resetsAt,
resetDescription: limit.resetsAt.flatMap { resetDescription?($0) }))
}
}
private static func nonEmpty(_ value: String?) -> String? {
let trimmed = value?.trimmingCharacters(in: .whitespacesAndNewlines)
return if let trimmed, !trimmed.isEmpty { trimmed } else { nil }
}
private static func slug(_ value: String) -> String {
var result = ""
var lastWasDash = false
for scalar in value.lowercased().unicodeScalars {
if CharacterSet.alphanumerics.contains(scalar) {
result.unicodeScalars.append(scalar)
lastWasDash = false
} else if !lastWasDash {
result.append("-")
lastWasDash = true
}
}
return result.trimmingCharacters(in: CharacterSet(charactersIn: "-"))
}
}
@@ -0,0 +1,224 @@
import Foundation
public struct ClaudeSourcePlanningInput: Equatable, Sendable {
public let runtime: ProviderRuntime
public let selectedDataSource: ClaudeUsageDataSource
public let webExtrasEnabled: Bool
public let hasWebSession: Bool
public let hasCLI: Bool
public let hasOAuthCredentials: Bool
public init(
runtime: ProviderRuntime,
selectedDataSource: ClaudeUsageDataSource,
webExtrasEnabled: Bool,
hasWebSession: Bool,
hasCLI: Bool,
hasOAuthCredentials: Bool)
{
self.runtime = runtime
self.selectedDataSource = selectedDataSource
self.webExtrasEnabled = webExtrasEnabled
self.hasWebSession = hasWebSession
self.hasCLI = hasCLI
self.hasOAuthCredentials = hasOAuthCredentials
}
}
public enum ClaudeSourcePlanReason: String, Equatable, Sendable {
case explicitSourceSelection = "explicit-source-selection"
case appAutoPreferredOAuth = "app-auto-preferred-oauth"
case appAutoFallbackCLI = "app-auto-fallback-cli"
case appAutoFallbackWeb = "app-auto-fallback-web"
case cliAutoPreferredWeb = "cli-auto-preferred-web"
case cliAutoFallbackCLI = "cli-auto-fallback-cli"
}
public struct ClaudeFetchPlanStep: Equatable, Sendable {
public let dataSource: ClaudeUsageDataSource
public let inclusionReason: ClaudeSourcePlanReason
public let isPlausiblyAvailable: Bool
public init(
dataSource: ClaudeUsageDataSource,
inclusionReason: ClaudeSourcePlanReason,
isPlausiblyAvailable: Bool)
{
self.dataSource = dataSource
self.inclusionReason = inclusionReason
self.isPlausiblyAvailable = isPlausiblyAvailable
}
}
public struct ClaudeFetchPlan: Equatable, Sendable {
public let input: ClaudeSourcePlanningInput
public let orderedSteps: [ClaudeFetchPlanStep]
public init(input: ClaudeSourcePlanningInput, orderedSteps: [ClaudeFetchPlanStep]) {
self.input = input
self.orderedSteps = orderedSteps
}
public var availableSteps: [ClaudeFetchPlanStep] {
self.orderedSteps.filter(\.isPlausiblyAvailable)
}
public var isNoSourceAvailable: Bool {
self.availableSteps.isEmpty
}
public var preferredStep: ClaudeFetchPlanStep? {
switch self.input.selectedDataSource {
case .auto:
self.availableSteps.first
case .api, .oauth, .web, .cli:
self.orderedSteps.first
}
}
public var executionSteps: [ClaudeFetchPlanStep] {
switch self.input.selectedDataSource {
case .auto:
self.availableSteps
case .api, .oauth, .web, .cli:
self.orderedSteps
}
}
public var compatibilityStrategy: ClaudeUsageStrategy? {
guard let preferredStep else { return nil }
let useWebExtras = self.input.runtime == .app
&& preferredStep.dataSource == .cli
&& self.input.webExtrasEnabled
return ClaudeUsageStrategy(
dataSource: preferredStep.dataSource,
useWebExtras: useWebExtras)
}
public var orderLabel: String {
self.orderedSteps.map(\.dataSource.sourceLabel).joined(separator: "")
}
public func debugLines() -> [String] {
var lines = ["planner_order=\(self.orderLabel)"]
lines.append("planner_selected=\(self.preferredStep?.dataSource.rawValue ?? "none")")
lines.append("planner_no_source=\(self.isNoSourceAvailable)")
for step in self.orderedSteps {
let availability = step.isPlausiblyAvailable ? "available" : "unavailable"
lines.append(
"planner_step.\(step.dataSource.rawValue)=\(availability) reason=\(step.inclusionReason.rawValue)")
}
return lines
}
}
public enum ClaudeCLIResolver {
#if DEBUG
@TaskLocal static var resolvedBinaryPathOverrideForTesting: String?
public static var currentResolvedBinaryPathOverrideForTesting: String? {
self.resolvedBinaryPathOverrideForTesting
}
public static func withResolvedBinaryPathOverrideForTesting<T>(
_ path: String?,
operation: () async throws -> T) async rethrows -> T
{
try await self.$resolvedBinaryPathOverrideForTesting.withValue(path) {
try await operation()
}
}
#endif
public static func resolvedBinaryPath(
environment: [String: String] = ProcessInfo.processInfo.environment,
loginPATH: [String]? = LoginShellPathCache.shared.current)
-> String?
{
#if DEBUG
if let override = self.resolvedBinaryPathOverrideForTesting {
return FileManager.default.isExecutableFile(atPath: override) ? override : nil
}
#endif
var normalizedEnvironment = environment
if let override = environment["CLAUDE_CLI_PATH"]?.trimmingCharacters(in: .whitespacesAndNewlines) {
if override.isEmpty {
normalizedEnvironment.removeValue(forKey: "CLAUDE_CLI_PATH")
} else {
normalizedEnvironment["CLAUDE_CLI_PATH"] = override
if FileManager.default.isExecutableFile(atPath: override) {
return override
}
}
}
return BinaryLocator.resolveClaudeBinary(
env: normalizedEnvironment,
loginPATH: loginPATH)
}
public static func isAvailable(environment: [String: String] = ProcessInfo.processInfo.environment) -> Bool {
self.resolvedBinaryPath(environment: environment) != nil
}
}
public enum ClaudeSourcePlanner {
public static func resolve(input: ClaudeSourcePlanningInput) -> ClaudeFetchPlan {
ClaudeFetchPlan(input: input, orderedSteps: self.makeSteps(input: input))
}
private static func makeSteps(input: ClaudeSourcePlanningInput) -> [ClaudeFetchPlanStep] {
switch input.selectedDataSource {
case .auto:
switch input.runtime {
case .app:
[
self.step(.oauth, reason: .appAutoPreferredOAuth, input: input),
self.step(.cli, reason: .appAutoFallbackCLI, input: input),
self.step(.web, reason: .appAutoFallbackWeb, input: input),
]
case .cli:
[
self.step(.web, reason: .cliAutoPreferredWeb, input: input),
self.step(.cli, reason: .cliAutoFallbackCLI, input: input),
]
}
case .api:
[self.step(.api, reason: .explicitSourceSelection, input: input)]
case .oauth:
[self.step(.oauth, reason: .explicitSourceSelection, input: input)]
case .web:
[self.step(.web, reason: .explicitSourceSelection, input: input)]
case .cli:
[self.step(.cli, reason: .explicitSourceSelection, input: input)]
}
}
private static func step(
_ dataSource: ClaudeUsageDataSource,
reason: ClaudeSourcePlanReason,
input: ClaudeSourcePlanningInput) -> ClaudeFetchPlanStep
{
ClaudeFetchPlanStep(
dataSource: dataSource,
inclusionReason: reason,
isPlausiblyAvailable: self.isPlausiblyAvailable(dataSource, input: input))
}
private static func isPlausiblyAvailable(
_ dataSource: ClaudeUsageDataSource,
input: ClaudeSourcePlanningInput) -> Bool
{
switch dataSource {
case .auto, .api:
false
case .oauth:
input.hasOAuthCredentials
case .web:
input.hasWebSession
case .cli:
input.hasCLI
}
}
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,216 @@
import CoreFoundation
import Foundation
/// Strictly parsed result of `cswap --list --json` (schema v1).
///
/// Only the fields allow-listed in `docs/claude-multi-account-and-status-items.md`
/// are decoded: slot number, display email, active state, usage status, and the
/// 5-hour/7-day windows (percent + reset timestamp). Everything else in the
/// payload is ignored; unknown schema versions and partial top-level shapes are
/// rejected.
public struct ClaudeSwapAccountList: Equatable, Sendable {
public let activeAccountNumber: Int?
public let accounts: [ClaudeSwapAccountRow]
public init(activeAccountNumber: Int?, accounts: [ClaudeSwapAccountRow]) {
self.activeAccountNumber = activeAccountNumber
self.accounts = accounts
}
}
public struct ClaudeSwapAccountRow: Equatable, Sendable {
public let number: Int
/// Display-only sensitive value; never logged or persisted.
public let email: String
public let isActive: Bool
public let usageStatus: ClaudeSwapUsageStatus
public let fiveHour: ClaudeSwapUsageWindow?
public let sevenDay: ClaudeSwapUsageWindow?
public init(
number: Int,
email: String,
isActive: Bool,
usageStatus: ClaudeSwapUsageStatus,
fiveHour: ClaudeSwapUsageWindow?,
sevenDay: ClaudeSwapUsageWindow?)
{
self.number = number
self.email = email
self.isActive = isActive
self.usageStatus = usageStatus
self.fiveHour = fiveHour
self.sevenDay = sevenDay
}
}
public struct ClaudeSwapUsageWindow: Equatable, Sendable {
public let usedPercent: Double
public let resetsAt: Date?
public init(usedPercent: Double, resetsAt: Date?) {
self.usedPercent = usedPercent
self.resetsAt = resetsAt
}
}
/// The sentinel statuses `cswap` emits per account. Unknown values from newer
/// claude-swap releases are preserved rather than failing the whole payload.
public enum ClaudeSwapUsageStatus: Equatable, Sendable {
case ok
case tokenExpired
case apiKey
case keychainUnavailable
case noCredentials
case unavailable
case unknown(String)
init(rawValue: String) {
switch rawValue {
case "ok": self = .ok
case "token_expired": self = .tokenExpired
case "api_key": self = .apiKey
case "keychain_unavailable": self = .keychainUnavailable
case "no_credentials": self = .noCredentials
case "unavailable": self = .unavailable
default: self = .unknown(rawValue)
}
}
}
public enum ClaudeSwapListParserError: LocalizedError, Equatable, Sendable {
case notJSONObject
case missingSchemaVersion
case unsupportedSchemaVersion(Int)
case reportedError(type: String, message: String)
case malformedShape(String)
public var errorDescription: String? {
switch self {
case .notJSONObject:
"claude-swap returned output that is not a JSON object."
case .missingSchemaVersion:
"claude-swap output has no schemaVersion field."
case let .unsupportedSchemaVersion(version):
"claude-swap output uses unsupported schema version \(version); CodexBar supports version 1."
case let .reportedError(type, message):
"claude-swap reported \(type): \(message)"
case let .malformedShape(details):
"claude-swap output is malformed: \(details)"
}
}
}
public enum ClaudeSwapListParser {
public static let supportedSchemaVersion = 1
public static func parse(_ data: Data) throws -> ClaudeSwapAccountList {
let raw: Any
do {
raw = try JSONSerialization.jsonObject(with: data)
} catch {
throw ClaudeSwapListParserError.notJSONObject
}
guard let object = raw as? [String: Any] else {
throw ClaudeSwapListParserError.notJSONObject
}
guard let schemaVersion = object["schemaVersion"] as? Int else {
throw ClaudeSwapListParserError.missingSchemaVersion
}
guard schemaVersion == self.supportedSchemaVersion else {
throw ClaudeSwapListParserError.unsupportedSchemaVersion(schemaVersion)
}
if let errorObject = object["error"] as? [String: Any] {
throw ClaudeSwapListParserError.reportedError(
type: errorObject["type"] as? String ?? "Error",
message: errorObject["message"] as? String ?? "unknown error")
}
guard let rawAccounts = object["accounts"] as? [Any] else {
throw ClaudeSwapListParserError.malformedShape("missing accounts array")
}
guard let rawActiveAccountNumber = object["activeAccountNumber"] else {
throw ClaudeSwapListParserError.malformedShape("missing activeAccountNumber")
}
let activeAccountNumber: Int? = switch rawActiveAccountNumber {
case is NSNull: nil
case let number as Int where number > 0: number
default:
throw ClaudeSwapListParserError.malformedShape("activeAccountNumber is not a numeric slot or null")
}
var seenSlots: Set<Int> = []
let accounts = try rawAccounts.map { rawRow -> ClaudeSwapAccountRow in
guard let row = rawRow as? [String: Any] else {
throw ClaudeSwapListParserError.malformedShape("account row is not an object")
}
let account = try self.parseRow(row)
guard seenSlots.insert(account.number).inserted else {
throw ClaudeSwapListParserError.malformedShape("duplicate account slot \(account.number)")
}
return account
}
let activeSlots = accounts.filter(\.isActive).map(\.number)
guard activeSlots == (activeAccountNumber.map { [$0] } ?? []) else {
throw ClaudeSwapListParserError.malformedShape("active account fields disagree")
}
return ClaudeSwapAccountList(activeAccountNumber: activeAccountNumber, accounts: accounts)
}
private static func parseRow(_ row: [String: Any]) throws -> ClaudeSwapAccountRow {
guard let number = row["number"] as? Int else {
throw ClaudeSwapListParserError.malformedShape("account row has no numeric slot")
}
guard number > 0 else {
throw ClaudeSwapListParserError.malformedShape("account slot must be positive")
}
guard let isActive = row["active"] as? Bool else {
throw ClaudeSwapListParserError.malformedShape("slot \(number) has no active flag")
}
guard let rawStatus = row["usageStatus"] as? String else {
throw ClaudeSwapListParserError.malformedShape("slot \(number) has no usageStatus")
}
let email = (row["email"] as? String ?? "").trimmingCharacters(in: .whitespacesAndNewlines)
let usage = row["usage"] as? [String: Any]
return try ClaudeSwapAccountRow(
number: number,
email: email,
isActive: isActive,
usageStatus: ClaudeSwapUsageStatus(rawValue: rawStatus),
fiveHour: self.parseWindow(usage?["fiveHour"], slot: number, name: "fiveHour"),
sevenDay: self.parseWindow(usage?["sevenDay"], slot: number, name: "sevenDay"))
}
private static func parseWindow(_ raw: Any?, slot: Int, name: String) throws -> ClaudeSwapUsageWindow? {
guard let raw else { return nil }
guard let window = raw as? [String: Any] else {
throw ClaudeSwapListParserError.malformedShape("slot \(slot) \(name) window is not an object")
}
guard let pct = Self.finiteDouble(window["pct"]) else {
throw ClaudeSwapListParserError.malformedShape("slot \(slot) \(name) percent is not a finite number")
}
var resetsAt: Date?
if let rawResetsAt = window["resetsAt"] {
guard let text = rawResetsAt as? String, let date = Self.parseTimestamp(text) else {
throw ClaudeSwapListParserError.malformedShape("slot \(slot) \(name) resetsAt is not a timestamp")
}
resetsAt = date
}
return ClaudeSwapUsageWindow(usedPercent: min(max(pct, 0), 100), resetsAt: resetsAt)
}
private static func finiteDouble(_ raw: Any?) -> Double? {
guard let number = raw as? NSNumber else { return nil }
// JSON booleans bridge to NSNumber too; only accept genuine numbers.
guard CFGetTypeID(number) != CFBooleanGetTypeID() else { return nil }
let value = number.doubleValue
return value.isFinite ? value : nil
}
private static func parseTimestamp(_ text: String) -> Date? {
let withFraction = ISO8601DateFormatter()
withFraction.formatOptions = [.withInternetDateTime, .withFractionalSeconds]
if let date = withFraction.date(from: text) { return date }
let plain = ISO8601DateFormatter()
plain.formatOptions = [.withInternetDateTime]
return plain.date(from: text)
}
}
@@ -0,0 +1,102 @@
import Foundation
/// Projects parsed `cswap --list --json` rows into the provider-neutral
/// account snapshot consumed by menus. Identity uses the source-issued numeric
/// slot (`claude-swap:<slot>`), never email or credential-derived values.
public enum ClaudeSwapAccountProjection {
public static let sourceName = "claude-swap"
public static let sourceLabel = "claude-swap"
static let fiveHourWindowMinutes = 5 * 60
static let sevenDayWindowMinutes = 7 * 24 * 60
public static func accountSnapshots(
from list: ClaudeSwapAccountList,
now: Date = Date()) -> [ProviderAccountUsageSnapshot]
{
let ordered = list.accounts.sorted { lhs, rhs in
if lhs.isActive != rhs.isActive { return lhs.isActive }
return lhs.number < rhs.number
}
return ordered.map { row in
ProviderAccountUsageSnapshot(
id: ProviderAccountIdentity(source: self.sourceName, opaqueID: String(row.number)),
provider: .claude,
displayLabel: self.displayLabel(for: row),
isActive: row.isActive,
canActivate: !row.isActive && self.canActivate(row),
snapshot: self.usageSnapshot(for: row, now: now),
error: self.errorText(for: row),
sourceLabel: self.sourceLabel)
}
}
public static func displayError(
accountError: String?,
adapterError: String?,
switchError: String? = nil) -> String?
{
switchError.map { "Account switch failed: \($0)" }
?? accountError
?? adapterError.map { "Showing the last successful update: \($0)" }
}
static func displayLabel(for row: ClaudeSwapAccountRow) -> String {
row.email.isEmpty ? "Account \(row.number)" : row.email
}
private static func usageSnapshot(for row: ClaudeSwapAccountRow, now: Date) -> UsageSnapshot? {
guard row.usageStatus == .ok else { return nil }
let primary = row.fiveHour.map { window in
RateWindow(
usedPercent: window.usedPercent,
windowMinutes: self.fiveHourWindowMinutes,
resetsAt: window.resetsAt,
resetDescription: nil)
}
let secondary = row.sevenDay.map { window in
RateWindow(
usedPercent: window.usedPercent,
windowMinutes: self.sevenDayWindowMinutes,
resetsAt: window.resetsAt,
resetDescription: nil)
}
guard primary != nil || secondary != nil else { return nil }
return UsageSnapshot(
primary: primary,
secondary: secondary,
updatedAt: now,
identity: ProviderIdentitySnapshot(
providerID: .claude,
accountEmail: self.displayLabel(for: row),
accountOrganization: nil,
loginMethod: self.sourceLabel))
}
private static func errorText(for row: ClaudeSwapAccountRow) -> String? {
switch row.usageStatus {
case .ok:
row.fiveHour == nil && row.sevenDay == nil ? "No usage windows reported." : nil
case .tokenExpired:
"Token expired. Switch to this account in claude-swap to refresh it."
case .apiKey:
"API-key account; subscription usage is unavailable."
case .keychainUnavailable:
"claude-swap could not read the active account's Keychain entry."
case .noCredentials:
"No stored credentials for this account slot."
case .unavailable:
"Usage fetch failed."
case let .unknown(raw):
"Unrecognized claude-swap status: \(raw)"
}
}
private static func canActivate(_ row: ClaudeSwapAccountRow) -> Bool {
switch row.usageStatus {
case .ok, .apiKey, .unavailable:
true
case .tokenExpired, .keychainUnavailable, .noCredentials, .unknown:
false
}
}
}
@@ -0,0 +1,141 @@
import Foundation
public enum ClaudeSwapAccountReaderError: LocalizedError, Sendable {
case executablePathNotConfigured
case outputTooLarge(byteCount: Int)
public var errorDescription: String? {
switch self {
case .executablePathNotConfigured:
"No claude-swap executable path is configured."
case let .outputTooLarge(byteCount):
"claude-swap produced \(byteCount) bytes of output; refusing to parse more than " +
"\(ClaudeSwapAccountReader.maxOutputBytes)."
}
}
}
/// Bounded adapter over the external `claude-swap` executable.
///
/// Executes only the fixed list and explicit switch argument arrays (never a
/// shell or config-defined passthrough arguments), per the contracts in
/// `docs/claude-multi-account-and-status-items.md`. CodexBar never reads
/// claude-swap or Claude Code credential storage; the subprocess is solely
/// responsible for its own credential access.
public enum ClaudeSwapAccountReader {
public static let maxOutputBytes = 262_144
public static let defaultTimeout: TimeInterval = 30
public static func readAccountList(
executablePath: String,
timeout: TimeInterval = ClaudeSwapAccountReader.defaultTimeout) async throws -> ClaudeSwapAccountList
{
// Handled claude-swap failures print a schema-v1 error envelope to stdout
// and exit non-zero, so non-zero exits still parse; the parser surfaces
// the envelope as `ClaudeSwapListParserError.reportedError`.
let result = try await self.run(
executablePath: executablePath,
arguments: ["--list", "--json"],
timeout: timeout,
acceptsNonZeroExit: true,
label: "claude-swap list")
return try ClaudeSwapListParser.parse(Data(result.utf8))
}
/// Activates one source-issued numeric slot through claude-swap.
///
/// The external tool owns the credential transaction. CodexBar supplies no
/// credentials and accepts no config-defined arguments.
public static func switchAccount(
executablePath: String,
accountNumber: Int) async throws
-> ClaudeSwapAccountSwitchResult
{
guard accountNumber > 0 else {
throw ClaudeSwapSwitchParserError.malformedShape("requested account slot must be positive")
}
let result = try await self.runCredentialTransaction(
executablePath: executablePath,
arguments: ["--switch-to", String(accountNumber), "--json"],
acceptsNonZeroExit: true,
label: "claude-swap switch")
let parsed = try ClaudeSwapSwitchParser.parse(Data(result.utf8))
guard parsed.toAccountNumber == accountNumber else {
throw ClaudeSwapSwitchParserError.mismatchedTarget(
expected: accountNumber,
actual: parsed.toAccountNumber)
}
return parsed
}
/// Best-effort version probe (`cswap --version` prints `cswap <version>`).
public static func readVersion(
executablePath: String,
timeout: TimeInterval = 10) async -> String?
{
guard let output = try? await self.run(
executablePath: executablePath,
arguments: ["--version"],
timeout: timeout,
label: "claude-swap version")
else {
return nil
}
let trimmed = output.trimmingCharacters(in: .whitespacesAndNewlines)
guard let version = trimmed.split(whereSeparator: \.isWhitespace).last, !version.isEmpty else {
return nil
}
return String(version)
}
public static func resolvedExecutablePath(_ configuredPath: String) throws -> String {
let trimmed = configuredPath.trimmingCharacters(in: .whitespacesAndNewlines)
guard !trimmed.isEmpty else {
throw ClaudeSwapAccountReaderError.executablePathNotConfigured
}
return (trimmed as NSString).expandingTildeInPath
}
private static func run(
executablePath: String,
arguments: [String],
timeout: TimeInterval,
acceptsNonZeroExit: Bool = false,
label: String) async throws -> String
{
try Task.checkCancellation()
let binary = try self.resolvedExecutablePath(executablePath)
let result = try await SubprocessRunner.run(
binary: binary,
arguments: arguments,
environment: ProcessInfo.processInfo.environment,
timeout: timeout,
acceptsNonZeroExit: acceptsNonZeroExit,
label: label)
guard result.stdout.utf8.count <= self.maxOutputBytes else {
throw ClaudeSwapAccountReaderError.outputTooLarge(byteCount: result.stdout.utf8.count)
}
return result.stdout
}
/// Once launched, a credential mutation must reach the external tool's
/// natural exit; caller cancellation and read-probe timeouts must not kill it.
private static func runCredentialTransaction(
executablePath: String,
arguments: [String],
acceptsNonZeroExit: Bool,
label: String) async throws -> String
{
let binary = try self.resolvedExecutablePath(executablePath)
let result = try await SubprocessRunner.runToCompletion(
binary: binary,
arguments: arguments,
environment: ProcessInfo.processInfo.environment,
acceptsNonZeroExit: acceptsNonZeroExit,
label: label)
guard result.stdout.utf8.count <= self.maxOutputBytes else {
throw ClaudeSwapAccountReaderError.outputTooLarge(byteCount: result.stdout.utf8.count)
}
return result.stdout
}
}
@@ -0,0 +1,117 @@
import CoreFoundation
import Foundation
/// Allow-listed result of `cswap --switch-to <slot> --json` (schema v1).
public struct ClaudeSwapAccountSwitchResult: Equatable, Sendable {
public let switched: Bool
public let fromAccountNumber: Int?
public let toAccountNumber: Int
public let reason: String
public init(switched: Bool, fromAccountNumber: Int?, toAccountNumber: Int, reason: String) {
self.switched = switched
self.fromAccountNumber = fromAccountNumber
self.toAccountNumber = toAccountNumber
self.reason = reason
}
}
public enum ClaudeSwapSwitchParserError: LocalizedError, Equatable, Sendable {
case notJSONObject
case missingSchemaVersion
case unsupportedSchemaVersion(Int)
case reportedError(type: String, message: String)
case malformedShape(String)
case mismatchedTarget(expected: Int, actual: Int)
public var errorDescription: String? {
switch self {
case .notJSONObject:
"claude-swap returned switch output that is not a JSON object."
case .missingSchemaVersion:
"claude-swap switch output has no schemaVersion field."
case let .unsupportedSchemaVersion(version):
"claude-swap switch output uses unsupported schema version \(version); CodexBar supports version 1."
case let .reportedError(type, message):
"claude-swap reported \(type): \(message)"
case let .malformedShape(details):
"claude-swap switch output is malformed: \(details)"
case let .mismatchedTarget(expected, actual):
"claude-swap reported account slot \(actual) after CodexBar requested slot \(expected)."
}
}
}
public enum ClaudeSwapSwitchParser {
public static func parse(_ data: Data) throws -> ClaudeSwapAccountSwitchResult {
let raw: Any
do {
raw = try JSONSerialization.jsonObject(with: data)
} catch {
throw ClaudeSwapSwitchParserError.notJSONObject
}
guard let object = raw as? [String: Any] else {
throw ClaudeSwapSwitchParserError.notJSONObject
}
guard let schemaVersion = self.integer(object["schemaVersion"]) else {
throw ClaudeSwapSwitchParserError.missingSchemaVersion
}
guard schemaVersion == ClaudeSwapListParser.supportedSchemaVersion else {
throw ClaudeSwapSwitchParserError.unsupportedSchemaVersion(schemaVersion)
}
if let errorObject = object["error"] as? [String: Any] {
throw ClaudeSwapSwitchParserError.reportedError(
type: errorObject["type"] as? String ?? "Error",
message: errorObject["message"] as? String ?? "unknown error")
}
guard let switched = object["switched"] as? Bool else {
throw ClaudeSwapSwitchParserError.malformedShape("missing switched flag")
}
guard let reason = (object["reason"] as? String)?.trimmingCharacters(in: .whitespacesAndNewlines),
!reason.isEmpty
else {
throw ClaudeSwapSwitchParserError.malformedShape("missing reason")
}
let fromAccountNumber = try self.accountNumber(
in: object["from"],
field: "from",
allowsNull: true)
guard let toAccountNumber = try self.accountNumber(
in: object["to"],
field: "to",
allowsNull: false)
else {
throw ClaudeSwapSwitchParserError.malformedShape("to account has no numeric slot")
}
return ClaudeSwapAccountSwitchResult(
switched: switched,
fromAccountNumber: fromAccountNumber,
toAccountNumber: toAccountNumber,
reason: reason)
}
private static func accountNumber(in raw: Any?, field: String, allowsNull: Bool) throws -> Int? {
if raw is NSNull, allowsNull { return nil }
guard let account = raw as? [String: Any] else {
throw ClaudeSwapSwitchParserError.malformedShape("missing \(field) account")
}
guard let rawNumber = account["number"] else {
throw ClaudeSwapSwitchParserError.malformedShape("\(field) account has no number")
}
if rawNumber is NSNull, allowsNull { return nil }
guard let number = self.integer(rawNumber), number > 0
else {
throw ClaudeSwapSwitchParserError.malformedShape("\(field) account number is not a positive slot")
}
return number
}
private static func integer(_ raw: Any?) -> Int? {
guard let number = raw as? NSNumber,
CFGetTypeID(number) != CFBooleanGetTypeID()
else {
return nil
}
return raw as? Int
}
}
@@ -0,0 +1,38 @@
import Foundation
public enum ClaudeUsageDataSource: String, CaseIterable, Identifiable, Sendable {
case auto
case api
case oauth
case web
case cli
public var id: String {
self.rawValue
}
public var displayName: String {
switch self {
case .auto: "Auto"
case .api: "API (Admin key)"
case .oauth: "OAuth API"
case .web: "Web API (cookies)"
case .cli: "CLI (PTY)"
}
}
public var sourceLabel: String {
switch self {
case .auto:
"auto"
case .api:
"api"
case .oauth:
"oauth"
case .web:
"web"
case .cli:
"cli"
}
}
}
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,123 @@
import Foundation
enum ClaudeWebExtraRateWindowParser {
private static let definitions: [(id: String, title: String, keys: [String])] = [
(
id: "claude-routines",
title: "Daily Routines",
keys: [
"seven_day_routines",
"seven_day_claude_routines",
"claude_routines",
"routines",
"routine",
"seven_day_cowork",
"cowork",
]),
]
static func parse(from json: [String: Any]) -> (windows: [NamedRateWindow], sourceKeys: [String: String]) {
var windows: [NamedRateWindow] = []
var sourceKeys: [String: String] = [:]
windows.reserveCapacity(Self.definitions.count)
for definition in Self.definitions {
if let foundWindow = Self.firstUsageWindow(in: json, keys: definition.keys) {
let rawWindow = foundWindow.window
guard let utilization = Self.percentValue(from: rawWindow["utilization"]) else { continue }
let resetsAt = (rawWindow["resets_at"] as? String).flatMap(Self.parseISO8601Date)
windows.append(Self.namedWindow(
id: definition.id,
title: definition.title,
usedPercent: utilization,
resetsAt: resetsAt))
sourceKeys[definition.id] = foundWindow.sourceKey
continue
}
// Some accounts expose the key with null payloads (for example `seven_day_cowork: null`).
// Preserve the bar in that case with a 0% window so the product section remains visible.
if let key = Self.firstUsageKey(in: json, keys: definition.keys) {
windows.append(Self.namedWindow(
id: definition.id,
title: definition.title,
usedPercent: 0,
resetsAt: nil))
sourceKeys[definition.id] = key
}
}
windows.append(contentsOf: Self.scopedWeeklyLimitWindows(from: json))
return (windows, sourceKeys)
}
private static func scopedWeeklyLimitWindows(from json: [String: Any]) -> [NamedRateWindow] {
guard let limits = json["limits"] as? [[String: Any]] else { return [] }
let mappedLimits = limits.map { entry in
let scope = entry["scope"] as? [String: Any]
let model = scope?["model"] as? [String: Any]
return ClaudeScopedWeeklyLimitMapper.Limit(
kind: entry["kind"] as? String,
group: entry["group"] as? String,
percent: Self.percentValue(from: entry["percent"]),
resetsAt: (entry["resets_at"] as? String).flatMap(Self.parseISO8601Date),
modelID: model?["id"] as? String,
modelName: model?["display_name"] as? String)
}
return ClaudeScopedWeeklyLimitMapper.extraRateWindows(from: mappedLimits)
}
private static func namedWindow(
id: String,
title: String,
usedPercent: Double,
resetsAt: Date?) -> NamedRateWindow
{
NamedRateWindow(
id: id,
title: title,
window: RateWindow(
usedPercent: usedPercent,
windowMinutes: 7 * 24 * 60,
resetsAt: resetsAt,
resetDescription: nil))
}
private static func firstUsageWindow(
in json: [String: Any],
keys: [String]) -> (window: [String: Any], sourceKey: String)?
{
for key in keys {
if let window = json[key] as? [String: Any] {
return (window, key)
}
}
return nil
}
private static func firstUsageKey(in json: [String: Any], keys: [String]) -> String? {
for key in keys where json.keys.contains(key) {
return key
}
return nil
}
private static func percentValue(from value: Any?) -> Double? {
if let intValue = value as? Int {
return Double(intValue)
}
if let doubleValue = value as? Double {
return doubleValue
}
return nil
}
private static func parseISO8601Date(_ string: String) -> Date? {
let formatter = ISO8601DateFormatter()
formatter.formatOptions = [.withInternetDateTime, .withFractionalSeconds]
if let date = formatter.date(from: string) {
return date
}
formatter.formatOptions = [.withInternetDateTime]
return formatter.date(from: string)
}
}
@@ -0,0 +1,62 @@
import Foundation
public enum ClawRouterProviderDescriptor {
public static let descriptor: ProviderDescriptor = Self.makeDescriptor()
static func makeDescriptor() -> ProviderDescriptor {
ProviderDescriptor(
id: .clawrouter,
metadata: ProviderMetadata(
id: .clawrouter,
displayName: "ClawRouter",
sessionLabel: "Monthly budget",
weeklyLabel: "Requests",
opusLabel: nil,
supportsOpus: false,
supportsCredits: false,
creditsHint: "",
toggleTitle: "Show ClawRouter usage",
cliName: "clawrouter",
defaultEnabled: false,
dashboardURL: "https://clawrouter.openclaw.ai/dashboard/access",
statusPageURL: nil),
branding: ProviderBranding(
iconStyle: .clawrouter,
iconResourceName: "ProviderIcon-clawrouter",
color: ProviderColor(red: 89 / 255, green: 110 / 255, blue: 246 / 255)),
tokenCost: ProviderTokenCostConfig(
supportsTokenCost: false,
noDataMessage: { "ClawRouter spend is reported by its usage API." }),
fetchPlan: ProviderFetchPlan(
sourceModes: [.auto, .api],
pipeline: ProviderFetchPipeline(resolveStrategies: { _ in [ClawRouterAPIFetchStrategy()] })),
cli: ProviderCLIConfig(
name: "clawrouter",
aliases: ["claw-router"],
versionDetector: nil))
}
}
struct ClawRouterAPIFetchStrategy: ProviderFetchStrategy {
let id = "clawrouter.api"
let kind: ProviderFetchKind = .apiToken
func isAvailable(_ context: ProviderFetchContext) async -> Bool {
ProviderTokenResolver.clawRouterToken(environment: context.env) != nil
}
func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult {
guard let apiKey = ProviderTokenResolver.clawRouterToken(environment: context.env) else {
throw ClawRouterUsageError.missingCredentials
}
try ClawRouterSettingsReader.validateEndpointOverride(environment: context.env)
let usage = try await ClawRouterUsageFetcher.fetchUsage(
apiKey: apiKey,
baseURL: ClawRouterSettingsReader.baseURL(environment: context.env))
return self.makeResult(usage: usage.toUsageSnapshot(), sourceLabel: "api")
}
func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool {
false
}
}
@@ -0,0 +1,55 @@
import Foundation
public enum ClawRouterSettingsError: LocalizedError, Equatable, Sendable {
case invalidEndpointOverride(String)
public var errorDescription: String? {
switch self {
case let .invalidEndpointOverride(key):
"ClawRouter endpoint override \(key) is invalid. Use an HTTPS URL without embedded credentials."
}
}
}
public enum ClawRouterSettingsReader {
public static let apiKeyEnvironmentKey = "CLAWROUTER_API_KEY"
public static let baseURLEnvironmentKey = "CLAWROUTER_BASE_URL"
public static let defaultBaseURL = URL(string: "https://clawrouter.openclaw.ai")!
public static func apiKey(
environment: [String: String] = ProcessInfo.processInfo.environment) -> String?
{
self.cleaned(environment[self.apiKeyEnvironmentKey])
}
public static func baseURL(
environment: [String: String] = ProcessInfo.processInfo.environment) -> URL
{
guard let raw = self.cleaned(environment[self.baseURLEnvironmentKey]) else {
return self.defaultBaseURL
}
return ProviderEndpointOverrideValidator.normalizedHTTPSURL(from: raw) ?? self.defaultBaseURL
}
public static func validateEndpointOverride(
environment: [String: String] = ProcessInfo.processInfo.environment) throws
{
guard let raw = self.cleaned(environment[self.baseURLEnvironmentKey]) else { return }
guard ProviderEndpointOverrideValidator.normalizedHTTPSURL(from: raw) != nil else {
throw ClawRouterSettingsError.invalidEndpointOverride(self.baseURLEnvironmentKey)
}
}
static func cleaned(_ raw: String?) -> String? {
guard var value = raw?.trimmingCharacters(in: .whitespacesAndNewlines), !value.isEmpty else {
return nil
}
if (value.hasPrefix("\"") && value.hasSuffix("\"")) ||
(value.hasPrefix("'") && value.hasSuffix("'"))
{
value = String(value.dropFirst().dropLast())
}
value = value.trimmingCharacters(in: .whitespacesAndNewlines)
return value.isEmpty ? nil : value
}
}
@@ -0,0 +1,255 @@
import Foundation
#if canImport(FoundationNetworking)
import FoundationNetworking
#endif
public enum ClawRouterUsageError: LocalizedError, Equatable, Sendable {
case missingCredentials
case invalidCredentials
case apiError(Int)
case parseFailed(String)
public var errorDescription: String? {
switch self {
case .missingCredentials:
"Missing ClawRouter API key. Add one in Settings or set CLAWROUTER_API_KEY."
case .invalidCredentials:
"ClawRouter rejected the API key. Check the key and its policy status."
case let .apiError(statusCode):
"ClawRouter API returned HTTP \(statusCode)."
case let .parseFailed(message):
"Could not parse ClawRouter usage: \(message)"
}
}
}
public struct ClawRouterUsageSnapshot: Codable, Sendable, Equatable {
public struct ProviderSummary: Codable, Sendable, Equatable {
public let provider: String
public let requestCount: Int
public let successCount: Int
public let errorCount: Int
public let totalTokens: Int
public let actualCostUSD: Double
}
public let budgetConfigured: Bool
public let budgetLedger: String
public let budgetLimitUSD: Double?
public let budgetSpentUSD: Double?
public let budgetRemainingUSD: Double?
public let budgetResetsAt: Date?
public let requestCount: Int
public let successCount: Int
public let errorCount: Int
public let inputTokens: Int
public let outputTokens: Int
public let totalTokens: Int
public let actualCostUSD: Double
public let providers: [ProviderSummary]
public let updatedAt: Date
public func toUsageSnapshot() -> UsageSnapshot {
let usedPercent: Double? = if let spent = self.budgetSpentUSD,
let limit = self.budgetLimitUSD,
limit > 0
{
min(100, max(0, spent / limit * 100))
} else {
nil
}
let providerCost: ProviderCostSnapshot? = if let spent = self.budgetSpentUSD,
let limit = self.budgetLimitUSD
{
ProviderCostSnapshot(
used: spent,
limit: limit,
currencyCode: "USD",
period: "This month",
resetsAt: self.budgetResetsAt,
updatedAt: self.updatedAt)
} else if self.actualCostUSD > 0 {
ProviderCostSnapshot(
used: self.actualCostUSD,
limit: 0,
currencyCode: "USD",
period: "This month",
resetsAt: self.budgetResetsAt,
updatedAt: self.updatedAt)
} else {
nil
}
return UsageSnapshot(
primary: usedPercent.map {
RateWindow(
usedPercent: $0,
windowMinutes: nil,
resetsAt: self.budgetResetsAt,
resetDescription: nil)
},
secondary: nil,
providerCost: providerCost,
clawRouterUsage: self,
updatedAt: self.updatedAt,
identity: ProviderIdentitySnapshot(
providerID: .clawrouter,
accountEmail: nil,
accountOrganization: "\(self.providers.count) routed providers",
loginMethod: self.budgetConfigured ? "Managed monthly budget" : "Unmetered"),
dataConfidence: .exact)
}
}
private struct ClawRouterUsageResponse: Decodable {
struct Budget: Decodable {
let configured: Bool
let ledger: String
let windowKey: String?
let limitMicros: Int64?
let spentMicros: Int64?
let remainingMicros: Int64?
}
struct Usage: Decodable {
struct Summary: Decodable {
let requestCount: Int
let successCount: Int
let errorCount: Int
let inputTokens: Int
let outputTokens: Int
let totalTokens: Int
let actualCostMicros: Int64
}
struct Provider: Decodable {
let provider: String
let requestCount: Int
let successCount: Int
let errorCount: Int
let totalTokens: Int
let actualCostMicros: Int64
}
let summary: Summary
let providers: [Provider]
}
let budget: Budget
let usage: Usage
}
public enum ClawRouterUsageFetcher {
public static func fetchUsage(
apiKey: String,
baseURL: URL,
transport: any ProviderHTTPTransport = ProviderHTTPClient.shared,
updatedAt: Date = Date()) async throws -> ClawRouterUsageSnapshot
{
guard !apiKey.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty else {
throw ClawRouterUsageError.missingCredentials
}
var request = URLRequest(url: self.usageURL(baseURL: baseURL))
request.httpMethod = "GET"
request.timeoutInterval = 15
request.setValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization")
request.setValue("application/json", forHTTPHeaderField: "Accept")
let response = try await transport.response(for: request)
if response.statusCode == 401 || response.statusCode == 403 {
throw ClawRouterUsageError.invalidCredentials
}
guard (200..<300).contains(response.statusCode) else {
throw ClawRouterUsageError.apiError(response.statusCode)
}
return try self.parseSnapshot(data: response.data, updatedAt: updatedAt)
}
public static func _parseSnapshotForTesting(
_ data: Data,
updatedAt: Date) throws -> ClawRouterUsageSnapshot
{
try self.parseSnapshot(data: data, updatedAt: updatedAt)
}
public static func _usageURLForTesting(baseURL: URL) -> URL {
self.usageURL(baseURL: baseURL)
}
private static func usageURL(baseURL: URL) -> URL {
let path = baseURL.path.trimmingCharacters(in: CharacterSet(charactersIn: "/"))
let versionedBaseURL = path.split(separator: "/").last == "v1"
? baseURL
: baseURL.appendingPathComponent("v1")
return versionedBaseURL.appendingPathComponent("usage")
}
private static func parseSnapshot(data: Data, updatedAt: Date) throws -> ClawRouterUsageSnapshot {
do {
let response = try JSONDecoder().decode(ClawRouterUsageResponse.self, from: data)
return ClawRouterUsageSnapshot(
budgetConfigured: response.budget.configured,
budgetLedger: response.budget.ledger,
budgetLimitUSD: self.dollars(response.budget.limitMicros),
budgetSpentUSD: self.dollars(response.budget.spentMicros),
budgetRemainingUSD: self.dollars(response.budget.remainingMicros),
budgetResetsAt: self.nextMonthlyReset(windowKey: response.budget.windowKey),
requestCount: response.usage.summary.requestCount,
successCount: response.usage.summary.successCount,
errorCount: response.usage.summary.errorCount,
inputTokens: response.usage.summary.inputTokens,
outputTokens: response.usage.summary.outputTokens,
totalTokens: response.usage.summary.totalTokens,
actualCostUSD: self.dollars(response.usage.summary.actualCostMicros),
providers: response.usage.providers.map {
ClawRouterUsageSnapshot.ProviderSummary(
provider: $0.provider,
requestCount: $0.requestCount,
successCount: $0.successCount,
errorCount: $0.errorCount,
totalTokens: $0.totalTokens,
actualCostUSD: self.dollars($0.actualCostMicros))
}.sorted {
if $0.actualCostUSD != $1.actualCostUSD { return $0.actualCostUSD > $1.actualCostUSD }
if $0.requestCount != $1.requestCount { return $0.requestCount > $1.requestCount }
return $0.provider < $1.provider
},
updatedAt: updatedAt)
} catch let error as ClawRouterUsageError {
throw error
} catch {
throw ClawRouterUsageError.parseFailed(error.localizedDescription)
}
}
private static func dollars(_ micros: Int64?) -> Double? {
micros.map(self.dollars)
}
private static func dollars(_ micros: Int64) -> Double {
Double(micros) / 1_000_000
}
private static func nextMonthlyReset(windowKey: String?) -> Date? {
guard let suffix = windowKey?.split(separator: "/").last else { return nil }
let components = suffix.split(separator: "-")
guard components.count == 2,
let year = Int(components[0]),
let month = Int(components[1]),
(1...12).contains(month)
else { return nil }
var nextYear = year
var nextMonth = month + 1
if nextMonth == 13 {
nextMonth = 1
nextYear += 1
}
return DateComponents(
calendar: Calendar(identifier: .gregorian),
timeZone: TimeZone(secondsFromGMT: 0),
year: nextYear,
month: nextMonth,
day: 1).date
}
}
@@ -0,0 +1,95 @@
import Foundation
public enum CodebuffProviderDescriptor {
public static let descriptor: ProviderDescriptor = Self.makeDescriptor()
static func makeDescriptor() -> ProviderDescriptor {
ProviderDescriptor(
id: .codebuff,
metadata: ProviderMetadata(
id: .codebuff,
displayName: "Codebuff",
sessionLabel: "Credits",
weeklyLabel: "Weekly",
opusLabel: nil,
supportsOpus: false,
supportsCredits: true,
creditsHint: "Credit balance from the Codebuff API",
toggleTitle: "Show Codebuff usage",
cliName: "codebuff",
defaultEnabled: false,
isPrimaryProvider: false,
usesAccountFallback: false,
browserCookieOrder: nil,
dashboardURL: "https://www.codebuff.com/usage",
statusPageURL: nil,
statusLinkURL: nil),
branding: ProviderBranding(
iconStyle: .codebuff,
iconResourceName: "ProviderIcon-codebuff",
color: ProviderColor(red: 68 / 255, green: 255 / 255, blue: 0 / 255)),
tokenCost: ProviderTokenCostConfig(
supportsTokenCost: false,
noDataMessage: { "Codebuff cost summary is not yet supported." }),
fetchPlan: ProviderFetchPlan(
sourceModes: [.auto, .api],
pipeline: ProviderFetchPipeline(resolveStrategies: { _ in [CodebuffAPIFetchStrategy()] })),
cli: ProviderCLIConfig(
name: "codebuff",
aliases: ["manicode"],
versionDetector: nil))
}
}
struct CodebuffAPIFetchStrategy: ProviderFetchStrategy {
let id: String = "codebuff.api"
let kind: ProviderFetchKind = .apiToken
func isAvailable(_ context: ProviderFetchContext) async -> Bool {
_ = context
// Keep the strategy available so missing-token surfaces as a user-friendly error
// instead of a generic "no strategy" outcome.
return true
}
func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult {
guard let resolution = Self.resolveToken(environment: context.env) else {
throw CodebuffUsageError.missingCredentials
}
let usage = try await CodebuffUsageFetcher.fetchUsage(
apiKey: resolution.token,
environment: context.env,
includeSubscription: Self.shouldFetchSubscription(for: resolution))
return self.makeResult(
usage: usage.toUsageSnapshot(),
sourceLabel: "api")
}
func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool {
false
}
static func shouldFetchSubscription(for resolution: ProviderTokenResolution) -> Bool {
resolution.source == .authFile
}
private static func resolveToken(environment: [String: String]) -> ProviderTokenResolution? {
ProviderTokenResolver.codebuffResolution(environment: environment)
}
}
/// Errors related to Codebuff settings.
public enum CodebuffSettingsError: LocalizedError, Sendable, Equatable {
case missingToken
case invalidEndpointOverride(String)
public var errorDescription: String? {
switch self {
case .missingToken:
"Codebuff API token not configured. Set CODEBUFF_API_KEY or run `codebuff login` to " +
"populate ~/.config/manicode/credentials.json."
case let .invalidEndpointOverride(key):
"Codebuff endpoint override \(key) must use HTTPS or a bare host."
}
}
}
@@ -0,0 +1,86 @@
import Foundation
/// Reads Codebuff settings from the environment or the local credentials file
/// that the `codebuff` CLI (formerly `manicode`) writes when the user logs in.
public enum CodebuffSettingsReader {
/// Environment variable key for the Codebuff API token.
public static let apiTokenKey = "CODEBUFF_API_KEY"
/// Returns the API token from environment if present and non-empty.
public static func apiKey(environment: [String: String] = ProcessInfo.processInfo.environment) -> String? {
self.cleaned(environment[self.apiTokenKey])
}
/// Returns the API base URL, defaulting to the production endpoint.
public static func apiURL(environment: [String: String] = ProcessInfo.processInfo.environment) -> URL {
if let override = self.validAPIURL(environment: environment) {
return override
}
return URL(string: "https://www.codebuff.com")!
}
public static func validateEndpointOverrides(
environment: [String: String] = ProcessInfo.processInfo.environment) throws
{
guard let raw = self.cleaned(environment["CODEBUFF_API_URL"]) else { return }
guard ProviderEndpointOverrideValidator.normalizedHTTPSURL(from: raw) == nil else { return }
throw CodebuffSettingsError.invalidEndpointOverride("CODEBUFF_API_URL")
}
/// Returns the auth token from the local credentials file if present.
public static func authToken(
authFileURL: URL? = nil,
homeDirectory: URL = FileManager.default.homeDirectoryForCurrentUser) -> String?
{
let fileURL = authFileURL ?? self.defaultAuthFileURL(homeDirectory: homeDirectory)
guard let data = try? Data(contentsOf: fileURL) else { return nil }
return self.parseAuthToken(data: data)
}
/// Default on-disk credentials path: `~/.config/manicode/credentials.json`.
static func defaultAuthFileURL(homeDirectory: URL) -> URL {
homeDirectory
.appendingPathComponent(".config", isDirectory: true)
.appendingPathComponent("manicode", isDirectory: true)
.appendingPathComponent("credentials.json", isDirectory: false)
}
static func parseAuthToken(data: Data) -> String? {
guard let payload = try? JSONDecoder().decode(CredentialsFile.self, from: data) else {
return nil
}
return self.cleaned(payload.default?.authToken) ?? self.cleaned(payload.authToken)
}
static func cleaned(_ raw: String?) -> String? {
guard var value = raw?.trimmingCharacters(in: .whitespacesAndNewlines), !value.isEmpty else {
return nil
}
if (value.hasPrefix("\"") && value.hasSuffix("\"")) ||
(value.hasPrefix("'") && value.hasSuffix("'"))
{
value = String(value.dropFirst().dropLast())
}
value = value.trimmingCharacters(in: .whitespacesAndNewlines)
return value.isEmpty ? nil : value
}
private static func validAPIURL(environment: [String: String]) -> URL? {
guard let raw = self.cleaned(environment["CODEBUFF_API_URL"]) else { return nil }
return ProviderEndpointOverrideValidator.normalizedHTTPSURL(from: raw)
}
}
private struct CredentialsFile: Decodable {
let `default`: CredentialsProfile?
let authToken: String?
}
private struct CredentialsProfile: Decodable {
let authToken: String?
let fingerprintId: String?
let email: String?
let name: String?
}
@@ -0,0 +1,40 @@
import Foundation
public enum CodebuffUsageError: LocalizedError, Sendable, Equatable {
case missingCredentials
case unauthorized
case endpointNotFound
case serviceUnavailable(Int)
case apiError(Int)
case networkError(String)
case parseFailed(String)
public static let missingToken: CodebuffUsageError = .missingCredentials
public var errorDescription: String? {
switch self {
case .missingCredentials:
"Codebuff API token not configured. Set CODEBUFF_API_KEY or run `codebuff login` to " +
"populate ~/.config/manicode/credentials.json."
case .unauthorized:
"Unauthorized. Please sign in to Codebuff again."
case .endpointNotFound:
"Codebuff usage endpoint not found."
case let .serviceUnavailable(status):
"Codebuff API is temporarily unavailable (status \(status))."
case let .apiError(status):
"Codebuff API returned an unexpected status (\(status))."
case let .networkError(message):
"Codebuff API error: \(message)"
case let .parseFailed(message):
"Could not parse Codebuff usage: \(message)"
}
}
public var isAuthRelated: Bool {
switch self {
case .unauthorized, .missingCredentials: true
default: false
}
}
}
@@ -0,0 +1,365 @@
import Foundation
#if canImport(FoundationNetworking)
import FoundationNetworking
#endif
/// Fetches live credit balance and subscription details from the Codebuff API.
/// Uses Bearer token auth against the public www.codebuff.com endpoints used by
/// the dashboard + the `codebuff` CLI.
public enum CodebuffUsageFetcher {
private static let requestTimeoutSeconds: TimeInterval = 15
/// Extra grace period to wait for the optional subscription endpoint after the
/// primary usage call returns. Keeps the menu responsive when `/api/user/subscription`
/// is slow or hangs while `/api/v1/usage` succeeds quickly.
private static let subscriptionGraceSeconds: TimeInterval = 2
public static func fetchUsage(
apiKey: String,
environment: [String: String] = ProcessInfo.processInfo.environment,
includeSubscription: Bool = true,
session transport: any ProviderHTTPTransport = ProviderHTTPClient.shared) async throws -> CodebuffUsageSnapshot
{
try await self.fetchUsage(
apiKey: apiKey,
environment: environment,
includeSubscription: includeSubscription,
transport: transport,
subscriptionGrace: .seconds(self.subscriptionGraceSeconds))
}
static func _fetchUsageForTesting(
apiKey: String,
environment: [String: String] = ProcessInfo.processInfo.environment,
includeSubscription: Bool = true,
transport: any ProviderHTTPTransport,
subscriptionGrace: Duration) async throws -> CodebuffUsageSnapshot
{
try await self.fetchUsage(
apiKey: apiKey,
environment: environment,
includeSubscription: includeSubscription,
transport: transport,
subscriptionGrace: subscriptionGrace)
}
private static func fetchUsage(
apiKey: String,
environment: [String: String],
includeSubscription: Bool,
transport: any ProviderHTTPTransport,
subscriptionGrace: Duration) async throws -> CodebuffUsageSnapshot
{
let trimmed = apiKey.trimmingCharacters(in: .whitespacesAndNewlines)
guard !trimmed.isEmpty else {
throw CodebuffUsageError.missingCredentials
}
try CodebuffSettingsReader.validateEndpointOverrides(environment: environment)
let baseURL = CodebuffSettingsReader.apiURL(environment: environment)
let (usageValues, subscriptionValues) = try await self.fetchPayloads(
apiKey: trimmed,
baseURL: baseURL,
includeSubscription: includeSubscription,
transport: transport,
subscriptionGrace: subscriptionGrace)
return CodebuffUsageSnapshot(
creditsUsed: usageValues.used,
creditsTotal: usageValues.total,
creditsRemaining: usageValues.remaining,
weeklyUsed: subscriptionValues?.weeklyUsed,
weeklyLimit: subscriptionValues?.weeklyLimit,
weeklyResetsAt: subscriptionValues?.weeklyResetsAt,
billingPeriodEnd: subscriptionValues?.billingPeriodEnd,
nextQuotaReset: usageValues.nextQuotaReset,
tier: subscriptionValues?.tier,
subscriptionStatus: subscriptionValues?.status,
autoTopUpEnabled: usageValues.autoTopupEnabled,
accountEmail: subscriptionValues?.email,
updatedAt: Date())
}
private static func fetchPayloads(
apiKey: String,
baseURL: URL,
includeSubscription: Bool,
transport: any ProviderHTTPTransport,
subscriptionGrace: Duration) async throws -> (UsagePayload, SubscriptionPayload?)
{
guard includeSubscription else {
return try await (
self.fetchUsagePayload(apiKey: apiKey, baseURL: baseURL, transport: transport),
nil)
}
let subscriptionTask = Task<SubscriptionPayload?, Error> {
try await self.fetchSubscriptionPayload(
apiKey: apiKey,
baseURL: baseURL,
transport: transport)
}
let usageValues: UsagePayload
do {
usageValues = try await withTaskCancellationHandler {
try await self.fetchUsagePayload(
apiKey: apiKey,
baseURL: baseURL,
transport: transport)
} onCancel: {
subscriptionTask.cancel()
}
} catch {
subscriptionTask.cancel()
throw error
}
do {
try Task.checkCancellation()
} catch {
subscriptionTask.cancel()
throw error
}
let race = BoundedTaskJoin(sourceTask: subscriptionTask)
switch await race.value(joinGrace: subscriptionGrace) {
case let .value(subscriptionValues):
try Task.checkCancellation()
return (usageValues, subscriptionValues)
case .timedOut:
try Task.checkCancellation()
return (usageValues, nil)
case .failure:
subscriptionTask.cancel()
try Task.checkCancellation()
return (usageValues, nil)
}
}
// MARK: - Endpoint helpers
struct UsagePayload {
let used: Double?
let total: Double?
let remaining: Double?
let nextQuotaReset: Date?
let autoTopupEnabled: Bool?
}
struct SubscriptionPayload {
let status: String?
let tier: String?
let billingPeriodEnd: Date?
let weeklyUsed: Double?
let weeklyLimit: Double?
let weeklyResetsAt: Date?
let email: String?
}
static func usageURL(baseURL: URL) -> URL {
baseURL.appendingPathComponent("/api/v1/usage")
}
static func subscriptionURL(baseURL: URL) -> URL {
baseURL.appendingPathComponent("/api/user/subscription")
}
static func statusError(for statusCode: Int) -> CodebuffUsageError? {
switch statusCode {
case 401, 403: .unauthorized
case 404: .endpointNotFound
case 500...599: .serviceUnavailable(statusCode)
default: nil
}
}
static func parseUsagePayload(_ data: Data) throws -> UsagePayload {
guard let root = try? JSONSerialization.jsonObject(with: data) as? [String: Any] else {
throw CodebuffUsageError.parseFailed("Invalid JSON")
}
let used = self.double(from: root["usage"]) ?? self.double(from: root["used"])
let total = self.double(from: root["quota"]) ?? self.double(from: root["limit"])
let remaining = self.double(from: root["remainingBalance"]) ?? self.double(from: root["remaining"])
let reset = self.date(from: root["next_quota_reset"])
let autoTopUp = root["autoTopupEnabled"] as? Bool ?? root["auto_topup_enabled"] as? Bool
return UsagePayload(
used: used,
total: total,
remaining: remaining,
nextQuotaReset: reset,
autoTopupEnabled: autoTopUp)
}
static func parseSubscriptionPayload(_ data: Data) throws -> SubscriptionPayload {
guard let root = try? JSONSerialization.jsonObject(with: data) as? [String: Any] else {
throw CodebuffUsageError.parseFailed("Invalid JSON")
}
let subscription = root["subscription"] as? [String: Any]
let rateLimit = root["rateLimit"] as? [String: Any]
let tier = self.string(from: subscription?["displayName"])
?? self.string(from: root["displayName"])
?? self.string(from: subscription?["tier"])
?? self.string(from: root["tier"])
?? self.string(from: subscription?["scheduledTier"])
let status = subscription?["status"] as? String
let email = root["email"] as? String ?? (root["user"] as? [String: Any])?["email"] as? String
let billingPeriodEnd = self.date(from: subscription?["billingPeriodEnd"])
?? self.date(from: subscription?["currentPeriodEnd"])
let weeklyUsed = self.double(from: rateLimit?["weeklyUsed"])
?? self.double(from: rateLimit?["used"])
let weeklyLimit = self.double(from: rateLimit?["weeklyLimit"])
?? self.double(from: rateLimit?["limit"])
let weeklyResetsAt = self.date(from: rateLimit?["weeklyResetsAt"])
return SubscriptionPayload(
status: status,
tier: tier,
billingPeriodEnd: billingPeriodEnd,
weeklyUsed: weeklyUsed,
weeklyLimit: weeklyLimit,
weeklyResetsAt: weeklyResetsAt,
email: email)
}
// MARK: - Test hooks
static func _parseUsagePayloadForTesting(_ data: Data) throws -> UsagePayload {
try self.parseUsagePayload(data)
}
static func _parseSubscriptionPayloadForTesting(_ data: Data) throws -> SubscriptionPayload {
try self.parseSubscriptionPayload(data)
}
static func _statusErrorForTesting(_ statusCode: Int) -> CodebuffUsageError? {
self.statusError(for: statusCode)
}
// MARK: - Networking
private static func fetchUsagePayload(
apiKey: String,
baseURL: URL,
transport: any ProviderHTTPTransport) async throws -> UsagePayload
{
var request = URLRequest(url: self.usageURL(baseURL: baseURL))
request.httpMethod = "POST"
request.timeoutInterval = self.requestTimeoutSeconds
request.setValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization")
request.setValue("application/json", forHTTPHeaderField: "Content-Type")
request.setValue("application/json", forHTTPHeaderField: "Accept")
request.httpBody = try? JSONSerialization.data(withJSONObject: ["fingerprintId": "codexbar-usage"])
let response = try await self.send(request: request, transport: transport)
if let err = self.statusError(for: response.statusCode) {
throw err
}
guard response.statusCode == 200 else {
throw CodebuffUsageError.apiError(response.statusCode)
}
return try self.parseUsagePayload(response.data)
}
private static func fetchSubscriptionPayload(
apiKey: String,
baseURL: URL,
transport: any ProviderHTTPTransport) async throws -> SubscriptionPayload
{
var request = URLRequest(url: self.subscriptionURL(baseURL: baseURL))
request.httpMethod = "GET"
request.timeoutInterval = self.requestTimeoutSeconds
request.setValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization")
request.setValue("application/json", forHTTPHeaderField: "Accept")
let response = try await self.send(request: request, transport: transport)
if let err = self.statusError(for: response.statusCode) {
throw err
}
guard response.statusCode == 200 else {
throw CodebuffUsageError.apiError(response.statusCode)
}
return try self.parseSubscriptionPayload(response.data)
}
private static func send(
request: URLRequest,
transport: any ProviderHTTPTransport) async throws -> ProviderHTTPResponse
{
do {
return try await transport.response(for: request)
} catch let error as CodebuffUsageError {
throw error
} catch let error as URLError where error.code == .badServerResponse {
throw CodebuffUsageError.networkError("Invalid response")
} catch {
throw CodebuffUsageError.networkError(error.localizedDescription)
}
}
// MARK: - Value parsing
private static func double(from value: Any?) -> Double? {
switch value {
case let number as NSNumber:
let raw = number.doubleValue
return raw.isFinite ? raw : nil
case let string as String:
let trimmed = string.trimmingCharacters(in: .whitespacesAndNewlines)
guard !trimmed.isEmpty, let raw = Double(trimmed), raw.isFinite else { return nil }
return raw
default:
return nil
}
}
private static func string(from value: Any?) -> String? {
switch value {
case let string as String:
let trimmed = string.trimmingCharacters(in: .whitespacesAndNewlines)
return trimmed.isEmpty ? nil : trimmed
case let number as NSNumber:
let raw = number.doubleValue
guard raw.isFinite else { return nil }
return number.stringValue
default:
return nil
}
}
private static func date(from value: Any?) -> Date? {
switch value {
case let string as String:
let trimmed = string.trimmingCharacters(in: .whitespacesAndNewlines)
guard !trimmed.isEmpty else { return nil }
let fractional = ISO8601DateFormatter()
fractional.formatOptions = [.withInternetDateTime, .withFractionalSeconds]
if let date = fractional.date(from: trimmed) {
return date
}
let plain = ISO8601DateFormatter()
plain.formatOptions = [.withInternetDateTime]
if let date = plain.date(from: trimmed) {
return date
}
if let interval = Double(trimmed), interval.isFinite {
return Self.dateFromNumeric(interval)
}
return nil
case let number as NSNumber:
let raw = number.doubleValue
return raw.isFinite ? Self.dateFromNumeric(raw) : nil
default:
return nil
}
}
private static func dateFromNumeric(_ value: Double) -> Date? {
if value > 10_000_000_000 {
return Date(timeIntervalSince1970: value / 1000)
}
return Date(timeIntervalSince1970: value)
}
}

Some files were not shown because too many files have changed in this diff Show More