chore: import upstream snapshot with attribution
This commit is contained in:
@@ -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)
|
||||
}
|
||||
}
|
||||
+23
@@ -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
|
||||
}
|
||||
}
|
||||
+137
@@ -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
|
||||
}
|
||||
}
|
||||
+111
@@ -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."
|
||||
}
|
||||
}
|
||||
}
|
||||
+18
@@ -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
|
||||
}
|
||||
}
|
||||
+375
@@ -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)
|
||||
}
|
||||
}
|
||||
+430
@@ -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
+580
@@ -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
|
||||
}
|
||||
}
|
||||
+184
@@ -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
|
||||
+169
@@ -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
|
||||
+177
@@ -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
|
||||
}
|
||||
+31
@@ -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
|
||||
+49
@@ -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
|
||||
}
|
||||
+35
@@ -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
|
||||
+142
@@ -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()
|
||||
}
|
||||
}
|
||||
+371
@@ -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
Reference in New Issue
Block a user