chore: import upstream snapshot with attribution

This commit is contained in:
wehub-resource-sync
2026-07-13 12:22:33 +08:00
commit fc4fcbab58
1848 changed files with 472303 additions and 0 deletions
+445
View File
@@ -0,0 +1,445 @@
import Foundation
public struct AgentSession: Codable, Equatable, Sendable, Identifiable {
public enum Provider: String, Codable, Sendable {
case codex
case claude
}
public enum Source: String, Codable, Sendable {
case cli
case desktopApp
case ide
case unknown
}
public enum State: String, Codable, Sendable {
case active
case idle
}
public var id: String
public var provider: Provider
public var source: Source
public var state: State
public var pid: Int32?
public var cwd: String?
public var projectName: String?
public var startedAt: Date?
public var lastActivityAt: Date?
public var transcriptPath: String?
public var host: String
public init(
id: String,
provider: Provider,
source: Source,
state: State,
pid: Int32?,
cwd: String?,
projectName: String?,
startedAt: Date?,
lastActivityAt: Date?,
transcriptPath: String?,
host: String)
{
self.id = id
self.provider = provider
self.source = source
self.state = state
self.pid = pid
self.cwd = cwd
self.projectName = projectName
self.startedAt = startedAt
self.lastActivityAt = lastActivityAt
self.transcriptPath = transcriptPath
self.host = host
}
}
public struct SessionScanConfig: Equatable, Sendable {
public var activeWindow: TimeInterval
public var fileOnlyWindow: TimeInterval
public init(activeWindow: TimeInterval = 120, fileOnlyWindow: TimeInterval = 30 * 60) {
self.activeWindow = activeWindow
self.fileOnlyWindow = fileOnlyWindow
}
public func state(lastActivityAt: Date?, now: Date, hasLiveProcess: Bool) -> AgentSession.State {
guard let lastActivityAt else { return hasLiveProcess ? .active : .idle }
return now.timeIntervalSince(lastActivityAt) <= self.activeWindow ? .active : .idle
}
}
public struct AgentProcessRecord: Equatable, Sendable {
public let pid: Int32
public let ppid: Int32
public let startedAt: Date?
public let command: String
public init(pid: Int32, ppid: Int32, startedAt: Date?, command: String) {
self.pid = pid
self.ppid = ppid
self.startedAt = startedAt
self.command = command
}
public var executableBasename: String {
let firstToken = self.command.split(whereSeparator: \ .isWhitespace).first.map(String.init) ?? ""
let firstBasename = URL(fileURLWithPath: firstToken).lastPathComponent
if firstBasename == "disclaimer" {
return firstBasename
}
if self.command.contains("Application Support/Claude/claude-code/claude") {
return "claude"
}
return firstBasename
}
}
public enum AgentPSOutputParser {
public static func parse(_ output: String) -> [AgentProcessRecord] {
let formatter = DateFormatter()
formatter.locale = Locale(identifier: "en_US_POSIX")
formatter.dateFormat = "EEE MMM d HH:mm:ss yyyy"
return output.split(whereSeparator: \ .isNewline).compactMap { rawLine -> AgentProcessRecord? in
let fields = rawLine.split(maxSplits: 7, omittingEmptySubsequences: true, whereSeparator: \ .isWhitespace)
guard fields.count == 8,
let pid = Int32(fields[0]),
let ppid = Int32(fields[1])
else { return nil }
let dateText = fields[2...6].joined(separator: " ")
return AgentProcessRecord(
pid: pid,
ppid: ppid,
startedAt: formatter.date(from: dateText),
command: String(fields[7]))
}
}
public static func agentProcesses(from records: [AgentProcessRecord]) -> [AgentProcessRecord] {
let candidates = records.filter { record in
let basename = record.executableBasename.lowercased()
if basename == "codex" {
let arguments = self.arguments(record.command)
return self.isCodexAgentExecutable(record.command) &&
!arguments.contains("app-server") &&
!arguments.contains("--help") &&
!arguments.contains("--version")
}
if basename == "claude" {
return self.isClaudeAgentExecutable(record.command) && !self.isObviousClaudeHelper(record.command)
}
return basename == "disclaimer" && record.command.contains("claude")
}
let recordsByPID = Dictionary(uniqueKeysWithValues: candidates.map { ($0.pid, $0) })
return candidates.filter { record in
guard record.executableBasename.lowercased() == "disclaimer" else { return true }
return !candidates.contains { child in
child.ppid == record.pid &&
child.executableBasename.lowercased() == "claude" &&
self.normalizedClaudeArguments(child.command) == self.normalizedClaudeArguments(record.command)
} && recordsByPID[record.ppid] == nil
}
}
public static func provider(for record: AgentProcessRecord) -> AgentSession.Provider? {
let basename = record.executableBasename.lowercased()
if basename == "codex" {
return .codex
}
if basename == "claude" || basename == "disclaimer" {
return .claude
}
return nil
}
public static func source(for record: AgentProcessRecord) -> AgentSession.Source {
guard self.provider(for: record) == .claude else { return .cli }
return record.command.contains("Application Support/Claude/claude-code") ? .desktopApp : .cli
}
public static func hasCodexAppServer(in records: [AgentProcessRecord]) -> Bool {
records.contains { record in
record.executableBasename.lowercased() == "codex" &&
self.isCodexAgentExecutable(record.command) &&
self.arguments(record.command).contains("app-server")
}
}
private static func arguments(_ command: String) -> [String] {
command.split(whereSeparator: \ .isWhitespace).dropFirst().map(String.init)
}
private static func normalizedClaudeArguments(_ command: String) -> [String] {
let arguments = self.arguments(command)
if let index = arguments.firstIndex(where: { URL(fileURLWithPath: $0).lastPathComponent == "claude" }) {
return Array(arguments.suffix(from: arguments.index(after: index)))
}
return arguments
}
private static func isObviousClaudeHelper(_ command: String) -> Bool {
let lowercased = command.lowercased()
return lowercased.contains("--version") ||
lowercased.contains("--help") ||
lowercased.contains("claude-code-acp")
}
private static func isCodexAgentExecutable(_ command: String) -> Bool {
let lowercased = command.lowercased()
guard lowercased.contains(".app/") else { return true }
return lowercased.hasPrefix("/applications/codex.app/contents/resources/codex ") ||
lowercased.hasPrefix("/applications/codex.app/contents/resources/codex\t")
}
private static func isClaudeAgentExecutable(_ command: String) -> Bool {
let lowercased = command.lowercased()
return !lowercased.contains(".app/") || lowercased.contains("application support/claude/claude-code/claude")
}
}
public enum LSOFCWDOutputParser {
public static func parse(_ output: String) -> [Int32: String] {
var result: [Int32: String] = [:]
var currentPID: Int32?
for line in output.split(whereSeparator: \ .isNewline).map(String.init) {
guard let prefix = line.first else { continue }
switch prefix {
case "p":
currentPID = Int32(line.dropFirst())
case "n":
if let currentPID {
result[currentPID] = String(line.dropFirst())
}
default:
continue
}
}
return result
}
}
public enum ClaudeSessionProjectMapper {
public struct Transcript: Equatable, Sendable {
public let url: URL
public let modifiedAt: Date
public init(url: URL, modifiedAt: Date) {
self.url = url
self.modifiedAt = modifiedAt
}
}
public static func escapedCWD(_ cwd: String) -> String {
String(cwd.unicodeScalars.map { scalar in
let value = scalar.value
let isASCIIAlphanumeric = (48...57).contains(value) ||
(65...90).contains(value) ||
(97...122).contains(value)
return isASCIIAlphanumeric ? Character(scalar) : "-"
})
}
public static func projectDirectories(
cwd: String,
homeDirectory: URL = FileManager.default.homeDirectoryForCurrentUser,
fileManager: FileManager = .default) -> [URL]
{
let standardRoot = homeDirectory
.appendingPathComponent(".claude", isDirectory: true)
.appendingPathComponent("projects", isDirectory: true)
return ([standardRoot] + ClaudeDesktopProjectsLocator.roots(
homeDirectory: homeDirectory,
fileManager: fileManager))
.map { $0.appendingPathComponent(self.escapedCWD(cwd), isDirectory: true) }
}
public static func newestTranscript(
cwd: String,
homeDirectory: URL = FileManager.default.homeDirectoryForCurrentUser,
fileManager: FileManager = .default) -> (url: URL, modifiedAt: Date)?
{
self.transcripts(cwd: cwd, homeDirectory: homeDirectory, fileManager: fileManager).first
.map { (url: $0.url, modifiedAt: $0.modifiedAt) }
}
public static func transcripts(
cwd: String,
homeDirectory: URL = FileManager.default.homeDirectoryForCurrentUser,
fileManager: FileManager = .default) -> [Transcript]
{
self.projectDirectories(cwd: cwd, homeDirectory: homeDirectory, fileManager: fileManager)
.flatMap { directory -> [(URL, Date)] in
guard let files = try? fileManager.contentsOfDirectory(
at: directory,
includingPropertiesForKeys: [.contentModificationDateKey],
options: [.skipsHiddenFiles])
else { return [] }
return files.compactMap { file in
guard file.pathExtension == "jsonl",
let modifiedAt = try? file.resourceValues(
forKeys: [.contentModificationDateKey]).contentModificationDate
else { return nil }
return (file, modifiedAt)
}
}
.sorted { $0.1 > $1.1 }
.map { Transcript(url: $0.0, modifiedAt: $0.1) }
}
}
public enum AgentSessionCorrelation {
public static func newestProcessesFirst(_ processes: [AgentProcessRecord]) -> [AgentProcessRecord] {
processes.sorted { lhs, rhs in
let lhsDate = lhs.startedAt ?? .distantPast
let rhsDate = rhs.startedAt ?? .distantPast
if lhsDate != rhsDate { return lhsDate > rhsDate }
return lhs.pid > rhs.pid
}
}
public static func assignClaudeTranscripts(
processes: [AgentProcessRecord],
cwdByPID: [Int32: String],
transcriptsByCWD: [String: [ClaudeSessionProjectMapper.Transcript]])
-> [Int32: ClaudeSessionProjectMapper.Transcript]
{
var assigned: [Int32: ClaudeSessionProjectMapper.Transcript] = [:]
var usedPaths = Set<String>()
let unambiguousPIDs = self.unambiguousProcessIDs(processes: processes, cwdByPID: cwdByPID)
for process in self.newestProcessesFirst(processes) {
guard unambiguousPIDs.contains(process.pid),
let cwd = cwdByPID[process.pid],
let candidates = transcriptsByCWD[cwd]
else { continue }
let candidate = candidates.first { transcript in
guard !usedPaths.contains(transcript.url.path) else { return false }
guard let startedAt = process.startedAt else { return true }
return transcript.modifiedAt >= startedAt
}
if let candidate {
assigned[process.pid] = candidate
usedPaths.insert(candidate.url.path)
}
}
return assigned
}
public static func unambiguousProcessIDs(
processes: [AgentProcessRecord],
cwdByPID: [Int32: String]) -> Set<Int32>
{
let grouped = Dictionary(grouping: processes.compactMap { process -> (Int32, String)? in
guard let cwd = cwdByPID[process.pid] else { return nil }
return (process.pid, URL(fileURLWithPath: cwd).standardizedFileURL.path)
}, by: { $0.1 })
return Set(grouped.values.compactMap { records in
records.count == 1 ? records[0].0 : nil
})
}
public static func codexWorkingDirectoriesMatch(_ lhs: String?, _ rhs: String?) -> Bool {
guard let lhs, let rhs else { return false }
return URL(fileURLWithPath: lhs).standardizedFileURL.path ==
URL(fileURLWithPath: rhs).standardizedFileURL.path
}
public static func fileOnlyCodexSource(
metadataSource: AgentSession.Source,
appServerPresent: Bool) -> AgentSession.Source
{
metadataSource == .unknown && appServerPresent ? .desktopApp : metadataSource
}
}
public struct CodexRolloutMetadata: Equatable, Sendable {
public let sessionID: String
public let cwd: String?
public let originator: String?
public let source: String?
public init(sessionID: String, cwd: String?, originator: String?, source: String?) {
self.sessionID = sessionID
self.cwd = cwd
self.originator = originator
self.source = source
}
public var sessionSource: AgentSession.Source {
let value = [self.originator, self.source]
.compactMap { $0?.lowercased() }
.joined(separator: " ")
if value.contains("desktop") || value.contains("app-server") {
return .desktopApp
}
if value.contains("ide") || value.contains("vscode") || value.contains("cursor") || value.contains("zed") {
return .ide
}
if value.contains("codex_exec") || value.contains("exec") || value.contains("cli") {
return .cli
}
return .unknown
}
}
public enum CodexRolloutFirstLineParser {
public static func parse(_ line: String) -> CodexRolloutMetadata? {
guard let data = line.data(using: .utf8),
let object = try? JSONSerialization.jsonObject(with: data) as? [String: Any],
object["type"] as? String == "session_meta",
let payload = object["payload"] as? [String: Any],
let sessionID = (payload["session_id"] as? String) ?? (payload["id"] as? String)
else { return nil }
return CodexRolloutMetadata(
sessionID: sessionID,
cwd: payload["cwd"] as? String,
originator: payload["originator"] as? String,
source: payload["source"] as? String)
}
public static func read(from url: URL) -> CodexRolloutMetadata? {
guard let handle = try? FileHandle(forReadingFrom: url) else { return nil }
defer { try? handle.close() }
var data = Data()
while data.count < 256 * 1024 {
guard let chunk = try? handle.read(upToCount: 4096), !chunk.isEmpty else { break }
if let newline = chunk.firstIndex(of: 0x0A) {
data.append(chunk.prefix(upTo: newline))
break
}
data.append(chunk)
}
guard let line = String(data: data, encoding: .utf8) else { return nil }
return self.parse(line)
}
public static func makeSession(
metadata: CodexRolloutMetadata,
transcriptURL: URL,
modifiedAt: Date,
pid: Int32? = nil,
startedAt: Date? = nil,
host: String,
config: SessionScanConfig = SessionScanConfig(),
now: Date = Date()) -> AgentSession?
{
guard pid != nil || now.timeIntervalSince(modifiedAt) <= config.fileOnlyWindow else { return nil }
let cwd = metadata.cwd
return AgentSession(
id: metadata.sessionID,
provider: .codex,
source: metadata.sessionSource,
state: config.state(lastActivityAt: modifiedAt, now: now, hasLiveProcess: pid != nil),
pid: pid,
cwd: cwd,
projectName: cwd.map { URL(fileURLWithPath: $0).lastPathComponent },
startedAt: startedAt,
lastActivityAt: modifiedAt,
transcriptPath: transcriptURL.path,
host: host)
}
}
+247
View File
@@ -0,0 +1,247 @@
import Foundation
#if os(macOS)
import Security
#endif
public enum AppGroupSupport {
public static let defaultTeamID = "Y5PE65HELJ"
public static let teamIDInfoKey = "CodexBarTeamID"
public static let legacyReleaseGroupID = "group.com.steipete.codexbar"
public static let legacyDebugGroupID = "group.com.steipete.codexbar.debug"
public static let widgetSnapshotFilename = "widget-snapshot.json"
public static let migrationVersion = 1
public static let migrationVersionKey = "appGroupMigrationVersion"
private static let sharedDefaultsMigrationKeys = [
"debugDisableKeychainAccess",
"widgetSelectedProvider",
]
public struct MigrationResult: Sendable {
public enum Status: String, Sendable {
case alreadyCompleted
case targetUnavailable
case noChangesNeeded
case migrated
}
public let status: Status
public let copiedSnapshot: Bool
public let copiedDefaults: Int
public init(status: Status, copiedSnapshot: Bool = false, copiedDefaults: Int = 0) {
self.status = status
self.copiedSnapshot = copiedSnapshot
self.copiedDefaults = copiedDefaults
}
}
public static func currentGroupID(for bundleID: String? = Bundle.main.bundleIdentifier) -> String {
self.currentGroupID(teamID: self.resolvedTeamID(), bundleID: bundleID)
}
static func currentGroupID(teamID: String, bundleID: String?) -> String {
let base = "\(teamID).com.steipete.codexbar"
return self.isDebugBundleID(bundleID) ? "\(base).debug" : base
}
public static func resolvedTeamID(bundle: Bundle = .main) -> String {
self.resolvedTeamID(
infoDictionaryOverride: bundle.infoDictionary,
bundleURLOverride: bundle.bundleURL)
}
static func resolvedTeamID(
infoDictionaryOverride: [String: Any]?,
bundleURLOverride: URL?) -> String
{
if let teamID = self.codeSignatureTeamID(bundleURL: bundleURLOverride) {
return teamID
}
if let teamID = infoDictionaryOverride?[self.teamIDInfoKey] as? String,
!teamID.isEmpty
{
return teamID
}
return self.defaultTeamID
}
public static func legacyGroupID(for bundleID: String? = Bundle.main.bundleIdentifier) -> String {
self.isDebugBundleID(bundleID) ? self.legacyDebugGroupID : self.legacyReleaseGroupID
}
public static func sharedDefaults(
bundleID: String? = Bundle.main.bundleIdentifier,
fileManager: FileManager = .default)
-> UserDefaults?
{
guard self.currentContainerURL(bundleID: bundleID, fileManager: fileManager) != nil else { return nil }
return UserDefaults(suiteName: self.currentGroupID(for: bundleID))
}
public static func currentContainerURL(
bundleID: String? = Bundle.main.bundleIdentifier,
fileManager: FileManager = .default)
-> URL?
{
#if os(macOS)
fileManager.containerURL(forSecurityApplicationGroupIdentifier: self.currentGroupID(for: bundleID))
#else
nil
#endif
}
public static func snapshotURL(
bundleID: String? = Bundle.main.bundleIdentifier,
fileManager: FileManager = .default,
homeDirectory: URL = FileManager.default.homeDirectoryForCurrentUser)
-> URL
{
if let container = self.currentContainerURL(bundleID: bundleID, fileManager: fileManager) {
return container.appendingPathComponent(self.widgetSnapshotFilename, isDirectory: false)
}
let directory = self.localFallbackDirectory(fileManager: fileManager, homeDirectory: homeDirectory)
return directory.appendingPathComponent(self.widgetSnapshotFilename, isDirectory: false)
}
public static func localFallbackDirectory(
fileManager: FileManager = .default,
homeDirectory _: URL = FileManager.default.homeDirectoryForCurrentUser)
-> URL
{
let base = fileManager.urls(for: .applicationSupportDirectory, in: .userDomainMask).first
?? fileManager.temporaryDirectory
let directory = base.appendingPathComponent("CodexBar", isDirectory: true)
try? fileManager.createDirectory(at: directory, withIntermediateDirectories: true)
return directory
}
public static func legacyContainerCandidateURL(
bundleID: String? = Bundle.main.bundleIdentifier,
homeDirectory: URL = FileManager.default.homeDirectoryForCurrentUser)
-> URL
{
homeDirectory
.appendingPathComponent("Library", isDirectory: true)
.appendingPathComponent("Group Containers", isDirectory: true)
.appendingPathComponent(self.legacyGroupID(for: bundleID), isDirectory: true)
}
public static func migrateLegacyDataIfNeeded(
bundleID: String? = Bundle.main.bundleIdentifier,
standardDefaults: UserDefaults = .standard,
fileManager: FileManager = .default,
homeDirectory: URL = FileManager.default.homeDirectoryForCurrentUser,
currentDefaultsOverride: UserDefaults? = nil,
legacyDefaultsOverride: UserDefaults? = nil,
currentSnapshotURLOverride: URL? = nil,
legacySnapshotURLOverride: URL? = nil)
-> MigrationResult
{
if standardDefaults.integer(forKey: self.migrationVersionKey) >= self.migrationVersion {
return MigrationResult(status: .alreadyCompleted)
}
guard let currentDefaults = currentDefaultsOverride ?? self.sharedDefaults(
bundleID: bundleID,
fileManager: fileManager)
else {
return MigrationResult(status: .targetUnavailable)
}
let legacyDefaults = legacyDefaultsOverride ?? UserDefaults(suiteName: self.legacyGroupID(for: bundleID))
let currentSnapshotURL = currentSnapshotURLOverride
?? self.currentContainerURL(bundleID: bundleID, fileManager: fileManager)?
.appendingPathComponent(self.widgetSnapshotFilename, isDirectory: false)
let legacySnapshotURL = legacySnapshotURLOverride
?? self.legacyContainerCandidateURL(bundleID: bundleID, homeDirectory: homeDirectory)
.appendingPathComponent(self.widgetSnapshotFilename, isDirectory: false)
let copiedSnapshot = {
guard let currentSnapshotURL else { return false }
guard !fileManager.fileExists(atPath: currentSnapshotURL.path),
fileManager.fileExists(atPath: legacySnapshotURL.path)
else {
return false
}
do {
try fileManager.createDirectory(
at: currentSnapshotURL.deletingLastPathComponent(),
withIntermediateDirectories: true)
try fileManager.copyItem(at: legacySnapshotURL, to: currentSnapshotURL)
return true
} catch {
return false
}
}()
let copiedDefaults = self.copyLegacySharedDefaults(
from: legacyDefaults,
to: currentDefaults)
let result = if copiedSnapshot || copiedDefaults > 0 {
MigrationResult(
status: .migrated,
copiedSnapshot: copiedSnapshot,
copiedDefaults: copiedDefaults)
} else {
MigrationResult(status: .noChangesNeeded)
}
standardDefaults.set(self.migrationVersion, forKey: self.migrationVersionKey)
return result
}
private static func copyLegacySharedDefaults(
from legacyDefaults: UserDefaults?,
to currentDefaults: UserDefaults) -> Int
{
guard let legacyDefaults else { return 0 }
var copied = 0
for key in self.sharedDefaultsMigrationKeys {
guard currentDefaults.object(forKey: key) == nil,
let legacyValue = legacyDefaults.object(forKey: key)
else {
continue
}
currentDefaults.set(legacyValue, forKey: key)
copied += 1
}
return copied
}
private static func isDebugBundleID(_ bundleID: String?) -> Bool {
guard let bundleID, !bundleID.isEmpty else { return false }
return bundleID.contains(".debug")
}
private static func codeSignatureTeamID(bundleURL: URL?) -> String? {
#if os(macOS)
guard let bundleURL else { return nil }
var staticCode: SecStaticCode?
guard SecStaticCodeCreateWithPath(bundleURL as CFURL, SecCSFlags(), &staticCode) == errSecSuccess,
let code = staticCode
else {
return nil
}
var infoCF: CFDictionary?
guard SecCodeCopySigningInformation(
code,
SecCSFlags(rawValue: kSecCSSigningInformation),
&infoCF) == errSecSuccess,
let info = infoCF as? [String: Any],
let teamID = info[kSecCodeInfoTeamIdentifier as String] as? String,
!teamID.isEmpty
else {
return nil
}
return teamID
#else
_ = bundleURL
return nil
#endif
}
}
@@ -0,0 +1,8 @@
import Foundation
#if os(Linux)
@discardableResult
func autoreleasepool<Result>(_ work: () throws -> Result) rethrows -> Result {
try work()
}
#endif
@@ -0,0 +1,81 @@
import Foundation
package enum BoundedTaskJoinOutcome<Value: Sendable> {
case value(Value)
case failure(any Error)
case timedOut
}
package final class BoundedTaskJoin<Value: Sendable>: @unchecked Sendable {
private let lock = NSLock()
private let sourceTask: Task<Value, Error>
private var outcome: BoundedTaskJoinOutcome<Value>?
private var continuation: CheckedContinuation<BoundedTaskJoinOutcome<Value>, Never>?
private var observerTask: Task<Void, Never>?
private var timeoutTask: Task<Void, Never>?
package init(sourceTask: Task<Value, Error>) {
self.sourceTask = sourceTask
}
package func value(joinGrace: Duration) async -> BoundedTaskJoinOutcome<Value> {
await withTaskCancellationHandler {
await withCheckedContinuation { continuation in
self.lock.lock()
if let outcome = self.outcome {
self.lock.unlock()
continuation.resume(returning: outcome)
return
}
self.continuation = continuation
let sourceTask = self.sourceTask
self.observerTask = Task { [weak self] in
do {
let value = try await sourceTask.value
self?.resolve(.value(value), cancelSource: false)
} catch {
self?.resolve(.failure(error), cancelSource: false)
}
}
self.timeoutTask = Task { [weak self] in
do {
if joinGrace > .zero {
try await Task.sleep(for: joinGrace)
}
self?.resolve(.timedOut, cancelSource: true)
} catch {
// The source completed or the caller canceled the race.
}
}
self.lock.unlock()
}
} onCancel: {
self.resolve(.failure(CancellationError()), cancelSource: true)
}
}
private func resolve(_ outcome: BoundedTaskJoinOutcome<Value>, cancelSource: Bool) {
self.lock.lock()
guard self.outcome == nil else {
self.lock.unlock()
return
}
self.outcome = outcome
let continuation = self.continuation
self.continuation = nil
let observerTask = self.observerTask
let timeoutTask = self.timeoutTask
self.observerTask = nil
self.timeoutTask = nil
self.lock.unlock()
if cancelSource {
self.sourceTask.cancel()
}
observerTask?.cancel()
timeoutTask?.cancel()
continuation?.resume(returning: outcome)
}
}
@@ -0,0 +1,355 @@
import Foundation
#if os(macOS)
import os.lock
import SweetCookieKit
enum BrowserCookieStoreAccessDecision: Equatable {
case allowed
case suppressed
}
struct BrowserCookieStoreAccessSuppressedError: LocalizedError {
var errorDescription: String? {
"Browser cookie store access is suppressed for this process."
}
}
public enum BrowserCookieAccessGate {
private struct State {
var loaded = false
var deniedUntilByBrowser: [String: Date] = [:]
var chromiumFamilyDeniedUntil: Date?
}
private final class ExplicitRetryScope: @unchecked Sendable {
private struct State {
var selectedBrowser: Browser?
var cookieReadClaimed = false
}
private let lock = OSAllocatedUnfairLock<State>(initialState: State())
func allows(_ browser: Browser) -> Bool {
self.lock.withLock { state in
if let selectedBrowser = state.selectedBrowser {
return selectedBrowser == browser && !state.cookieReadClaimed
}
state.selectedBrowser = browser
return true
}
}
func contains(_ browser: Browser) -> Bool {
self.lock.withLock { $0.selectedBrowser == browser }
}
func claimCookieRead(for browser: Browser) -> Bool {
self.lock.withLock { state in
guard state.selectedBrowser == browser, !state.cookieReadClaimed else { return false }
state.cookieReadClaimed = true
return true
}
}
}
private static let lock = OSAllocatedUnfairLock<State>(initialState: State())
private static let defaultsKey = "browserCookieAccessDeniedUntil"
private static let chromiumFamilyDefaultsKey = "__chromiumFamily__"
private static let cooldownInterval: TimeInterval = 60 * 60 * 6
private static let log = CodexBarLog.logger(LogCategories.browserCookieGate)
@TaskLocal private static var explicitRetryScope: ExplicitRetryScope?
@TaskLocal private static var deniedBrowsersForTesting: [Browser]?
static let allowTestCookieAccessEnvironmentKey = "CODEXBAR_ALLOW_TEST_BROWSER_COOKIE_ACCESS"
static func cookieStoreAccessDecision(
homeDirectories: [URL],
processName: String = ProcessInfo.processInfo.processName,
environment: [String: String] = ProcessInfo.processInfo.environment) -> BrowserCookieStoreAccessDecision
{
guard KeychainTestSafety.isRunningUnderTests(processName: processName, environment: environment),
environment[self.allowTestCookieAccessEnvironmentKey] != "1"
else {
return .allowed
}
let defaultHomes = Set(BrowserCookieClient.defaultHomeDirectories().map(Self.normalizedPath))
let usesDefaultHome = homeDirectories.contains { defaultHomes.contains(Self.normalizedPath($0)) }
return usesDefaultHome ? .suppressed : .allowed
}
public static func shouldAttempt(_ browser: Browser, now: Date = Date()) -> Bool {
guard browser.usesKeychainForCookieDecryption else { return true }
guard !KeychainAccessGate.isDisabled else { return false }
if self.deniedBrowsersForTesting?.contains(browser) == true {
return self.isExplicitRetryAllowed(for: browser)
}
let shouldCheckKeychain = self.lock.withLock { state in
self.loadIfNeeded(&state)
if let blockedUntil = state.deniedUntilByBrowser[browser.rawValue] {
if blockedUntil > now {
if self.isExplicitRetryAllowed(for: browser) {
self.log.info(
"Explicit browser cookie retry allowed",
metadata: ["browser": browser.displayName])
return true
}
self.log.debug(
"Cookie access blocked",
metadata: ["browser": browser.displayName, "until": "\(blockedUntil.timeIntervalSince1970)"])
return false
}
state.deniedUntilByBrowser.removeValue(forKey: browser.rawValue)
self.persist(state)
}
if let blockedUntil = state.chromiumFamilyDeniedUntil {
if blockedUntil > now {
self.log.debug(
"Chromium-family cookie access blocked",
metadata: ["browser": browser.displayName, "until": "\(blockedUntil.timeIntervalSince1970)"])
return false
}
state.chromiumFamilyDeniedUntil = nil
self.persist(state)
}
return true
}
guard shouldCheckKeychain else { return false }
let requiresInteraction = self.chromiumKeychainRequiresInteraction(for: browser)
if requiresInteraction {
self.recordDenied(for: browser, now: now)
if self.isExplicitRetryAllowed(for: browser) {
self.log.info(
"Browser cookie interaction allowed for explicit retry",
metadata: ["browser": browser.displayName])
return true
}
self.log.info(
"Cookie access requires keychain interaction; suppressing Chromium family",
metadata: ["browser": browser.displayName])
return false
}
self.log.debug("Cookie access allowed", metadata: ["browser": browser.displayName])
return true
}
public static func withExplicitRetry<T>(_ operation: () throws -> T) rethrows -> T {
try self.$explicitRetryScope.withValue(ExplicitRetryScope()) {
try operation()
}
}
public static func withExplicitRetry<T>(
isolation: isolated (any Actor)? = #isolation,
operation: () async throws -> T) async rethrows -> T
{
try await self.$explicitRetryScope.withValue(ExplicitRetryScope()) {
try await operation()
}
}
static func withDeniedBrowsersForTesting<T>(
_ browsers: [Browser],
isolation _: isolated (any Actor)? = #isolation,
operation: () async throws -> T) async rethrows -> T
{
try await self.$deniedBrowsersForTesting.withValue(browsers) {
try await operation()
}
}
static func operationPreservingAccessContext<T: Sendable>(
_ operation: @escaping @Sendable () throws -> T) -> @Sendable () throws -> T
{
let interaction = ProviderInteractionContext.current
let retryScope = self.explicitRetryScope
return {
try ProviderInteractionContext.$current.withValue(interaction) {
try self.$explicitRetryScope.withValue(retryScope) {
try operation()
}
}
}
}
public static func recordIfNeeded(_ error: Error, now: Date = Date()) {
guard let error = error as? BrowserCookieError else { return }
guard case .accessDenied = error else { return }
self.recordDenied(for: error.browser, now: now)
}
public static func recordDenied(for browser: Browser, now: Date = Date()) {
guard browser.usesKeychainForCookieDecryption else { return }
let blockedUntil = now.addingTimeInterval(self.cooldownInterval)
self.lock.withLock { state in
self.loadIfNeeded(&state)
state.deniedUntilByBrowser[browser.rawValue] = blockedUntil
state.chromiumFamilyDeniedUntil = blockedUntil
self.persist(state)
}
self.log
.info(
"Browser cookie access denied; suppressing Chromium family",
metadata: [
"browser": browser.displayName,
"until": "\(blockedUntil.timeIntervalSince1970)",
])
}
static func recordAllowed(for browser: Browser) {
guard browser.usesKeychainForCookieDecryption,
ProviderInteractionContext.current == .userInitiated,
self.explicitRetryScope?.contains(browser) == true
else { return }
let clearedCooldown = self.lock.withLock { state in
self.loadIfNeeded(&state)
guard state.deniedUntilByBrowser[browser.rawValue] != nil,
state.chromiumFamilyDeniedUntil != nil
else { return false }
state.deniedUntilByBrowser.removeValue(forKey: browser.rawValue)
state.chromiumFamilyDeniedUntil = state.deniedUntilByBrowser.values.max()
self.persist(state)
return true
}
guard clearedCooldown else { return }
self.log.info(
"Explicit browser cookie retry succeeded",
metadata: ["browser": browser.displayName])
}
static func claimExplicitRetryCookieReadIfNeeded(for browser: Browser) -> Bool {
guard browser.usesKeychainForCookieDecryption,
ProviderInteractionContext.current == .userInitiated,
let retryScope = self.explicitRetryScope,
retryScope.contains(browser)
else { return true }
let hasActiveBrowserCooldown = self.lock.withLock { state in
self.loadIfNeeded(&state)
return state.deniedUntilByBrowser[browser.rawValue] != nil
}
guard hasActiveBrowserCooldown else { return true }
return retryScope.claimCookieRead(for: browser)
}
public static func resetForTesting() {
self.lock.withLock { state in
state.loaded = true
state.deniedUntilByBrowser.removeAll()
state.chromiumFamilyDeniedUntil = nil
UserDefaults.standard.removeObject(forKey: self.defaultsKey)
}
}
private static func isExplicitRetryAllowed(for browser: Browser) -> Bool {
ProviderInteractionContext.current == .userInitiated && self.explicitRetryScope?.allows(browser) == true
}
private static func chromiumKeychainRequiresInteraction(for browser: Browser) -> Bool {
let labels = browser.safeStorageLabels.isEmpty ? self.safeStorageLabels : browser.safeStorageLabels
for label in labels {
switch KeychainAccessPreflight.checkGenericPassword(service: label.service, account: label.account) {
case .allowed:
return false
case .interactionRequired:
return true
case .notFound, .failure:
continue
}
}
return false
}
private static let safeStorageLabels: [(service: String, account: String)] = Browser.safeStorageLabels
private static func normalizedPath(_ url: URL) -> String {
url.standardizedFileURL.resolvingSymlinksInPath().path
}
private static func loadIfNeeded(_ state: inout State) {
guard !state.loaded else { return }
state.loaded = true
guard let raw = UserDefaults.standard.dictionary(forKey: self.defaultsKey) as? [String: Double] else {
return
}
state.chromiumFamilyDeniedUntil = raw[self.chromiumFamilyDefaultsKey].map(Date.init(timeIntervalSince1970:))
state.deniedUntilByBrowser = raw
.filter { $0.key != self.chromiumFamilyDefaultsKey }
.mapValues { Date(timeIntervalSince1970: $0) }
if state.chromiumFamilyDeniedUntil == nil {
// Existing per-browser cooldowns become family-wide on upgrade.
state.chromiumFamilyDeniedUntil = state.deniedUntilByBrowser.values.max()
}
}
private static func persist(_ state: State) {
var raw = state.deniedUntilByBrowser.mapValues { $0.timeIntervalSince1970 }
raw[self.chromiumFamilyDefaultsKey] = state.chromiumFamilyDeniedUntil?.timeIntervalSince1970
UserDefaults.standard.set(raw, forKey: self.defaultsKey)
}
}
extension BrowserCookieClient {
public func codexBarStores(for browser: Browser) throws -> [BrowserCookieStore] {
guard BrowserCookieAccessGate.cookieStoreAccessDecision(
homeDirectories: self.configuration.homeDirectories) == .allowed
else {
throw BrowserCookieStoreAccessSuppressedError()
}
guard BrowserCookieAccessGate.shouldAttempt(browser) else { return [] }
return self.stores(for: browser)
}
public func codexBarRecords(
matching query: BrowserCookieQuery,
in browser: Browser,
logger: ((String) -> Void)? = nil) throws -> [BrowserCookieStoreRecords]
{
guard BrowserCookieAccessGate.cookieStoreAccessDecision(
homeDirectories: self.configuration.homeDirectories) == .allowed
else {
throw BrowserCookieStoreAccessSuppressedError()
}
guard BrowserCookieAccessGate.shouldAttempt(browser) else { return [] }
guard BrowserCookieAccessGate.claimExplicitRetryCookieReadIfNeeded(for: browser) else { return [] }
do {
let records = try self.records(matching: query, in: browser, logger: logger)
BrowserCookieAccessGate.recordAllowed(for: browser)
return records
} catch {
BrowserCookieAccessGate.recordIfNeeded(error)
throw error
}
}
}
#else
public enum BrowserCookieAccessGate {
public static func shouldAttempt(_ browser: Browser, now: Date = Date()) -> Bool {
true
}
public static func withExplicitRetry<T>(_ operation: () throws -> T) rethrows -> T {
try operation()
}
public static func withExplicitRetry<T>(
isolation: isolated (any Actor)? = #isolation,
operation: () async throws -> T) async rethrows -> T
{
try await operation()
}
static func withDeniedBrowsersForTesting<T>(
_: [Browser],
isolation _: isolated (any Actor)? = #isolation,
operation: () async throws -> T) async rethrows -> T
{
try await operation()
}
public static func recordIfNeeded(_ error: Error, now: Date = Date()) {}
public static func recordDenied(for browser: Browser, now: Date = Date()) {}
public static func resetForTesting() {}
}
#endif
@@ -0,0 +1,62 @@
#if os(macOS)
import SweetCookieKit
public typealias BrowserCookieImportOrder = [Browser]
#else
public struct Browser: Sendable, Hashable {
public init() {}
}
public typealias BrowserCookieImportOrder = [Browser]
#endif
extension [Browser] {
/// Filters a browser list to sources worth attempting for cookie imports.
///
/// This is intentionally stricter than "app installed": it aims to avoid unnecessary Keychain prompts.
public func cookieImportCandidates(using detection: BrowserDetection) -> [Browser] {
let candidates = self.filter { browser in
if KeychainAccessGate.isDisabled, browser.usesKeychainForCookieDecryption {
return false
}
return detection.isCookieSourceAvailable(browser)
}
return candidates.filter { BrowserCookieAccessGate.shouldAttempt($0) }
}
/// Filters a browser list to sources with usable profile data on disk.
public func browsersWithProfileData(using detection: BrowserDetection) -> [Browser] {
self.filter { detection.hasUsableProfileData($0) }
}
}
#if os(macOS)
extension Browser {
var usesKeychainForCookieDecryption: Bool {
switch self {
case .safari, .firefox, .zen:
return false
case .chrome, .chromeBeta, .chromeCanary,
.arc, .arcBeta, .arcCanary,
.chatgptAtlas,
.chromium,
.brave, .braveBeta, .braveNightly,
.edge, .edgeBeta, .edgeCanary,
.helium,
.vivaldi,
.dia,
.yandex,
.comet:
return true
@unknown default:
return true
}
}
}
#else
extension Browser {
var usesKeychainForCookieDecryption: Bool {
false
}
}
#endif
+378
View File
@@ -0,0 +1,378 @@
import Foundation
#if os(macOS)
@preconcurrency import AppKit
import Darwin
import os.lock
import SweetCookieKit
enum BrowserProfileAccessIssue: Equatable {
case accessDenied
case unreadable
}
/// Browser presence + profile heuristics.
///
/// Primary goal: avoid triggering unnecessary Keychain prompts (e.g. Chromium Safe Storage) by skipping
/// cookie imports from browsers that have no profile data on disk.
public final class BrowserDetection: Sendable {
public static let defaultCacheTTL: TimeInterval = 60 * 10
private let cache = OSAllocatedUnfairLock<[CacheKey: CachedResult]>(initialState: [:])
private let homeDirectory: String
private let cacheTTL: TimeInterval
private let now: @Sendable () -> Date
private let fileExists: @Sendable (String) -> Bool
private let directoryContents: @Sendable (String) -> [String]?
private let applicationURLs: @Sendable (String) -> [URL]
private let profileAccessIssue: @Sendable (String) -> BrowserProfileAccessIssue?
private struct CachedResult {
let value: Bool
let timestamp: Date
}
private enum ProbeKind: Int, Hashable {
case appInstalled
case usableProfileData
case usableCookieStore
}
private struct CacheKey: Hashable {
let browser: Browser
let kind: ProbeKind
}
public convenience init(
homeDirectory: String = FileManager.default.homeDirectoryForCurrentUser.path,
cacheTTL: TimeInterval = BrowserDetection.defaultCacheTTL,
now: @escaping @Sendable () -> Date = Date.init,
fileExists: @escaping @Sendable (String) -> Bool = { path in FileManager.default.fileExists(atPath: path) },
directoryContents: @escaping @Sendable (String) -> [String]? = { path in
try? FileManager.default.contentsOfDirectory(atPath: path)
})
{
self.init(
homeDirectory: homeDirectory,
cacheTTL: cacheTTL,
now: now,
fileExists: fileExists,
directoryContents: directoryContents,
applicationURLs: Self.registeredApplicationURLs,
profileAccessIssue: Self.probeProfileAccessIssue)
}
init(
homeDirectory: String,
cacheTTL: TimeInterval,
now: @escaping @Sendable () -> Date,
fileExists: @escaping @Sendable (String) -> Bool,
directoryContents: @escaping @Sendable (String) -> [String]?,
applicationURLs: @escaping @Sendable (String) -> [URL],
profileAccessIssue: @escaping @Sendable (String) -> BrowserProfileAccessIssue?)
{
self.homeDirectory = homeDirectory
self.cacheTTL = cacheTTL
self.now = now
self.fileExists = fileExists
self.directoryContents = directoryContents
self.applicationURLs = applicationURLs
self.profileAccessIssue = profileAccessIssue
}
public func isAppInstalled(_ browser: Browser) -> Bool {
// Safari is always available on macOS.
if browser == .safari {
return true
}
return self.cachedBool(browser: browser, kind: .appInstalled) {
self.detectAppInstalled(for: browser)
}
}
/// Returns true when a cookie import attempt for this browser should be allowed.
///
/// This is intentionally stricter than `isAppInstalled`: non-Safari browsers must still be installed,
/// and Chromium browsers must have profile data (to avoid stale sources and unnecessary Keychain prompts).
public func isCookieSourceAvailable(_ browser: Browser) -> Bool {
let homeURL = URL(fileURLWithPath: self.homeDirectory, isDirectory: true)
guard BrowserCookieAccessGate.cookieStoreAccessDecision(homeDirectories: [homeURL]) == .allowed else {
return false
}
// Safari does not need Keychain decryption and can still yield cookies if its storage path changes.
if browser == .safari {
return true
}
// Do not cache app presence here: uninstalling a browser must remove it from the next import attempt.
guard self.detectAppInstalled(for: browser) else { return false }
// For browsers that typically require keychain-backed decryption, ensure an actual cookie store exists.
if self.requiresProfileValidation(browser) {
return self.hasUsableCookieStore(browser)
}
return self.hasUsableProfileData(browser)
}
func cookieSourceProfileAccessIssue(_ browser: Browser) -> BrowserProfileAccessIssue? {
let homeURL = URL(fileURLWithPath: self.homeDirectory, isDirectory: true)
guard BrowserCookieAccessGate.cookieStoreAccessDecision(homeDirectories: [homeURL]) == .allowed,
browser != .safari,
self.detectAppInstalled(for: browser),
let profilePath = self.profilePath(for: browser, homeDirectory: self.homeDirectory)
else {
return nil
}
return self.profileAccessIssue(profilePath)
}
public func hasUsableProfileData(_ browser: Browser) -> Bool {
self.cachedBool(browser: browser, kind: .usableProfileData) {
self.detectUsableProfileData(for: browser)
}
}
private func hasUsableCookieStore(_ browser: Browser) -> Bool {
self.cachedBool(browser: browser, kind: .usableCookieStore) {
self.detectUsableCookieStore(for: browser)
}
}
public func clearCache() {
self.cache.withLock { cache in
cache.removeAll()
}
}
// MARK: - Detection Logic
private func cachedBool(browser: Browser, kind: ProbeKind, compute: () -> Bool) -> Bool {
let now = self.now()
let key = CacheKey(browser: browser, kind: kind)
if let cached = self.cache.withLock({ cache in cache[key] }) {
if now.timeIntervalSince(cached.timestamp) < self.cacheTTL {
return cached.value
}
}
let result = compute()
self.cache.withLock { cache in
cache[key] = CachedResult(value: result, timestamp: now)
}
return result
}
private func detectAppInstalled(for browser: Browser) -> Bool {
let appPaths = self.applicationPaths(for: browser)
for path in appPaths where self.fileExists(path) {
return true
}
guard let appName = self.applicationName(for: browser) else { return false }
return self.applicationURLs(appName).contains { self.fileExists($0.path) }
}
private func detectUsableProfileData(for browser: Browser) -> Bool {
guard let profilePath = self.profilePath(for: browser, homeDirectory: self.homeDirectory) else {
return false
}
guard self.fileExists(profilePath) else {
return false
}
// For Chromium-based browsers (and Firefox), verify actual profile data exists.
if self.requiresProfileValidation(browser) {
return self.hasValidProfileDirectory(for: browser, at: profilePath)
}
return true
}
private func detectUsableCookieStore(for browser: Browser) -> Bool {
guard let profilePath = self.profilePath(for: browser, homeDirectory: self.homeDirectory) else {
return false
}
guard self.fileExists(profilePath) else {
return false
}
return self.hasValidCookieStore(for: browser, at: profilePath)
}
private func applicationPaths(for browser: Browser) -> [String] {
guard let appName = self.applicationName(for: browser) else { return [] }
return [
"/Applications/\(appName).app",
"\(self.homeDirectory)/Applications/\(appName).app",
]
}
private func applicationName(for browser: Browser) -> String? {
browser.appBundleName
}
private static func registeredApplicationURLs(named appName: String) -> [URL] {
let probeURL = URL(string: "https://chatgpt.com")!
let bundleName = "\(appName).app"
return NSWorkspace.shared.urlsForApplications(toOpen: probeURL)
.filter { $0.lastPathComponent == bundleName }
}
private func profilePath(for browser: Browser, homeDirectory: String) -> String? {
if browser == .safari {
return "\(homeDirectory)/Library/Cookies/Cookies.binarycookies"
}
if let relativePath = browser.chromiumProfileRelativePath {
return "\(homeDirectory)/Library/Application Support/\(relativePath)"
}
if let geckoFolder = browser.geckoProfilesFolder {
return "\(homeDirectory)/Library/Application Support/\(geckoFolder)/Profiles"
}
return nil
}
private func requiresProfileValidation(_ browser: Browser) -> Bool {
// Chromium-based browsers should have Default/ or Profile*/ subdirectories
if browser == .safari {
return false
}
if browser == .helium {
// Helium doesn't use the Default/Profile* pattern
return false
}
if browser.usesGeckoProfileStore {
// Firefox should have at least one *.default* directory
return true
}
if browser.usesChromiumProfileStore {
return true
}
return false
}
private func hasValidProfileDirectory(for browser: Browser, at profilePath: String) -> Bool {
guard let contents = self.directoryContents(profilePath) else { return false }
// Check for Default/ or Profile*/ subdirectories for Chromium browsers
let hasProfile = contents.contains { name in
name == "Default" || name.hasPrefix("Profile ") || name.hasPrefix("user-")
}
if browser.usesGeckoProfileStore {
return contents.contains { name in
name.range(of: ".default", options: [.caseInsensitive]) != nil
}
}
return hasProfile
}
private func hasValidCookieStore(for browser: Browser, at profilePath: String) -> Bool {
guard let contents = self.directoryContents(profilePath) else { return false }
if browser.usesGeckoProfileStore {
for name in contents where name.range(of: ".default", options: [.caseInsensitive]) != nil {
let cookieDB = "\(profilePath)/\(name)/cookies.sqlite"
if self.fileExists(cookieDB) {
return true
}
}
return false
}
for name in contents where name == "Default" || name.hasPrefix("Profile ") || name.hasPrefix("user-") {
let cookieDBLegacy = "\(profilePath)/\(name)/Cookies"
let cookieDBNetwork = "\(profilePath)/\(name)/Network/Cookies"
if self.fileExists(cookieDBLegacy) || self.fileExists(cookieDBNetwork) {
return true
}
}
return false
}
private static func probeProfileAccessIssue(_ path: String) -> BrowserProfileAccessIssue? {
do {
_ = try FileManager.default.contentsOfDirectory(atPath: path)
return nil
} catch {
if self.isPermissionError(error) {
return .accessDenied
}
let nsError = error as NSError
if nsError.domain == NSCocoaErrorDomain,
nsError.code == CocoaError.fileNoSuchFile.rawValue
{
return nil
}
return .unreadable
}
}
private static func isPermissionError(_ error: Error) -> Bool {
let nsError = error as NSError
if nsError.domain == NSCocoaErrorDomain,
nsError.code == CocoaError.fileReadNoPermission.rawValue
{
return true
}
if nsError.domain == NSPOSIXErrorDomain,
nsError.code == Int(EACCES) || nsError.code == Int(EPERM)
{
return true
}
guard let underlying = nsError.userInfo[NSUnderlyingErrorKey] as? Error else { return false }
return Self.isPermissionError(underlying)
}
}
#else
// MARK: - Non-macOS stub
public struct BrowserDetection: Sendable {
public static let defaultCacheTTL: TimeInterval = 0
public init(
homeDirectory: String = "",
cacheTTL: TimeInterval = BrowserDetection.defaultCacheTTL,
now: @escaping @Sendable () -> Date = Date.init,
fileExists: @escaping @Sendable (String) -> Bool = { _ in false },
directoryContents: @escaping @Sendable (String) -> [String]? = { _ in nil })
{
_ = homeDirectory
_ = cacheTTL
_ = now
_ = fileExists
_ = directoryContents
}
public func isAppInstalled(_ browser: Browser) -> Bool {
false
}
public func isCookieSourceAvailable(_ browser: Browser) -> Bool {
false
}
public func hasUsableProfileData(_ browser: Browser) -> Bool {
false
}
public func clearCache() {}
}
#endif
+42
View File
@@ -0,0 +1,42 @@
import Foundation
public enum CodexHomeScope {
public static func normalizedHomePath(
_ rawPath: String?,
fileManager: FileManager = .default)
-> String?
{
guard var path = rawPath?.trimmingCharacters(in: .whitespacesAndNewlines), !path.isEmpty else {
return nil
}
if path == "~" {
path = fileManager.homeDirectoryForCurrentUser.path
} else if path.hasPrefix("~/") {
path = fileManager.homeDirectoryForCurrentUser
.appendingPathComponent(String(path.dropFirst(2)), isDirectory: true)
.path
} else if path.hasPrefix("~") {
return nil
}
guard (path as NSString).isAbsolutePath else { return nil }
return URL(fileURLWithPath: path, isDirectory: true).standardizedFileURL.path
}
public static func ambientHomeURL(
env: [String: String],
fileManager: FileManager = .default)
-> URL
{
if let raw = env["CODEX_HOME"]?.trimmingCharacters(in: .whitespacesAndNewlines), !raw.isEmpty {
return URL(fileURLWithPath: raw, isDirectory: true)
}
return fileManager.homeDirectoryForCurrentUser.appendingPathComponent(".codex", isDirectory: true)
}
public static func scopedEnvironment(base: [String: String], codexHome: String?) -> [String: String] {
guard let codexHome, !codexHome.isEmpty else { return base }
var env = base
env["CODEX_HOME"] = codexHome
return env
}
}
@@ -0,0 +1,169 @@
import Crypto
import Foundation
public struct ManagedCodexAccount: Codable, Identifiable, Sendable {
public let id: UUID
public let email: String
public let providerAccountID: String?
public let workspaceLabel: String?
public let workspaceAccountID: String?
public let authFingerprint: String?
public let managedHomePath: String
public let createdAt: TimeInterval
public let updatedAt: TimeInterval
public let lastAuthenticatedAt: TimeInterval?
public init(
id: UUID,
email: String,
providerAccountID: String? = nil,
workspaceLabel: String? = nil,
workspaceAccountID: String? = nil,
authFingerprint: String? = nil,
managedHomePath: String,
createdAt: TimeInterval,
updatedAt: TimeInterval,
lastAuthenticatedAt: TimeInterval?)
{
self.id = id
self.email = Self.normalizeEmail(email)
self.providerAccountID = Self.normalizeProviderAccountID(providerAccountID)
self.workspaceLabel = Self.normalizeWorkspaceLabel(workspaceLabel)
self.workspaceAccountID = Self.normalizeWorkspaceAccountID(workspaceAccountID)
self.authFingerprint = CodexAuthFingerprint.normalize(authFingerprint)
self.managedHomePath = managedHomePath
self.createdAt = createdAt
self.updatedAt = updatedAt
self.lastAuthenticatedAt = lastAuthenticatedAt
}
static func normalizeEmail(_ email: String) -> String {
email.trimmingCharacters(in: .whitespacesAndNewlines).lowercased()
}
public static func normalizeProviderAccountID(_ providerAccountID: String?) -> String? {
CodexIdentityResolver.normalizeAccountID(providerAccountID)
}
public static func normalizeWorkspaceLabel(_ workspaceLabel: String?) -> String? {
guard let trimmed = workspaceLabel?.trimmingCharacters(in: .whitespacesAndNewlines), !trimmed.isEmpty else {
return nil
}
return trimmed
}
public static func normalizeWorkspaceAccountID(_ workspaceAccountID: String?) -> String? {
guard let trimmed = workspaceAccountID?.trimmingCharacters(in: .whitespacesAndNewlines), !trimmed.isEmpty else {
return nil
}
return trimmed.lowercased()
}
public init(from decoder: any Decoder) throws {
let container = try decoder.container(keyedBy: CodingKeys.self)
try self.init(
id: container.decode(UUID.self, forKey: .id),
email: container.decode(String.self, forKey: .email),
providerAccountID: container.decodeIfPresent(String.self, forKey: .providerAccountID),
workspaceLabel: container.decodeIfPresent(String.self, forKey: .workspaceLabel),
workspaceAccountID: container.decodeIfPresent(String.self, forKey: .workspaceAccountID),
authFingerprint: container.decodeIfPresent(String.self, forKey: .authFingerprint),
managedHomePath: container.decode(String.self, forKey: .managedHomePath),
createdAt: container.decode(TimeInterval.self, forKey: .createdAt),
updatedAt: container.decode(TimeInterval.self, forKey: .updatedAt),
lastAuthenticatedAt: container.decodeIfPresent(TimeInterval.self, forKey: .lastAuthenticatedAt))
}
}
public enum CodexAuthFingerprint {
public static func authFileURL(homePath: String) -> URL {
URL(fileURLWithPath: homePath, isDirectory: true)
.appendingPathComponent("auth.json", isDirectory: false)
}
public static func fingerprint(homePath: String, fileManager: FileManager = .default) -> String? {
let url = self.authFileURL(homePath: homePath)
guard fileManager.fileExists(atPath: url.path),
let data = try? Data(contentsOf: url)
else {
return nil
}
return self.fingerprint(data: data)
}
public static func fingerprint(env: [String: String], fileManager: FileManager = .default) -> String? {
self.fingerprint(
homePath: CodexHomeScope.ambientHomeURL(env: env, fileManager: fileManager).path,
fileManager: fileManager)
}
public static func fingerprint(data: Data) -> String {
SHA256.hash(data: data).map { String(format: "%02x", $0) }.joined()
}
public static func normalize(_ fingerprint: String?) -> String? {
guard let trimmed = fingerprint?.trimmingCharacters(in: .whitespacesAndNewlines).lowercased(),
!trimmed.isEmpty
else {
return nil
}
return trimmed
}
}
public struct ManagedCodexAccountSet: Codable, Sendable {
public let version: Int
public let accounts: [ManagedCodexAccount]
public init(version: Int, accounts: [ManagedCodexAccount]) {
self.version = version
self.accounts = Self.sanitizedAccounts(accounts)
}
public func account(id: UUID) -> ManagedCodexAccount? {
self.accounts.first { $0.id == id }
}
public func account(email: String, providerAccountID: String? = nil) -> ManagedCodexAccount? {
let normalizedEmail = ManagedCodexAccount.normalizeEmail(email)
if let normalizedProviderAccountID = ManagedCodexAccount.normalizeProviderAccountID(providerAccountID),
let exactMatch = self.accounts.first(where: {
$0.email == normalizedEmail && $0.providerAccountID == normalizedProviderAccountID
})
{
return exactMatch
}
if providerAccountID != nil {
return self.accounts.first { $0.email == normalizedEmail && $0.providerAccountID == nil }
}
return self.accounts.first { $0.email == normalizedEmail }
}
public init(from decoder: any Decoder) throws {
let container = try decoder.container(keyedBy: CodingKeys.self)
self.version = try container.decode(Int.self, forKey: .version)
self.accounts = try container.decode([ManagedCodexAccount].self, forKey: .accounts)
}
private static func sanitizedAccounts(_ accounts: [ManagedCodexAccount]) -> [ManagedCodexAccount] {
var seenIDs: Set<UUID> = []
var seenProviderAccountKeys: Set<String> = []
var seenLegacyEmails: Set<String> = []
var sanitized: [ManagedCodexAccount] = []
sanitized.reserveCapacity(accounts.count)
for account in accounts {
guard seenIDs.insert(account.id).inserted else { continue }
if let providerAccountID = account.providerAccountID {
guard seenProviderAccountKeys.insert("\(account.email)\u{0}\(providerAccountID)").inserted else {
continue
}
} else {
guard seenLegacyEmails.insert(account.email).inserted else { continue }
}
sanitized.append(account)
}
return sanitized
}
}
@@ -0,0 +1,59 @@
import Foundation
public enum CodexActiveSource: Codable, Equatable, Sendable {
case liveSystem
case managedAccount(id: UUID)
case profileHome(path: String)
private enum CodingKeys: String, CodingKey {
case kind
case accountID
case homePath
}
private enum Kind: String, Codable {
case liveSystem
case managedAccount
case profileHome
}
public init(from decoder: any Decoder) throws {
let container = try decoder.container(keyedBy: CodingKeys.self)
switch try container.decode(Kind.self, forKey: .kind) {
case .liveSystem:
if let path = try container.decodeIfPresent(String.self, forKey: .homePath),
!path.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty
{
self = .profileHome(path: path)
} else {
self = .liveSystem
}
case .managedAccount:
let id = try container.decode(UUID.self, forKey: .accountID)
self = .managedAccount(id: id)
case .profileHome:
let path = try container.decode(String.self, forKey: .homePath)
if path.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty {
self = .liveSystem
} else {
self = .profileHome(path: path)
}
}
}
public func encode(to encoder: any Encoder) throws {
var container = encoder.container(keyedBy: CodingKeys.self)
switch self {
case .liveSystem:
try container.encode(Kind.liveSystem, forKey: .kind)
case let .managedAccount(id):
try container.encode(Kind.managedAccount, forKey: .kind)
try container.encode(id, forKey: .accountID)
case let .profileHome(path):
// Released builds only understand liveSystem/managedAccount. Keep the envelope
// downgrade-readable while newer builds recover the profile selection from homePath.
try container.encode(Kind.liveSystem, forKey: .kind)
try container.encode(path, forKey: .homePath)
}
}
}
@@ -0,0 +1,332 @@
import Foundation
public struct CodexBarConfig: Codable, Sendable {
public static let currentVersion = 1
public var version: Int
public var providers: [ProviderConfig]
public init(version: Int = Self.currentVersion, providers: [ProviderConfig]) {
self.version = version
self.providers = providers
}
public static func makeDefault(
metadata: [UsageProvider: ProviderMetadata] = ProviderDescriptorRegistry.metadata) -> CodexBarConfig
{
let providers = UsageProvider.allCases.map { provider in
Self.defaultProviderConfig(
provider,
metadata: metadata,
alibabaTokenPlanRegion: .international)
}
return CodexBarConfig(version: Self.currentVersion, providers: providers)
}
/// Alphabetical provider ordering with enabled providers on top: enabled first, then disabled,
/// each group sorted case-insensitively by display name. Used by the Providers settings pane's
/// alphabetical sort toggle; it never mutates the user's stored manual order.
public static func alphabeticalProviderOrder(
metadata: [UsageProvider: ProviderMetadata] = ProviderDescriptorRegistry.metadata,
enablement: (UsageProvider) -> Bool) -> [UsageProvider]
{
UsageProvider.allCases.sorted { lhs, rhs in
let lhsEnabled = enablement(lhs)
let rhsEnabled = enablement(rhs)
if lhsEnabled != rhsEnabled { return lhsEnabled }
let lhsName = metadata[lhs]?.displayName ?? lhs.rawValue
let rhsName = metadata[rhs]?.displayName ?? rhs.rawValue
switch lhsName.localizedCaseInsensitiveCompare(rhsName) {
case .orderedAscending: return true
case .orderedDescending: return false
case .orderedSame: return lhs.rawValue < rhs.rawValue
}
}
}
public func normalized(
metadata: [UsageProvider: ProviderMetadata] = ProviderDescriptorRegistry.metadata) -> CodexBarConfig
{
var seen: Set<UsageProvider> = []
var normalized: [ProviderConfig] = []
normalized.reserveCapacity(max(self.providers.count, UsageProvider.allCases.count))
for provider in self.providers {
guard !seen.contains(provider.id) else { continue }
seen.insert(provider.id)
normalized.append(provider)
}
for provider in UsageProvider.allCases where !seen.contains(provider) {
normalized.append(Self.defaultProviderConfig(
provider,
metadata: metadata,
alibabaTokenPlanRegion: .chinaMainland))
}
return CodexBarConfig(
version: Self.currentVersion,
providers: normalized)
}
public func orderedProviders() -> [UsageProvider] {
self.providers.map(\.id)
}
public func enabledProviders(
metadata: [UsageProvider: ProviderMetadata] = ProviderDescriptorRegistry.metadata) -> [UsageProvider]
{
self.providers.compactMap { config in
let enabled = config.enabled ?? metadata[config.id]?.defaultEnabled ?? false
return enabled ? config.id : nil
}
}
public func providerConfig(for id: UsageProvider) -> ProviderConfig? {
self.providers.first(where: { $0.id == id })
}
public mutating func setProviderConfig(_ config: ProviderConfig) {
if let index = self.providers.firstIndex(where: { $0.id == config.id }) {
self.providers[index] = config
} else {
self.providers.append(config)
}
}
private static func defaultProviderConfig(
_ provider: UsageProvider,
metadata: [UsageProvider: ProviderMetadata],
alibabaTokenPlanRegion: AlibabaTokenPlanAPIRegion) -> ProviderConfig
{
ProviderConfig(
id: provider,
enabled: metadata[provider]?.defaultEnabled,
region: provider == .alibabatokenplan ? alibabaTokenPlanRegion.rawValue : nil)
}
}
public struct ProviderConfig: Codable, Sendable, Identifiable {
public let id: UsageProvider
public var enabled: Bool?
public var source: ProviderSourceMode?
public var extrasEnabled: Bool?
public var apiKey: String?
public var secretKey: String?
public var cookieHeader: String?
public var cookieSource: ProviderCookieSource?
public var region: String?
public var workspaceID: String?
public var enterpriseHost: String?
public var tokenAccounts: ProviderTokenAccountData?
public var claudeSwapEnabled: Bool?
public var claudeSwapExecutablePath: String?
public var codexActiveSource: CodexActiveSource?
public var codexProfileHomePaths: [String]?
public var quotaWarnings: QuotaWarningConfig?
public var kiloKnownOrganizations: [KiloOrganization]?
public var kiloEnabledOrganizationIDs: [String]?
public var awsProfile: String?
public var awsAuthMode: String?
public init(
id: UsageProvider,
enabled: Bool? = nil,
source: ProviderSourceMode? = nil,
extrasEnabled: Bool? = nil,
apiKey: String? = nil,
secretKey: String? = nil,
cookieHeader: String? = nil,
cookieSource: ProviderCookieSource? = nil,
region: String? = nil,
workspaceID: String? = nil,
enterpriseHost: String? = nil,
tokenAccounts: ProviderTokenAccountData? = nil,
claudeSwapEnabled: Bool? = nil,
claudeSwapExecutablePath: String? = nil,
codexActiveSource: CodexActiveSource? = nil,
codexProfileHomePaths: [String]? = nil,
quotaWarnings: QuotaWarningConfig? = nil,
kiloKnownOrganizations: [KiloOrganization]? = nil,
kiloEnabledOrganizationIDs: [String]? = nil,
awsProfile: String? = nil,
awsAuthMode: String? = nil)
{
self.id = id
self.enabled = enabled
self.source = source
self.extrasEnabled = extrasEnabled
self.apiKey = apiKey
self.secretKey = secretKey
self.cookieHeader = cookieHeader
self.cookieSource = cookieSource
self.region = region
self.workspaceID = workspaceID
self.enterpriseHost = enterpriseHost
self.tokenAccounts = tokenAccounts
self.claudeSwapEnabled = claudeSwapEnabled
self.claudeSwapExecutablePath = claudeSwapExecutablePath
self.codexActiveSource = codexActiveSource
self.codexProfileHomePaths = codexProfileHomePaths
self.quotaWarnings = quotaWarnings
self.kiloKnownOrganizations = kiloKnownOrganizations
self.kiloEnabledOrganizationIDs = kiloEnabledOrganizationIDs
self.awsProfile = awsProfile
self.awsAuthMode = awsAuthMode
}
public var sanitizedAPIKey: String? {
Self.clean(self.apiKey)
}
public var sanitizedSecretKey: String? {
Self.clean(self.secretKey)
}
public var sanitizedCookieHeader: String? {
Self.clean(self.cookieHeader)
}
public var sanitizedRegion: String? {
Self.clean(self.region)
}
public var sanitizedWorkspaceID: String? {
Self.clean(self.workspaceID)
}
public var sanitizedEnterpriseHost: String? {
Self.clean(self.enterpriseHost)
}
public var sanitizedClaudeSwapExecutablePath: String? {
Self.clean(self.claudeSwapExecutablePath)
}
public var sanitizedAWSProfile: String? {
Self.clean(self.awsProfile)
}
public var sanitizedAWSAuthMode: String? {
Self.clean(self.awsAuthMode)
}
private static func clean(_ 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 QuotaWarningWindow: String, Codable, Sendable, CaseIterable {
case session
case weekly
public var displayName: String {
switch self {
case .session:
"session"
case .weekly:
"weekly"
}
}
}
public struct QuotaWarningWindowConfig: Codable, Sendable, Equatable {
public var thresholds: [Int]?
public var enabled: Bool?
public init(thresholds: [Int]? = nil, enabled: Bool? = nil) {
self.thresholds = thresholds.map(QuotaWarningThresholds.sanitized)
self.enabled = enabled
}
public var hasOverride: Bool {
self.thresholds != nil || self.enabled != nil
}
public func isEnabled(global: Bool) -> Bool {
self.enabled ?? (self.thresholds != nil ? true : global)
}
}
public struct QuotaWarningConfig: Codable, Sendable, Equatable {
public var session: QuotaWarningWindowConfig?
public var weekly: QuotaWarningWindowConfig?
public init(
session: QuotaWarningWindowConfig? = nil,
weekly: QuotaWarningWindowConfig? = nil)
{
self.session = session
self.weekly = weekly
}
public func thresholds(for window: QuotaWarningWindow, global: [Int]) -> [Int] {
switch window {
case .session:
QuotaWarningThresholds.sanitized(self.session?.thresholds ?? global)
case .weekly:
QuotaWarningThresholds.sanitized(self.weekly?.thresholds ?? global)
}
}
public func isEnabled(for window: QuotaWarningWindow, global: Bool) -> Bool {
switch window {
case .session:
self.session?.isEnabled(global: global) ?? global
case .weekly:
self.weekly?.isEnabled(global: global) ?? global
}
}
public func hasOverride(for window: QuotaWarningWindow) -> Bool {
switch window {
case .session:
self.session?.hasOverride ?? false
case .weekly:
self.weekly?.hasOverride ?? false
}
}
public var isEmpty: Bool {
self.session?.hasOverride != true && self.weekly?.hasOverride != true
}
}
public enum QuotaWarningThresholds {
public static let defaults = [50, 20]
public static let allowedRange = 0...99
public static func sanitized(_ raw: [Int]) -> [Int] {
guard !raw.isEmpty else { return self.defaults }
let unique = Set(raw.map(self.clamped))
let sorted = unique.sorted(by: >)
return sorted.isEmpty ? self.defaults : sorted
}
public static func active(_ raw: [Int]) -> [Int] {
self.sanitized(raw).filter { $0 > 0 }
}
public static func resolved(upper: Int?, lower: Int?) -> [Int] {
guard upper != nil || lower != nil else { return self.defaults }
let resolvedUpper = self.clamped(upper ?? self.defaults[0])
let lowerDefault = resolvedUpper < self.defaults[1] ? 0 : self.defaults[1]
let resolvedLower = self.clamped(lower ?? lowerDefault)
return self.sanitized([resolvedUpper, resolvedLower])
}
public static func clamped(_ value: Int) -> Int {
min(max(value, self.allowedRange.lowerBound), self.allowedRange.upperBound)
}
}
@@ -0,0 +1,125 @@
import Foundation
public enum CodexBarConfigStoreError: LocalizedError {
case invalidURL
case decodeFailed(String)
case encodeFailed(String)
public var errorDescription: String? {
switch self {
case .invalidURL:
"Invalid CodexBar config path."
case let .decodeFailed(details):
"Failed to decode CodexBar config: \(details)"
case let .encodeFailed(details):
"Failed to encode CodexBar config: \(details)"
}
}
}
public struct CodexBarConfigStore: @unchecked Sendable {
public static let pathEnvironmentKey = "CODEXBAR_CONFIG"
public static let xdgConfigHomeEnvironmentKey = "XDG_CONFIG_HOME"
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 -> CodexBarConfig? {
guard self.fileManager.fileExists(atPath: self.fileURL.path) else { return nil }
let data = try Data(contentsOf: self.fileURL)
let decoder = JSONDecoder()
do {
let decoded = try decoder.decode(CodexBarConfig.self, from: data)
return decoded.normalized()
} catch {
throw CodexBarConfigStoreError.decodeFailed(error.localizedDescription)
}
}
public func loadOrCreateDefault() throws -> CodexBarConfig {
if let existing = try self.load() {
return existing
}
let config = CodexBarConfig.makeDefault()
try self.save(config)
return config
}
public func save(_ config: CodexBarConfig) throws {
let normalized = config.normalized()
let encoder = JSONEncoder()
encoder.outputFormatting = [.prettyPrinted, .sortedKeys]
let data: Data
do {
data = try encoder.encode(normalized)
} catch {
throw CodexBarConfigStoreError.encodeFailed(error.localizedDescription)
}
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 {
guard self.fileManager.fileExists(atPath: self.fileURL.path) else { return }
try self.fileManager.removeItem(at: self.fileURL)
}
public static func defaultURL(
home: URL = FileManager.default.homeDirectoryForCurrentUser,
environment: [String: String] = ProcessInfo.processInfo.environment,
fileManager: FileManager = .default) -> URL
{
if let override = environment[pathEnvironmentKey]?.trimmingCharacters(in: .whitespacesAndNewlines),
!override.isEmpty
{
let expanded = (override as NSString).expandingTildeInPath
return URL(fileURLWithPath: expanded)
}
if let xdgConfigHome = environment[xdgConfigHomeEnvironmentKey]?
.trimmingCharacters(in: .whitespacesAndNewlines),
!xdgConfigHome.isEmpty
{
let expanded = (xdgConfigHome as NSString).expandingTildeInPath
if (expanded as NSString).isAbsolutePath {
return URL(fileURLWithPath: expanded, isDirectory: true)
.appendingPathComponent("codexbar", isDirectory: true)
.appendingPathComponent("config.json")
}
}
let xdgDefault = home
.appendingPathComponent(".config", isDirectory: true)
.appendingPathComponent("codexbar", isDirectory: true)
.appendingPathComponent("config.json")
if fileManager.fileExists(atPath: xdgDefault.path) {
return xdgDefault
}
let legacy = home
.appendingPathComponent(".codexbar", isDirectory: true)
.appendingPathComponent("config.json")
if fileManager.fileExists(atPath: legacy.path) {
return legacy
}
return xdgDefault
}
private func applySecurePermissionsIfNeeded() throws {
#if os(macOS) || os(Linux)
try self.fileManager.setAttributes([
.posixPermissions: NSNumber(value: Int16(0o600)),
], ofItemAtPath: self.fileURL.path)
#endif
}
}
@@ -0,0 +1,359 @@
import Foundation
public enum CodexBarConfigIssueSeverity: String, Codable, Sendable {
case warning
case error
}
public struct CodexBarConfigIssue: Codable, Sendable, Equatable {
public let severity: CodexBarConfigIssueSeverity
public let provider: UsageProvider?
public let field: String?
public let code: String
public let message: String
public init(
severity: CodexBarConfigIssueSeverity,
provider: UsageProvider?,
field: String?,
code: String,
message: String)
{
self.severity = severity
self.provider = provider
self.field = field
self.code = code
self.message = message
}
}
public enum CodexBarConfigValidator {
private static let enterpriseHostProviders: [UsageProvider] = [
.azureopenai,
.clawrouter,
.copilot,
.kimi,
.litellm,
.llmproxy,
.sub2api,
.wayfinder,
]
private static let workspaceIDProviders: [UsageProvider] = [
.azureopenai,
.openai,
.opencode,
.opencodego,
.devin,
.deepgram,
]
public static func validate(_ config: CodexBarConfig) -> [CodexBarConfigIssue] {
var issues: [CodexBarConfigIssue] = []
if config.version != CodexBarConfig.currentVersion {
issues.append(CodexBarConfigIssue(
severity: .error,
provider: nil,
field: "version",
code: "version_mismatch",
message: "Unsupported config version \(config.version)."))
}
for entry in config.providers {
self.validateProvider(entry, issues: &issues)
}
return issues
}
private static func validateProvider(_ entry: ProviderConfig, issues: inout [CodexBarConfigIssue]) {
let provider = entry.id
let descriptor = ProviderDescriptorRegistry.descriptor(for: provider)
let supportedSources = descriptor.fetchPlan.sourceModes
let supportsWeb = supportedSources.contains(.auto) || supportedSources.contains(.web)
let supportsAPI = supportedSources.contains(.api)
if let source = entry.source, !supportedSources.contains(source) {
issues.append(CodexBarConfigIssue(
severity: .error,
provider: provider,
field: "source",
code: "unsupported_source",
message: "Source \(source.rawValue) is not supported for \(provider.rawValue)."))
}
if let apiKey = entry.apiKey, !apiKey.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty, !supportsAPI {
issues.append(CodexBarConfigIssue(
severity: .warning,
provider: provider,
field: "apiKey",
code: "api_key_unused",
message: "apiKey is set but \(provider.rawValue) does not support api source."))
}
if let source = entry.source, source == .api, !supportsAPI {
issues.append(CodexBarConfigIssue(
severity: .error,
provider: provider,
field: "source",
code: "api_source_unsupported",
message: "Source api is not supported for \(provider.rawValue)."))
}
if let source = entry.source, source == .api,
self.providerRequiresAPIKey(provider),
!self.hasConfiguredAPICredential(entry)
{
issues.append(CodexBarConfigIssue(
severity: .warning,
provider: provider,
field: "apiKey",
code: "api_key_missing",
message: "Source api is selected but apiKey is missing for \(provider.rawValue)."))
}
if entry.cookieSource != nil, !supportsWeb {
issues.append(CodexBarConfigIssue(
severity: .warning,
provider: provider,
field: "cookieSource",
code: "cookie_source_unused",
message: "cookieSource is set but \(provider.rawValue) does not use web cookies."))
}
if let cookieHeader = entry.cookieHeader,
!cookieHeader.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty,
!supportsWeb
{
issues.append(CodexBarConfigIssue(
severity: .warning,
provider: provider,
field: "cookieHeader",
code: "cookie_header_unused",
message: "cookieHeader is set but \(provider.rawValue) does not use web cookies."))
}
if let cookieSource = entry.cookieSource,
cookieSource == .manual,
entry.cookieHeader?.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty ?? true
{
issues.append(CodexBarConfigIssue(
severity: .warning,
provider: provider,
field: "cookieHeader",
code: "cookie_header_missing",
message: "cookieSource manual is set but cookieHeader is missing for \(provider.rawValue)."))
}
self.validateSecretKey(entry, issues: &issues)
self.validateSub2APIBaseURL(entry, issues: &issues)
self.validateRegion(entry, issues: &issues)
self.validateZaiTeamContext(entry, issues: &issues)
if let workspaceID = entry.workspaceID,
!workspaceID.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty,
!self.providerSupportsWorkspaceID(provider)
{
issues.append(CodexBarConfigIssue(
severity: .warning,
provider: provider,
field: "workspaceID",
code: "workspace_unused",
message: "workspaceID is set but only \(self.workspaceIDProviderList) support workspaceID."))
}
if let enterpriseHost = entry.enterpriseHost,
!enterpriseHost.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty,
!self.providerSupportsEnterpriseHost(provider)
{
issues.append(CodexBarConfigIssue(
severity: .warning,
provider: provider,
field: "enterpriseHost",
code: "enterprise_host_unused",
message: "enterpriseHost is set but only \(self.enterpriseHostProviderList) support enterpriseHost."))
}
if let tokenAccounts = entry.tokenAccounts, !tokenAccounts.accounts.isEmpty,
TokenAccountSupportCatalog.support(for: provider) == nil
{
issues.append(CodexBarConfigIssue(
severity: .warning,
provider: provider,
field: "tokenAccounts",
code: "token_accounts_unused",
message: "tokenAccounts are set but \(provider.rawValue) does not support token accounts."))
}
}
private static func validateSecretKey(_ entry: ProviderConfig, issues: inout [CodexBarConfigIssue]) {
guard let secretKey = entry.secretKey,
!secretKey.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty,
entry.id != .bedrock,
entry.id != .doubao
else {
return
}
issues.append(CodexBarConfigIssue(
severity: .warning,
provider: entry.id,
field: "secretKey",
code: "secret_key_unused",
message: "secretKey is set but only bedrock and doubao use secretKey."))
}
private static func validateSub2APIBaseURL(_ entry: ProviderConfig, issues: inout [CodexBarConfigIssue]) {
guard entry.id == .sub2api,
let raw = entry.enterpriseHost?.trimmingCharacters(in: .whitespacesAndNewlines),
!raw.isEmpty,
Sub2APISettingsReader.baseURL(environment: [Sub2APISettingsReader.baseURLEnvironmentKey: raw]) == nil
else {
return
}
issues.append(CodexBarConfigIssue(
severity: .error,
provider: .sub2api,
field: "enterpriseHost",
code: "invalid_enterprise_host",
message: Sub2APISettingsError.invalidBaseURL.errorDescription ?? "Invalid sub2api base URL."))
}
private static func validateZaiTeamContext(_ entry: ProviderConfig, issues: inout [CodexBarConfigIssue]) {
guard entry.id == .zai else { return }
guard let tokenAccounts = entry.tokenAccounts else { return }
for account in tokenAccounts.accounts
where account.sanitizedUsageScope?.lowercased() == ZaiUsageScope.team.rawValue
{
if account.sanitizedOrganizationID == nil || account.sanitizedWorkspaceID == nil {
issues.append(self.zaiMissingTeamContextIssue(field: "tokenAccounts"))
return
}
}
}
private static func zaiMissingTeamContextIssue(field: String) -> CodexBarConfigIssue {
CodexBarConfigIssue(
severity: .warning,
provider: .zai,
field: field,
code: "zai_team_context_missing",
message: "z.ai Team mode requires both organizationID and workspaceID.")
}
private static func providerSupportsWorkspaceID(_ provider: UsageProvider) -> Bool {
self.workspaceIDProviders.contains(provider)
}
private static var workspaceIDProviderList: String {
self.formattedProviderList(self.workspaceIDProviders)
}
private static func formattedProviderList(_ providers: [UsageProvider]) -> String {
let names = providers.map(\.rawValue)
guard let last = names.last else { return "" }
guard names.count > 1 else { return last }
return "\(names.dropLast().joined(separator: ", ")), and \(last)"
}
private static func providerSupportsEnterpriseHost(_ provider: UsageProvider) -> Bool {
self.enterpriseHostProviders.contains(provider)
}
private static func providerRequiresAPIKey(_ provider: UsageProvider) -> Bool {
provider != .wayfinder
}
private static func hasConfiguredAPICredential(_ entry: ProviderConfig) -> Bool {
if let apiKey = entry.apiKey?.trimmingCharacters(in: .whitespacesAndNewlines), !apiKey.isEmpty {
return true
}
return entry.tokenAccounts?.accounts.contains(where: { account in
let token = account.token.trimmingCharacters(in: .whitespacesAndNewlines)
return !token.isEmpty &&
TokenAccountSupportCatalog.envOverride(for: entry.id, token: token)?.isEmpty == false
}) == true
}
private static var enterpriseHostProviderList: String {
self.formattedProviderList(self.enterpriseHostProviders)
}
private static func validateRegion(_ entry: ProviderConfig, issues: inout [CodexBarConfigIssue]) {
let provider = entry.id
guard let region = entry.region?.trimmingCharacters(in: .whitespacesAndNewlines),
!region.isEmpty
else {
return
}
switch provider {
case .minimax:
self.validateKnownRegion(
region,
provider: provider,
isValid: MiniMaxAPIRegion(rawValue: region) != nil,
displayName: "MiniMax",
issues: &issues)
case .zai:
self.validateKnownRegion(
region,
provider: provider,
isValid: ZaiAPIRegion(rawValue: region) != nil,
displayName: "z.ai",
issues: &issues)
case .alibaba:
self.validateKnownRegion(
region,
provider: provider,
isValid: AlibabaCodingPlanAPIRegion(rawValue: region) != nil,
displayName: "Alibaba Coding Plan",
issues: &issues)
case .alibabatokenplan:
self.validateKnownRegion(
region,
provider: provider,
isValid: AlibabaTokenPlanAPIRegion(rawValue: region) != nil,
displayName: "Alibaba Token Plan",
issues: &issues)
case .moonshot:
self.validateKnownRegion(
region,
provider: provider,
isValid: MoonshotRegion(rawValue: region) != nil,
displayName: "Moonshot",
issues: &issues)
case .bedrock, .doubao:
break
default:
issues.append(CodexBarConfigIssue(
severity: .warning,
provider: provider,
field: "region",
code: "region_unused",
message: "region is set but \(provider.rawValue) does not use regions."))
}
}
private static func validateKnownRegion(
_ region: String,
provider: UsageProvider,
isValid: Bool,
displayName: String,
issues: inout [CodexBarConfigIssue])
{
guard !isValid else { return }
issues.append(CodexBarConfigIssue(
severity: .error,
provider: provider,
field: "region",
code: "invalid_region",
message: "Region \(region) is not a valid \(displayName) region."))
}
}
@@ -0,0 +1,416 @@
import Foundation
public enum ProviderConfigEnvironment {
public static func applyAPIKeyOverride(
base: [String: String],
provider: UsageProvider,
config: ProviderConfig?) -> [String: String]
{
if let env = self.applyDedicatedProviderOverrides(base: base, provider: provider, config: config) {
return env
}
guard let apiKey = config?.sanitizedAPIKey, !apiKey.isEmpty else { return base }
var env = base
if let key = self.directAPIKeyEnvironmentKey(for: provider) {
env[key] = apiKey
return env
}
switch provider {
case .copilot:
env["COPILOT_API_TOKEN"] = apiKey
case .kimik2:
if let key = KimiK2SettingsReader.apiKeyEnvironmentKeys.first {
env[key] = apiKey
}
case .warp:
if let key = WarpSettingsReader.apiKeyEnvironmentKeys.first {
env[key] = apiKey
}
case .codebuff:
// Preserve a token already present in the process environment so that
// runtime/CI overrides win over a key saved in Settings (matches the
// precedence used by `ProviderTokenResolver.codebuffResolution`).
if CodebuffSettingsReader.apiKey(environment: base) == nil {
env[CodebuffSettingsReader.apiTokenKey] = apiKey
}
case .crof:
if CrofSettingsReader.apiKey(environment: base) == nil,
let key = CrofSettingsReader.apiKeyEnvironmentKeys.first
{
env[key] = apiKey
}
case .doubao:
if let key = DoubaoSettingsReader.apiKeyEnvironmentKeys.first {
env[key] = apiKey
}
default:
break
}
return env
}
public static func supportsAPIKeyOverride(for provider: UsageProvider) -> Bool {
if self.directAPIKeyEnvironmentKey(for: provider) != nil { return true }
switch provider {
case .copilot, .kimik2, .warp, .codebuff, .crof, .doubao:
return true
case .azureopenai:
return true
default:
return false
}
}
private static func baseURLEnvironmentKey(for provider: UsageProvider) -> String? {
switch provider {
case .llmproxy:
LLMProxySettingsReader.baseURLEnvironmentKey
case .litellm:
LiteLLMSettingsReader.baseURLEnvironmentKey
case .clawrouter:
ClawRouterSettingsReader.baseURLEnvironmentKey
case .sub2api:
Sub2APISettingsReader.baseURLEnvironmentKey
case .wayfinder:
WayfinderSettingsReader.baseURLEnvironmentKey
default:
nil
}
}
private static func supportsAPIKeyAndBaseURLOverride(_ provider: UsageProvider) -> Bool {
self.baseURLEnvironmentKey(for: provider) != nil
}
private static func applyDedicatedProviderOverrides(
base: [String: String],
provider: UsageProvider,
config: ProviderConfig?) -> [String: String]?
{
if self.supportsAPIKeyAndBaseURLOverride(provider) {
return self.applyAPIKeyAndBaseURLOverrides(base: base, provider: provider, config: config)
}
if let env = self.applyMultiFieldCredentialOverrides(base: base, provider: provider, config: config) {
return env
}
return switch provider {
case .deepgram:
self.applyDeepgramOverrides(base: base, config: config)
case .azureopenai:
self.applyAzureOpenAIOverrides(base: base, config: config)
case .kimi:
self.applyKimiOverrides(base: base, config: config)
case .doubao:
self.applyDoubaoOverrides(base: base, config: config)
case .sakana:
self.applySakanaOverrides(base: base, config: config)
default:
nil
}
}
private static func applyMultiFieldCredentialOverrides(
base: [String: String],
provider: UsageProvider,
config: ProviderConfig?) -> [String: String]?
{
switch provider {
case .openai:
self.applyOpenAIOverrides(base: base, config: config)
case .bedrock:
self.applyBedrockOverrides(base: base, config: config)
default:
nil
}
}
private static func directAPIKeyEnvironmentKey(for provider: UsageProvider) -> String? {
switch provider {
case .amp:
AmpSettingsReader.apiTokenKey
case .openai:
OpenAIAPISettingsReader.adminAPIKeyEnvironmentKey
case .azureopenai:
AzureOpenAISettingsReader.apiKeyEnvironmentKey
case .claude:
ClaudeAdminAPISettingsReader.adminAPIKeyEnvironmentKey
case .zai:
ZaiSettingsReader.apiTokenKey
case .minimax:
MiniMaxAPISettingsReader.apiTokenKey
case .alibaba:
AlibabaCodingPlanSettingsReader.apiTokenKey
case .kilo:
KiloSettingsReader.apiTokenKey
case .synthetic:
SyntheticSettingsReader.apiKeyKey
case .openrouter:
OpenRouterSettingsReader.envKey
case .elevenlabs:
ElevenLabsSettingsReader.apiKeyEnvironmentKey
case .moonshot:
MoonshotSettingsReader.apiKeyEnvironmentKeys.first
case .kimi:
KimiSettingsReader.apiKeyEnvironmentKeys.first
case .ollama:
OllamaAPISettingsReader.apiKeyEnvironmentKeys.first
case .venice:
VeniceSettingsReader.apiKeyEnvironmentKey
case .deepgram:
DeepgramSettingsReader.apiKeyEnvironmentKey
case .groq:
GroqSettingsReader.apiKeyEnvironmentKey
case .llmproxy:
LLMProxySettingsReader.apiKeyEnvironmentKey
case .chutes, .poe, .litellm, .crossmodel, .clawrouter, .factory, .sub2api:
self.additionalAPIKeyEnvironmentKey(for: provider)
default:
nil
}
}
private static func additionalAPIKeyEnvironmentKey(for provider: UsageProvider) -> String? {
switch provider {
case .chutes:
ChutesSettingsReader.apiKeyEnvironmentKey
case .poe:
PoeSettingsReader.apiKeyEnvironmentKey
case .litellm:
LiteLLMSettingsReader.apiKeyEnvironmentKey
case .crossmodel:
CrossModelSettingsReader.envKey
case .clawrouter:
ClawRouterSettingsReader.apiKeyEnvironmentKey
case .sub2api:
Sub2APISettingsReader.apiKeyEnvironmentKey
case .factory:
FactorySettingsReader.apiTokenKey
default:
nil
}
}
private static func applyOpenAIOverrides(
base: [String: String],
config: ProviderConfig?) -> [String: String]
{
guard let config else { return base }
var env = base
if let apiKey = config.sanitizedAPIKey {
env[OpenAIAPISettingsReader.adminAPIKeyEnvironmentKey] = apiKey
}
if let projectID = config.sanitizedWorkspaceID {
env[OpenAIAPISettingsReader.projectIDEnvironmentKey] = projectID
}
return env
}
private static func applyBedrockOverrides(
base: [String: String],
config: ProviderConfig?) -> [String: String]
{
guard let config else { return base }
var env = base
// Only project an explicit auth-mode selection. When the config does not
// specify one, leave the base environment untouched so an env-driven setup
// (AWS_PROFILE or CODEXBAR_BEDROCK_AUTH_MODE from the launch environment) is
// still inferred by BedrockSettingsReader instead of being forced to `keys`.
let configMode = config.sanitizedAWSAuthMode.flatMap(BedrockAuthMode.init(rawValue:))
if let configMode {
env[BedrockSettingsReader.authModeKey] = configMode.rawValue
}
let baseMode = BedrockSettingsReader
.cleaned(base[BedrockSettingsReader.authModeKey])
.flatMap { BedrockAuthMode(rawValue: $0.lowercased()) }
let mergedAccessKey = config.sanitizedAPIKey ?? BedrockSettingsReader.accessKeyID(environment: base)
let mergedSecretKey = config.sanitizedSecretKey ?? BedrockSettingsReader.secretAccessKey(environment: base)
let hasMergedStaticKeys = mergedAccessKey != nil && mergedSecretKey != nil
let effectiveMode: BedrockAuthMode = if let configMode {
configMode
} else if let baseMode {
baseMode
} else if hasMergedStaticKeys {
// Upgrade path: a config saved before auth modes existed keeps using
// static credentials (including env+config layering) even if AWS_PROFILE
// is present in the base environment, so existing users are never
// silently switched to a profile/account.
.keys
} else {
BedrockSettingsReader.authMode(environment: base)
}
switch effectiveMode {
case .profile:
if let profile = config.sanitizedAWSProfile {
env[BedrockSettingsReader.profileKey] = profile
}
case .keys:
if let accessKeyID = config.sanitizedAPIKey {
env[BedrockSettingsReader.accessKeyIDKey] = accessKeyID
}
if let secretAccessKey = config.sanitizedSecretKey {
env[BedrockSettingsReader.secretAccessKeyKey] = secretAccessKey
}
}
if let region = config.sanitizedRegion {
env[BedrockSettingsReader.regionKeys[0]] = region
}
return env
}
private static func applyDeepgramOverrides(
base: [String: String],
config: ProviderConfig?) -> [String: String]
{
guard let config else { return base }
var env = base
if let apiKey = config.sanitizedAPIKey {
env[DeepgramSettingsReader.apiKeyEnvironmentKey] = apiKey
}
if let projectID = config.sanitizedWorkspaceID {
env[DeepgramSettingsReader.projectIDEnvironmentKey] = projectID
}
return env
}
private static func applyAPIKeyAndBaseURLOverrides(
base: [String: String],
provider: UsageProvider,
config: ProviderConfig?) -> [String: String]
{
var env = base
if let apiKey = config?.sanitizedAPIKey,
let key = self.directAPIKeyEnvironmentKey(for: provider)
{
env[key] = apiKey
}
if let baseURL = config?.sanitizedEnterpriseHost,
let key = self.baseURLEnvironmentKey(for: provider)
{
env[key] = baseURL
}
return env
}
private static func applyKimiOverrides(
base: [String: String],
config: ProviderConfig?) -> [String: String]
{
guard let config else { return base }
var env = base
if let apiKey = config.sanitizedAPIKey,
let key = KimiSettingsReader.apiKeyEnvironmentKeys.first
{
env[key] = apiKey
}
if let baseURL = config.sanitizedEnterpriseHost {
env[KimiSettingsReader.codeAPIBaseURLEnvironmentKey] = baseURL
}
return env
}
private static func applyDoubaoOverrides(
base: [String: String],
config: ProviderConfig?) -> [String: String]
{
guard let config else { return base }
var env = base
let apiKey = config.sanitizedAPIKey
let secretKey = config.sanitizedSecretKey
if let apiKey, self.doubaoAccessKeyID(from: apiKey) == nil {
self.clearDoubaoCodingPlanCredentialKeys(in: &env)
env[DoubaoSettingsReader.apiKeyEnvironmentKeys[0]] = apiKey
if let region = config.sanitizedRegion {
env[DoubaoSettingsReader.regionEnvironmentKeys[0]] = region
}
return env
}
let accessKeyID = self.doubaoAccessKeyID(from: apiKey) ?? DoubaoSettingsReader.accessKeyID(environment: base)
let secretAccessKey = secretKey ?? DoubaoSettingsReader.secretAccessKey(environment: base)
if let accessKeyID, let secretAccessKey {
env[DoubaoSettingsReader.accessKeyIDEnvironmentKeys[0]] = accessKeyID
env[DoubaoSettingsReader.secretAccessKeyEnvironmentKeys[0]] = secretAccessKey
if let region = config.sanitizedRegion ?? self.firstDoubaoRegionValue(in: base) {
env[DoubaoSettingsReader.regionEnvironmentKeys[0]] = region
}
return env
}
if let region = config.sanitizedRegion {
env[DoubaoSettingsReader.regionEnvironmentKeys[0]] = region
}
return env
}
private static func applySakanaOverrides(
base: [String: String],
config: ProviderConfig?) -> [String: String]
{
guard let config else { return base }
var env = base
if let cookieHeader = config.sanitizedCookieHeader {
env[SakanaSettingsReader.cookieHeaderKey] = cookieHeader
}
return env
}
private static func doubaoAccessKeyID(from apiKey: String?) -> String? {
guard let apiKey, apiKey.hasPrefix("AKLT") else { return nil }
return apiKey
}
private static func clearDoubaoCodingPlanCredentialKeys(in environment: inout [String: String]) {
for key in DoubaoSettingsReader.accessKeyIDEnvironmentKeys {
environment.removeValue(forKey: key)
}
for key in DoubaoSettingsReader.secretAccessKeyEnvironmentKeys {
environment.removeValue(forKey: key)
}
}
private static func firstDoubaoRegionValue(in environment: [String: String]) -> String? {
for key in DoubaoSettingsReader.regionEnvironmentKeys {
guard let value = DoubaoSettingsReader.cleaned(environment[key]) else { continue }
return value
}
return nil
}
private static func applyAzureOpenAIOverrides(
base: [String: String],
config: ProviderConfig?) -> [String: String]
{
guard let config else { return base }
var env = base
if let apiKey = config.sanitizedAPIKey {
env[AzureOpenAISettingsReader.apiKeyEnvironmentKey] = apiKey
}
if let endpoint = config.sanitizedEnterpriseHost {
env[AzureOpenAISettingsReader.endpointEnvironmentKey] = endpoint
}
if let deploymentName = config.sanitizedWorkspaceID {
env[AzureOpenAISettingsReader.deploymentNameEnvironmentKey] = deploymentName
}
return env
}
public static func applyProviderConfigOverrides(
base: [String: String],
provider: UsageProvider,
config: ProviderConfig?) -> [String: String]
{
self.applyAPIKeyOverride(base: base, provider: provider, config: config)
}
}
@@ -0,0 +1,982 @@
import Foundation
#if canImport(CryptoKit)
import CryptoKit
#endif
#if canImport(Darwin)
import Darwin
#elseif canImport(Glibc)
import Glibc
#elseif canImport(Musl)
import Musl
#endif
public enum CookieHeaderCache {
public enum Scope: Sendable, Equatable {
case managedAccount(UUID)
case managedStoreUnreadable
case profileHome(String)
case providerVariant(String)
public var isolationIdentifier: String {
switch self {
case let .managedAccount(accountID):
"managed.\(accountID.uuidString.lowercased())"
case .managedStoreUnreadable:
"managed-store-unreadable"
case let .profileHome(path):
"profile-home.\(Self.profileHomeDigest(path))"
case let .providerVariant(variant):
"provider-variant.\(Self.providerVariantDigest(variant))"
}
}
fileprivate var keychainIdentifier: String {
self.isolationIdentifier
}
private static func profileHomeDigest(_ path: String) -> String {
let standardized = URL(fileURLWithPath: path, isDirectory: true).standardizedFileURL.path
#if canImport(CryptoKit)
return SHA256.hash(data: Data(standardized.utf8))
.map { String(format: "%02x", $0) }
.joined()
#else
let digest = standardized.utf8.reduce(UInt64(14_695_981_039_346_656_037)) { partial, byte in
(partial ^ UInt64(byte)) &* 1_099_511_628_211
}
return String(digest, radix: 16)
#endif
}
private static func providerVariantDigest(_ variant: String) -> String {
#if canImport(CryptoKit)
return SHA256.hash(data: Data(variant.utf8))
.map { String(format: "%02x", $0) }
.joined()
#else
let digest = variant.utf8.reduce(UInt64(14_695_981_039_346_656_037)) { partial, byte in
(partial ^ UInt64(byte)) &* 1_099_511_628_211
}
return String(digest, radix: 16)
#endif
}
}
public struct Entry: Codable, Sendable {
public let cookieHeader: String
public let storedAt: Date
public let sourceLabel: String
public init(cookieHeader: String, storedAt: Date, sourceLabel: String) {
self.cookieHeader = cookieHeader
self.storedAt = storedAt
self.sourceLabel = sourceLabel
}
}
public struct ClearSummary: Equatable, Sendable {
public let clearedCount: Int
public let failedCount: Int
public init(clearedCount: Int, failedCount: Int) {
self.clearedCount = clearedCount
self.failedCount = failedCount
}
}
private static let log = CodexBarLog.logger(LogCategories.cookieCache)
private static let legacyBaseURLOverrideLock = NSLock()
private nonisolated(unsafe) static var legacyBaseURLOverride: URL?
private struct DisplaySnapshot {
let entry: Entry?
let refreshAfter: Date
}
private enum LoadOutcome {
case authoritative(Entry?, loadedFromLegacy: Bool)
case temporarilyUnavailable
}
private static let legacyMutationLock = NSLock()
private static let displayCacheLock = NSLock()
private nonisolated(unsafe) static var displayCache: [KeychainCacheStore.Key: DisplaySnapshot] = [:]
private nonisolated(unsafe) static var displayGenerations: [KeychainCacheStore.Key: UInt64] = [:]
private nonisolated(unsafe) static var displayRevalidationsInFlight: Set<KeychainCacheStore.Key> = []
private nonisolated(unsafe) static var legacyMigrationsInFlight: Set<UsageProvider> = []
private nonisolated(unsafe) static var displayStalenessIntervalOverride: TimeInterval?
private nonisolated(unsafe) static var displayUnavailableRetryIntervalOverride: TimeInterval?
private static let displayStalenessInterval: TimeInterval = 30
private static let displayUnavailableRetryInterval: TimeInterval = 1
#if DEBUG
@TaskLocal private static var taskLegacyBaseURLOverride: URL?
@TaskLocal private static var legacyRemovalFailureOverride = false
#endif
private enum LegacyRemovalResult: Equatable {
case removed
case missing
case failed
}
/// Settings rows render the "Cached: " cookie label inside SwiftUI body evaluations, which
/// run repeatedly within a single AppKit layout pass. Each `load` pays a synchronous
/// securityd round-trip and decrypt, so display paths use this memoized variant instead: it
/// returns the last known entry immediately and revalidates a stale snapshot off the calling
/// path. In-process `store` and `clear` calls update the snapshot synchronously; only the
/// first lookup per key pays the keychain read.
public static func loadForDisplay(provider: UsageProvider, scope: Scope? = nil) -> Entry? {
let key = self.key(for: provider, scope: scope)
let (cached, generation) = self.beginDisplayRead(key: key)
guard let cached else {
switch self.loadOutcome(provider: provider, scope: scope, migrateLegacy: false) {
case let .authoritative(entry, loadedFromLegacy):
let committed = self.commitDisplaySnapshotIfCurrent(key: key, entry: entry, generation: generation)
if loadedFromLegacy {
self.scheduleLegacyMigration(provider: provider)
}
return committed
case .temporarilyUnavailable:
return self.commitTemporaryDisplaySnapshotIfCurrent(key: key, generation: generation)
}
}
if Date() >= cached.refreshAfter {
self.scheduleDisplayRevalidation(provider: provider, scope: scope, key: key, generation: generation)
}
return cached.entry
}
/// Registers the key before the Keychain read starts so `clearAll` can invalidate an
/// in-flight first population even when no display snapshot exists yet.
private static func beginDisplayRead(key: KeychainCacheStore.Key) -> (DisplaySnapshot?, UInt64) {
self.displayCacheLock.lock()
defer { self.displayCacheLock.unlock() }
let generation = self.displayGenerations[key] ?? 0
self.displayGenerations[key] = generation
return (self.displayCache[key], generation)
}
private static func scheduleDisplayRevalidation(
provider: UsageProvider,
scope: Scope?,
key: KeychainCacheStore.Key,
generation: UInt64)
{
self.displayCacheLock.lock()
let inserted = self.displayRevalidationsInFlight.insert(key).inserted
self.displayCacheLock.unlock()
guard inserted else { return }
Task(priority: .utility) {
self.revalidateDisplaySnapshot(provider: provider, scope: scope, key: key, generation: generation)
}
}
private static func revalidateDisplaySnapshot(
provider: UsageProvider,
scope: Scope?,
key: KeychainCacheStore.Key,
generation: UInt64)
{
switch self.loadOutcome(provider: provider, scope: scope, migrateLegacy: false) {
case let .authoritative(entry, loadedFromLegacy):
_ = self.commitDisplaySnapshotIfCurrent(key: key, entry: entry, generation: generation)
if loadedFromLegacy {
self.scheduleLegacyMigration(provider: provider)
}
case .temporarilyUnavailable:
self.deferDisplayRetryIfCurrent(key: key, generation: generation)
}
self.displayCacheLock.lock()
self.displayRevalidationsInFlight.remove(key)
self.displayCacheLock.unlock()
}
private static func scheduleLegacyMigration(provider: UsageProvider) {
self.displayCacheLock.lock()
let inserted = self.legacyMigrationsInFlight.insert(provider).inserted
self.displayCacheLock.unlock()
guard inserted else { return }
Task(priority: .utility) {
_ = self.migrateLegacyEntryIfNeeded(provider: provider)
_ = self.displayCacheLock.withLock {
self.legacyMigrationsInFlight.remove(provider)
}
}
}
/// Keychain reads for the display cache happen outside the lock, so a concurrent `store` or
/// `clear` can publish newer state before the read commits. Each mutation bumps the per-key
/// generation; a read only commits if the generation it started from is still current, and
/// otherwise returns whatever newer snapshot won the race.
private static func commitDisplaySnapshotIfCurrent(
key: KeychainCacheStore.Key,
entry: Entry?,
generation: UInt64) -> Entry?
{
self.displayCacheLock.lock()
defer { self.displayCacheLock.unlock() }
guard self.displayGenerations[key, default: 0] == generation else {
return self.displayCache[key]?.entry
}
self.displayCache[key] = self.authoritativeDisplaySnapshot(entry: entry)
return entry
}
private static func commitTemporaryDisplaySnapshotIfCurrent(
key: KeychainCacheStore.Key,
generation: UInt64) -> Entry?
{
self.displayCacheLock.lock()
defer { self.displayCacheLock.unlock() }
guard self.displayGenerations[key, default: 0] == generation else {
return self.displayCache[key]?.entry
}
if let current = self.displayCache[key] {
return current.entry
}
self.displayCache[key] = self.temporaryDisplaySnapshot(entry: nil)
return nil
}
private static func deferDisplayRetryIfCurrent(key: KeychainCacheStore.Key, generation: UInt64) {
self.displayCacheLock.lock()
defer { self.displayCacheLock.unlock() }
guard self.displayGenerations[key, default: 0] == generation,
let current = self.displayCache[key]
else { return }
self.displayCache[key] = self.temporaryDisplaySnapshot(entry: current.entry)
}
private static func updateDisplaySnapshot(key: KeychainCacheStore.Key, entry: Entry?) {
self.displayCacheLock.lock()
self.displayCache[key] = self.authoritativeDisplaySnapshot(entry: entry)
self.displayGenerations[key, default: 0] += 1
self.displayCacheLock.unlock()
}
private static func invalidateDisplaySnapshot(key: KeychainCacheStore.Key) {
self.displayCacheLock.lock()
self.displayCache.removeValue(forKey: key)
self.displayGenerations[key, default: 0] += 1
self.displayCacheLock.unlock()
}
private static func currentDisplayEntry(key: KeychainCacheStore.Key) -> Entry? {
self.displayCacheLock.lock()
defer { self.displayCacheLock.unlock() }
return self.displayCache[key]?.entry
}
private static func authoritativeDisplaySnapshot(entry: Entry?) -> DisplaySnapshot {
DisplaySnapshot(entry: entry, refreshAfter: Date().addingTimeInterval(self.currentDisplayStalenessInterval))
}
private static func temporaryDisplaySnapshot(entry: Entry?) -> DisplaySnapshot {
DisplaySnapshot(
entry: entry,
refreshAfter: Date().addingTimeInterval(self.currentDisplayUnavailableRetryInterval))
}
private static var currentDisplayStalenessInterval: TimeInterval {
self.displayStalenessIntervalOverride ?? self.displayStalenessInterval
}
private static var currentDisplayUnavailableRetryInterval: TimeInterval {
self.displayUnavailableRetryIntervalOverride ?? self.displayUnavailableRetryInterval
}
static func setDisplayStalenessIntervalOverrideForTesting(_ interval: TimeInterval?) {
self.displayStalenessIntervalOverride = interval
}
static func setDisplayUnavailableRetryIntervalOverrideForTesting(_ interval: TimeInterval?) {
self.displayUnavailableRetryIntervalOverride = interval
}
static func resetDisplayCacheForTesting() {
self.displayCacheLock.lock()
self.displayCache.removeAll()
self.displayGenerations.removeAll()
self.displayRevalidationsInFlight.removeAll()
self.legacyMigrationsInFlight.removeAll()
self.displayCacheLock.unlock()
}
static func beginDisplayReadGenerationForTesting(provider: UsageProvider, scope: Scope? = nil) -> UInt64 {
self.beginDisplayRead(key: self.key(for: provider, scope: scope)).1
}
static func currentDisplayEntryForTesting(provider: UsageProvider, scope: Scope? = nil) -> Entry? {
self.currentDisplayEntry(key: self.key(for: provider, scope: scope))
}
#if DEBUG
static func withLegacyRemovalFailureForTesting<T>(_ operation: () throws -> T) rethrows -> T {
try self.$legacyRemovalFailureOverride.withValue(true) {
try operation()
}
}
#endif
@discardableResult
static func commitDisplaySnapshotIfCurrentForTesting(
provider: UsageProvider,
scope: Scope? = nil,
entry: Entry?,
generation: UInt64) -> Entry?
{
self.commitDisplaySnapshotIfCurrent(
key: self.key(for: provider, scope: scope),
entry: entry,
generation: generation)
}
public static func load(provider: UsageProvider, scope: Scope? = nil) -> Entry? {
switch self.loadOutcome(provider: provider, scope: scope, migrateLegacy: true) {
case let .authoritative(entry, _):
entry
case .temporarilyUnavailable:
nil
}
}
static func loadSerialized(provider: UsageProvider, scope: Scope? = nil) -> Entry? {
do {
return try self.withLegacyMutationLock {
let key = self.key(for: provider, scope: scope)
switch KeychainCacheStore.load(key: key, as: Entry.self) {
case let .found(entry):
return entry
case .temporarilyUnavailable:
return nil
case .invalid:
KeychainCacheStore.clear(key: key)
case .missing:
break
}
guard scope == nil else { return nil }
return self.migrateLegacyEntryIfNeededLocked(provider: provider)
}
} catch {
self.log.error("Cookie cache serialized load failed: \(error)")
return nil
}
}
private static func loadOutcome(
provider: UsageProvider,
scope: Scope?,
migrateLegacy: Bool) -> LoadOutcome
{
let key = self.key(for: provider, scope: scope)
switch KeychainCacheStore.load(key: key, as: Entry.self) {
case let .found(entry):
self.log.debug("Cookie cache hit", metadata: ["provider": provider.rawValue])
return .authoritative(entry, loadedFromLegacy: false)
case .temporarilyUnavailable:
self.log.debug("Cookie cache temporarily unavailable", metadata: ["provider": provider.rawValue])
return .temporarilyUnavailable
case .invalid:
self.log.warning("Cookie cache invalid; clearing", metadata: ["provider": provider.rawValue])
KeychainCacheStore.clear(key: key)
case .missing:
self.log.debug("Cookie cache miss", metadata: ["provider": provider.rawValue])
}
guard scope == nil else { return .authoritative(nil, loadedFromLegacy: false) }
if migrateLegacy {
return .authoritative(
self.migrateLegacyEntryIfNeeded(provider: provider),
loadedFromLegacy: false)
}
guard let legacy = self.loadLegacyEntry(for: provider) else {
return .authoritative(nil, loadedFromLegacy: false)
}
return .authoritative(legacy, loadedFromLegacy: true)
}
/// Re-reads both stores while serialized with global cookie mutations. A display-triggered
/// migration may be queued before a clear, so using the captured legacy entry here would
/// otherwise allow the delayed task to restore credentials the user just removed.
private static func migrateLegacyEntryIfNeeded(provider: UsageProvider) -> Entry? {
do {
return try self.withLegacyMutationLock {
self.migrateLegacyEntryIfNeededLocked(provider: provider)
}
} catch {
self.log.error("Cookie cache migration lock failed: \(error)")
return nil
}
}
private static func migrateLegacyEntryIfNeededLocked(provider: UsageProvider) -> Entry? {
let key = self.key(for: provider, scope: nil)
switch KeychainCacheStore.load(key: key, as: Entry.self) {
case let .found(entry):
_ = self.removeLegacyEntry(for: provider)
return entry
case .temporarilyUnavailable:
return nil
case .invalid:
KeychainCacheStore.clear(key: key)
case .missing:
break
}
guard let legacy = self.loadLegacyEntry(for: provider) else { return nil }
if KeychainCacheStore.storeResult(key: key, entry: legacy),
self.removeLegacyEntry(for: provider) == .removed
{
self.log.debug("Cookie cache migrated from legacy store", metadata: ["provider": provider.rawValue])
}
return legacy
}
public static func store(
provider: UsageProvider,
scope: Scope? = nil,
cookieHeader: String,
sourceLabel: String,
now: Date = Date())
{
let trimmed = cookieHeader.trimmingCharacters(in: .whitespacesAndNewlines)
guard let normalized = CookieHeaderNormalizer.normalize(trimmed), !normalized.isEmpty else {
self.clear(provider: provider, scope: scope)
return
}
let entry = Entry(cookieHeader: normalized, storedAt: now, sourceLabel: sourceLabel)
do {
try self.withLegacyMutationLock {
_ = self.store(entry: entry, provider: provider, scope: scope, sourceLabel: sourceLabel)
}
} catch {
self.log.error("Cookie cache store lock failed: \(error)")
}
}
/// Stores only while the cache still matches the entry observed before an asynchronous refresh.
/// A nil expected entry means the cache must still be empty.
@discardableResult
static func storeIfCurrent(
provider: UsageProvider,
scope: Scope? = nil,
expected: Entry?,
cookieHeader: String,
sourceLabel: String,
now: Date = Date()) -> Bool
{
let trimmed = cookieHeader.trimmingCharacters(in: .whitespacesAndNewlines)
guard let normalized = CookieHeaderNormalizer.normalize(trimmed), !normalized.isEmpty else { return false }
let entry = Entry(cookieHeader: normalized, storedAt: now, sourceLabel: sourceLabel)
do {
return try self.withLegacyMutationLock {
guard self.currentEntryMatches(expected, provider: provider, scope: scope) else { return false }
return self.store(entry: entry, provider: provider, scope: scope, sourceLabel: sourceLabel)
}
} catch {
self.log.error("Cookie cache conditional store lock failed: \(error)")
return false
}
}
/// Clears only while the cache still matches the entry that produced a failed asynchronous refresh.
@discardableResult
static func clearIfCurrent(
provider: UsageProvider,
scope: Scope? = nil,
expected: Entry?) -> Bool
{
do {
return try self.withLegacyMutationLock {
guard self.currentEntryMatches(expected, provider: provider, scope: scope) else { return false }
// Keep the expected Keychain row intact when legacy cleanup fails so fallback can replace it.
if scope == nil, self.removeLegacyEntry(for: provider) == .failed {
return false
}
let key = self.key(for: provider, scope: scope)
let result = KeychainCacheStore.clearResult(key: key)
guard result != .failed else { return false }
self.updateDisplaySnapshot(key: key, entry: nil)
return true
}
} catch {
self.log.error("Cookie cache conditional clear lock failed: \(error)")
return false
}
}
private static func currentEntryMatches(
_ expected: Entry?,
provider: UsageProvider,
scope: Scope?) -> Bool
{
let key = self.key(for: provider, scope: scope)
switch KeychainCacheStore.load(key: key, as: Entry.self) {
case let .found(current):
return self.entriesMatch(current, expected)
case .missing:
if scope == nil, let legacy = self.loadLegacyEntry(for: provider) {
return self.entriesMatch(legacy, expected)
}
return expected == nil
case .invalid, .temporarilyUnavailable:
return false
}
}
private static func entriesMatch(_ current: Entry, _ expected: Entry?) -> Bool {
guard let expected else { return false }
return current.cookieHeader == expected.cookieHeader
&& current.storedAt == expected.storedAt
&& current.sourceLabel == expected.sourceLabel
}
@discardableResult
private static func store(
entry: Entry,
provider: UsageProvider,
scope: Scope?,
sourceLabel: String) -> Bool
{
let key = self.key(for: provider, scope: scope)
guard KeychainCacheStore.storeResult(key: key, entry: entry) else { return false }
self.updateDisplaySnapshot(key: key, entry: entry)
if scope == nil {
_ = self.removeLegacyEntry(for: provider)
}
self.log.debug("Cookie cache stored", metadata: ["provider": provider.rawValue, "source": sourceLabel])
return true
}
@discardableResult
public static func clear(provider: UsageProvider, scope: Scope? = nil) -> Int {
self.clearDetailed(provider: provider, scope: scope).clearedCount
}
public static func clearDetailed(provider: UsageProvider, scope: Scope? = nil) -> ClearSummary {
do {
return try self.withLegacyMutationLock {
self.clearDetailedLocked(provider: provider, scope: scope)
}
} catch {
self.log.error("Cookie cache clear lock failed: \(error)")
return ClearSummary(clearedCount: 0, failedCount: 1)
}
}
private static func clearDetailedLocked(provider: UsageProvider, scope: Scope?) -> ClearSummary {
let key = self.key(for: provider, scope: scope)
let result = KeychainCacheStore.clearResult(key: key)
var cleared = result == .removed ? 1 : 0
var failed = result == .failed ? 1 : 0
if result != .failed {
self.updateDisplaySnapshot(key: key, entry: nil)
}
let legacyResult: LegacyRemovalResult = if scope == nil {
self.removeLegacyEntry(for: provider)
} else {
.missing
}
if legacyResult == .removed {
cleared += 1
if result == .failed {
self.invalidateDisplaySnapshot(key: key)
}
} else if legacyResult == .failed {
failed += 1
}
self.log.debug("Cookie cache cleared", metadata: ["provider": provider.rawValue])
return ClearSummary(clearedCount: cleared, failedCount: failed)
}
/// Clears all cookie cache scopes for one provider, including managed Codex account scopes.
/// Returns keychain/legacy removal and failure counts.
@discardableResult
public static func clearAllScopes(provider: UsageProvider) -> Int {
self.clearAllScopesDetailed(provider: provider).clearedCount
}
public static func clearAllScopesDetailed(provider: UsageProvider) -> ClearSummary {
do {
return try self.withLegacyMutationLock {
self.clearAllScopesDetailedLocked(provider: provider)
}
} catch {
self.log.error("Cookie cache clearAllScopes lock failed: \(error)")
return ClearSummary(clearedCount: 0, failedCount: 1)
}
}
private static func clearAllScopesDetailedLocked(provider: UsageProvider) -> ClearSummary {
let (keys, enumerationFailed) = self.cookieKeysResult(for: provider)
var cleared = 0
var failedKeys = Set<KeychainCacheStore.Key>()
for key in keys {
let result = KeychainCacheStore.clearResult(key: key)
if result == .removed {
cleared += 1
}
if result != .failed {
self.updateDisplaySnapshot(key: key, entry: nil)
} else {
failedKeys.insert(key)
}
}
let legacyResult = self.removeLegacyEntry(for: provider)
if legacyResult == .removed {
cleared += 1
let globalKey = self.key(for: provider, scope: nil)
if failedKeys.contains(globalKey) {
self.invalidateDisplaySnapshot(key: globalKey)
}
}
let failedCount = failedKeys.count + (enumerationFailed ? 1 : 0) + (legacyResult == .failed ? 1 : 0)
self.log.debug("Cookie cache clearAllScopes completed", metadata: [
"provider": provider.rawValue,
"cleared": "\(cleared)",
])
return ClearSummary(clearedCount: cleared, failedCount: failedCount)
}
/// Clears cookie caches for all providers, including corrupt/invalid entries.
/// Returns keychain/legacy removal and failure counts.
@discardableResult
public static func clearAll() -> Int {
self.clearAllDetailed().clearedCount
}
public static func clearAllDetailed() -> ClearSummary {
do {
return try self.withLegacyMutationLock {
self.clearAllDetailedLocked()
}
} catch {
self.log.error("Cookie cache clearAll lock failed: \(error)")
return ClearSummary(clearedCount: 0, failedCount: 1)
}
}
private static func clearAllDetailedLocked() -> ClearSummary {
self.displayCacheLock.lock()
let knownDisplayKeys = Set(self.displayCache.keys).union(self.displayGenerations.keys)
self.displayCacheLock.unlock()
let enumeratedKeys: [KeychainCacheStore.Key]
let enumerationFailed: Bool
switch KeychainCacheStore.keysResult(category: "cookie") {
case let .found(keys):
enumeratedKeys = keys
enumerationFailed = false
case .temporarilyUnavailable, .failed:
enumeratedKeys = []
enumerationFailed = true
}
let keys = Set(enumeratedKeys).union(knownDisplayKeys)
var cleared = 0
var failedKeys = Set<KeychainCacheStore.Key>()
for key in keys {
let result = KeychainCacheStore.clearResult(key: key)
if result == .removed {
cleared += 1
}
if result != .failed {
self.updateDisplaySnapshot(key: key, entry: nil)
} else {
failedKeys.insert(key)
}
}
var legacyFailures = 0
for provider in UsageProvider.allCases {
switch self.removeLegacyEntry(for: provider) {
case .removed:
cleared += 1
let globalKey = self.key(for: provider, scope: nil)
if failedKeys.contains(globalKey) {
self.invalidateDisplaySnapshot(key: globalKey)
}
case .failed:
legacyFailures += 1
case .missing:
break
}
}
self.log.debug("Cookie cache clearAll completed", metadata: ["cleared": "\(cleared)"])
return ClearSummary(
clearedCount: cleared,
failedCount: failedKeys.count + (enumerationFailed ? 1 : 0) + legacyFailures)
}
private static func cookieKeysResult(for provider: UsageProvider) -> ([KeychainCacheStore.Key], Bool) {
let exactIdentifier = provider.rawValue
let scopedPrefix = "\(provider.rawValue)."
var seen = Set<KeychainCacheStore.Key>()
var keys: [KeychainCacheStore.Key] = []
let enumeratedKeys: [KeychainCacheStore.Key]
let enumerationFailed: Bool
switch KeychainCacheStore.keysResult(category: "cookie") {
case let .found(keys):
enumeratedKeys = keys
enumerationFailed = false
case .temporarilyUnavailable, .failed:
enumeratedKeys = []
enumerationFailed = true
}
for key in enumeratedKeys {
guard key.identifier == exactIdentifier || key.identifier.hasPrefix(scopedPrefix) else {
continue
}
if seen.insert(key).inserted {
keys.append(key)
}
}
let global = self.key(for: provider, scope: nil)
if seen.insert(global).inserted {
keys.append(global)
}
return (keys, enumerationFailed)
}
static func load(from url: URL) -> Entry? {
guard let data = try? Data(contentsOf: url) else { return nil }
let decoder = JSONDecoder()
decoder.dateDecodingStrategy = .iso8601
return try? decoder.decode(Entry.self, from: data)
}
static func store(_ entry: Entry, to url: URL) {
do {
let dir = url.deletingLastPathComponent()
try FileManager.default.createDirectory(at: dir, withIntermediateDirectories: true)
let encoder = JSONEncoder()
encoder.dateEncodingStrategy = .iso8601
let data = try encoder.encode(entry)
try data.write(to: url, options: [.atomic])
} catch {
self.log.error("Failed to persist cookie cache: \(error)")
}
}
static func setLegacyBaseURLOverrideForTesting(_ url: URL?) {
self.legacyBaseURLOverrideLock.withLock {
self.legacyBaseURLOverride = url
}
}
#if DEBUG
static func withLegacyBaseURLOverrideForTesting<T>(
_ url: URL?,
operation: () throws -> T) rethrows -> T
{
try self.$taskLegacyBaseURLOverride.withValue(url) {
try operation()
}
}
static func withLegacyBaseURLOverrideForTesting<T>(
_ url: URL?,
operation: () async throws -> T) async rethrows -> T
{
try await self.$taskLegacyBaseURLOverride.withValue(url) {
try await operation()
}
}
#endif
static func hasLegacyEntryForTesting(provider: UsageProvider) -> Bool {
self.loadLegacyEntry(for: provider) != nil
}
static func legacyURLForTesting(provider: UsageProvider) -> URL {
self.legacyURL(for: provider)
}
private static func hasKeychainEntry(provider: UsageProvider, scope: Scope?) -> Bool {
let key = self.key(for: provider, scope: scope)
switch KeychainCacheStore.load(key: key, as: Entry.self) {
case .found, .invalid:
return true
case .missing, .temporarilyUnavailable:
return false
}
}
static func hasKeychainEntryForTesting(provider: UsageProvider, scope: Scope? = nil) -> Bool {
self.hasKeychainEntry(provider: provider, scope: scope)
}
static func migrateLegacyEntryIfNeededForTesting(provider: UsageProvider) -> Entry? {
self.migrateLegacyEntryIfNeeded(provider: provider)
}
private static func withLegacyMutationLock<T>(_ operation: () throws -> T) throws -> T {
try self.legacyMutationLock.withLock {
let lockURL = self.legacyMutationLockURL
try FileManager.default.createDirectory(
at: lockURL.deletingLastPathComponent(),
withIntermediateDirectories: true)
let fd = open(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()
}
}
private static var legacyMutationLockURL: URL {
let base = self.currentLegacyBaseURLOverride ?? self.defaultLegacyBaseURL
return base.appendingPathComponent("cookie-cache.lock")
}
private static func removeLegacyEntry(for provider: UsageProvider) -> LegacyRemovalResult {
let url = self.legacyURL(for: provider)
#if DEBUG
if self.legacyRemovalFailureOverride {
return .failed
}
#endif
let existed = FileManager.default.fileExists(atPath: url.path)
do {
try FileManager.default.removeItem(at: url)
return existed ? .removed : .missing
} catch {
if (error as NSError).code != NSFileNoSuchFileError {
Self.log.error("Failed to remove cookie cache (\(provider.rawValue)): \(error)")
return .failed
}
return .missing
}
}
private static func loadLegacyEntry(for provider: UsageProvider) -> Entry? {
self.load(from: self.legacyURL(for: provider))
}
private static func legacyURL(for provider: UsageProvider) -> URL {
if let override = self.currentLegacyBaseURLOverride {
return override.appendingPathComponent("\(provider.rawValue)-cookie.json")
}
return self.defaultLegacyBaseURL.appendingPathComponent("\(provider.rawValue)-cookie.json")
}
private static var currentLegacyBaseURLOverride: URL? {
#if DEBUG
if let taskOverride = self.taskLegacyBaseURLOverride {
return taskOverride
}
#endif
return self.legacyBaseURLOverrideLock.withLock {
self.legacyBaseURLOverride
}
}
private static var defaultLegacyBaseURL: URL {
let fm = FileManager.default
let base = fm.urls(for: .applicationSupportDirectory, in: .userDomainMask).first
?? fm.temporaryDirectory
return base.appendingPathComponent("CodexBar", isDirectory: true)
}
private static func key(for provider: UsageProvider, scope: Scope?) -> KeychainCacheStore.Key {
KeychainCacheStore.Key.cookie(provider: provider, scopeIdentifier: scope?.keychainIdentifier)
}
}
extension CookieHeaderCache {
enum ConditionalMutationObservation {
case authoritative(Entry?)
case keychainTemporarilyUnavailable(legacyEntry: Entry?)
var entry: Entry? {
switch self {
case let .authoritative(entry): entry
case .keychainTemporarilyUnavailable: nil
}
}
}
/// Captures enough state to conditionally persist an asynchronous refresh after a transient
/// Keychain read failure without mistaking an untouched legacy entry for a concurrent write.
static func observeForConditionalMutation(
provider: UsageProvider,
scope: Scope? = nil) -> ConditionalMutationObservation
{
do {
return try self.withLegacyMutationLock {
let key = self.key(for: provider, scope: scope)
switch KeychainCacheStore.load(key: key, as: Entry.self) {
case let .found(entry):
return .authoritative(entry)
case .temporarilyUnavailable:
let legacyEntry = scope == nil ? self.loadLegacyEntry(for: provider) : nil
return .keychainTemporarilyUnavailable(legacyEntry: legacyEntry)
case .invalid:
KeychainCacheStore.clear(key: key)
case .missing:
break
}
guard scope == nil else { return .authoritative(nil) }
return .authoritative(self.migrateLegacyEntryIfNeededLocked(provider: provider))
}
} catch {
self.log.error("Cookie cache observation lock failed: \(error)")
return .keychainTemporarilyUnavailable(legacyEntry: nil)
}
}
/// Stores against a state captured by `observeForConditionalMutation`. A transient initial
/// Keychain failure may proceed only after Keychain recovers as missing and legacy state is unchanged.
@discardableResult
static func storeIfObservationCurrent(
provider: UsageProvider,
scope: Scope? = nil,
expected: ConditionalMutationObservation,
cookieHeader: String,
sourceLabel: String,
now: Date = Date()) -> Bool
{
let trimmed = cookieHeader.trimmingCharacters(in: .whitespacesAndNewlines)
guard let normalized = CookieHeaderNormalizer.normalize(trimmed), !normalized.isEmpty else { return false }
let entry = Entry(cookieHeader: normalized, storedAt: now, sourceLabel: sourceLabel)
do {
return try self.withLegacyMutationLock {
guard self.currentStateMatches(expected, provider: provider, scope: scope) else { return false }
return self.store(entry: entry, provider: provider, scope: scope, sourceLabel: sourceLabel)
}
} catch {
self.log.error("Cookie cache observed store lock failed: \(error)")
return false
}
}
private static func currentStateMatches(
_ expected: ConditionalMutationObservation,
provider: UsageProvider,
scope: Scope?) -> Bool
{
switch expected {
case let .authoritative(entry):
return self.currentEntryMatches(entry, provider: provider, scope: scope)
case let .keychainTemporarilyUnavailable(expectedLegacyEntry):
let key = self.key(for: provider, scope: scope)
guard case .missing = KeychainCacheStore.load(key: key, as: Entry.self) else { return false }
guard scope == nil else { return true }
return self.optionalEntriesMatch(self.loadLegacyEntry(for: provider), expectedLegacyEntry)
}
}
private static func optionalEntriesMatch(_ current: Entry?, _ expected: Entry?) -> Bool {
switch (current, expected) {
case (nil, nil): true
case let (current?, expected?): self.entriesMatch(current, expected)
default: false
}
}
}
@@ -0,0 +1,92 @@
import Foundation
public enum CookieHeaderNormalizer {
private static let headerPatterns: [String] = [
#"(?i)-H\s*'Cookie:\s*([^']+)'"#,
#"(?i)-H\s*\"Cookie:\s*([^\"]+)\""#,
#"(?i)\bcookie:\s*'([^']+)'"#,
#"(?i)\bcookie:\s*\"([^\"]+)\""#,
#"(?i)\bcookie:\s*([^\r\n]+)"#,
#"(?i)(?:^|\s)(?:--cookie|-b)\s*'([^']+)'"#,
#"(?i)(?:^|\s)(?:--cookie|-b)\s*\"([^\"]+)\""#,
#"(?i)(?:^|\s)-b([^\s=]+=[^\s]+)"#,
#"(?i)(?:^|\s)(?:--cookie|-b)\s+([^\s]+)"#,
]
public static func normalize(_ raw: String?) -> String? {
guard var value = raw?.trimmingCharacters(in: .whitespacesAndNewlines), !value.isEmpty else {
return nil
}
if let extracted = self.extractHeader(from: value) {
value = extracted
}
value = self.stripCookiePrefix(value)
value = self.stripWrappingQuotes(value)
value = value.trimmingCharacters(in: .whitespacesAndNewlines)
return value.isEmpty ? nil : value
}
public static func pairs(from raw: String) -> [(name: String, value: String)] {
guard let normalized = self.normalize(raw) else { return [] }
var results: [(name: String, value: String)] = []
results.reserveCapacity(6)
for part in normalized.split(separator: ";") {
let trimmed = part.trimmingCharacters(in: .whitespacesAndNewlines)
guard !trimmed.isEmpty,
let equalsIndex = trimmed.firstIndex(of: "=")
else {
continue
}
let name = trimmed[..<equalsIndex].trimmingCharacters(in: .whitespacesAndNewlines)
let value = trimmed[trimmed.index(after: equalsIndex)...]
.trimmingCharacters(in: .whitespacesAndNewlines)
guard !name.isEmpty else { continue }
results.append((name: String(name), value: String(value)))
}
return results
}
public static func filteredHeader(from raw: String?, allowedNames: Set<String>) -> String? {
let filtered = self.pairs(from: raw ?? "").filter { allowedNames.contains($0.name) }
guard !filtered.isEmpty else { return nil }
return filtered.map { "\($0.name)=\($0.value)" }.joined(separator: "; ")
}
private static func extractHeader(from raw: String) -> String? {
for pattern in self.headerPatterns {
guard let regex = try? NSRegularExpression(pattern: pattern, options: []) else { continue }
let range = NSRange(raw.startIndex..<raw.endIndex, in: raw)
guard let match = regex.firstMatch(in: raw, options: [], range: range),
match.numberOfRanges >= 2,
let captureRange = Range(match.range(at: 1), in: raw)
else {
continue
}
let captured = raw[captureRange].trimmingCharacters(in: .whitespacesAndNewlines)
if !captured.isEmpty { return String(captured) }
}
return nil
}
private static func stripCookiePrefix(_ raw: String) -> String {
let trimmed = raw.trimmingCharacters(in: .whitespacesAndNewlines)
guard trimmed.lowercased().hasPrefix("cookie:") else { return trimmed }
let idx = trimmed.index(trimmed.startIndex, offsetBy: "cookie:".count)
return String(trimmed[idx...]).trimmingCharacters(in: .whitespacesAndNewlines)
}
private static func stripWrappingQuotes(_ raw: String) -> String {
guard raw.count >= 2 else { return raw }
if (raw.hasPrefix("\"") && raw.hasSuffix("\"")) ||
(raw.hasPrefix("'") && raw.hasSuffix("'"))
{
return String(raw.dropFirst().dropLast())
}
return raw
}
}
@@ -0,0 +1,343 @@
import Foundation
public struct CopilotUsageResponse: Sendable, Decodable {
private struct AnyCodingKey: CodingKey {
let stringValue: String
let intValue: Int?
init?(stringValue: String) {
self.stringValue = stringValue
self.intValue = nil
}
init?(intValue: Int) {
self.stringValue = String(intValue)
self.intValue = intValue
}
}
public struct QuotaSnapshot: Sendable, Decodable {
public let entitlement: Double
public let remaining: Double
public let percentRemaining: Double
public let quotaId: String
public let hasPercentRemaining: Bool
public let unlimited: Bool
private let entitlementWasDecoded: Bool
private let remainingWasDecoded: Bool
public var usedPercent: Double {
max(0, 100 - self.percentRemaining)
}
public var overQuotaUsedPercent: Double? {
self.usedPercent > 100 ? self.usedPercent : nil
}
public var isPlaceholder: Bool {
if self.unlimited {
return false
}
if self.entitlement == 0,
self.remaining == 0,
self.percentRemaining == 0,
!self.hasPercentRemaining
{
return true
}
// An explicit zero-entitlement, zero-remaining snapshot carries no usable quota signal.
// GitHub returns this shape for token-based billing / Copilot Business seats,
// sometimes as percent_remaining=100 with a non-empty quota_id, which would
// otherwise render as a misleading "0% used" (100 - 100). Treat it as a
// placeholder so the usual handling drops it instead of showing fake usage.
return self.entitlementWasDecoded && self.remainingWasDecoded && self.entitlement == 0 && self
.remaining == 0
}
private enum CodingKeys: String, CodingKey {
case entitlement
case remaining
case percentRemaining = "percent_remaining"
case quotaId = "quota_id"
case unlimited
}
public init(
entitlement: Double,
remaining: Double,
percentRemaining: Double,
quotaId: String,
hasPercentRemaining: Bool = true,
unlimited: Bool = false)
{
self.entitlement = entitlement
self.remaining = remaining
self.percentRemaining = unlimited ? 100 : percentRemaining
self.quotaId = quotaId
self.hasPercentRemaining = unlimited || hasPercentRemaining
self.unlimited = unlimited
self.entitlementWasDecoded = true
self.remainingWasDecoded = true
}
public init(from decoder: any Decoder) throws {
let container = try decoder.container(keyedBy: CodingKeys.self)
let decodedEntitlement = Self.decodeNumberIfPresent(container: container, key: .entitlement)
let decodedRemaining = Self.decodeNumberIfPresent(container: container, key: .remaining)
self.entitlement = decodedEntitlement ?? 0
self.remaining = decodedRemaining ?? 0
self.entitlementWasDecoded = decodedEntitlement != nil
self.remainingWasDecoded = decodedRemaining != nil
let decodedUnlimited = try container.decodeIfPresent(Bool.self, forKey: .unlimited) ?? false
let decodedPercent = Self.decodeNumberIfPresent(container: container, key: .percentRemaining)
if decodedUnlimited {
self.percentRemaining = 100
self.hasPercentRemaining = true
} else if let decodedPercent {
self.percentRemaining = decodedPercent
self.hasPercentRemaining = true
} else if let entitlement = decodedEntitlement,
entitlement > 0,
let remaining = decodedRemaining
{
let derived = (remaining / entitlement) * 100
self.percentRemaining = derived
self.hasPercentRemaining = true
} else {
// Without percent_remaining and both inputs for derivation, the percent is unknown.
self.percentRemaining = 0
self.hasPercentRemaining = false
}
self.quotaId = try container.decodeIfPresent(String.self, forKey: .quotaId) ?? ""
self.unlimited = decodedUnlimited
}
private static func decodeNumberIfPresent(
container: KeyedDecodingContainer<CodingKeys>,
key: CodingKeys) -> Double?
{
if let value = try? container.decodeIfPresent(Double.self, forKey: key) {
return value
}
if let value = try? container.decodeIfPresent(Int.self, forKey: key) {
return Double(value)
}
if let value = try? container.decodeIfPresent(String.self, forKey: key) {
return Double(value)
}
return nil
}
}
public struct QuotaCounts: Sendable, Decodable {
public let chat: Double?
public let completions: Double?
private enum CodingKeys: String, CodingKey {
case chat
case completions
}
public init(chat: Double?, completions: Double?) {
self.chat = chat
self.completions = completions
}
public init(from decoder: any Decoder) throws {
let container = try decoder.container(keyedBy: CodingKeys.self)
self.chat = Self.decodeNumberIfPresent(container: container, key: .chat)
self.completions = Self.decodeNumberIfPresent(container: container, key: .completions)
}
private static func decodeNumberIfPresent(
container: KeyedDecodingContainer<CodingKeys>,
key: CodingKeys) -> Double?
{
if let value = try? container.decodeIfPresent(Double.self, forKey: key) {
return value
}
if let value = try? container.decodeIfPresent(Int.self, forKey: key) {
return Double(value)
}
if let value = try? container.decodeIfPresent(String.self, forKey: key) {
return Double(value)
}
return nil
}
}
public struct QuotaSnapshots: Sendable, Decodable {
public let premiumInteractions: QuotaSnapshot?
public let chat: QuotaSnapshot?
private enum CodingKeys: String, CodingKey {
case premiumInteractions = "premium_interactions"
case chat
}
public init(premiumInteractions: QuotaSnapshot?, chat: QuotaSnapshot?) {
self.premiumInteractions = premiumInteractions
self.chat = chat
}
public init(from decoder: any Decoder) throws {
let container = try decoder.container(keyedBy: CodingKeys.self)
var premium = try container.decodeIfPresent(QuotaSnapshot.self, forKey: .premiumInteractions)
var chat = try container.decodeIfPresent(QuotaSnapshot.self, forKey: .chat)
if premium?.isPlaceholder == true {
premium = nil
}
if chat?.isPlaceholder == true {
chat = nil
}
if premium == nil || chat == nil {
let dynamic = try decoder.container(keyedBy: AnyCodingKey.self)
var fallbackPremium: QuotaSnapshot?
var fallbackChat: QuotaSnapshot?
var firstUsable: QuotaSnapshot?
for key in dynamic.allKeys {
let value: QuotaSnapshot
do {
guard let decoded = try dynamic.decodeIfPresent(QuotaSnapshot.self, forKey: key) else {
continue
}
guard !decoded.isPlaceholder else { continue }
value = decoded
} catch {
continue
}
let name = key.stringValue.lowercased()
if firstUsable == nil {
firstUsable = value
}
if fallbackChat == nil, name.contains("chat") {
fallbackChat = value
continue
}
if fallbackPremium == nil,
name.contains("premium") || name.contains("completion") || name.contains("code")
{
fallbackPremium = value
}
}
if premium == nil {
premium = fallbackPremium
}
if chat == nil {
chat = fallbackChat
}
if premium == nil, chat == nil {
// If keys are unfamiliar, still expose one usable quota instead of failing.
chat = firstUsable
}
}
self.premiumInteractions = premium
self.chat = chat
}
}
public let quotaSnapshots: QuotaSnapshots
public let copilotPlan: String
public let tokenBasedBilling: Bool
public let assignedDate: String?
public let quotaResetDate: String?
private enum CodingKeys: String, CodingKey {
case quotaSnapshots = "quota_snapshots"
case copilotPlan = "copilot_plan"
case tokenBasedBilling = "token_based_billing"
case assignedDate = "assigned_date"
case quotaResetDate = "quota_reset_date"
case monthlyQuotas = "monthly_quotas"
case limitedUserQuotas = "limited_user_quotas"
}
public init(
quotaSnapshots: QuotaSnapshots,
copilotPlan: String,
tokenBasedBilling: Bool = false,
assignedDate: String?,
quotaResetDate: String?)
{
self.quotaSnapshots = quotaSnapshots
self.copilotPlan = copilotPlan
self.tokenBasedBilling = tokenBasedBilling
self.assignedDate = assignedDate
self.quotaResetDate = quotaResetDate
}
public init(from decoder: any Decoder) throws {
let container = try decoder.container(keyedBy: CodingKeys.self)
let directSnapshots = try container.decodeIfPresent(QuotaSnapshots.self, forKey: .quotaSnapshots)
let monthlyQuotas = try container.decodeIfPresent(QuotaCounts.self, forKey: .monthlyQuotas)
let limitedUserQuotas = try container.decodeIfPresent(QuotaCounts.self, forKey: .limitedUserQuotas)
let monthlyLimitedSnapshots = Self.makeQuotaSnapshots(monthly: monthlyQuotas, limited: limitedUserQuotas)
let premium = Self.usableQuotaSnapshot(from: directSnapshots?.premiumInteractions) ??
Self.usableQuotaSnapshot(from: monthlyLimitedSnapshots?.premiumInteractions)
let chat = Self.usableQuotaSnapshot(from: directSnapshots?.chat) ??
Self.usableQuotaSnapshot(from: monthlyLimitedSnapshots?.chat)
if premium != nil || chat != nil {
self.quotaSnapshots = QuotaSnapshots(premiumInteractions: premium, chat: chat)
} else {
self.quotaSnapshots = directSnapshots ?? QuotaSnapshots(premiumInteractions: nil, chat: nil)
}
self.copilotPlan = try container.decodeIfPresent(String.self, forKey: .copilotPlan) ?? "unknown"
self.tokenBasedBilling = try container.decodeIfPresent(Bool.self, forKey: .tokenBasedBilling) ?? false
self.assignedDate = try container.decodeIfPresent(String.self, forKey: .assignedDate)
self.quotaResetDate = try container.decodeIfPresent(String.self, forKey: .quotaResetDate)
}
private static func makeQuotaSnapshots(monthly: QuotaCounts?, limited: QuotaCounts?) -> QuotaSnapshots? {
let premium = Self.makeQuotaSnapshot(
monthly: monthly?.completions,
limited: limited?.completions,
quotaID: "completions")
let chat = Self.makeQuotaSnapshot(
monthly: monthly?.chat,
limited: limited?.chat,
quotaID: "chat")
guard premium != nil || chat != nil else { return nil }
return QuotaSnapshots(premiumInteractions: premium, chat: chat)
}
private static func makeQuotaSnapshot(monthly: Double?, limited: Double?, quotaID: String) -> QuotaSnapshot? {
guard monthly != nil || limited != nil else { return nil }
guard let monthly else {
// Without a monthly denominator, avoid fabricating a misleading percentage.
return nil
}
guard let limited else {
// Without the limited/remaining value, usage is unknown.
return nil
}
let entitlement = max(0, monthly)
guard entitlement > 0 else {
// A zero denominator cannot produce a meaningful percentage.
return nil
}
let remaining = max(0, limited)
let percentRemaining = max(0, min(100, (remaining / entitlement) * 100))
return QuotaSnapshot(
entitlement: entitlement,
remaining: remaining,
percentRemaining: percentRemaining,
quotaId: quotaID)
}
private static func usableQuotaSnapshot(from snapshot: QuotaSnapshot?) -> QuotaSnapshot? {
guard let snapshot, !snapshot.isPlaceholder, snapshot.hasPercentRemaining else {
return nil
}
return snapshot
}
}
+738
View File
@@ -0,0 +1,738 @@
import Foundation
public enum CostUsageError: LocalizedError, Sendable {
case unsupportedProvider(UsageProvider)
case timedOut(seconds: Int)
public var errorDescription: String? {
switch self {
case let .unsupportedProvider(provider):
return "Cost summary is not supported for \(provider.rawValue)."
case let .timedOut(seconds):
if seconds >= 60, seconds % 60 == 0 {
return "Cost refresh timed out after \(seconds / 60)m."
}
return "Cost refresh timed out after \(seconds)s."
}
}
}
public struct CostUsageFetcher: Sendable {
package struct CachedCodexTokenSnapshotResult: Sendable {
package let snapshot: CostUsageTokenSnapshot
package let lastRefreshAt: Date?
}
private let scannerOptions: CostUsageScanner.Options?
public init(cacheRoot: URL? = nil) {
self.scannerOptions = cacheRoot.map { CostUsageScanner.Options(cacheRoot: $0) }
}
init(scannerOptions: CostUsageScanner.Options) {
self.scannerOptions = scannerOptions
}
public func loadCachedCodexTokenSnapshot(
now: Date = Date(),
codexHomePath: String? = nil,
historyDays: Int = 30) async -> CostUsageTokenSnapshot?
{
await Self.loadCachedCodexTokenSnapshot(
now: now,
codexHomePath: codexHomePath,
historyDays: historyDays,
scannerOptions: self.scannerOptionsOverride())
}
package func loadCachedCodexTokenSnapshotResult(
now: Date = Date(),
codexHomePath: String? = nil,
historyDays: Int = 30) async -> CachedCodexTokenSnapshotResult?
{
await Self.loadCachedCodexTokenSnapshotResult(
now: now,
codexHomePath: codexHomePath,
historyDays: historyDays,
scannerOptions: self.scannerOptionsOverride())
}
public func loadTokenSnapshot(
provider: UsageProvider,
environment: [String: String] = ProcessInfo.processInfo.environment,
now: Date = Date(),
forceRefresh: Bool = false,
allowVertexClaudeFallback: Bool = false,
codexHomePath: String? = nil,
historyDays: Int = 30,
refreshPricingInBackground: Bool = true) async throws -> CostUsageTokenSnapshot
{
try await Self.loadTokenSnapshot(
provider: provider,
environment: environment,
now: now,
forceRefresh: forceRefresh,
allowVertexClaudeFallback: allowVertexClaudeFallback,
codexHomePath: codexHomePath,
historyDays: historyDays,
refreshPricingInBackground: refreshPricingInBackground,
bypassScannerDebounce: false,
scannerOptions: self.scannerOptionsOverride())
}
package func loadTokenSnapshot(
provider: UsageProvider,
environment: [String: String] = ProcessInfo.processInfo.environment,
now: Date = Date(),
forceRefresh: Bool = false,
allowVertexClaudeFallback: Bool = false,
codexHomePath: String? = nil,
historyDays: Int = 30,
refreshPricingInBackground: Bool = true,
bypassScannerDebounce: Bool) async throws -> CostUsageTokenSnapshot
{
try await Self.loadTokenSnapshot(
provider: provider,
environment: environment,
now: now,
forceRefresh: forceRefresh,
allowVertexClaudeFallback: allowVertexClaudeFallback,
codexHomePath: codexHomePath,
historyDays: historyDays,
refreshPricingInBackground: refreshPricingInBackground,
bypassScannerDebounce: bypassScannerDebounce,
scannerOptions: self.scannerOptionsOverride())
}
@available(*, deprecated, message: "Codex token-cost scans are uncapped; this limit is ignored.")
public func loadTokenSnapshot(
provider: UsageProvider,
environment: [String: String] = ProcessInfo.processInfo.environment,
now: Date = Date(),
forceRefresh: Bool = false,
allowVertexClaudeFallback: Bool = false,
codexHomePath: String? = nil,
historyDays: Int = 30,
refreshPricingInBackground: Bool = true,
automaticCodexScanByteLimit _: Int64?) async throws -> CostUsageTokenSnapshot
{
try await self.loadTokenSnapshot(
provider: provider,
environment: environment,
now: now,
forceRefresh: forceRefresh,
allowVertexClaudeFallback: allowVertexClaudeFallback,
codexHomePath: codexHomePath,
historyDays: historyDays,
refreshPricingInBackground: refreshPricingInBackground)
}
private func scannerOptionsOverride() -> CostUsageScanner.Options? {
self.scannerOptions
}
static func loadTokenSnapshot(
provider: UsageProvider,
environment: [String: String] = ProcessInfo.processInfo.environment,
now: Date = Date(),
forceRefresh: Bool = false,
allowVertexClaudeFallback: Bool = false,
codexHomePath: String? = nil,
historyDays: Int = 30,
refreshPricingInBackground: Bool = true,
bypassScannerDebounce: Bool = false,
scannerOptions overrideScannerOptions: CostUsageScanner.Options? = nil,
piScannerOptions overridePiScannerOptions: PiSessionCostScanner
.Options? = nil,
modelsDevClient: ModelsDevClient = ModelsDevClient(),
retryUnknownPricing: Bool = true) async throws -> CostUsageTokenSnapshot
{
guard provider == .codex || provider == .claude || provider == .vertexai || provider == .bedrock else {
throw CostUsageError.unsupportedProvider(provider)
}
let until = now
let clampedHistoryDays = max(1, min(365, historyDays))
// Rolling window is inclusive, so a 30-day display starts 29 days before `now`.
let since = Calendar.current.date(byAdding: .day, value: -(clampedHistoryDays - 1), to: now) ?? now
if provider == .bedrock {
let daily = try await Self.loadBedrockDailyReport(
environment: environment,
since: since,
until: until)
return Self.tokenSnapshot(
from: daily,
now: now,
historyDays: clampedHistoryDays,
useCurrentLocalDayForSession: false)
}
var options = overrideScannerOptions ?? CostUsageScanner.Options()
if provider == .codex,
let codexHomePath = codexHomePath?.trimmingCharacters(in: .whitespacesAndNewlines),
!codexHomePath.isEmpty
{
options.codexSessionsRoot = URL(fileURLWithPath: codexHomePath, isDirectory: true)
.appendingPathComponent("sessions", isDirectory: true)
}
if retryUnknownPricing, provider == .codex || provider == .claude {
let pricingCacheRoot = options.cacheRoot
if refreshPricingInBackground {
Task.detached(priority: .utility) {
await ModelsDevPricingPipeline.refreshIfNeeded(
now: now,
cacheRoot: pricingCacheRoot,
client: modelsDevClient)
}
} else {
await ModelsDevPricingPipeline.refreshIfNeeded(
now: now,
cacheRoot: pricingCacheRoot,
client: modelsDevClient)
}
}
if provider == .vertexai {
options.claudeLogProviderFilter = allowVertexClaudeFallback ? .all : .vertexAIOnly
} else if provider == .claude {
options.claudeLogProviderFilter = .excludeVertexAI
}
if forceRefresh || bypassScannerDebounce {
options.refreshMinIntervalSeconds = 0
}
var resolvedPiOptions = overridePiScannerOptions ?? PiSessionCostScanner.Options()
if resolvedPiOptions.cacheRoot == nil {
resolvedPiOptions.cacheRoot = options.cacheRoot
}
if forceRefresh || bypassScannerDebounce {
resolvedPiOptions.refreshMinIntervalSeconds = 0
}
let piOptions = resolvedPiOptions
try Task.checkCancellation()
// The corpus scans below are synchronous and can run for minutes on large session
// archives. They execute on the dedicated scan queue so they never occupy a cooperative
// pool thread; CostUsageScanExecutor bridges this task's cancellation into the
// scanner-level checks.
let scanOptions = options
let scanResult = try await CostUsageScanExecutor.run { checkCancellation in
var daily = try CostUsageScanner.loadDailyReportCancellable(
provider: provider,
since: since,
until: until,
now: now,
options: scanOptions,
checkCancellation: checkCancellation)
try checkCancellation()
if provider == .vertexai,
!allowVertexClaudeFallback,
scanOptions.claudeLogProviderFilter == .vertexAIOnly,
daily.data.isEmpty
{
var fallback = scanOptions
fallback.claudeLogProviderFilter = .all
daily = try CostUsageScanner.loadDailyReportCancellable(
provider: provider,
since: since,
until: until,
now: now,
options: fallback,
checkCancellation: checkCancellation)
try checkCancellation()
}
var projects: [CostUsageProjectBreakdown] = []
var piDaily: CostUsageDailyReport?
if provider == .codex {
let cache = CostUsageCacheIO.load(provider: .codex, cacheRoot: scanOptions.cacheRoot)
projects = CostUsageScanner.buildCodexProjectBreakdownsFromCache(
cache: cache,
range: CostUsageScanner.CostUsageDayRange(since: since, until: until),
modelsDevCacheRoot: scanOptions.cacheRoot)
}
if provider == .codex || provider == .claude {
let piReport = try PiSessionCostScanner.loadDailyReportCancellable(
provider: provider,
since: since,
until: until,
now: now,
options: piOptions,
checkCancellation: checkCancellation)
try checkCancellation()
if provider == .codex {
piDaily = piReport
}
daily = CostUsageDailyReport.merged([daily, piReport])
}
if provider == .codex {
projects = Self.mergedProjectBreakdowns(
projects + [piDaily.flatMap(Self.unknownProjectBreakdown(from:))].compactMap(\.self))
}
return (daily: daily, projects: projects)
}
if retryUnknownPricing,
let request = Self.unknownPricingRefreshRequest(
provider: provider,
daily: scanResult.daily,
now: now,
cacheRoot: options.cacheRoot,
client: modelsDevClient),
await Self.refreshUnknownPricingIfNeeded(request, inBackground: refreshPricingInBackground)
{
return try await self.loadTokenSnapshot(
provider: provider,
environment: environment,
now: now,
forceRefresh: forceRefresh,
allowVertexClaudeFallback: allowVertexClaudeFallback,
codexHomePath: codexHomePath,
historyDays: historyDays,
refreshPricingInBackground: false,
scannerOptions: options,
piScannerOptions: piOptions,
modelsDevClient: modelsDevClient,
retryUnknownPricing: false)
}
return Self.tokenSnapshot(
from: scanResult.daily,
now: now,
historyDays: clampedHistoryDays,
projects: scanResult.projects)
}
private struct UnknownPricingRefreshRequest: Sendable {
let providerID: String
let modelIDs: Set<String>
let now: Date
let cacheRoot: URL?
let client: ModelsDevClient
}
private static func unknownPricingRefreshRequest(
provider: UsageProvider,
daily: CostUsageDailyReport,
now: Date,
cacheRoot: URL?,
client: ModelsDevClient) -> UnknownPricingRefreshRequest?
{
guard provider == .codex || provider == .claude else { return nil }
let unknownModelIDs = Set(daily.data.flatMap { entry in
entry.modelBreakdowns?.compactMap { breakdown -> String? in
guard breakdown.costUSD == nil else { return nil }
if provider == .codex,
CostUsagePricing.isCodexUnattributedModel(breakdown.modelName)
{
return nil
}
return breakdown.modelName
} ?? []
})
guard !unknownModelIDs.isEmpty else { return nil }
return UnknownPricingRefreshRequest(
providerID: provider == .codex ? "openai" : "anthropic",
modelIDs: unknownModelIDs,
now: now,
cacheRoot: cacheRoot,
client: client)
}
private static func refreshUnknownPricingIfNeeded(
_ request: UnknownPricingRefreshRequest,
inBackground: Bool) async -> Bool
{
if inBackground {
Task.detached(priority: .utility) {
_ = await ModelsDevPricingPipeline.refreshForUnknownModelsIfNeeded(
providerID: request.providerID,
modelIDs: request.modelIDs,
now: request.now,
cacheRoot: request.cacheRoot,
client: request.client)
}
return false
}
return await ModelsDevPricingPipeline.refreshForUnknownModelsIfNeeded(
providerID: request.providerID,
modelIDs: request.modelIDs,
now: request.now,
cacheRoot: request.cacheRoot,
client: request.client) == .pricingAvailable
}
static func loadCachedCodexTokenSnapshot(
now: Date = Date(),
codexHomePath: String? = nil,
historyDays: Int = 30,
scannerOptions overrideScannerOptions: CostUsageScanner.Options? = nil) async -> CostUsageTokenSnapshot?
{
await self.loadCachedCodexTokenSnapshotResult(
now: now,
codexHomePath: codexHomePath,
historyDays: historyDays,
scannerOptions: overrideScannerOptions)?.snapshot
}
static func loadCachedCodexTokenSnapshotResult(
now: Date = Date(),
codexHomePath: String? = nil,
historyDays: Int = 30,
scannerOptions overrideScannerOptions: CostUsageScanner.Options? = nil) async
-> CachedCodexTokenSnapshotResult?
{
if let codexHomePath = codexHomePath?.trimmingCharacters(in: .whitespacesAndNewlines),
!codexHomePath.isEmpty
{
return nil
}
// Decoding the persisted scan cache parses multi-megabyte JSON; keep it off the
// cooperative pool alongside the scans themselves.
let cachedSnapshot: CachedCodexTokenSnapshotResult?? = try? await CostUsageScanExecutor.run { _ in
let clampedHistoryDays = max(1, min(365, historyDays))
let until = now
let since = Calendar.current.date(byAdding: .day, value: -(clampedHistoryDays - 1), to: now) ?? now
let range = CostUsageScanner.CostUsageDayRange(since: since, until: until)
let options = overrideScannerOptions ?? CostUsageScanner.Options()
let cache = CostUsageCacheIO.load(provider: .codex, cacheRoot: options.cacheRoot)
var reports: [CostUsageDailyReport] = []
var projects: [CostUsageProjectBreakdown] = []
// Raw inputs for the derived result fields below: the native cache's own scan
// time, every constituent scan time, and whether a second source joined the merge.
var nativeScanAt: Date?
var scanTimes: [Date] = []
var piMerged = false
if !cache.days.isEmpty,
cache.roots == CostUsageScanner.codexRootsFingerprint(options: options),
!CostUsageScanner.requestedWindowExpandsCache(range: range, cache: cache)
{
let daily = CostUsageScanner.buildCodexReportFromCache(
cache: cache,
range: range,
modelsDevCacheRoot: options.cacheRoot)
if !daily.data.isEmpty {
reports.append(daily)
if cache.lastScanUnixMs > 0 {
let scanAt = Date(timeIntervalSince1970: TimeInterval(cache.lastScanUnixMs) / 1000)
nativeScanAt = scanAt
scanTimes.append(scanAt)
}
if cache.codexProjectMetadataVersion == CostUsageScanner.codexProjectMetadataVersion {
projects.append(contentsOf: CostUsageScanner.buildCodexProjectBreakdownsFromCache(
cache: cache,
range: range,
modelsDevCacheRoot: options.cacheRoot))
}
}
}
if let piResult = PiSessionCostScanner.loadCachedDailyReportResult(
provider: .codex,
since: since,
until: until,
now: now,
cacheRoot: options.cacheRoot)
{
reports.append(piResult.report)
piMerged = true
if let piLastScanAt = piResult.lastScanAt {
scanTimes.append(piLastScanAt)
}
if let piProject = Self.unknownProjectBreakdown(from: piResult.report) {
projects.append(piProject)
}
}
guard !reports.isEmpty else { return nil }
// updatedAt keeps the caches' real (oldest) scan time; stamping the hydration time
// would let stale token rows inherit app-start freshness (#1964). lastRefreshAt
// drives TTL suppression and stays native-only: a merged load must never delay a
// rescan on the strength of another source's scan.
return CachedCodexTokenSnapshotResult(
snapshot: Self.tokenSnapshot(
from: CostUsageDailyReport.merged(reports),
now: now,
historyDays: clampedHistoryDays,
projects: Self.mergedProjectBreakdowns(projects),
updatedAt: scanTimes.min()),
lastRefreshAt: piMerged ? nil : nativeScanAt)
}
return cachedSnapshot.flatMap(\.self)
}
private static func loadBedrockDailyReport(
environment: [String: String],
since: Date,
until: Date) async throws -> CostUsageDailyReport
{
let resolved = try await BedrockCredentialResolver.resolve(environment: environment)
return try await BedrockUsageFetcher.fetchDailyReport(
credentials: resolved.credentials,
since: since,
until: until,
environment: environment)
}
static func tokenSnapshot(
from daily: CostUsageDailyReport,
now: Date,
historyDays: Int = 30,
useCurrentLocalDayForSession: Bool = true,
projects: [CostUsageProjectBreakdown] = [],
updatedAt: Date? = nil) -> CostUsageTokenSnapshot
{
let sessionEntry = useCurrentLocalDayForSession
? CostUsageTokenSnapshot.entry(in: daily.data, forLocalDayContaining: now)
: CostUsageTokenSnapshot.latestEntry(in: daily.data)
let hasHistoricalRows = !daily.data.isEmpty
let sessionTokens: Int? = if let sessionEntry {
sessionEntry.totalTokens
} else if hasHistoricalRows {
0
} else {
nil
}
let sessionCostUSD: Double? = if let sessionEntry {
sessionEntry.costUSD
} else if hasHistoricalRows {
0
} else {
nil
}
// Prefer summary totals when present; fall back to summing daily entries.
let totalFromSummary = daily.summary?.totalCostUSD
let totalFromEntries = daily.data.compactMap(\.costUSD).reduce(0, +)
let last30DaysCostUSD = totalFromSummary ?? (totalFromEntries > 0 ? totalFromEntries : nil)
let totalTokensFromSummary = daily.summary?.totalTokens
let totalTokensFromEntries = daily.data.compactMap(\.totalTokens).reduce(0, +)
let last30DaysTokens = totalTokensFromSummary ?? (totalTokensFromEntries > 0 ? totalTokensFromEntries : nil)
return CostUsageTokenSnapshot(
sessionTokens: sessionTokens,
sessionCostUSD: sessionCostUSD,
last30DaysTokens: last30DaysTokens,
last30DaysCostUSD: last30DaysCostUSD,
historyDays: historyDays,
daily: daily.data,
projects: projects,
updatedAt: updatedAt ?? now)
}
private static func unknownProjectBreakdown(from daily: CostUsageDailyReport) -> CostUsageProjectBreakdown? {
guard !daily.data.isEmpty else { return nil }
return CostUsageProjectBreakdown(
name: CostUsageProjectBreakdown.unknownProjectName,
path: nil,
totalTokens: daily.summary?.totalTokens,
totalCostUSD: daily.summary?.totalCostUSD,
daily: daily.data,
modelBreakdowns: self.projectModelBreakdowns(from: daily.data),
sources: [
CostUsageProjectSourceBreakdown(
name: CostUsageProjectBreakdown.unknownProjectName,
path: nil,
totalTokens: daily.summary?.totalTokens,
totalCostUSD: daily.summary?.totalCostUSD,
daily: daily.data,
modelBreakdowns: self.projectModelBreakdowns(from: daily.data)),
])
}
private static func mergedProjectBreakdowns(
_ projects: [CostUsageProjectBreakdown]) -> [CostUsageProjectBreakdown]
{
var dailyByPath: [String: [CostUsageDailyReport]] = [:]
var namesByPath: [String: String] = [:]
var sourceDailyByProjectPath: [String: [String: [CostUsageDailyReport]]] = [:]
var sourceNamesByProjectPath: [String: [String: String]] = [:]
for project in projects {
let key = project.path ?? ""
namesByPath[key] = project.name
dailyByPath[key, default: []].append(CostUsageDailyReport(data: project.daily, summary: nil))
let sources = project.sources.isEmpty
? [
CostUsageProjectSourceBreakdown(
name: project.name,
path: project.path,
totalTokens: project.totalTokens,
totalCostUSD: project.totalCostUSD,
daily: project.daily,
modelBreakdowns: project.modelBreakdowns),
]
: project.sources
for source in sources {
let sourceKey = source.path ?? ""
sourceNamesByProjectPath[key, default: [:]][sourceKey] = source.name
sourceDailyByProjectPath[key, default: [:]][sourceKey, default: []]
.append(CostUsageDailyReport(data: source.daily, summary: nil))
}
}
return dailyByPath.map { key, reports in
let merged = CostUsageDailyReport.merged(reports)
return CostUsageProjectBreakdown(
name: namesByPath[key] ?? CostUsageProjectBreakdown.unknownProjectName,
path: key.isEmpty ? nil : key,
totalTokens: merged.summary?.totalTokens,
totalCostUSD: merged.summary?.totalCostUSD,
daily: merged.data,
modelBreakdowns: Self.projectModelBreakdowns(from: merged.data),
sources: Self.mergedProjectSources(
sourceDailyByPath: sourceDailyByProjectPath[key] ?? [:],
sourceNamesByPath: sourceNamesByProjectPath[key] ?? [:]))
}
.sorted { lhs, rhs in
let lhsCost = lhs.totalCostUSD ?? -1
let rhsCost = rhs.totalCostUSD ?? -1
if lhsCost != rhsCost {
return lhsCost > rhsCost
}
let lhsTokens = lhs.totalTokens ?? -1
let rhsTokens = rhs.totalTokens ?? -1
if lhsTokens != rhsTokens {
return lhsTokens > rhsTokens
}
return lhs.name.localizedStandardCompare(rhs.name) == .orderedAscending
}
}
private static func mergedProjectSources(
sourceDailyByPath: [String: [CostUsageDailyReport]],
sourceNamesByPath: [String: String]) -> [CostUsageProjectSourceBreakdown]
{
sourceDailyByPath.map { key, reports in
let merged = CostUsageDailyReport.merged(reports)
return CostUsageProjectSourceBreakdown(
name: sourceNamesByPath[key] ?? CostUsageProjectBreakdown.unknownProjectName,
path: key.isEmpty ? nil : key,
totalTokens: merged.summary?.totalTokens,
totalCostUSD: merged.summary?.totalCostUSD,
daily: merged.data,
modelBreakdowns: Self.projectModelBreakdowns(from: merged.data))
}
.sorted { lhs, rhs in
let lhsCost = lhs.totalCostUSD ?? -1
let rhsCost = rhs.totalCostUSD ?? -1
if lhsCost != rhsCost {
return lhsCost > rhsCost
}
let lhsTokens = lhs.totalTokens ?? -1
let rhsTokens = rhs.totalTokens ?? -1
if lhsTokens != rhsTokens {
return lhsTokens > rhsTokens
}
return lhs.name.localizedStandardCompare(rhs.name) == .orderedAscending
}
}
private struct ProjectBreakdownAccumulator {
var totalTokens = 0
var sawTotalTokens = false
var costUSD: Double = 0
var sawCost = false
mutating func add(_ breakdown: CostUsageDailyReport.ModelBreakdown) {
if let totalTokens = breakdown.totalTokens {
self.totalTokens += totalTokens
self.sawTotalTokens = true
}
if let costUSD = breakdown.costUSD {
self.costUSD += costUSD
self.sawCost = true
}
}
func build(modelName: String) -> CostUsageDailyReport.ModelBreakdown {
CostUsageDailyReport.ModelBreakdown(
modelName: modelName,
costUSD: self.sawCost ? self.costUSD : nil,
totalTokens: self.sawTotalTokens ? self.totalTokens : nil)
}
}
private static func projectModelBreakdowns(
from entries: [CostUsageDailyReport.Entry]) -> [CostUsageDailyReport.ModelBreakdown]?
{
var accumulators: [String: ProjectBreakdownAccumulator] = [:]
for entry in entries {
for breakdown in entry.modelBreakdowns ?? [] {
var accumulator = accumulators[breakdown.modelName] ?? ProjectBreakdownAccumulator()
accumulator.add(breakdown)
accumulators[breakdown.modelName] = accumulator
}
}
guard !accumulators.isEmpty else { return nil }
return accumulators.map { modelName, accumulator in
accumulator.build(modelName: modelName)
}
.sorted { lhs, rhs in
let lhsCost = lhs.costUSD ?? -1
let rhsCost = rhs.costUSD ?? -1
if lhsCost != rhsCost {
return lhsCost > rhsCost
}
let lhsTokens = lhs.totalTokens ?? -1
let rhsTokens = rhs.totalTokens ?? -1
if lhsTokens != rhsTokens {
return lhsTokens > rhsTokens
}
return lhs.modelName > rhs.modelName
}
}
static func selectCurrentSession(from sessions: [CostUsageSessionReport.Entry])
-> CostUsageSessionReport.Entry?
{
if sessions.isEmpty {
return nil
}
return sessions.max { lhs, rhs in
let lDate = CostUsageDateParser.parse(lhs.lastActivity) ?? .distantPast
let rDate = CostUsageDateParser.parse(rhs.lastActivity) ?? .distantPast
if lDate != rDate {
return lDate < rDate
}
let lCost = lhs.costUSD ?? -1
let rCost = rhs.costUSD ?? -1
if lCost != rCost {
return lCost < rCost
}
let lTokens = lhs.totalTokens ?? -1
let rTokens = rhs.totalTokens ?? -1
if lTokens != rTokens {
return lTokens < rTokens
}
return lhs.session < rhs.session
}
}
static func selectMostRecentMonth(from months: [CostUsageMonthlyReport.Entry])
-> CostUsageMonthlyReport.Entry?
{
if months.isEmpty {
return nil
}
return months.max { lhs, rhs in
let lDate = CostUsageDateParser.parseMonth(lhs.month) ?? .distantPast
let rDate = CostUsageDateParser.parseMonth(rhs.month) ?? .distantPast
if lDate != rDate {
return lDate < rDate
}
let lCost = lhs.costUSD ?? -1
let rCost = rhs.costUSD ?? -1
if lCost != rCost {
return lCost < rCost
}
let lTokens = lhs.totalTokens ?? -1
let rTokens = rhs.totalTokens ?? -1
if lTokens != rTokens {
return lTokens < rTokens
}
return lhs.month < rhs.month
}
}
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,130 @@
import Foundation
/// Cost-usage scans read and parse the full local session corpus synchronously and can run for
/// minutes on large archives. Executing that work inline on Swift's cooperative thread pool
/// starves every other async task in the process menus freeze while the main thread sits idle
/// and overlapping provider scans multiply both the pool pressure and the disk load. This
/// executor pins all corpus scans to a single serial utility queue off the cooperative pool, so
/// long scans cost one dedicated thread instead of the app's async runtime.
public enum CostUsageScanExecutor {
public static let queueLabel = "com.steipete.codexbar.cost-usage-scan"
private static let queue = DispatchQueue(label: queueLabel, qos: .utility)
private final class RunState<Value: Sendable>: @unchecked Sendable {
private enum Phase {
case initial
case queued
case running
case completed
}
private let lock = NSLock()
private var phase: Phase = .initial
private var cancellationRequested = false
private var continuation: CheckedContinuation<Value, Error>?
func install(_ continuation: CheckedContinuation<Value, Error>) -> Bool {
let shouldEnqueue: Bool
let shouldResumeCancellation: Bool
self.lock.lock()
if self.cancellationRequested {
self.phase = .completed
shouldEnqueue = false
shouldResumeCancellation = true
} else {
self.phase = .queued
self.continuation = continuation
shouldEnqueue = true
shouldResumeCancellation = false
}
self.lock.unlock()
if shouldResumeCancellation {
continuation.resume(throwing: CancellationError())
}
return shouldEnqueue
}
func begin() -> Bool {
self.lock.lock()
defer { self.lock.unlock() }
guard self.phase == .queued else { return false }
self.phase = .running
return true
}
func cancel() {
let continuation: CheckedContinuation<Value, Error>?
self.lock.lock()
self.cancellationRequested = true
if self.phase == .queued {
self.phase = .completed
continuation = self.continuation
self.continuation = nil
} else {
continuation = nil
}
self.lock.unlock()
continuation?.resume(throwing: CancellationError())
}
func checkCancellation() throws {
self.lock.lock()
let cancellationRequested = self.cancellationRequested
self.lock.unlock()
if cancellationRequested {
throw CancellationError()
}
}
func complete(with result: Result<Value, Error>) {
let continuation: CheckedContinuation<Value, Error>?
let resolvedResult: Result<Value, Error>
self.lock.lock()
guard self.phase == .running else {
self.lock.unlock()
return
}
self.phase = .completed
continuation = self.continuation
self.continuation = nil
resolvedResult = self.cancellationRequested ? .failure(CancellationError()) : result
self.lock.unlock()
continuation?.resume(with: resolvedResult)
}
}
/// Runs `work` on the serial scan queue and bridges Swift task cancellation into the
/// scanner's cooperative `checkCancellation` callbacks. Work that is still queued when the
/// awaiting task is cancelled resumes immediately with `CancellationError` instead of
/// waiting behind an in-flight scan.
public static func run<T: Sendable>(
_ work: @escaping @Sendable (_ checkCancellation: @escaping @Sendable () throws -> Void) throws -> T)
async throws -> T
{
try await self.run(on: self.queue, work)
}
static func run<T: Sendable>(
on queue: DispatchQueue,
_ work: @escaping @Sendable (_ checkCancellation: @escaping @Sendable () throws -> Void) throws -> T)
async throws -> T
{
let state = RunState<T>()
let checkCancellation: @Sendable () throws -> Void = {
try state.checkCancellation()
}
return try await withTaskCancellationHandler {
try await withCheckedThrowingContinuation { continuation in
guard state.install(continuation) else { return }
queue.async {
guard state.begin() else { return }
state.complete(with: Result { try work(checkCancellation) })
}
}
} onCancel: {
state.cancel()
}
}
}
+320
View File
@@ -0,0 +1,320 @@
import Foundation
#if canImport(CryptoKit)
import CryptoKit
#else
import Crypto
#endif
public struct CreditEvent: Identifiable, Equatable, Codable, Sendable {
public var id: UUID
public let date: Date
public let service: String
public let creditsUsed: Double
public init(id: UUID = UUID(), date: Date, service: String, creditsUsed: Double) {
self.id = id
self.date = date
self.service = service
self.creditsUsed = creditsUsed
}
private enum CodingKeys: String, CodingKey {
case id
case date
case service
case creditsUsed
}
public init(from decoder: Decoder) throws {
let container = try decoder.container(keyedBy: CodingKeys.self)
self.id = try container.decodeIfPresent(UUID.self, forKey: .id) ?? UUID()
self.date = try container.decode(Date.self, forKey: .date)
self.service = try container.decode(String.self, forKey: .service)
self.creditsUsed = try container.decode(Double.self, forKey: .creditsUsed)
}
public func encode(to encoder: Encoder) throws {
var container = encoder.container(keyedBy: CodingKeys.self)
try container.encode(self.id, forKey: .id)
try container.encode(self.date, forKey: .date)
try container.encode(self.service, forKey: .service)
try container.encode(self.creditsUsed, forKey: .creditsUsed)
}
}
public struct CreditsSnapshot: Equatable, Codable, Sendable {
public let remaining: Double
public let events: [CreditEvent]
public let updatedAt: Date
public let codexCreditLimit: CodexCreditLimitSnapshot?
public init(
remaining: Double,
events: [CreditEvent],
updatedAt: Date,
codexCreditLimit: CodexCreditLimitSnapshot? = nil)
{
self.remaining = remaining
self.events = events
self.updatedAt = updatedAt
self.codexCreditLimit = codexCreditLimit
}
}
public struct CodexCreditLimitSnapshot: Equatable, Codable, Sendable {
public let title: String
public let used: Double
public let limit: Double
public let remaining: Double
public let remainingPercent: Double
public let resetsAt: Date?
public let updatedAt: Date
public init(
title: String = "Monthly credit limit",
used: Double,
limit: Double,
remainingPercent: Double,
resetsAt: Date?,
updatedAt: Date)
{
self.title = title.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty
? "Monthly credit limit"
: title.trimmingCharacters(in: .whitespacesAndNewlines)
self.used = max(0, used)
self.limit = max(0, limit)
self.remaining = max(0, self.limit - self.used)
self.remainingPercent = min(100, max(0, remainingPercent))
self.resetsAt = resetsAt
self.updatedAt = updatedAt
}
public var usedPercent: Double {
min(100, max(0, 100 - self.remainingPercent))
}
}
public struct CodexRateLimitResetCreditsSnapshot: Equatable, Codable, Sendable {
public let credits: [CodexRateLimitResetCredit]
public let availableCount: Int
public let updatedAt: Date
public init(credits: [CodexRateLimitResetCredit], availableCount: Int, updatedAt: Date) {
self.credits = credits
self.availableCount = availableCount
self.updatedAt = updatedAt
}
public var nextExpiringAvailableCredit: CodexRateLimitResetCredit? {
self.availableInventory(at: self.updatedAt).nextExpiringCredit
}
public func availableInventory(at date: Date) -> CodexRateLimitResetCreditInventory {
CodexRateLimitResetCreditInventory(credits: self.credits, at: date)
}
public func availableCredits(at date: Date) -> [CodexRateLimitResetCredit] {
self.availableInventory(at: date).credits
}
}
public struct CodexRateLimitResetCreditInventory: Equatable, Sendable {
public let credits: [CodexRateLimitResetCredit]
public var count: Int {
self.credits.count
}
public var nextExpiringCredit: CodexRateLimitResetCredit? {
self.credits.first { $0.expiresAt != nil }
}
public init(credits: [CodexRateLimitResetCredit], at date: Date) {
self.credits = credits
.filter { credit in
credit.status == .available && (credit.expiresAt.map { $0 > date } ?? true)
}
.sorted { lhs, rhs in
switch (lhs.expiresAt, rhs.expiresAt) {
case let (lhsDate?, rhsDate?):
if lhsDate != rhsDate { return lhsDate < rhsDate }
case (_?, nil):
return true
case (nil, _?):
return false
case (nil, nil):
break
}
return lhs.id < rhs.id
}
}
}
public struct CodexRateLimitResetCredit: Equatable, Codable, Sendable, Identifiable {
private static let stableIDPrefix = "codex-reset-credit-v1-"
public let id: String
public let resetType: String
public let status: CodexRateLimitResetCreditStatus
public let grantedAt: Date
public let expiresAt: Date?
public let redeemStartedAt: Date?
public let redeemedAt: Date?
public let title: String?
public let description: String?
private enum CodingKeys: String, CodingKey {
case id
case resetType = "reset_type"
case status
case grantedAt = "granted_at"
case expiresAt = "expires_at"
case redeemStartedAt = "redeem_started_at"
case redeemedAt = "redeemed_at"
case title
case description
}
public init(
id: String,
resetType: String,
status: CodexRateLimitResetCreditStatus,
grantedAt: Date,
expiresAt: Date?,
redeemStartedAt: Date?,
redeemedAt: Date?,
title: String?,
description: String?)
{
self.id = Self.stableID(forProviderID: id)
self.resetType = resetType
self.status = status
self.grantedAt = grantedAt
self.expiresAt = expiresAt
self.redeemStartedAt = redeemStartedAt
self.redeemedAt = redeemedAt
self.title = title
self.description = description
}
public init(from decoder: Decoder) throws {
let container = try decoder.container(keyedBy: CodingKeys.self)
let decodedID = try container.decode(String.self, forKey: .id)
let resetType = try container.decode(String.self, forKey: .resetType)
let status = try container.decode(CodexRateLimitResetCreditStatus.self, forKey: .status)
let grantedAt = try container.decode(Date.self, forKey: .grantedAt)
let expiresAt = try container.decodeIfPresent(Date.self, forKey: .expiresAt)
let redeemStartedAt = try container.decodeIfPresent(Date.self, forKey: .redeemStartedAt)
let redeemedAt = try container.decodeIfPresent(Date.self, forKey: .redeemedAt)
let title = try container.decodeIfPresent(String.self, forKey: .title)
let description = try container.decodeIfPresent(String.self, forKey: .description)
if Self.isCanonicalStableID(decodedID) {
self.init(
persistedStableID: decodedID,
resetType: resetType,
status: status,
grantedAt: grantedAt,
expiresAt: expiresAt,
redeemStartedAt: redeemStartedAt,
redeemedAt: redeemedAt,
title: title,
description: description)
} else {
self.init(
id: decodedID,
resetType: resetType,
status: status,
grantedAt: grantedAt,
expiresAt: expiresAt,
redeemStartedAt: redeemStartedAt,
redeemedAt: redeemedAt,
title: title,
description: description)
}
}
static func stableID(forProviderID providerID: String) -> String {
let domainSeparatedValue = "com.steipete.CodexBar.reset-credit-id.v1\0\(providerID)"
let digest = SHA256.hash(data: Data(domainSeparatedValue.utf8))
.map { String(format: "%02x", $0) }
.joined()
return Self.stableIDPrefix + digest
}
private init(
persistedStableID: String,
resetType: String,
status: CodexRateLimitResetCreditStatus,
grantedAt: Date,
expiresAt: Date?,
redeemStartedAt: Date?,
redeemedAt: Date?,
title: String?,
description: String?)
{
precondition(Self.isCanonicalStableID(persistedStableID))
self.id = persistedStableID
self.resetType = resetType
self.status = status
self.grantedAt = grantedAt
self.expiresAt = expiresAt
self.redeemStartedAt = redeemStartedAt
self.redeemedAt = redeemedAt
self.title = title
self.description = description
}
private static func isCanonicalStableID(_ id: String) -> Bool {
guard id.hasPrefix(self.stableIDPrefix) else { return false }
let digest = id.dropFirst(Self.stableIDPrefix.count)
return digest.utf8.count == 64 && digest.utf8.allSatisfy { byte in
(48...57).contains(byte) || (97...102).contains(byte)
}
}
}
public enum CodexRateLimitResetCreditStatus: Equatable, Codable, Sendable {
case available
case redeeming
case redeemed
case expired
case unknown(String)
public init(from decoder: Decoder) throws {
let container = try decoder.singleValueContainer()
let value = try container.decode(String.self)
switch value {
case "available":
self = .available
case "redeeming":
self = .redeeming
case "redeemed":
self = .redeemed
case "expired":
self = .expired
default:
self = .unknown(value)
}
}
public func encode(to encoder: Encoder) throws {
var container = encoder.singleValueContainer()
try container.encode(self.rawValue)
}
public var rawValue: String {
switch self {
case .available:
"available"
case .redeeming:
"redeeming"
case .redeemed:
"redeemed"
case .expired:
"expired"
case let .unknown(value):
value
}
}
}
@@ -0,0 +1,5 @@
extension Double {
public func clamped(to range: ClosedRange<Double>) -> Double {
min(range.upperBound, max(range.lowerBound, self))
}
}
@@ -0,0 +1,5 @@
// Generated by Scripts/regenerate-codex-parser-hash.sh. Do not edit by hand.
enum CodexParserHash {
static let value = "cdef6eb9658a43e2"
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,61 @@
#if canImport(Darwin)
import Darwin
#elseif canImport(Glibc)
import Glibc
#elseif canImport(Musl)
import Musl
#endif
import Foundation
enum PosixSpawnFileActionsCloseFrom {
enum CloseFromError: LocalizedError {
case descriptorEnumerationFailed(String)
case actionFailed(Int32)
var errorDescription: String? {
switch self {
case let .descriptorEnumerationFailed(details):
"Could not enumerate /proc/self/fd: \(details)"
case let .actionFailed(code):
"Could not configure descriptor cleanup: \(String(cString: strerror(code))) (\(code))"
}
}
}
static func descriptorsToClose(
startingAt minimumFileDescriptor: Int32,
contentsOfDirectory: (String) throws -> [String] = FileManager.default.contentsOfDirectory(atPath:)) throws
-> [Int32]
{
let entries: [String]
do {
entries = try contentsOfDirectory("/proc/self/fd")
} catch {
throw CloseFromError.descriptorEnumerationFailed(error.localizedDescription)
}
return entries.compactMap(Int32.init)
.filter { $0 >= minimumFileDescriptor }
.sorted()
}
#if canImport(Glibc) || canImport(Musl)
static func addCloseFrom(
_ fileActions: inout posix_spawn_file_actions_t,
startingAt minimumFileDescriptor: Int32) throws
{
#if canImport(Glibc)
try self.check(posix_spawn_file_actions_addclosefrom_np(&fileActions, minimumFileDescriptor))
#else
for descriptor in try self.descriptorsToClose(startingAt: minimumFileDescriptor) {
try self.check(posix_spawn_file_actions_addclose(&fileActions, descriptor))
}
#endif
}
private static func check(_ result: Int32) throws {
guard result == 0 else {
throw CloseFromError.actionFailed(result)
}
}
#endif
}
@@ -0,0 +1,150 @@
import Foundation
package final class ProcessPipeCapture: @unchecked Sendable {
package static let defaultMaxBytes = 1 * 1024 * 1024
private let handle: FileHandle
private let onData: (@Sendable () -> Void)?
private let maxBytes: Int
private let condition = NSCondition()
private var data = Data()
private var activeCallbacks = 0
private var isFinished = false
private var didReachEOF = false
private var isStopping = false
private var continuation: CheckedContinuation<Void, Never>?
package init(
pipe: Pipe,
maxBytes: Int = ProcessPipeCapture.defaultMaxBytes,
onData: (@Sendable () -> Void)? = nil)
{
self.handle = pipe.fileHandleForReading
self.maxBytes = max(0, maxBytes)
self.onData = onData
}
package func start() {
self.handle.readabilityHandler = { [weak self] handle in
self?.handleReadableData(from: handle)
}
}
package func finish(timeout: Duration) async -> Data {
let drainTask = Task<Void, Error> {
await self.waitUntilFinished()
}
let join = BoundedTaskJoin(sourceTask: drainTask)
_ = await join.value(joinGrace: timeout)
return self.stopAndSnapshot()
}
package func finishSynchronously(timeout: TimeInterval) -> Data {
let deadline = Date().addingTimeInterval(max(0, timeout))
self.condition.lock()
while !self.isFinished, !self.isStopping {
guard self.condition.wait(until: deadline) else { break }
}
self.condition.unlock()
return self.stopAndSnapshot()
}
/// Waits only for the first complete output line. Useful for helpers whose descendants may inherit stdout
/// after the helper itself exits, preventing EOF even though the caller already has its complete answer.
package func finishFirstLineSynchronously(timeout: TimeInterval) -> Data {
let deadline = Date().addingTimeInterval(max(0, timeout))
self.condition.lock()
while !self.isFinished, !self.isStopping, !self.data.contains(0x0A) {
guard self.condition.wait(until: deadline) else { break }
}
self.condition.unlock()
return self.stopAndSnapshot()
}
package func stop() {
_ = self.stopAndSnapshot()
}
package var reachedEOF: Bool {
self.condition.lock()
defer { self.condition.unlock() }
return self.didReachEOF
}
package static func decodeUTF8(_ data: Data) -> String {
// A byte cap can split the final scalar; lossy decoding preserves the valid captured prefix.
// swiftlint:disable:next optional_data_string_conversion
String(decoding: data, as: UTF8.self)
}
private func handleReadableData(from handle: FileHandle) {
self.condition.lock()
guard !self.isStopping else {
self.condition.unlock()
return
}
self.activeCallbacks += 1
self.condition.unlock()
let chunk = handle.availableData
var continuation: CheckedContinuation<Void, Never>?
self.condition.lock()
if chunk.isEmpty {
self.isFinished = true
self.didReachEOF = true
continuation = self.continuation
self.continuation = nil
} else {
let remainingBytes = max(0, self.maxBytes - self.data.count)
if remainingBytes > 0 {
self.data.append(chunk.prefix(remainingBytes))
}
}
self.activeCallbacks -= 1
if self.activeCallbacks == 0 {
self.condition.broadcast()
}
self.condition.unlock()
if chunk.isEmpty {
handle.readabilityHandler = nil
} else {
self.onData?()
}
continuation?.resume()
}
private func waitUntilFinished() async {
await withCheckedContinuation { continuation in
self.condition.lock()
if self.isFinished || self.isStopping {
self.condition.unlock()
continuation.resume()
return
}
self.continuation = continuation
self.condition.unlock()
}
}
private func stopAndSnapshot() -> Data {
self.handle.readabilityHandler = nil
let continuation: CheckedContinuation<Void, Never>?
let snapshot: Data
self.condition.lock()
self.isStopping = true
while self.activeCallbacks > 0 {
self.condition.wait()
}
self.isFinished = true
continuation = self.continuation
self.continuation = nil
snapshot = self.data
self.condition.unlock()
continuation?.resume()
return snapshot
}
}
@@ -0,0 +1,809 @@
#if canImport(Darwin)
import Darwin
#elseif canImport(Glibc)
import Glibc
#elseif canImport(Musl)
import Musl
#endif
import Foundation
package final class SpawnedProcessGroup: @unchecked Sendable {
package enum LaunchError: LocalizedError {
case setupFailed(String)
case spawnFailed(String)
package var errorDescription: String? {
switch self {
case let .setupFailed(details):
"Failed to prepare process: \(details)"
case let .spawnFailed(details):
"Failed to launch process: \(details)"
}
}
}
private final class TerminationState: @unchecked Sendable {
private let condition = NSCondition()
private var exitObserved = false
private var reapRequested = false
private var status: Int32?
var hasObservedExit: Bool {
self.condition.withLock { self.exitObserved }
}
var value: Int32? {
self.condition.withLock { self.status }
}
func observeExit() {
self.condition.withLock {
self.exitObserved = true
self.condition.broadcast()
}
}
func requestReap() {
self.condition.withLock {
self.reapRequested = true
self.condition.broadcast()
}
}
func waitForReapRequest(timeout: TimeInterval) {
let deadline = Date().addingTimeInterval(timeout)
self.condition.lock()
while !self.reapRequested, self.condition.wait(until: deadline) {}
self.condition.unlock()
}
func resolve(_ status: Int32) {
self.condition.withLock {
guard self.status == nil else { return }
self.status = status
self.condition.broadcast()
}
}
}
private final class ProcessIdentityState: @unchecked Sendable {
private let lock = NSLock()
private var identities: Set<TTYProcessTreeTerminator.ProcessIdentity> = []
var snapshot: Set<TTYProcessTreeTerminator.ProcessIdentity> {
self.lock.withLock { self.identities }
}
func formUnion(_ identities: Set<TTYProcessTreeTerminator.ProcessIdentity>) {
self.lock.withLock {
self.identities.formUnion(identities)
}
}
}
private struct OutputPipeIdentity: Hashable {
#if canImport(Darwin)
let firstHandle: UInt64
let secondHandle: UInt64
#else
let inode: UInt64
#endif
static func resolve(fileDescriptor: Int32) -> OutputPipeIdentity? {
#if canImport(Darwin)
var info = pipe_fdinfo()
let byteCount = proc_pidfdinfo(
getpid(),
fileDescriptor,
PROC_PIDFDPIPEINFO,
&info,
Int32(MemoryLayout<pipe_fdinfo>.size))
guard byteCount == MemoryLayout<pipe_fdinfo>.size else { return nil }
let handles = [info.pipeinfo.pipe_handle, info.pipeinfo.pipe_peerhandle].sorted()
guard handles[0] != 0, handles[1] != 0 else { return nil }
return OutputPipeIdentity(firstHandle: handles[0], secondHandle: handles[1])
#else
var info = stat()
guard fstat(fileDescriptor, &info) == 0 else { return nil }
return OutputPipeIdentity(inode: UInt64(info.st_ino))
#endif
}
static func holderPIDs(for pipes: Set<OutputPipeIdentity>) -> Set<pid_t> {
guard !pipes.isEmpty else { return [] }
#if canImport(Darwin)
return Set(SpawnedProcessGroup.allProcessIDs().filter { self.process(pid: $0, holdsAny: pipes) })
#else
let targets = Set(pipes.map { "pipe:[\($0.inode)]" })
return Set(SpawnedProcessGroup.allProcessIDs().filter { pid in
let directory = "/proc/\(pid)/fd"
guard let descriptors = try? FileManager.default.contentsOfDirectory(atPath: directory) else {
return false
}
return descriptors.contains { descriptor in
let path = "\(directory)/\(descriptor)"
guard let target = try? FileManager.default.destinationOfSymbolicLink(atPath: path) else {
return false
}
return targets.contains(target)
}
})
#endif
}
#if canImport(Darwin)
private static func process(pid: pid_t, holdsAny pipes: Set<OutputPipeIdentity>) -> Bool {
let requiredBytes = proc_pidinfo(pid, PROC_PIDLISTFDS, 0, nil, 0)
guard requiredBytes > 0 else { return false }
let stride = MemoryLayout<proc_fdinfo>.stride
var descriptors = [proc_fdinfo](
repeating: proc_fdinfo(),
count: Int(requiredBytes) / stride + 8)
let actualBytes = descriptors.withUnsafeMutableBytes { buffer in
proc_pidinfo(
pid,
PROC_PIDLISTFDS,
0,
buffer.baseAddress,
Int32(buffer.count))
}
guard actualBytes > 0 else { return false }
for descriptor in descriptors.prefix(Int(actualBytes) / stride)
where descriptor.proc_fdtype == PROX_FDTYPE_PIPE
{
var info = pipe_fdinfo()
let byteCount = proc_pidfdinfo(
pid,
descriptor.proc_fd,
PROC_PIDFDPIPEINFO,
&info,
Int32(MemoryLayout<pipe_fdinfo>.size))
guard byteCount == MemoryLayout<pipe_fdinfo>.size else { continue }
let handles = [info.pipeinfo.pipe_handle, info.pipeinfo.pipe_peerhandle].sorted()
let identity = OutputPipeIdentity(firstHandle: handles[0], secondHandle: handles[1])
if pipes.contains(identity) {
return true
}
}
return false
}
#endif
}
private struct OutputTTYIdentity: Hashable {
let device: UInt64
let inode: UInt64
let rawDevice: UInt64
static func resolve(fileDescriptor: Int32) -> OutputTTYIdentity? {
var info = stat()
guard fstat(fileDescriptor, &info) == 0 else { return nil }
#if canImport(Darwin)
let device = SpawnedProcessGroup.darwinDeviceIdentifier(info.st_dev)
let rawDevice = SpawnedProcessGroup.darwinDeviceIdentifier(info.st_rdev)
#else
let device = UInt64(info.st_dev)
let rawDevice = UInt64(info.st_rdev)
#endif
return OutputTTYIdentity(
device: device,
inode: UInt64(info.st_ino),
rawDevice: rawDevice)
}
static func holderPIDs(for terminals: Set<OutputTTYIdentity>) -> Set<pid_t> {
guard !terminals.isEmpty else { return [] }
#if canImport(Darwin)
return Set(SpawnedProcessGroup.allProcessIDs().filter { self.process(pid: $0, holdsAny: terminals) })
#else
return Set(SpawnedProcessGroup.allProcessIDs().filter { pid in
let directory = "/proc/\(pid)/fd"
guard let descriptors = try? FileManager.default.contentsOfDirectory(atPath: directory) else {
return false
}
return descriptors.contains { descriptor in
var info = stat()
let path = "\(directory)/\(descriptor)"
guard path.withCString({ fstatat(AT_FDCWD, $0, &info, 0) }) == 0 else { return false }
let identity = OutputTTYIdentity(
device: UInt64(info.st_dev),
inode: UInt64(info.st_ino),
rawDevice: UInt64(info.st_rdev))
return terminals.contains(identity)
}
})
#endif
}
#if canImport(Darwin)
private static func process(pid: pid_t, holdsAny terminals: Set<OutputTTYIdentity>) -> Bool {
let requiredBytes = proc_pidinfo(pid, PROC_PIDLISTFDS, 0, nil, 0)
guard requiredBytes > 0 else { return false }
let stride = MemoryLayout<proc_fdinfo>.stride
var descriptors = [proc_fdinfo](
repeating: proc_fdinfo(),
count: Int(requiredBytes) / stride + 8)
let actualBytes = descriptors.withUnsafeMutableBytes { buffer in
proc_pidinfo(
pid,
PROC_PIDLISTFDS,
0,
buffer.baseAddress,
Int32(buffer.count))
}
guard actualBytes > 0 else { return false }
for descriptor in descriptors.prefix(Int(actualBytes) / stride)
where descriptor.proc_fdtype == PROX_FDTYPE_VNODE
{
var info = vnode_fdinfo()
let byteCount = proc_pidfdinfo(
pid,
descriptor.proc_fd,
PROC_PIDFDVNODEINFO,
&info,
Int32(MemoryLayout<vnode_fdinfo>.size))
guard byteCount == MemoryLayout<vnode_fdinfo>.size else { continue }
let stats = info.pvi.vi_stat
let identity = OutputTTYIdentity(
device: UInt64(stats.vst_dev),
inode: stats.vst_ino,
rawDevice: UInt64(stats.vst_rdev))
if terminals.contains(identity) {
return true
}
}
return false
}
#endif
}
private static func allProcessIDs() -> [pid_t] {
#if canImport(Darwin)
return self.processIDs(type: UInt32(PROC_ALL_PIDS), typeInfo: 0)
#else
guard let entries = try? FileManager.default.contentsOfDirectory(atPath: "/proc") else { return [] }
return entries.compactMap(pid_t.init)
#endif
}
#if canImport(Darwin)
/// Darwin exposes `stat` device IDs as signed values while vnode inspection uses the same bits unsigned.
package static func darwinDeviceIdentifier(_ value: Int32) -> UInt64 {
UInt64(UInt32(bitPattern: value))
}
#endif
private static func processIDs(inProcessGroup processGroup: pid_t) -> [pid_t] {
#if canImport(Darwin)
return self.processIDs(type: UInt32(PROC_PGRP_ONLY), typeInfo: UInt32(bitPattern: processGroup))
#else
return self.allProcessIDs().filter { getpgid($0) == processGroup }
#endif
}
#if canImport(Darwin)
private static func processIDs(type: UInt32, typeInfo: UInt32) -> [pid_t] {
let requiredBytes = proc_listpids(type, typeInfo, nil, 0)
guard requiredBytes > 0 else { return [] }
let stride = MemoryLayout<pid_t>.stride
var pids = [pid_t](repeating: 0, count: Int(requiredBytes) / stride + 32)
let actualBytes = pids.withUnsafeMutableBytes { buffer in
proc_listpids(
type,
typeInfo,
buffer.baseAddress,
Int32(buffer.count))
}
guard actualBytes > 0 else { return [] }
return Array(pids.prefix(Int(actualBytes) / stride)).filter { $0 > 0 }
}
#endif
package let pid: pid_t
package let processGroup: pid_t
private let termination = TerminationState()
private let observedProcessGroupMembers = ProcessIdentityState()
private let outputPipes: Set<OutputPipeIdentity>
private let outputTTYs: Set<OutputTTYIdentity>
private let rootIdentity: TTYProcessTreeTerminator.ProcessIdentity?
private init(
pid: pid_t,
outputPipes: Set<OutputPipeIdentity>,
outputTTYs: Set<OutputTTYIdentity> = [])
{
self.pid = pid
self.processGroup = pid
self.outputPipes = outputPipes
self.outputTTYs = outputTTYs
self.rootIdentity = TTYProcessTreeTerminator.processIdentity(for: pid)
self.startWaiter()
}
package static func launch(
binary: String,
arguments: [String],
environment: [String: String],
stdoutPipe: Pipe,
stderrPipe: Pipe) throws -> SpawnedProcessGroup
{
#if canImport(Darwin)
var fileActions: posix_spawn_file_actions_t?
#else
var fileActions = posix_spawn_file_actions_t()
#endif
guard posix_spawn_file_actions_init(&fileActions) == 0 else {
throw LaunchError.setupFailed("posix_spawn_file_actions_init")
}
defer { posix_spawn_file_actions_destroy(&fileActions) }
let stdoutRead = stdoutPipe.fileHandleForReading.fileDescriptor
let stdoutWrite = stdoutPipe.fileHandleForWriting.fileDescriptor
let stderrRead = stderrPipe.fileHandleForReading.fileDescriptor
let stderrWrite = stderrPipe.fileHandleForWriting.fileDescriptor
let outputPipes = Set(
[stdoutRead, stderrRead].compactMap(OutputPipeIdentity.resolve(fileDescriptor:)))
var fileActionResults = [
posix_spawn_file_actions_addopen(&fileActions, STDIN_FILENO, "/dev/null", O_RDONLY, 0),
posix_spawn_file_actions_adddup2(&fileActions, stdoutWrite, STDOUT_FILENO),
posix_spawn_file_actions_adddup2(&fileActions, stderrWrite, STDERR_FILENO),
]
for descriptor in Self.pipeDescriptorsToClose([stdoutRead, stdoutWrite, stderrRead, stderrWrite]) {
fileActionResults.append(posix_spawn_file_actions_addclose(&fileActions, descriptor))
}
#if canImport(Glibc) || canImport(Musl)
do {
try PosixSpawnFileActionsCloseFrom.addCloseFrom(
&fileActions,
startingAt: STDERR_FILENO + 1)
} catch {
throw LaunchError.setupFailed(error.localizedDescription)
}
#endif
guard fileActionResults.allSatisfy({ $0 == 0 }) else {
throw LaunchError.setupFailed("posix_spawn file actions")
}
#if canImport(Darwin)
var attributes: posix_spawnattr_t?
#else
var attributes = posix_spawnattr_t()
#endif
guard posix_spawnattr_init(&attributes) == 0 else {
throw LaunchError.setupFailed("posix_spawnattr_init")
}
defer { posix_spawnattr_destroy(&attributes) }
#if canImport(Darwin)
let flags = POSIX_SPAWN_SETPGROUP | POSIX_SPAWN_CLOEXEC_DEFAULT
#else
let flags = POSIX_SPAWN_SETPGROUP
#endif
guard posix_spawnattr_setflags(&attributes, Int16(flags)) == 0,
posix_spawnattr_setpgroup(&attributes, 0) == 0
else {
throw LaunchError.setupFailed("posix_spawn process group")
}
var cArguments: [UnsafeMutablePointer<CChar>?] = ([binary] + arguments).map { strdup($0) }
cArguments.append(nil)
defer {
for argument in cArguments {
free(argument)
}
}
var cEnvironment: [UnsafeMutablePointer<CChar>?] = environment.map { key, value in
strdup("\(key)=\(value)")
}
cEnvironment.append(nil)
defer {
for entry in cEnvironment {
free(entry)
}
}
var pid: pid_t = 0
let spawnResult = binary.withCString { path in
posix_spawn(&pid, path, &fileActions, &attributes, cArguments, cEnvironment)
}
stdoutPipe.fileHandleForWriting.closeFile()
stderrPipe.fileHandleForWriting.closeFile()
guard spawnResult == 0 else {
throw LaunchError.spawnFailed(String(cString: strerror(spawnResult)))
}
return SpawnedProcessGroup(pid: pid, outputPipes: outputPipes)
}
package static func launchPTY(
binary: String,
arguments: [String],
environment: [String: String],
workingDirectory: URL?,
fileDescriptors: (primary: Int32, secondary: Int32)) throws -> SpawnedProcessGroup
{
let primaryFD = fileDescriptors.primary
let secondaryFD = fileDescriptors.secondary
guard let outputTTY = OutputTTYIdentity.resolve(fileDescriptor: secondaryFD) else {
throw LaunchError.setupFailed("resolve PTY identity")
}
#if canImport(Darwin)
var fileActions: posix_spawn_file_actions_t?
#else
var fileActions = posix_spawn_file_actions_t()
#endif
guard posix_spawn_file_actions_init(&fileActions) == 0 else {
throw LaunchError.setupFailed("posix_spawn_file_actions_init")
}
defer { posix_spawn_file_actions_destroy(&fileActions) }
var fileActionResults = [
posix_spawn_file_actions_adddup2(&fileActions, secondaryFD, STDIN_FILENO),
posix_spawn_file_actions_adddup2(&fileActions, secondaryFD, STDOUT_FILENO),
posix_spawn_file_actions_adddup2(&fileActions, secondaryFD, STDERR_FILENO),
]
for descriptor in Self.pipeDescriptorsToClose([primaryFD, secondaryFD]) {
fileActionResults.append(posix_spawn_file_actions_addclose(&fileActions, descriptor))
}
if let workingDirectory {
fileActionResults.append(workingDirectory.path.withCString { path in
posix_spawn_file_actions_addchdir_np(&fileActions, path)
})
}
#if canImport(Glibc) || canImport(Musl)
do {
try PosixSpawnFileActionsCloseFrom.addCloseFrom(
&fileActions,
startingAt: STDERR_FILENO + 1)
} catch {
throw LaunchError.setupFailed(error.localizedDescription)
}
#endif
guard fileActionResults.allSatisfy({ $0 == 0 }) else {
throw LaunchError.setupFailed("posix_spawn PTY file actions")
}
#if canImport(Darwin)
var attributes: posix_spawnattr_t?
#else
var attributes = posix_spawnattr_t()
#endif
guard posix_spawnattr_init(&attributes) == 0 else {
throw LaunchError.setupFailed("posix_spawnattr_init")
}
defer { posix_spawnattr_destroy(&attributes) }
#if canImport(Darwin)
let flags = POSIX_SPAWN_SETPGROUP | POSIX_SPAWN_CLOEXEC_DEFAULT
#else
let flags = POSIX_SPAWN_SETPGROUP
#endif
guard posix_spawnattr_setflags(&attributes, Int16(flags)) == 0,
posix_spawnattr_setpgroup(&attributes, 0) == 0
else {
throw LaunchError.setupFailed("posix_spawn PTY process group")
}
var cArguments: [UnsafeMutablePointer<CChar>?] = ([binary] + arguments).map { strdup($0) }
cArguments.append(nil)
defer {
for argument in cArguments {
free(argument)
}
}
var cEnvironment: [UnsafeMutablePointer<CChar>?] = environment.map { key, value in
strdup("\(key)=\(value)")
}
cEnvironment.append(nil)
defer {
for entry in cEnvironment {
free(entry)
}
}
var pid: pid_t = 0
let spawnResult = binary.withCString { path in
posix_spawn(&pid, path, &fileActions, &attributes, cArguments, cEnvironment)
}
guard spawnResult == 0 else {
throw LaunchError.spawnFailed(String(cString: strerror(spawnResult)))
}
return SpawnedProcessGroup(pid: pid, outputPipes: [], outputTTYs: [outputTTY])
}
package var isRunning: Bool {
!self.termination.hasObservedExit
}
package var terminationStatus: Int32? {
self.termination.value
}
package var hasResidualProcessGroup: Bool {
Self.processGroupExists(self.processGroup)
}
@discardableResult
package func terminateSynchronously(grace: TimeInterval = 0.4) -> Int32? {
let deadline = Date().addingTimeInterval(max(0, grace))
var processIdentities = self.currentResidualProcessIdentities(includeDescendants: true)
processIdentities.formUnion(self.currentProcessGroupMemberIdentities())
if let rootIdentity = self.rootIdentity {
processIdentities.insert(rootIdentity)
}
Self.signal(processIdentities: processIdentities, signal: SIGTERM)
while processIdentities.contains(where: TTYProcessTreeTerminator.isCurrent(_:)),
Date() < deadline
{
usleep(20000)
}
processIdentities.formUnion(self.currentResidualProcessIdentities(includeDescendants: self.isRunning))
processIdentities.formUnion(self.currentProcessGroupMemberIdentities())
if self.isRunning, let rootIdentity = self.rootIdentity {
processIdentities.insert(rootIdentity)
}
Self.signal(processIdentities: processIdentities, signal: SIGKILL)
let killDeadline = Date().addingTimeInterval(max(0, grace))
while processIdentities.contains(where: TTYProcessTreeTerminator.isCurrent(_:)),
Date() < killDeadline
{
usleep(20000)
}
return self.finishSynchronously()
}
@discardableResult
package func finishSynchronously(timeout: TimeInterval = 1) -> Int32? {
self.termination.requestReap()
let deadline = Date().addingTimeInterval(max(0, timeout))
while self.terminationStatus == nil, Date() < deadline {
usleep(10000)
}
return self.terminationStatus
}
@discardableResult
package func terminate(grace: TimeInterval = 0.4) async -> Int32? {
if self.isRunning {
let killDeadline = Date().addingTimeInterval(max(0, grace))
var processIdentities = self.currentResidualProcessIdentities(includeDescendants: true)
processIdentities.formUnion(self.currentProcessGroupMemberIdentities())
if let rootIdentity = TTYProcessTreeTerminator.processIdentity(for: self.pid) {
processIdentities.insert(rootIdentity)
}
Self.signal(processIdentities: processIdentities, signal: SIGTERM)
_ = await self.waitForExit(timeout: max(0, killDeadline.timeIntervalSinceNow))
while processIdentities.contains(where: TTYProcessTreeTerminator.isCurrent(_:)),
Date() < killDeadline
{
try? await Task.sleep(for: .milliseconds(20))
}
processIdentities.formUnion(self.currentResidualProcessIdentities(includeDescendants: false))
processIdentities.formUnion(self.currentProcessGroupMemberIdentities())
if self.isRunning {
processIdentities.formUnion(self.currentResidualProcessIdentities(includeDescendants: true))
if let rootIdentity = TTYProcessTreeTerminator.processIdentity(for: self.pid) {
processIdentities.insert(rootIdentity)
}
Self.signal(processIdentities: processIdentities, signal: SIGKILL)
_ = await self.waitForExit(timeout: grace)
} else {
Self.signal(processIdentities: processIdentities, signal: SIGKILL)
}
_ = await self.waitForResidualProcessesExit(processIdentities, timeout: grace)
await self.finish()
return self.terminationStatus
}
await self.terminateResidualProcesses(grace: grace)
await self.finish()
return self.terminationStatus
}
package func terminateResidualProcesses(grace: TimeInterval = 0.4) async {
let deadline = Date().addingTimeInterval(max(0, grace))
var processIdentities = self.currentResidualProcessIdentities(includeDescendants: false)
processIdentities.formUnion(self.currentProcessGroupMemberIdentities())
if self.isRunning {
processIdentities.formUnion(self.currentResidualProcessIdentities(includeDescendants: true))
if let rootIdentity = TTYProcessTreeTerminator.processIdentity(for: self.pid) {
processIdentities.insert(rootIdentity)
}
}
Self.signal(processIdentities: processIdentities, signal: SIGTERM)
while processIdentities.contains(where: TTYProcessTreeTerminator.isCurrent(_:)) {
guard Date() < deadline else { break }
try? await Task.sleep(for: .milliseconds(20))
}
processIdentities.formUnion(self.currentResidualProcessIdentities(includeDescendants: false))
processIdentities.formUnion(self.currentProcessGroupMemberIdentities())
if self.isRunning {
processIdentities.formUnion(self.currentResidualProcessIdentities(includeDescendants: true))
if let rootIdentity = TTYProcessTreeTerminator.processIdentity(for: self.pid) {
processIdentities.insert(rootIdentity)
}
}
guard processIdentities.contains(where: TTYProcessTreeTerminator.isCurrent(_:)) else {
return
}
Self.signal(processIdentities: processIdentities, signal: SIGKILL)
_ = await self.waitForResidualProcessesExit(processIdentities, timeout: grace)
}
package func finish() async {
self.termination.requestReap()
_ = await self.waitForTerminationStatus(timeout: 1)
}
private func startWaiter() {
let pid = self.pid
let processGroup = self.processGroup
let observedProcessGroupMembers = self.observedProcessGroupMembers
let termination = self.termination
DispatchQueue.global(qos: .userInitiated).async {
var info = siginfo_t()
var waitResult: Int32
repeat {
waitResult = waitid(P_PID, id_t(pid), &info, WEXITED | WNOWAIT)
} while waitResult == -1 && errno == EINTR
guard waitResult == 0 else {
termination.observeExit()
termination.resolve(1)
return
}
// The exited root remains unreaped here, so its PID cannot yet be reused as
// an unrelated process-group ID.
observedProcessGroupMembers.formUnion(
Self.processGroupMemberIdentities(processGroup: processGroup, excluding: pid))
termination.observeExit()
termination.waitForReapRequest(timeout: 30)
var rawStatus: Int32 = 0
var result: pid_t
repeat {
result = waitpid(pid, &rawStatus, 0)
} while result == -1 && errno == EINTR
let status = result == pid ? Self.exitStatus(from: rawStatus) : 1
termination.resolve(status)
}
}
private func waitForExit(timeout: TimeInterval) async -> Int32? {
let deadline = Date().addingTimeInterval(max(0, timeout))
while self.isRunning, Date() < deadline {
try? await Task.sleep(for: .milliseconds(20))
}
return self.terminationStatus
}
private func waitForTerminationStatus(timeout: TimeInterval) async -> Int32? {
let deadline = Date().addingTimeInterval(max(0, timeout))
while self.terminationStatus == nil, Date() < deadline {
try? await Task.sleep(for: .milliseconds(20))
}
return self.terminationStatus
}
private func waitForResidualProcessesExit(
_ processIdentities: Set<TTYProcessTreeTerminator.ProcessIdentity>,
timeout: TimeInterval) async -> Bool
{
let deadline = Date().addingTimeInterval(max(0, timeout))
while Date() < deadline {
guard processIdentities.contains(where: TTYProcessTreeTerminator.isCurrent(_:)) else {
return true
}
try? await Task.sleep(for: .milliseconds(20))
}
return !processIdentities.contains(where: TTYProcessTreeTerminator.isCurrent(_:))
}
private func currentResidualProcessIdentities(
includeDescendants: Bool) -> Set<TTYProcessTreeTerminator.ProcessIdentity>
{
var identities = self.currentOutputHolderIdentities()
identities.formUnion(self.observedProcessGroupMembers.snapshot)
if includeDescendants {
identities.formUnion(
TTYProcessTreeTerminator.descendantPIDs(of: self.pid)
.compactMap(TTYProcessTreeTerminator.processIdentity(for:)))
}
return identities
}
private func currentOutputHolderIdentities() -> Set<TTYProcessTreeTerminator.ProcessIdentity> {
let excludedPIDs: Set<pid_t> = [getpid(), self.pid]
var holderPIDs = OutputPipeIdentity.holderPIDs(for: self.outputPipes)
holderPIDs.formUnion(OutputTTYIdentity.holderPIDs(for: self.outputTTYs))
return Set(holderPIDs.subtracting(excludedPIDs).compactMap(TTYProcessTreeTerminator.processIdentity(for:)))
}
private func currentProcessGroupMemberIdentities() -> Set<TTYProcessTreeTerminator.ProcessIdentity> {
if self.termination.hasObservedExit, self.termination.value == nil {
let identities = Self.processGroupMemberIdentities(
processGroup: self.processGroup,
excluding: self.pid)
self.observedProcessGroupMembers.formUnion(identities)
return identities
}
guard let rootIdentity = self.rootIdentity
else {
return []
}
let identities = Self.processGroupMemberIdentities(
processGroup: self.processGroup,
rootIdentity: rootIdentity,
excluding: self.pid)
self.observedProcessGroupMembers.formUnion(identities)
return identities
}
private static func processGroupMemberIdentities(
processGroup: pid_t,
rootIdentity: TTYProcessTreeTerminator.ProcessIdentity,
excluding excludedPID: pid_t)
-> Set<TTYProcessTreeTerminator.ProcessIdentity>
{
guard TTYProcessTreeTerminator.isCurrent(rootIdentity) else { return [] }
let identities = Self.processGroupMemberIdentities(
processGroup: processGroup,
excluding: excludedPID)
guard TTYProcessTreeTerminator.isCurrent(rootIdentity) else { return [] }
return identities
}
private static func processGroupMemberIdentities(
processGroup: pid_t,
excluding excludedPID: pid_t)
-> Set<TTYProcessTreeTerminator.ProcessIdentity>
{
Set(self.processIDs(inProcessGroup: processGroup)
.compactMap { pid -> TTYProcessTreeTerminator.ProcessIdentity? in
guard pid != getpid(),
pid != excludedPID,
let identity = TTYProcessTreeTerminator.processIdentity(for: pid),
getpgid(pid) == processGroup,
TTYProcessTreeTerminator.isCurrent(identity)
else {
return nil
}
return identity
})
}
private static func processGroupExists(_ processGroup: pid_t) -> Bool {
errno = 0
return kill(-processGroup, 0) == 0 || errno == EPERM
}
private static func signal(
processIdentities: Set<TTYProcessTreeTerminator.ProcessIdentity>,
signal: Int32)
{
for identity in processIdentities where TTYProcessTreeTerminator.isCurrent(identity) {
_ = kill(identity.pid, signal)
}
}
package static func pipeDescriptorsToClose(_ descriptors: [Int32]) -> [Int32] {
Array(Set(descriptors.filter { $0 > STDERR_FILENO })).sorted()
}
private static func exitStatus(from rawStatus: Int32) -> Int32 {
let signal = rawStatus & 0x7F
return signal == 0 ? (rawStatus >> 8) & 0xFF : signal
}
}
@@ -0,0 +1,321 @@
#if canImport(Darwin)
import Darwin
#elseif canImport(Glibc)
import Glibc
#elseif canImport(Musl)
import Musl
#endif
import Foundation
public enum SubprocessRunnerError: LocalizedError, Sendable {
case binaryNotFound(String)
case launchFailed(String)
case timedOut(String)
case nonZeroExit(code: Int32, stderr: String)
public var errorDescription: String? {
switch self {
case let .binaryNotFound(binary):
return "Missing CLI '\(binary)'. Install it and restart CodexBar."
case let .launchFailed(details):
return "Failed to launch process: \(details)"
case let .timedOut(label):
return "Command timed out: \(label)"
case let .nonZeroExit(code, stderr):
let trimmed = stderr.trimmingCharacters(in: .whitespacesAndNewlines)
if trimmed.isEmpty {
return "Command failed with exit code \(code)."
}
return "Command failed (\(code)): \(trimmed)"
}
}
}
public struct SubprocessResult: Sendable {
public let stdout: String
public let stderr: String
}
public enum SubprocessRunner {
private static let log = CodexBarLog.logger(LogCategories.subprocess)
private static let timeoutQueue = DispatchQueue(
label: "com.steipete.codexbar.subprocess.timeout",
qos: .userInitiated,
attributes: .concurrent)
/// Thread-safe flag for communicating between concurrent tasks (e.g. timeout caller).
private final class KillFlag: @unchecked Sendable {
private let lock = NSLock()
private var value = false
func set() {
self.lock.withLock { self.value = true }
}
var isSet: Bool {
self.lock.withLock { self.value }
}
}
private final class TimeoutTimer: @unchecked Sendable {
private let timer: any DispatchSourceTimer
init(timer: any DispatchSourceTimer) {
self.timer = timer
}
func cancel() {
self.timer.cancel()
}
}
private static func timeoutInterval(_ timeout: TimeInterval) -> DispatchTimeInterval {
guard timeout.isFinite else {
return .seconds(Int.max)
}
let nanoseconds = max(0, min(timeout * 1_000_000_000, Double(Int.max)))
return .nanoseconds(Int(nanoseconds))
}
private final class ProcessTermination: @unchecked Sendable {
private let lock = NSLock()
private var status: Int32?
private var continuation: CheckedContinuation<Int32, Never>?
func resolve(_ status: Int32) {
let continuation: CheckedContinuation<Int32, Never>?
self.lock.lock()
self.status = status
continuation = self.continuation
self.continuation = nil
self.lock.unlock()
continuation?.resume(returning: status)
}
func wait() async -> Int32 {
await withCheckedContinuation { continuation in
let status: Int32?
self.lock.lock()
status = self.status
if status == nil {
self.continuation = continuation
}
self.lock.unlock()
if let status {
continuation.resume(returning: status)
}
}
}
}
/// Terminates a process and its process group, escalating from SIGTERM to SIGKILL.
/// Returns `true` if the process was actually killed, `false` if it had already exited.
@discardableResult
package static func terminateProcess(_ process: Process, processGroup: pid_t?) -> Bool {
guard process.isRunning else { return false }
let descendants = TTYProcessTreeTerminator.descendantPIDs(of: process.processIdentifier)
let descendantIdentities = descendants.compactMap(TTYProcessTreeTerminator.processIdentity(for:))
TTYProcessTreeTerminator.terminateProcessTree(
rootPID: process.processIdentifier,
processGroup: processGroup,
signal: SIGTERM,
knownDescendants: descendants)
let killDeadline = Date().addingTimeInterval(0.4)
while process.isRunning, Date() < killDeadline {
usleep(50000)
}
if process.isRunning {
let currentDescendants = descendantIdentities
.filter(TTYProcessTreeTerminator.isCurrent(_:))
.map(\.pid)
TTYProcessTreeTerminator.terminateProcessTree(
rootPID: process.processIdentifier,
processGroup: processGroup,
signal: SIGKILL,
knownDescendants: currentDescendants)
let reapDeadline = Date().addingTimeInterval(0.4)
while process.isRunning, Date() < reapDeadline {
usleep(50000)
}
} else {
for identity in descendantIdentities where TTYProcessTreeTerminator.isCurrent(identity) {
kill(identity.pid, SIGKILL)
}
}
return true
}
// MARK: - Public API
/// Runs a process to natural exit without letting caller cancellation or a
/// forced timeout interrupt an external mutation after launch.
public static func runToCompletion(
binary: String,
arguments: [String],
environment: [String: String],
currentDirectoryURL: URL? = nil,
acceptsNonZeroExit: Bool = false,
label: String) async throws -> SubprocessResult
{
let task = Task.detached(priority: .userInitiated) {
try await self.run(
binary: binary,
arguments: arguments,
environment: environment,
timeout: .infinity,
currentDirectoryURL: currentDirectoryURL,
acceptsNonZeroExit: acceptsNonZeroExit,
label: label)
}
return try await task.value
}
public static func run(
binary: String,
arguments: [String],
environment: [String: String],
timeout: TimeInterval,
standardInput: Any? = nil,
currentDirectoryURL: URL? = nil,
acceptsNonZeroExit: Bool = false,
label: String) async throws -> SubprocessResult
{
guard FileManager.default.isExecutableFile(atPath: binary) else {
throw SubprocessRunnerError.binaryNotFound(binary)
}
let start = Date()
let binaryName = URL(fileURLWithPath: binary).lastPathComponent
self.log.debug(
"Subprocess start",
metadata: ["label": label, "binary": binaryName, "timeout": "\(timeout)"])
let process = Process()
process.executableURL = URL(fileURLWithPath: binary)
process.arguments = arguments
process.environment = environment
process.currentDirectoryURL = currentDirectoryURL
let stdoutPipe = Pipe()
let stderrPipe = Pipe()
process.standardOutput = stdoutPipe
process.standardError = stderrPipe
process.standardInput = standardInput
let stdoutCapture = ProcessPipeCapture(pipe: stdoutPipe)
let stderrCapture = ProcessPipeCapture(pipe: stderrPipe)
let termination = ProcessTermination()
process.terminationHandler = { process in
termination.resolve(process.terminationStatus)
}
do {
try process.run()
} catch {
process.terminationHandler = nil
stdoutCapture.stop()
stdoutPipe.fileHandleForWriting.closeFile()
stderrCapture.stop()
stderrPipe.fileHandleForWriting.closeFile()
throw SubprocessRunnerError.launchFailed(error.localizedDescription)
}
stdoutCapture.start()
stderrCapture.start()
let pid = process.processIdentifier
let processGroup: pid_t? = setpgid(pid, pid) == 0 ? pid : nil
let exitCodeTask = Task<Int32, Never> {
await termination.wait()
}
let killedByTimeout = KillFlag()
let timeoutTimerBox: TimeoutTimer? = if timeout.isFinite {
{
let timeoutTimer = DispatchSource.makeTimerSource(queue: self.timeoutQueue)
timeoutTimer.schedule(deadline: .now() + self.timeoutInterval(timeout))
timeoutTimer.setEventHandler {
guard process.isRunning else { return }
killedByTimeout.set()
self.terminateProcess(process, processGroup: processGroup)
}
timeoutTimer.resume()
return TimeoutTimer(timer: timeoutTimer)
}()
} else {
nil
}
do {
let exitCode = try await withTaskCancellationHandler {
try Task.checkCancellation()
let code = await exitCodeTask.value
try Task.checkCancellation()
return code
} onCancel: {
timeoutTimerBox?.cancel()
self.terminateProcess(process, processGroup: processGroup)
}
timeoutTimerBox?.cancel()
let duration = Date().timeIntervalSince(start)
// Race guard: the timeout timer may kill the process just before the
// exit code arrives. Key off the explicit kill flag so a completed
// process is not misclassified when the awaiting task resumes late.
if killedByTimeout.isSet {
self.log.warning(
"Subprocess timed out",
metadata: [
"label": label,
"binary": binaryName,
"duration_ms": "\(Int(duration * 1000))",
])
throw SubprocessRunnerError.timedOut(label)
}
async let stdoutData = stdoutCapture.finish(timeout: .seconds(1))
async let stderrData = stderrCapture.finish(timeout: .seconds(1))
let stdout = await ProcessPipeCapture.decodeUTF8(stdoutData)
let stderr = await ProcessPipeCapture.decodeUTF8(stderrData)
if exitCode != 0, !acceptsNonZeroExit {
let duration = Date().timeIntervalSince(start)
self.log.warning(
"Subprocess failed",
metadata: [
"label": label,
"binary": binaryName,
"status": "\(exitCode)",
"duration_ms": "\(Int(duration * 1000))",
])
throw SubprocessRunnerError.nonZeroExit(code: exitCode, stderr: stderr)
}
self.log.debug(
"Subprocess exit",
metadata: [
"label": label,
"binary": binaryName,
"status": "\(exitCode)",
"duration_ms": "\(Int(duration * 1000))",
])
return SubprocessResult(stdout: stdout, stderr: stderr)
} catch {
let duration = Date().timeIntervalSince(start)
self.log.warning(
"Subprocess error",
metadata: [
"label": label,
"binary": binaryName,
"duration_ms": "\(Int(duration * 1000))",
])
// Safety net: ensure the process is dead (may already be killed by timeout timer).
self.terminateProcess(process, processGroup: processGroup)
exitCodeTask.cancel()
stdoutCapture.stop()
stderrCapture.stop()
throw error
}
}
}
@@ -0,0 +1,100 @@
import Foundation
#if canImport(SweetCookieKit)
import SweetCookieKit
#endif
public enum KeychainAccessGate {
private static let flagKey = "debugDisableKeychainAccess"
static let disableAccessEnvironmentKey = "CODEXBAR_DISABLE_KEYCHAIN_ACCESS"
@TaskLocal private static var taskOverrideValue: Bool?
private nonisolated(unsafe) static var overrideValue: Bool?
private static let processForceDisabledLock = NSLock()
private nonisolated(unsafe) static var processForceDisabledReason: String?
public nonisolated(unsafe) static var isDisabled: Bool {
get {
if let taskOverrideValue { return taskOverrideValue }
if self.isDisabledByEnvironment() { return true }
#if DEBUG
if Self.forcesDisabledUnderTests {
return true
}
#endif
if self.processDisableReason != nil { return true }
if let overrideValue { return overrideValue }
if UserDefaults.standard.bool(forKey: Self.flagKey) { return true }
if let shared = AppGroupSupport.sharedDefaults(), shared.bool(forKey: Self.flagKey) {
return true
}
return false
}
set {
overrideValue = newValue
#if os(macOS) && canImport(SweetCookieKit)
BrowserCookieKeychainAccessGate.isDisabled = self.isDisabled
#endif
}
}
static func isDisabledByEnvironment(
_ environment: [String: String] = ProcessInfo.processInfo.environment) -> Bool
{
environment[self.disableAccessEnvironmentKey] == "1"
}
public static func forceDisabledForProcess(reason: String) {
self.processForceDisabledLock.lock()
self.processForceDisabledReason = reason
self.processForceDisabledLock.unlock()
#if os(macOS) && canImport(SweetCookieKit)
BrowserCookieKeychainAccessGate.isDisabled = self.isDisabled
#endif
}
public static var processDisableReason: String? {
self.processForceDisabledLock.lock()
defer { self.processForceDisabledLock.unlock() }
return self.processForceDisabledReason
}
#if DEBUG
private nonisolated(unsafe) static var forcesDisabledUnderTests: Bool {
KeychainTestSafety.shouldBlockRealKeychainAccess()
}
#endif
static func withTaskOverrideForTesting<T>(
_ disabled: Bool?,
operation: () throws -> T) rethrows -> T
{
try self.$taskOverrideValue.withValue(disabled) {
try operation()
}
}
static func withTaskOverrideForTesting<T>(
_ disabled: Bool?,
isolation _: isolated (any Actor)? = #isolation,
operation: () async throws -> T) async rethrows -> T
{
try await self.$taskOverrideValue.withValue(disabled) {
try await operation()
}
}
static var currentOverrideForTesting: Bool? {
self.taskOverrideValue ?? self.overrideValue
}
#if DEBUG
static func resetOverrideForTesting() {
self.overrideValue = nil
self.processForceDisabledLock.lock()
self.processForceDisabledReason = nil
self.processForceDisabledLock.unlock()
#if os(macOS) && canImport(SweetCookieKit)
BrowserCookieKeychainAccessGate.isDisabled = self.isDisabled
#endif
}
#endif
}
@@ -0,0 +1,196 @@
#if os(macOS)
import LocalAuthentication
import Security
#endif
public struct KeychainPromptContext: Sendable {
public enum Kind: Sendable {
case claudeOAuth
case codexCookie
case claudeCookie
case cursorCookie
case opencodeCookie
case factoryCookie
case zaiToken
case syntheticToken
case copilotToken
case kimiToken
case kimiK2Token
case minimaxCookie
case minimaxToken
case augmentCookie
case ampCookie
}
public let kind: Kind
public let service: String
public let account: String?
public init(kind: Kind, service: String, account: String?) {
self.kind = kind
self.service = service
self.account = account
}
}
public enum KeychainPromptHandler {
final class HandlerStore: @unchecked Sendable {
let handler: (KeychainPromptContext) -> Void
init(handler: @escaping (KeychainPromptContext) -> Void) {
self.handler = handler
}
}
@TaskLocal private static var taskHandlerStore: HandlerStore?
public nonisolated(unsafe) static var handler: ((KeychainPromptContext) -> Void)?
public static func notify(_ context: KeychainPromptContext) {
_ = self.notifyIfHandled(context)
}
@discardableResult
static func notifyIfHandled(_ context: KeychainPromptContext) -> Bool {
if let taskHandlerStore {
taskHandlerStore.handler(context)
return true
}
guard let handler else { return false }
handler(context)
return true
}
#if DEBUG
static func withHandlerForTesting<T>(
_ handler: ((KeychainPromptContext) -> Void)?,
operation: () throws -> T) rethrows -> T
{
try self.$taskHandlerStore.withValue(handler.map(HandlerStore.init(handler:))) {
try operation()
}
}
static func withHandlerForTesting<T>(
_ handler: ((KeychainPromptContext) -> Void)?,
operation: () async throws -> T) async rethrows -> T
{
try await self.$taskHandlerStore.withValue(handler.map(HandlerStore.init(handler:))) {
try await operation()
}
}
#endif
}
public enum KeychainAccessPreflight {
public enum Outcome: Sendable {
case allowed
case interactionRequired
case notFound
case failure(Int)
}
private static let log = CodexBarLog.logger(LogCategories.keychainPreflight)
#if DEBUG
final class CheckGenericPasswordOverrideStore: @unchecked Sendable {
let check: (String, String?) -> Outcome
init(check: @escaping (String, String?) -> Outcome) {
self.check = check
}
}
@TaskLocal private static var taskCheckGenericPasswordOverrideStore: CheckGenericPasswordOverrideStore?
private nonisolated(unsafe) static var checkGenericPasswordOverride: ((String, String?) -> Outcome)?
static func setCheckGenericPasswordOverrideForTesting(_ override: ((String, String?) -> Outcome)?) {
self.checkGenericPasswordOverride = override
}
static var hasCheckGenericPasswordOverrideForTesting: Bool {
self.taskCheckGenericPasswordOverrideStore != nil || self.checkGenericPasswordOverride != nil
}
static func withCheckGenericPasswordOverrideForTesting<T>(
_ override: ((String, String?) -> Outcome)?,
operation: () throws -> T) rethrows -> T
{
try self.$taskCheckGenericPasswordOverrideStore.withValue(
override.map(CheckGenericPasswordOverrideStore.init(check:)))
{
try operation()
}
}
static func withCheckGenericPasswordOverrideForTesting<T>(
_ override: ((String, String?) -> Outcome)?,
isolation _: isolated (any Actor)? = #isolation,
operation: () async throws -> T) async rethrows -> T
{
try await self.$taskCheckGenericPasswordOverrideStore.withValue(
override.map(CheckGenericPasswordOverrideStore.init(check:)))
{
try await operation()
}
}
#endif
public static func checkGenericPassword(service: String, account: String?) -> Outcome {
#if os(macOS)
#if DEBUG
if let override = self.taskCheckGenericPasswordOverrideStore {
return override.check(service, account)
}
if let override = self.checkGenericPasswordOverride {
return override(service, account)
}
#endif
guard !KeychainAccessGate.isDisabled else { return .notFound }
let query = self.makeGenericPasswordPreflightQuery(service: service, account: account)
var result: AnyObject?
let status = KeychainSecurity.copyMatching(query as CFDictionary, &result)
switch status {
case errSecSuccess:
self.log.debug("Keychain preflight allowed", metadata: ["service": service])
return .allowed
case errSecItemNotFound:
self.log.debug(
"Keychain preflight not found",
metadata: ["service": service])
return .notFound
case errSecInteractionNotAllowed:
self.log.info(
"Keychain preflight requires interaction",
metadata: ["service": service])
return .interactionRequired
default:
self.log.warning(
"Keychain preflight failed",
metadata: ["service": service, "status": "\(status)"])
return .failure(Int(status))
}
#else
return .notFound
#endif
}
#if os(macOS)
static func makeGenericPasswordPreflightQuery(service: String, account: String?) -> [String: Any] {
var query: [String: Any] = [
kSecClass as String: kSecClassGenericPassword,
kSecAttrService as String: service,
kSecMatchLimit as String: kSecMatchLimitOne,
// Preflight should never trigger UI. Avoid requesting the secret payload (`kSecReturnData`) because
// some macOS configurations have been observed to show the legacy keychain prompt unless the query
// is strictly non-interactive.
kSecReturnAttributes as String: true,
]
KeychainNoUIQuery.apply(to: &query)
if let account {
query[kSecAttrAccount as String] = account
}
return query
}
#endif
}
@@ -0,0 +1,665 @@
import Foundation
#if os(macOS)
import Darwin
import Security
#endif
public enum KeychainCacheStore {
public struct Key: Hashable, Sendable {
public let category: String
public let identifier: String
public init(category: String, identifier: String) {
self.category = category
self.identifier = identifier
}
var account: String {
"\(self.category).\(self.identifier)"
}
}
public enum LoadResult<Entry> {
case found(Entry)
case missing
case temporarilyUnavailable
case invalid
}
public enum ClearResult: Equatable, Sendable {
case removed
case missing
case failed
}
public enum KeysResult: Equatable, Sendable {
case found([Key])
case temporarilyUnavailable
case failed
}
private static let log = CodexBarLog.logger(LogCategories.keychainCache)
private static let cacheService = "com.steipete.codexbar.cache"
private static let cacheLabel = "CodexBar Cache"
private nonisolated(unsafe) static var globalServiceOverride: String?
@TaskLocal private static var serviceOverride: String?
#if DEBUG
@TaskLocal private static var operationRecorder: OperationRecorder?
enum Operation: Equatable, Sendable {
case load
case store
case clear
}
final class OperationRecorder: @unchecked Sendable {
private let lock = NSLock()
private var recordedOperations: [Operation] = []
var operations: [Operation] {
self.lock.withLock { self.recordedOperations }
}
func record(_ operation: Operation) {
self.lock.withLock {
self.recordedOperations.append(operation)
}
}
}
#endif
#if DEBUG && os(macOS)
@TaskLocal private static var loadFailureStatusOverride: OSStatus?
@TaskLocal private static var storeFailureStatusOverride: OSStatus?
@TaskLocal private static var clearFailureStatusOverride: OSStatus?
@TaskLocal private static var keysFailureStatusOverride: OSStatus?
#endif
private static let testStoreLock = NSLock()
private struct TestStoreKey: Hashable {
let service: String
let account: String
}
private nonisolated(unsafe) static var testStore: [TestStoreKey: Data]?
private nonisolated(unsafe) static var implicitTestStore: [TestStoreKey: Data] = [:]
private nonisolated(unsafe) static var testStoreRefCount = 0
public static func load<Entry: Codable>(
key: Key,
as type: Entry.Type = Entry.self) -> LoadResult<Entry>
{
#if DEBUG
self.operationRecorder?.record(.load)
#endif
#if DEBUG && os(macOS)
if let status = self.loadFailureStatusOverride {
return self.loadResultForKeychainReadFailure(status: status, key: key)
}
#endif
if let testResult = loadFromTestStore(key: key, as: type) {
return testResult
}
guard self.canUseRealKeychain else { return .missing }
#if os(macOS)
var query: [String: Any] = [
kSecClass as String: kSecClassGenericPassword,
kSecAttrService as String: self.serviceName,
kSecAttrAccount as String: key.account,
kSecMatchLimit as String: kSecMatchLimitOne,
kSecReturnData as String: true,
]
KeychainNoUIQuery.apply(to: &query)
var result: AnyObject?
let status = KeychainSecurity.copyMatching(query as CFDictionary, &result)
switch status {
case errSecSuccess:
guard let data = result as? Data, !data.isEmpty else {
self.log.error("Keychain cache item was empty (\(key.account))")
return .invalid
}
let decoder = Self.makeDecoder()
guard let decoded = try? decoder.decode(Entry.self, from: data) else {
self.log.error("Failed to decode keychain cache (\(key.account))")
return .invalid
}
return .found(decoded)
default:
return self.loadResultForKeychainReadFailure(status: status, key: key)
}
#else
return .missing
#endif
}
public static func store(key: Key, entry: some Codable) {
_ = self.storeResult(key: key, entry: entry)
}
@discardableResult
public static func storeResult(key: Key, entry: some Codable) -> Bool {
#if DEBUG
self.operationRecorder?.record(.store)
#endif
#if DEBUG && os(macOS)
if let status = self.storeFailureStatusOverride {
self.log.error("Keychain cache store failed (\(key.account)): \(status)")
return false
}
#endif
if let stored = self.storeInTestStore(key: key, entry: entry) {
return stored
}
guard self.canUseRealKeychain else { return false }
#if os(macOS)
let encoder = Self.makeEncoder()
guard let data = try? encoder.encode(entry) else {
self.log.error("Failed to encode keychain cache (\(key.account))")
return false
}
var query: [String: Any] = [
kSecClass as String: kSecClassGenericPassword,
kSecAttrService as String: self.serviceName,
kSecAttrAccount as String: key.account,
]
KeychainNoUIQuery.apply(to: &query)
let updateStatus = KeychainSecurity.update(
query as CFDictionary,
[kSecValueData as String: data] as CFDictionary)
if updateStatus == errSecSuccess {
return true
}
if updateStatus != errSecItemNotFound {
self.log.error("Keychain cache update failed (\(key.account)): \(updateStatus)")
return false
}
var addQuery = query
addQuery[kSecValueData as String] = data
addQuery[kSecAttrLabel as String] = self.cacheLabel
addQuery[kSecAttrAccessible as String] = kSecAttrAccessibleAfterFirstUnlockThisDeviceOnly
if let access = self.cacheAccessControl() {
addQuery[kSecAttrAccess as String] = access
}
let addStatus = KeychainSecurity.add(addQuery as CFDictionary, nil)
if addStatus != errSecSuccess {
self.log.error("Keychain cache add failed (\(key.account)): \(addStatus)")
}
return addStatus == errSecSuccess
#else
return false
#endif
}
@discardableResult
public static func clear(key: Key) -> Bool {
self.clearResult(key: key) == .removed
}
public static func clearResult(key: Key) -> ClearResult {
#if DEBUG
self.operationRecorder?.record(.clear)
#endif
#if DEBUG && os(macOS)
if let status = self.clearFailureStatusOverride {
return self.clearResultForKeychainDeleteStatus(status, key: key)
}
#endif
if let removed = self.clearTestStore(key: key) {
return removed ? .removed : .missing
}
guard self.canUseRealKeychain else { return .failed }
#if os(macOS)
var query: [String: Any] = [
kSecClass as String: kSecClassGenericPassword,
kSecAttrService as String: self.serviceName,
kSecAttrAccount as String: key.account,
]
KeychainNoUIQuery.apply(to: &query)
return self.clearResultForKeychainDeleteStatus(KeychainSecurity.delete(query as CFDictionary), key: key)
#else
return .failed
#endif
}
public static func keys(category: String) -> [Key] {
switch self.keysResult(category: category) {
case let .found(keys):
keys
case .temporarilyUnavailable, .failed:
[]
}
}
public static func keysResult(category: String) -> KeysResult {
#if DEBUG && os(macOS)
if let status = self.keysFailureStatusOverride {
return self.keysResultForKeychainStatus(status, category: category, result: nil)
}
#endif
if let keys = self.keysFromTestStore(category: category) {
return .found(keys)
}
guard self.canUseRealKeychain else { return .failed }
#if os(macOS)
var query: [String: Any] = [
kSecClass as String: kSecClassGenericPassword,
kSecAttrService as String: self.serviceName,
kSecMatchLimit as String: kSecMatchLimitAll,
kSecReturnAttributes as String: true,
]
KeychainNoUIQuery.apply(to: &query)
var result: AnyObject?
let status = KeychainSecurity.copyMatching(query as CFDictionary, &result)
return self.keysResultForKeychainStatus(status, category: category, result: result)
#else
return .failed
#endif
}
static func setServiceOverrideForTesting(_ service: String?) {
self.globalServiceOverride = service
}
public static func withServiceOverrideForTesting<T>(
_ service: String?,
operation: () throws -> T) rethrows -> T
{
try self.$serviceOverride.withValue(service) {
try operation()
}
}
public static func withServiceOverrideForTesting<T>(
_ service: String?,
operation: () async throws -> T) async rethrows -> T
{
try await self.$serviceOverride.withValue(service) {
try await operation()
}
}
public static func withCurrentServiceOverrideForTesting<T>(
operation: () async throws -> T) async rethrows -> T
{
let service = self.serviceOverride
return try await self.$serviceOverride.withValue(service) {
try await operation()
}
}
public static var currentServiceOverrideForTesting: String? {
self.serviceOverride
}
#if DEBUG
static func withOperationRecorderForTesting<T>(
_ recorder: OperationRecorder?,
operation: () throws -> T) rethrows -> T
{
try self.$operationRecorder.withValue(recorder) {
try operation()
}
}
static func withOperationRecorderForTesting<T>(
_ recorder: OperationRecorder?,
operation: () async throws -> T) async rethrows -> T
{
try await self.$operationRecorder.withValue(recorder) {
try await operation()
}
}
static var currentOperationRecorderForTesting: OperationRecorder? {
self.operationRecorder
}
#endif
static var canUseRealKeychainForTesting: Bool {
self.canUseRealKeychain
}
static var canEnumerateOrDeleteRealKeychainForTesting: Bool {
self.canUseRealKeychain
}
#if DEBUG && os(macOS)
public static func withLoadFailureStatusOverrideForTesting<T>(
_ status: OSStatus?,
operation: () throws -> T) rethrows -> T
{
try self.$loadFailureStatusOverride.withValue(status) {
try operation()
}
}
public static func withStoreFailureStatusOverrideForTesting<T>(
_ status: OSStatus?,
operation: () throws -> T) rethrows -> T
{
try self.$storeFailureStatusOverride.withValue(status) {
try operation()
}
}
public static func withClearFailureStatusOverrideForTesting<T>(
_ status: OSStatus?,
operation: () throws -> T) rethrows -> T
{
try self.$clearFailureStatusOverride.withValue(status) {
try operation()
}
}
public static func withClearFailureStatusOverrideForTesting<T>(
_ status: OSStatus?,
operation: () async throws -> T) async rethrows -> T
{
try await self.$clearFailureStatusOverride.withValue(status) {
try await operation()
}
}
public static func withKeysFailureStatusOverrideForTesting<T>(
_ status: OSStatus?,
operation: () throws -> T) rethrows -> T
{
try self.$keysFailureStatusOverride.withValue(status) {
try operation()
}
}
#endif
static func setTestStoreForTesting(_ enabled: Bool) {
self.testStoreLock.lock()
defer { self.testStoreLock.unlock() }
if enabled {
self.testStoreRefCount += 1
if self.testStoreRefCount == 1 {
self.testStore = [:]
}
} else {
self.testStoreRefCount = max(0, self.testStoreRefCount - 1)
if self.testStoreRefCount == 0 {
self.testStore = nil
}
}
}
private static var serviceName: String {
serviceOverride ?? self.globalServiceOverride ?? self.cacheService
}
private static var canUseRealKeychain: Bool {
!KeychainAccessGate.isDisabled
}
#if DEBUG
private static var shouldUseImplicitTestStore: Bool {
KeychainTestSafety.isRunningUnderTests(
processName: ProcessInfo.processInfo.processName,
environment: ProcessInfo.processInfo.environment) && !self.canUseRealKeychain
}
#else
private static var shouldUseImplicitTestStore: Bool {
false
}
#endif
private static func makeEncoder() -> JSONEncoder {
let encoder = JSONEncoder()
encoder.dateEncodingStrategy = .iso8601
return encoder
}
private static func makeDecoder() -> JSONDecoder {
let decoder = JSONDecoder()
decoder.dateDecodingStrategy = .iso8601
return decoder
}
#if os(macOS)
static func loadResultForKeychainReadFailure<Entry>(
status: OSStatus,
key: Key) -> LoadResult<Entry>
{
switch status {
case errSecItemNotFound:
return .missing
case errSecInteractionNotAllowed:
// Keychain is temporarily locked, e.g. immediately after wake from sleep.
self.log.info("Keychain cache temporarily locked (\(key.account)), will retry on next access")
return .temporarilyUnavailable
default:
self.log.error("Keychain cache read failed (\(key.account)): \(status)")
return .invalid
}
}
static func clearResultForKeychainDeleteStatus(_ status: OSStatus, key: Key) -> ClearResult {
switch status {
case errSecSuccess:
return .removed
case errSecItemNotFound:
return .missing
case errSecInteractionNotAllowed:
self.log.info("Keychain cache delete temporarily unavailable (\(key.account))")
return .failed
default:
self.log.error("Keychain cache delete failed (\(key.account)): \(status)")
return .failed
}
}
private static func keysResultForKeychainStatus(
_ status: OSStatus,
category: String,
result: AnyObject?) -> KeysResult
{
switch status {
case errSecSuccess:
guard let rows = result as? [[String: Any]] else { return .failed }
let keys: [Key] = rows.compactMap { row in
guard let account = row[kSecAttrAccount as String] as? String else { return nil }
return self.key(fromAccount: account, category: category)
}
return .found(keys)
case errSecItemNotFound:
return .found([])
case errSecInteractionNotAllowed:
self.log.info("Keychain cache keys temporarily unavailable (\(category))")
return .temporarilyUnavailable
default:
self.log.error("Keychain cache key listing failed (\(category)): \(status)")
return .failed
}
}
static func trustedApplicationPathsForCacheAccess(
bundleURL: URL = Bundle.main.bundleURL,
executableURL: URL? = Bundle.main.executableURL,
fileExists: (String) -> Bool = { FileManager.default.fileExists(atPath: $0) }) -> [String]
{
var paths: [String] = []
func append(_ path: String) {
guard !path.isEmpty, fileExists(path), !paths.contains(path) else { return }
paths.append(path)
}
let appBundle = self.appBundleURL(containing: bundleURL)
?? executableURL.flatMap(self.appBundleURL(containing:))
if let appBundle {
append(appBundle.path)
append(appBundle.appendingPathComponent("Contents/Helpers/CodexBarCLI").path)
}
if let executableURL {
append(executableURL.path)
}
return paths
}
private static func appBundleURL(containing url: URL) -> URL? {
var current = url.standardizedFileURL
while current.path != "/" {
if current.pathExtension == "app" {
return current
}
current.deleteLastPathComponent()
}
return nil
}
private static func cacheAccessControl() -> SecAccess? {
let trustedPaths = self.trustedApplicationPathsForCacheAccess()
guard !trustedPaths.isEmpty else { return nil }
var trustedApplications: [SecTrustedApplication] = []
for path in trustedPaths {
let (status, application) = self.createTrustedApplication(path: path)
if status == errSecSuccess, let application {
trustedApplications.append(application)
} else {
self.log.error("Keychain cache trusted app creation failed (\(path)): \(status)")
}
}
guard !trustedApplications.isEmpty else { return nil }
let (status, access) = self.createAccessControl(trustedApplications: trustedApplications)
if status != errSecSuccess {
self.log.error("Keychain cache access control creation failed: \(status)")
return nil
}
return access
}
private typealias SecTrustedApplicationCreateFromPathFunction = @convention(c) (
UnsafePointer<CChar>?,
UnsafeMutablePointer<SecTrustedApplication?>?) -> OSStatus
private typealias SecAccessCreateFunction = @convention(c) (
CFString,
CFArray,
UnsafeMutablePointer<SecAccess?>?) -> OSStatus
private static func createTrustedApplication(path: String) -> (OSStatus, SecTrustedApplication?) {
guard let symbol = self.securitySymbol(named: "SecTrustedApplicationCreateFromPath") else {
return (errSecInternalComponent, nil)
}
let function = unsafeBitCast(symbol, to: SecTrustedApplicationCreateFromPathFunction.self)
var application: SecTrustedApplication?
let status = path.withCString { cPath in
function(cPath, &application)
}
return (status, application)
}
private static func createAccessControl(trustedApplications: [SecTrustedApplication]) -> (OSStatus, SecAccess?) {
guard let symbol = self.securitySymbol(named: "SecAccessCreate") else {
return (errSecInternalComponent, nil)
}
let function = unsafeBitCast(symbol, to: SecAccessCreateFunction.self)
var access: SecAccess?
let status = function(self.cacheLabel as CFString, trustedApplications as CFArray, &access)
return (status, access)
}
private nonisolated(unsafe) static let securityFrameworkHandle: UnsafeMutableRawPointer? = {
let securityPath = "/System/Library/Frameworks/Security.framework/Security"
return dlopen(securityPath, RTLD_NOW)
}()
private static func securitySymbol(named name: String) -> UnsafeMutableRawPointer? {
// Resolve deprecated SecKeychain ACL helpers at runtime so release builds stay warning-free
// while still granting the app bundle and bundled CLI prompt-free access to cache entries.
guard let securityFrameworkHandle else { return nil }
return dlsym(securityFrameworkHandle, name)
}
#endif
private static func loadFromTestStore<Entry: Codable>(
key: Key,
as type: Entry.Type) -> LoadResult<Entry>?
{
self.testStoreLock.lock()
defer { self.testStoreLock.unlock() }
guard let store = self.testStore ?? (self.shouldUseImplicitTestStore ? self.implicitTestStore : nil)
else { return nil }
let testKey = TestStoreKey(service: self.serviceName, account: key.account)
guard let data = store[testKey] else { return .missing }
let decoder = Self.makeDecoder()
guard let decoded = try? decoder.decode(Entry.self, from: data) else {
return .invalid
}
return .found(decoded)
}
private static func storeInTestStore(key: Key, entry: some Codable) -> Bool? {
self.testStoreLock.lock()
defer { self.testStoreLock.unlock() }
let encoder = Self.makeEncoder()
guard let data = try? encoder.encode(entry) else { return false }
let testKey = TestStoreKey(service: self.serviceName, account: key.account)
if var store = self.testStore {
store[testKey] = data
self.testStore = store
return true
}
if self.shouldUseImplicitTestStore {
self.implicitTestStore[testKey] = data
return true
}
return nil
}
private static func clearTestStore(key: Key) -> Bool? {
self.testStoreLock.lock()
defer { self.testStoreLock.unlock() }
let testKey = TestStoreKey(service: self.serviceName, account: key.account)
if var store = self.testStore {
let removed = store.removeValue(forKey: testKey) != nil
self.testStore = store
return removed
}
if self.shouldUseImplicitTestStore {
return self.implicitTestStore.removeValue(forKey: testKey) != nil
}
return nil
}
private static func keysFromTestStore(category: String) -> [Key]? {
self.testStoreLock.lock()
defer { self.testStoreLock.unlock() }
guard let store = self.testStore ?? (self.shouldUseImplicitTestStore ? self.implicitTestStore : nil)
else { return nil }
return store.keys
.filter { $0.service == self.serviceName }
.compactMap { self.key(fromAccount: $0.account, category: category) }
.sorted { $0.identifier < $1.identifier }
}
private static func key(fromAccount account: String, category: String) -> Key? {
let prefix = "\(category)."
guard account.hasPrefix(prefix) else { return nil }
let identifier = String(account.dropFirst(prefix.count))
guard !identifier.isEmpty else { return nil }
return Key(category: category, identifier: identifier)
}
}
extension KeychainCacheStore.Key {
public static func cookie(provider: UsageProvider, scopeIdentifier: String? = nil) -> Self {
let identifier: String = if let scopeIdentifier, !scopeIdentifier.isEmpty {
"\(provider.rawValue).\(scopeIdentifier)"
} else {
provider.rawValue
}
return Self(category: "cookie", identifier: identifier)
}
public static func oauth(provider: UsageProvider) -> Self {
Self(category: "oauth", identifier: provider.rawValue)
}
}
@@ -0,0 +1,41 @@
import Foundation
#if os(macOS)
import Darwin
import LocalAuthentication
import Security
enum KeychainNoUIQuery {
private static let uiFailPolicy = KeychainNoUIQuery.resolveUIFailPolicy()
static func apply(to query: inout [String: Any]) {
let context = LAContext()
context.interactionNotAllowed = true
query[kSecUseAuthenticationContext as String] = context
// Keep explicit UI-fail policy for legacy keychain behavior on macOS where
// `interactionNotAllowed` alone can still surface Allow/Deny prompts.
query[kSecUseAuthenticationUI as String] = self.uiFailPolicy as CFString
}
static func uiFailPolicyForTesting() -> String {
self.uiFailPolicy
}
private static func resolveUIFailPolicy() -> String {
// Resolve the Security symbol at runtime to preserve the true constant value
// without directly referencing deprecated API at compile time.
let securityPath = "/System/Library/Frameworks/Security.framework/Security"
guard let handle = dlopen(securityPath, RTLD_NOW) else {
return "u_AuthUIF"
}
defer { dlclose(handle) }
guard let symbol = dlsym(handle, "kSecUseAuthenticationUIFail") else {
return "u_AuthUIF"
}
let valuePointer = symbol.assumingMemoryBound(to: CFString?.self)
return (valuePointer.pointee as String?) ?? "u_AuthUIF"
}
}
#endif
@@ -0,0 +1,74 @@
import Foundation
enum KeychainTestSafety {
static let suppressAccessEnvironmentKey = "CODEXBAR_SUPPRESS_TEST_KEYCHAIN_ACCESS"
static let allowAccessEnvironmentKey = "CODEXBAR_ALLOW_TEST_KEYCHAIN_ACCESS"
static func shouldBlockRealKeychainAccess(
processName: String = ProcessInfo.processInfo.processName,
environment: [String: String] = ProcessInfo.processInfo.environment) -> Bool
{
if environment[self.allowAccessEnvironmentKey] == "1" { return false }
if environment[self.suppressAccessEnvironmentKey] == "1" { return true }
if environment[KeychainAccessGate.disableAccessEnvironmentKey] == "1" { return true }
return self.isRunningUnderTests(processName: processName, environment: environment)
}
static func isRunningUnderTests(
processName: String,
environment: [String: String]) -> Bool
{
processName == "swiftpm-testing-helper"
|| processName.hasSuffix("PackageTests")
|| processName.hasSuffix(".xctest")
|| environment["XCTestConfigurationFilePath"] != nil
|| environment["XCTestBundlePath"] != nil
|| environment["XCTestSessionIdentifier"] != nil
|| environment["TESTING_LIBRARY_VERSION"] != nil
|| environment["SWIFT_TESTING"] != nil
|| environment["SWIFT_TESTING_ENABLED"] != nil
}
}
#if os(macOS)
import Security
/// The only first-party entry point for Security.framework item operations.
/// Test processes fail closed before touching the user's Keychain, even when a test enables
/// higher-level Keychain logic with `KeychainAccessGate.withTaskOverrideForTesting(false)`.
public enum KeychainSecurity {
public static func copyMatching(
_ query: CFDictionary,
_ result: UnsafeMutablePointer<CFTypeRef?>?) -> OSStatus
{
guard !KeychainTestSafety.shouldBlockRealKeychainAccess() else {
return errSecInteractionNotAllowed
}
return SecItemCopyMatching(query, result)
}
public static func update(_ query: CFDictionary, _ attributesToUpdate: CFDictionary) -> OSStatus {
guard !KeychainTestSafety.shouldBlockRealKeychainAccess() else {
return errSecInteractionNotAllowed
}
return SecItemUpdate(query, attributesToUpdate)
}
public static func add(
_ attributes: CFDictionary,
_ result: UnsafeMutablePointer<CFTypeRef?>?) -> OSStatus
{
guard !KeychainTestSafety.shouldBlockRealKeychainAccess() else {
return errSecInteractionNotAllowed
}
return SecItemAdd(attributes, result)
}
public static func delete(_ query: CFDictionary) -> OSStatus {
guard !KeychainTestSafety.shouldBlockRealKeychainAccess() else {
return errSecInteractionNotAllowed
}
return SecItemDelete(query)
}
}
#endif
@@ -0,0 +1,229 @@
import Foundation
public struct LocalAgentSessionScanner: Sendable {
private struct Rollout: Sendable {
let url: URL
let modifiedAt: Date
let metadata: CodexRolloutMetadata
}
private struct ScanContext: Sendable {
let homeDirectory: URL
let host: String
let now: Date
}
public let config: SessionScanConfig
public init(config: SessionScanConfig = SessionScanConfig()) {
self.config = config
}
public func scan(
now: Date = Date(),
environment: [String: String] = ProcessInfo.processInfo.environment) async -> [AgentSession]
{
let processOutput = await self.processOutput(environment: environment)
let allProcesses = AgentPSOutputParser.parse(processOutput)
let processes = AgentSessionCorrelation.newestProcessesFirst(
AgentPSOutputParser.agentProcesses(from: allProcesses))
let codexAppServerPresent = AgentPSOutputParser.hasCodexAppServer(in: allProcesses)
let cwdByPID = await self.cwdByPID(processes.map(\ .pid), environment: environment)
let homeDirectory = URL(fileURLWithPath: environment["HOME"] ?? NSHomeDirectory(), isDirectory: true)
let host = ProcessInfo.processInfo.hostName
let rollouts = self.codexRollouts(
now: now,
environment: environment,
homeDirectory: homeDirectory)
return self.sessions(
processes: processes,
cwdByPID: cwdByPID,
rollouts: rollouts,
codexAppServerPresent: codexAppServerPresent,
context: ScanContext(homeDirectory: homeDirectory, host: host, now: now))
}
private func sessions(
processes: [AgentProcessRecord],
cwdByPID: [Int32: String],
rollouts: [Rollout],
codexAppServerPresent: Bool,
context: ScanContext) -> [AgentSession]
{
var sessions: [AgentSession] = []
var matchedRolloutPaths = Set<String>()
let claudeProcesses = processes.filter { AgentPSOutputParser.provider(for: $0) == .claude }
let claudeCWDs = Set(claudeProcesses.compactMap { cwdByPID[$0.pid] })
let claudeTranscriptsByCWD = Dictionary(uniqueKeysWithValues: claudeCWDs.map { cwd in
(cwd, ClaudeSessionProjectMapper.transcripts(cwd: cwd, homeDirectory: context.homeDirectory))
})
let claudeTranscripts = AgentSessionCorrelation.assignClaudeTranscripts(
processes: claudeProcesses,
cwdByPID: cwdByPID,
transcriptsByCWD: claudeTranscriptsByCWD)
for process in processes {
guard let provider = AgentPSOutputParser.provider(for: process) else { continue }
let cwd = cwdByPID[process.pid]
switch provider {
case .claude:
let transcript = claudeTranscripts[process.pid]
sessions.append(AgentSession(
id: transcript?.url.deletingPathExtension().lastPathComponent ?? "pid:\(process.pid)",
provider: .claude,
source: AgentPSOutputParser.source(for: process),
state: self.config.state(
lastActivityAt: transcript?.modifiedAt,
now: context.now,
hasLiveProcess: true),
pid: process.pid,
cwd: cwd,
projectName: Self.projectName(cwd),
startedAt: process.startedAt,
lastActivityAt: transcript?.modifiedAt,
transcriptPath: transcript?.url.path,
host: context.host))
case .codex:
let rollout = rollouts.first { candidate in
!matchedRolloutPaths.contains(candidate.url.path) &&
AgentSessionCorrelation.codexWorkingDirectoriesMatch(candidate.metadata.cwd, cwd)
}
if let rollout {
matchedRolloutPaths.insert(rollout.url.path)
}
let rolloutSource = rollout?.metadata.sessionSource
sessions.append(AgentSession(
id: rollout?.metadata.sessionID ?? "pid:\(process.pid)",
provider: .codex,
source: rolloutSource == nil || rolloutSource == .unknown ? .cli : rolloutSource ?? .cli,
state: self.config.state(
lastActivityAt: rollout?.modifiedAt,
now: context.now,
hasLiveProcess: true),
pid: process.pid,
cwd: cwd ?? rollout?.metadata.cwd,
projectName: Self.projectName(cwd ?? rollout?.metadata.cwd),
startedAt: process.startedAt,
lastActivityAt: rollout?.modifiedAt,
transcriptPath: rollout?.url.path,
host: context.host))
}
}
for rollout in rollouts where !matchedRolloutPaths.contains(rollout.url.path) {
guard var session = CodexRolloutFirstLineParser.makeSession(
metadata: rollout.metadata,
transcriptURL: rollout.url,
modifiedAt: rollout.modifiedAt,
host: context.host,
config: self.config,
now: context.now)
else { continue }
session.source = AgentSessionCorrelation.fileOnlyCodexSource(
metadataSource: session.source,
appServerPresent: codexAppServerPresent)
sessions.append(session)
}
var seen = Set<String>()
return sessions
.sorted { lhs, rhs in
if lhs.state != rhs.state {
return lhs.state == .active
}
return (lhs.lastActivityAt ?? lhs.startedAt ?? .distantPast) >
(rhs.lastActivityAt ?? rhs.startedAt ?? .distantPast)
}
.filter { seen.insert("\($0.host):\($0.id)").inserted }
}
private func processOutput(environment: [String: String]) async -> String {
let binary = ["/bin/ps", "/usr/bin/ps"].first { FileManager.default.isExecutableFile(atPath: $0) }
guard let binary,
let result = try? await SubprocessRunner.run(
binary: binary,
arguments: ["-axo", "pid=,ppid=,lstart=,command="],
environment: environment,
timeout: 5,
label: "agent session process scan")
else { return "" }
return result.stdout
}
private func cwdByPID(_ pids: [Int32], environment: [String: String]) async -> [Int32: String] {
guard !pids.isEmpty else { return [:] }
if let lsof = self.findExecutable("lsof", environment: environment) {
let joinedPIDs = pids.map(String.init).joined(separator: ",")
if let result = try? await SubprocessRunner.run(
binary: lsof,
arguments: ["-a", "-d", "cwd", "-Fn", "-p", joinedPIDs],
environment: environment,
timeout: 5,
acceptsNonZeroExit: true,
label: "agent session cwd scan")
{
return LSOFCWDOutputParser.parse(result.stdout)
}
}
#if os(Linux)
return Dictionary(uniqueKeysWithValues: pids.compactMap { pid in
let path = "/proc/\(pid)/cwd"
guard let destination = try? FileManager.default.destinationOfSymbolicLink(atPath: path) else { return nil }
return (pid, destination)
})
#else
return [:]
#endif
}
private func codexRollouts(
now: Date,
environment: [String: String],
homeDirectory: URL) -> [Rollout]
{
let root = URL(
fileURLWithPath: environment["CODEX_HOME"] ?? homeDirectory.appendingPathComponent(".codex").path,
isDirectory: true)
.appendingPathComponent("sessions", isDirectory: true)
let calendar = Calendar(identifier: .gregorian)
let days = [now, calendar.date(byAdding: .day, value: -1, to: now)].compactMap(\.self)
let formatter = DateFormatter()
formatter.locale = Locale(identifier: "en_US_POSIX")
formatter.dateFormat = "yyyy/MM/dd"
let fileManager = FileManager.default
return days.flatMap { day -> [Rollout] in
let directory = root.appendingPathComponent(formatter.string(from: day), isDirectory: true)
guard let files = try? fileManager.contentsOfDirectory(
at: directory,
includingPropertiesForKeys: [.contentModificationDateKey],
options: [.skipsHiddenFiles])
else { return [] }
return files.compactMap { file in
guard file.lastPathComponent.hasPrefix("rollout-"), file.pathExtension == "jsonl",
let modifiedAt = try? file.resourceValues(
forKeys: [.contentModificationDateKey]).contentModificationDate,
let metadata = CodexRolloutFirstLineParser.read(from: file)
else { return nil }
return Rollout(url: file, modifiedAt: modifiedAt, metadata: metadata)
}
}.sorted { $0.modifiedAt > $1.modifiedAt }
}
private func findExecutable(_ name: String, environment: [String: String]) -> String? {
let path = environment["PATH"] ?? "/usr/local/bin:/usr/bin:/bin:/opt/homebrew/bin"
return path.split(separator: ":")
.map { String($0) + "/" + name }
.first { FileManager.default.isExecutableFile(atPath: $0) }
}
private static func standardized(_ path: String?) -> String? {
path.map { URL(fileURLWithPath: $0).standardizedFileURL.path }
}
private static func projectName(_ cwd: String?) -> String? {
guard let cwd, !cwd.isEmpty else { return nil }
return URL(fileURLWithPath: cwd).lastPathComponent
}
}
@@ -0,0 +1,229 @@
import Foundation
import Logging
public enum CodexBarLog {
public enum Destination: Sendable {
case stderr
case discard
case oslog(subsystem: String)
}
public enum Level: String, CaseIterable, Identifiable, Sendable {
case trace
case verbose
case debug
case info
case warning
case error
case critical
public var id: String {
self.rawValue
}
public var displayName: String {
switch self {
case .trace: "Trace"
case .verbose: "Verbose"
case .debug: "Debug"
case .info: "Info"
case .warning: "Warning"
case .error: "Error"
case .critical: "Critical"
}
}
public var rank: Int {
switch self {
case .trace: 0
case .verbose: 1
case .debug: 2
case .info: 3
case .warning: 4
case .error: 5
case .critical: 6
}
}
public var asSwiftLogLevel: Logger.Level {
switch self {
case .trace: .trace
case .verbose: .debug
case .debug: .debug
case .info: .info
case .warning: .warning
case .error: .error
case .critical: .critical
}
}
}
public struct Configuration: Sendable {
public let destination: Destination
public let level: Level
public let json: Bool
public init(destination: Destination, level: Level, json: Bool) {
self.destination = destination
self.level = level
self.json = json
}
}
private static let lock = NSLock()
private static let levelLock = NSLock()
private nonisolated(unsafe) static var isBootstrapped = false
private nonisolated(unsafe) static var currentLevel: Level = .info
public static func bootstrapIfNeeded(_ config: Configuration) {
self.lock.lock()
defer { lock.unlock() }
guard !self.isBootstrapped else { return }
self.currentLevel = config.level
let baseFactory: @Sendable (String) -> any LogHandler = { label in
switch config.destination {
case .stderr:
if config.json { return JSONStderrLogHandler(label: label) }
return StreamLogHandler.standardError(label: label)
case .discard:
return DiscardLogHandler()
case let .oslog(subsystem):
#if canImport(os)
return OSLogLogHandler(label: label, subsystem: subsystem)
#else
if config.json { return JSONStderrLogHandler(label: label) }
return StreamLogHandler.standardError(label: label)
#endif
}
}
LoggingSystem.bootstrap { label in
let primary = baseFactory(label)
let fileHandler = FileLogHandler(label: label)
var handler = CompositeLogHandler(primary: primary, secondary: fileHandler)
handler.logLevel = .trace
return handler
}
self.isBootstrapped = true
}
public static func logger(_ category: String) -> CodexBarLogger {
let logger = Logger(label: "com.steipete.codexbar.\(category)")
return CodexBarLogger(isEnabled: { level in
self.shouldLog(level)
}, log: { level, message, metadata in
let swiftLogLevel = level.asSwiftLogLevel
let safeMessage = LogRedactor.redact(message)
let meta = metadata?.reduce(into: Logger.Metadata()) { partial, entry in
partial[entry.key] = .string(LogRedactor.redact(entry.value))
}
logger.log(level: swiftLogLevel, "\(safeMessage)", metadata: meta)
})
}
public static func parseLevel(_ raw: String?) -> Level? {
guard let raw, !raw.isEmpty else { return nil }
return Level(rawValue: raw.trimmingCharacters(in: .whitespacesAndNewlines).lowercased())
}
public static func setLogLevel(_ level: Level) {
self.levelLock.lock()
self.currentLevel = level
self.levelLock.unlock()
let logger = self.logger(LogCategories.logging)
logger.info("Log level set to \(level.rawValue)")
}
public static func currentLogLevel() -> Level {
self.levelLock.lock()
defer { self.levelLock.unlock() }
return self.currentLevel
}
private static func shouldLog(_ level: Level) -> Bool {
level.rank >= self.currentLogLevel().rank
}
public static var fileLogURL: URL {
FileLogSink.defaultURL
}
public static func setFileLoggingEnabled(_ enabled: Bool) {
FileLogSink.shared.setEnabled(enabled, fileURL: self.fileLogURL)
let state = enabled ? "enabled" : "disabled"
let logger = self.logger(LogCategories.logging)
logger.info("File logging \(state)", metadata: ["path": self.fileLogURL.path])
}
}
private struct DiscardLogHandler: LogHandler {
var metadata: Logger.Metadata = [:]
var logLevel: Logger.Level = .critical
subscript(metadataKey metadataKey: String) -> Logger.Metadata.Value? {
get { self.metadata[metadataKey] }
set { self.metadata[metadataKey] = newValue }
}
func log(event _: LogEvent) {}
}
public struct CodexBarLogger: Sendable {
private let isEnabled: @Sendable (CodexBarLog.Level) -> Bool
private let logFn: @Sendable (CodexBarLog.Level, String, [String: String]?) -> Void
fileprivate init(
isEnabled: @escaping @Sendable (CodexBarLog.Level) -> Bool,
log: @escaping @Sendable (CodexBarLog.Level, String, [String: String]?) -> Void)
{
self.isEnabled = isEnabled
self.logFn = log
}
init(
minimumLevel: CodexBarLog.Level,
log: @escaping @Sendable (CodexBarLog.Level, String, [String: String]?) -> Void)
{
self.isEnabled = { level in level.rank >= minimumLevel.rank }
self.logFn = log
}
private func log(
_ level: CodexBarLog.Level,
_ message: @autoclosure () -> String,
metadata: [String: String]?)
{
guard self.isEnabled(level) else { return }
self.logFn(level, message(), metadata)
}
public func trace(_ message: @autoclosure () -> String, metadata: [String: String]? = nil) {
self.log(.trace, message(), metadata: metadata)
}
public func verbose(_ message: @autoclosure () -> String, metadata: [String: String]? = nil) {
self.log(.verbose, message(), metadata: metadata)
}
public func debug(_ message: @autoclosure () -> String, metadata: [String: String]? = nil) {
self.log(.debug, message(), metadata: metadata)
}
public func info(_ message: @autoclosure () -> String, metadata: [String: String]? = nil) {
self.log(.info, message(), metadata: metadata)
}
public func warning(_ message: @autoclosure () -> String, metadata: [String: String]? = nil) {
self.log(.warning, message(), metadata: metadata)
}
public func error(_ message: @autoclosure () -> String, metadata: [String: String]? = nil) {
self.log(.error, message(), metadata: metadata)
}
public func critical(_ message: @autoclosure () -> String, metadata: [String: String]? = nil) {
self.log(.critical, message(), metadata: metadata)
}
}
@@ -0,0 +1,40 @@
import Logging
struct CompositeLogHandler: LogHandler {
private var primary: any LogHandler
private var secondary: any LogHandler
init(primary: any LogHandler, secondary: any LogHandler) {
self.primary = primary
self.secondary = secondary
}
var metadata: Logger.Metadata {
get { self.primary.metadata }
set {
self.primary.metadata = newValue
self.secondary.metadata = newValue
}
}
var logLevel: Logger.Level {
get { self.primary.logLevel }
set {
self.primary.logLevel = newValue
self.secondary.logLevel = newValue
}
}
subscript(metadataKey metadataKey: String) -> Logger.Metadata.Value? {
get { self.primary[metadataKey: metadataKey] }
set {
self.primary[metadataKey: metadataKey] = newValue
self.secondary[metadataKey: metadataKey] = newValue
}
}
func log(event: LogEvent) {
self.primary.log(event: event)
self.secondary.log(event: event)
}
}
@@ -0,0 +1,150 @@
import Foundation
import Logging
final class FileLogSink: @unchecked Sendable {
static let shared = FileLogSink()
static let defaultURL: URL = {
let base = FileManager.default.urls(for: .libraryDirectory, in: .userDomainMask).first
?? FileManager.default.homeDirectoryForCurrentUser.appendingPathComponent("Library", isDirectory: true)
return base
.appendingPathComponent("Logs", isDirectory: true)
.appendingPathComponent("CodexBar", isDirectory: true)
.appendingPathComponent("CodexBar.log")
}()
private let queue = DispatchQueue(label: "com.steipete.codexbar.filelog", qos: .utility)
private let fileManager: FileManager
private var isEnabled = false
private var fileHandle: FileHandle?
private var fileURL: URL = FileLogSink.defaultURL
private let maxBytes: Int64 = 10 * 1024 * 1024
init(fileManager: FileManager = .default) {
self.fileManager = fileManager
}
func setEnabled(_ enabled: Bool, fileURL: URL = FileLogSink.defaultURL) {
self.queue.async {
self.isEnabled = enabled
self.fileURL = fileURL
if !enabled {
self.closeHandle()
return
}
_ = self.openHandleIfNeeded()
}
}
func currentURL() -> URL {
self.queue.sync { self.fileURL }
}
func isWriteEnabled() -> Bool {
self.queue.sync { self.isEnabled }
}
func write(_ text: String) {
self.queue.async {
guard self.isEnabled else { return }
guard let data = text.data(using: .utf8) else { return }
guard let handle = self.openHandleIfNeeded() else { return }
handle.write(data)
}
}
private func openHandleIfNeeded() -> FileHandle? {
if let handle = self.fileHandle { return handle }
do {
try self.prepareFile(at: self.fileURL)
let handle = try FileHandle(forWritingTo: self.fileURL)
handle.seekToEndOfFile()
self.fileHandle = handle
return handle
} catch {
return nil
}
}
private func prepareFile(at url: URL) throws {
let directory = url.deletingLastPathComponent()
if !self.fileManager.fileExists(atPath: directory.path) {
try self.fileManager.createDirectory(at: directory, withIntermediateDirectories: true)
}
if self.fileManager.fileExists(atPath: url.path) {
let attributes = try self.fileManager.attributesOfItem(atPath: url.path)
let size = (attributes[.size] as? NSNumber)?.int64Value ?? 0
if size > self.maxBytes {
try self.fileManager.removeItem(at: url)
}
}
if !self.fileManager.fileExists(atPath: url.path) {
_ = self.fileManager.createFile(atPath: url.path, contents: nil)
}
}
private func closeHandle() {
if let handle = self.fileHandle {
try? handle.close()
}
self.fileHandle = nil
}
}
struct FileLogHandler: LogHandler {
var metadata: Logger.Metadata = [:]
var logLevel: Logger.Level = .info
private let label: String
private let sink: FileLogSink
init(label: String, sink: FileLogSink = .shared) {
self.label = label
self.sink = sink
}
subscript(metadataKey metadataKey: String) -> Logger.Metadata.Value? {
get { self.metadata[metadataKey] }
set { self.metadata[metadataKey] = newValue }
}
func log(event: LogEvent) {
guard self.sink.isWriteEnabled() else { return }
let ts = Self.timestamp()
var combined = self.metadata
if let metadata = event.metadata { combined.merge(metadata, uniquingKeysWith: { _, new in new }) }
var metaText = ""
if !combined.isEmpty {
let pairs = combined
.sorted(by: { $0.key < $1.key })
.map { key, value in
let rendered = Self.renderMetadataValue(value)
let safeValue = LogRedactor.redact(rendered)
return "\(key)=\(safeValue)"
}
.joined(separator: " ")
metaText = " \(pairs)"
}
let safeMessage = LogRedactor.redact("\(event.message)")
let lineText = "[\(ts)] [\(event.level.rawValue.uppercased())] \(self.label): \(safeMessage)\(metaText)\n"
_ = event.source
_ = event.file
_ = event.function
_ = event.line
self.sink.write(lineText)
}
private static func timestamp() -> String {
let formatter = ISO8601DateFormatter()
formatter.formatOptions = [.withInternetDateTime, .withFractionalSeconds]
return formatter.string(from: Date())
}
private static func renderMetadataValue(_ value: Logger.Metadata.Value) -> String {
switch value {
case let .string(text):
text
default:
String(describing: value)
}
}
}
@@ -0,0 +1,72 @@
import Foundation
import Logging
#if canImport(Darwin)
import Darwin
#elseif canImport(Glibc)
import Glibc
#elseif canImport(Musl)
import Musl
#endif
struct JSONStderrLogHandler: LogHandler {
var metadata: Logger.Metadata = [:]
var logLevel: Logger.Level = .info
private let label: String
private let encoder: JSONEncoder
init(label: String) {
self.label = label
self.encoder = JSONEncoder()
self.encoder.dateEncodingStrategy = .iso8601
self.encoder.outputFormatting = [.sortedKeys]
}
subscript(metadataKey metadataKey: String) -> Logger.Metadata.Value? {
get { self.metadata[metadataKey] }
set { self.metadata[metadataKey] = newValue }
}
func log(event: LogEvent) {
let ts = Date()
var combined = self.metadata
if let metadata = event.metadata { combined.merge(metadata, uniquingKeysWith: { _, new in new }) }
let payload = JSONLogLine(
timestamp: ts,
level: event.level.rawValue,
label: self.label,
message: event.message.description,
source: event.source,
file: event.file,
function: event.function,
line: event.line,
metadata: combined.isEmpty ? nil : combined.mapValues(\.description))
guard let data = try? self.encoder.encode(payload),
let text = String(data: data, encoding: .utf8) else { return }
Self.writeStderr(text + "\n")
}
}
private struct JSONLogLine: Encodable {
let timestamp: Date
let level: String
let label: String
let message: String
let source: String
let file: String
let function: String
let line: UInt
let metadata: [String: String]?
}
extension JSONStderrLogHandler {
private static func writeStderr(_ text: String) {
let bytes = Array(text.utf8)
bytes.withUnsafeBytes { buffer in
guard let baseAddress = buffer.baseAddress else { return }
_ = write(STDERR_FILENO, baseAddress, buffer.count)
}
}
}
@@ -0,0 +1,96 @@
public enum LogCategories {
public static let abacusCookie = "abacus-cookie"
public static let abacusUsage = "abacus-usage"
public static let adaptiveRefresh = "adaptive-refresh"
public static let amp = "amp"
public static let antigravity = "antigravity"
public static let app = "app"
public static let auggieCLI = "auggie-cli"
public static let augment = "augment"
public static let augmentKeepalive = "augment-keepalive"
public static let bedrockUsage = "bedrock-usage"
public static let browserCookieGate = "browser-cookie-gate"
public static let claudeCLI = "claude-cli"
public static let claudeProbe = "claude-probe"
public static let claudeUsage = "claude-usage"
public static let codexRPC = "codex-rpc"
public static let commandcodeCookie = "commandcode-cookie"
public static let commandcodeUsage = "commandcode-usage"
public static let configMigration = "config-migration"
public static let configStore = "config-store"
public static let confetti = "confetti"
public static let cookieCache = "cookie-cache"
public static let cookieHeaderStore = "cookie-header-store"
public static let copilotTokenStore = "copilot-token-store"
public static let creditsPurchase = "creditsPurchase"
public static let crossModelUsage = "crossmodel-usage"
public static let cursorLogin = "cursor-login"
public static let deepSeekSettings = "deepseek-settings"
public static let deepSeekUsage = "deepseek-usage"
public static let deepgramUsage = "deepgram-usage"
public static let devin = "devin"
public static let doubaoUsage = "doubao-usage"
public static let elevenLabsUsage = "elevenlabs-usage"
public static let geminiProbe = "gemini-probe"
public static let grok = "grok"
public static let keychainCache = "keychain-cache"
public static let keychainMigration = "keychain-migration"
public static let keychainPreflight = "keychain-preflight"
public static let keychainPrompt = "keychain-prompt"
public static let kimiAPI = "kimi-api"
public static let kimiCookie = "kimi-cookie"
public static let kimiK2TokenStore = "kimi-k2-token-store"
public static let kimiK2Usage = "kimi-k2-usage"
public static let kimiTokenStore = "kimi-token-store"
public static let kimiWeb = "kimi-web"
public static let kiro = "kiro"
public static let launchAtLogin = "launch-at-login"
public static let login = "login"
public static let logging = "logging"
public static let manusAPI = "manus-api"
public static let manusCookie = "manus-cookie"
public static let manusWeb = "manus-web"
public static let memoryPressure = "memory-pressure"
public static let minimaxAPITokenStore = "minimax-api-token-store"
public static let minimaxCookie = "minimax-cookie"
public static let minimaxCookieStore = "minimax-cookie-store"
public static let minimaxUsage = "minimax-usage"
public static let minimaxWeb = "minimax-web"
public static let moonshotUsage = "moonshot-usage"
public static let notifications = "notifications"
public static let openAIWeb = "openai-web"
public static let openAIWebview = "openai-webview"
public static let ollama = "ollama"
public static let opencodeUsage = "opencode-usage"
public static let opencodeGoUsage = "opencode-go-usage"
public static let openRouterUsage = "openrouter-usage"
public static let perplexityAPI = "perplexity-api"
public static let perplexityCookie = "perplexity-cookie"
public static let perplexityWeb = "perplexity-web"
public static let poeUsage = "poe-usage"
public static let providerDetection = "provider-detection"
public static let providers = "providers"
public static let qoderCookie = "qoder-cookie"
public static let qoderUsage = "qoder-usage"
public static let quotaWarningNotifications = "quotaWarningNotifications"
public static let sessionQuota = "sessionQuota"
public static let sessionQuotaNotifications = "sessionQuotaNotifications"
public static let settings = "settings"
public static let subprocess = "subprocess"
public static let syntheticTokenStore = "synthetic-token-store"
public static let syntheticUsage = "synthetic-usage"
public static let t3chat = "t3chat"
public static let terminal = "terminal"
public static let tokenAccounts = "token-accounts"
public static let tokenCost = "token-cost"
public static let ttyRunner = "tty-runner"
public static let veniceUsage = "venice-usage"
public static let vertexAIFetcher = "vertexai-fetcher"
public static let warpUsage = "warp-usage"
public static let zed = "zed"
public static let webkitTeardown = "webkit-teardown"
public static let zaiSettings = "zai-settings"
public static let zaiTokenStore = "zai-token-store"
public static let zaiUsage = "zai-usage"
public static let stepfunUsage = "stepfun-usage"
}
@@ -0,0 +1,10 @@
import Foundation
public enum LogMetadata {
public static func secretSummary(_ value: String?) -> [String: String] {
let trimmed = value?.trimmingCharacters(in: .whitespacesAndNewlines) ?? ""
let state = trimmed.isEmpty ? "cleared" : "set"
let length = trimmed.count
return ["state": state, "length": "\(length)"]
}
}
@@ -0,0 +1,61 @@
import Foundation
public enum LogRedactor {
private static let fallbackRegex: NSRegularExpression = {
do {
return try NSRegularExpression(pattern: "$^", options: [])
} catch {
fatalError("Failed to build fallback regex: \(error)")
}
}()
private static let emailRegex = Self.makeRegex(
pattern: #"[A-Z0-9._%+-]+@[A-Z0-9.-]+\.[A-Z]{2,}"#,
options: [.caseInsensitive])
private static let cookieHeaderRegex = Self.makeRegex(
pattern: #"(?i)(cookie\s*:\s*)([^\r\n]+)"#)
private static let authorizationRegex = Self.makeRegex(
pattern: #"(?i)(authorization\s*:\s*)([^\r\n]+)"#)
private static let bearerRegex = Self.makeRegex(
pattern: #"(?i)\bbearer\s+[a-z0-9._\-]+=*\b"#)
private static let minimaxCodingPlanTokenRegex = Self.makeRegex(
pattern: #"sk-cp-[^\s"'`;,)>\]]+"#)
private static let minimaxApiTokenRegex = Self.makeRegex(
pattern: #"sk-api-[^\s"'`;,)>\]]+"#)
public static func redact(_ text: String) -> String {
guard self.mayContainSensitiveValue(text) else { return text }
var output = text
// Email is broad and safe first
output = self.replace(self.emailRegex, in: output, with: "<redacted-email>")
// MiniMax tokens before broader rules catch them
output = self.replace(self.minimaxCodingPlanTokenRegex, in: output, with: "<redacted-minimax-token>")
output = self.replace(self.minimaxApiTokenRegex, in: output, with: "<redacted-minimax-token>")
// Bearer catches "bearer <token>" before authorization wraps it
output = self.replace(self.bearerRegex, in: output, with: "Bearer <redacted>")
// Authorization catches the rest (already-redacted content)
output = self.replace(self.cookieHeaderRegex, in: output, with: "$1<redacted>")
output = self.replace(self.authorizationRegex, in: output, with: "$1<redacted>")
return output
}
private static func mayContainSensitiveValue(_ text: String) -> Bool {
if text.range(of: "@") != nil { return true }
if text.range(of: "sk-cp-", options: [.caseInsensitive]) != nil { return true }
if text.range(of: "sk-api-", options: [.caseInsensitive]) != nil { return true }
if text.range(of: "bearer", options: [.caseInsensitive]) != nil { return true }
if text.range(of: "cookie", options: [.caseInsensitive]) != nil { return true }
if text.range(of: "authorization", options: [.caseInsensitive]) != nil { return true }
return false
}
private static func makeRegex(pattern: String, options: NSRegularExpression.Options = []) -> NSRegularExpression {
(try? NSRegularExpression(pattern: pattern, options: options)) ?? self.fallbackRegex
}
private static func replace(_ regex: NSRegularExpression, in text: String, with template: String) -> String {
let range = NSRange(text.startIndex..<text.endIndex, in: text)
return regex.stringByReplacingMatches(in: text, options: [], range: range, withTemplate: template)
}
}
@@ -0,0 +1,77 @@
#if canImport(os)
import Foundation
import Logging
import os
struct OSLogLogHandler: LogHandler {
var metadata: Logging.Logger.Metadata = [:]
var logLevel: Logging.Logger.Level = .info
private let label: String
private let subsystem: String
private let logger: os.Logger
init(label: String, subsystem: String) {
self.label = label
self.subsystem = subsystem
let category = Self.category(from: label, subsystem: subsystem)
self.logger = os.Logger(subsystem: subsystem, category: category)
}
subscript(metadataKey metadataKey: String) -> Logging.Logger.Metadata.Value? {
get { self.metadata[metadataKey] }
set { self.metadata[metadataKey] = newValue }
}
func log(event: LogEvent) {
let msg = Self.decorate(
message: event.message.description,
label: self.label,
subsystem: self.subsystem,
metadata: self.metadata,
extraMetadata: event.metadata)
switch event.level {
case .trace:
self.logger.debug("\(msg, privacy: .public)")
case .debug:
self.logger.debug("\(msg, privacy: .public)")
case .info, .notice:
self.logger.info("\(msg, privacy: .public)")
case .warning:
self.logger.warning("\(msg, privacy: .public)")
case .error:
self.logger.error("\(msg, privacy: .public)")
case .critical:
self.logger.fault("\(msg, privacy: .public)")
}
}
private static func category(from label: String, subsystem: String) -> String {
let prefix = subsystem + "."
guard label.hasPrefix(prefix) else { return label }
return String(label.dropFirst(prefix.count))
}
private static func decorate(
message: String,
label: String,
subsystem: String,
metadata: Logging.Logger.Metadata,
extraMetadata: Logging.Logger.Metadata?)
-> String
{
var merged = metadata
if let extraMetadata { merged.merge(extraMetadata, uniquingKeysWith: { _, new in new }) }
guard !merged.isEmpty else { return message }
let suffix = merged
.sorted(by: { $0.key < $1.key })
.map { "\($0.key)=\($0.value)" }
.joined(separator: " ")
_ = label
_ = subsystem
return "\(message) (\(suffix))"
}
}
#endif
@@ -0,0 +1,27 @@
public enum ProviderLogging {
public static func logStartupState(
logger: CodexBarLogger,
providers: [UsageProvider],
isEnabled: (UsageProvider) -> Bool,
modeSnapshot: [String: String])
{
let ordered = providers.sorted { $0.rawValue < $1.rawValue }
let states = ordered
.map { provider -> String in
let enabled = isEnabled(provider)
return "\(provider.rawValue)=\(enabled ? "1" : "0")"
}
.joined(separator: ",")
let enabledProviders = ordered
.filter { isEnabled($0) }
.map(\.rawValue)
.joined(separator: ",")
logger.info(
"Provider enablement at startup",
metadata: [
"states": states,
"enabled": enabledProviders.isEmpty ? "none" : enabledProviders,
])
logger.info("Provider mode snapshot", metadata: modeSnapshot)
}
}
@@ -0,0 +1,115 @@
import Foundation
public enum FileManagedCodexAccountStoreError: Error, Equatable, Sendable {
case unsupportedVersion(Int)
}
public protocol ManagedCodexAccountStoring: Sendable {
func loadAccounts() throws -> ManagedCodexAccountSet
func storeAccounts(_ accounts: ManagedCodexAccountSet) throws
func ensureFileExists() throws -> URL
}
public struct FileManagedCodexAccountStore: ManagedCodexAccountStoring, @unchecked Sendable {
public static let currentVersion = 3
private let fileURL: URL
private let fileManager: FileManager
public init(fileURL: URL = Self.defaultURL(), fileManager: FileManager = .default) {
self.fileURL = fileURL
self.fileManager = fileManager
}
public func loadAccounts() throws -> ManagedCodexAccountSet {
guard self.fileManager.fileExists(atPath: self.fileURL.path) else {
return Self.emptyAccountSet()
}
let data = try Data(contentsOf: self.fileURL)
let decoder = JSONDecoder()
let accounts = try decoder.decode(ManagedCodexAccountSet.self, from: data)
guard (1...Self.currentVersion).contains(accounts.version) else {
throw FileManagedCodexAccountStoreError.unsupportedVersion(accounts.version)
}
if accounts.version == Self.currentVersion {
return ManagedCodexAccountSet(version: Self.currentVersion, accounts: accounts.accounts)
}
return self.migrateLegacyAccounts(accounts)
}
public func storeAccounts(_ accounts: ManagedCodexAccountSet) throws {
let normalizedAccounts = ManagedCodexAccountSet(
version: Self.currentVersion,
accounts: accounts.accounts)
let encoder = JSONEncoder()
encoder.outputFormatting = [.prettyPrinted, .sortedKeys]
let data = try encoder.encode(normalizedAccounts)
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 ensureFileExists() throws -> URL {
if self.fileManager.fileExists(atPath: self.fileURL.path) { return self.fileURL }
try self.storeAccounts(Self.emptyAccountSet())
return self.fileURL
}
private func applySecurePermissionsIfNeeded() throws {
#if os(macOS)
try self.fileManager.setAttributes([
.posixPermissions: NSNumber(value: Int16(0o600)),
], ofItemAtPath: self.fileURL.path)
#endif
}
private static func emptyAccountSet() -> ManagedCodexAccountSet {
ManagedCodexAccountSet(version: self.currentVersion, accounts: [])
}
private func migrateLegacyAccounts(_ accounts: ManagedCodexAccountSet) -> ManagedCodexAccountSet {
let migratedAccounts = accounts.accounts.map { account in
let hydratedProviderAccountID = account.providerAccountID ?? self.hydrateProviderAccountID(for: account)
return ManagedCodexAccount(
id: account.id,
email: account.email,
providerAccountID: hydratedProviderAccountID,
workspaceLabel: account.workspaceLabel,
workspaceAccountID: account.workspaceAccountID,
authFingerprint: account.authFingerprint ?? CodexAuthFingerprint.fingerprint(
homePath: account.managedHomePath,
fileManager: self.fileManager),
managedHomePath: account.managedHomePath,
createdAt: account.createdAt,
updatedAt: account.updatedAt,
lastAuthenticatedAt: account.lastAuthenticatedAt)
}
return ManagedCodexAccountSet(version: Self.currentVersion, accounts: migratedAccounts)
}
private func hydrateProviderAccountID(for account: ManagedCodexAccount) -> String? {
guard let credentials = try? CodexOAuthCredentialsStore.load(
env: ["CODEX_HOME": account.managedHomePath])
else {
return nil
}
let payload = credentials.idToken.flatMap(UsageFetcher.parseJWT)
let authDict = payload?["https://api.openai.com/auth"] as? [String: Any]
let providerAccountID = credentials.accountId
?? (authDict?["chatgpt_account_id"] as? String)
?? (payload?["chatgpt_account_id"] as? String)
return ManagedCodexAccount.normalizeProviderAccountID(providerAccountID)
}
public static func defaultURL() -> URL {
let base = FileManager.default.urls(for: .applicationSupportDirectory, in: .userDomainMask).first
?? FileManager.default.homeDirectoryForCurrentUser
return base
.appendingPathComponent("CodexBar", isDirectory: true)
.appendingPathComponent("managed-codex-accounts.json")
}
}
@@ -0,0 +1,367 @@
import Foundation
public struct OpenAIDashboardSnapshot: Codable, Equatable, Sendable {
public let signedInEmail: String?
public let codeReviewRemainingPercent: Double?
public let codeReviewLimit: RateWindow?
public let creditEvents: [CreditEvent]
public let dailyBreakdown: [OpenAIDashboardDailyBreakdown]
/// Usage breakdown time series from the Codex dashboard chart ("Usage breakdown", 30 days).
///
/// This is distinct from `dailyBreakdown`, which is derived from `creditEvents` (credits usage history table).
public let usageBreakdown: [OpenAIDashboardDailyBreakdown]
public let creditsPurchaseURL: String?
public let primaryLimit: RateWindow?
public let secondaryLimit: RateWindow?
/// Named model-specific limits (e.g. Codex Spark) decoded from the dashboard
/// `wham/usage` response's `additional_rate_limits` array.
public let extraRateWindows: [NamedRateWindow]?
public let creditsRemaining: Double?
public let codexCreditLimit: CodexCreditLimitSnapshot?
public let accountPlan: String?
public let updatedAt: Date
public init(
signedInEmail: String?,
codeReviewRemainingPercent: Double?,
codeReviewLimit: RateWindow? = nil,
creditEvents: [CreditEvent],
dailyBreakdown: [OpenAIDashboardDailyBreakdown],
usageBreakdown: [OpenAIDashboardDailyBreakdown],
creditsPurchaseURL: String?,
primaryLimit: RateWindow? = nil,
secondaryLimit: RateWindow? = nil,
extraRateWindows: [NamedRateWindow]? = nil,
creditsRemaining: Double? = nil,
codexCreditLimit: CodexCreditLimitSnapshot? = nil,
accountPlan: String? = nil,
updatedAt: Date)
{
self.signedInEmail = signedInEmail
self.codeReviewRemainingPercent = codeReviewRemainingPercent
self.codeReviewLimit = codeReviewLimit
self.creditEvents = creditEvents
self.dailyBreakdown = dailyBreakdown
self.usageBreakdown = OpenAIDashboardDailyBreakdown.removingSkillUsageServices(from: usageBreakdown)
self.creditsPurchaseURL = creditsPurchaseURL
self.primaryLimit = primaryLimit
self.secondaryLimit = secondaryLimit
self.extraRateWindows = extraRateWindows
self.creditsRemaining = creditsRemaining
self.codexCreditLimit = codexCreditLimit
self.accountPlan = accountPlan
self.updatedAt = updatedAt
}
private enum CodingKeys: String, CodingKey {
case signedInEmail
case codeReviewRemainingPercent
case codeReviewLimit
case creditEvents
case dailyBreakdown
case usageBreakdown
case creditsPurchaseURL
case primaryLimit
case secondaryLimit
case extraRateWindows
case creditsRemaining
case codexCreditLimit
case accountPlan
case updatedAt
}
public init(from decoder: Decoder) throws {
let container = try decoder.container(keyedBy: CodingKeys.self)
self.signedInEmail = try container.decodeIfPresent(String.self, forKey: .signedInEmail)
self.codeReviewRemainingPercent = try container.decodeIfPresent(
Double.self,
forKey: .codeReviewRemainingPercent)
self.codeReviewLimit = try container.decodeIfPresent(RateWindow.self, forKey: .codeReviewLimit)
self.creditEvents = try container.decodeIfPresent([CreditEvent].self, forKey: .creditEvents) ?? []
self.dailyBreakdown = try container.decodeIfPresent(
[OpenAIDashboardDailyBreakdown].self,
forKey: .dailyBreakdown)
?? Self.makeDailyBreakdown(from: self.creditEvents, maxDays: 30)
let decodedUsageBreakdown = try container.decodeIfPresent(
[OpenAIDashboardDailyBreakdown].self,
forKey: .usageBreakdown) ?? []
self.usageBreakdown = OpenAIDashboardDailyBreakdown.removingSkillUsageServices(
from: decodedUsageBreakdown)
self.creditsPurchaseURL = try container.decodeIfPresent(String.self, forKey: .creditsPurchaseURL)
self.primaryLimit = try container.decodeIfPresent(RateWindow.self, forKey: .primaryLimit)
self.secondaryLimit = try container.decodeIfPresent(RateWindow.self, forKey: .secondaryLimit)
// Backward-compatible: older cached snapshots simply lack the key and decode to nil.
self.extraRateWindows = try container.decodeIfPresent(
[NamedRateWindow].self,
forKey: .extraRateWindows)
self.creditsRemaining = try container.decodeIfPresent(Double.self, forKey: .creditsRemaining)
self.codexCreditLimit = try container.decodeIfPresent(CodexCreditLimitSnapshot.self, forKey: .codexCreditLimit)
self.accountPlan = try container.decodeIfPresent(String.self, forKey: .accountPlan)
self.updatedAt = try container.decode(Date.self, forKey: .updatedAt)
}
public static func makeDailyBreakdown(from events: [CreditEvent], maxDays: Int) -> [OpenAIDashboardDailyBreakdown] {
guard !events.isEmpty else { return [] }
let calendar = Calendar(identifier: .gregorian)
let formatter = DateFormatter()
formatter.locale = Locale(identifier: "en_US_POSIX")
formatter.calendar = calendar
formatter.timeZone = TimeZone.current
formatter.dateFormat = "yyyy-MM-dd"
var totals: [String: [String: Double]] = [:] // day -> service -> credits
for event in events {
let day = formatter.string(from: event.date)
totals[day, default: [:]][event.service, default: 0] += event.creditsUsed
}
let sortedDays = totals.keys.sorted(by: >).prefix(maxDays)
return sortedDays.map { day in
let serviceTotals = totals[day] ?? [:]
let services = serviceTotals
.map { OpenAIDashboardServiceUsage(service: $0.key, creditsUsed: $0.value) }
.sorted { lhs, rhs in
if lhs.creditsUsed == rhs.creditsUsed { return lhs.service < rhs.service }
return lhs.creditsUsed > rhs.creditsUsed
}
let total = services.reduce(0) { $0 + $1.creditsUsed }
return OpenAIDashboardDailyBreakdown(day: day, services: services, totalCreditsUsed: total)
}
}
}
extension OpenAIDashboardSnapshot {
public func toUsageSnapshot(
provider: UsageProvider = .codex,
accountEmail: String? = nil,
accountPlan: String? = nil) -> UsageSnapshot?
{
CodexReconciledState.fromAttachedDashboard(
snapshot: self,
provider: provider,
accountEmail: accountEmail,
accountPlan: accountPlan)?
.toUsageSnapshot()
}
public func toCreditsSnapshot() -> CreditsSnapshot? {
guard self.creditsRemaining != nil || self.codexCreditLimit != nil else { return nil }
return CreditsSnapshot(
remaining: self.creditsRemaining ?? 0,
events: self.creditEvents,
updatedAt: self.updatedAt,
codexCreditLimit: self.codexCreditLimit)
}
}
public struct OpenAIDashboardDailyBreakdown: Codable, Equatable, Sendable {
/// Day key in `yyyy-MM-dd` (local time).
public let day: String
public let services: [OpenAIDashboardServiceUsage]
public let totalCreditsUsed: Double
public init(day: String, services: [OpenAIDashboardServiceUsage], totalCreditsUsed: Double) {
self.day = day
self.services = services
self.totalCreditsUsed = totalCreditsUsed
}
public static func isSkillUsageService(_ service: String) -> Bool {
service
.trimmingCharacters(in: .whitespacesAndNewlines)
.lowercased()
.hasPrefix("skillusage:")
}
public static func removingSkillUsageServices(
from breakdown: [OpenAIDashboardDailyBreakdown])
-> [OpenAIDashboardDailyBreakdown]
{
breakdown.compactMap { day in
guard !day.services.isEmpty else {
return day.totalCreditsUsed > 0 ? day : nil
}
let services = day.services.filter { !self.isSkillUsageService($0.service) }
guard !services.isEmpty else { return nil }
let total = services.reduce(0) { $0 + $1.creditsUsed }
return OpenAIDashboardDailyBreakdown(
day: day.day,
services: services,
totalCreditsUsed: total)
}
}
public static func recentUsageSummary(
from breakdown: [OpenAIDashboardDailyBreakdown],
historyDays: Int = 30,
now: Date = Date(),
calendar: Calendar = .current) -> OpenAIDashboardUsageBreakdownSummary
{
let days = max(1, min(historyDays, 365))
var dayCalendar = Calendar(identifier: .gregorian)
dayCalendar.timeZone = calendar.timeZone
let today = dayCalendar.startOfDay(for: now)
let start = dayCalendar.date(byAdding: .day, value: -(days - 1), to: today) ?? today
let startKey = Self.dayKey(from: start, calendar: dayCalendar)
let todayKey = Self.dayKey(from: today, calendar: dayCalendar)
let recent = self.removingSkillUsageServices(from: breakdown)
.compactMap { self.sanitized($0, startKey: startKey, todayKey: todayKey, calendar: dayCalendar) }
.sorted { $0.day < $1.day }
let todayCredits = recent.first(where: { $0.day == todayKey })?.totalCreditsUsed
?? (recent.isEmpty ? nil : 0)
let totalCredits = self.finiteSum(recent.map(\.totalCreditsUsed))
return OpenAIDashboardUsageBreakdownSummary(
historyDays: days,
todayCredits: todayCredits,
totalCredits: totalCredits,
daily: recent)
}
private static func sanitized(
_ day: OpenAIDashboardDailyBreakdown,
startKey: String,
todayKey: String,
calendar: Calendar) -> OpenAIDashboardDailyBreakdown?
{
guard self.date(fromDayKey: day.day, calendar: calendar) != nil,
day.day >= startKey,
day.day <= todayKey
else { return nil }
if day.services.isEmpty {
guard day.totalCreditsUsed.isFinite, day.totalCreditsUsed > 0 else { return nil }
return day
}
let services = day.services.filter { $0.creditsUsed.isFinite && $0.creditsUsed > 0 }
guard !services.isEmpty,
let total = self.finiteSum(services.map(\.creditsUsed)),
total > 0
else { return nil }
return OpenAIDashboardDailyBreakdown(day: day.day, services: services, totalCreditsUsed: total)
}
private static func finiteSum(_ values: [Double]) -> Double? {
var total = 0.0
for value in values {
let next = total + value
guard next.isFinite else { return nil }
total = next
}
return values.isEmpty ? nil : total
}
private static func date(fromDayKey key: String, calendar: Calendar) -> Date? {
let parts = key.split(separator: "-", omittingEmptySubsequences: false)
guard parts.count == 3,
let year = Int(parts[0]),
let month = Int(parts[1]),
let day = Int(parts[2])
else { return nil }
let components = DateComponents(
calendar: calendar,
timeZone: calendar.timeZone,
year: year,
month: month,
day: day,
hour: 12)
guard let date = calendar.date(from: components) else { return nil }
let resolved = calendar.dateComponents([.year, .month, .day], from: date)
guard resolved.year == year, resolved.month == month, resolved.day == day else { return nil }
return date
}
private static func dayKey(from date: Date, calendar: Calendar) -> String {
let components = calendar.dateComponents([.year, .month, .day], from: date)
return String(
format: "%04d-%02d-%02d",
components.year ?? 0,
components.month ?? 0,
components.day ?? 0)
}
}
public struct OpenAIDashboardUsageBreakdownSummary: Equatable, Sendable {
public let historyDays: Int
public let todayCredits: Double?
public let totalCredits: Double?
public let daily: [OpenAIDashboardDailyBreakdown]
public init(
historyDays: Int,
todayCredits: Double?,
totalCredits: Double?,
daily: [OpenAIDashboardDailyBreakdown])
{
self.historyDays = historyDays
self.todayCredits = todayCredits
self.totalCredits = totalCredits
self.daily = daily
}
}
public struct OpenAIDashboardServiceUsage: Codable, Equatable, Sendable {
public let service: String
public let creditsUsed: Double
public init(service: String, creditsUsed: Double) {
self.service = service
self.creditsUsed = creditsUsed
}
}
public struct OpenAIDashboardCache: Codable, Equatable, Sendable {
public let accountEmail: String
public let snapshot: OpenAIDashboardSnapshot
public init(accountEmail: String, snapshot: OpenAIDashboardSnapshot) {
self.accountEmail = accountEmail
self.snapshot = snapshot
}
}
public enum OpenAIDashboardCacheStore {
@TaskLocal static var cacheURLOverride: URL?
public static func load() -> OpenAIDashboardCache? {
guard let url = self.cacheURL else { return nil }
guard let data = try? Data(contentsOf: url) else { return nil }
let decoder = JSONDecoder()
decoder.dateDecodingStrategy = .iso8601
return try? decoder.decode(OpenAIDashboardCache.self, from: data)
}
public static func save(_ cache: OpenAIDashboardCache) {
guard let url = self.cacheURL else { return }
do {
try FileManager.default.createDirectory(
at: url.deletingLastPathComponent(),
withIntermediateDirectories: true)
let encoder = JSONEncoder()
encoder.dateEncodingStrategy = .iso8601
let data = try encoder.encode(cache)
try data.write(to: url, options: [.atomic])
} catch {
// Best-effort cache only; ignore errors.
}
}
public static func clear() {
guard let url = self.cacheURL else { return }
try? FileManager.default.removeItem(at: url)
}
private static var cacheURL: URL? {
if let cacheURLOverride {
return cacheURLOverride
}
guard let root = FileManager.default.urls(for: .applicationSupportDirectory, in: .userDomainMask).first else {
return nil
}
let dir = root.appendingPathComponent("com.steipete.codexbar", isDirectory: true)
return dir.appendingPathComponent("openai-dashboard.json")
}
}
@@ -0,0 +1,190 @@
#if os(macOS)
import Foundation
extension OpenAIDashboardBrowserCookieImporter {
private struct PendingCookieStoreMutation {
let token: UUID
let task: Task<Void, Never>
}
@MainActor private static var pendingCookieStoreMutations: [ObjectIdentifier: PendingCookieStoreMutation] = [:]
private nonisolated static let cookieCacheQueue = DispatchQueue(
label: "com.steipete.codexbar.openai-cookie-cache")
private nonisolated static let deadlineQueue = DispatchQueue(
label: "com.steipete.codexbar.openai-cookie-deadline",
qos: .userInitiated)
private final class CookieLoadCompletion: @unchecked Sendable {
private let lock = NSLock()
private var didFinish = false
func finish(_ action: () -> Void) {
let shouldFinish = self.lock.withLock {
guard !self.didFinish else { return false }
self.didFinish = true
return true
}
if shouldFinish { action() }
}
}
nonisolated static func remainingTimeout(
until deadline: Date?,
cappedAt localLimit: TimeInterval? = nil,
now: Date = Date()) throws -> TimeInterval
{
guard let deadline else {
return localLimit.map(OpenAIDashboardFetcher.sanitizedTimeout) ?? .greatestFiniteMagnitude
}
let remaining = deadline.timeIntervalSince(now)
guard remaining > 0 else { throw URLError(.timedOut) }
guard let localLimit else { return remaining }
return min(OpenAIDashboardFetcher.sanitizedTimeout(localLimit), remaining)
}
nonisolated static func runBoundedCookieLoad<T: Sendable>(
deadline: Date?,
timeoutObserver: (@Sendable () -> Void)? = nil,
operation: @escaping @Sendable () throws -> T) async throws -> T
{
// The detached/GCD loader does not inherit TaskLocal prompt policy or retry selection.
let contextualOperation = BrowserCookieAccessGate.operationPreservingAccessContext(operation)
guard let deadline else {
return try await Task.detached(priority: .userInitiated, operation: contextualOperation).value
}
let timeout = try self.remainingTimeout(until: deadline)
let completion = CookieLoadCompletion()
return try await withCheckedThrowingContinuation { continuation in
DispatchQueue.global(qos: .userInitiated).async {
let result = Result(catching: contextualOperation)
completion.finish { continuation.resume(with: result) }
}
self.deadlineQueue.asyncAfter(deadline: .now() + timeout) {
completion.finish {
timeoutObserver?()
continuation.resume(throwing: URLError(.timedOut))
}
}
}
}
nonisolated static func runBoundedCookieCacheOperation<T: Sendable>(
deadline: Date?,
operation: @escaping @Sendable () throws -> T) async throws -> T
{
guard let deadline else {
return try await withCheckedThrowingContinuation { continuation in
self.cookieCacheQueue.async {
continuation.resume(with: Result(catching: operation))
}
}
}
let timeout = try self.remainingTimeout(until: deadline)
let completion = CookieLoadCompletion()
return try await withCheckedThrowingContinuation { continuation in
self.cookieCacheQueue.async {
let result = Result(catching: operation)
completion.finish { continuation.resume(with: result) }
}
self.deadlineQueue.asyncAfter(deadline: .now() + timeout) {
completion.finish { continuation.resume(throwing: URLError(.timedOut)) }
}
}
}
static func runBoundedCallback(
deadline: Date?,
timeoutObserver: (@Sendable () -> Void)? = nil,
start: (@escaping @Sendable () -> Void) -> Void) async throws
{
let completion = CookieLoadCompletion()
guard let deadline else {
await withCheckedContinuation { continuation in
start {
completion.finish { continuation.resume() }
}
}
return
}
let timeout = try self.remainingTimeout(until: deadline)
try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<Void, Error>) in
start {
completion.finish { continuation.resume() }
}
self.deadlineQueue.asyncAfter(deadline: .now() + timeout) {
completion.finish {
timeoutObserver?()
continuation.resume(throwing: URLError(.timedOut))
}
}
}
}
static func runBoundedValueCallback<T: Sendable>(
deadline: Date?,
timeoutObserver: (@Sendable () -> Void)? = nil,
start: (@escaping @Sendable (T) -> Void) -> Void) async throws -> T
{
let completion = CookieLoadCompletion()
guard let deadline else {
return await withCheckedContinuation { continuation in
start { value in
completion.finish { continuation.resume(returning: value) }
}
}
}
let timeout = try self.remainingTimeout(until: deadline)
return try await withCheckedThrowingContinuation { continuation in
start { value in
completion.finish { continuation.resume(returning: value) }
}
self.deadlineQueue.asyncAfter(deadline: .now() + timeout) {
completion.finish {
timeoutObserver?()
continuation.resume(throwing: URLError(.timedOut))
}
}
}
}
static func runSerializedCallback(
key: ObjectIdentifier,
deadline: Date?,
start: @escaping (@escaping @Sendable () -> Void) -> Void) async throws
{
while let pending = self.pendingCookieStoreMutations[key] {
try await self.waitForMutation(pending.task, deadline: deadline)
if self.pendingCookieStoreMutations[key]?.token == pending.token {
self.pendingCookieStoreMutations[key] = nil
}
}
let token = UUID()
let task = Task { @MainActor in
await withCheckedContinuation { continuation in
start { continuation.resume() }
}
}
self.pendingCookieStoreMutations[key] = PendingCookieStoreMutation(token: token, task: task)
Task { @MainActor in
await task.value
if self.pendingCookieStoreMutations[key]?.token == token {
self.pendingCookieStoreMutations[key] = nil
}
}
try await self.waitForMutation(task, deadline: deadline)
}
private static func waitForMutation(_ task: Task<Void, Never>, deadline: Date?) async throws {
try await self.runBoundedCallback(deadline: deadline) { completion in
Task { @MainActor in
await task.value
completion()
}
}
}
}
#endif
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,30 @@
#if os(macOS)
import Foundation
extension OpenAIDashboardFetcher {
struct ReturnableDashboardDataInput {
let codeReview: Double?
let events: [CreditEvent]
let usageBreakdown: [OpenAIDashboardDailyBreakdown]
let hasUsageLimits: Bool
let creditsRemaining: Double?
let codexCreditLimit: CodexCreditLimitSnapshot?
}
nonisolated static func hasReturnableDashboardData(_ input: ReturnableDashboardDataInput) -> Bool {
input.codeReview != nil
|| !input.events.isEmpty
|| !input.usageBreakdown.isEmpty
|| input.hasUsageLimits
|| input.creditsRemaining != nil
|| input.codexCreditLimit != nil
}
nonisolated static func hasAnyDashboardSignal(
hasReturnableData: Bool,
creditsHeaderPresent: Bool) -> Bool
{
hasReturnableData || creditsHeaderPresent
}
}
#endif
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,101 @@
#if os(macOS)
import Foundation
import WebKit
// MARK: - Navigation helper (revived from the old credits scraper)
@MainActor
final class NavigationDelegate: NSObject, WKNavigationDelegate {
private let completion: (Result<Void, Error>) -> Void
private var hasCompleted: Bool = false
private var timeoutWorkItem: DispatchWorkItem?
private var postCommitWorkItem: DispatchWorkItem?
static var associationKey: UInt8 = 0
nonisolated static let postCommitSuccessDelay: TimeInterval = 0.75
init(completion: @escaping (Result<Void, Error>) -> Void) {
self.completion = completion
}
func armTimeout(seconds: TimeInterval) {
self.timeoutWorkItem?.cancel()
let delay = max(seconds, 0)
let workItem = DispatchWorkItem { [weak self] in
MainActor.assumeIsolated {
self?.completeOnce(.failure(URLError(.timedOut)))
}
}
self.timeoutWorkItem = workItem
DispatchQueue.main.asyncAfter(deadline: .now() + delay, execute: workItem)
}
func cancel() {
self.completeOnce(.failure(CancellationError()))
}
func webView(_ webView: WKWebView, didFinish navigation: WKNavigation!) {
self.completeOnce(.success(()))
}
func webView(_ webView: WKWebView, didCommit navigation: WKNavigation!) {
guard !self.hasCompleted else { return }
self.postCommitWorkItem?.cancel()
let workItem = DispatchWorkItem { [weak self] in
MainActor.assumeIsolated {
self?.completeOnce(.success(()))
}
}
self.postCommitWorkItem = workItem
DispatchQueue.main.asyncAfter(deadline: .now() + Self.postCommitSuccessDelay, execute: workItem)
}
func webView(_ webView: WKWebView, didFail navigation: WKNavigation!, withError error: Error) {
if Self.shouldIgnoreNavigationError(error) { return }
self.completeOnce(.failure(error))
}
func webView(_ webView: WKWebView, didFailProvisionalNavigation navigation: WKNavigation!, withError error: Error) {
if Self.shouldIgnoreNavigationError(error) { return }
self.completeOnce(.failure(error))
}
nonisolated static func shouldIgnoreNavigationError(_ error: Error) -> Bool {
let nsError = error as NSError
if nsError.domain == NSURLErrorDomain, nsError.code == NSURLErrorCancelled {
return true
}
if nsError.domain == "WebKitErrorDomain", nsError.code == 102 {
return true
}
return false
}
private func completeOnce(_ result: Result<Void, Error>) {
guard !self.hasCompleted else { return }
self.hasCompleted = true
self.timeoutWorkItem?.cancel()
self.timeoutWorkItem = nil
self.postCommitWorkItem?.cancel()
self.postCommitWorkItem = nil
self.completion(result)
}
}
extension WKWebView {
var codexNavigationDelegate: NavigationDelegate? {
get {
objc_getAssociatedObject(self, &NavigationDelegate.associationKey) as? NavigationDelegate
}
set {
objc_setAssociatedObject(
self,
&NavigationDelegate.associationKey,
newValue,
.OBJC_ASSOCIATION_RETAIN_NONATOMIC)
}
}
}
#endif
@@ -0,0 +1,551 @@
import Foundation
public enum OpenAIDashboardParser {
/// Extracts the signed-in email from the embedded `client-bootstrap` JSON payload, if present.
///
/// The Codex usage dashboard currently ships a JSON blob in:
/// `<script type="application/json" id="client-bootstrap"></script>`.
/// WebKit `document.body.innerText` often does not include the email, so we parse it from HTML.
public static func parseSignedInEmailFromClientBootstrap(html: String) -> String? {
guard let data = self.clientBootstrapJSONData(fromHTML: html) else { return nil }
guard let json = try? JSONSerialization.jsonObject(with: data, options: []) else { return nil }
// Fast path: common structure.
if let dict = json as? [String: Any] {
if let session = dict["session"] as? [String: Any],
let user = session["user"] as? [String: Any],
let email = user["email"] as? String,
email.contains("@")
{
return email.trimmingCharacters(in: .whitespacesAndNewlines)
}
if let user = dict["user"] as? [String: Any],
let email = user["email"] as? String,
email.contains("@")
{
return email.trimmingCharacters(in: .whitespacesAndNewlines)
}
}
// Fallback: BFS scan for an email key/value.
var queue: [Any] = [json]
var seen = 0
while !queue.isEmpty, seen < 4000 {
let cur = queue.removeFirst()
seen += 1
if let dict = cur as? [String: Any] {
for (k, v) in dict {
if k.lowercased() == "email", let email = v as? String, email.contains("@") {
return email.trimmingCharacters(in: .whitespacesAndNewlines)
}
queue.append(v)
}
} else if let arr = cur as? [Any] {
queue.append(contentsOf: arr)
}
}
return nil
}
/// Extracts the auth status from `client-bootstrap`, if present.
/// Expected values include `logged_in` and `logged_out`.
public static func parseAuthStatusFromClientBootstrap(html: String) -> String? {
guard let data = self.clientBootstrapJSONData(fromHTML: html) else { return nil }
guard let json = try? JSONSerialization.jsonObject(with: data, options: []) else { return nil }
guard let dict = json as? [String: Any] else { return nil }
if let authStatus = dict["authStatus"] as? String, !authStatus.isEmpty {
return authStatus.trimmingCharacters(in: .whitespacesAndNewlines)
}
return nil
}
public static func parseCodeReviewRemainingPercent(bodyText: String) -> Double? {
let cleaned = bodyText.replacingOccurrences(of: "\r", with: "\n")
for regex in self.codeReviewRegexes {
let range = NSRange(cleaned.startIndex..<cleaned.endIndex, in: cleaned)
guard let match = regex.firstMatch(in: cleaned, options: [], range: range),
match.numberOfRanges >= 2,
let r = Range(match.range(at: 1), in: cleaned)
else { continue }
if let val = Double(cleaned[r]) { return min(100, max(0, val)) }
}
return nil
}
public static func parseCreditsRemaining(bodyText: String) -> Double? {
let cleaned = bodyText.replacingOccurrences(of: "\r", with: "\n")
let patterns = [
#"credits\s*remaining[^0-9]*([0-9][0-9.,]*)"#,
#"remaining\s*credits[^0-9]*([0-9][0-9.,]*)"#,
#"credit\s*balance[^0-9]*([0-9][0-9.,]*)"#,
]
for pattern in patterns {
if let val = TextParsing.firstNumber(pattern: pattern, text: cleaned) { return val }
}
return nil
}
public static func parseRateLimits(
bodyText: String,
now: Date = .init()) -> (primary: RateWindow?, secondary: RateWindow?)
{
let cleaned = bodyText.replacingOccurrences(of: "\r", with: "\n")
let lines = cleaned
.split(whereSeparator: \.isNewline)
.map { $0.trimmingCharacters(in: .whitespacesAndNewlines) }
.filter { !$0.isEmpty }
let primary = self.parseRateWindow(
lines: lines,
match: self.isFiveHourLimitLine,
windowMinutes: 5 * 60,
now: now)
let secondary = self.parseRateWindow(
lines: lines,
match: self.isWeeklyLimitLine,
windowMinutes: 7 * 24 * 60,
now: now)
return (primary, secondary)
}
public static func parseCodeReviewLimit(bodyText: String, now: Date = .init()) -> RateWindow? {
let cleaned = bodyText.replacingOccurrences(of: "\r", with: "\n")
let lines = cleaned
.split(whereSeparator: \.isNewline)
.map { $0.trimmingCharacters(in: .whitespacesAndNewlines) }
.filter { !$0.isEmpty }
return self.parseRateWindow(
lines: lines,
match: self.isCodeReviewLimitLine,
windowMinutes: nil,
now: now)
}
public static func parsePlanFromHTML(html: String) -> String? {
if let data = self.clientBootstrapJSONData(fromHTML: html),
let plan = self.findPlan(in: data)
{
return plan
}
if let data = self.nextDataJSONData(fromHTML: html),
let plan = self.findPlan(in: data)
{
return plan
}
return nil
}
public static func parseCreditEvents(rows: [[String]]) -> [CreditEvent] {
let formatter = self.creditDateFormatter()
return rows.compactMap { row in
guard row.count >= 3 else { return nil }
let dateString = row[0]
let service = row[1].trimmingCharacters(in: .whitespacesAndNewlines)
let amountString = row[2]
guard let date = formatter.date(from: dateString) else { return nil }
let creditsUsed = Self.parseCreditsUsed(amountString)
return CreditEvent(date: date, service: service, creditsUsed: creditsUsed)
}
.sorted { $0.date > $1.date }
}
private static func parseCreditsUsed(_ text: String) -> Double {
guard let raw = self.firstNumberToken(in: text) else { return 0 }
let token = raw
.replacingOccurrences(of: "\u{00A0}", with: "")
.replacingOccurrences(of: "\u{202F}", with: "")
.replacingOccurrences(of: " ", with: "")
let hasComma = token.contains(",")
let hasDot = token.contains(".")
if hasComma, hasDot {
return TextParsing.firstNumber(pattern: #"([0-9][0-9.,\s\p{Zs}]*)"#, text: token) ?? 0
}
if hasComma {
if self.usesLocalizedDecimalCommaCreditLabel(text) {
return Double(token.replacingOccurrences(of: ",", with: ".")) ?? 0
}
if token.range(of: #"^\d{1,3}(,\d{3})+$"#, options: .regularExpression) != nil {
return Double(token.replacingOccurrences(of: ",", with: "")) ?? 0
}
return Double(token.replacingOccurrences(of: ",", with: ".")) ?? 0
}
return Double(token) ?? 0
}
private static func usesLocalizedDecimalCommaCreditLabel(_ text: String) -> Bool {
text
.lowercased()
.contains("crédit")
}
private static func firstNumberToken(in text: String) -> String? {
guard let regex = try? NSRegularExpression(
pattern: #"([0-9][0-9.,\s\p{Zs}]*)"#,
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 >= 2,
let tokenRange = Range(match.range(at: 1), in: text)
else {
return nil
}
return String(text[tokenRange])
}
// MARK: - Private
private static let codeReviewRegexes: [NSRegularExpression] = {
let patterns = [
#"Code\s*review[^0-9%]*([0-9]{1,3})%\s*remaining"#,
#"Core\s*review[^0-9%]*([0-9]{1,3})%\s*remaining"#,
]
return patterns.compactMap { pattern in
try? NSRegularExpression(pattern: pattern, options: [.caseInsensitive])
}
}()
private static let creditDateFormatterKey = "OpenAIDashboardParser.creditDateFormatter"
private static let clientBootstrapNeedle = Data("id=\"client-bootstrap\"".utf8)
private static let nextDataNeedle = Data("id=\"__NEXT_DATA__\"".utf8)
private static let scriptCloseNeedle = Data("</script>".utf8)
private static func creditDateFormatter() -> DateFormatter {
let threadDict = Thread.current.threadDictionary
if let cached = threadDict[self.creditDateFormatterKey] as? DateFormatter {
return cached
}
let formatter = DateFormatter()
formatter.locale = Locale(identifier: "en_US_POSIX")
formatter.dateFormat = "MMM d, yyyy"
threadDict[self.creditDateFormatterKey] = formatter
return formatter
}
private static func clientBootstrapJSONData(fromHTML html: String) -> Data? {
let data = Data(html.utf8)
guard let idRange = data.range(of: self.clientBootstrapNeedle) else { return nil }
guard let openTagEnd = data[idRange.upperBound...].firstIndex(of: UInt8(ascii: ">")) else { return nil }
let contentStart = data.index(after: openTagEnd)
guard let closeRange = data.range(
of: self.scriptCloseNeedle,
options: [],
in: contentStart..<data.endIndex)
else {
return nil
}
let rawData = data[contentStart..<closeRange.lowerBound]
let trimmed = self.trimASCIIWhitespace(Data(rawData))
return trimmed.isEmpty ? nil : trimmed
}
private static func nextDataJSONData(fromHTML html: String) -> Data? {
let data = Data(html.utf8)
guard let idRange = data.range(of: self.nextDataNeedle) else { return nil }
guard let openTagEnd = data[idRange.upperBound...].firstIndex(of: UInt8(ascii: ">")) else { return nil }
let contentStart = data.index(after: openTagEnd)
guard let closeRange = data.range(
of: self.scriptCloseNeedle,
options: [],
in: contentStart..<data.endIndex)
else {
return nil
}
let rawData = data[contentStart..<closeRange.lowerBound]
let trimmed = self.trimASCIIWhitespace(Data(rawData))
return trimmed.isEmpty ? nil : trimmed
}
private static func trimASCIIWhitespace(_ data: Data) -> Data {
guard !data.isEmpty else { return data }
var start = data.startIndex
var end = data.endIndex
while start < end, data[start].isASCIIWhitespace {
start = data.index(after: start)
}
while end > start {
let prev = data.index(before: end)
if data[prev].isASCIIWhitespace {
end = prev
} else {
break
}
}
return data.subdata(in: start..<end)
}
private static func parseRateWindow(
lines: [String],
match: (String) -> Bool,
windowMinutes: Int?,
now: Date) -> RateWindow?
{
for idx in lines.indices where match(lines[idx]) {
let end = min(lines.count - 1, idx + 5)
let windowLines = Array(lines[idx...end])
var percentValue: Double?
var isRemaining = true
for line in windowLines {
if let percent = self.parsePercent(from: line) {
percentValue = percent.value
isRemaining = percent.isRemaining
break
}
}
guard let percentValue else { continue }
let usedPercent = isRemaining ? max(0, min(100, 100 - percentValue)) : max(0, min(100, percentValue))
let resetLine = windowLines.first { $0.localizedCaseInsensitiveContains("reset") }
let resetDescription = resetLine?.trimmingCharacters(in: .whitespacesAndNewlines)
let resetsAt = resetLine.flatMap { self.parseResetDate(from: $0, now: now) }
let fallbackDescription = resetsAt.map { UsageFormatter.resetDescription(from: $0) }
return RateWindow(
usedPercent: usedPercent,
windowMinutes: windowMinutes,
resetsAt: resetsAt,
resetDescription: resetDescription ?? fallbackDescription)
}
return nil
}
private static func parsePercent(from line: String) -> (value: Double, isRemaining: Bool)? {
guard let percent = TextParsing.firstNumber(pattern: #"([0-9]{1,3})\s*%"#, text: line) else { return nil }
let lower = line.lowercased()
let isRemaining = lower.contains("remaining") || lower.contains("left")
let isUsed = lower.contains("used") || lower.contains("spent") || lower.contains("consumed")
if isUsed { return (percent, false) }
if isRemaining { return (percent, true) }
return (percent, true)
}
private static func isFiveHourLimitLine(_ line: String) -> Bool {
let lower = line.lowercased()
if lower.contains("5h") { return true }
if lower.range(of: #"\b5\s*h\b"#, options: .regularExpression) != nil { return true }
if lower.contains("5-hour") { return true }
if lower.contains("5 hour") { return true }
return false
}
private static func isWeeklyLimitLine(_ line: String) -> Bool {
let lower = line.lowercased()
if lower.contains("weekly") { return true }
if lower.contains("7-day") { return true }
if lower.contains("7 day") { return true }
if lower.contains("7d") { return true }
if lower.range(of: #"\b7\s*d\b"#, options: .regularExpression) != nil { return true }
return false
}
private static func isCodeReviewLimitLine(_ line: String) -> Bool {
let lower = line.lowercased()
guard lower.contains("code review") || lower.contains("core review") else { return false }
if lower.contains("github code review") { return false }
return true
}
private static func parseResetDate(from line: String, now: Date) -> Date? {
var raw = line.trimmingCharacters(in: .whitespacesAndNewlines)
raw = raw.replacingOccurrences(of: #"(?i)^resets?:?\s*"#, with: "", options: .regularExpression)
raw = raw.replacingOccurrences(of: " at ", with: " ", options: .caseInsensitive)
raw = raw.replacingOccurrences(of: " on ", with: " ", options: .caseInsensitive)
raw = raw.replacingOccurrences(of: #"\s+"#, with: " ", options: .regularExpression)
.trimmingCharacters(in: .whitespacesAndNewlines)
let calendar = Calendar(identifier: .gregorian)
let monthDayFormatter = DateFormatter()
monthDayFormatter.locale = Locale(identifier: "en_US_POSIX")
monthDayFormatter.timeZone = TimeZone.current
monthDayFormatter.dateFormat = "MMM d"
var candidate = raw
let lower = candidate.lowercased()
var usedRelativeDay = false
if lower.contains("today") {
usedRelativeDay = true
let dateText = monthDayFormatter.string(from: now)
candidate = candidate.replacingOccurrences(of: "today", with: dateText, options: .caseInsensitive)
} else if lower.contains("tomorrow") {
usedRelativeDay = true
if let tomorrow = calendar.date(byAdding: .day, value: 1, to: now) {
let dateText = monthDayFormatter.string(from: tomorrow)
candidate = candidate.replacingOccurrences(of: "tomorrow", with: dateText, options: .caseInsensitive)
}
}
if let weekdayMatch = self.weekdayMatch(in: candidate) {
usedRelativeDay = true
let target = self.nextWeekdayDate(weekday: weekdayMatch.weekday, now: now, calendar: calendar)
let dateText = monthDayFormatter.string(from: target)
candidate = candidate.replacingOccurrences(
of: weekdayMatch.matched,
with: dateText,
options: .caseInsensitive)
}
let formatter = DateFormatter()
formatter.locale = Locale(identifier: "en_US_POSIX")
formatter.timeZone = TimeZone.current
formatter.defaultDate = now
let formats = [
"MMM d h:mma",
"MMM d, h:mma",
"MMM d h:mm a",
"MMM d, h:mm a",
"MMM d HH:mm",
"MMM d, HH:mm",
"MMM d",
"M/d h:mma",
"M/d h:mm a",
"M/d/yyyy h:mm a",
"M/d/yy h:mm a",
"M/d",
"yyyy-MM-dd HH:mm",
"yyyy-MM-dd h:mm a",
"yyyy-MM-dd",
]
for format in formats {
formatter.dateFormat = format
if let date = formatter.date(from: candidate) {
if usedRelativeDay, date < now {
if lower.contains("today"),
let bumped = calendar.date(byAdding: .day, value: 1, to: date)
{
return bumped
}
if let bumped = calendar.date(byAdding: .day, value: 7, to: date) {
return bumped
}
}
return date
}
}
return nil
}
private struct WeekdayMatch {
let matched: String
let weekday: Int
}
private static func weekdayMatch(in text: String) -> WeekdayMatch? {
// The optional "(day)?" suffix requires each alternative to be long enough that
// adding "day" reproduces the full weekday name. "wed"+"day"="wedday" and
// "sat"+"day"="satday", so the longer prefixes "wednes" and "satur" are needed
// to cover the full "Wednesday" and "Saturday" forms mirroring how the existing
// "tue|tues" and "thu|thur|thurs" alternatives layer up to "tuesday"/"thursday".
let pattern = #"\b(mon|tue|tues|wed|wednes|thu|thur|thurs|fri|sat|satur|sun)(day)?\b"#
guard let regex = try? NSRegularExpression(pattern: pattern, options: [.caseInsensitive]) else { return nil }
let range = NSRange(text.startIndex..<text.endIndex, in: text)
guard let match = regex.firstMatch(in: text, options: [], range: range),
let r = Range(match.range(at: 0), in: text)
else { return nil }
let matched = String(text[r])
let lower = matched.lowercased()
let weekday = switch lower.prefix(3) {
case "mon": 2
case "tue": 3
case "wed": 4
case "thu": 5
case "fri": 6
case "sat": 7
default: 1
}
return WeekdayMatch(matched: matched, weekday: weekday)
}
private static func nextWeekdayDate(weekday: Int, now: Date, calendar: Calendar) -> Date {
let currentWeekday = calendar.component(.weekday, from: now)
var delta = weekday - currentWeekday
if delta < 0 { delta += 7 }
guard let next = calendar.date(byAdding: .day, value: delta, to: calendar.startOfDay(for: now)) else {
return now
}
return next
}
private static func findPlan(in data: Data) -> String? {
guard let json = try? JSONSerialization.jsonObject(with: data, options: []) else { return nil }
return self.findPlan(in: json)
}
private static func findPlan(in json: Any) -> String? {
var queue: [Any] = [json]
var seen = 0
while !queue.isEmpty, seen < 6000 {
let cur = queue.removeFirst()
seen += 1
if let dict = cur as? [String: Any] {
for (k, v) in dict {
if let plan = self.planCandidate(forKey: k, value: v) { return plan }
queue.append(v)
}
} else if let arr = cur as? [Any] {
queue.append(contentsOf: arr)
}
}
return nil
}
private static func planCandidate(forKey key: String, value: Any) -> String? {
guard self.isPlanKey(key) else { return nil }
if let str = value as? String {
return self.normalizePlanValue(str)
}
if let dict = value as? [String: Any] {
if let name = dict["name"] as? String, let plan = self.normalizePlanValue(name) { return plan }
if let display = dict["displayName"] as? String, let plan = self.normalizePlanValue(display) { return plan }
if let tier = dict["tier"] as? String, let plan = self.normalizePlanValue(tier) { return plan }
}
return nil
}
private static func isPlanKey(_ key: String) -> Bool {
let lower = key.lowercased()
return lower.contains("plan") || lower.contains("tier") || lower.contains("subscription")
}
private static func normalizePlanValue(_ value: String) -> String? {
let trimmed = value.trimmingCharacters(in: .whitespacesAndNewlines)
guard !trimmed.isEmpty else { return nil }
let lower = trimmed.lowercased()
let allowed = [
"free",
"plus",
"pro",
"team",
"enterprise",
"business",
"edu",
"education",
"gov",
"premium",
"essential",
]
guard allowed.contains(where: { lower.contains($0) }) else { return nil }
return CodexPlanFormatting.displayName(trimmed) ?? UsageFormatter.cleanPlanName(trimmed)
}
}
extension UInt8 {
fileprivate var isASCIIWhitespace: Bool {
switch self {
case 9, 10, 13, 32: true
default: false
}
}
}
@@ -0,0 +1,907 @@
#if os(macOS)
let openAIDashboardScrapeScript = """
(() => {
const textOf = el => {
const raw = el && (el.innerText || el.textContent) ? String(el.innerText || el.textContent) : '';
return raw.trim();
};
const parseHexColor = (color) => {
if (!color) return null;
const c = String(color).trim().toLowerCase();
if (c.startsWith('#')) {
if (c.length === 4) {
return '#' + c[1] + c[1] + c[2] + c[2] + c[3] + c[3];
}
if (c.length === 7) return c;
return c;
}
const m = c.match(/^rgba?\\(([^)]+)\\)$/);
if (m) {
const parts = m[1].split(',').map(x => parseFloat(x.trim())).filter(x => Number.isFinite(x));
if (parts.length >= 3) {
const r = Math.max(0, Math.min(255, Math.round(parts[0])));
const g = Math.max(0, Math.min(255, Math.round(parts[1])));
const b = Math.max(0, Math.min(255, Math.round(parts[2])));
const toHex = n => n.toString(16).padStart(2, '0');
return '#' + toHex(r) + toHex(g) + toHex(b);
}
}
return c;
};
const reactPropsOf = (el) => {
if (!el) return null;
try {
const keys = Object.keys(el);
const propsKey = keys.find(k => k.startsWith('__reactProps$'));
if (propsKey) return el[propsKey] || null;
const fiberKey = keys.find(k => k.startsWith('__reactFiber$'));
if (fiberKey) {
const fiber = el[fiberKey];
return (fiber && (fiber.memoizedProps || fiber.pendingProps)) || null;
}
} catch {}
return null;
};
const reactFiberOf = (el) => {
if (!el) return null;
try {
const keys = Object.keys(el);
const fiberKey = keys.find(k => k.startsWith('__reactFiber$'));
return fiberKey ? (el[fiberKey] || null) : null;
} catch {
return null;
}
};
const nestedBarMetaOf = (root) => {
if (!root || typeof root !== 'object') return null;
const queue = [root];
const seen = typeof WeakSet !== 'undefined' ? new WeakSet() : null;
let steps = 0;
while (queue.length && steps < 250) {
const cur = queue.shift();
steps++;
if (!cur || typeof cur !== 'object') continue;
if (seen) {
if (seen.has(cur)) continue;
seen.add(cur);
}
if (cur.payload && (cur.dataKey || cur.name || cur.value !== undefined)) return cur;
const values = Array.isArray(cur) ? cur : Object.values(cur);
for (const v of values) {
if (v && typeof v === 'object') queue.push(v);
}
}
return null;
};
const barMetaFromElement = (el) => {
const direct = reactPropsOf(el);
if (direct && direct.payload && (direct.dataKey || direct.name || direct.value !== undefined)) return direct;
const fiber = reactFiberOf(el);
if (fiber) {
let cur = fiber;
for (let i = 0; i < 10 && cur; i++) {
const props = (cur.memoizedProps || cur.pendingProps) || null;
if (props && props.payload && (props.dataKey || props.name || props.value !== undefined)) return props;
const nested = props ? nestedBarMetaOf(props) : null;
if (nested) return nested;
cur = cur.return || null;
}
}
if (direct) {
const nested = nestedBarMetaOf(direct);
if (nested) return nested;
}
return null;
};
const normalizeHref = (raw) => {
if (!raw) return null;
const href = String(raw).trim();
if (!href) return null;
if (href.startsWith('http://') || href.startsWith('https://')) return href;
if (href.startsWith('//')) return window.location.protocol + href;
if (href.startsWith('/')) return window.location.origin + href;
return window.location.origin + '/' + href;
};
const isLikelyCreditsURL = (raw) => {
if (!raw) return false;
try {
const url = new URL(raw, window.location.origin);
if (!url.host || !url.host.includes('chatgpt.com')) return false;
const path = url.pathname.toLowerCase();
return (
path.includes('settings') ||
path.includes('usage') ||
path.includes('billing') ||
path.includes('credits')
);
} catch {
return false;
}
};
const purchaseTextMatches = (text) => {
const lower = String(text || '').trim().toLowerCase();
if (!lower) return false;
if (lower.includes('add more')) return true;
if (!lower.includes('credit')) return false;
return (
lower.includes('buy') ||
lower.includes('add') ||
lower.includes('purchase') ||
lower.includes('top up') ||
lower.includes('top-up')
);
};
const elementLabel = (el) => {
if (!el) return '';
return (
textOf(el) ||
el.getAttribute('aria-label') ||
el.getAttribute('title') ||
''
);
};
const urlFromProps = (props) => {
if (!props || typeof props !== 'object') return null;
const candidates = [
props.href,
props.to,
props.url,
props.link,
props.destination,
props.navigateTo
];
for (const candidate of candidates) {
if (typeof candidate === 'string' && candidate.trim()) {
return normalizeHref(candidate);
}
}
return null;
};
const purchaseURLFromElement = (el) => {
if (!el) return null;
const isAnchor = el.tagName && el.tagName.toLowerCase() === 'a';
const anchor = isAnchor ? el : (el.closest ? el.closest('a') : null);
const anchorHref = anchor ? anchor.getAttribute('href') : null;
const dataHref = el.getAttribute
? (el.getAttribute('data-href') ||
el.getAttribute('data-url') ||
el.getAttribute('data-link') ||
el.getAttribute('data-destination'))
: null;
const propHref = urlFromProps(reactPropsOf(el)) || urlFromProps(reactPropsOf(anchor));
const normalized = normalizeHref(anchorHref || dataHref || propHref);
return normalized && isLikelyCreditsURL(normalized) ? normalized : null;
};
const cleanPlanName = (raw) => String(raw || '')
.replace(/\\b(claude|codex|account|plan)\\b/gi, ' ')
.replace(/_/g, ' ')
.replace(/-/g, ' ')
.replace(/\\s+/g, ' ')
.trim();
const codexPlanDisplayName = (raw) => {
const trimmed = String(raw || '').trim();
if (!trimmed) return null;
const lower = trimmed.toLowerCase();
const exact = {
pro: 'Pro 20x',
prolite: 'Pro 5x',
'pro_lite': 'Pro 5x',
'pro-lite': 'Pro 5x',
'pro lite': 'Pro 5x'
};
if (exact[lower]) return exact[lower];
const cleaned = cleanPlanName(trimmed);
if (!cleaned) return trimmed;
if (exact[cleaned.toLowerCase()]) return exact[cleaned.toLowerCase()];
return cleaned.split(' ')
.filter(Boolean)
.map(word => {
const wordLower = word.toLowerCase();
if (wordLower === 'cbp' || wordLower === 'k12') return wordLower.toUpperCase();
if (word === word.toUpperCase() && /[a-z]/i.test(word)) return word;
return word.charAt(0).toUpperCase() + word.slice(1);
})
.join(' ') || cleaned;
};
const normalizePlanValue = (value) => {
const trimmed = String(value || '').trim();
if (!trimmed) return null;
const lower = trimmed.toLowerCase();
const allowed = [
'free',
'plus',
'pro',
'team',
'enterprise',
'business',
'edu',
'education',
'gov',
'premium',
'essential'
];
if (!allowed.some(token => lower.includes(token))) return null;
return codexPlanDisplayName(trimmed) || cleanPlanName(trimmed);
};
const planCandidate = (key, value) => {
const lower = String(key || '').toLowerCase();
if (!lower.includes('plan') && !lower.includes('tier') && !lower.includes('subscription')) return null;
if (typeof value === 'string') return normalizePlanValue(value);
if (value && typeof value === 'object' && !Array.isArray(value)) {
return normalizePlanValue(value.name) ||
normalizePlanValue(value.displayName) ||
normalizePlanValue(value.tier);
}
return null;
};
const findPlan = (root) => {
if (!root || typeof root !== 'object') return null;
const queue = [root];
const seenObjects = typeof WeakSet !== 'undefined' ? new WeakSet() : null;
let index = 0;
let seen = 0;
while (index < queue.length && seen < 6000) {
const cur = queue[index++];
seen++;
if (!cur || typeof cur !== 'object') continue;
if (seenObjects) {
if (seenObjects.has(cur)) continue;
seenObjects.add(cur);
}
if (Array.isArray(cur)) {
for (const v of cur) {
if (v && typeof v === 'object') queue.push(v);
}
continue;
}
for (const [k, v] of Object.entries(cur)) {
const plan = planCandidate(k, v);
if (plan) return plan;
if (v && typeof v === 'object') queue.push(v);
}
}
return null;
};
const parseJSONScript = (id) => {
try {
const node = document.getElementById(id);
const raw = node && node.textContent ? String(node.textContent) : '';
return raw ? JSON.parse(raw) : null;
} catch {
return null;
}
};
const pickLikelyPurchaseButton = (buttons) => {
if (!buttons || buttons.length === 0) return null;
const labeled = buttons.find(btn => {
const label = elementLabel(btn);
if (purchaseTextMatches(label)) return true;
const aria = String(btn.getAttribute('aria-label') || '').toLowerCase();
return aria.includes('credit') || aria.includes('buy') || aria.includes('add');
});
return labeled || buttons[0];
};
const findCreditsPurchaseButton = () => {
const nodes = Array.from(document.querySelectorAll('h1,h2,h3,div,span,p'));
const labelMatch = nodes.find(node => {
const lower = textOf(node).toLowerCase();
return lower === 'credits remaining' || (lower.includes('credits') && lower.includes('remaining'));
});
if (!labelMatch) return null;
let cur = labelMatch;
for (let i = 0; i < 6 && cur; i++) {
const buttons = Array.from(cur.querySelectorAll('button, a'));
const picked = pickLikelyPurchaseButton(buttons);
if (picked) return picked;
cur = cur.parentElement;
}
return null;
};
const dayKeyFromPayload = (payload) => {
if (!payload || typeof payload !== 'object') return null;
const localDayKeyForDate = (date) => {
const year = date.getFullYear();
const month = String(date.getMonth() + 1).padStart(2, '0');
const day = String(date.getDate()).padStart(2, '0');
return `${year}-${month}-${day}`;
};
const keys = ['day', 'date', 'name', 'label', 'x', 'time', 'timestamp'];
for (const k of keys) {
const v = payload[k];
if (typeof v === 'string') {
const s = v.trim();
if (/^\\d{4}-\\d{2}-\\d{2}$/.test(s)) return s;
const iso = s.match(/^(\\d{4}-\\d{2}-\\d{2})/);
if (iso) return iso[1];
}
if (typeof v === 'number' && Number.isFinite(v) && (k === 'timestamp' || k === 'time' || k === 'x')) {
try {
const d = new Date(v);
if (!isNaN(d.getTime())) return localDayKeyForDate(d);
} catch {}
}
}
return null;
};
const isSkillUsageServiceKey = (raw) => {
const key = raw === null || raw === undefined ? '' : String(raw).trim().toLowerCase();
return key.startsWith('skillusage:');
};
const displayNameForUsageServiceKey = (raw) => {
const key = raw === null || raw === undefined ? '' : String(raw).trim();
if (!key) return key;
if (isSkillUsageServiceKey(key)) return null;
if (key.toUpperCase() === key && key.length <= 6) return key;
const lower = key.toLowerCase();
if (lower === 'cli') return 'CLI';
if (lower.includes('github') && lower.includes('review')) return 'GitHub Code Review';
const words = lower.replace(/[_-]+/g, ' ').split(' ').filter(Boolean);
return words.map(w => w.length <= 2 ? w.toUpperCase() : w.charAt(0).toUpperCase() + w.slice(1)).join(' ');
};
const isLikelyCodexUsageService = (raw) => {
const service = raw === null || raw === undefined ? '' : String(raw).trim().toLowerCase();
return (
service === 'cli' ||
service === 'desktop' ||
service === 'desktop app' ||
service === 'vscode' ||
service === 'vs code' ||
service === 'unknown' ||
(service.includes('github') && service.includes('review'))
);
};
const usageChartRootForPath = (path) => {
if (!path || !path.closest) return null;
return (
path.closest('.recharts-wrapper') ||
path.closest('svg.recharts-surface') ||
path.closest('section') ||
path.parentElement ||
null
);
};
const uniqueUsageChartRoots = (paths) => {
const roots = [];
for (const path of paths) {
const root = usageChartRootForPath(path);
if (root && !roots.includes(root)) roots.push(root);
}
return roots;
};
const usageBreakdownTitleScore = (title) => {
const lower = String(title || '').trim().toLowerCase().replace(/\\s+/g, ' ');
if (!lower) return 0;
if (lower === 'usage breakdown') return 1000000;
if (lower.includes('usage breakdown')) return 900000;
if (lower === 'personal usage') return 800000;
if (lower.includes('threads') ||
lower.includes('turns') ||
lower.includes('client') ||
lower.includes('skill') ||
lower.includes('invocation')) return -1000000;
return 0;
};
const titleLikeElements = (scope) => {
try {
return Array.from(scope.querySelectorAll('h1,h2,h3,[role=\"heading\"],div,span,p'))
.filter(el => {
const title = textOf(el);
const lower = title.toLowerCase();
const tag = el.tagName ? el.tagName.toLowerCase() : '';
const isHeading = tag === 'h1' ||
tag === 'h2' ||
tag === 'h3' ||
String(el.getAttribute('role') || '').toLowerCase() === 'heading';
return title.length > 0 &&
title.length <= 80 &&
(
isHeading ||
usageBreakdownTitleScore(title) !== 0 ||
lower.includes('usage breakdown') ||
lower.includes('threads') ||
lower.includes('turns') ||
lower.includes('client') ||
lower.includes('skill') ||
lower.includes('invocation')
);
});
} catch {
return [];
}
};
const titleNodePrecedesRoot = (titleNode, root) => {
if (!titleNode || titleNode === root || root.contains(titleNode) || titleNode.contains(root)) return false;
const relation = titleNode.compareDocumentPosition(root);
return Boolean(relation & Node.DOCUMENT_POSITION_FOLLOWING);
};
const nearestScoredChartTitleInScope = (scope, root) => {
let best = null;
for (const titleNode of titleLikeElements(scope)) {
if (!titleNodePrecedesRoot(titleNode, root)) continue;
const title = textOf(titleNode);
const score = usageBreakdownTitleScore(title);
if (score === 0) continue;
if (!best || score >= best.score) best = { title, score };
}
return best ? best.title : '';
};
const chartTitleBoundaryForRoot = (root) => {
if (!root) return null;
try {
return root.closest('section,[role=\"region\"],article') || root.parentElement || null;
} catch {
return root.parentElement || null;
}
};
const nearestTitleTextInScope = (scope, root) => {
if (!scope) return '';
let nearest = null;
for (const titleNode of titleLikeElements(scope)) {
if (titleNodePrecedesRoot(titleNode, root)) nearest = titleNode;
}
return textOf(nearest);
};
const nearestChartTitleTextForRoot = (root) => {
if (!root) return '';
try {
const boundary = chartTitleBoundaryForRoot(root) || root.parentElement || null;
let ancestor = root.parentElement || null;
for (let i = 0; i < 8 && ancestor; i++) {
const scoredTitle = nearestScoredChartTitleInScope(ancestor, root);
if (scoredTitle) return scoredTitle;
if (ancestor === boundary) break;
ancestor = ancestor.parentElement || null;
}
return nearestTitleTextInScope(boundary, root);
} catch {
return '';
}
};
const legendMapForUsageChartRoot = (root) => {
const legendMap = {};
const scopes = [
root,
root && root.parentElement,
root && root.closest ? root.closest('section') : null
].filter(Boolean);
for (const scope of scopes) {
try {
const legendItems = Array.from(scope.querySelectorAll('div[title]'));
for (const item of legendItems) {
const title = item.getAttribute('title') ? String(item.getAttribute('title')).trim() : '';
const square = item.querySelector('div[style*=\"background-color\"]');
const color = (square && square.style && square.style.backgroundColor)
? square.style.backgroundColor
: null;
const hex = parseHexColor(color);
if (title && hex) legendMap[hex] = title;
}
} catch {}
if (Object.keys(legendMap).length > 0) break;
}
return legendMap;
};
const parseUsageBreakdownFromChartPaths = (paths, legendMap) => {
const totalsByDay = {}; // day -> service -> value
const addValue = (day, service, value) => {
if (!day || !service || isSkillUsageServiceKey(service)) return false;
if (typeof value !== 'number' || !Number.isFinite(value) || value <= 0) return false;
if (!totalsByDay[day]) totalsByDay[day] = {};
totalsByDay[day][service] = (totalsByDay[day][service] || 0) + value;
return true;
};
let pointCount = 0;
for (const path of paths) {
const meta = barMetaFromElement(path) || barMetaFromElement(path.parentElement) || null;
if (!meta) continue;
const payload = meta.payload || null;
const day = dayKeyFromPayload(payload);
if (!day) continue;
const valuesObj = (payload && payload.values && typeof payload.values === 'object') ? payload.values : null;
if (valuesObj) {
for (const [k, v] of Object.entries(valuesObj)) {
const service = displayNameForUsageServiceKey(k);
if (addValue(day, service, v)) pointCount++;
}
continue;
}
let value = null;
if (typeof meta.value === 'number' && Number.isFinite(meta.value)) value = meta.value;
if (value === null && typeof meta.value === 'string') {
const v = parseFloat(meta.value.replace(/,/g, ''));
if (Number.isFinite(v)) value = v;
}
if (value === null) continue;
const fill = parseHexColor(meta.fill || path.getAttribute('fill'));
const service =
(fill && legendMap[fill]) ||
(typeof meta.name === 'string' && meta.name) ||
null;
if (addValue(day, service, value)) pointCount++;
}
const dayKeys = Object.keys(totalsByDay)
.filter(day => Object.keys(totalsByDay[day] || {}).length > 0)
.sort((a, b) => b.localeCompare(a))
.slice(0, 30);
const breakdown = dayKeys.map(day => {
const servicesMap = totalsByDay[day] || {};
const services = Object.keys(servicesMap).map(service => ({
service,
creditsUsed: servicesMap[service]
})).sort((a, b) => {
if (a.creditsUsed === b.creditsUsed) return a.service.localeCompare(b.service);
return b.creditsUsed - a.creditsUsed;
});
const totalCreditsUsed = services.reduce((sum, s) => sum + (Number(s.creditsUsed) || 0), 0);
return { day, services, totalCreditsUsed };
});
const services = Array.from(new Set(breakdown.flatMap(day => day.services.map(service => service.service))));
const totalCreditsUsed = breakdown.reduce((sum, day) => sum + (Number(day.totalCreditsUsed) || 0), 0);
const likelyCodexServiceCount = services.filter(isLikelyCodexUsageService).length;
return {
breakdown,
pointCount,
services,
totalCreditsUsed,
likelyCodexServiceCount,
score: likelyCodexServiceCount * 1000 + services.length * 100 + pointCount + totalCreditsUsed / 1000
};
};
const usageBreakdownJSON = (() => {
try {
if (window.__codexbarUsageBreakdownJSON) return window.__codexbarUsageBreakdownJSON;
const paths = Array.from(document.querySelectorAll('g.recharts-bar-rectangle path.recharts-rectangle'));
let debug = {
pathCount: paths.length,
chartCount: 0,
eligibleCandidateCount: 0,
selectedCandidateTitle: null,
candidateSummaries: [],
sampleReactKeys: null,
sampleMetaKeys: null,
samplePayloadKeys: null,
sampleValuesKeys: null,
sampleDayKey: null
};
try {
const sample = paths[0] || null;
if (sample) {
const names = Object.getOwnPropertyNames(sample);
debug.sampleReactKeys = names.filter(k => k.includes('react')).slice(0, 10);
const metaSample = barMetaFromElement(sample) || barMetaFromElement(sample.parentElement) || null;
if (metaSample) {
debug.sampleMetaKeys = Object.keys(metaSample).slice(0, 12);
const payload = metaSample.payload || null;
if (payload && typeof payload === 'object') {
debug.samplePayloadKeys = Object.keys(payload).slice(0, 12);
debug.sampleDayKey = dayKeyFromPayload(payload);
const values = payload.values || null;
if (values && typeof values === 'object') {
debug.sampleValuesKeys = Object.keys(values).slice(0, 12);
}
}
}
}
} catch {}
const roots = uniqueUsageChartRoots(paths);
debug.chartCount = roots.length;
const candidates = roots.map(root => {
const chartPaths = paths.filter(path => usageChartRootForPath(path) === root);
const title = nearestChartTitleTextForRoot(root);
const titleScore = usageBreakdownTitleScore(title);
const parsed = parseUsageBreakdownFromChartPaths(chartPaths, legendMapForUsageChartRoot(root));
return {
root,
title,
titleScore,
pathCount: chartPaths.length,
...parsed,
score: titleScore + parsed.score
};
}).filter(candidate => candidate.breakdown.length > 0);
const rejectedTitleCandidates = candidates.filter(candidate => candidate.titleScore < 0);
const titledCandidates = candidates.filter(candidate => candidate.titleScore > 0);
const unknownTitleCandidates = candidates.filter(candidate => candidate.titleScore === 0);
const eligibleCandidates = titledCandidates;
eligibleCandidates.sort((a, b) => b.score - a.score);
debug.eligibleCandidateCount = eligibleCandidates.length;
debug.selectedCandidateTitle = eligibleCandidates[0] ? eligibleCandidates[0].title : null;
if (eligibleCandidates.length === 0 && candidates.length > 0) {
if (unknownTitleCandidates.length > 0) {
debug.error = 'No English usage breakdown chart title found. Candidate titles: ' +
candidates.map(candidate => candidate.title || 'Untitled chart').join(', ');
} else if (rejectedTitleCandidates.length > 0) {
debug.error = 'Only non-usage chart candidates found: ' +
rejectedTitleCandidates.map(candidate => candidate.title || 'Untitled chart').join(', ');
}
}
debug.candidateSummaries = candidates.slice(0, 6).map(candidate => ({
title: candidate.title,
titleScore: candidate.titleScore,
pathCount: candidate.pathCount,
dayCount: candidate.breakdown.length,
pointCount: candidate.pointCount,
serviceCount: candidate.services.length,
likelyCodexServiceCount: candidate.likelyCodexServiceCount,
services: candidate.services.slice(0, 8)
}));
const breakdown = eligibleCandidates[0] ? eligibleCandidates[0].breakdown : [];
const json = (breakdown.length > 0) ? JSON.stringify(breakdown) : null;
window.__codexbarUsageBreakdownJSON = json;
window.__codexbarUsageBreakdownDebug = json ? null : JSON.stringify(debug);
return json;
} catch {
return null;
}
})();
const usageBreakdownDebug = (() => {
try {
return window.__codexbarUsageBreakdownDebug || null;
} catch {
return null;
}
})();
const usageBreakdownError = (() => {
try {
if (!usageBreakdownDebug) return null;
const parsed = JSON.parse(usageBreakdownDebug);
return parsed && parsed.error ? String(parsed.error) : null;
} catch {
return null;
}
})();
const bodyText = document.body ? String(document.body.innerText || '').trim() : '';
const href = window.location ? String(window.location.href || '') : '';
const workspacePicker = bodyText.includes('Select a workspace');
const title = document.title ? String(document.title || '') : '';
const cloudflareInterstitial =
title.toLowerCase().includes('just a moment') ||
bodyText.toLowerCase().includes('checking your browser') ||
bodyText.toLowerCase().includes('cloudflare');
const authSelector = [
'input[type="email"]',
'input[type="password"]',
'input[name="username"]'
].join(', ');
const hasAuthInputs = !!document.querySelector(authSelector);
const lower = bodyText.toLowerCase();
const loginCTA =
lower.includes('sign in') ||
lower.includes('log in') ||
lower.includes('continue with google') ||
lower.includes('continue with apple') ||
lower.includes('continue with microsoft');
const loginRequired =
href.includes('/auth/') ||
href.includes('/login') ||
(hasAuthInputs && loginCTA) ||
(!hasAuthInputs && loginCTA && href.includes('chatgpt.com'));
const scrollY = (typeof window.scrollY === 'number') ? window.scrollY : 0;
const scrollHeight = document.documentElement ? (document.documentElement.scrollHeight || 0) : 0;
const viewportHeight = (typeof window.innerHeight === 'number') ? window.innerHeight : 0;
let creditsHeaderPresent = false;
let creditsHeaderInViewport = false;
let didScrollToCredits = false;
let rows = [];
try {
const looksLikeCreditsEventRow = (cells) => {
if (!cells || cells.length < 3) return false;
const first = String(cells[0] || '');
const amount = String(cells[2] || '');
return /\\d{4}|\\d{1,2}[\\/.\\-]\\d{1,2}/.test(first) && /\\d/.test(amount);
};
const allTableRows = () => Array.from(document.querySelectorAll('tbody tr')).map(tr => {
const cells = Array.from(tr.querySelectorAll('td')).map(td => textOf(td));
return cells;
}).filter(looksLikeCreditsEventRow);
const headings = Array.from(document.querySelectorAll('h1,h2,h3'));
const header = headings.find(h => textOf(h).toLowerCase() === 'credits usage history');
if (header) {
creditsHeaderPresent = true;
const rect = header.getBoundingClientRect();
creditsHeaderInViewport = rect.top >= 0 && rect.top <= viewportHeight;
// Only scrape rows from the *credits usage history* table. The page can contain other tables,
// and treating any <table> as credits history can prevent our scroll-to-load logic from running.
const container = header.closest('section') || header.parentElement || document;
const table = container.querySelector('table') || null;
const scope = table || container;
rows = Array.from(scope.querySelectorAll('tbody tr')).map(tr => {
const cells = Array.from(tr.querySelectorAll('td')).map(td => textOf(td));
return cells;
}).filter(r => r.length >= 3);
if (rows.length === 0) {
rows = allTableRows();
}
if (rows.length === 0 && !window.__codexbarDidScrollToCredits) {
window.__codexbarDidScrollToCredits = true;
// If the table is virtualized/lazy-loaded, we need to scroll to trigger rendering even if the
// header is already in view.
header.scrollIntoView({ block: 'start', inline: 'nearest' });
if (creditsHeaderInViewport) {
window.scrollBy(0, Math.max(220, viewportHeight * 0.6));
}
didScrollToCredits = true;
}
} else if (rows.length === 0 && !window.__codexbarDidScrollToCredits && scrollHeight > viewportHeight * 1.5) {
rows = allTableRows();
if (rows.length > 0) {
creditsHeaderPresent = true;
creditsHeaderInViewport = true;
}
}
if (rows.length === 0 && !window.__codexbarDidScrollToCredits && scrollHeight > viewportHeight * 1.5) {
// The credits history section often isn't part of the DOM until you scroll down. Nudge the page
// once so subsequent scrapes can find the header and rows.
window.__codexbarDidScrollToCredits = true;
window.scrollTo(0, Math.max(0, scrollHeight - viewportHeight - 40));
didScrollToCredits = true;
}
} catch {}
let creditsPurchaseURL = null;
try {
const creditsButton = findCreditsPurchaseButton();
if (creditsButton) {
const url = purchaseURLFromElement(creditsButton);
if (url) creditsPurchaseURL = url;
}
const candidates = Array.from(document.querySelectorAll('a, button'));
for (const node of candidates) {
const label = elementLabel(node);
if (!purchaseTextMatches(label)) continue;
const url = purchaseURLFromElement(node);
if (url) {
creditsPurchaseURL = url;
break;
}
}
if (!creditsPurchaseURL) {
const anchors = Array.from(document.querySelectorAll('a[href]'));
for (const anchor of anchors) {
const label = elementLabel(anchor);
const href = anchor.getAttribute('href') || '';
const hrefLooksRelevant = /credits|billing/i.test(href);
if (!hrefLooksRelevant && !purchaseTextMatches(label)) continue;
const url = normalizeHref(href);
if (url) {
creditsPurchaseURL = url;
break;
}
}
}
} catch {}
let signedInEmail = null;
let authStatus = null;
let accountPlan = null;
try {
const next = window.__NEXT_DATA__ || null;
const props = (next && next.props && next.props.pageProps) ? next.props.pageProps : null;
const userEmail = (props && props.user) ? props.user.email : null;
const sessionEmail = (props && props.session && props.session.user) ? props.session.user.email : null;
signedInEmail = userEmail || sessionEmail || null;
} catch {}
const clientBootstrap = parseJSONScript('client-bootstrap');
if (clientBootstrap) {
try {
authStatus = typeof clientBootstrap.authStatus === 'string' ? clientBootstrap.authStatus : null;
if (!signedInEmail) {
const session = clientBootstrap.session || null;
const user = (session && session.user) || clientBootstrap.user || null;
const email = user && typeof user.email === 'string' ? user.email : null;
if (email && email.includes('@')) signedInEmail = email;
}
if (!accountPlan) accountPlan = findPlan(clientBootstrap);
} catch {}
}
if (!accountPlan) {
try {
accountPlan = findPlan(window.__NEXT_DATA__ || parseJSONScript('__NEXT_DATA__'));
} catch {}
}
if (!signedInEmail) {
try {
const obj = parseJSONScript('__NEXT_DATA__');
if (obj) {
const queue = [obj];
let seen = 0;
while (queue.length && seen < 2000 && !signedInEmail) {
const cur = queue.shift();
seen++;
if (!cur) continue;
if (typeof cur === 'string') {
if (cur.includes('@')) signedInEmail = cur;
continue;
}
if (typeof cur !== 'object') continue;
for (const [k, v] of Object.entries(cur)) {
if (signedInEmail) break;
if (k === 'email' && typeof v === 'string' && v.includes('@')) {
signedInEmail = v;
break;
}
if (typeof v === 'object' && v) queue.push(v);
}
}
}
} catch {}
}
if (!signedInEmail) {
try {
const emailRe = /[A-Z0-9._%+-]+@[A-Z0-9.-]+\\.[A-Z]{2,}/ig;
const found = (bodyText.match(emailRe) || []).map(x => String(x).trim().toLowerCase());
const unique = Array.from(new Set(found));
if (unique.length === 1) {
signedInEmail = unique[0];
} else if (unique.length > 1) {
signedInEmail = unique[0];
}
} catch {}
}
if (!signedInEmail) {
// Last resort: open the account menu so the email becomes part of the DOM text.
try {
const emailRe = /[A-Z0-9._%+-]+@[A-Z0-9.-]+\\.[A-Z]{2,}/ig;
const hasMenu = Boolean(document.querySelector('[role="menu"]'));
if (!hasMenu) {
const button =
document.querySelector('button[aria-haspopup="menu"]') ||
document.querySelector('button[aria-expanded]');
if (button && !button.disabled) {
button.click();
}
}
const afterText = document.body ? String(document.body.innerText || '').trim() : '';
const found = (afterText.match(emailRe) || []).map(x => String(x).trim().toLowerCase());
const unique = Array.from(new Set(found));
if (unique.length === 1) {
signedInEmail = unique[0];
} else if (unique.length > 1) {
signedInEmail = unique[0];
}
} catch {}
}
return {
loginRequired,
workspacePicker,
cloudflareInterstitial,
href,
bodyText,
signedInEmail,
authStatus,
accountPlan,
creditsPurchaseURL,
rows,
usageBreakdownJSON,
usageBreakdownDebug,
usageBreakdownError,
scrollY,
scrollHeight,
viewportHeight,
creditsHeaderPresent,
creditsHeaderInViewport,
didScrollToCredits
};
})();
"""
#endif
@@ -0,0 +1,782 @@
#if os(macOS)
import AppKit
import Foundation
import WebKit
struct OpenAIDashboardWebViewLease {
let webView: WKWebView
let log: (String) -> Void
let setPreserveLoadedPageOnRelease: (Bool) -> Void
let release: () -> Void
}
@MainActor
final class OpenAIDashboardWebViewCache {
static let shared = OpenAIDashboardWebViewCache()
fileprivate static let log = CodexBarLog.logger(LogCategories.openAIWebview)
private final class ReleaseState {
var preserveLoadedPageOnRelease: Bool
init(preserveLoadedPageOnRelease: Bool) {
self.preserveLoadedPageOnRelease = preserveLoadedPageOnRelease
}
}
private struct AcquireOptions {
let allowTimeoutRetry: Bool
let preserveLoadedPageOnRelease: Bool
}
@MainActor
private final class NavigationCancellationState {
private weak var webView: WKWebView?
private var delegate: NavigationDelegate?
private var isCancelled = false
func install(webView: WKWebView, delegate: NavigationDelegate) {
self.webView = webView
self.delegate = delegate
if self.isCancelled {
self.cancel()
}
}
func cancel() {
self.isCancelled = true
guard let webView, let delegate else { return }
delegate.cancel()
if webView.codexNavigationDelegate === delegate {
webView.stopLoading()
webView.navigationDelegate = nil
webView.codexNavigationDelegate = nil
}
self.delegate = nil
self.webView = nil
}
}
private final class Entry {
let webView: WKWebView
let host: OffscreenWebViewHost
var lastUsedAt: Date
var isBusy: Bool
var preservedPageExpiresAt: Date?
var preservedPageExpiryWorkItem: DispatchWorkItem?
init(
webView: WKWebView,
host: OffscreenWebViewHost,
lastUsedAt: Date,
isBusy: Bool,
preservedPageExpiresAt: Date? = nil)
{
self.webView = webView
self.host = host
self.lastUsedAt = lastUsedAt
self.isBusy = isBusy
self.preservedPageExpiresAt = preservedPageExpiresAt
}
func armPreservedPage(until expiry: Date) {
self.preservedPageExpiresAt = expiry
}
func setPreservedPageExpiryWorkItem(_ workItem: DispatchWorkItem?) {
self.preservedPageExpiryWorkItem?.cancel()
self.preservedPageExpiryWorkItem = workItem
}
func clearPreservedPage() {
self.preservedPageExpiresAt = nil
self.preservedPageExpiryWorkItem?.cancel()
self.preservedPageExpiryWorkItem = nil
}
func consumePreservedPageReuseIfAvailable(now: Date) -> Bool {
guard let preservedPageExpiresAt else { return false }
self.preservedPageExpiresAt = nil
self.preservedPageExpiryWorkItem?.cancel()
self.preservedPageExpiryWorkItem = nil
return preservedPageExpiresAt > now
}
func hasExpiredPreservedPage(now: Date) -> Bool {
guard let preservedPageExpiresAt else { return false }
return preservedPageExpiresAt <= now
}
}
private var entries: [ObjectIdentifier: Entry] = [:]
/// Keep the WebView alive only long enough for immediate retries/menu reopens.
/// Long-lived hidden ChatGPT tabs still consume noticeable energy on some setups.
private let idleTimeout: TimeInterval
private var idlePruneWorkItem: DispatchWorkItem?
private var idlePruneGeneration = 0
#if DEBUG
private(set) var idlePruneDeadlineForTesting: Date?
#endif
/// Reuse the validated analytics page only for the immediate next handoff.
private let preservedPageHandoffTimeout: TimeInterval = 5
private let blankURL = URL(string: "about:blank")!
private let idlePageClearScript = """
(() => {
try {
document.documentElement.innerHTML = '';
return true;
} catch {
return false;
}
})();
"""
private let reusablePageResetScript = """
(() => {
try {
delete window.__codexbarDidScrollToCredits;
delete window.__codexbarUsageBreakdownJSON;
delete window.__codexbarUsageBreakdownDebug;
return true;
} catch {
return false;
}
})();
"""
private let preferredLanguageScript = """
(() => {
const define = (target, name, value) => {
try {
Object.defineProperty(target, name, {
get: () => value,
configurable: true
});
} catch {}
};
define(Navigator.prototype, 'language', 'en-US');
define(Navigator.prototype, 'languages', ['en-US', 'en']);
define(navigator, 'language', 'en-US');
define(navigator, 'languages', ['en-US', 'en']);
})();
"""
init(idleTimeout: TimeInterval = 60) {
self.idleTimeout = idleTimeout
}
nonisolated static func remainingNavigationTimeout(
until deadline: Date,
now: Date = Date()) throws -> TimeInterval
{
let remaining = deadline.timeIntervalSince(now)
guard remaining > 0 else { throw URLError(.timedOut) }
return remaining
}
private func releaseCachedEntry(_ entry: Entry, preserveLoadedPage: Bool) {
entry.isBusy = false
let now = Date()
entry.lastUsedAt = now
self.updatePreservedPageState(for: entry, preserveLoadedPage: preserveLoadedPage)
self.prepareCachedWebViewForIdle(
entry.webView,
host: entry.host,
preserveLoadedPage: preserveLoadedPage)
self.prune(now: now)
self.scheduleNextIdlePrune(now: now)
}
private func releaseNewEntry(_ entry: Entry, webView: WKWebView, preserveLoadedPage: Bool) {
entry.isBusy = false
let now = Date()
entry.lastUsedAt = now
self.updatePreservedPageState(for: entry, preserveLoadedPage: preserveLoadedPage)
self.prepareCachedWebViewForIdle(
webView,
host: entry.host,
preserveLoadedPage: preserveLoadedPage)
self.prune(now: now)
self.scheduleNextIdlePrune(now: now)
}
// MARK: - Testing support
#if DEBUG
/// Number of cached WebView entries (for testing).
var entryCount: Int {
self.entries.count
}
/// Check if a WebView is cached for the given data store (for testing).
func hasCachedEntry(for websiteDataStore: WKWebsiteDataStore) -> Bool {
let key = ObjectIdentifier(websiteDataStore)
return self.entries[key] != nil
}
/// Force prune with a custom "now" timestamp (for testing idle timeout).
func pruneForTesting(now: Date) {
self.prune(now: now)
}
var idleTimeoutForTesting: TimeInterval {
self.idleTimeout
}
var preservedPageHandoffTimeoutForTesting: TimeInterval {
self.preservedPageHandoffTimeout
}
func hasPreservedPageForTesting(for websiteDataStore: WKWebsiteDataStore) -> Bool {
let key = ObjectIdentifier(websiteDataStore)
return self.entries[key]?.preservedPageExpiresAt != nil
}
func markPreservedPageForTesting(
websiteDataStore: WKWebsiteDataStore,
expiresAt: Date = .init().addingTimeInterval(5))
{
let key = ObjectIdentifier(websiteDataStore)
guard let entry = self.entries[key] else { return }
entry.armPreservedPage(until: expiresAt)
self.schedulePreservedPageExpiry(for: key, entry: entry, expiresAt: expiresAt)
}
func consumePreservedPageForTesting(websiteDataStore: WKWebsiteDataStore, now: Date = Date()) -> Bool {
let key = ObjectIdentifier(websiteDataStore)
guard let entry = self.entries[key] else { return false }
return entry.consumePreservedPageReuseIfAvailable(now: now)
}
/// Seed a cached entry without navigating a real page (for test stability).
@discardableResult
func cacheEntryForTesting(
websiteDataStore: WKWebsiteDataStore,
lastUsedAt: Date = Date(),
isBusy: Bool = false) -> WKWebView
{
let key = ObjectIdentifier(websiteDataStore)
if let existing = self.entries.removeValue(forKey: key) {
existing.host.close()
}
let (webView, host) = self.makeWebView(websiteDataStore: websiteDataStore)
let entry = Entry(webView: webView, host: host, lastUsedAt: lastUsedAt, isBusy: isBusy)
self.entries[key] = entry
if isBusy {
host.show()
} else {
host.hide()
}
return webView
}
/// Clear all cached entries (for test isolation).
func clearAllForTesting() {
self.cancelIdlePrune()
for (_, entry) in self.entries {
entry.clearPreservedPage()
entry.host.close()
}
self.entries.removeAll()
}
func resetReusablePageStateForTesting(_ webView: WKWebView) async -> Bool {
await self.resetReusablePageState(webView)
}
#endif
func acquire(
websiteDataStore: WKWebsiteDataStore,
usageURL: URL,
logger: ((String) -> Void)?,
navigationTimeout: TimeInterval = 15,
allowTimeoutRetry: Bool = true,
preserveLoadedPageOnRelease: Bool = false) async throws -> OpenAIDashboardWebViewLease
{
let deadline = Date().addingTimeInterval(max(navigationTimeout, 0.01))
return try await self.acquire(
websiteDataStore: websiteDataStore,
usageURL: usageURL,
logger: logger,
deadline: deadline,
options: .init(
allowTimeoutRetry: allowTimeoutRetry,
preserveLoadedPageOnRelease: preserveLoadedPageOnRelease))
}
private func acquire(
websiteDataStore: WKWebsiteDataStore,
usageURL: URL,
logger: ((String) -> Void)?,
deadline: Date,
options: AcquireOptions) async throws -> OpenAIDashboardWebViewLease
{
let now = Date()
self.prune(now: now)
let log: (String) -> Void = { message in
logger?("[webview] \(message)")
}
let key = ObjectIdentifier(websiteDataStore)
let remainingTimeout = try Self.remainingNavigationTimeout(until: deadline, now: now)
if let entry = self.entries[key] {
if entry.isBusy {
log("Cached WebView busy; using a temporary WebView.")
let (webView, host) = self.makeWebView(websiteDataStore: websiteDataStore)
host.show()
do {
try await self.prepareWebView(
webView,
usageURL: usageURL,
timeout: remainingTimeout,
canReuseLoadedPage: false)
} catch {
if options.allowTimeoutRetry, Self.isPrepareTimeout(error) {
host.close()
log("Temporary OpenAI WebView timed out; retrying with a fresh WebView.")
return try await self.acquireTemporaryWebView(
websiteDataStore: websiteDataStore,
usageURL: usageURL,
log: log,
deadline: deadline)
}
host.close()
throw error
}
return OpenAIDashboardWebViewLease(
webView: webView,
log: log,
setPreserveLoadedPageOnRelease: { _ in },
release: { host.close() })
}
entry.isBusy = true
entry.lastUsedAt = now
let canReuseLoadedPage = entry.consumePreservedPageReuseIfAvailable(now: now)
let releaseState = ReleaseState(preserveLoadedPageOnRelease: options.preserveLoadedPageOnRelease)
entry.host.show()
do {
try await self.prepareWebView(
entry.webView,
usageURL: usageURL,
timeout: remainingTimeout,
canReuseLoadedPage: canReuseLoadedPage)
} catch {
if options.allowTimeoutRetry, Self.isPrepareTimeout(error) {
entry.isBusy = false
entry.lastUsedAt = Date()
entry.clearPreservedPage()
entry.host.close()
self.entries.removeValue(forKey: key)
log("Cached OpenAI WebView timed out; recreating it.")
return try await self.acquire(
websiteDataStore: websiteDataStore,
usageURL: usageURL,
logger: logger,
deadline: deadline,
options: .init(
allowTimeoutRetry: false,
preserveLoadedPageOnRelease: options.preserveLoadedPageOnRelease))
}
entry.isBusy = false
entry.lastUsedAt = Date()
entry.clearPreservedPage()
entry.host.close()
self.entries.removeValue(forKey: key)
Self.log.warning("OpenAI webview prepare failed")
throw error
}
return OpenAIDashboardWebViewLease(
webView: entry.webView,
log: log,
setPreserveLoadedPageOnRelease: { preserveLoadedPageOnRelease in
releaseState.preserveLoadedPageOnRelease = preserveLoadedPageOnRelease
},
release: { [weak self, weak entry] in
guard let self, let entry else { return }
self.releaseCachedEntry(
entry,
preserveLoadedPage: releaseState.preserveLoadedPageOnRelease)
})
}
let (webView, host) = self.makeWebView(websiteDataStore: websiteDataStore)
let entry = Entry(webView: webView, host: host, lastUsedAt: now, isBusy: true)
self.entries[key] = entry
host.show()
let releaseState = ReleaseState(preserveLoadedPageOnRelease: options.preserveLoadedPageOnRelease)
do {
try await self.prepareWebView(
webView,
usageURL: usageURL,
timeout: remainingTimeout,
canReuseLoadedPage: false)
} catch {
if options.allowTimeoutRetry, Self.isPrepareTimeout(error) {
self.entries.removeValue(forKey: key)
host.close()
log("OpenAI WebView timed out during prepare; retrying once.")
return try await self.acquire(
websiteDataStore: websiteDataStore,
usageURL: usageURL,
logger: logger,
deadline: deadline,
options: .init(
allowTimeoutRetry: false,
preserveLoadedPageOnRelease: options.preserveLoadedPageOnRelease))
}
self.entries.removeValue(forKey: key)
host.close()
Self.log.warning("OpenAI webview prepare failed")
throw error
}
return OpenAIDashboardWebViewLease(
webView: webView,
log: log,
setPreserveLoadedPageOnRelease: { preserveLoadedPageOnRelease in
releaseState.preserveLoadedPageOnRelease = preserveLoadedPageOnRelease
},
release: { [weak self, weak entry] in
guard let self, let entry else { return }
self.releaseNewEntry(
entry,
webView: webView,
preserveLoadedPage: releaseState.preserveLoadedPageOnRelease)
})
}
func evict(websiteDataStore: WKWebsiteDataStore) {
let key = ObjectIdentifier(websiteDataStore)
guard let entry = self.entries.removeValue(forKey: key) else { return }
entry.clearPreservedPage()
Self.log.debug("OpenAI webview evicted")
entry.host.close()
self.scheduleNextIdlePrune()
}
func evictAll() {
self.cancelIdlePrune()
let existing = self.entries
self.entries.removeAll()
for (_, entry) in existing {
entry.clearPreservedPage()
entry.host.close()
}
if !existing.isEmpty {
Self.log.debug("OpenAI webview evicted all")
}
}
func evictIdle() {
let idleEntries = self.entries.filter { _, entry in
!entry.isBusy
}
guard !idleEntries.isEmpty else { return }
for (key, entry) in idleEntries {
entry.clearPreservedPage()
entry.host.close()
self.entries.removeValue(forKey: key)
}
Self.log.debug("OpenAI idle webviews evicted", metadata: ["count": "\(idleEntries.count)"])
self.scheduleNextIdlePrune()
}
private func prepareCachedWebViewForIdle(
_ webView: WKWebView,
host: OffscreenWebViewHost,
preserveLoadedPage: Bool)
{
webView.navigationDelegate = nil
webView.codexNavigationDelegate = nil
if preserveLoadedPage {
host.hide()
return
}
// Detach the heavyweight ChatGPT SPA as soon as a scrape completes. Keeping the WebView object around
// still helps with immediate reuse, but letting chatgpt.com remain the active document is too expensive.
webView.stopLoading()
webView.evaluateJavaScript(self.idlePageClearScript, completionHandler: nil)
_ = webView.load(URLRequest(url: self.blankURL))
host.hide()
}
/// Schedule against the oldest idle entry so later releases cannot postpone its eviction.
private func scheduleNextIdlePrune(now: Date = Date()) {
self.cancelIdlePrune()
guard let nextExpiry = self.entries.values
.filter({ !$0.isBusy })
.map({ $0.lastUsedAt.addingTimeInterval(self.idleTimeout) })
.min()
else { return }
let generation = self.idlePruneGeneration
let workItem = DispatchWorkItem { [weak self] in
MainActor.assumeIsolated {
guard let self, self.idlePruneGeneration == generation else { return }
self.idlePruneWorkItem = nil
#if DEBUG
self.idlePruneDeadlineForTesting = nil
#endif
let pruneTime = Date()
self.prune(now: pruneTime)
self.scheduleNextIdlePrune(now: pruneTime)
}
}
self.idlePruneWorkItem = workItem
#if DEBUG
self.idlePruneDeadlineForTesting = nextExpiry
#endif
let delay = max(0, nextExpiry.timeIntervalSince(now)) + 0.01
DispatchQueue.main.asyncAfter(deadline: .now() + delay, execute: workItem)
}
private func cancelIdlePrune() {
self.idlePruneGeneration &+= 1
self.idlePruneWorkItem?.cancel()
self.idlePruneWorkItem = nil
#if DEBUG
self.idlePruneDeadlineForTesting = nil
#endif
}
private func prune(now: Date) {
for entry in self.entries.values where !entry.isBusy && entry.hasExpiredPreservedPage(now: now) {
entry.clearPreservedPage()
self.prepareCachedWebViewForIdle(
entry.webView,
host: entry.host,
preserveLoadedPage: false)
Self.log.debug("OpenAI webview preserved page expired")
}
let expired = self.entries.filter { _, entry in
!entry.isBusy && now.timeIntervalSince(entry.lastUsedAt) >= self.idleTimeout
}
for (key, entry) in expired {
entry.host.close()
self.entries.removeValue(forKey: key)
Self.log.debug("OpenAI webview pruned")
}
}
private func makeWebView(websiteDataStore: WKWebsiteDataStore) -> (WKWebView, OffscreenWebViewHost) {
let config = WKWebViewConfiguration()
config.websiteDataStore = websiteDataStore
let userContentController = WKUserContentController()
userContentController.addUserScript(WKUserScript(
source: self.preferredLanguageScript,
injectionTime: .atDocumentStart,
forMainFrameOnly: false))
config.userContentController = userContentController
if #available(macOS 14.0, *) {
config.preferences.inactiveSchedulingPolicy = .suspend
}
let webView = WKWebView(frame: .zero, configuration: config)
let host = OffscreenWebViewHost(webView: webView)
return (webView, host)
}
private func prepareWebView(
_ webView: WKWebView,
usageURL: URL,
timeout: TimeInterval,
canReuseLoadedPage: Bool) async throws
{
#if DEBUG
if usageURL.absoluteString == "about:blank" {
_ = webView.loadHTMLString("", baseURL: nil)
return
}
#endif
if canReuseLoadedPage,
let currentURL = webView.url?.absoluteString,
OpenAIDashboardFetcher.isUsageRoute(currentURL)
{
if await self.resetReusablePageState(webView) {
return
}
Self.log.debug("OpenAI preserved page reset failed; reloading usage URL")
}
try Task.checkCancellation()
let cancellationState = NavigationCancellationState()
try await withTaskCancellationHandler {
try await withCheckedThrowingContinuation { (cont: CheckedContinuation<Void, Error>) in
let delegate = NavigationDelegate { result in
cont.resume(with: result)
}
webView.navigationDelegate = delegate
webView.codexNavigationDelegate = delegate
cancellationState.install(webView: webView, delegate: delegate)
delegate.armTimeout(seconds: timeout)
_ = webView.load(OpenAIDashboardFetcher.usageURLRequest(url: usageURL))
if Task.isCancelled {
cancellationState.cancel()
}
}
} onCancel: {
Task { @MainActor in
cancellationState.cancel()
}
}
}
private func acquireTemporaryWebView(
websiteDataStore: WKWebsiteDataStore,
usageURL: URL,
log: @escaping (String) -> Void,
deadline: Date) async throws -> OpenAIDashboardWebViewLease
{
let remainingTimeout = try Self.remainingNavigationTimeout(until: deadline)
let (webView, host) = self.makeWebView(websiteDataStore: websiteDataStore)
host.show()
do {
try await self.prepareWebView(
webView,
usageURL: usageURL,
timeout: remainingTimeout,
canReuseLoadedPage: false)
} catch {
host.close()
throw error
}
return OpenAIDashboardWebViewLease(
webView: webView,
log: log,
setPreserveLoadedPageOnRelease: { _ in },
release: { host.close() })
}
private static func isPrepareTimeout(_ error: Error) -> Bool {
let nsError = error as NSError
return nsError.domain == NSURLErrorDomain && nsError.code == NSURLErrorTimedOut
}
private func updatePreservedPageState(for entry: Entry, preserveLoadedPage: Bool) {
if preserveLoadedPage {
let expiresAt = Date().addingTimeInterval(self.preservedPageHandoffTimeout)
entry.armPreservedPage(until: expiresAt)
if let key = self.entries.first(where: { $0.value === entry })?.key {
self.schedulePreservedPageExpiry(for: key, entry: entry, expiresAt: expiresAt)
}
} else {
entry.clearPreservedPage()
}
}
private func schedulePreservedPageExpiry(
for key: ObjectIdentifier,
entry: Entry,
expiresAt: Date)
{
let delay = max(0, expiresAt.timeIntervalSinceNow)
let workItem = DispatchWorkItem { [weak self] in
MainActor.assumeIsolated {
self?.expirePreservedPageIfNeeded(for: key, expectedExpiry: expiresAt)
}
}
entry.setPreservedPageExpiryWorkItem(workItem)
DispatchQueue.main.asyncAfter(deadline: .now() + delay, execute: workItem)
}
private func expirePreservedPageIfNeeded(for key: ObjectIdentifier, expectedExpiry: Date) {
guard let entry = self.entries[key],
!entry.isBusy,
let preservedPageExpiresAt = entry.preservedPageExpiresAt,
preservedPageExpiresAt == expectedExpiry,
preservedPageExpiresAt <= Date()
else {
return
}
entry.clearPreservedPage()
self.prepareCachedWebViewForIdle(
entry.webView,
host: entry.host,
preserveLoadedPage: false)
Self.log.debug("OpenAI webview preserved page expired")
self.prune(now: Date())
}
private func resetReusablePageState(_ webView: WKWebView) async -> Bool {
do {
let any = try await webView.evaluateJavaScript(self.reusablePageResetScript)
return (any as? Bool) ?? true
} catch {
return false
}
}
}
@MainActor
private final class OffscreenWebViewHost {
private let window: NSWindow
private weak var webView: WKWebView?
init(webView: WKWebView) {
// WebKit throttles timers/RAF aggressively when a WKWebView is not considered "visible".
// The Codex usage page uses streaming SSR + client hydration; if RAF is throttled, the
// dashboard never becomes part of the visible DOM and `document.body.innerText` stays tiny.
//
// Keep a transparent (mouse-ignoring) window technically "on-screen" while scraping, but
// place it almost entirely off-screen so we never ghost-render dashboard UI over the desktop.
let visibleFrame = NSScreen.main?.visibleFrame ?? NSRect(x: 0, y: 0, width: 900, height: 700)
let frame = OpenAIDashboardFetcher.offscreenHostWindowFrame(for: visibleFrame)
let window = NSWindow(
contentRect: frame,
styleMask: [.borderless],
backing: .buffered,
defer: false)
window.isReleasedWhenClosed = false
window.backgroundColor = .clear
window.isOpaque = false
// Keep it effectively invisible, but non-zero alpha so WebKit treats it as "visible" and doesn't
// stall hydration (we've observed a head-only HTML shell for minutes at alpha=0).
window.alphaValue = OpenAIDashboardFetcher.offscreenHostAlphaValue()
window.hasShadow = false
window.ignoresMouseEvents = true
window.level = .floating
window.collectionBehavior = [.canJoinAllSpaces, .fullScreenAuxiliary]
window.isExcludedFromWindowsMenu = true
window.contentView = webView
self.window = window
self.webView = webView
}
func show() {
OpenAIDashboardWebViewCache.log.debug("OpenAI webview show")
self.window.alphaValue = OpenAIDashboardFetcher.offscreenHostAlphaValue()
self.window.orderFrontRegardless()
}
func hide() {
// Set alpha to 0 so WebKit recognizes the page as inactive and applies
// its scheduling policy (throttle/suspend), reducing CPU when idle.
OpenAIDashboardWebViewCache.log.debug("OpenAI webview hide")
self.window.alphaValue = 0.0
self.window.orderOut(nil)
}
func close() {
OpenAIDashboardWebViewCache.log.debug("OpenAI webview close")
WebKitTeardown.scheduleCleanup(
owner: self,
window: self.window,
webView: self.webView,
closeWindow: { [window] in
window.orderOut(nil)
window.close()
})
}
}
#endif
@@ -0,0 +1,118 @@
#if os(macOS)
import CryptoKit
import Foundation
import WebKit
/// Per-account persistent `WKWebsiteDataStore` for the OpenAI dashboard scrape.
///
/// Why: `WKWebsiteDataStore.default()` is a single shared cookie jar. If the user switches Codex accounts,
/// we want to keep multiple signed-in dashboard sessions around (one per email) without clearing cookies.
///
/// Implementation detail: macOS 14+ supports `WKWebsiteDataStore.dataStore(forIdentifier:)`, which creates
/// persistent isolated stores keyed by an identifier. We derive a stable UUID from the email and optional
/// Codex source scope so distinct profiles with the same email never share a cookie store.
///
/// Important: We cache the `WKWebsiteDataStore` instances so the same object is returned for the same
/// account email. This ensures `OpenAIDashboardWebViewCache` can use object identity for cache lookups.
@MainActor
public enum OpenAIDashboardWebsiteDataStore {
/// Cached data store instances keyed by normalized email and optional Codex source scope.
/// Using the same instance ensures stable object identity for WebView cache lookups.
private static var cachedStores: [String: WKWebsiteDataStore] = [:]
public static func store(
forAccountEmail email: String?,
scope: CookieHeaderCache.Scope? = nil) -> WKWebsiteDataStore
{
guard let normalized = normalizeEmail(email) else { return .default() }
let storageKey = self.storageKey(normalizedEmail: normalized, scope: scope)
// Return cached instance if available to maintain stable object identity
if let cached = cachedStores[storageKey] {
return cached
}
let id = Self.identifier(forStorageKey: storageKey)
let store = WKWebsiteDataStore(forIdentifier: id)
self.cachedStores[storageKey] = store
return store
}
/// Clears the persistent cookie store for a single account email.
///
/// Note: this does *not* impact other accounts, and is safe to use when the stored session is "stuck"
/// or signed in to a different account than expected.
public static func clearStore(
forAccountEmail email: String?,
scope: CookieHeaderCache.Scope? = nil) async
{
// Clear only ChatGPT/OpenAI domain data for the per-account store.
// Avoid deleting the entire persistent store (WebKit requires all WKWebViews using it to be released).
let store = self.store(forAccountEmail: email, scope: scope)
await withCheckedContinuation { cont in
store.fetchDataRecords(ofTypes: WKWebsiteDataStore.allWebsiteDataTypes()) { records in
let filtered = records.filter { record in
let name = record.displayName.lowercased()
return name.contains("chatgpt.com") || name.contains("openai.com")
}
store.removeData(ofTypes: WKWebsiteDataStore.allWebsiteDataTypes(), for: filtered) {
cont.resume()
}
}
}
// Remove from cache so a fresh instance is created on next access
if let normalized = normalizeEmail(email) {
self.cachedStores.removeValue(forKey: self.storageKey(normalizedEmail: normalized, scope: scope))
}
}
#if DEBUG
/// Clear all cached store instances (for test isolation).
public static func clearCacheForTesting() {
self.cachedStores.removeAll()
}
#endif
// MARK: - Private
private static func normalizeEmail(_ email: String?) -> String? {
guard let raw = email?.trimmingCharacters(in: .whitespacesAndNewlines), !raw.isEmpty else { return nil }
return raw.lowercased()
}
private static func storageKey(normalizedEmail: String, scope: CookieHeaderCache.Scope?) -> String {
guard let scope else { return normalizedEmail }
return "\(normalizedEmail)|\(scope.isolationIdentifier)"
}
private static func identifier(forStorageKey storageKey: String) -> UUID {
let digest = SHA256.hash(data: Data(storageKey.utf8))
var bytes = Array(digest.prefix(16))
// Make it a well-formed UUID (v4 + RFC4122 variant) while staying deterministic.
bytes[6] = (bytes[6] & 0x0F) | 0x40
bytes[8] = (bytes[8] & 0x3F) | 0x80
let uuidBytes: uuid_t = (
bytes[0],
bytes[1],
bytes[2],
bytes[3],
bytes[4],
bytes[5],
bytes[6],
bytes[7],
bytes[8],
bytes[9],
bytes[10],
bytes[11],
bytes[12],
bytes[13],
bytes[14],
bytes[15])
return UUID(uuid: uuidBytes)
}
}
#endif
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,97 @@
import Foundation
enum PiSessionCostCacheIO {
/// Artifact schema version. Pricing changes are tracked separately by `pricingKey`.
private static let artifactVersion = 5
private static func defaultCacheRoot() -> URL {
let root = FileManager.default.urls(for: .cachesDirectory, in: .userDomainMask).first!
return root.appendingPathComponent("CodexBar", isDirectory: true)
}
static func cacheFileURL(cacheRoot: URL? = nil) -> URL {
let root = cacheRoot ?? self.defaultCacheRoot()
return root
.appendingPathComponent("cost-usage", isDirectory: true)
.appendingPathComponent("pi-sessions-v\(Self.artifactVersion).json", isDirectory: false)
}
static func load(cacheRoot: URL? = nil) -> PiSessionCostCache {
let url = self.cacheFileURL(cacheRoot: cacheRoot)
guard let data = try? Data(contentsOf: url),
let decoded = try? JSONDecoder().decode(PiSessionCostCache.self, from: data),
decoded.version == Self.artifactVersion
else {
return PiSessionCostCache(version: Self.artifactVersion)
}
return decoded
}
static func save(cache: PiSessionCostCache, cacheRoot: URL? = nil) {
let url = self.cacheFileURL(cacheRoot: cacheRoot)
let dir = url.deletingLastPathComponent()
try? FileManager.default.createDirectory(at: dir, withIntermediateDirectories: true)
let tmp = dir.appendingPathComponent(".tmp-\(UUID().uuidString).json", isDirectory: false)
let data = (try? JSONEncoder().encode(cache)) ?? Data()
do {
try data.write(to: tmp, options: [.atomic])
if FileManager.default.fileExists(atPath: url.path) {
_ = try FileManager.default.replaceItemAt(url, withItemAt: tmp)
} else {
try FileManager.default.moveItem(at: tmp, to: url)
}
} catch {
try? FileManager.default.removeItem(at: tmp)
}
}
}
struct PiSessionCostCache: Codable {
var version: Int
var lastScanUnixMs: Int64 = 0
var scanSinceKey: String?
var scanUntilKey: String?
var pricingKey: String?
var daysByProvider: [String: [String: [String: PiPackedUsage]]] = [:]
var files: [String: PiSessionFileUsage] = [:]
init(version: Int = 5) {
self.version = version
}
}
struct PiSessionFileUsage: Codable {
var mtimeUnixMs: Int64
var size: Int64
var parsedBytes: Int64
var lastModelContext: PiModelContext?
var contributions: [String: [String: [String: PiPackedUsage]]]
}
struct PiModelContext: Codable, Equatable {
var providerRawValue: String
var modelName: String
}
struct PiPackedUsage: Codable, Equatable {
var inputTokens: Int = 0
var cacheReadTokens: Int = 0
var cacheWriteTokens: Int = 0
var outputTokens: Int = 0
var totalTokens: Int = 0
var costNanos: Int64 = 0
var costSampleCount: Int = 0
var usageSampleCount: Int?
var isZero: Bool {
self.inputTokens == 0
&& self.cacheReadTokens == 0
&& self.cacheWriteTokens == 0
&& self.outputTokens == 0
&& self.totalTokens == 0
&& self.costNanos == 0
&& self.costSampleCount == 0
&& (self.usageSampleCount ?? 0) == 0
}
}
@@ -0,0 +1,943 @@
import Foundation
private final class PiSessionISO8601FormatterBox: @unchecked Sendable {
let lock = NSLock()
let withFractional: ISO8601DateFormatter = {
let formatter = ISO8601DateFormatter()
formatter.formatOptions = [.withInternetDateTime, .withFractionalSeconds]
return formatter
}()
let plain: ISO8601DateFormatter = {
let formatter = ISO8601DateFormatter()
formatter.formatOptions = [.withInternetDateTime]
return formatter
}()
}
enum PiSessionCostScanner {
struct Options {
var piSessionsRoot: URL?
var cacheRoot: URL?
var refreshMinIntervalSeconds: TimeInterval = 60
var forceRescan: Bool = false
init(
piSessionsRoot: URL? = nil,
cacheRoot: URL? = nil,
refreshMinIntervalSeconds: TimeInterval = 60,
forceRescan: Bool = false)
{
self.piSessionsRoot = piSessionsRoot
self.cacheRoot = cacheRoot
self.refreshMinIntervalSeconds = refreshMinIntervalSeconds
self.forceRescan = forceRescan
}
}
private struct ParseResult {
let contributions: [String: [String: [String: PiPackedUsage]]]
let parsedBytes: Int64
let lastModelContext: PiModelContext?
}
private struct AssistantIdentity {
let provider: UsageProvider
let modelName: String
}
private struct ModelsDevPricingContext {
let catalog: ModelsDevCatalog?
let cacheRoot: URL?
let pricingKey: String
}
private struct ScanContext {
let range: CostUsageScanner.CostUsageDayRange
let forceRescan: Bool
let pricingContext: ModelsDevPricingContext
let checkCancellation: CostUsageScanner.CancellationCheck?
}
private static let costScale = 1_000_000_000.0
/// Bump for Pi-only cost formula changes not represented by the parser or pricing fingerprints.
private static let costFormulaVersion = 1
private static let maxLineBytes = 16 * 1024 * 1024
private static let maxSafeRoundedInt = Double(Int.max) - 1
private static let sessionStartFilenameRegex = try? NSRegularExpression(
pattern: "^(\\d{4}-\\d{2}-\\d{2})T(\\d{2})-(\\d{2})-(\\d{2})-(\\d{3})Z_")
private static let isoFormatterBox = PiSessionISO8601FormatterBox()
static func loadDailyReport(
provider: UsageProvider,
since: Date,
until: Date,
now: Date = Date(),
options: Options = Options()) -> CostUsageDailyReport
{
(
try? self.loadDailyReportCancellable(
provider: provider,
since: since,
until: until,
now: now,
options: options,
checkCancellation: nil)) ?? CostUsageDailyReport(data: [], summary: nil)
}
static func loadDailyReportCancellable(
provider: UsageProvider,
since: Date,
until: Date,
now: Date = Date(),
options: Options = Options(),
checkCancellation: CostUsageScanner.CancellationCheck?) throws -> CostUsageDailyReport
{
guard provider == .codex || provider == .claude else {
return CostUsageDailyReport(data: [], summary: nil)
}
let range = CostUsageScanner.CostUsageDayRange(since: since, until: until)
var cache = PiSessionCostCacheIO.load(cacheRoot: options.cacheRoot)
let nowMs = Int64(now.timeIntervalSince1970 * 1000)
let refreshMs = Int64(max(0, options.refreshMinIntervalSeconds) * 1000)
let pricingContext = self.pricingContext(now: now, cacheRoot: options.cacheRoot)
let windowExpanded = self.requestedWindowExpandsCache(range: range, cache: cache)
let pricingChanged = cache.pricingKey != pricingContext.pricingKey
let shouldRefresh = options.forceRescan
|| windowExpanded
|| pricingChanged
|| refreshMs == 0
|| cache.lastScanUnixMs == 0
|| nowMs - cache.lastScanUnixMs > refreshMs
if shouldRefresh {
try checkCancellation?()
let root = self.defaultPiSessionsRoot(options: options)
let startCutoff = self.dateFromDayKey(range.scanSinceKey) ?? since
let files = self.listPiSessionFiles(root: root, startCutoffLocal: startCutoff)
let filePathsInScan = Set(files.map(\.path))
for fileURL in files {
try self.scanPiSessionFile(
fileURL: fileURL,
cache: &cache,
context: ScanContext(
range: range,
forceRescan: options.forceRescan || windowExpanded || pricingChanged,
pricingContext: pricingContext,
checkCancellation: checkCancellation))
}
try checkCancellation?()
for key in cache.files.keys where !filePathsInScan.contains(key) {
if let old = cache.files[key] {
self.applyContributions(
daysByProvider: &cache.daysByProvider,
contributions: old.contributions,
sign: -1)
}
cache.files.removeValue(forKey: key)
}
cache.scanSinceKey = range.scanSinceKey
cache.scanUntilKey = range.scanUntilKey
cache.pricingKey = pricingContext.pricingKey
cache.lastScanUnixMs = nowMs
try checkCancellation?()
PiSessionCostCacheIO.save(cache: cache, cacheRoot: options.cacheRoot)
}
return self.buildReport(
provider: provider,
cache: cache,
range: range,
pricingContext: pricingContext)
}
struct CachedDailyReportResult {
let report: CostUsageDailyReport
let lastScanAt: Date?
}
static func loadCachedDailyReport(
provider: UsageProvider,
since: Date,
until: Date,
now: Date = Date(),
cacheRoot: URL? = nil) -> CostUsageDailyReport?
{
self.loadCachedDailyReportResult(
provider: provider,
since: since,
until: until,
now: now,
cacheRoot: cacheRoot)?.report
}
static func loadCachedDailyReportResult(
provider: UsageProvider,
since: Date,
until: Date,
now: Date = Date(),
cacheRoot: URL? = nil) -> CachedDailyReportResult?
{
guard provider == .codex || provider == .claude else { return nil }
let range = CostUsageScanner.CostUsageDayRange(since: since, until: until)
let cache = PiSessionCostCacheIO.load(cacheRoot: cacheRoot)
guard !cache.daysByProvider.isEmpty else { return nil }
guard !self.requestedWindowExpandsCache(range: range, cache: cache) else { return nil }
let pricingContext = self.pricingContext(now: now, cacheRoot: cacheRoot)
guard cache.pricingKey == pricingContext.pricingKey else { return nil }
let report = self.buildReport(
provider: provider,
cache: cache,
range: range,
pricingContext: pricingContext)
guard !report.data.isEmpty else { return nil }
let lastScanAt = cache.lastScanUnixMs > 0
? Date(timeIntervalSince1970: TimeInterval(cache.lastScanUnixMs) / 1000)
: nil
return CachedDailyReportResult(report: report, lastScanAt: lastScanAt)
}
private static func pricingContext(now: Date, cacheRoot: URL?) -> ModelsDevPricingContext {
let modelsDevArtifact = ModelsDevCache.load(now: now, cacheRoot: cacheRoot).artifact
return ModelsDevPricingContext(
catalog: modelsDevArtifact?.catalog,
cacheRoot: cacheRoot,
pricingKey: CostUsagePricingKey.codex(
modelsDevArtifact: modelsDevArtifact,
formulaVersion: Self.costFormulaVersion,
parserHash: CodexParserHash.value,
modelsDevProviderIDs: ["anthropic", "openai"]))
}
private static func requestedWindowExpandsCache(
range: CostUsageScanner.CostUsageDayRange,
cache: PiSessionCostCache) -> Bool
{
guard let cachedSince = cache.scanSinceKey,
let cachedUntil = cache.scanUntilKey
else {
return true
}
if range.scanSinceKey < cachedSince {
return true
}
if range.scanUntilKey > cachedUntil {
return true
}
return false
}
private static func defaultPiSessionsRoot(options: Options) -> URL {
if let override = options.piSessionsRoot { return override }
return FileManager.default.homeDirectoryForCurrentUser
.appendingPathComponent(".pi", isDirectory: true)
.appendingPathComponent("agent", isDirectory: true)
.appendingPathComponent("sessions", isDirectory: true)
}
private static func listPiSessionFiles(root: URL, startCutoffLocal: Date) -> [URL] {
guard FileManager.default.fileExists(atPath: root.path) else { return [] }
let keys: Set<URLResourceKey> = [.isRegularFileKey, .contentModificationDateKey]
guard let enumerator = FileManager.default.enumerator(
at: root,
includingPropertiesForKeys: Array(keys),
options: [.skipsHiddenFiles])
else {
return []
}
var output: [URL] = []
while let item = enumerator.nextObject() as? URL {
guard item.pathExtension.lowercased() == "jsonl" else { continue }
let values = try? item.resourceValues(forKeys: keys)
guard values?.isRegularFile == true else { continue }
let startedAt = self.parseSessionStartFromFilename(item.lastPathComponent)
let modifiedAt = values?.contentModificationDate
if self
.shouldIncludeFile(startedAt: startedAt, modifiedAt: modifiedAt, startCutoffLocal: startCutoffLocal)
{
output.append(item)
}
}
return output.sorted(by: { $0.path < $1.path })
}
private static func shouldIncludeFile(
startedAt: Date?,
modifiedAt: Date?,
startCutoffLocal: Date) -> Bool
{
if let modifiedAt, self.localMidnight(modifiedAt) >= startCutoffLocal {
return true
}
if let startedAt, self.localMidnight(startedAt) >= startCutoffLocal {
return true
}
return false
}
private static func scanPiSessionFile(
fileURL: URL,
cache: inout PiSessionCostCache,
context: ScanContext)
throws
{
try context.checkCancellation?()
let path = fileURL.path
let attrs = (try? FileManager.default.attributesOfItem(atPath: path)) ?? [:]
let mtime = (attrs[.modificationDate] as? Date)?.timeIntervalSince1970 ?? 0
let size = (attrs[.size] as? NSNumber)?.int64Value ?? 0
let mtimeMs = Int64(mtime * 1000)
func storeFileUsage(_ usage: PiSessionFileUsage) {
cache.files[path] = usage
}
let cached = cache.files[path]
if !context.forceRescan,
let cached,
cached.mtimeUnixMs == mtimeMs,
cached.size == size
{
return
}
if !context.forceRescan,
let cached,
size > cached.size,
cached.parsedBytes > 0,
cached.parsedBytes <= size
{
let delta = try self.parsePiSessionFile(
fileURL: fileURL,
range: context.range,
startOffset: cached.parsedBytes,
initialModelContext: cached.lastModelContext,
pricingContext: context.pricingContext,
checkCancellation: context.checkCancellation)
if !delta.contributions.isEmpty {
self.applyContributions(
daysByProvider: &cache.daysByProvider,
contributions: delta.contributions,
sign: 1)
}
let merged = self.mergedContributions(existing: cached.contributions, delta: delta.contributions)
storeFileUsage(PiSessionFileUsage(
mtimeUnixMs: mtimeMs,
size: size,
parsedBytes: delta.parsedBytes,
lastModelContext: delta.lastModelContext,
contributions: merged))
return
}
if let cached {
self.applyContributions(
daysByProvider: &cache.daysByProvider,
contributions: cached.contributions,
sign: -1)
}
let parsed = try self.parsePiSessionFile(
fileURL: fileURL,
range: context.range,
pricingContext: context.pricingContext,
checkCancellation: context.checkCancellation)
if !parsed.contributions.isEmpty {
self.applyContributions(daysByProvider: &cache.daysByProvider, contributions: parsed.contributions, sign: 1)
}
storeFileUsage(PiSessionFileUsage(
mtimeUnixMs: mtimeMs,
size: size,
parsedBytes: parsed.parsedBytes,
lastModelContext: parsed.lastModelContext,
contributions: parsed.contributions))
}
private static func parsePiSessionFile(
fileURL: URL,
range: CostUsageScanner.CostUsageDayRange,
startOffset: Int64 = 0,
initialModelContext: PiModelContext? = nil,
pricingContext: ModelsDevPricingContext? = nil,
checkCancellation: CostUsageScanner.CancellationCheck? = nil) throws -> ParseResult
{
var currentModelContext = initialModelContext
var contributions: [String: [String: [String: PiPackedUsage]]] = [:]
func add(provider: UsageProvider, dayKey: String, modelName: String, usage: PiPackedUsage) {
guard !usage.isZero else { return }
guard CostUsageScanner.CostUsageDayRange.isInRange(
dayKey: dayKey,
since: range.scanSinceKey,
until: range.scanUntilKey)
else {
return
}
let providerKey = provider.rawValue
var providerDays = contributions[providerKey] ?? [:]
var dayModels = providerDays[dayKey] ?? [:]
let merged = self.addPacked(a: dayModels[modelName] ?? PiPackedUsage(), b: usage, sign: 1)
if merged.isZero {
dayModels.removeValue(forKey: modelName)
} else {
dayModels[modelName] = merged
}
if dayModels.isEmpty {
providerDays.removeValue(forKey: dayKey)
} else {
providerDays[dayKey] = dayModels
}
if providerDays.isEmpty {
contributions.removeValue(forKey: providerKey)
} else {
contributions[providerKey] = providerDays
}
}
let parsedBytes: Int64
do {
parsedBytes = try CostUsageJsonl.scan(
fileURL: fileURL,
offset: startOffset,
maxLineBytes: Self.maxLineBytes,
prefixBytes: Self.maxLineBytes,
checkCancellation: checkCancellation,
onLine: { line in
guard !line.bytes.isEmpty, !line.wasTruncated else { return }
autoreleasepool {
guard let object = (try? JSONSerialization.jsonObject(with: line.bytes)) as? [String: Any]
else { return }
guard let type = object["type"] as? String else { return }
if type == "model_change" {
currentModelContext = self.modelContext(from: object)
return
}
guard type == "message", let message = object["message"] as? [String: Any] else { return }
guard (message["role"] as? String) == "assistant" else { return }
let identity = self.resolveAssistantIdentity(
entry: object,
message: message,
fallback: currentModelContext)
guard let identity else { return }
guard let date = self.timestampDate(entry: object, message: message) else { return }
let dayKey = CostUsageScanner.CostUsageDayRange.dayKey(from: date)
let usage = self.extractUsage(
provider: identity.provider,
modelName: identity.modelName,
message: message,
pricingDate: date,
pricingContext: pricingContext)
add(provider: identity.provider, dayKey: dayKey, modelName: identity.modelName, usage: usage)
}
})
} catch is CancellationError {
throw CancellationError()
} catch {
parsedBytes = startOffset
}
return ParseResult(
contributions: contributions,
parsedBytes: parsedBytes,
lastModelContext: currentModelContext)
}
private static func modelContext(from object: [String: Any]) -> PiModelContext? {
guard let providerText = object["provider"] as? String,
let provider = self.mappedProvider(fromPiProvider: providerText)
else {
return nil
}
let rawModelName = (object["modelId"] as? String)?.trimmingCharacters(in: .whitespacesAndNewlines) ?? ""
guard let modelName = self.normalizeModelName(rawModelName, provider: provider) else { return nil }
return PiModelContext(providerRawValue: provider.rawValue, modelName: modelName)
}
private static func resolveAssistantIdentity(
entry: [String: Any],
message: [String: Any],
fallback: PiModelContext?) -> AssistantIdentity?
{
let explicitProviderText = self.extractProviderText(entry: entry, message: message)
let explicitProvider = explicitProviderText.flatMap(self.mappedProvider(fromPiProvider:))
let explicitModelText = self.extractModelText(entry: entry, message: message)
if explicitProviderText != nil, explicitProvider == nil {
return nil
}
if let explicitProvider,
let explicitModelText,
let explicitModel = self.normalizeModelName(explicitModelText, provider: explicitProvider)
{
return AssistantIdentity(provider: explicitProvider, modelName: explicitModel)
}
if let explicitProvider,
let fallback,
fallback.providerRawValue == explicitProvider.rawValue
{
return AssistantIdentity(provider: explicitProvider, modelName: fallback.modelName)
}
if explicitProviderText == nil,
let explicitModelText,
let fallbackProvider = fallback.flatMap({ UsageProvider(rawValue: $0.providerRawValue) }),
let explicitModel = self.normalizeModelName(explicitModelText, provider: fallbackProvider)
{
return AssistantIdentity(provider: fallbackProvider, modelName: explicitModel)
}
if explicitProviderText == nil,
let fallback,
let provider = UsageProvider(rawValue: fallback.providerRawValue)
{
return AssistantIdentity(provider: provider, modelName: fallback.modelName)
}
return nil
}
private static func extractProviderText(entry: [String: Any], message: [String: Any]) -> String? {
if let provider = (message["provider"] as? String)?.trimmingCharacters(in: .whitespacesAndNewlines),
!provider.isEmpty
{
return provider
}
if let provider = (entry["provider"] as? String)?.trimmingCharacters(in: .whitespacesAndNewlines),
!provider.isEmpty
{
return provider
}
return nil
}
private static func extractModelText(entry: [String: Any], message: [String: Any]) -> String? {
for value in [message["model"], entry["model"], message["modelId"], entry["modelId"]] {
if let model = (value as? String)?.trimmingCharacters(in: .whitespacesAndNewlines), !model.isEmpty {
return model
}
}
return nil
}
private static func normalizeModelName(_ raw: String, provider: UsageProvider) -> String? {
let trimmed = raw.trimmingCharacters(in: .whitespacesAndNewlines)
guard !trimmed.isEmpty else { return nil }
return switch provider {
case .codex:
CostUsagePricing.normalizeCodexModel(trimmed)
case .claude:
CostUsagePricing.normalizeClaudeModel(trimmed)
default:
trimmed
}
}
private static func timestampDate(entry: [String: Any], message: [String: Any]) -> Date? {
self.parseTimestampValue(message["timestamp"])
?? self.parseTimestampValue(entry["timestamp"])
}
private static func parseTimestampValue(_ value: Any?) -> Date? {
if let number = value as? NSNumber {
let raw = number.doubleValue
guard raw.isFinite else { return nil }
if raw > 1_000_000_000_000 {
return Date(timeIntervalSince1970: raw / 1000)
}
return Date(timeIntervalSince1970: raw)
}
if let string = value as? String {
if let numeric = Double(string), numeric.isFinite {
if numeric > 1_000_000_000_000 {
return Date(timeIntervalSince1970: numeric / 1000)
}
return Date(timeIntervalSince1970: numeric)
}
return self.parseISO(string)
}
return nil
}
private static func extractUsage(
provider: UsageProvider,
modelName: String,
message: [String: Any],
pricingDate: Date? = nil,
pricingContext: ModelsDevPricingContext? = nil) -> PiPackedUsage
{
let usage = (message["usage"] as? [String: Any]) ?? [:]
let input = self.readNonNegativeInt(
usage["input"]
?? usage["inputTokens"]
?? usage["input_tokens"]
?? usage["promptTokens"]
?? usage["prompt_tokens"])
let cacheRead = self.readNonNegativeInt(
usage["cacheRead"]
?? usage["cacheReadTokens"]
?? usage["cache_read"]
?? usage["cache_read_tokens"]
?? usage["cacheReadInputTokens"]
?? usage["cache_read_input_tokens"])
let cacheWrite = self.readNonNegativeInt(
usage["cacheWrite"]
?? usage["cacheWriteTokens"]
?? usage["cache_write"]
?? usage["cache_write_tokens"]
?? usage["cacheCreationTokens"]
?? usage["cache_creation_tokens"]
?? usage["cacheCreationInputTokens"]
?? usage["cache_creation_input_tokens"])
let output = self.readNonNegativeInt(
usage["output"]
?? usage["outputTokens"]
?? usage["output_tokens"]
?? usage["completionTokens"]
?? usage["completion_tokens"])
let directTotal = self.readNonNegativeInt(
usage["totalTokens"]
?? usage["total_tokens"]
?? usage["tokenCount"]
?? usage["token_count"]
?? usage["tokens"])
let derivedTotal = input + cacheRead + cacheWrite + output
let totalTokens = max(directTotal, derivedTotal)
let rawUsage = PiPackedUsage(
inputTokens: input,
cacheReadTokens: cacheRead,
cacheWriteTokens: cacheWrite,
outputTokens: output,
totalTokens: totalTokens)
// Pi JSONL does not record Anthropic cache retention, so use Pi's persisted default tariff.
let costUSD = self.computedCostUSD(
provider: provider,
modelName: modelName,
usage: rawUsage,
pricingDate: pricingDate,
pricingContext: pricingContext)
let costNanos = costUSD.map { Int64(($0 * self.costScale).rounded()) } ?? 0
return PiPackedUsage(
inputTokens: rawUsage.inputTokens,
cacheReadTokens: rawUsage.cacheReadTokens,
cacheWriteTokens: rawUsage.cacheWriteTokens,
outputTokens: rawUsage.outputTokens,
totalTokens: rawUsage.totalTokens,
costNanos: costNanos,
costSampleCount: costUSD == nil ? 0 : 1,
usageSampleCount: 1)
}
private static func computedCostUSD(
provider: UsageProvider,
modelName: String,
usage: PiPackedUsage,
pricingDate: Date? = nil,
pricingContext: ModelsDevPricingContext? = nil) -> Double?
{
switch provider {
case .codex:
// Pi records input, cache reads, and cache writes as disjoint counts. Codex pricing
// expects cached/write tokens to be subsets of total input, so reconstruct that total
// here and pass writes separately (1.25x input for GPT-5.6 when rates are known).
CostUsagePricing.codexCostUSD(
model: modelName,
inputTokens: usage.inputTokens + usage.cacheReadTokens + usage.cacheWriteTokens,
cachedInputTokens: usage.cacheReadTokens,
outputTokens: usage.outputTokens,
cacheWriteInputTokens: usage.cacheWriteTokens,
modelsDevCatalog: pricingContext?.catalog,
modelsDevCacheRoot: pricingContext?.cacheRoot)
case .claude:
CostUsagePricing.claudeCostUSD(
model: modelName,
inputTokens: usage.inputTokens,
cacheReadInputTokens: usage.cacheReadTokens,
cacheCreationInputTokens: usage.cacheWriteTokens,
outputTokens: usage.outputTokens,
pricingDate: pricingDate,
modelsDevCatalog: pricingContext?.catalog,
modelsDevCacheRoot: pricingContext?.cacheRoot)
default:
nil
}
}
private static func readNonNegativeInt(_ value: Any?) -> Int {
if let number = value as? NSNumber {
let numeric = number.doubleValue
guard numeric.isFinite, numeric >= 0, numeric <= self.maxSafeRoundedInt else { return 0 }
return Int(numeric.rounded())
}
if let string = value as? String,
let numeric = Double(string),
numeric.isFinite,
numeric >= 0,
numeric <= self.maxSafeRoundedInt
{
return Int(numeric.rounded())
}
return 0
}
}
extension PiSessionCostScanner {
private static func mappedProvider(fromPiProvider provider: String) -> UsageProvider? {
switch provider.lowercased() {
case "openai-codex":
.codex
case "anthropic":
.claude
default:
nil
}
}
private static func buildReport(
provider: UsageProvider,
cache: PiSessionCostCache,
range: CostUsageScanner.CostUsageDayRange,
pricingContext: ModelsDevPricingContext? = nil) -> CostUsageDailyReport
{
guard let providerDays = cache.daysByProvider[provider.rawValue] else {
return CostUsageDailyReport(data: [], summary: nil)
}
let dayKeys = providerDays.keys.sorted().filter {
CostUsageScanner.CostUsageDayRange.isInRange(dayKey: $0, since: range.sinceKey, until: range.untilKey)
}
var entries: [CostUsageDailyReport.Entry] = []
var totalInput = 0
var totalOutput = 0
var totalCacheRead = 0
var totalCacheWrite = 0
var totalTokens = 0
var totalCostNanos: Int64 = 0
var totalCostSamples = 0
for dayKey in dayKeys {
guard let models = providerDays[dayKey] else { continue }
let modelNames = models.keys.sorted()
var dayInput = 0
var dayOutput = 0
var dayCacheRead = 0
var dayCacheWrite = 0
var dayTotalTokens = 0
var dayCostNanos: Int64 = 0
var dayCostSamples = 0
var breakdown: [CostUsageDailyReport.ModelBreakdown] = []
for modelName in modelNames {
let packed = models[modelName] ?? PiPackedUsage()
let modelTotalTokens = max(
packed.totalTokens,
packed.inputTokens + packed.cacheReadTokens + packed.cacheWriteTokens + packed.outputTokens)
let currentPricingCost = self.computedCostUSD(
provider: provider,
modelName: modelName,
usage: packed,
pricingContext: pricingContext)
let usageSampleCount = packed.usageSampleCount
let hasCompleteCachedCost = (usageSampleCount ?? 0) > 0
&& packed.costSampleCount == usageSampleCount
// Cached costs are accumulated per message, which preserves Claude long-context threshold boundaries.
let costNanos = hasCompleteCachedCost
? packed.costNanos
: currentPricingCost.map { Int64(($0 * self.costScale).rounded()) }
breakdown.append(CostUsageDailyReport.ModelBreakdown(
modelName: modelName,
costUSD: costNanos.map { Double($0) / Self.costScale },
totalTokens: modelTotalTokens > 0 ? modelTotalTokens : nil))
dayInput += packed.inputTokens
dayOutput += packed.outputTokens
dayCacheRead += packed.cacheReadTokens
dayCacheWrite += packed.cacheWriteTokens
dayTotalTokens += modelTotalTokens
if let costNanos {
dayCostNanos += costNanos
dayCostSamples += 1
}
}
let sortedBreakdown = self.sortedModelBreakdowns(breakdown)
entries.append(CostUsageDailyReport.Entry(
date: dayKey,
inputTokens: dayInput > 0 ? dayInput : nil,
outputTokens: dayOutput > 0 ? dayOutput : nil,
cacheReadTokens: dayCacheRead > 0 ? dayCacheRead : nil,
cacheCreationTokens: dayCacheWrite > 0 ? dayCacheWrite : nil,
totalTokens: dayTotalTokens > 0 ? dayTotalTokens : nil,
costUSD: dayCostSamples > 0 ? Double(dayCostNanos) / Self.costScale : nil,
modelsUsed: modelNames,
modelBreakdowns: sortedBreakdown))
totalInput += dayInput
totalOutput += dayOutput
totalCacheRead += dayCacheRead
totalCacheWrite += dayCacheWrite
totalTokens += dayTotalTokens
totalCostNanos += dayCostNanos
totalCostSamples += dayCostSamples
}
guard !entries.isEmpty else { return CostUsageDailyReport(data: [], summary: nil) }
return CostUsageDailyReport(
data: entries,
summary: CostUsageDailyReport.Summary(
totalInputTokens: totalInput > 0 ? totalInput : nil,
totalOutputTokens: totalOutput > 0 ? totalOutput : nil,
cacheReadTokens: totalCacheRead > 0 ? totalCacheRead : nil,
cacheCreationTokens: totalCacheWrite > 0 ? totalCacheWrite : nil,
totalTokens: totalTokens > 0 ? totalTokens : nil,
totalCostUSD: totalCostSamples > 0 ? Double(totalCostNanos) / Self.costScale : nil))
}
private static func mergedContributions(
existing: [String: [String: [String: PiPackedUsage]]],
delta: [String: [String: [String: PiPackedUsage]]]) -> [String: [String: [String: PiPackedUsage]]]
{
var merged = existing
self.applyContributions(daysByProvider: &merged, contributions: delta, sign: 1)
return merged
}
private static func applyContributions(
daysByProvider: inout [String: [String: [String: PiPackedUsage]]],
contributions: [String: [String: [String: PiPackedUsage]]],
sign: Int)
{
for (providerKey, providerDays) in contributions {
var mergedProviderDays = daysByProvider[providerKey] ?? [:]
for (dayKey, dayModels) in providerDays {
var mergedDayModels = mergedProviderDays[dayKey] ?? [:]
for (modelName, packed) in dayModels {
let updated = self.addPacked(
a: mergedDayModels[modelName] ?? PiPackedUsage(),
b: packed,
sign: sign)
if updated.isZero {
mergedDayModels.removeValue(forKey: modelName)
} else {
mergedDayModels[modelName] = updated
}
}
if mergedDayModels.isEmpty {
mergedProviderDays.removeValue(forKey: dayKey)
} else {
mergedProviderDays[dayKey] = mergedDayModels
}
}
if mergedProviderDays.isEmpty {
daysByProvider.removeValue(forKey: providerKey)
} else {
daysByProvider[providerKey] = mergedProviderDays
}
}
}
private static func addPacked(a: PiPackedUsage, b: PiPackedUsage, sign: Int) -> PiPackedUsage {
let aUsageSampleCount = a.usageSampleCount ?? (a.isZero ? 0 : nil)
let bUsageSampleCount = b.usageSampleCount ?? (b.isZero ? 0 : nil)
let usageSampleCount: Int? = if let aCount = aUsageSampleCount, let bCount = bUsageSampleCount {
max(0, aCount + sign * bCount)
} else {
nil
}
return PiPackedUsage(
inputTokens: max(0, a.inputTokens + sign * b.inputTokens),
cacheReadTokens: max(0, a.cacheReadTokens + sign * b.cacheReadTokens),
cacheWriteTokens: max(0, a.cacheWriteTokens + sign * b.cacheWriteTokens),
outputTokens: max(0, a.outputTokens + sign * b.outputTokens),
totalTokens: max(0, a.totalTokens + sign * b.totalTokens),
costNanos: max(0, a.costNanos + Int64(sign) * b.costNanos),
costSampleCount: max(0, a.costSampleCount + sign * b.costSampleCount),
usageSampleCount: usageSampleCount)
}
private static func parseSessionStartFromFilename(_ filename: String) -> Date? {
guard let regex = self.sessionStartFilenameRegex else { return nil }
let range = NSRange(filename.startIndex..<filename.endIndex, in: filename)
guard let match = regex.firstMatch(in: filename, range: range) else { return nil }
guard (1...5).allSatisfy({ Range(match.range(at: $0), in: filename) != nil }) else { return nil }
let date = String(filename[Range(match.range(at: 1), in: filename)!])
let hour = String(filename[Range(match.range(at: 2), in: filename)!])
let minute = String(filename[Range(match.range(at: 3), in: filename)!])
let second = String(filename[Range(match.range(at: 4), in: filename)!])
let millis = String(filename[Range(match.range(at: 5), in: filename)!])
return self.parseISO("\(date)T\(hour):\(minute):\(second).\(millis)Z")
}
private static func parseISO(_ text: String) -> Date? {
self.isoFormatterBox.lock.lock()
defer { self.isoFormatterBox.lock.unlock() }
return self.isoFormatterBox.withFractional.date(from: text)
?? self.isoFormatterBox.plain.date(from: text)
}
private static func localMidnight(_ date: Date) -> Date {
let components = Calendar.current.dateComponents([.year, .month, .day], from: date)
return Calendar.current.date(from: components) ?? date
}
private static func dateFromDayKey(_ key: String) -> Date? {
let parts = key.split(separator: "-")
guard parts.count == 3,
let year = Int(parts[0]),
let month = Int(parts[1]),
let day = Int(parts[2]) else { return nil }
var components = DateComponents()
components.calendar = Calendar.current
components.timeZone = TimeZone.current
components.year = year
components.month = month
components.day = day
components.hour = 0
return components.date
}
private static func sortedModelBreakdowns(_ breakdowns: [CostUsageDailyReport.ModelBreakdown])
-> [CostUsageDailyReport.ModelBreakdown]
{
breakdowns.sorted { lhs, rhs in
let lhsCost = lhs.costUSD ?? -1
let rhsCost = rhs.costUSD ?? -1
if lhsCost != rhsCost {
return lhsCost > rhsCost
}
let lhsTokens = lhs.totalTokens ?? -1
let rhsTokens = rhs.totalTokens ?? -1
if lhsTokens != rhsTokens {
return lhsTokens > rhsTokens
}
return lhs.modelName > rhs.modelName
}
}
}
@@ -0,0 +1,55 @@
import Foundation
/// Stable identity for one account row surfaced by a multi-account source adapter.
///
/// `source` names the adapter (for example `claude-swap`) and `opaqueID` is the
/// source-issued identifier (for example a numeric slot). Identity never derives
/// from emails or credential material, per
/// `docs/claude-multi-account-and-status-items.md`.
public struct ProviderAccountIdentity: Hashable, Sendable {
public let source: String
public let opaqueID: String
public init(source: String, opaqueID: String) {
self.source = source
self.opaqueID = opaqueID
}
}
/// Provider-neutral projection of one account's usage, consumed by menus (and,
/// later, per-account status items) without teaching UI code about any specific
/// credential source.
public struct ProviderAccountUsageSnapshot: Identifiable, Sendable {
public let id: ProviderAccountIdentity
public let provider: UsageProvider
/// Display-only label (may contain personal data such as an email); UI is
/// responsible for privacy redaction. Never logged or persisted.
public let displayLabel: String
public let isActive: Bool
/// Whether the source can make this inactive account the provider's active account.
/// Activation remains source-owned; CodexBar never handles credential material.
public let canActivate: Bool
public let snapshot: UsageSnapshot?
public let error: String?
public let sourceLabel: String?
public init(
id: ProviderAccountIdentity,
provider: UsageProvider,
displayLabel: String,
isActive: Bool,
canActivate: Bool = false,
snapshot: UsageSnapshot?,
error: String?,
sourceLabel: String?)
{
self.id = id
self.provider = provider
self.displayLabel = displayLabel
self.isActive = isActive
self.canActivate = canActivate
self.snapshot = snapshot
self.error = error
self.sourceLabel = sourceLabel
}
}
@@ -0,0 +1,38 @@
import Foundation
/// Provider-specific spend/budget snapshot (e.g. Claude "Extra usage" monthly spend vs limit).
public struct ProviderCostSnapshot: Equatable, Codable, Sendable {
public let used: Double
public let limit: Double
public let currencyCode: String
/// Human-friendly period label (e.g. "Monthly"). Optional; some providers don't expose a period.
public let period: String?
/// Optional renewal/reset timestamp for the period.
public let resetsAt: Date?
/// Optional amount restored on the next regeneration tick for providers with rolling credit recovery.
public let nextRegenAmount: Double?
/// This account's own contribution when `used`/`limit` describe a shared/pooled budget
/// (e.g. Cursor team on-demand pool). nil when the budget is already personal.
public let personalUsed: Double?
public let updatedAt: Date
public init(
used: Double,
limit: Double,
currencyCode: String,
period: String? = nil,
resetsAt: Date? = nil,
nextRegenAmount: Double? = nil,
personalUsed: Double? = nil,
updatedAt: Date)
{
self.used = used
self.limit = limit
self.currencyCode = currencyCode
self.period = period
self.resetsAt = resetsAt
self.nextRegenAmount = nextRegenAmount
self.personalUsed = personalUsed
self.updatedAt = updatedAt
}
}
@@ -0,0 +1,170 @@
import Foundation
enum ProviderEndpointOverrideError: LocalizedError, Equatable {
case minimax(String)
case alibabaCodingPlan(String)
var errorDescription: String? {
switch self {
case let .minimax(key):
"MiniMax endpoint override \(key) is not allowed. " +
"Use an HTTPS endpoint without user info or encoded host tricks. " +
"If MINIMAX_REQUIRE_PROVIDER_ENDPOINT_OVERRIDES=true is set, the endpoint must also be MiniMax-owned."
case let .alibabaCodingPlan(key):
"Alibaba Coding Plan endpoint override \(key) is not allowed. " +
"Use an HTTPS endpoint without user info or encoded host tricks. " +
"If ALIBABA_CODING_PLAN_REQUIRE_PROVIDER_ENDPOINT_OVERRIDES=true is set, " +
"the endpoint must also be Alibaba-owned."
}
}
}
struct ProviderEndpointOverrideValidator {
enum HostPolicy {
case allowAnyHTTPSHost
case providerOwnedOnly
}
private let allowedHosts: Set<String>
private let allowedDomainSuffixes: Set<String>
init(allowedHosts: [String] = [], allowedDomainSuffixes: [String] = []) {
self.allowedHosts = Set(allowedHosts.map { $0.lowercased() })
self.allowedDomainSuffixes = Set(allowedDomainSuffixes.map { $0.lowercased() })
}
func validatedHost(_ raw: String?, policy: HostPolicy = .allowAnyHTTPSHost) -> String? {
guard let raw,
let url = self.url(from: raw),
let host = self.validatedDecodedHost(for: url, policy: policy)
else { return nil }
return self.hostAuthority(host: host, port: url.port)
}
func validatedURL(_ raw: String?, policy: HostPolicy = .allowAnyHTTPSHost) -> URL? {
guard let raw,
let url = self.url(from: raw),
self.validatedDecodedHost(for: url, policy: policy) != nil
else { return nil }
return url
}
func validatedURLAllowingLoopbackHTTP(_ raw: String?) -> URL? {
guard let raw,
Self.hasExplicitURLScheme(raw),
let url = URL(string: raw),
let scheme = url.scheme?.lowercased(),
scheme == "https" || scheme == "http",
url.user == nil,
url.password == nil,
let host = self.validatedDecodedHost(for: url, policy: .allowAnyHTTPSHost),
scheme == "https" || Self.isLoopbackHost(host)
else { return nil }
return url
}
static func normalizedHTTPSURL(from raw: String) -> URL? {
let url = if Self.hasExplicitURLScheme(raw) {
URL(string: raw)
} else {
URL(string: "https://\(raw)")
}
guard let url else { return nil }
guard let scheme = url.scheme?.lowercased(), scheme == "https" else { return nil }
guard url.user == nil, url.password == nil else { return nil }
guard let decodedHost = url.host(percentEncoded: false)?.lowercased(),
!decodedHost.isEmpty,
!decodedHost.contains("%"),
decodedHost.rangeOfCharacter(from: .whitespacesAndNewlines) == nil,
decodedHost.rangeOfCharacter(from: .controlCharacters) == nil,
let encodedHost = url.host(percentEncoded: true)?.lowercased(),
Self.hostHasNoEncodedDelimiters(encodedHost, decodedHost: decodedHost, url: url)
else { return nil }
return url
}
private func url(from raw: String) -> URL? {
Self.normalizedHTTPSURL(from: raw)
}
private static func hasExplicitURLScheme(_ raw: String) -> Bool {
guard let colonIndex = raw.firstIndex(of: ":") else { return false }
if raw[colonIndex...].hasPrefix("://") { return true }
if let authorityEnd = raw.firstIndex(where: { ["/", "?", "#"].contains($0) }),
colonIndex > authorityEnd
{
return false
}
let afterColon = raw.index(after: colonIndex)
guard afterColon < raw.endIndex else { return true }
let portEnd = raw[afterColon...].firstIndex { Set<Character>(["/", "?", "#"]).contains($0) } ?? raw.endIndex
let suffix = raw[afterColon..<portEnd]
if !suffix.isEmpty, suffix.allSatisfy(\.isNumber) {
return false
}
let scheme = raw[..<colonIndex]
guard let first = scheme.first, first.isLetter else { return false }
return scheme.dropFirst().allSatisfy { $0.isLetter || $0.isNumber || ["+", "-", "."].contains($0) }
}
private static func isLoopbackHost(_ host: String) -> Bool {
if host == "localhost" || host == "::1" { return true }
let octets = host.split(separator: ".", omittingEmptySubsequences: false)
guard octets.count == 4,
let first = UInt8(octets[0]),
octets.dropFirst().allSatisfy({ UInt8($0) != nil })
else { return false }
return first == 127
}
private func hostAuthority(host: String, port: Int?) -> String {
let authorityHost = host.contains(":") ? "[\(host)]" : host
guard let port else { return authorityHost }
return "\(authorityHost):\(port)"
}
private func validatedDecodedHost(for url: URL, policy: HostPolicy) -> String? {
guard let decodedHost = url.host(percentEncoded: false)?.lowercased(),
!decodedHost.isEmpty,
!decodedHost.contains("%"),
decodedHost.rangeOfCharacter(from: .whitespacesAndNewlines) == nil,
decodedHost.rangeOfCharacter(from: .controlCharacters) == nil,
let encodedHost = url.host(percentEncoded: true)?.lowercased(),
Self.hostHasNoEncodedDelimiters(encodedHost, decodedHost: decodedHost, url: url)
else { return nil }
switch policy {
case .allowAnyHTTPSHost:
return decodedHost
case .providerOwnedOnly:
let isAllowedHost = self.allowedHosts.contains(decodedHost)
let isAllowedSuffix = self.allowedDomainSuffixes.contains { suffix in
decodedHost == suffix || decodedHost.hasSuffix(".\(suffix)")
}
guard isAllowedHost || isAllowedSuffix else { return nil }
return decodedHost
}
}
private static func hostHasNoEncodedDelimiters(_ encodedHost: String, decodedHost: String, url: URL) -> Bool {
if decodedHost.contains(":") {
guard encodedHost == decodedHost,
let componentHost = URLComponents(url: url, resolvingAgainstBaseURL: false)?.host,
componentHost.hasPrefix("["),
componentHost.hasSuffix("]")
else { return false }
let address = componentHost.dropFirst().dropLast()
return !address.isEmpty && address.allSatisfy { $0.isHexDigit || $0 == ":" || $0 == "." }
}
let decodedDelimiters = CharacterSet(charactersIn: "/\\?#@:")
guard decodedHost.rangeOfCharacter(from: decodedDelimiters) == nil else { return false }
let encodedDelimiters = ["%2f", "%5c", "%3f", "%23", "%40", "%3a"]
return !encodedDelimiters.contains { encodedHost.contains($0) }
}
}
@@ -0,0 +1,253 @@
import Foundation
#if canImport(FoundationNetworking)
import FoundationNetworking
#endif
public protocol ProviderHTTPTransport: Sendable {
func data(for request: URLRequest) async throws -> (Data, URLResponse)
}
#if !os(Linux)
extension URLSession: ProviderHTTPTransport {}
#endif
extension URLSession {
public func response(for request: URLRequest) async throws -> ProviderHTTPResponse {
let (data, response) = try await self.data(for: request)
guard let httpResponse = response as? HTTPURLResponse else {
throw URLError(.badServerResponse)
}
return ProviderHTTPResponse(data: data, response: httpResponse)
}
}
public struct ProviderHTTPResponse: Sendable {
public let data: Data
public let response: HTTPURLResponse
public init(data: Data, response: HTTPURLResponse) {
self.data = data
self.response = response
}
public var statusCode: Int {
self.response.statusCode
}
}
public struct ProviderHTTPRetryPolicy: Sendable {
public let maxRetries: Int
public let retryableStatusCodes: Set<Int>
public let retryableURLErrorCodes: Set<URLError.Code>
public let retryableMethods: Set<String>
public let baseDelaySeconds: TimeInterval
public let maxDelaySeconds: TimeInterval
public init(
maxRetries: Int,
retryableStatusCodes: Set<Int> = [408, 429, 500, 502, 503, 504],
retryableURLErrorCodes: Set<URLError.Code> = [
.timedOut,
.networkConnectionLost,
.cannotConnectToHost,
.cannotFindHost,
.dnsLookupFailed,
],
retryableMethods: Set<String> = ["GET", "HEAD", "OPTIONS"],
baseDelaySeconds: TimeInterval = 1,
maxDelaySeconds: TimeInterval = 10)
{
self.maxRetries = max(0, maxRetries)
self.retryableStatusCodes = retryableStatusCodes
self.retryableURLErrorCodes = retryableURLErrorCodes
self.retryableMethods = retryableMethods
self.baseDelaySeconds = max(0, baseDelaySeconds)
self.maxDelaySeconds = max(0, maxDelaySeconds)
}
public static let disabled = ProviderHTTPRetryPolicy(
maxRetries: 0,
retryableStatusCodes: [],
retryableURLErrorCodes: [],
baseDelaySeconds: 0,
maxDelaySeconds: 0)
public static let transientIdempotent = ProviderHTTPRetryPolicy(maxRetries: 1)
func shouldRetry(request: URLRequest, attempt: Int, statusCode: Int) -> Bool {
self.canRetry(request: request, attempt: attempt)
&& self.retryableStatusCodes.contains(statusCode)
}
func shouldRetry(request: URLRequest, attempt: Int, error: Error) -> Bool {
guard self.canRetry(request: request, attempt: attempt) else { return false }
guard let urlError = error as? URLError else { return false }
return self.retryableURLErrorCodes.contains(urlError.code)
}
func delaySeconds(attempt: Int, response: HTTPURLResponse?) -> TimeInterval {
if let retryAfter = response?.value(forHTTPHeaderField: "Retry-After"),
let seconds = TimeInterval(retryAfter.trimmingCharacters(in: .whitespacesAndNewlines)),
seconds >= 0
{
return min(seconds, self.maxDelaySeconds)
}
guard self.baseDelaySeconds > 0 else { return 0 }
let multiplier = pow(2, Double(max(0, attempt)))
return min(self.baseDelaySeconds * multiplier, self.maxDelaySeconds)
}
private func canRetry(request: URLRequest, attempt: Int) -> Bool {
guard attempt < self.maxRetries else { return false }
let method = request.httpMethod?.uppercased() ?? "GET"
return self.retryableMethods.contains(method)
}
}
public struct ProviderHTTPTransportHandler: ProviderHTTPTransport {
private let handler: @Sendable (URLRequest) async throws -> (Data, URLResponse)
public init(_ handler: @escaping @Sendable (URLRequest) async throws -> (Data, URLResponse)) {
self.handler = handler
}
public func data(for request: URLRequest) async throws -> (Data, URLResponse) {
try await self.handler(request)
}
}
extension ProviderHTTPTransport {
public func response(for request: URLRequest) async throws -> ProviderHTTPResponse {
try await self.response(for: request, retryPolicy: .disabled)
}
public func response(
for request: URLRequest,
retryPolicy: ProviderHTTPRetryPolicy) async throws -> ProviderHTTPResponse
{
var attempt = 0
while true {
do {
let (data, response) = try await self.data(for: request)
guard let httpResponse = response as? HTTPURLResponse else {
throw URLError(.badServerResponse)
}
let providerResponse = ProviderHTTPResponse(data: data, response: httpResponse)
guard retryPolicy.shouldRetry(
request: request,
attempt: attempt,
statusCode: providerResponse.statusCode)
else {
return providerResponse
}
try await Self.sleepBeforeRetry(policy: retryPolicy, attempt: attempt, response: httpResponse)
attempt += 1
} catch {
guard retryPolicy.shouldRetry(request: request, attempt: attempt, error: error) else {
throw error
}
try await Self.sleepBeforeRetry(policy: retryPolicy, attempt: attempt, response: nil)
attempt += 1
}
}
}
private static func sleepBeforeRetry(
policy: ProviderHTTPRetryPolicy,
attempt: Int,
response: HTTPURLResponse?) async throws
{
let delay = policy.delaySeconds(attempt: attempt, response: response)
guard delay > 0 else { return }
try await Task.sleep(nanoseconds: UInt64(delay * 1_000_000_000))
}
}
public final class ProviderHTTPClient: ProviderHTTPTransport, @unchecked Sendable {
public static let shared = ProviderHTTPClient(session: ProviderHTTPClient.sharedSession())
private let session: URLSession
public init(session: URLSession? = nil) {
self.session = session ?? Self.redirectGuardedSession()
}
static func defaultConfiguration() -> URLSessionConfiguration {
let configuration = URLSessionConfiguration.default
configuration.timeoutIntervalForRequest = 30
configuration.timeoutIntervalForResource = 90
#if !os(Linux)
configuration.waitsForConnectivity = false
#endif
return configuration
}
private static func sharedSession() -> URLSession {
if self.isRunningTests {
// XCTest URLProtocol.registerClass stubs only intercept URLSession.shared on macOS.
return .shared
}
return self.redirectGuardedSession()
}
static func redirectGuardedSession(
configuration: URLSessionConfiguration = ProviderHTTPClient.defaultConfiguration()) -> URLSession
{
URLSession(
configuration: configuration,
delegate: ProviderHTTPRedirectGuardDelegate(),
delegateQueue: nil)
}
private static var isRunningTests: Bool {
let environment = ProcessInfo.processInfo.environment
if environment["XCTestConfigurationFilePath"] != nil || environment["XCTestBundlePath"] != nil {
return true
}
if ProcessInfo.processInfo.processName.lowercased().contains("xctest") {
return true
}
return CommandLine.arguments.contains { $0.lowercased().contains(".xctest") }
}
public func data(for request: URLRequest) async throws -> (Data, URLResponse) {
try await self.session.data(for: request)
}
}
final class ProviderHTTPRedirectGuardDelegate: NSObject, URLSessionTaskDelegate, @unchecked Sendable {
func urlSession(
_: URLSession,
task: URLSessionTask,
willPerformHTTPRedirection _: HTTPURLResponse,
newRequest request: URLRequest,
completionHandler: @escaping @Sendable (URLRequest?) -> Void)
{
completionHandler(Self.guardedRedirectRequest(originalURL: task.originalRequest?.url, redirectRequest: request))
}
static func guardedRedirectRequest(originalURL: URL?, redirectRequest request: URLRequest) -> URLRequest? {
guard let originalURL, let redirectedURL = request.url else { return nil }
guard originalURL.scheme?.caseInsensitiveCompare("https") == .orderedSame else { return nil }
guard redirectedURL.scheme?.caseInsensitiveCompare("https") == .orderedSame else { return nil }
guard self.isSameOrigin(originalURL, redirectedURL) else { return nil }
return request
}
private static func isSameOrigin(_ lhs: URL, _ rhs: URL) -> Bool {
lhs.scheme?.lowercased() == rhs.scheme?.lowercased()
&& lhs.host?.lowercased() == rhs.host?.lowercased()
&& self.normalizedPort(lhs) == self.normalizedPort(rhs)
}
private static func normalizedPort(_ url: URL) -> Int? {
if let port = url.port { return port }
switch url.scheme?.lowercased() {
case "http": return 80
case "https": return 443
default: return nil
}
}
}
@@ -0,0 +1,543 @@
import Foundation
public struct ProviderStorageFootprint: Sendable, Equatable {
public struct Component: Sendable, Equatable, Identifiable {
public let id: String
public let path: String
public let totalBytes: Int64
public init(path: String, totalBytes: Int64) {
self.id = path
self.path = path
self.totalBytes = totalBytes
}
public var name: String {
let url = URL(fileURLWithPath: self.path)
let last = url.lastPathComponent
if last.isEmpty { return self.path }
return last
}
}
public let provider: UsageProvider
public let totalBytes: Int64
public let paths: [String]
public let missingPaths: [String]
public let unreadablePaths: [String]
public let components: [Component]
public let updatedAt: Date
public init(
provider: UsageProvider,
totalBytes: Int64,
paths: [String],
missingPaths: [String],
unreadablePaths: [String],
components: [Component] = [],
updatedAt: Date)
{
self.provider = provider
self.totalBytes = totalBytes
self.paths = paths
self.missingPaths = missingPaths
self.unreadablePaths = unreadablePaths
self.components = components
self.updatedAt = updatedAt
}
public var hasLocalData: Bool {
self.totalBytes > 0
}
/// Value equality that ignores `updatedAt`. Two scans of identical on-disk data differ only by
/// their scan timestamp, so callers use this to avoid re-publishing observable state (and the
/// menu-invalidation churn that follows) when nothing the user sees has actually changed.
public func hasSameContents(as other: ProviderStorageFootprint) -> Bool {
self.provider == other.provider &&
self.totalBytes == other.totalBytes &&
self.paths == other.paths &&
self.missingPaths == other.missingPaths &&
self.unreadablePaths == other.unreadablePaths &&
self.components == other.components
}
public var cleanupRecommendations: [ProviderStorageRecommendation] {
ProviderStorageRecommendation.recommendations(for: self)
}
public func replacingProvider(_ provider: UsageProvider) -> ProviderStorageFootprint {
ProviderStorageFootprint(
provider: provider,
totalBytes: self.totalBytes,
paths: self.paths,
missingPaths: self.missingPaths,
unreadablePaths: self.unreadablePaths,
components: self.components,
updatedAt: self.updatedAt)
}
}
public struct ProviderStorageRecommendation: Sendable, Equatable, Identifiable {
public enum RiskLevel: String, Sendable {
case informational
case manualCleanup
}
public let id: String
public let provider: UsageProvider
public let path: String
public let bytes: Int64
public let title: String
public let riskLevel: RiskLevel
public let consequence: String
public let sortPriority: Int
public init(
provider: UsageProvider,
path: String,
bytes: Int64,
title: String,
riskLevel: RiskLevel,
consequence: String,
sortPriority: Int)
{
self.id = path
self.provider = provider
self.path = path
self.bytes = bytes
self.title = title
self.riskLevel = riskLevel
self.consequence = consequence
self.sortPriority = sortPriority
}
public static func recommendations(for footprint: ProviderStorageFootprint) -> [ProviderStorageRecommendation] {
let candidates: [ProviderStorageRecommendation] = footprint.components.compactMap { component in
switch footprint.provider {
case .claude:
self.claudeRecommendation(for: component)
case .codex:
self.codexRecommendation(for: component, roots: footprint.paths)
default:
nil
}
}
return candidates.sorted { lhs, rhs in
if lhs.sortPriority == rhs.sortPriority {
if lhs.bytes == rhs.bytes {
return lhs.path.localizedCaseInsensitiveCompare(rhs.path) == .orderedAscending
}
return lhs.bytes > rhs.bytes
}
return lhs.sortPriority < rhs.sortPriority
}
}
private static func claudeRecommendation(
for component: ProviderStorageFootprint.Component)
-> ProviderStorageRecommendation?
{
switch component.name {
case "projects":
self.make(
provider: .claude,
component: component,
title: "Manual cleanup: past sessions",
consequence: "Clearing removes past resume, continue, and rewind history.",
priority: 10)
case "file-history":
self.make(
provider: .claude,
component: component,
title: "Manual cleanup: file checkpoints",
consequence: "Clearing removes checkpoint restore data for previous edits.",
priority: 20)
case "plans":
self.make(
provider: .claude,
component: component,
title: "Manual cleanup: saved plans",
consequence: "Clearing removes old plan-mode files.",
priority: 30)
case "debug":
self.make(
provider: .claude,
component: component,
title: "Manual cleanup: debug logs",
consequence: "Clearing removes past debug logs.",
priority: 40)
case "paste-cache", "image-cache":
self.make(
provider: .claude,
component: component,
title: "Manual cleanup: attachment cache",
consequence: "Clearing removes cached large pastes or attached images.",
priority: 50)
case "session-env":
self.make(
provider: .claude,
component: component,
title: "Manual cleanup: session metadata",
consequence: "Clearing removes per-session environment metadata.",
priority: 60)
case "shell-snapshots":
self.make(
provider: .claude,
component: component,
title: "Manual cleanup: shell snapshots",
consequence: "Clearing removes leftover runtime shell snapshot files.",
priority: 70)
case "todos":
self.make(
provider: .claude,
component: component,
title: "Manual cleanup: legacy todos",
consequence: "Clearing removes legacy per-session task lists.",
priority: 80)
default:
nil
}
}
private static func codexRecommendation(
for component: ProviderStorageFootprint.Component,
roots: [String])
-> ProviderStorageRecommendation?
{
guard self.path(component.path, isContainedIn: roots) else { return nil }
return switch component.name {
case "sessions":
self.make(
provider: .codex,
component: component,
title: "Manual cleanup: sessions",
consequence: "Clearing removes past Codex session history.",
priority: 10)
case "archived_sessions":
self.make(
provider: .codex,
component: component,
title: "Manual cleanup: archived sessions",
consequence: "Clearing removes archived Codex session history.",
priority: 20)
case "cache", "caches", "Cache", "Caches":
self.make(
provider: .codex,
component: component,
title: "Manual cleanup: cache",
consequence: "Clearing removes provider-owned cached data.",
priority: 30)
case "log", "logs", "debug":
self.make(
provider: .codex,
component: component,
title: "Manual cleanup: logs",
consequence: "Clearing removes local diagnostic logs.",
priority: 40)
case let name where name.hasPrefix("logs_") && name.hasSuffix(".sqlite"):
self.make(
provider: .codex,
component: component,
title: "Manual cleanup: logs",
consequence: "Clearing removes local diagnostic logs.",
priority: 40)
case "file-history":
self.make(
provider: .codex,
component: component,
title: "Manual cleanup: file history",
consequence: "Clearing removes local edit checkpoint history.",
priority: 50)
case "paste-cache", "image-cache", "session-env", "shell-snapshots", "shell_snapshots", "tmp", "temp", ".tmp":
self.make(
provider: .codex,
component: component,
title: "Manual cleanup: temporary data",
consequence: "Clearing removes local temporary provider data.",
priority: 60)
default:
nil
}
}
private static func make(
provider: UsageProvider,
component: ProviderStorageFootprint.Component,
title: String,
consequence: String,
priority: Int)
-> ProviderStorageRecommendation
{
ProviderStorageRecommendation(
provider: provider,
path: component.path,
bytes: component.totalBytes,
title: title,
riskLevel: .manualCleanup,
consequence: consequence,
sortPriority: priority)
}
private static func path(_ path: String, isContainedIn roots: [String]) -> Bool {
let standardizedPath = URL(fileURLWithPath: path).standardizedFileURL.path
return roots.contains { root in
let standardizedRoot = URL(fileURLWithPath: root, isDirectory: true).standardizedFileURL.path
return standardizedPath == standardizedRoot || standardizedPath.hasPrefix(standardizedRoot + "/")
}
}
}
public enum ProviderStoragePathCatalog {
public static func candidatePaths(
for provider: UsageProvider,
environment: [String: String],
managedCodexAccounts: [ManagedCodexAccount] = [],
fileManager: FileManager = .default)
-> [String]
{
let home = fileManager.homeDirectoryForCurrentUser
func homePath(_ relativePath: String) -> String {
home.appendingPathComponent(relativePath, isDirectory: true).path
}
let candidates: [String] = switch provider {
case .codex:
[CodexHomeScope.ambientHomeURL(env: environment, fileManager: fileManager).path] +
managedCodexAccounts.map(\.managedHomePath)
case .claude:
[
homePath(".claude"),
homePath(".config/claude"),
home
.appendingPathComponent("Library/Application Support/CodexBar/ClaudeProbe", isDirectory: true)
.path,
]
case .gemini:
[
homePath(".gemini"),
homePath(".config/gemini"),
]
case .opencode, .opencodego:
[
homePath(".config/opencode"),
]
case .copilot:
[
homePath(".config/github-copilot"),
]
case .cursor:
[
homePath("Library/Application Support/Cursor"),
homePath("Library/Application Support/Caches/cursor-updater"),
homePath(".cursor"),
homePath("Library/Caches/Cursor"),
homePath("Library/Caches/com.todesktop.230313mzl4w4u92"),
homePath("Library/Caches/com.todesktop.230313mzl4w4u92.ShipIt"),
homePath("Library/Caches/cursor-compile-cache"),
homePath("Library/HTTPStorages/com.todesktop.230313mzl4w4u92"),
]
default:
[]
}
return Self.uniqueStandardizedPaths(candidates)
}
private static func uniqueStandardizedPaths(_ paths: [String]) -> [String] {
var seen: Set<String> = []
var result: [String] = []
for path in paths {
let trimmed = path.trimmingCharacters(in: .whitespacesAndNewlines)
guard !trimmed.isEmpty else { continue }
let standardized = URL(fileURLWithPath: trimmed, isDirectory: true).standardizedFileURL.path
guard seen.insert(standardized).inserted else { continue }
result.append(standardized)
}
return result
}
}
public struct ProviderStorageScanner: @unchecked Sendable {
private struct DirectoryScanResult {
var bytes: Int64 = 0
var unreadablePaths: [String] = []
var componentBytes: [String: Int64] = [:]
}
private let fileManager: FileManager
public init(fileManager: FileManager = .default) {
self.fileManager = fileManager
}
public func scan(
provider: UsageProvider,
candidatePaths: [String],
now: Date = Date())
-> ProviderStorageFootprint
{
var totalBytes: Int64 = 0
var existingPaths: [String] = []
var missingPaths: [String] = []
var unreadablePaths: [String] = []
var components: [ProviderStorageFootprint.Component] = []
for path in candidatePaths {
if Task.isCancelled { break }
var isDirectory: ObjCBool = false
guard self.fileManager.fileExists(atPath: path, isDirectory: &isDirectory) else {
missingPaths.append(path)
continue
}
existingPaths.append(path)
let url = URL(fileURLWithPath: path, isDirectory: isDirectory.boolValue)
if self.isSymbolicLink(at: url) {
continue
}
if isDirectory.boolValue {
let result = self.scanDirectory(at: url)
if Task.isCancelled { break }
totalBytes += result.bytes
unreadablePaths.append(contentsOf: result.unreadablePaths)
components.append(contentsOf: result.componentBytes.map {
ProviderStorageFootprint.Component(path: $0.key, totalBytes: $0.value)
})
} else {
let result = self.sizeOfFile(at: url)
totalBytes += result.bytes
unreadablePaths.append(contentsOf: result.unreadablePaths)
if result.bytes > 0 {
components.append(.init(path: url.path, totalBytes: result.bytes))
}
}
}
return ProviderStorageFootprint(
provider: provider,
totalBytes: totalBytes,
paths: existingPaths,
missingPaths: missingPaths,
unreadablePaths: unreadablePaths,
components: components.sorted { lhs, rhs in
if lhs.totalBytes == rhs.totalBytes {
return lhs.path.localizedCaseInsensitiveCompare(rhs.path) == .orderedAscending
}
return lhs.totalBytes > rhs.totalBytes
},
updatedAt: now)
}
private func isSymbolicLink(at url: URL) -> Bool {
(try? url.resourceValues(forKeys: [.isSymbolicLinkKey]).isSymbolicLink) == true
}
private func sizeOfFile(at url: URL) -> (bytes: Int64, unreadablePaths: [String]) {
if Task.isCancelled { return (0, []) }
let keys: Set<URLResourceKey> = [
.isRegularFileKey,
.isSymbolicLinkKey,
.fileSizeKey,
]
guard let values = try? url.resourceValues(forKeys: keys) else {
return (0, [url.path])
}
if values.isSymbolicLink == true {
return (0, [])
}
if values.isRegularFile == true {
return (Int64(values.fileSize ?? 0), [])
}
return (0, [])
}
private func scanDirectory(at url: URL) -> DirectoryScanResult {
if Task.isCancelled { return DirectoryScanResult() }
let keys: Set<URLResourceKey> = [
.isDirectoryKey,
.isRegularFileKey,
.isSymbolicLinkKey,
.fileSizeKey,
]
let unreadableCollector = ProviderStorageUnreadablePathCollector()
guard let enumerator = self.fileManager.enumerator(
at: url,
includingPropertiesForKeys: Array(keys),
options: [.skipsPackageDescendants],
errorHandler: { url, _ in
unreadableCollector.append(url.path)
return true
})
else {
return DirectoryScanResult(unreadablePaths: [url.path])
}
var result = DirectoryScanResult()
let rootPath = url.standardizedFileURL.path
for case let itemURL as URL in enumerator {
if Task.isCancelled {
enumerator.skipDescendants()
break
}
guard let itemValues = try? itemURL.resourceValues(forKeys: keys) else {
unreadableCollector.append(itemURL.path)
continue
}
if itemValues.isSymbolicLink == true {
if itemValues.isDirectory == true {
enumerator.skipDescendants()
}
continue
}
if itemValues.isRegularFile == true {
let bytes = Int64(itemValues.fileSize ?? 0)
result.bytes += bytes
if bytes > 0, let componentPath = self.topLevelComponentPath(for: itemURL, rootPath: rootPath) {
result.componentBytes[componentPath, default: 0] += bytes
}
}
}
result.unreadablePaths = unreadableCollector.paths
return result
}
private func topLevelComponentPath(for url: URL, rootPath: String) -> String? {
let itemPath = url.standardizedFileURL.path
let pathPrefix = rootPath.hasSuffix("/") ? rootPath : "\(rootPath)/"
guard itemPath.hasPrefix(pathPrefix) else { return nil }
let suffix = itemPath.dropFirst(pathPrefix.count)
let relative = suffix.drop { $0 == "/" }
guard let first = relative.split(separator: "/", maxSplits: 1, omittingEmptySubsequences: true).first else {
return nil
}
return URL(fileURLWithPath: rootPath, isDirectory: true)
.appendingPathComponent(String(first))
.path
}
}
private final class ProviderStorageUnreadablePathCollector: @unchecked Sendable {
private let lock = NSLock()
private var storage: [String] = []
var paths: [String] {
self.lock.lock()
defer { self.lock.unlock() }
return self.storage
}
func append(_ path: String) {
self.lock.lock()
defer { self.lock.unlock() }
self.storage.append(path)
}
}
@@ -0,0 +1,72 @@
import Foundation
public struct APITokenFetchStrategy: ProviderFetchStrategy {
public typealias TokenResolver = @Sendable ([String: String]) -> String?
public typealias MissingCredentialsError = @Sendable () -> (any Error & Sendable)
public typealias UsageLoader = @Sendable (String, ProviderFetchContext) async throws -> UsageSnapshot
public let id: String
public let kind: ProviderFetchKind = .apiToken
private let sourceLabel: String
private let reportsMissingCredentials: Bool
private let resolveToken: TokenResolver
private let missingCredentialsError: MissingCredentialsError
private let loadUsage: UsageLoader
public init(
id: String,
sourceLabel: String = "api",
reportsMissingCredentials: Bool = false,
resolveToken: @escaping TokenResolver,
missingCredentialsError: @escaping MissingCredentialsError,
loadUsage: @escaping UsageLoader)
{
self.id = id
self.sourceLabel = sourceLabel
self.reportsMissingCredentials = reportsMissingCredentials
self.resolveToken = resolveToken
self.missingCredentialsError = missingCredentialsError
self.loadUsage = loadUsage
}
public func isAvailable(_ context: ProviderFetchContext) async -> Bool {
self.reportsMissingCredentials || self.resolveToken(context.env) != nil
}
public func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult {
guard let token = self.resolveToken(context.env) else {
throw self.missingCredentialsError()
}
return try await self.makeResult(
usage: self.loadUsage(token, context),
sourceLabel: self.sourceLabel)
}
public func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool {
false
}
}
extension ProviderFetchPlan {
public static func apiToken(
strategyID: String,
sourceLabel: String = "api",
reportsMissingCredentials: Bool = false,
resolveToken: @escaping APITokenFetchStrategy.TokenResolver,
missingCredentialsError: @escaping APITokenFetchStrategy.MissingCredentialsError,
loadUsage: @escaping APITokenFetchStrategy.UsageLoader) -> ProviderFetchPlan
{
ProviderFetchPlan(
sourceModes: [.auto, .api],
pipeline: ProviderFetchPipeline(resolveStrategies: { _ in
[APITokenFetchStrategy(
id: strategyID,
sourceLabel: sourceLabel,
reportsMissingCredentials: reportsMissingCredentials,
resolveToken: resolveToken,
missingCredentialsError: missingCredentialsError,
loadUsage: loadUsage)]
}))
}
}
@@ -0,0 +1,134 @@
import Foundation
#if os(macOS)
import SweetCookieKit
// MARK: - Abacus Cookie Importer
public enum AbacusCookieImporter {
private static let log = CodexBarLog.logger(LogCategories.abacusCookie)
private static let cookieClient = BrowserCookieClient()
private static let cookieDomains = ["abacus.ai", "apps.abacus.ai"]
private static let cookieImportOrder: BrowserCookieImportOrder =
ProviderDefaults.metadata[.abacus]?.browserCookieOrder ?? Browser.defaultImportOrder
/// Exact cookie names known to carry Abacus session state.
/// CSRF tokens are deliberately excluded they are present in anonymous
/// jars and do not indicate an authenticated session.
private static let knownSessionCookieNames: Set<String> = [
"sessionid", "session_id", "session_token",
"auth_token", "access_token",
]
/// Substrings that indicate a session or auth cookie (applied only when
/// no exact-name match is found). Deliberately excludes overly broad
/// patterns like "id" and "token" that match analytics/CSRF cookies.
private static let sessionCookieSubstrings = ["session", "auth", "sid", "jwt"]
/// Cookie name prefixes that indicate a non-session cookie even when a
/// substring match would otherwise accept it (e.g. "csrftoken").
private static let excludedCookiePrefixes = ["csrf", "_ga", "_gid", "tracking", "analytics"]
public struct SessionInfo: Sendable {
public let cookies: [HTTPCookie]
public let sourceLabel: String
public init(cookies: [HTTPCookie], sourceLabel: String) {
self.cookies = cookies
self.sourceLabel = sourceLabel
}
public var cookieHeader: String {
self.cookies.map { "\($0.name)=\($0.value)" }.joined(separator: "; ")
}
}
/// Returns all candidate sessions across browsers/profiles, ordered by
/// import priority. Callers should try each in turn so that a stale
/// session in the first source doesn't block a valid one further down.
///
/// Defaults to Chrome-only per AGENTS.md guideline. Pass an empty
/// `preferredBrowsers` list to fall back to the full descriptor-defined
/// import order (Safari, Firefox, etc.) when Chrome has no cookies.
public static func importSessions(
browserDetection: BrowserDetection = BrowserDetection(),
preferredBrowsers: [Browser] = [.chrome],
logger: ((String) -> Void)? = nil) throws -> [SessionInfo]
{
var candidates: [SessionInfo] = []
let installedBrowsers = preferredBrowsers.isEmpty
? self.cookieImportOrder.cookieImportCandidates(using: browserDetection)
: preferredBrowsers.cookieImportCandidates(using: browserDetection)
for browserSource in installedBrowsers {
do {
let query = BrowserCookieQuery(domains: self.cookieDomains)
let sources = try Self.cookieClient.codexBarRecords(
matching: query,
in: browserSource,
logger: { msg in self.emit(msg, logger: logger) })
for source in sources where !source.records.isEmpty {
let httpCookies = BrowserCookieClient.makeHTTPCookies(source.records, origin: query.origin)
guard !httpCookies.isEmpty else { continue }
guard Self.containsSessionCookie(httpCookies) else {
self.emit(
"Skipping \(source.label): no session cookie found",
logger: logger)
continue
}
self.emit(
"Found \(httpCookies.count) session cookies in \(source.label)",
logger: logger)
candidates.append(SessionInfo(cookies: httpCookies, sourceLabel: source.label))
}
} catch {
BrowserCookieAccessGate.recordIfNeeded(error)
self.emit(
"\(browserSource.displayName) cookie import failed: \(error.localizedDescription)",
logger: logger)
}
}
guard !candidates.isEmpty else {
throw AbacusUsageError.noSessionCookie
}
return candidates
}
/// Cheap check for whether any browser has an Abacus session cookie,
/// used by the fetch strategy's `isAvailable()`.
public static func hasSession(
browserDetection: BrowserDetection = BrowserDetection(),
preferredBrowsers: [Browser] = [.chrome],
logger: ((String) -> Void)? = nil) -> Bool
{
do {
return try !self.importSessions(
browserDetection: browserDetection,
preferredBrowsers: preferredBrowsers,
logger: logger).isEmpty
} catch {
return false
}
}
/// Returns `true` if the cookie set contains at least one cookie whose name
/// indicates session or authentication state. Checks exact known names
/// first, then falls back to conservative substring matching.
private static func containsSessionCookie(_ cookies: [HTTPCookie]) -> Bool {
cookies.contains { cookie in
let lower = cookie.name.lowercased()
if self.knownSessionCookieNames.contains(lower) { return true }
if self.excludedCookiePrefixes.contains(where: { lower.hasPrefix($0) }) { return false }
return self.sessionCookieSubstrings.contains { lower.contains($0) }
}
}
private static func emit(_ message: String, logger: ((String) -> Void)?) {
logger?("[abacus-cookie] \(message)")
self.log.debug(message)
}
}
#endif
@@ -0,0 +1,91 @@
import Foundation
#if os(macOS)
import SweetCookieKit
#endif
public enum AbacusProviderDescriptor {
public static let descriptor: ProviderDescriptor = Self.makeDescriptor()
static func makeDescriptor() -> ProviderDescriptor {
ProviderDescriptor(
id: .abacus,
metadata: ProviderMetadata(
id: .abacus,
displayName: "Abacus AI",
sessionLabel: "Credits",
weeklyLabel: "Weekly",
opusLabel: nil,
supportsOpus: false,
supportsCredits: false,
creditsHint: "",
toggleTitle: "Show Abacus AI usage",
cliName: "abacusai",
defaultEnabled: false,
isPrimaryProvider: false,
usesAccountFallback: false,
browserCookieOrder: ProviderBrowserCookieDefaults.defaultImportOrder,
dashboardURL: "https://apps.abacus.ai/chatllm/admin/compute-points-usage",
statusPageURL: nil,
statusLinkURL: nil),
branding: ProviderBranding(
iconStyle: .abacus,
iconResourceName: "ProviderIcon-abacus",
color: ProviderColor(red: 56 / 255, green: 189 / 255, blue: 248 / 255)),
tokenCost: ProviderTokenCostConfig(
supportsTokenCost: false,
noDataMessage: { "Abacus AI cost summary is not supported." }),
fetchPlan: ProviderFetchPlan(
sourceModes: [.auto, .web],
pipeline: ProviderFetchPipeline(resolveStrategies: { _ in
[AbacusWebFetchStrategy()]
})),
cli: ProviderCLIConfig(
name: "abacusai",
aliases: ["abacus-ai"],
versionDetector: nil))
}
}
// MARK: - Fetch Strategy
struct AbacusWebFetchStrategy: ProviderFetchStrategy {
let id: String = "abacus.web"
let kind: ProviderFetchKind = .web
func isAvailable(_ context: ProviderFetchContext) async -> Bool {
context.settings?.abacus?.cookieSource != .off
}
func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult {
let manual: String?
if context.settings?.abacus?.cookieSource == .manual {
guard let header = Self.manualCookieHeader(from: context) else {
throw AbacusUsageError.noSessionCookie
}
manual = header
} else {
manual = nil
}
let logger: ((String) -> Void)? = context.verbose
? { msg in CodexBarLog.logger(LogCategories.abacusUsage).verbose(msg) }
: nil
let snap = try await AbacusUsageFetcher.fetchUsage(
cookieHeaderOverride: manual,
browserDetection: context.browserDetection,
timeout: context.webTimeout,
logger: logger)
return self.makeResult(
usage: snap.toUsageSnapshot(),
sourceLabel: "web")
}
func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool {
false
}
private static func manualCookieHeader(from context: ProviderFetchContext) -> String? {
guard context.settings?.abacus?.cookieSource == .manual else { return nil }
return CookieHeaderNormalizer.normalize(context.settings?.abacus?.manualCookieHeader)
}
}
@@ -0,0 +1,67 @@
import Foundation
// MARK: - Abacus Usage Error
public enum AbacusUsageError: LocalizedError, Sendable, Equatable {
case noSessionCookie
case sessionExpired
case networkError(String)
case parseFailed(String)
case unauthorized
/// Whether this error indicates an authentication/session problem that
/// should trigger cache eviction.
public var isRecoverable: Bool {
switch self {
case .unauthorized, .sessionExpired: true
default: false
}
}
public var isAuthRelated: Bool {
self.isRecoverable
}
/// Whether browser-import scanning should continue to later sessions after
/// this failure. Imported sessions can differ by profile/browser, so we keep
/// scanning on per-session fetch failures and surface the first one only if
/// every candidate is exhausted.
var shouldTryNextImportedSession: Bool {
switch self {
case .unauthorized, .sessionExpired, .networkError, .parseFailed: true
case .noSessionCookie: false
}
}
/// Whether a cached cookie header should be evicted before falling back to
/// a fresh browser import. Parse/auth failures usually indicate that the
/// cached session is stale or no longer accepted.
var shouldClearCachedCookie: Bool {
switch self {
case .unauthorized, .sessionExpired, .parseFailed: true
case .networkError, .noSessionCookie: false
}
}
public var errorDescription: String? {
switch self {
case .noSessionCookie:
"No Abacus AI session found. Please log in to apps.abacus.ai in your browser "
+ "or paste a Cookie header in manual mode."
case .sessionExpired:
"Abacus AI session expired. Please log in again."
case let .networkError(msg):
"Abacus AI API error: \(msg)"
case let .parseFailed(msg):
"Could not parse Abacus AI usage: \(msg)"
case .unauthorized:
"Unauthorized. Please log in to Abacus AI."
}
}
}
#if !os(macOS)
extension AbacusUsageError {
public static let notSupported = AbacusUsageError.networkError("Abacus AI is only supported on macOS.")
}
#endif
@@ -0,0 +1,336 @@
import Foundation
#if canImport(FoundationNetworking)
import FoundationNetworking
#endif
#if os(macOS)
import SweetCookieKit
// MARK: - Abacus Usage Fetcher
public enum AbacusUsageFetcher {
private struct BrowserFetchRequest {
let browserDetection: BrowserDetection
let preferredBrowsers: [Browser]
let label: String
let timeout: TimeInterval
let logger: ((String) -> Void)?
}
/// Parsed JSON dictionaries are treated as immutable snapshots here and are
/// only moved between sibling fetch tasks before being consumed locally.
private struct JSONDictionaryBox: @unchecked Sendable {
let value: [String: Any]
}
private static let log = CodexBarLog.logger(LogCategories.abacusUsage)
private static let computePointsURL =
URL(string: "https://apps.abacus.ai/api/_getOrganizationComputePoints")!
private static let billingInfoURL =
URL(string: "https://apps.abacus.ai/api/_getBillingInfo")!
public static func fetchUsage(
cookieHeaderOverride: String? = nil,
browserDetection: BrowserDetection = BrowserDetection(),
timeout: TimeInterval = 15.0,
logger: ((String) -> Void)? = nil) async throws -> AbacusUsageSnapshot
{
// Manual cookie header no fallback, errors propagate directly
if let override = CookieHeaderNormalizer.normalize(cookieHeaderOverride) {
self.emit("Using manual cookie header", logger: logger)
return try await self.fetchWithCookieHeader(override, timeout: timeout, logger: logger)
}
// Cached cookie header fall back to a fresh browser import when the
// cached session is rejected or looks stale.
if let cached = CookieHeaderCache.load(provider: .abacus),
!cached.cookieHeader.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty
{
self.emit("Using cached cookie header from \(cached.sourceLabel)", logger: logger)
do {
return try await self.fetchWithCookieHeader(
cached.cookieHeader, timeout: timeout, logger: logger)
} catch let error as AbacusUsageError where error.shouldTryNextImportedSession {
if error.shouldClearCachedCookie {
CookieHeaderCache.clear(provider: .abacus)
self.emit(
"Cached cookie failed (\(error.localizedDescription)); cleared, trying fresh import",
logger: logger)
} else {
self.emit(
"Cached cookie failed (\(error.localizedDescription)); trying fresh import",
logger: logger)
}
}
}
// Fresh browser import try Chrome first (AGENTS.md default), then broaden
// to all browsers if Chrome has no sessions OR if every imported Chrome
// session is exhausted without a successful fetch.
var lastError: AbacusUsageError = .noSessionCookie
if let snapshot = try await self.tryFetchFromBrowsers(
BrowserFetchRequest(
browserDetection: browserDetection,
preferredBrowsers: [.chrome],
label: "Chrome",
timeout: timeout,
logger: logger),
lastError: &lastError)
{
return snapshot
}
self.emit("Chrome sessions exhausted; falling back to all browsers", logger: logger)
if let snapshot = try await self.tryFetchFromBrowsers(
BrowserFetchRequest(
browserDetection: browserDetection,
preferredBrowsers: [],
label: "all browsers",
timeout: timeout,
logger: logger),
lastError: &lastError)
{
return snapshot
}
throw lastError
}
/// Tries to import sessions from `preferredBrowsers` and fetch usage. Returns
/// the snapshot on success, nil if no sessions were available or every
/// imported session was exhausted without success.
private static func tryFetchFromBrowsers(
_ request: BrowserFetchRequest,
lastError: inout AbacusUsageError) async throws -> AbacusUsageSnapshot?
{
let sessions: [AbacusCookieImporter.SessionInfo]
do {
sessions = try AbacusCookieImporter.importSessions(
browserDetection: request.browserDetection,
preferredBrowsers: request.preferredBrowsers,
logger: request.logger)
} catch {
BrowserCookieAccessGate.recordIfNeeded(error)
self.emit(
"\(request.label) cookie import failed: \(error.localizedDescription)",
logger: request.logger)
return nil
}
for session in sessions {
self.emit("Trying cookies from \(session.sourceLabel)", logger: request.logger)
do {
let snapshot = try await self.fetchWithCookieHeader(
session.cookieHeader,
timeout: request.timeout,
logger: request.logger)
CookieHeaderCache.store(
provider: .abacus,
cookieHeader: session.cookieHeader,
sourceLabel: session.sourceLabel)
return snapshot
} catch let error as AbacusUsageError where error.shouldTryNextImportedSession {
self.emit(
"\(session.sourceLabel): \(error.localizedDescription), trying next source",
logger: request.logger)
lastError = error
continue
} catch {
self.emit(
"\(session.sourceLabel): \(error.localizedDescription), trying next source",
logger: request.logger)
lastError = .networkError(error.localizedDescription)
continue
}
}
return nil
}
// MARK: - API Requests
private static func fetchWithCookieHeader(
_ cookieHeader: String,
timeout: TimeInterval,
logger: ((String) -> Void)? = nil) async throws -> AbacusUsageSnapshot
{
enum FetchPart: Sendable {
case computePoints(JSONDictionaryBox)
case billingInfoSuccess(JSONDictionaryBox)
case billingInfoFailure(String)
}
// Fetch compute points (required, full timeout) and billing info
// (optional, shorter budget) concurrently. Billing is bounded so a
// slow/flaky billing endpoint can't delay credit rendering.
let billingBudget = min(timeout, 5.0)
var computePointsResult: [String: Any]?
var billingInfoResult: [String: Any] = [:]
try await withThrowingTaskGroup(of: FetchPart.self) { group in
group.addTask {
let result = try await self.fetchJSON(
url: self.computePointsURL,
method: "GET",
cookieHeader: cookieHeader,
timeout: timeout)
return .computePoints(JSONDictionaryBox(value: result))
}
group.addTask {
do {
let result = try await self.fetchJSON(
url: self.billingInfoURL,
method: "POST",
cookieHeader: cookieHeader,
timeout: billingBudget)
return .billingInfoSuccess(JSONDictionaryBox(value: result))
} catch {
return .billingInfoFailure(error.localizedDescription)
}
}
while let result = try await group.next() {
switch result {
case let .computePoints(value):
computePointsResult = value.value
case let .billingInfoSuccess(value):
billingInfoResult = value.value
case let .billingInfoFailure(message):
self.emit(
"Billing info fetch failed: \(message); credits shown without plan/reset",
logger: logger)
}
}
}
guard let computePointsResult else {
throw AbacusUsageError.networkError("Abacus compute points fetch did not complete")
}
return try self.parseResults(computePoints: computePointsResult, billingInfo: billingInfoResult)
}
private static func fetchJSON(
url: URL, method: String, cookieHeader: String, timeout: TimeInterval) async throws -> [String: Any]
{
var request = URLRequest(url: url)
request.httpMethod = method
request.timeoutInterval = timeout
request.setValue("application/json", forHTTPHeaderField: "Content-Type")
request.setValue("application/json", forHTTPHeaderField: "Accept")
request.setValue(cookieHeader, forHTTPHeaderField: "Cookie")
if method == "POST" {
request.httpBody = Data("{}".utf8)
}
let response = try await ProviderHTTPClient.shared.response(for: request)
let data = response.data
let statusCode = response.statusCode
if statusCode == 401 || statusCode == 403 {
throw AbacusUsageError.unauthorized
}
guard statusCode == 200 else {
let body = String(data: data.prefix(200), encoding: .utf8) ?? ""
throw AbacusUsageError.networkError("HTTP \(statusCode): \(body)")
}
let parsed: Any
do {
parsed = try JSONSerialization.jsonObject(with: data)
} catch {
let preview = String(data: data.prefix(200), encoding: .utf8) ?? "<non-UTF8>"
throw AbacusUsageError.parseFailed(
"\(url.lastPathComponent): \(error.localizedDescription) — preview: \(preview)")
}
guard let root = parsed as? [String: Any] else {
throw AbacusUsageError.parseFailed("\(url.lastPathComponent): top-level JSON is not a dictionary")
}
guard root["success"] as? Bool == true,
let result = root["result"] as? [String: Any]
else {
let errorMsg = (root["error"] as? String ?? "Unknown error").lowercased()
if errorMsg.contains("expired") || errorMsg.contains("session")
|| errorMsg.contains("login") || errorMsg.contains("authenticate")
|| errorMsg.contains("unauthorized") || errorMsg.contains("unauthenticated")
|| errorMsg.contains("forbidden")
{
throw AbacusUsageError.unauthorized
}
throw AbacusUsageError.parseFailed("\(url.lastPathComponent): \(errorMsg)")
}
return result
}
// MARK: - Parsing
private static func parseResults(
computePoints: [String: Any], billingInfo: [String: Any]) throws -> AbacusUsageSnapshot
{
let totalCredits = self.double(from: computePoints["totalComputePoints"])
let creditsLeft = self.double(from: computePoints["computePointsLeft"])
guard let totalCredits, let creditsLeft else {
let keys = computePoints.keys.sorted().joined(separator: ", ")
throw AbacusUsageError.parseFailed(
"Missing credit fields in compute points response. Keys: [\(keys)]")
}
let creditsUsed = totalCredits - creditsLeft
let nextBillingDate = billingInfo["nextBillingDate"] as? String
let currentTier = billingInfo["currentTier"] as? String
let resetsAt = self.parseDate(nextBillingDate)
return AbacusUsageSnapshot(
creditsUsed: creditsUsed,
creditsTotal: totalCredits,
resetsAt: resetsAt,
planName: currentTier)
}
private static func double(from value: Any?) -> Double? {
if let d = value as? Double { return d }
if let i = value as? Int { return Double(i) }
if let n = value as? NSNumber { return n.doubleValue }
return nil
}
private static func parseDate(_ isoString: String?) -> Date? {
guard let isoString else { return nil }
let formatter = ISO8601DateFormatter()
formatter.formatOptions = [.withInternetDateTime, .withFractionalSeconds]
if let date = formatter.date(from: isoString) { return date }
formatter.formatOptions = [.withInternetDateTime]
return formatter.date(from: isoString)
}
// MARK: - Logging
private static func emit(_ message: String, logger: ((String) -> Void)?) {
logger?("[abacus] \(message)")
self.log.debug(message)
}
}
#else
// MARK: - Abacus (Unsupported)
public enum AbacusUsageFetcher {
public static func fetchUsage(
cookieHeaderOverride _: String? = nil,
browserDetection _: BrowserDetection = BrowserDetection(),
timeout _: TimeInterval = 15.0,
logger _: ((String) -> Void)? = nil) async throws -> AbacusUsageSnapshot
{
throw AbacusUsageError.notSupported
}
}
#endif
@@ -0,0 +1,77 @@
import Foundation
// MARK: - Abacus Usage Snapshot
public struct AbacusUsageSnapshot: Sendable {
public let creditsUsed: Double?
public let creditsTotal: Double?
public let resetsAt: Date?
public let planName: String?
public init(
creditsUsed: Double? = nil,
creditsTotal: Double? = nil,
resetsAt: Date? = nil,
planName: String? = nil)
{
self.creditsUsed = creditsUsed
self.creditsTotal = creditsTotal
self.resetsAt = resetsAt
self.planName = planName
}
public func toUsageSnapshot() -> UsageSnapshot {
let percentUsed: Double = if let used = self.creditsUsed, let total = self.creditsTotal, total > 0 {
(used / total) * 100.0
} else {
0
}
let resetDesc: String? = if let used = self.creditsUsed, let total = self.creditsTotal {
"\(Self.formatCredits(used)) / \(Self.formatCredits(total)) credits"
} else {
nil
}
// Derive window from actual billing cycle when possible.
// Assume the cycle started one calendar month before resetsAt.
let windowMinutes: Int = if let resetDate = self.resetsAt,
let cycleStart = Calendar.current.date(byAdding: .month, value: -1, to: resetDate)
{
max(1, Int(resetDate.timeIntervalSince(cycleStart) / 60))
} else {
30 * 24 * 60
}
let primary = RateWindow(
usedPercent: percentUsed,
windowMinutes: windowMinutes,
resetsAt: self.resetsAt,
resetDescription: resetDesc)
let identity = ProviderIdentitySnapshot(
providerID: .abacus,
accountEmail: nil,
accountOrganization: nil,
loginMethod: self.planName)
return UsageSnapshot(
primary: primary,
secondary: nil,
tertiary: nil,
providerCost: nil,
updatedAt: Date(),
identity: identity)
}
// MARK: - Formatting
/// Thread-safe credit formatting allocates per call to avoid shared mutable state.
private static func formatCredits(_ value: Double) -> String {
let formatter = NumberFormatter()
formatter.numberStyle = .decimal
formatter.locale = Locale(identifier: "en_US")
formatter.maximumFractionDigits = value >= 1000 ? 0 : 1
return formatter.string(from: NSNumber(value: value)) ?? String(format: "%.0f", value)
}
}
@@ -0,0 +1,133 @@
import Foundation
public enum AlibabaCodingPlanAPIRegion: String, CaseIterable, Sendable {
case international = "intl"
case chinaMainland = "cn"
private static let endpointPath = "data/api.json"
public var displayName: String {
switch self {
case .international:
"International (modelstudio.console.alibabacloud.com)"
case .chinaMainland:
"China mainland (bailian.console.aliyun.com)"
}
}
public var gatewayBaseURLString: String {
switch self {
case .international:
"https://modelstudio.console.alibabacloud.com"
case .chinaMainland:
"https://bailian.console.aliyun.com"
}
}
public var dashboardURL: URL {
switch self {
case .international:
URL(
string: "https://modelstudio.console.alibabacloud.com/ap-southeast-1/?tab=coding-plan#/efm/coding_plan")!
case .chinaMainland:
URL(string: "https://bailian.console.aliyun.com/cn-beijing/?tab=model#/efm/coding_plan")!
}
}
public var consoleDomain: String {
switch self {
case .international:
"modelstudio.console.alibabacloud.com"
case .chinaMainland:
"bailian.console.aliyun.com"
}
}
public var consoleSite: String {
switch self {
case .international:
"MODELSTUDIO_ALIBABACLOUD"
case .chinaMainland:
"BAILIAN_ALIYUN"
}
}
public var consoleRefererURL: URL {
switch self {
case .international:
URL(string: "https://modelstudio.console.alibabacloud.com/ap-southeast-1/?tab=coding-plan")!
case .chinaMainland:
URL(string: "https://bailian.console.aliyun.com/cn-beijing/?tab=model")!
}
}
public var quotaURL: URL {
var components = URLComponents(string: self.gatewayBaseURLString)!
components.path = "/" + Self.endpointPath
components.queryItems = [
URLQueryItem(
name: "action",
value: "zeldaEasy.broadscope-bailian.codingPlan.queryCodingPlanInstanceInfoV2"),
URLQueryItem(name: "product", value: "broadscope-bailian"),
URLQueryItem(name: "api", value: "queryCodingPlanInstanceInfoV2"),
URLQueryItem(name: "currentRegionId", value: self.currentRegionID),
]
return components.url!
}
public var consoleRPCBaseURLString: String {
switch self {
case .international:
"https://bailian-singapore-cs.alibabacloud.com"
case .chinaMainland:
"https://bailian-cs.console.aliyun.com"
}
}
public var consoleRPCURL: URL {
var components = URLComponents(string: self.consoleRPCBaseURLString)!
components.path = "/" + Self.endpointPath
components.queryItems = [
URLQueryItem(name: "action", value: self.consoleRPCAction),
URLQueryItem(name: "product", value: self.consoleRPCProduct),
URLQueryItem(name: "api", value: self.consoleQuotaAPIName),
URLQueryItem(name: "_v", value: "undefined"),
]
return components.url!
}
public var consoleRPCAction: String {
switch self {
case .international:
"IntlBroadScopeAspnGateway"
case .chinaMainland:
"BroadScopeAspnGateway"
}
}
public var consoleRPCProduct: String {
"sfm_bailian"
}
public var consoleQuotaAPIName: String {
"zeldaEasy.broadscope-bailian.codingPlan.queryCodingPlanInstanceInfoV2"
}
public var commodityCode: String {
switch self {
case .international:
"sfm_codingplan_public_intl"
case .chinaMainland:
"sfm_codingplan_public_cn"
}
}
public var currentRegionID: String {
switch self {
case .international:
"ap-southeast-1"
case .chinaMainland:
"cn-beijing"
}
}
}
@@ -0,0 +1,510 @@
import Foundation
#if os(macOS)
import CommonCrypto
import Security
import SQLite3
import SweetCookieKit
private let alibabaCookieImportOrder: BrowserCookieImportOrder =
ProviderDefaults.metadata[.alibaba]?.browserCookieOrder ?? Browser.defaultImportOrder
public enum AlibabaCodingPlanCookieImporter {
private static let cookieClient = BrowserCookieClient()
private static let cookieDomains = [
"bailian-singapore-cs.alibabacloud.com",
"bailian-cs.console.aliyun.com",
"bailian-beijing-cs.aliyuncs.com",
"modelstudio.console.alibabacloud.com",
"bailian.console.aliyun.com",
"free.aliyun.com",
"account.aliyun.com",
"signin.aliyun.com",
"passport.alibabacloud.com",
"console.alibabacloud.com",
"console.aliyun.com",
"alibabacloud.com",
"aliyun.com",
]
public struct SessionInfo: Sendable {
public let cookies: [HTTPCookie]
public let sourceLabel: String
public init(cookies: [HTTPCookie], sourceLabel: String) {
self.cookies = cookies
self.sourceLabel = sourceLabel
}
public var cookieHeader: String {
var byName: [String: HTTPCookie] = [:]
byName.reserveCapacity(self.cookies.count)
for cookie in self.cookies {
if let expiry = cookie.expiresDate, expiry < Date() {
continue
}
guard !cookie.value.isEmpty else { continue }
if let existing = byName[cookie.name] {
let existingExpiry = existing.expiresDate ?? .distantPast
let candidateExpiry = cookie.expiresDate ?? .distantPast
if candidateExpiry >= existingExpiry {
byName[cookie.name] = cookie
}
} else {
byName[cookie.name] = cookie
}
}
return byName.keys.sorted().compactMap { name in
guard let cookie = byName[name] else { return nil }
return "\(cookie.name)=\(cookie.value)"
}.joined(separator: "; ")
}
}
nonisolated(unsafe) static var importSessionOverrideForTesting:
((BrowserDetection, ((String) -> Void)?) throws -> SessionInfo)?
public static func importSession(
browserDetection: BrowserDetection,
logger: ((String) -> Void)? = nil) throws -> SessionInfo
{
if let override = self.importSessionOverrideForTesting {
return try override(browserDetection, logger)
}
let log: (String) -> Void = { msg in logger?("[alibaba-cookie] \(msg)") }
var accessDeniedHints: [String] = []
var failureDetails: [String] = []
let installedBrowsers = self.cookieImportCandidates(browserDetection: browserDetection)
log("Cookie import candidates: \(installedBrowsers.map(\.displayName).joined(separator: ", "))")
for browserSource in installedBrowsers {
do {
log("Checking \(browserSource.displayName)")
let query = BrowserCookieQuery(domains: self.cookieDomains)
let sources = try Self.cookieClient.codexBarRecords(
matching: query,
in: browserSource,
logger: log)
if sources.isEmpty {
log("No matching cookie records in \(browserSource.displayName)")
if let fallbackSession = try Self.importChromiumFallbackSession(
browser: browserSource,
logger: log)
{
return fallbackSession
}
}
for source in sources where !source.records.isEmpty {
let httpCookies = BrowserCookieClient.makeHTTPCookies(source.records, origin: query.origin)
if self.isAuthenticatedSession(cookies: httpCookies) {
log("Found \(httpCookies.count) Alibaba cookies in \(source.label)")
return SessionInfo(cookies: httpCookies, sourceLabel: source.label)
}
let cookieNames = Set(httpCookies.map(\.name))
let hasTicket = cookieNames.contains("login_aliyunid_ticket")
let hasAccount =
cookieNames.contains("login_aliyunid_pk") ||
cookieNames.contains("login_current_pk") ||
cookieNames.contains("login_aliyunid")
log("Skipping \(source.label): missing auth cookies (ticket=\(hasTicket), account=\(hasAccount))")
}
if let fallbackSession = try Self.importChromiumFallbackSession(browser: browserSource, logger: log) {
return fallbackSession
}
} catch let error as BrowserCookieError {
BrowserCookieAccessGate.recordIfNeeded(error)
if let hint = error.accessDeniedHint {
accessDeniedHints.append(hint)
}
failureDetails.append("\(browserSource.displayName): \(error.localizedDescription)")
log("\(browserSource.displayName) cookie import failed: \(error.localizedDescription)")
} catch {
failureDetails.append("\(browserSource.displayName): \(error.localizedDescription)")
log("\(browserSource.displayName) cookie import failed: \(error.localizedDescription)")
}
}
let details = (Array(Set(accessDeniedHints)).sorted() + Array(Set(failureDetails)).sorted())
.joined(separator: " ")
throw AlibabaCodingPlanSettingsError.missingCookie(details: details.isEmpty ? nil : details)
}
private static func isAuthenticatedSession(cookies: [HTTPCookie]) -> Bool {
guard !cookies.isEmpty else { return false }
let names = Set(cookies.map(\.name))
let hasTicket = names.contains("login_aliyunid_ticket")
let hasAccount =
names.contains("login_aliyunid_pk") ||
names.contains("login_current_pk") ||
names.contains("login_aliyunid")
return hasTicket && hasAccount
}
public static func hasSession(
browserDetection: BrowserDetection,
logger: ((String) -> Void)? = nil) -> Bool
{
do {
_ = try self.importSession(browserDetection: browserDetection, logger: logger)
return true
} catch {
return false
}
}
private static func importChromiumFallbackSession(
browser: Browser,
logger: ((String) -> Void)? = nil) throws -> SessionInfo?
{
guard browser.usesChromiumProfileStore else { return nil }
return try AlibabaChromiumCookieFallbackImporter.importSession(
browser: browser,
domains: self.cookieDomains,
logger: logger)
}
static func cookieImportCandidates(
browserDetection: BrowserDetection,
importOrder: BrowserCookieImportOrder = alibabaCookieImportOrder) -> [Browser]
{
importOrder.cookieImportCandidates(using: browserDetection)
}
static func matchesCookieDomain(_ domain: String, patterns: [String] = Self.cookieDomains) -> Bool {
let normalized = self.normalizeCookieDomain(domain)
return patterns.contains { pattern in
let normalizedPattern = self.normalizeCookieDomain(pattern)
return normalized == normalizedPattern || normalized.hasSuffix(".\(normalizedPattern)")
}
}
static func normalizeCookieDomain(_ domain: String) -> String {
let trimmed = domain.trimmingCharacters(in: .whitespacesAndNewlines)
let normalized = trimmed.hasPrefix(".") ? String(trimmed.dropFirst()) : trimmed
return normalized.lowercased()
}
}
enum AlibabaChromiumCookieFallbackImporter {
private struct ChromiumCookieRecord {
let domain: String
let name: String
let path: String
let value: String
let expires: Date?
let isSecure: Bool
}
enum ImportError: LocalizedError {
case keyUnavailable(browser: Browser)
case keychainDenied(browser: Browser)
case sqliteFailed(label: String, details: String)
var errorDescription: String? {
switch self {
case let .keyUnavailable(browser):
"\(browser.displayName) Safe Storage key not found."
case let .keychainDenied(browser):
"macOS Keychain denied access to \(browser.displayName) Safe Storage."
case let .sqliteFailed(label, details):
"\(label) cookie fallback failed: \(details)"
}
}
}
static func importSession(
browser: Browser,
domains: [String],
cookieClient: BrowserCookieClient = BrowserCookieClient(),
logger: ((String) -> Void)? = nil) throws -> AlibabaCodingPlanCookieImporter.SessionInfo?
{
let stores = try cookieClient.codexBarStores(for: browser).filter { $0.databaseURL != nil }
guard !stores.isEmpty else { return nil }
logger?("[alibaba-cookie] Trying \(browser.displayName) Chromium fallback")
let keys = try self.derivedKeys(for: browser)
for store in stores {
let cookies = try self.loadCookies(from: store, domains: domains, keys: keys)
guard !cookies.isEmpty else { continue }
if self.isAuthenticatedSession(cookies) {
logger?("[alibaba-cookie] Found \(cookies.count) Alibaba cookies via \(store.label) fallback")
return AlibabaCodingPlanCookieImporter.SessionInfo(cookies: cookies, sourceLabel: store.label)
}
}
return nil
}
private static func isAuthenticatedSession(_ cookies: [HTTPCookie]) -> Bool {
let names = Set(cookies.map(\.name))
let hasTicket = names.contains("login_aliyunid_ticket")
let hasAccount =
names.contains("login_aliyunid_pk") ||
names.contains("login_current_pk") ||
names.contains("login_aliyunid")
return hasTicket && hasAccount
}
private static func loadCookies(
from store: BrowserCookieStore,
domains: [String],
keys: [Data]) throws -> [HTTPCookie]
{
guard let sourceDB = store.databaseURL else { return [] }
let records = try self.readCookiesFromLockedDB(
sourceDB: sourceDB,
domains: domains,
keys: keys,
label: store.label)
return records.compactMap(self.makeCookie)
}
private static func readCookiesFromLockedDB(
sourceDB: URL,
domains: [String],
keys: [Data],
label: String) throws -> [ChromiumCookieRecord]
{
let tempDir = FileManager.default.temporaryDirectory
.appendingPathComponent("alibaba-chromium-cookies-\(UUID().uuidString)", isDirectory: true)
try FileManager.default.createDirectory(at: tempDir, withIntermediateDirectories: true)
let copiedDB = tempDir.appendingPathComponent("Cookies")
try FileManager.default.copyItem(at: sourceDB, to: copiedDB)
for suffix in ["-wal", "-shm"] {
let src = URL(fileURLWithPath: sourceDB.path + suffix)
if FileManager.default.fileExists(atPath: src.path) {
let dst = URL(fileURLWithPath: copiedDB.path + suffix)
try? FileManager.default.copyItem(at: src, to: dst)
}
}
defer { try? FileManager.default.removeItem(at: tempDir) }
return try self.readCookies(fromDB: copiedDB.path, domains: domains, keys: keys, label: label)
}
private static func readCookies(
fromDB path: String,
domains: [String],
keys: [Data],
label: String) throws -> [ChromiumCookieRecord]
{
var db: OpaquePointer?
guard sqlite3_open_v2(path, &db, SQLITE_OPEN_READONLY, nil) == SQLITE_OK else {
throw ImportError.sqliteFailed(label: label, details: String(cString: sqlite3_errmsg(db)))
}
defer { sqlite3_close(db) }
let sql = "SELECT host_key, name, path, expires_utc, is_secure, value, encrypted_value FROM cookies"
var stmt: OpaquePointer?
guard sqlite3_prepare_v2(db, sql, -1, &stmt, nil) == SQLITE_OK else {
throw ImportError.sqliteFailed(label: label, details: String(cString: sqlite3_errmsg(db)))
}
defer { sqlite3_finalize(stmt) }
var records: [ChromiumCookieRecord] = []
while sqlite3_step(stmt) == SQLITE_ROW {
guard let hostKey = self.readText(stmt, index: 0), self.matches(domain: hostKey, patterns: domains) else {
continue
}
guard let name = self.readText(stmt, index: 1), let path = self.readText(stmt, index: 2) else {
continue
}
let value: String? = if let plain = self.readText(stmt, index: 5), !plain.isEmpty {
plain
} else if let encrypted = self.readBlob(stmt, index: 6) {
self.decrypt(encrypted, usingAnyOf: keys)
} else {
nil
}
guard let value, !value.isEmpty else { continue }
records.append(ChromiumCookieRecord(
domain: AlibabaCodingPlanCookieImporter.normalizeCookieDomain(hostKey),
name: name,
path: path,
value: value,
expires: self.chromiumExpiry(sqlite3_column_int64(stmt, 3)),
isSecure: sqlite3_column_int(stmt, 4) != 0))
}
return records.filter { record in
guard let expires = record.expires else { return true }
return expires >= Date()
}
}
private static func derivedKeys(for browser: Browser) throws -> [Data] {
var keys: [Data] = []
var sawDenied = false
for label in browser.safeStorageLabels {
switch KeychainAccessPreflight.checkGenericPassword(service: label.service, account: label.account) {
case .interactionRequired:
sawDenied = true
continue
case .allowed, .notFound, .failure:
break
}
if let password = self.safeStoragePassword(service: label.service, account: label.account) {
keys.append(self.deriveKey(from: password))
}
}
if !keys.isEmpty {
return keys
}
if sawDenied {
throw ImportError.keychainDenied(browser: browser)
}
throw ImportError.keyUnavailable(browser: browser)
}
private static func safeStoragePassword(service: String, account: String) -> String? {
// The preflight classifies prompt-requiring items as .interactionRequired, but its
// .notFound (gate disabled) and .failure outcomes still reach this read. Honor the
// access gate and keep the read strictly non-interactive so it can never prompt.
guard !KeychainAccessGate.isDisabled else { return nil }
var query: [String: Any] = [
kSecClass as String: kSecClassGenericPassword,
kSecAttrService as String: service,
kSecAttrAccount as String: account,
kSecMatchLimit as String: kSecMatchLimitOne,
kSecReturnData as String: true,
]
KeychainNoUIQuery.apply(to: &query)
var result: AnyObject?
let status = KeychainSecurity.copyMatching(query as CFDictionary, &result)
guard status == errSecSuccess, let data = result as? Data else { return nil }
return String(data: data, encoding: .utf8)
}
private static func deriveKey(from password: String) -> Data {
let salt = Data("saltysalt".utf8)
var key = Data(count: kCCKeySizeAES128)
let keyLength = key.count
_ = key.withUnsafeMutableBytes { keyBytes in
password.utf8CString.withUnsafeBytes { passBytes in
salt.withUnsafeBytes { saltBytes in
CCKeyDerivationPBKDF(
CCPBKDFAlgorithm(kCCPBKDF2),
passBytes.bindMemory(to: Int8.self).baseAddress,
passBytes.count - 1,
saltBytes.bindMemory(to: UInt8.self).baseAddress,
salt.count,
CCPseudoRandomAlgorithm(kCCPRFHmacAlgSHA1),
1003,
keyBytes.bindMemory(to: UInt8.self).baseAddress,
keyLength)
}
}
}
return key
}
private static func decrypt(_ encryptedValue: Data, usingAnyOf keys: [Data]) -> String? {
for key in keys {
if let value = self.decrypt(encryptedValue, key: key) {
return value
}
}
return nil
}
private static func decrypt(_ encryptedValue: Data, key: Data) -> String? {
guard encryptedValue.count > 3 else { return nil }
let prefix = String(data: encryptedValue.prefix(3), encoding: .utf8)
guard prefix == "v10" else { return nil }
let payload = Data(encryptedValue.dropFirst(3))
let iv = Data(repeating: 0x20, count: kCCBlockSizeAES128)
var outLength = 0
var out = Data(count: payload.count + kCCBlockSizeAES128)
let outCapacity = out.count
let status = out.withUnsafeMutableBytes { outBytes in
payload.withUnsafeBytes { payloadBytes in
key.withUnsafeBytes { keyBytes in
iv.withUnsafeBytes { ivBytes in
CCCrypt(
CCOperation(kCCDecrypt),
CCAlgorithm(kCCAlgorithmAES),
CCOptions(kCCOptionPKCS7Padding),
keyBytes.baseAddress,
key.count,
ivBytes.baseAddress,
payloadBytes.baseAddress,
payload.count,
outBytes.baseAddress,
outCapacity,
&outLength)
}
}
}
}
guard status == kCCSuccess else { return nil }
out.count = outLength
if let value = String(data: out, encoding: .utf8), !value.isEmpty {
return value
}
if out.count > 32 {
let trimmed = out.dropFirst(32)
if let value = String(data: trimmed, encoding: .utf8), !value.isEmpty {
return value
}
}
return nil
}
private static func makeCookie(from record: ChromiumCookieRecord) -> HTTPCookie? {
var properties: [HTTPCookiePropertyKey: Any] = [
.domain: record.domain,
.path: record.path,
.name: record.name,
.value: record.value,
]
if record.isSecure {
properties[.secure] = true
}
if let expires = record.expires {
properties[.expires] = expires
}
return HTTPCookie(properties: properties)
}
private static func readText(_ stmt: OpaquePointer?, index: Int32) -> String? {
guard sqlite3_column_type(stmt, index) != SQLITE_NULL,
let value = sqlite3_column_text(stmt, index)
else {
return nil
}
return String(cString: value)
}
private static func readBlob(_ stmt: OpaquePointer?, index: Int32) -> Data? {
guard sqlite3_column_type(stmt, index) != SQLITE_NULL,
let bytes = sqlite3_column_blob(stmt, index)
else {
return nil
}
return Data(bytes: bytes, count: Int(sqlite3_column_bytes(stmt, index)))
}
private static func matches(domain: String, patterns: [String]) -> Bool {
AlibabaCodingPlanCookieImporter.matchesCookieDomain(domain, patterns: patterns)
}
private static func chromiumExpiry(_ expiresUTC: Int64) -> Date? {
guard expiresUTC > 0 else { return nil }
let seconds = (Double(expiresUTC) / 1_000_000.0) - 11_644_473_600.0
guard seconds > 0 else { return nil }
return Date(timeIntervalSince1970: seconds)
}
}
#endif
@@ -0,0 +1,266 @@
import Foundation
#if os(macOS)
import SweetCookieKit
#endif
public enum AlibabaCodingPlanProviderDescriptor {
public static let descriptor: ProviderDescriptor = Self.makeDescriptor()
static func makeDescriptor() -> ProviderDescriptor {
#if os(macOS)
let browserOrder: BrowserCookieImportOrder = [
.chrome,
.chromeBeta,
.brave,
.edge,
.arc,
.firefox,
.safari,
]
#else
let browserOrder: BrowserCookieImportOrder? = nil
#endif
return ProviderDescriptor(
id: .alibaba,
metadata: ProviderMetadata(
id: .alibaba,
displayName: "Alibaba",
sessionLabel: "5-hour",
weeklyLabel: "Weekly",
opusLabel: "Monthly",
supportsOpus: true,
supportsCredits: false,
creditsHint: "",
toggleTitle: "Show Alibaba usage",
cliName: "alibaba-coding-plan",
defaultEnabled: false,
isPrimaryProvider: false,
usesAccountFallback: false,
browserCookieOrder: browserOrder,
dashboardURL: AlibabaCodingPlanAPIRegion.international.dashboardURL.absoluteString,
statusPageURL: nil,
statusLinkURL: "https://status.aliyun.com"),
branding: ProviderBranding(
iconStyle: .alibaba,
iconResourceName: "ProviderIcon-alibaba",
color: ProviderColor(red: 1.0, green: 106 / 255, blue: 0)),
tokenCost: ProviderTokenCostConfig(
supportsTokenCost: false,
noDataMessage: { "Alibaba Coding Plan cost summary is not supported." }),
fetchPlan: ProviderFetchPlan(
sourceModes: [.auto, .web, .api],
pipeline: ProviderFetchPipeline(resolveStrategies: self.resolveStrategies)),
cli: ProviderCLIConfig(
name: "alibaba-coding-plan",
aliases: ["alibaba", "bailian"],
versionDetector: nil))
}
private static func resolveStrategies(context: ProviderFetchContext) async -> [any ProviderFetchStrategy] {
switch context.sourceMode {
case .web:
return [AlibabaCodingPlanWebFetchStrategy()]
case .api:
return [AlibabaCodingPlanAPIFetchStrategy()]
case .cli, .oauth:
return []
case .auto:
break
}
if context.settings?.alibaba?.cookieSource == .off {
return [AlibabaCodingPlanAPIFetchStrategy()]
}
return [AlibabaCodingPlanWebFetchStrategy(), AlibabaCodingPlanAPIFetchStrategy()]
}
}
struct AlibabaCodingPlanWebFetchStrategy: ProviderFetchStrategy {
let id: String = "alibaba-coding-plan.web"
let kind: ProviderFetchKind = .web
func isAvailable(_ context: ProviderFetchContext) async -> Bool {
guard context.settings?.alibaba?.cookieSource != .off else { return false }
if AlibabaCodingPlanSettingsReader.cookieHeader(environment: context.env) != nil {
return true
}
if let settings = context.settings?.alibaba,
settings.cookieSource == .manual
{
return CookieHeaderNormalizer.normalize(settings.manualCookieHeader) != nil
}
#if os(macOS)
if let cached = CookieHeaderCache.load(provider: .alibaba),
!cached.cookieHeader.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty
{
return true
}
if AlibabaCodingPlanCookieImporter.hasSession(browserDetection: context.browserDetection) {
return true
}
return false
#else
return false
#endif
}
func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult {
let cookieSource = context.settings?.alibaba?.cookieSource ?? .auto
let cookieHeader = try Self.resolveCookieHeader(context: context, allowCached: true)
do {
let region = context.settings?.alibaba?.apiRegion ?? .international
let usage = try await AlibabaCodingPlanUsageFetcher.fetchUsage(
cookieHeader: cookieHeader,
region: region,
environment: context.env)
return self.makeResult(usage: usage.toUsageSnapshot(), sourceLabel: "web")
} catch let error as AlibabaCodingPlanUsageError
where (error == .invalidCredentials || error == .loginRequired) && cookieSource != .manual
{
#if os(macOS)
CookieHeaderCache.clear(provider: .alibaba)
let refreshedHeader = try Self.resolveCookieHeader(context: context, allowCached: false)
let region = context.settings?.alibaba?.apiRegion ?? .international
let usage = try await AlibabaCodingPlanUsageFetcher.fetchUsage(
cookieHeader: refreshedHeader,
region: region,
environment: context.env)
return self.makeResult(usage: usage.toUsageSnapshot(), sourceLabel: "web")
#else
throw AlibabaCodingPlanUsageError.invalidCredentials
#endif
}
}
func shouldFallback(on error: Error, context: ProviderFetchContext) -> Bool {
guard context.sourceMode == .auto else { return false }
if let urlError = error as? URLError {
switch urlError.code {
case .timedOut,
.cannotFindHost,
.cannotConnectToHost,
.networkConnectionLost,
.dnsLookupFailed,
.secureConnectionFailed,
.serverCertificateHasBadDate,
.serverCertificateUntrusted,
.serverCertificateHasUnknownRoot,
.serverCertificateNotYetValid,
.clientCertificateRejected,
.clientCertificateRequired,
.cannotLoadFromNetwork,
.internationalRoamingOff,
.callIsActive,
.dataNotAllowed,
.requestBodyStreamExhausted,
.resourceUnavailable,
.notConnectedToInternet:
return true
default:
break
}
}
if let settingsError = error as? AlibabaCodingPlanSettingsError {
switch settingsError {
case .missingCookie, .invalidCookie:
return true
case .missingToken:
return false
}
}
guard let alibabaError = error as? AlibabaCodingPlanUsageError else { return false }
switch alibabaError {
case .loginRequired:
return true
case .invalidCredentials:
return true
case .apiKeyUnavailableInRegion:
return false
case let .apiError(message):
return message.contains("HTTP 404") || message.contains("HTTP 403")
case .networkError:
return true
case .parseFailed:
return false
}
}
static func resolveCookieHeader(context: ProviderFetchContext, allowCached: Bool) throws -> String {
if let settings = context.settings?.alibaba,
settings.cookieSource == .manual
{
guard let header = CookieHeaderNormalizer.normalize(settings.manualCookieHeader) else {
throw AlibabaCodingPlanSettingsError.invalidCookie
}
return header
}
if let envCookie = AlibabaCodingPlanSettingsReader.cookieHeader(environment: context.env),
let normalized = CookieHeaderNormalizer.normalize(envCookie)
{
return normalized
}
#if os(macOS)
if allowCached,
let cached = CookieHeaderCache.load(provider: .alibaba),
!cached.cookieHeader.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty
{
return cached.cookieHeader
}
do {
let session = try AlibabaCodingPlanCookieImporter.importSession(browserDetection: context.browserDetection)
CookieHeaderCache.store(
provider: .alibaba,
cookieHeader: session.cookieHeader,
sourceLabel: session.sourceLabel)
return session.cookieHeader
} catch {
throw error
}
#else
throw AlibabaCodingPlanSettingsError.missingCookie()
#endif
}
}
struct AlibabaCodingPlanAPIFetchStrategy: ProviderFetchStrategy {
let id: String = "alibaba-coding-plan.api"
let kind: ProviderFetchKind = .apiToken
func isAvailable(_ context: ProviderFetchContext) async -> Bool {
Self.resolveToken(environment: context.env) != nil
}
func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult {
guard let apiKey = Self.resolveToken(environment: context.env) else {
throw AlibabaCodingPlanSettingsError.missingToken
}
let region = context.settings?.alibaba?.apiRegion ?? .international
let usage = try await AlibabaCodingPlanUsageFetcher.fetchUsage(
apiKey: apiKey,
region: region,
environment: context.env)
return self.makeResult(
usage: usage.toUsageSnapshot(),
sourceLabel: "api")
}
func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool {
false
}
private static func resolveToken(environment: [String: String]) -> String? {
ProviderTokenResolver.alibabaToken(environment: environment)
}
}
@@ -0,0 +1,118 @@
import Foundation
public struct AlibabaCodingPlanSettingsReader: Sendable {
private static let endpointValidator = ProviderEndpointOverrideValidator(
allowedHosts: [
"modelstudio.console.alibabacloud.com",
"bailian-singapore-cs.alibabacloud.com",
"bailian.console.aliyun.com",
"bailian-cs.console.aliyun.com",
"bailian-beijing-cs.aliyuncs.com",
])
public static let apiTokenKey = "ALIBABA_CODING_PLAN_API_KEY"
public static let qwenAPITokenKey = "ALIBABA_QWEN_API_KEY"
public static let dashScopeAPITokenKey = "DASHSCOPE_API_KEY"
public static let apiTokenEnvironmentKeys = [
Self.apiTokenKey,
Self.qwenAPITokenKey,
Self.dashScopeAPITokenKey,
]
public static let cookieHeaderKey = "ALIBABA_CODING_PLAN_COOKIE"
public static let hostKey = "ALIBABA_CODING_PLAN_HOST"
public static let quotaURLKey = "ALIBABA_CODING_PLAN_QUOTA_URL"
public static let requireProviderEndpointOverridesKey = "ALIBABA_CODING_PLAN_REQUIRE_PROVIDER_ENDPOINT_OVERRIDES"
private static let endpointOverrideKeys = [
Self.hostKey,
Self.quotaURLKey,
]
public static func apiToken(
environment: [String: String] = ProcessInfo.processInfo.environment) -> String?
{
for key in self.apiTokenEnvironmentKeys {
if let token = self.cleaned(environment[key]) { return token }
}
return nil
}
public static func hostOverride(
environment: [String: String] = ProcessInfo.processInfo.environment) -> String?
{
self.endpointValidator.validatedHost(
self.cleaned(environment[self.hostKey]),
policy: self.endpointOverrideHostPolicy(environment: environment))
}
public static func rejectedEndpointOverrideKey(
environment: [String: String] = ProcessInfo.processInfo.environment) -> String?
{
let policy = self.endpointOverrideHostPolicy(environment: environment)
return self.endpointOverrideKeys.first { key in
guard let value = self.cleaned(environment[key]) else { return false }
if key == Self.hostKey {
return self.endpointValidator.validatedHost(value, policy: policy) == nil
}
return self.endpointValidator.validatedURL(value, policy: policy) == nil
}
}
public static func cookieHeader(
environment: [String: String] = ProcessInfo.processInfo.environment) -> String?
{
self.cleaned(environment[self.cookieHeaderKey])
}
public static func quotaURL(
environment: [String: String] = ProcessInfo.processInfo.environment) -> URL?
{
self.endpointValidator.validatedURL(
self.cleaned(environment[self.quotaURLKey]),
policy: self.endpointOverrideHostPolicy(environment: environment))
}
static func endpointOverrideHostPolicy(environment: [String: String]) -> ProviderEndpointOverrideValidator
.HostPolicy {
guard let value = self.cleaned(environment[self.requireProviderEndpointOverridesKey])?.lowercased(),
["1", "true", "yes", "on"].contains(value)
else { return .allowAnyHTTPSHost }
return .providerOwnedOnly
}
static func cleaned(_ raw: String?) -> String? {
guard var value = raw?.trimmingCharacters(in: .whitespacesAndNewlines), !value.isEmpty else {
return nil
}
if (value.hasPrefix("\"") && value.hasSuffix("\"")) ||
(value.hasPrefix("'") && value.hasSuffix("'"))
{
value = String(value.dropFirst().dropLast())
}
value = value.trimmingCharacters(in: .whitespacesAndNewlines)
return value.isEmpty ? nil : value
}
}
public enum AlibabaCodingPlanSettingsError: LocalizedError, Sendable {
case missingToken
case missingCookie(details: String? = nil)
case invalidCookie
public var errorDescription: String? {
switch self {
case .missingToken:
return "Alibaba Coding Plan API key not found. " +
"Set apiKey in ~/.codexbar/config.json, ALIBABA_CODING_PLAN_API_KEY, " +
"ALIBABA_QWEN_API_KEY, or DASHSCOPE_API_KEY."
case let .missingCookie(details):
let base = "No Alibaba Coding Plan session cookies found in browsers. " +
"If you use Safari, enable Full Disk Access for CodexBar/Terminal or paste a manual Cookie header."
guard let details, !details.isEmpty else { return base }
return "\(base) \(details)"
case .invalidCookie:
return "Alibaba Coding Plan cookie header is invalid."
}
}
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,126 @@
import Foundation
public struct AlibabaCodingPlanUsageSnapshot: Sendable {
public let planName: String?
public let fiveHourUsedQuota: Int?
public let fiveHourTotalQuota: Int?
public let fiveHourNextRefreshTime: Date?
public let weeklyUsedQuota: Int?
public let weeklyTotalQuota: Int?
public let weeklyNextRefreshTime: Date?
public let monthlyUsedQuota: Int?
public let monthlyTotalQuota: Int?
public let monthlyNextRefreshTime: Date?
public let updatedAt: Date
public init(
planName: String?,
fiveHourUsedQuota: Int?,
fiveHourTotalQuota: Int?,
fiveHourNextRefreshTime: Date?,
weeklyUsedQuota: Int?,
weeklyTotalQuota: Int?,
weeklyNextRefreshTime: Date?,
monthlyUsedQuota: Int?,
monthlyTotalQuota: Int?,
monthlyNextRefreshTime: Date?,
updatedAt: Date)
{
self.planName = planName
self.fiveHourUsedQuota = fiveHourUsedQuota
self.fiveHourTotalQuota = fiveHourTotalQuota
self.fiveHourNextRefreshTime = fiveHourNextRefreshTime
self.weeklyUsedQuota = weeklyUsedQuota
self.weeklyTotalQuota = weeklyTotalQuota
self.weeklyNextRefreshTime = weeklyNextRefreshTime
self.monthlyUsedQuota = monthlyUsedQuota
self.monthlyTotalQuota = monthlyTotalQuota
self.monthlyNextRefreshTime = monthlyNextRefreshTime
self.updatedAt = updatedAt
}
}
extension AlibabaCodingPlanUsageSnapshot {
public func toUsageSnapshot() -> UsageSnapshot {
let primaryPercent = Self.usedPercent(used: self.fiveHourUsedQuota, total: self.fiveHourTotalQuota)
let secondaryPercent = Self.usedPercent(used: self.weeklyUsedQuota, total: self.weeklyTotalQuota)
let tertiaryPercent = Self.usedPercent(used: self.monthlyUsedQuota, total: self.monthlyTotalQuota)
let primary: RateWindow? = if let primaryPercent {
RateWindow(
usedPercent: primaryPercent,
windowMinutes: 5 * 60,
resetsAt: Self.normalizedFiveHourReset(
self.fiveHourNextRefreshTime,
updatedAt: self.updatedAt),
resetDescription: Self.usageDetail(used: self.fiveHourUsedQuota, total: self.fiveHourTotalQuota))
} else {
nil
}
let secondary: RateWindow? = if let secondaryPercent {
RateWindow(
usedPercent: secondaryPercent,
windowMinutes: 7 * 24 * 60,
resetsAt: self.weeklyNextRefreshTime,
resetDescription: Self.usageDetail(used: self.weeklyUsedQuota, total: self.weeklyTotalQuota))
} else {
nil
}
let tertiary: RateWindow? = if let tertiaryPercent {
RateWindow(
usedPercent: tertiaryPercent,
windowMinutes: 30 * 24 * 60,
resetsAt: self.monthlyNextRefreshTime,
resetDescription: Self.usageDetail(used: self.monthlyUsedQuota, total: self.monthlyTotalQuota))
} else {
nil
}
let loginMethod = self.planName?.trimmingCharacters(in: .whitespacesAndNewlines)
let identity = ProviderIdentitySnapshot(
providerID: .alibaba,
accountEmail: nil,
accountOrganization: nil,
loginMethod: loginMethod)
return UsageSnapshot(
primary: primary,
secondary: secondary,
tertiary: tertiary,
providerCost: nil,
updatedAt: self.updatedAt,
identity: identity)
}
private static func usedPercent(used: Int?, total: Int?) -> Double? {
guard let used, let total, total > 0 else { return nil }
let normalizedUsed = max(0, min(used, total))
return Double(normalizedUsed) / Double(total) * 100
}
private static func limitDescription(total: Int?, label: String) -> String? {
guard let total, total > 0 else { return nil }
return "\(total) requests / \(label)"
}
private static func usageDetail(used: Int?, total: Int?) -> String? {
guard let used, let total, total > 0 else { return nil }
return "\(used) / \(total) used"
}
private static func normalizedFiveHourReset(_ raw: Date?, updatedAt: Date) -> Date? {
guard let raw else { return nil }
if raw.timeIntervalSince(updatedAt) >= 60 {
return raw
}
let shifted = raw.addingTimeInterval(TimeInterval(5 * 60 * 60))
if shifted.timeIntervalSince(updatedAt) >= 60 {
return shifted
}
return updatedAt.addingTimeInterval(TimeInterval(5 * 60 * 60))
}
}
@@ -0,0 +1,60 @@
import Foundation
public enum AlibabaTokenPlanAPIRegion: String, CaseIterable, Sendable {
case international = "intl"
case chinaMainland = "cn"
public var displayName: String {
switch self {
case .international:
"International (modelstudio.console.alibabacloud.com)"
case .chinaMainland:
"China mainland (bailian.console.aliyun.com)"
}
}
public var gatewayBaseURLString: String {
switch self {
case .international:
"https://modelstudio.console.alibabacloud.com"
case .chinaMainland:
"https://bailian.console.aliyun.com"
}
}
public var dashboardOriginURLString: String {
self.gatewayBaseURLString
}
public var dashboardURL: URL {
switch self {
case .international:
URL(
string: "https://modelstudio.console.alibabacloud.com/ap-southeast-1/?tab=plan#/efm/subscription/token-plan")!
case .chinaMainland:
URL(string: "https://bailian.console.aliyun.com/cn-beijing?tab=plan#/efm/subscription/token-plan")!
}
}
public var currentRegionID: String {
switch self {
case .international:
"ap-southeast-1"
case .chinaMainland:
"cn-beijing"
}
}
public var tokenPlanProductCode: String {
switch self {
case .international:
"sfm_tokenplanteams_dp_intl"
case .chinaMainland:
"sfm_tokenplanteams_dp_cn"
}
}
public var cookieCacheScope: CookieHeaderCache.Scope {
.providerVariant("alibaba-token-plan.\(self.rawValue)")
}
}
@@ -0,0 +1,163 @@
import Foundation
#if canImport(FoundationNetworking)
import FoundationNetworking
#endif
struct AlibabaTokenPlanCookieHeaders {
private static let cachedAPIHeaderName = "__codexbar_alibaba_token_plan_api"
private static let cachedDashboardHeaderName = "__codexbar_alibaba_token_plan_dashboard"
let apiCookieHeader: String
let dashboardCookieHeader: String
init(apiCookieHeader: String, dashboardCookieHeader: String) {
self.apiCookieHeader = apiCookieHeader
self.dashboardCookieHeader = dashboardCookieHeader
}
init?(singleHeader raw: String?) {
guard let normalized = CookieHeaderNormalizer.normalize(raw) else { return nil }
self.apiCookieHeader = normalized
self.dashboardCookieHeader = normalized
}
init?(cachedHeader raw: String?) {
var valuesByName: [String: String] = [:]
for pair in CookieHeaderNormalizer.pairs(from: raw ?? "") {
valuesByName[pair.name] = pair.value
}
if let encodedAPI = valuesByName[Self.cachedAPIHeaderName],
let encodedDashboard = valuesByName[Self.cachedDashboardHeaderName],
let apiHeader = Self.decodeCachedHeader(encodedAPI),
let dashboardHeader = Self.decodeCachedHeader(encodedDashboard),
let normalizedAPI = CookieHeaderNormalizer.normalize(apiHeader),
let normalizedDashboard = CookieHeaderNormalizer.normalize(dashboardHeader)
{
self.init(apiCookieHeader: normalizedAPI, dashboardCookieHeader: normalizedDashboard)
return
}
self.init(singleHeader: raw)
}
var cacheCookieHeader: String {
[
"\(Self.cachedAPIHeaderName)=\(Self.encodeCachedHeader(self.apiCookieHeader))",
"\(Self.cachedDashboardHeaderName)=\(Self.encodeCachedHeader(self.dashboardCookieHeader))",
].joined(separator: "; ")
}
var apiCookieNames: [String] {
Self.cookieNames(from: self.apiCookieHeader)
}
var dashboardCookieNames: [String] {
Self.cookieNames(from: self.dashboardCookieHeader)
}
func hasCookie(named name: String) -> Bool {
Self.cookieNames(from: self.apiCookieHeader).contains(name) ||
Self.cookieNames(from: self.dashboardCookieHeader).contains(name)
}
private static func cookieNames(from header: String) -> [String] {
CookieHeaderNormalizer.pairs(from: header)
.map(\.name)
.filter { !$0.isEmpty }
.uniquedSorted()
}
private static func encodeCachedHeader(_ header: String) -> String {
Data(header.utf8).base64EncodedString()
}
private static func decodeCachedHeader(_ encoded: String) -> String? {
guard let data = Data(base64Encoded: encoded) else { return nil }
return String(data: data, encoding: .utf8)
}
}
enum AlibabaTokenPlanCookieHeader {
static func headers(
from cookies: [HTTPCookie],
region: AlibabaTokenPlanAPIRegion = .international,
environment: [String: String] = ProcessInfo.processInfo.environment) -> AlibabaTokenPlanCookieHeaders?
{
guard let apiHeader = self.header(
from: cookies,
targetURL: AlibabaTokenPlanUsageFetcher.resolveQuotaURL(region: region, environment: environment)),
let dashboardHeader = self.header(
from: cookies,
targetURL: AlibabaTokenPlanUsageFetcher.dashboardURL(region: region, environment: environment))
else {
return nil
}
return AlibabaTokenPlanCookieHeaders(apiCookieHeader: apiHeader, dashboardCookieHeader: dashboardHeader)
}
static func header(from cookies: [HTTPCookie], targetURL: URL) -> String? {
var byName: [String: HTTPCookie] = [:]
for cookie in cookies {
guard !cookie.name.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty else { continue }
guard !cookie.value.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty else { continue }
if let expiry = cookie.expiresDate, expiry < Date() { continue }
guard self.matchesRequestURL(cookie: cookie, url: targetURL) else { continue }
if let existing = byName[cookie.name] {
if self.cookieSortKey(for: cookie) >= self.cookieSortKey(for: existing) {
byName[cookie.name] = cookie
}
} else {
byName[cookie.name] = cookie
}
}
guard !byName.isEmpty else { return nil }
return byName.keys.sorted().compactMap { name in
guard let cookie = byName[name] else { return nil }
return "\(cookie.name)=\(cookie.value)"
}.joined(separator: "; ")
}
private static func matchesRequestURL(cookie: HTTPCookie, url: URL) -> Bool {
guard let host = url.host?.lowercased() else { return false }
let normalizedDomain = cookie.domain.lowercased().trimmingCharacters(in: CharacterSet(charactersIn: "."))
guard !normalizedDomain.isEmpty else { return false }
guard host == normalizedDomain || host.hasSuffix(".\(normalizedDomain)") else { return false }
let cookiePath = cookie.path.isEmpty ? "/" : cookie.path
let requestPath = url.path.isEmpty ? "/" : url.path
if requestPath == cookiePath {
return true
}
guard requestPath.hasPrefix(cookiePath) else { return false }
guard cookiePath != "/" else { return true }
if cookiePath.hasSuffix("/") {
return true
}
guard
let boundaryIndex = requestPath.index(
requestPath.startIndex,
offsetBy: cookiePath.count,
limitedBy: requestPath.endIndex),
boundaryIndex < requestPath.endIndex
else {
return true
}
return requestPath[boundaryIndex] == "/"
}
private static func cookieSortKey(for cookie: HTTPCookie) -> (Int, Int, Date) {
let pathLength = cookie.path.count
let normalizedDomain = cookie.domain.lowercased().trimmingCharacters(in: CharacterSet(charactersIn: "."))
let domainLength = normalizedDomain.count
let expiry = cookie.expiresDate ?? .distantPast
return (pathLength, domainLength, expiry)
}
}
extension [String] {
fileprivate func uniquedSorted() -> [String] {
Array(Set(self)).sorted()
}
}
@@ -0,0 +1,289 @@
import Foundation
#if os(macOS)
import SweetCookieKit
#endif
public enum AlibabaTokenPlanProviderDescriptor {
public static let descriptor: ProviderDescriptor = Self.makeDescriptor()
static func makeDescriptor() -> ProviderDescriptor {
#if os(macOS)
let browserOrder: BrowserCookieImportOrder = [
.chrome,
.chromeBeta,
.brave,
.edge,
.arc,
.firefox,
.safari,
]
#else
let browserOrder: BrowserCookieImportOrder? = nil
#endif
return ProviderDescriptor(
id: .alibabatokenplan,
metadata: ProviderMetadata(
id: .alibabatokenplan,
displayName: "Alibaba Token Plan",
sessionLabel: "Credits",
weeklyLabel: "Usage",
opusLabel: nil,
supportsOpus: false,
supportsCredits: false,
creditsHint: "",
toggleTitle: "Show Alibaba Token Plan usage",
cliName: "alibaba-token-plan",
defaultEnabled: false,
isPrimaryProvider: false,
usesAccountFallback: false,
browserCookieOrder: browserOrder,
dashboardURL: AlibabaTokenPlanUsageFetcher.dashboardURL.absoluteString,
statusPageURL: nil,
statusLinkURL: "https://status.aliyun.com"),
branding: ProviderBranding(
iconStyle: .alibaba,
iconResourceName: "ProviderIcon-alibaba",
color: ProviderColor(red: 1.0, green: 106 / 255, blue: 0)),
tokenCost: ProviderTokenCostConfig(
supportsTokenCost: false,
noDataMessage: { "Alibaba Token Plan cost summary is not supported." }),
fetchPlan: ProviderFetchPlan(
sourceModes: [.auto, .web],
pipeline: ProviderFetchPipeline(resolveStrategies: self.resolveStrategies)),
cli: ProviderCLIConfig(
name: "alibaba-token-plan",
aliases: ["alibaba-token", "bailian-token-plan"],
versionDetector: nil))
}
private static func resolveStrategies(context: ProviderFetchContext) async -> [any ProviderFetchStrategy] {
guard context.settings?.alibabaTokenPlan?.cookieSource != .off else { return [] }
switch context.sourceMode {
case .auto, .web:
return [AlibabaTokenPlanWebFetchStrategy()]
case .api, .cli, .oauth:
return []
}
}
}
struct AlibabaTokenPlanWebFetchStrategy: ProviderFetchStrategy {
private static let log = CodexBarLog.logger("alibaba-token-plan")
let id: String = "alibaba-token-plan.web"
let kind: ProviderFetchKind = .web
func isAvailable(_ context: ProviderFetchContext) async -> Bool {
guard context.settings?.alibabaTokenPlan?.cookieSource != .off else { return false }
let region = context.settings?.alibabaTokenPlan?.apiRegion ?? .international
if AlibabaTokenPlanSettingsReader.cookieHeader(environment: context.env) != nil {
return true
}
if let settings = context.settings?.alibabaTokenPlan,
settings.cookieSource == .manual
{
return CookieHeaderNormalizer.normalize(settings.manualCookieHeader) != nil
}
#if os(macOS)
if let cached = Self.cachedCookieEntry(region: region),
!cached.cookieHeader.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty
{
return true
}
return true
#else
return false
#endif
}
func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult {
let cookieSource = context.settings?.alibabaTokenPlan?.cookieSource ?? .auto
let region = context.settings?.alibabaTokenPlan?.apiRegion ?? .international
let cookieHeaders = try Self.resolveCookieHeaders(context: context, allowCached: true, region: region)
do {
let usage = try await AlibabaTokenPlanUsageFetcher.fetchUsage(
apiCookieHeader: cookieHeaders.apiCookieHeader,
dashboardCookieHeader: cookieHeaders.dashboardCookieHeader,
region: region,
environment: context.env)
return self.makeResult(usage: usage.toUsageSnapshot(), sourceLabel: "web")
} catch let error as AlibabaTokenPlanUsageError
where error.isCredentialFailure && cookieSource != .manual
{
#if os(macOS)
CookieHeaderCache.clear(provider: .alibabatokenplan, scope: region.cookieCacheScope)
let refreshedHeaders = try Self.resolveCookieHeaders(context: context, allowCached: false, region: region)
let usage = try await AlibabaTokenPlanUsageFetcher.fetchUsage(
apiCookieHeader: refreshedHeaders.apiCookieHeader,
dashboardCookieHeader: refreshedHeaders.dashboardCookieHeader,
region: region,
environment: context.env)
return self.makeResult(usage: usage.toUsageSnapshot(), sourceLabel: "web")
#else
throw error
#endif
}
}
func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool {
false
}
static func resolveCookieHeader(context: ProviderFetchContext, allowCached: Bool) throws -> String {
try self.resolveCookieHeaders(context: context, allowCached: allowCached, region: .international)
.apiCookieHeader
}
static func resolveCookieHeaders(
context: ProviderFetchContext,
allowCached: Bool,
region: AlibabaTokenPlanAPIRegion = .international) throws -> AlibabaTokenPlanCookieHeaders
{
if let settings = context.settings?.alibabaTokenPlan,
settings.cookieSource == .manual
{
guard let headers = AlibabaTokenPlanCookieHeaders(singleHeader: settings.manualCookieHeader) else {
self.log.warning("Alibaba Token Plan manual cookie header is invalid")
throw AlibabaTokenPlanSettingsError.invalidCookie
}
Self.log.info(
"Alibaba Token Plan using manual cookie header",
metadata: [
"apiCookieNames": headers.apiCookieNames.joined(separator: ","),
"dashboardCookieNames": headers.dashboardCookieNames.joined(separator: ","),
"hasSecToken": headers.hasCookie(named: "sec_token") ? "1" : "0",
])
return headers
}
if let envCookie = AlibabaTokenPlanSettingsReader.cookieHeader(environment: context.env),
let headers = AlibabaTokenPlanCookieHeaders(singleHeader: envCookie)
{
Self.log.info(
"Alibaba Token Plan using environment cookie header",
metadata: [
"apiCookieNames": headers.apiCookieNames.joined(separator: ","),
"dashboardCookieNames": headers.dashboardCookieNames.joined(separator: ","),
"hasSecToken": headers.hasCookie(named: "sec_token") ? "1" : "0",
])
return headers
}
#if os(macOS)
if allowCached,
let cached = Self.cachedCookieEntry(region: region),
let headers = AlibabaTokenPlanCookieHeaders(cachedHeader: cached.cookieHeader)
{
Self.log.info(
"Alibaba Token Plan using cached browser cookie header",
metadata: [
"source": cached.sourceLabel,
"apiCookieNames": headers.apiCookieNames.joined(separator: ","),
"dashboardCookieNames": headers.dashboardCookieNames.joined(separator: ","),
"hasSecToken": headers.hasCookie(named: "sec_token") ? "1" : "0",
])
return headers
}
do {
var importLog: [String] = []
let session = try AlibabaCodingPlanCookieImporter.importSession(
browserDetection: context.browserDetection,
logger: { importLog.append($0) })
let rawCookieNames = session.cookies.map(\.name).filter { !$0.isEmpty }.uniquedSorted()
guard let headers = AlibabaTokenPlanCookieHeader.headers(
from: session.cookies,
region: region,
environment: context.env)
else {
Self.log.warning(
"Alibaba Token Plan browser cookie header was empty",
metadata: [
"source": session.sourceLabel,
"rawCookieNames": rawCookieNames.joined(separator: ","),
])
throw AlibabaTokenPlanSettingsError.missingCookie(
details: "No Alibaba Token Plan browser cookies were available after import.")
}
CookieHeaderCache.store(
provider: .alibabatokenplan,
scope: region.cookieCacheScope,
cookieHeader: headers.cacheCookieHeader,
sourceLabel: session.sourceLabel)
Self.log.info(
"Alibaba Token Plan imported browser cookies",
metadata: [
"source": session.sourceLabel,
"rawCookieNames": rawCookieNames.joined(separator: ","),
"apiCookieNames": headers.apiCookieNames.joined(separator: ","),
"dashboardCookieNames": headers.dashboardCookieNames.joined(separator: ","),
"hasSecToken": headers.hasCookie(named: "sec_token") ? "1" : "0",
"importLogLines": "\(importLog.count)",
])
return headers
} catch {
Self.log.warning(
"Alibaba Token Plan cookie resolution failed",
metadata: ["error": error.localizedDescription])
throw AlibabaTokenPlanSettingsError.missingCookie(details: Self.missingCookieDetails(from: error))
}
#else
throw AlibabaTokenPlanSettingsError.missingCookie()
#endif
}
#if os(macOS)
/// The former unscoped cache only ever represented the China gateway. Never expose it to
/// International requests; migrate it into the China scope after a successful scoped write.
private static func cachedCookieEntry(region: AlibabaTokenPlanAPIRegion) -> CookieHeaderCache.Entry? {
if let scoped = CookieHeaderCache.load(provider: .alibabatokenplan, scope: region.cookieCacheScope) {
return scoped
}
guard region == .chinaMainland,
let legacy = CookieHeaderCache.load(provider: .alibabatokenplan)
else { return nil }
CookieHeaderCache.store(
provider: .alibabatokenplan,
scope: region.cookieCacheScope,
cookieHeader: legacy.cookieHeader,
sourceLabel: legacy.sourceLabel,
now: legacy.storedAt)
if let migrated = CookieHeaderCache.load(provider: .alibabatokenplan, scope: region.cookieCacheScope) {
CookieHeaderCache.clear(provider: .alibabatokenplan)
return migrated
}
return legacy
}
#endif
private static func missingCookieDetails(from error: Error) -> String? {
if case let AlibabaCodingPlanSettingsError.missingCookie(details) = error {
return details
}
let message = error.localizedDescription.trimmingCharacters(in: .whitespacesAndNewlines)
return message.isEmpty ? nil : message
}
}
extension [String] {
fileprivate func uniquedSorted() -> [String] {
Array(Set(self)).sorted()
}
}
extension AlibabaTokenPlanUsageError {
fileprivate var isCredentialFailure: Bool {
switch self {
case .loginRequired, .invalidCredentials:
true
case .apiError, .networkError, .parseFailed:
false
}
}
}
@@ -0,0 +1,67 @@
import Foundation
public struct AlibabaTokenPlanSettingsReader: Sendable {
public static let cookieHeaderKey = "ALIBABA_TOKEN_PLAN_COOKIE"
public static let hostKey = "ALIBABA_TOKEN_PLAN_HOST"
public static let quotaURLKey = "ALIBABA_TOKEN_PLAN_QUOTA_URL"
private static let endpointValidator = ProviderEndpointOverrideValidator(
allowedHosts: [
"modelstudio.console.alibabacloud.com",
"bailian.console.aliyun.com",
])
public static func cookieHeader(
environment: [String: String] = ProcessInfo.processInfo.environment) -> String?
{
self.cleaned(environment[self.cookieHeaderKey])
}
public static func hostOverride(
environment: [String: String] = ProcessInfo.processInfo.environment) -> String?
{
self.endpointValidator.validatedHost(
self.cleaned(environment[self.hostKey]),
policy: .allowAnyHTTPSHost)
}
public static func quotaURL(
environment: [String: String] = ProcessInfo.processInfo.environment) -> URL?
{
self.endpointValidator.validatedURL(
self.cleaned(environment[self.quotaURLKey]),
policy: .allowAnyHTTPSHost)
}
static func cleaned(_ raw: String?) -> String? {
guard var value = raw?.trimmingCharacters(in: .whitespacesAndNewlines), !value.isEmpty else {
return nil
}
if (value.hasPrefix("\"") && value.hasSuffix("\"")) ||
(value.hasPrefix("'") && value.hasSuffix("'"))
{
value = String(value.dropFirst().dropLast())
}
value = value.trimmingCharacters(in: .whitespacesAndNewlines)
return value.isEmpty ? nil : value
}
}
public enum AlibabaTokenPlanSettingsError: LocalizedError, Sendable {
case missingCookie(details: String? = nil)
case invalidCookie
public var errorDescription: String? {
switch self {
case let .missingCookie(details):
let base = "No Alibaba Token Plan session cookies found in browsers. " +
"Sign in to Model Studio/Bailian in Chrome, " +
"allow CodexBar to access Chrome Safe Storage in Keychain Access, " +
"or paste a manual Cookie header."
guard let details, !details.isEmpty else { return base }
return "\(base) \(details)"
case .invalidCookie:
return "Alibaba Token Plan cookie header is invalid."
}
}
}
@@ -0,0 +1,997 @@
import Foundation
#if canImport(FoundationNetworking)
import FoundationNetworking
#endif
public enum AlibabaTokenPlanUsageError: LocalizedError, Sendable, Equatable {
case loginRequired
case invalidCredentials
case apiError(String)
case networkError(String)
case parseFailed(String)
public var errorDescription: String? {
switch self {
case .loginRequired:
"Alibaba Token Plan login required."
case .invalidCredentials:
"Alibaba Token Plan credentials are invalid."
case let .apiError(message):
"Alibaba Token Plan API error: \(message)"
case let .networkError(message):
"Alibaba Token Plan network error: \(message)"
case let .parseFailed(message):
"Could not parse Alibaba Token Plan usage: \(message)"
}
}
}
// swiftlint:disable:next type_body_length
public struct AlibabaTokenPlanUsageFetcher: Sendable {
private static let log = CodexBarLog.logger("alibaba-token-plan")
private static let bssServiceCode = "BssOpenAPI-V3"
private static let subscriptionSummaryAction = "GetSubscriptionSummary"
private static let browserLikeUserAgent =
"Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) " +
"AppleWebKit/537.36 (KHTML, like Gecko) Chrome/143.0.0.0 Safari/537.36"
private static let safariLikeUserAgent =
"Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) " +
"AppleWebKit/605.1.15 (KHTML, like Gecko) Version/26.3 Safari/605.1.15"
public static var dashboardURL: URL {
Self.dashboardURL(region: .international, environment: ProcessInfo.processInfo.environment)
}
public static func dashboardURL(
region: AlibabaTokenPlanAPIRegion,
environment: [String: String] = ProcessInfo.processInfo.environment) -> URL
{
if let host = AlibabaTokenPlanSettingsReader.hostOverride(environment: environment),
let base = ProviderEndpointOverrideValidator.normalizedHTTPSURL(from: host),
var components = URLComponents(url: base, resolvingAgainstBaseURL: false),
let dashboardComponents = URLComponents(url: region.dashboardURL, resolvingAgainstBaseURL: false)
{
components.path = dashboardComponents.path
components.percentEncodedQuery = dashboardComponents.percentEncodedQuery
components.fragment = dashboardComponents.fragment
return components.url ?? region.dashboardURL
}
return region.dashboardURL
}
public static func fetchUsage(
cookieHeader: String,
region: AlibabaTokenPlanAPIRegion = .international,
environment: [String: String] = ProcessInfo.processInfo.environment,
now: Date = Date()) async throws -> AlibabaTokenPlanUsageSnapshot
{
guard let headers = AlibabaTokenPlanCookieHeaders(singleHeader: cookieHeader) else {
throw AlibabaTokenPlanSettingsError.invalidCookie
}
return try await self.fetchUsage(
apiCookieHeader: headers.apiCookieHeader,
dashboardCookieHeader: headers.dashboardCookieHeader,
region: region,
environment: environment,
now: now)
}
static func fetchUsage(
apiCookieHeader: String,
dashboardCookieHeader: String,
region: AlibabaTokenPlanAPIRegion = .international,
environment: [String: String] = ProcessInfo.processInfo.environment,
now: Date = Date(),
session overrideSession: URLSession? = nil) async throws -> AlibabaTokenPlanUsageSnapshot
{
guard let normalizedAPIHeader = CookieHeaderNormalizer.normalize(apiCookieHeader),
let normalizedDashboardHeader = CookieHeaderNormalizer.normalize(dashboardCookieHeader)
else {
throw AlibabaTokenPlanSettingsError.invalidCookie
}
let url = self.resolveQuotaURL(region: region, environment: environment)
let apiRedirectDiagnostics = RedirectDiagnostics(cookieHeader: normalizedAPIHeader)
let dashboardRedirectDiagnostics: RedirectDiagnostics?
let apiSession: URLSession
let dashboardSession: URLSession
if let overrideSession {
apiSession = overrideSession
dashboardSession = overrideSession
dashboardRedirectDiagnostics = nil
} else {
let dashboardDiagnostics = RedirectDiagnostics(cookieHeader: normalizedDashboardHeader)
apiSession = URLSession(
configuration: .default,
delegate: apiRedirectDiagnostics,
delegateQueue: nil)
dashboardSession = URLSession(
configuration: .default,
delegate: dashboardDiagnostics,
delegateQueue: nil)
dashboardRedirectDiagnostics = dashboardDiagnostics
}
defer {
if overrideSession == nil {
apiSession.invalidateAndCancel()
dashboardSession.invalidateAndCancel()
}
}
let secToken = await self.resolveSECToken(
dashboardCookieHeader: normalizedDashboardHeader,
apiCookieHeader: normalizedAPIHeader,
region: region,
environment: environment,
session: dashboardSession)
Self.log.info(
"Fetching Alibaba Token Plan usage",
metadata: [
"apiHost": url.host ?? "unknown",
"region": region.rawValue,
"apiCookieNames": self.cookieNamesDescription(from: normalizedAPIHeader),
"dashboardCookieNames": self.cookieNamesDescription(from: normalizedDashboardHeader),
"hasCSRF": self.hasCSRF(in: normalizedAPIHeader) ? "1" : "0",
"secTokenSource": secToken == nil ? "missing" : "resolved",
])
var request = URLRequest(url: url)
request.httpMethod = "POST"
request.timeoutInterval = 20
request.httpBody = self.subscriptionSummaryRequestBody(region: region, secToken: secToken)
request.setValue("application/x-www-form-urlencoded", forHTTPHeaderField: "Content-Type")
request.setValue("*/*", forHTTPHeaderField: "Accept")
request.setValue(normalizedAPIHeader, forHTTPHeaderField: "Cookie")
if let csrf = self.extractCookieValue(name: "login_aliyunid_csrf", from: normalizedAPIHeader) ??
self.extractCookieValue(name: "csrf", from: normalizedAPIHeader)
{
request.setValue(csrf, forHTTPHeaderField: "x-xsrf-token")
request.setValue(csrf, forHTTPHeaderField: "x-csrf-token")
}
request.setValue("XMLHttpRequest", forHTTPHeaderField: "X-Requested-With")
request.setValue(Self.browserLikeUserAgent, forHTTPHeaderField: "User-Agent")
request.setValue(region.dashboardOriginURLString, forHTTPHeaderField: "Origin")
request.setValue(
Self.dashboardURL(region: region, environment: environment).absoluteString,
forHTTPHeaderField: "Referer")
let data: Data
let response: URLResponse
do {
(data, response) = try await apiSession.data(for: request)
} catch {
Self.log.error(
"Alibaba Token Plan request failed",
metadata: [
"apiHost": url.host ?? "unknown",
"error": error.localizedDescription,
])
throw AlibabaTokenPlanUsageError.networkError(error.localizedDescription)
}
if let dashboardRedirectDiagnostics, !dashboardRedirectDiagnostics.redirects.isEmpty {
Self.log.info(
"Alibaba Token Plan dashboard redirects",
metadata: [
"count": "\(dashboardRedirectDiagnostics.redirects.count)",
"items": dashboardRedirectDiagnostics.redirects.joined(separator: " | "),
])
}
if !apiRedirectDiagnostics.redirects.isEmpty {
Self.log.info(
"Alibaba Token Plan redirects",
metadata: [
"count": "\(apiRedirectDiagnostics.redirects.count)",
"items": apiRedirectDiagnostics.redirects.joined(separator: " | "),
])
}
guard let httpResponse = response as? HTTPURLResponse else {
Self.log.error("Alibaba Token Plan response was not HTTP")
throw AlibabaTokenPlanUsageError.networkError("Invalid response")
}
Self.log.info(
"Alibaba Token Plan HTTP response",
metadata: [
"status": "\(httpResponse.statusCode)",
"bodyBytes": "\(data.count)",
"contentType": httpResponse.value(forHTTPHeaderField: "Content-Type") ?? "none",
])
guard httpResponse.statusCode == 200 else {
if httpResponse.statusCode == 401 || httpResponse.statusCode == 403 {
throw AlibabaTokenPlanUsageError.loginRequired
}
Self.log.error("Alibaba Token Plan returned HTTP \(httpResponse.statusCode)")
throw AlibabaTokenPlanUsageError.apiError("HTTP \(httpResponse.statusCode)")
}
return try self.parseUsageSnapshot(from: data, now: now)
}
static func resolveQuotaURL(
region: AlibabaTokenPlanAPIRegion,
environment: [String: String]) -> URL
{
if let override = AlibabaTokenPlanSettingsReader.quotaURL(environment: environment) {
return override
}
if let host = AlibabaTokenPlanSettingsReader.hostOverride(environment: environment),
let hostURL = self.quotaURL(from: host)
{
return hostURL
}
return self.defaultQuotaURL(region: region)
}
static var defaultQuotaURL: URL {
self.defaultQuotaURL(region: .international)
}
static func defaultQuotaURL(region: AlibabaTokenPlanAPIRegion) -> URL {
var components = URLComponents(string: region.gatewayBaseURLString)!
components.path = "/data/api.json"
components.queryItems = [
URLQueryItem(name: "action", value: Self.subscriptionSummaryAction),
URLQueryItem(name: "product", value: Self.bssServiceCode),
URLQueryItem(name: "_tag", value: ""),
]
return components.url!
}
static func parseUsageSnapshot(from data: Data, now: Date = Date()) throws -> AlibabaTokenPlanUsageSnapshot {
guard !data.isEmpty else {
throw AlibabaTokenPlanUsageError.parseFailed("Empty response body")
}
let object: Any
do {
object = try JSONSerialization.jsonObject(with: data, options: [])
} catch {
if self.isLikelyLoginHTML(data) {
throw AlibabaTokenPlanUsageError.loginRequired
}
throw AlibabaTokenPlanUsageError.parseFailed("Invalid JSON response")
}
let expanded = self.expandedJSON(object)
guard let dictionary = expanded as? [String: Any] else {
throw AlibabaTokenPlanUsageError.parseFailed("Unexpected payload")
}
try self.throwIfErrorPayload(dictionary)
let summary = self.findSubscriptionSummary(in: dictionary) ?? dictionary
let total = self.anyDouble(for: Self.totalQuotaKeys, in: summary)
let remaining = self.anyDouble(for: Self.remainingQuotaKeys, in: summary)
let used = self.anyDouble(for: Self.usedQuotaKeys, in: summary) ??
total.flatMap { total in remaining.map { max(0, total - $0) } }
let resetsAt = self.findResetDate(in: summary) ?? self.findResetDate(in: dictionary)
let totalCount = self.anyDouble(for: Self.subscriptionCountKeys, in: summary)
let planName = self.findPlanName(in: summary) ?? ((totalCount ?? 0) > 0 || total != nil ? "TOKEN PLAN" : nil)
if planName == nil, total == nil, used == nil, remaining == nil, totalCount == nil {
let diagnostics = self.payloadDiagnostics(payload: dictionary)
Self.log.error("Alibaba Token Plan payload missing expected fields: \(diagnostics)")
throw AlibabaTokenPlanUsageError.parseFailed("Missing token plan data (\(diagnostics))")
}
return AlibabaTokenPlanUsageSnapshot(
planName: planName,
usedQuota: used,
totalQuota: total,
remainingQuota: remaining,
resetsAt: resetsAt,
updatedAt: now)
}
private static func subscriptionSummaryRequestBody(region: AlibabaTokenPlanAPIRegion, secToken: String?) -> Data {
let paramsObject = ["ProductCode": region.tokenPlanProductCode]
guard let paramsData = try? JSONSerialization.data(withJSONObject: paramsObject, options: []),
let paramsString = String(data: paramsData, encoding: .utf8)
else {
return Data()
}
var components = URLComponents()
var queryItems = [
URLQueryItem(name: "product", value: Self.bssServiceCode),
URLQueryItem(name: "action", value: Self.subscriptionSummaryAction),
URLQueryItem(name: "params", value: paramsString),
URLQueryItem(name: "region", value: region.currentRegionID),
]
if let secToken, !secToken.isEmpty {
queryItems.append(URLQueryItem(name: "sec_token", value: secToken))
}
components.queryItems = queryItems
return Data((components.percentEncodedQuery ?? "").utf8)
}
private static func resolveSECToken(
dashboardCookieHeader: String,
apiCookieHeader: String,
region: AlibabaTokenPlanAPIRegion,
environment: [String: String],
session: URLSession) async -> String?
{
let cookieSECToken = self.extractCookieValue(name: "sec_token", from: dashboardCookieHeader) ??
self.extractCookieValue(name: "sec_token", from: apiCookieHeader)
var request = URLRequest(url: self.dashboardURL(region: region, environment: environment))
request.httpMethod = "GET"
request.timeoutInterval = 10
request.setValue(dashboardCookieHeader, forHTTPHeaderField: "Cookie")
request.setValue(Self.safariLikeUserAgent, forHTTPHeaderField: "User-Agent")
request.setValue(
"text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8",
forHTTPHeaderField: "Accept")
if let (data, response) = try? await session.data(for: request),
let httpResponse = response as? HTTPURLResponse,
httpResponse.statusCode == 200,
let html = String(data: data, encoding: .utf8),
let token = self.extractSECToken(from: html)
{
Self.log.info(
"Resolved Alibaba Token Plan sec_token from dashboard HTML",
metadata: [
"dashboardHost": request.url?.host ?? "unknown",
"htmlBytes": "\(data.count)",
])
return token
}
if let token = await self.fetchSECTokenFromUserInfo(
cookieHeader: dashboardCookieHeader,
region: region,
environment: environment,
session: session)
{
return token
}
if let cookieSECToken, !cookieSECToken.isEmpty {
Self.log.info("Resolved Alibaba Token Plan sec_token from cookies")
return cookieSECToken
}
Self.log.info(
"Alibaba Token Plan sec_token missing; continuing with cookie-only request",
metadata: [
"dashboardCookieNames": self.cookieNamesDescription(from: dashboardCookieHeader),
"apiCookieNames": self.cookieNamesDescription(from: apiCookieHeader),
])
return nil
}
private static func fetchSECTokenFromUserInfo(
cookieHeader: String,
region: AlibabaTokenPlanAPIRegion,
environment: [String: String],
session: URLSession) async -> String?
{
let baseURL = self.consoleBaseURL(region: region, environment: environment)
let userInfoURL = baseURL.appendingPathComponent("tool/user/info.json")
var request = URLRequest(url: userInfoURL)
request.httpMethod = "GET"
request.timeoutInterval = 10
request.setValue(cookieHeader, forHTTPHeaderField: "Cookie")
request.setValue(Self.safariLikeUserAgent, forHTTPHeaderField: "User-Agent")
request.setValue("application/json, text/plain, */*", forHTTPHeaderField: "Accept")
let referer = baseURL.absoluteString.hasSuffix("/") ? baseURL.absoluteString : "\(baseURL.absoluteString)/"
request.setValue(referer, forHTTPHeaderField: "Referer")
guard let (data, response) = try? await session.data(for: request),
let httpResponse = response as? HTTPURLResponse,
httpResponse.statusCode == 200,
let object = try? JSONSerialization.jsonObject(with: data, options: [])
else {
return nil
}
let expanded = self.expandedJSON(object)
guard let token = self.findFirstString(forKeys: ["secToken", "sec_token"], in: expanded),
!token.isEmpty
else {
return nil
}
Self.log.info(
"Resolved Alibaba Token Plan sec_token from user info",
metadata: [
"userInfoHost": userInfoURL.host ?? "unknown",
"bodyBytes": "\(data.count)",
])
return token
}
private static func consoleBaseURL(
region: AlibabaTokenPlanAPIRegion,
environment: [String: String]) -> URL
{
let dashboard = self.dashboardURL(region: region, environment: environment)
var components = URLComponents()
components.scheme = dashboard.scheme
components.host = dashboard.host
components.port = dashboard.port
return components.url ?? URL(string: region.dashboardOriginURLString)!
}
private static func quotaURL(from rawHost: String) -> URL? {
let cleaned = AlibabaTokenPlanSettingsReader.cleaned(rawHost)
guard let cleaned else { return nil }
guard let base = ProviderEndpointOverrideValidator.normalizedHTTPSURL(from: cleaned) else { return nil }
var components = URLComponents(url: base, resolvingAgainstBaseURL: false)
let defaultComponents = URLComponents(
url: Self.defaultQuotaURL(region: .international),
resolvingAgainstBaseURL: false)
components?.path = "/data/api.json"
components?.queryItems = defaultComponents?.queryItems
return components?.url
}
private final class RedirectDiagnostics: NSObject, URLSessionTaskDelegate, @unchecked Sendable {
private let cookieHeader: String
var redirects: [String] = []
init(cookieHeader: String) {
self.cookieHeader = cookieHeader
}
func urlSession(
_: URLSession,
task _: URLSessionTask,
willPerformHTTPRedirection response: HTTPURLResponse,
newRequest request: URLRequest,
completionHandler: @escaping (URLRequest?) -> Void)
{
let from = AlibabaTokenPlanUsageFetcher.redactedURLDescription(response.url)
let to = AlibabaTokenPlanUsageFetcher.redactedURLDescription(request.url)
self.redirects.append("\(response.statusCode) \(from) -> \(to)")
completionHandler(AlibabaTokenPlanUsageFetcher.redirectedRequest(
response: response,
request: request,
cookieHeader: self.cookieHeader))
}
}
private static func redactedURLDescription(_ url: URL?) -> String {
guard let url else { return "unknown" }
var components = URLComponents(url: url, resolvingAgainstBaseURL: false)
components?.query = nil
components?.fragment = nil
return components?.string ?? "\(url.scheme ?? "unknown")://\(url.host ?? "unknown")"
}
static func redirectedRequest(
response: HTTPURLResponse,
request: URLRequest,
cookieHeader: String) -> URLRequest?
{
guard request.url?.scheme?.lowercased() == "https" else {
return nil
}
var updated = request
if self.shouldForwardRedirectCookies(from: response.url, to: request.url) {
updated.setValue(cookieHeader, forHTTPHeaderField: "Cookie")
} else {
updated.setValue(nil, forHTTPHeaderField: "Cookie")
}
return updated
}
private static func shouldForwardRedirectCookies(from sourceURL: URL?, to targetURL: URL?) -> Bool {
guard let sourceHost = sourceURL?.host?.lowercased(),
let targetHost = targetURL?.host?.lowercased()
else {
return false
}
return sourceHost == targetHost
}
private static func throwIfErrorPayload(_ dictionary: [String: Any]) throws {
if self.parseBool(dictionary["successResponse"]) == false {
if let statusCode = self.findFirstInt(forKeys: ["statusCode", "status_code", "code"], in: dictionary),
statusCode == 401 || statusCode == 403
{
throw AlibabaTokenPlanUsageError.invalidCredentials
}
let code = self.findFirstString(forKeys: ["code", "status", "statusCode"], in: dictionary)
let message = self.findFirstString(forKeys: ["message", "msg", "statusMessage"], in: dictionary) ??
code ??
"request was not successful"
if self.isLoginOrTokenError(code: code, message: message) {
throw AlibabaTokenPlanUsageError.loginRequired
}
throw AlibabaTokenPlanUsageError.apiError(message)
}
if self.findBoolValues(forKeys: ["Success", "success"], in: dictionary).contains(false) {
let code = self.findFirstString(forKeys: ["Code", "code"], in: dictionary)
let message = self.findFirstString(forKeys: ["Message", "message", "msg", "Code", "code"], in: dictionary)
?? "request was not successful"
if self.isLoginOrTokenError(code: code, message: message) {
throw AlibabaTokenPlanUsageError.loginRequired
}
throw AlibabaTokenPlanUsageError.apiError(message)
}
if let statusCode = self.findFirstInt(forKeys: ["statusCode", "status_code", "code"], in: dictionary),
statusCode != 0,
statusCode != 200
{
let message = self.findFirstString(
forKeys: ["statusMessage", "status_msg", "message", "msg"],
in: dictionary)
?? "status code \(statusCode)"
if statusCode == 401 || statusCode == 403 {
throw AlibabaTokenPlanUsageError.invalidCredentials
}
throw AlibabaTokenPlanUsageError.apiError(message)
}
let codeText = self.findFirstString(forKeys: ["code", "status", "statusCode"], in: dictionary)?.lowercased()
let messageText = self.findFirstString(forKeys: ["message", "msg", "statusMessage"], in: dictionary)?
.lowercased()
if self.isLoginOrTokenError(code: codeText, message: messageText) {
throw AlibabaTokenPlanUsageError.loginRequired
}
}
private static func isLoginOrTokenError(code: String?, message: String?) -> Bool {
let combined = [code, message]
.compactMap { $0?.lowercased() }
.joined(separator: " ")
return combined.contains("needlogin") ||
combined.contains("login") ||
combined.contains("postonlyortokenerror") ||
combined.contains("tokenerror") ||
combined.contains("request has expired") ||
combined.contains("refresh page") ||
combined.contains("请求已经过期")
}
private static let planNameKeys = [
"planName",
"plan_name",
"packageName",
"package_name",
"commodityName",
"commodity_name",
"instanceName",
"instance_name",
"displayName",
"display_name",
"ProductName",
"productName",
"name",
"title",
"planType",
"plan_type",
]
private static let usedQuotaKeys = [
"usedQuota",
"used_quota",
"usedCredits",
"usedCredit",
"consumedCredits",
"usage",
"used",
"usedAmount",
"consumeAmount",
"usedValue",
"UsedValue",
"consumedValue",
"ConsumedValue",
]
private static let totalQuotaKeys = [
"totalQuota",
"total_quota",
"totalCredits",
"totalCredit",
"quota",
"creditLimit",
"creditsTotal",
"monthlyTotalQuota",
"amount",
"totalValue",
"TotalValue",
]
private static let remainingQuotaKeys = [
"remainingQuota",
"remainQuota",
"remainingCredits",
"remainingCredit",
"availableCredits",
"balance",
"remaining",
"availableAmount",
"remainAmount",
"totalSurplusValue",
"TotalSurplusValue",
"surplusValue",
"SurplusValue",
]
private static let subscriptionCountKeys = [
"totalCount",
"TotalCount",
"subscriptionTotalNumber",
"SubscriptionTotalNumber",
]
private static let resetDateKeys = [
"nextRefreshTime",
"resetTime",
"periodEndTime",
"billingCycleEnd",
"billCycleEndTime",
"expireTime",
"expirationTime",
"endTime",
"validEndTime",
"instanceEndTime",
"nearestExpireDate",
"NearestExpireDate",
]
private static func findSubscriptionSummary(in payload: [String: Any]) -> [String: Any]? {
if let data = self.findFirstDictionary(
forKeys: ["Data", "data", "successResponse", "success_response"],
in: payload),
self.containsSubscriptionSummaryFields(data)
{
return data
}
return self.findFirstDictionary(
matchingAnyKey: Self.usedQuotaKeys + Self.totalQuotaKeys + Self.remainingQuotaKeys +
Self.subscriptionCountKeys,
in: payload)
}
private static func containsSubscriptionSummaryFields(_ payload: [String: Any]) -> Bool {
let keys = self.usedQuotaKeys + self.totalQuotaKeys + self.remainingQuotaKeys + self.subscriptionCountKeys
return keys.contains { payload[$0] != nil }
}
private static func findPlanName(in payload: [String: Any]) -> String? {
self.anyString(for: self.planNameKeys, in: payload) ??
self.findFirstString(forKeys: self.planNameKeys, in: payload)
}
private static func findResetDate(in payload: [String: Any]) -> Date? {
self.anyDate(for: self.resetDateKeys, in: payload) ??
self.findFirstDate(forKeys: self.resetDateKeys, in: payload)
}
private static func payloadDiagnostics(payload: [String: Any]) -> String {
let topKeys = payload.keys.sorted()
let dataDict = self.findFirstDictionary(
forKeys: ["Data", "data", "successResponse", "success_response"],
in: payload)
let dataKeys = dataDict?.keys.sorted() ?? []
return "topKeys=\(topKeys.joined(separator: ",")) dataKeys=\(dataKeys.joined(separator: ","))"
}
private static func isLikelyLoginHTML(_ data: Data) -> Bool {
guard let text = String(data: data, encoding: .utf8)?.lowercased() else { return false }
return text.contains("<html") &&
(text.contains("login") || text.contains("sign in") || text.contains("signin"))
}
private static func findFirstDictionary(forKeys keys: [String], in value: Any) -> [String: Any]? {
if let dict = value as? [String: Any] {
for key in keys {
if let nested = dict[key] as? [String: Any] {
return nested
}
}
for nestedValue in dict.values {
if let nested = self.findFirstDictionary(forKeys: keys, in: nestedValue) {
return nested
}
}
return nil
}
if let array = value as? [Any] {
for item in array {
if let nested = self.findFirstDictionary(forKeys: keys, in: item) {
return nested
}
}
}
return nil
}
private static func findFirstDictionary(matchingAnyKey keys: [String], in value: Any) -> [String: Any]? {
if let dict = value as? [String: Any] {
if keys.contains(where: { dict[$0] != nil }) {
return dict
}
for nestedValue in dict.values {
if let nested = self.findFirstDictionary(matchingAnyKey: keys, in: nestedValue) {
return nested
}
}
return nil
}
if let array = value as? [Any] {
for item in array {
if let nested = self.findFirstDictionary(matchingAnyKey: keys, in: item) {
return nested
}
}
}
return nil
}
private static func findFirstString(forKeys keys: [String], in value: Any) -> String? {
if let dict = value as? [String: Any] {
for key in keys {
if let parsed = self.parseString(dict[key]) {
return parsed
}
}
for nestedValue in dict.values {
if let parsed = self.findFirstString(forKeys: keys, in: nestedValue) {
return parsed
}
}
return nil
}
if let array = value as? [Any] {
for item in array {
if let parsed = self.findFirstString(forKeys: keys, in: item) {
return parsed
}
}
}
return nil
}
private static func findBoolValues(forKeys keys: [String], in value: Any) -> [Bool] {
if let dict = value as? [String: Any] {
let directValues = keys.compactMap { self.parseBool(dict[$0]) }
let nestedValues = dict.values.flatMap { self.findBoolValues(forKeys: keys, in: $0) }
return directValues + nestedValues
}
if let array = value as? [Any] {
return array.flatMap { self.findBoolValues(forKeys: keys, in: $0) }
}
return []
}
private static func findFirstInt(forKeys keys: [String], in value: Any) -> Int? {
if let dict = value as? [String: Any] {
for key in keys {
if let parsed = self.parseInt(dict[key]) {
return parsed
}
}
for nestedValue in dict.values {
if let parsed = self.findFirstInt(forKeys: keys, in: nestedValue) {
return parsed
}
}
return nil
}
if let array = value as? [Any] {
for item in array {
if let parsed = self.findFirstInt(forKeys: keys, in: item) {
return parsed
}
}
}
return nil
}
private static func findFirstDate(forKeys keys: [String], in value: Any) -> Date? {
if let dict = value as? [String: Any] {
for key in keys {
if let parsed = self.parseDate(dict[key]) {
return parsed
}
}
for nestedValue in dict.values {
if let parsed = self.findFirstDate(forKeys: keys, in: nestedValue) {
return parsed
}
}
return nil
}
if let array = value as? [Any] {
for item in array {
if let parsed = self.findFirstDate(forKeys: keys, in: item) {
return parsed
}
}
}
return nil
}
private static func expandedJSON(_ value: Any) -> Any {
if let dict = value as? [String: Any] {
var expanded: [String: Any] = [:]
expanded.reserveCapacity(dict.count)
for (key, nested) in dict {
expanded[key] = self.expandedJSON(nested)
}
return expanded
}
if let array = value as? [Any] {
return array.map { self.expandedJSON($0) }
}
if let string = value as? String,
let data = string.data(using: .utf8),
let nested = try? JSONSerialization.jsonObject(with: data, options: []),
nested is [String: Any] || nested is [Any]
{
return self.expandedJSON(nested)
}
return value
}
private static func anyString(for keys: [String], in dict: [String: Any]) -> String? {
for key in keys {
if let value = self.parseString(dict[key]) {
return value
}
}
return nil
}
private static func anyDouble(for keys: [String], in dict: [String: Any]) -> Double? {
for key in keys {
if let value = self.parseDouble(dict[key]) {
return value
}
}
return nil
}
private static func anyDate(for keys: [String], in dict: [String: Any]) -> Date? {
for key in keys {
if let value = self.parseDate(dict[key]) {
return value
}
}
return nil
}
private static func anyBool(for keys: [String], in dict: [String: Any]) -> Bool? {
for key in keys {
if let value = self.parseBool(dict[key]) {
return value
}
}
return nil
}
private static func parseInt(_ raw: Any?) -> Int? {
if let value = raw as? Int { return value }
if let value = raw as? Int64 { return Int(value) }
if let value = raw as? Double { return Int(value) }
if let value = raw as? NSNumber { return value.intValue }
if let value = self.parseString(raw) { return Int(value) }
return nil
}
private static func parseDouble(_ raw: Any?) -> Double? {
if let value = raw as? Double { return value }
if let value = raw as? Int { return Double(value) }
if let value = raw as? Int64 { return Double(value) }
if let value = raw as? NSNumber { return value.doubleValue }
if let value = self.parseString(raw) {
let cleaned = value.replacingOccurrences(of: ",", with: "")
return Double(cleaned)
}
return nil
}
private static func parseString(_ raw: Any?) -> String? {
guard let value = raw as? String else { return nil }
let trimmed = value.trimmingCharacters(in: .whitespacesAndNewlines)
return trimmed.isEmpty ? nil : trimmed
}
private static func parseDate(_ raw: Any?) -> Date? {
if let intValue = self.parseInt(raw) {
if intValue > 1_000_000_000_000 {
return Date(timeIntervalSince1970: TimeInterval(intValue) / 1000)
}
if intValue > 1_000_000_000 {
return Date(timeIntervalSince1970: TimeInterval(intValue))
}
}
if let string = self.parseString(raw) {
let formatter = ISO8601DateFormatter()
if let date = formatter.date(from: string) {
return date
}
let dateFormatter = DateFormatter()
dateFormatter.locale = Locale(identifier: "en_US_POSIX")
for format in ["yyyy-MM-dd", "yyyy-MM-dd HH:mm", "yyyy-MM-dd HH:mm:ss"] {
dateFormatter.dateFormat = format
if let date = dateFormatter.date(from: string) {
return date
}
}
}
return nil
}
private static func parseBool(_ raw: Any?) -> Bool? {
if let value = raw as? Bool { return value }
if let number = raw as? NSNumber { return number.boolValue }
guard let string = self.parseString(raw)?.lowercased() else { return nil }
switch string {
case "true", "1", "yes", "active", "valid", "normal":
return true
case "false", "0", "no", "inactive", "invalid", "expired":
return false
default:
return nil
}
}
private static func extractCookieValue(name: String, from cookieHeader: String) -> String? {
cookieHeader
.split(separator: ";")
.compactMap { part -> (String, String)? in
let pieces = part.split(separator: "=", maxSplits: 1).map {
$0.trimmingCharacters(in: .whitespacesAndNewlines)
}
guard pieces.count == 2 else { return nil }
return (pieces[0], pieces[1])
}
.first { $0.0 == name }?
.1
}
private static func hasCSRF(in cookieHeader: String) -> Bool {
self.extractCookieValue(name: "login_aliyunid_csrf", from: cookieHeader) != nil ||
self.extractCookieValue(name: "csrf", from: cookieHeader) != nil
}
static func cookieNames(from cookieHeader: String) -> [String] {
CookieHeaderNormalizer.pairs(from: cookieHeader)
.map(\.name)
.filter { !$0.isEmpty }
.uniquedSorted()
}
static func cookieNamesDescription(from cookieHeader: String) -> String {
let names = self.cookieNames(from: cookieHeader)
return names.isEmpty ? "none" : names.joined(separator: ",")
}
private static func extractSECToken(from html: String) -> String? {
let patterns = [
#""secToken"\s*:\s*"([^"]+)""#,
#""sec_token"\s*:\s*"([^"]+)""#,
#"secToken['"]?\s*[:=]\s*['"]([^'"]+)['"]"#,
#"sec_token['"]?\s*[:=]\s*['"]([^'"]+)['"]"#,
]
for pattern in patterns {
if let token = self.matchFirstGroup(pattern: pattern, in: html), !token.isEmpty {
return token
}
}
return nil
}
private static func matchFirstGroup(pattern: String, in text: String) -> String? {
guard let regex = try? NSRegularExpression(pattern: pattern, options: []) else {
return nil
}
let range = NSRange(text.startIndex..<text.endIndex, in: text)
guard let match = regex.firstMatch(in: text, options: [], range: range),
match.numberOfRanges > 1,
let valueRange = Range(match.range(at: 1), in: text)
else {
return nil
}
let value = text[valueRange].trimmingCharacters(in: .whitespacesAndNewlines)
return value.isEmpty ? nil : String(value)
}
}
extension [String] {
fileprivate func uniquedSorted() -> [String] {
Array(Set(self)).sorted()
}
}
@@ -0,0 +1,96 @@
import Foundation
public struct AlibabaTokenPlanUsageSnapshot: Sendable {
public let planName: String?
public let usedQuota: Double?
public let totalQuota: Double?
public let remainingQuota: Double?
public let resetsAt: Date?
public let updatedAt: Date
public init(
planName: String?,
usedQuota: Double?,
totalQuota: Double?,
remainingQuota: Double?,
resetsAt: Date?,
updatedAt: Date)
{
self.planName = planName
self.usedQuota = usedQuota
self.totalQuota = totalQuota
self.remainingQuota = remainingQuota
self.resetsAt = resetsAt
self.updatedAt = updatedAt
}
}
extension AlibabaTokenPlanUsageSnapshot {
public func toUsageSnapshot() -> UsageSnapshot {
let primary: RateWindow? = Self.usedPercent(
used: self.usedQuota,
total: self.totalQuota,
remaining: self.remainingQuota).map {
RateWindow(
usedPercent: $0,
windowMinutes: 30 * 24 * 60,
resetsAt: self.resetsAt,
resetDescription: Self.quotaDetail(
used: self.usedQuota,
total: self.totalQuota,
remaining: self.remainingQuota))
}
let planName = self.planName?.trimmingCharacters(in: .whitespacesAndNewlines)
let loginMethod = (planName?.isEmpty ?? true) ? nil : planName
let identity = ProviderIdentitySnapshot(
providerID: .alibabatokenplan,
accountEmail: nil,
accountOrganization: nil,
loginMethod: loginMethod)
return UsageSnapshot(
primary: primary,
secondary: nil,
tertiary: nil,
providerCost: nil,
updatedAt: self.updatedAt,
identity: identity)
}
private static func usedPercent(used: Double?, total: Double?, remaining: Double?) -> Double? {
guard let total, total > 0 else { return nil }
let usedValue: Double? = if let used {
used
} else if let remaining {
total - remaining
} else {
nil
}
guard let usedValue else { return nil }
let normalizedUsed = max(0, min(usedValue, total))
return normalizedUsed / total * 100
}
private static func quotaDetail(used: Double?, total: Double?, remaining: Double?) -> String? {
if let used, let total, total > 0 {
return "\(self.format(used)) / \(self.format(total)) credits used"
}
if let remaining, let total, total > 0 {
return "\(Self.format(remaining)) / \(Self.format(total)) credits left"
}
if let remaining {
return "\(Self.format(remaining)) credits left"
}
return nil
}
private static func format(_ value: Double) -> String {
let formatter = NumberFormatter()
formatter.numberStyle = .decimal
formatter.usesGroupingSeparator = true
formatter.maximumFractionDigits = value.rounded() == value ? 0 : 2
formatter.locale = Locale(identifier: "en_US_POSIX")
return formatter.string(from: NSNumber(value: value)) ?? String(format: "%.2f", value)
}
}
@@ -0,0 +1,46 @@
import Foundation
public struct AmpCLIProbe: Sendable {
private static let commandTimeout: TimeInterval = 15
private let arguments: [String]
public init() {
self.arguments = ["usage"]
}
init(arguments: [String]) {
self.arguments = arguments
}
public func fetch(
environment: [String: String] = ProcessInfo.processInfo.environment,
now: Date = Date()) async throws -> AmpUsageSnapshot
{
let loginPATH = LoginShellPathCache.shared.current
guard let executable = BinaryLocator.resolveAmpBinary(env: environment, loginPATH: loginPATH) else {
throw SubprocessRunnerError.binaryNotFound("amp")
}
var commandEnvironment = environment
commandEnvironment["NO_COLOR"] = "1"
commandEnvironment["PATH"] = PathBuilder.effectivePATH(
purposes: [.tty, .nodeTooling],
env: environment,
loginPATH: loginPATH)
let result = try await SubprocessRunner.run(
binary: executable,
arguments: self.arguments,
environment: commandEnvironment,
timeout: Self.commandTimeout,
standardInput: FileHandle.nullDevice,
label: "amp-usage")
let output = result.stdout.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty
? result.stderr
: result.stdout
guard !output.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty else {
throw AmpUsageError.parseFailed("The Amp CLI returned no usage data.")
}
return try AmpUsageParser.parse(displayText: output, now: now)
}
}
@@ -0,0 +1,169 @@
import Foundation
public enum AmpProviderDescriptor {
public static let descriptor: ProviderDescriptor = Self.makeDescriptor()
static func makeDescriptor() -> ProviderDescriptor {
ProviderDescriptor(
id: .amp,
metadata: ProviderMetadata(
id: .amp,
displayName: "Amp",
sessionLabel: "Amp Free",
weeklyLabel: "Balance",
opusLabel: nil,
supportsOpus: false,
supportsCredits: true,
creditsHint: "Individual and workspace credit balances from Amp.",
toggleTitle: "Show Amp usage",
cliName: "amp",
defaultEnabled: false,
isPrimaryProvider: false,
usesAccountFallback: false,
browserCookieOrder: ProviderBrowserCookieDefaults.defaultImportOrder,
dashboardURL: "https://ampcode.com/settings/usage",
statusPageURL: nil),
branding: ProviderBranding(
iconStyle: .amp,
iconResourceName: "ProviderIcon-amp",
color: ProviderColor(red: 220 / 255, green: 38 / 255, blue: 38 / 255)),
tokenCost: ProviderTokenCostConfig(
supportsTokenCost: false,
noDataMessage: { "Amp cost summary is not supported." }),
fetchPlan: ProviderFetchPlan(
sourceModes: [.auto, .api, .web, .cli],
pipeline: ProviderFetchPipeline(resolveStrategies: self.resolveStrategies)),
cli: ProviderCLIConfig(
name: "amp",
versionDetector: nil))
}
private static func resolveStrategies(context: ProviderFetchContext) async -> [any ProviderFetchStrategy] {
switch context.sourceMode {
case .auto:
[AmpCLIFetchStrategy(), AmpAPIFetchStrategy(), AmpStatusFetchStrategy()]
case .cli:
[AmpCLIFetchStrategy()]
case .api:
[AmpAPIFetchStrategy()]
case .web:
[AmpStatusFetchStrategy()]
case .oauth:
[]
}
}
}
struct AmpCLIFetchStrategy: ProviderFetchStrategy {
let id: String = "amp.cli"
let kind: ProviderFetchKind = .cli
func isAvailable(_ context: ProviderFetchContext) async -> Bool {
BinaryLocator.resolveAmpBinary(
env: context.env,
loginPATH: LoginShellPathCache.shared.current) != nil
}
func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult {
let snapshot = try await AmpCLIProbe().fetch(environment: context.env)
return self.makeResult(
usage: snapshot.toUsageSnapshot(now: snapshot.updatedAt),
sourceLabel: "cli")
}
func shouldFallback(on error: Error, context: ProviderFetchContext) -> Bool {
if error is CancellationError || (error as? URLError)?.code == .cancelled {
return false
}
return context.sourceMode == .auto
}
}
struct AmpAPIFetchStrategy: ProviderFetchStrategy {
let id: String = "amp.api"
let kind: ProviderFetchKind = .apiToken
func isAvailable(_ context: ProviderFetchContext) async -> Bool {
_ = context
return true
}
func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult {
guard let token = ProviderTokenResolver.ampToken(environment: context.env) else {
throw AmpUsageError.missingAPIToken
}
let logger: ((String) -> Void)? = context.verbose
? { msg in CodexBarLog.logger(LogCategories.amp).verbose(msg) }
: nil
let snapshot = try await AmpUsageFetcher(browserDetection: context.browserDetection)
.fetch(apiToken: token, logger: logger)
return self.makeResult(
usage: snapshot.toUsageSnapshot(now: snapshot.updatedAt),
sourceLabel: "api")
}
func shouldFallback(on error: Error, context: ProviderFetchContext) -> Bool {
guard context.sourceMode == .auto else { return false }
return !(error is CancellationError) && (error as? URLError)?.code != .cancelled
}
}
struct AmpStatusFetchStrategy: ProviderFetchStrategy {
let id: String = "amp.web"
let kind: ProviderFetchKind = .web
func isAvailable(_ context: ProviderFetchContext) async -> Bool {
guard context.sourceMode.usesWeb else { return false }
#if os(macOS)
let canImportBrowserCookies = true
#else
let canImportBrowserCookies = false
#endif
return Self.canUseWebFallback(
settings: context.settings?.amp,
canImportBrowserCookies: canImportBrowserCookies)
}
func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult {
let fetcher = AmpUsageFetcher(browserDetection: context.browserDetection)
let manual = Self.manualCookieHeader(from: context)
let logger: ((String) -> Void)? = context.verbose
? { msg in CodexBarLog.logger(LogCategories.amp).verbose(msg) }
: nil
let snap = try await fetcher.fetch(cookieHeaderOverride: manual, logger: logger)
return self.makeResult(
usage: snap.toUsageSnapshot(now: snap.updatedAt),
sourceLabel: "web")
}
func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool {
false
}
static func canUseWebFallback(
settings: ProviderSettingsSnapshot.AmpProviderSettings?,
canImportBrowserCookies: Bool) -> Bool
{
guard let settings else { return canImportBrowserCookies }
switch settings.cookieSource {
case .auto:
return canImportBrowserCookies
case .manual:
return Self.manualCookieHeader(from: settings.manualCookieHeader) != nil
case .off:
return false
}
}
private static func manualCookieHeader(from context: ProviderFetchContext) -> String? {
guard let settings = context.settings?.amp, settings.cookieSource == .manual else { return nil }
return Self.manualCookieHeader(from: settings.manualCookieHeader)
}
private static func manualCookieHeader(from rawHeader: String?) -> String? {
guard let header = CookieHeaderNormalizer.normalize(rawHeader),
CookieHeaderNormalizer.pairs(from: header).contains(where: { $0.name == "session" })
else { return nil }
return header
}
}
@@ -0,0 +1,22 @@
import Foundation
public enum AmpSettingsReader {
public static let apiTokenKey = "AMP_API_KEY"
public static func apiToken(environment: [String: String] = ProcessInfo.processInfo.environment) -> String? {
self.cleaned(environment[self.apiTokenKey])
}
static func cleaned(_ raw: String?) -> String? {
guard var value = raw?.trimmingCharacters(in: .whitespacesAndNewlines), !value.isEmpty else {
return nil
}
if (value.hasPrefix("\"") && value.hasSuffix("\"")) ||
(value.hasPrefix("'") && value.hasSuffix("'"))
{
value = String(value.dropFirst().dropLast())
}
value = value.trimmingCharacters(in: .whitespacesAndNewlines)
return value.isEmpty ? nil : value
}
}
@@ -0,0 +1,491 @@
import Foundation
#if canImport(FoundationNetworking)
import FoundationNetworking
#endif
#if os(macOS)
import SweetCookieKit
#endif
public enum AmpUsageError: LocalizedError, Sendable {
case notLoggedIn
case invalidCredentials
case missingAPIToken
case invalidAPIToken
case parseFailed(String)
case networkError(String)
case noSessionCookie
public var errorDescription: String? {
switch self {
case .notLoggedIn:
"Not logged in to Amp. Please log in via ampcode.com."
case .invalidCredentials:
"Amp session cookie expired. Please log in again."
case .missingAPIToken:
"Amp access token not configured. Set AMP_API_KEY or add it in Settings."
case .invalidAPIToken:
"Amp access token is invalid or expired."
case let .parseFailed(message):
"Could not parse Amp usage: \(message)"
case let .networkError(message):
"Amp request failed: \(message)"
case .noSessionCookie:
"No Amp session cookie found. Please log in to ampcode.com in your browser."
}
}
}
#if os(macOS)
private let ampCookieImportOrder: BrowserCookieImportOrder =
ProviderDefaults.metadata[.amp]?.browserCookieOrder ?? Browser.defaultImportOrder
public enum AmpCookieImporter {
private static let cookieClient = BrowserCookieClient()
private static let cookieDomains = ["ampcode.com", "www.ampcode.com"]
private static let sessionCookieNames: Set<String> = [
"session",
]
public struct SessionInfo: Sendable {
public let cookies: [HTTPCookie]
public let sourceLabel: String
public init(cookies: [HTTPCookie], sourceLabel: String) {
self.cookies = cookies
self.sourceLabel = sourceLabel
}
public var cookieHeader: String {
self.cookies.map { "\($0.name)=\($0.value)" }.joined(separator: "; ")
}
}
public static func importSession(
browserDetection: BrowserDetection,
logger: ((String) -> Void)? = nil) throws -> SessionInfo
{
let log: (String) -> Void = { msg in logger?("[amp-cookie] \(msg)") }
let installed = ampCookieImportOrder.cookieImportCandidates(using: browserDetection)
for browserSource in installed {
do {
let query = BrowserCookieQuery(domains: self.cookieDomains)
let sources = try Self.cookieClient.codexBarRecords(
matching: query,
in: browserSource,
logger: log)
for source in sources where !source.records.isEmpty {
let cookies = BrowserCookieClient.makeHTTPCookies(source.records, origin: query.origin)
guard !cookies.isEmpty else { continue }
let names = cookies.map(\.name).joined(separator: ", ")
log("\(source.label) cookies: \(names)")
let sessionCookies = cookies.filter { Self.sessionCookieNames.contains($0.name) }
if !sessionCookies.isEmpty {
log("Found Amp session cookie in \(source.label)")
return SessionInfo(cookies: sessionCookies, sourceLabel: source.label)
}
log("\(source.label) cookies found, but no Amp session cookie present")
log("Expected one of: \(Self.sessionCookieNames.joined(separator: ", "))")
}
} catch {
BrowserCookieAccessGate.recordIfNeeded(error)
log("\(browserSource.displayName) cookie import failed: \(error.localizedDescription)")
}
}
throw AmpUsageError.noSessionCookie
}
}
#endif
public struct AmpUsageFetcher: Sendable {
private static let settingsURL = URL(string: "https://ampcode.com/settings")!
static let usageURL = URL(string: "https://ampcode.com/api/internal?userDisplayBalanceInfo")!
@MainActor private static var recentDumps: [String] = []
public let browserDetection: BrowserDetection
public init(browserDetection: BrowserDetection) {
self.browserDetection = browserDetection
}
public func fetch(
cookieHeaderOverride: String? = nil,
logger: ((String) -> Void)? = nil,
now: Date = Date()) async throws -> AmpUsageSnapshot
{
let log: (String) -> Void = { msg in logger?("[amp] \(msg)") }
let cookieHeader = try await self.resolveCookieHeader(override: cookieHeaderOverride, logger: log)
if let logger {
let names = self.cookieNames(from: cookieHeader)
if !names.isEmpty {
logger("[amp] Cookie names: \(names.joined(separator: ", "))")
}
let diagnostics = RedirectDiagnostics(cookieHeader: cookieHeader, logger: logger)
do {
let (html, responseInfo) = try await self.fetchLegacyHTMLWithDiagnostics(
cookieHeader: cookieHeader,
diagnostics: diagnostics)
self.logDiagnostics(responseInfo: responseInfo, diagnostics: diagnostics, logger: logger)
return try AmpUsageParser.parse(html: html, now: now)
} catch {
self.logDiagnostics(responseInfo: nil, diagnostics: diagnostics, logger: logger)
logger("[amp] Fetch failed: \(error.localizedDescription)")
throw error
}
}
let diagnostics = RedirectDiagnostics(cookieHeader: cookieHeader, logger: nil)
let (html, _) = try await self.fetchLegacyHTMLWithDiagnostics(
cookieHeader: cookieHeader,
diagnostics: diagnostics)
return try AmpUsageParser.parse(html: html, now: now)
}
public func fetch(
apiToken: String,
logger: ((String) -> Void)? = nil,
now: Date = Date()) async throws -> AmpUsageSnapshot
{
guard let token = AmpSettingsReader.cleaned(apiToken) else {
throw AmpUsageError.missingAPIToken
}
let request = try Self.makeUsageAPIRequest(apiToken: token)
let diagnostics = APIRedirectDiagnostics(logger: logger)
let session = URLSession(configuration: .ephemeral, delegate: diagnostics, delegateQueue: nil)
let httpResponse = try await session.response(for: request)
logger?("[amp] API response: \(httpResponse.statusCode) " +
"\(httpResponse.response.url?.absoluteString ?? "unknown")")
try Self.validateAPIResponse(httpResponse)
return try Self.parseUsageAPIResponse(httpResponse.data, now: now)
}
public func debugRawProbe(cookieHeaderOverride: String? = nil) async -> String {
let stamp = ISO8601DateFormatter().string(from: Date())
var lines: [String] = []
lines.append("=== Amp Debug Probe @ \(stamp) ===")
lines.append("")
do {
let cookieHeader = try await self.resolveCookieHeader(
override: cookieHeaderOverride,
logger: { msg in lines.append("[cookie] \(msg)") })
let diagnostics = RedirectDiagnostics(cookieHeader: cookieHeader, logger: nil)
let cookieNames = CookieHeaderNormalizer.pairs(from: cookieHeader).map(\.name)
lines.append("Cookie names: \(cookieNames.joined(separator: ", "))")
let (html, responseInfo) = try await self.fetchLegacyHTMLWithDiagnostics(
cookieHeader: cookieHeader,
diagnostics: diagnostics)
let snapshot = try AmpUsageParser.parse(html: html)
lines.append("")
lines.append("Fetch Success")
lines.append("Status: \(responseInfo.statusCode) \(responseInfo.url)")
if !diagnostics.redirects.isEmpty {
lines.append("")
lines.append("Redirects:")
for entry in diagnostics.redirects {
lines.append(" \(entry)")
}
}
lines.append("")
lines.append("Amp Free:")
lines.append(" quota=\(snapshot.freeQuota?.description ?? "nil")")
lines.append(" used=\(snapshot.freeUsed?.description ?? "nil")")
lines.append(" hourlyReplenishment=\(snapshot.hourlyReplenishment?.description ?? "nil")")
lines.append(" windowHours=\(snapshot.windowHours?.description ?? "nil")")
lines.append(" individualCredits=\(snapshot.individualCredits?.description ?? "nil")")
for workspace in snapshot.workspaceBalances {
lines.append(" workspace[\(workspace.name)]=\(workspace.remaining)")
}
let output = lines.joined(separator: "\n")
Task { @MainActor in Self.recordDump(output) }
return output
} catch {
lines.append("")
lines.append("Probe Failed: \(error.localizedDescription)")
let output = lines.joined(separator: "\n")
Task { @MainActor in Self.recordDump(output) }
return output
}
}
public static func latestDumps() async -> String {
await MainActor.run {
let result = Self.recentDumps.joined(separator: "\n\n---\n\n")
return result.isEmpty ? "No Amp probe dumps captured yet." : result
}
}
private func resolveCookieHeader(
override: String?,
logger: ((String) -> Void)?) async throws -> String
{
if let override = CookieHeaderNormalizer.normalize(override) {
if let sessionHeader = self.sessionCookieHeader(from: override) {
logger?("[amp] Using manual session cookie")
return sessionHeader
}
throw AmpUsageError.noSessionCookie
}
#if os(macOS)
let session = try AmpCookieImporter.importSession(browserDetection: self.browserDetection, logger: logger)
logger?("[amp] Using cookies from \(session.sourceLabel)")
return session.cookieHeader
#else
throw AmpUsageError.noSessionCookie
#endif
}
static func makeUsageAPIRequest(apiToken: String) throws -> URLRequest {
var request = URLRequest(url: Self.usageURL)
request.httpMethod = "POST"
request.httpBody = try JSONSerialization.data(withJSONObject: [
"method": "userDisplayBalanceInfo",
"params": [:],
])
request.setValue("Bearer \(apiToken)", forHTTPHeaderField: "Authorization")
request.setValue("application/json", forHTTPHeaderField: "accept")
request.setValue("application/json", forHTTPHeaderField: "content-type")
return request
}
private func fetchLegacyHTMLWithDiagnostics(
cookieHeader: String,
diagnostics: RedirectDiagnostics) async throws -> (String, ResponseInfo)
{
var request = URLRequest(url: Self.settingsURL)
request.httpMethod = "GET"
request.setValue(cookieHeader, forHTTPHeaderField: "Cookie")
request.setValue(
"text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8",
forHTTPHeaderField: "accept")
Self.applyBrowserHeaders(to: &request)
let session = URLSession(configuration: .ephemeral, delegate: diagnostics, delegateQueue: nil)
let httpResponse = try await session.response(for: request)
let responseInfo = ResponseInfo(
statusCode: httpResponse.statusCode,
url: httpResponse.response.url?.absoluteString ?? "unknown")
try Self.validateBrowserResponse(response: httpResponse, diagnostics: diagnostics)
let html = String(data: httpResponse.data, encoding: .utf8) ?? ""
return (html, responseInfo)
}
static func parseUsageAPIResponse(_ data: Data, now: Date = Date()) throws -> AmpUsageSnapshot {
let response: UsageAPIResponse
do {
response = try JSONDecoder().decode(UsageAPIResponse.self, from: data)
} catch {
throw AmpUsageError.parseFailed("Invalid Amp usage API response.")
}
guard response.ok else {
if response.error?.code == "auth-required" {
throw AmpUsageError.invalidAPIToken
}
throw AmpUsageError.networkError(response.error?.message ?? "Amp usage API returned an error.")
}
guard let displayText = response.result?.displayText, !displayText.isEmpty else {
throw AmpUsageError.parseFailed("Missing Amp usage display text.")
}
return try AmpUsageParser.parse(displayText: displayText, now: now)
}
private static func applyBrowserHeaders(to request: inout URLRequest) {
request.setValue(
"Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) " +
"AppleWebKit/537.36 (KHTML, like Gecko) Chrome/143.0.0.0 Safari/537.36",
forHTTPHeaderField: "user-agent")
request.setValue("en-US,en;q=0.9", forHTTPHeaderField: "accept-language")
request.setValue("https://ampcode.com", forHTTPHeaderField: "origin")
request.setValue(self.settingsURL.absoluteString, forHTTPHeaderField: "referer")
}
private static func validateBrowserResponse(
response: ProviderHTTPResponse,
diagnostics: RedirectDiagnostics) throws
{
guard response.statusCode == 200 else {
if response.statusCode == 401 || response.statusCode == 403 || diagnostics.detectedLoginRedirect {
throw AmpUsageError.invalidCredentials
}
throw AmpUsageError.networkError("HTTP \(response.statusCode)")
}
}
private static func validateAPIResponse(_ response: ProviderHTTPResponse) throws {
guard response.statusCode == 200 else {
if response.statusCode == 401 || response.statusCode == 403 {
throw AmpUsageError.invalidAPIToken
}
throw AmpUsageError.networkError("HTTP \(response.statusCode)")
}
}
@MainActor private static func recordDump(_ text: String) {
if self.recentDumps.count >= 5 { self.recentDumps.removeFirst() }
self.recentDumps.append(text)
}
private final class RedirectDiagnostics: NSObject, URLSessionTaskDelegate, @unchecked Sendable {
private let cookieHeader: String
private let logger: ((String) -> Void)?
var redirects: [String] = []
private(set) var detectedLoginRedirect = false
init(cookieHeader: String, logger: ((String) -> Void)?) {
self.cookieHeader = cookieHeader
self.logger = logger
}
func urlSession(
_: URLSession,
task: URLSessionTask,
willPerformHTTPRedirection response: HTTPURLResponse,
newRequest request: URLRequest,
completionHandler: @escaping (URLRequest?) -> Void)
{
let from = response.url?.absoluteString ?? "unknown"
let to = request.url?.absoluteString ?? "unknown"
self.redirects.append("\(response.statusCode) \(from) -> \(to)")
if let toURL = request.url, AmpUsageFetcher.isLoginRedirect(toURL) {
if let logger {
logger("[amp] Detected login redirect, aborting (invalid session)")
}
self.detectedLoginRedirect = true
completionHandler(nil)
return
}
var updated = request
if AmpUsageFetcher.shouldAttachCookie(to: request.url), !self.cookieHeader.isEmpty {
updated.setValue(self.cookieHeader, forHTTPHeaderField: "Cookie")
} else {
updated.setValue(nil, forHTTPHeaderField: "Cookie")
}
if let referer = response.url?.absoluteString {
updated.setValue(referer, forHTTPHeaderField: "referer")
}
if let logger {
logger("[amp] Redirect \(response.statusCode) \(from) -> \(to)")
}
completionHandler(updated)
}
}
/// Amp's balance RPC should not redirect. Refusing redirects guarantees the bearer token cannot cross hosts.
private final class APIRedirectDiagnostics: NSObject, URLSessionTaskDelegate, @unchecked Sendable {
private let logger: ((String) -> Void)?
init(logger: ((String) -> Void)?) {
self.logger = logger
}
func urlSession(
_: URLSession,
task _: URLSessionTask,
willPerformHTTPRedirection response: HTTPURLResponse,
newRequest request: URLRequest,
completionHandler: @escaping (URLRequest?) -> Void)
{
let from = response.url?.absoluteString ?? "unknown"
let to = request.url?.absoluteString ?? "unknown"
self.logger?("[amp] API redirect blocked: \(response.statusCode) \(from) -> \(to)")
completionHandler(nil)
}
}
private struct ResponseInfo {
let statusCode: Int
let url: String
}
private struct UsageAPIResponse: Decodable {
let ok: Bool
let result: Result?
let error: APIError?
struct Result: Decodable {
let displayText: String
}
struct APIError: Decodable {
let code: String?
let message: String?
}
}
private func logDiagnostics(
responseInfo: ResponseInfo?,
diagnostics: RedirectDiagnostics,
logger: (String) -> Void)
{
if let responseInfo {
logger("[amp] Response: \(responseInfo.statusCode) \(responseInfo.url)")
}
if !diagnostics.redirects.isEmpty {
logger("[amp] Redirects:")
for entry in diagnostics.redirects {
logger("[amp] \(entry)")
}
}
}
private func cookieNames(from header: String) -> [String] {
header.split(separator: ";").compactMap { part in
let trimmed = part.trimmingCharacters(in: .whitespacesAndNewlines)
guard let idx = trimmed.firstIndex(of: "=") else { return nil }
let name = trimmed[..<idx].trimmingCharacters(in: .whitespacesAndNewlines)
return name.isEmpty ? nil : String(name)
}
}
private func sessionCookieHeader(from header: String) -> String? {
let pairs = CookieHeaderNormalizer.pairs(from: header)
let sessionPairs = pairs.filter { $0.name == "session" }
guard !sessionPairs.isEmpty else { return nil }
return sessionPairs.map { "\($0.name)=\($0.value)" }.joined(separator: "; ")
}
static func shouldAttachCookie(to url: URL?) -> Bool {
guard url?.scheme?.lowercased() == "https" else { return false }
return self.isAmpHost(url)
}
private static func isAmpHost(_ url: URL?) -> Bool {
guard let host = url?.host?.lowercased() else { return false }
if host == "ampcode.com" || host == "www.ampcode.com" { return true }
return host.hasSuffix(".ampcode.com")
}
static func isLoginRedirect(_ url: URL) -> Bool {
guard self.isAmpHost(url) else { return false }
if url.host?.lowercased() == "auth.ampcode.com" { return true }
let path = url.path.lowercased()
let components = path.split(separator: "/").map(String.init)
if components.contains("login") { return true }
if components.contains("signin") { return true }
if components.contains("sign-in") { return true }
// Amp currently redirects to /auth/sign-in?returnTo=... when session is invalid. Keep this slightly broader
// than one exact path so we keep working if Amp changes auth routes.
if components.contains("auth") {
let query = url.query?.lowercased() ?? ""
if query.contains("returnto=") { return true }
if query.contains("redirect=") { return true }
if query.contains("redirectto=") { return true }
}
return false
}
}
@@ -0,0 +1,221 @@
import Foundation
enum AmpUsageParser {
static func parse(html: String, now: Date = Date()) throws -> AmpUsageSnapshot {
guard let usage = self.parseFreeTierUsage(html) else {
if self.looksSignedOut(html) {
throw AmpUsageError.notLoggedIn
}
throw AmpUsageError.parseFailed("Missing Amp Free usage data.")
}
return AmpUsageSnapshot(
freeQuota: usage.quota,
freeUsed: usage.used,
hourlyReplenishment: usage.hourlyReplenishment,
windowHours: usage.windowHours,
updatedAt: now)
}
static func parse(displayText: String, now: Date = Date()) throws -> AmpUsageSnapshot {
let text = TextParsing.stripANSICodes(displayText)
let identityPattern = #"(?im)^\s*Signed in as\s+([^\s(]+)(?:\s+\(([^\r\n)]+)\))?\s*$"#
let identity = self.captures(in: text, pattern: identityPattern)
if identity == nil, self.looksSignedOut(text) {
throw AmpUsageError.notLoggedIn
}
let amountPattern = #"([0-9][0-9,]*(?:\.[0-9]+)?)"#
let freePattern = #"(?im)^\s*Amp Free:\s*\$?"# + amountPattern +
#"\s*/\s*\$?"# + amountPattern +
#"\s+remaining(?:\s*\(replenishes\s*\+\$?"# + amountPattern + #"\s*/\s*hour\))?"#
let freePercentPattern = #"(?im)^\s*Amp Free:\s*"# + amountPattern +
#"\s*%\s+remaining(?:\s+today)?(?:\s*\(resets\s+daily\))?"#
let creditsPattern = #"(?im)^\s*Individual credits:\s*\$?"# + amountPattern + #"\s+remaining"#
let individualCredits = self.captures(in: text, pattern: creditsPattern)?.first
.flatMap(self.number(from:))
let workspacePattern = #"(?im)^\s*Workspace\s+(.+?):\s*\$?"# + amountPattern + #"\s+remaining"#
let workspaceBalances: [AmpWorkspaceBalance] = self.allCaptures(
in: text,
pattern: workspacePattern).compactMap { captures -> AmpWorkspaceBalance? in
guard captures.count == 2,
let name = self.nonEmpty(captures[0]),
let remaining = self.number(from: captures[1])
else { return nil }
return AmpWorkspaceBalance(name: name, remaining: remaining)
}
let freeUsage: FreeTierUsage? = {
guard let free = self.captures(in: text, pattern: freePattern),
let remaining = self.number(from: free[0]),
let quota = self.number(from: free[1])
else { return nil }
let hourlyReplenishment = self.number(from: free[2]) ?? 0
let windowHours = hourlyReplenishment > 0
? max(1, (quota / hourlyReplenishment).rounded())
: nil
return FreeTierUsage(
quota: quota,
used: max(0, quota - remaining),
hourlyReplenishment: hourlyReplenishment,
windowHours: windowHours,
resetDescription: nil)
}()
let freePercentUsage: FreeTierUsage? = {
guard let remainingText = self.captures(in: text, pattern: freePercentPattern)?.first,
let remaining = self.number(from: remainingText)
else { return nil }
let clampedRemaining = min(100, max(0, remaining))
return FreeTierUsage(
quota: 100,
used: 100 - clampedRemaining,
hourlyReplenishment: 0,
windowHours: 24,
resetDescription: "resets daily")
}()
let resolvedFreeUsage = freeUsage ?? freePercentUsage
guard resolvedFreeUsage != nil || individualCredits != nil || !workspaceBalances.isEmpty else {
throw AmpUsageError.parseFailed("Missing Amp usage data.")
}
return AmpUsageSnapshot(
freeQuota: resolvedFreeUsage?.quota,
freeUsed: resolvedFreeUsage?.used,
hourlyReplenishment: resolvedFreeUsage?.hourlyReplenishment,
windowHours: resolvedFreeUsage?.windowHours,
individualCredits: individualCredits,
workspaceBalances: workspaceBalances,
accountEmail: self.nonEmpty(identity?[0]),
accountOrganization: self.nonEmpty(identity?[1]),
updatedAt: now,
freeResetDescription: resolvedFreeUsage?.resetDescription)
}
private struct FreeTierUsage {
let quota: Double
let used: Double
let hourlyReplenishment: Double
let windowHours: Double?
let resetDescription: String?
}
private static func parseFreeTierUsage(_ html: String) -> FreeTierUsage? {
let tokens = ["freeTierUsage", "getFreeTierUsage"]
for token in tokens {
if let object = self.extractObject(named: token, in: html),
let usage = self.parseFreeTierUsageObject(object)
{
return usage
}
}
return nil
}
private static func parseFreeTierUsageObject(_ object: String) -> FreeTierUsage? {
guard let quota = self.number(for: "quota", in: object),
let used = self.number(for: "used", in: object),
let hourly = self.number(for: "hourlyReplenishment", in: object)
else { return nil }
let windowHours = self.number(for: "windowHours", in: object)
return FreeTierUsage(
quota: quota,
used: used,
hourlyReplenishment: hourly,
windowHours: windowHours,
resetDescription: nil)
}
private static func extractObject(named token: String, in text: String) -> String? {
guard let tokenRange = text.range(of: token) else { return nil }
guard let braceIndex = text[tokenRange.upperBound...].firstIndex(of: "{") else { return nil }
var depth = 0
var inString = false
var isEscaped = false
var index = braceIndex
while index < text.endIndex {
let char = text[index]
if inString {
if isEscaped {
isEscaped = false
} else if char == "\\" {
isEscaped = true
} else if char == "\"" {
inString = false
}
} else {
if char == "\"" {
inString = true
} else if char == "{" {
depth += 1
} else if char == "}" {
depth -= 1
if depth == 0 {
return String(text[braceIndex...index])
}
}
}
index = text.index(after: index)
}
return nil
}
private static func number(for key: String, in text: String) -> Double? {
let pattern = "\\b\(NSRegularExpression.escapedPattern(for: key))\\b\\s*:\\s*([0-9]+(?:\\.[0-9]+)?)"
guard let regex = try? NSRegularExpression(pattern: pattern) else { return nil }
let range = NSRange(text.startIndex..<text.endIndex, in: text)
guard let match = regex.firstMatch(in: text, options: [], range: range),
match.numberOfRanges > 1,
let valueRange = Range(match.range(at: 1), in: text)
else { return nil }
return Double(text[valueRange])
}
private static func captures(in text: String, pattern: String) -> [String]? {
guard let regex = try? NSRegularExpression(pattern: pattern) else { return nil }
let range = NSRange(text.startIndex..<text.endIndex, in: text)
guard let match = regex.firstMatch(in: text, options: [], range: range) else { return nil }
return self.captures(in: text, match: match)
}
private static func allCaptures(in text: String, pattern: String) -> [[String]] {
guard let regex = try? NSRegularExpression(pattern: pattern) else { return [] }
let range = NSRange(text.startIndex..<text.endIndex, in: text)
return regex.matches(in: text, options: [], range: range).map { match in
self.captures(in: text, match: match)
}
}
private static func captures(in text: String, match: NSTextCheckingResult) -> [String] {
(1..<match.numberOfRanges).map { index in
let captureRange = match.range(at: index)
guard captureRange.location != NSNotFound,
let range = Range(captureRange, in: text)
else { return "" }
return String(text[range]).trimmingCharacters(in: .whitespacesAndNewlines)
}
}
private static func number(from text: String) -> Double? {
Double(text.replacingOccurrences(of: ",", with: ""))
}
private static func nonEmpty(_ text: String?) -> String? {
guard let text, !text.isEmpty else { return nil }
return text
}
private static func looksSignedOut(_ html: String) -> Bool {
let lower = html.lowercased()
if lower.contains("sign in") || lower.contains("log in") || lower.contains("login") {
return true
}
if lower.contains("/login") || lower.contains("ampcode.com/login") {
return true
}
return false
}
}
@@ -0,0 +1,109 @@
import Foundation
public struct AmpWorkspaceBalance: Codable, Equatable, Sendable {
public let name: String
public let remaining: Double
public init(name: String, remaining: Double) {
self.name = name
self.remaining = remaining
}
}
public struct AmpUsageDetails: Codable, Equatable, Sendable {
public let individualCredits: Double?
public let workspaceBalances: [AmpWorkspaceBalance]
public init(individualCredits: Double?, workspaceBalances: [AmpWorkspaceBalance]) {
self.individualCredits = individualCredits
self.workspaceBalances = workspaceBalances
}
}
public struct AmpUsageSnapshot: Sendable {
public let freeQuota: Double?
public let freeUsed: Double?
public let hourlyReplenishment: Double?
public let windowHours: Double?
public let individualCredits: Double?
public let workspaceBalances: [AmpWorkspaceBalance]
public let accountEmail: String?
public let accountOrganization: String?
public let updatedAt: Date
public let freeResetDescription: String?
public init(
freeQuota: Double?,
freeUsed: Double?,
hourlyReplenishment: Double?,
windowHours: Double?,
individualCredits: Double? = nil,
workspaceBalances: [AmpWorkspaceBalance] = [],
accountEmail: String? = nil,
accountOrganization: String? = nil,
updatedAt: Date,
freeResetDescription: String? = nil)
{
self.freeQuota = freeQuota
self.freeUsed = freeUsed
self.hourlyReplenishment = hourlyReplenishment
self.windowHours = windowHours
self.individualCredits = individualCredits
self.workspaceBalances = workspaceBalances
self.accountEmail = accountEmail
self.accountOrganization = accountOrganization
self.updatedAt = updatedAt
self.freeResetDescription = freeResetDescription
}
}
extension AmpUsageSnapshot {
public func toUsageSnapshot(now: Date = Date()) -> UsageSnapshot {
let primary: RateWindow? = if let freeQuota, let freeUsed {
{
let quota = max(0, freeQuota)
let used = max(0, freeUsed)
let percent = quota > 0 ? min(100, (used / quota) * 100) : 0
let windowMinutes: Int? = if let hours = self.windowHours, hours > 0 {
Int((hours * 60).rounded())
} else {
nil
}
let resetsAt: Date? = {
guard quota > 0, let hourlyReplenishment, hourlyReplenishment > 0 else { return nil }
return now.addingTimeInterval(max(0, used / hourlyReplenishment * 3600))
}()
return RateWindow(
usedPercent: percent,
windowMinutes: windowMinutes,
resetsAt: resetsAt,
resetDescription: self.freeResetDescription)
}()
} else {
nil
}
let identity = ProviderIdentitySnapshot(
providerID: .amp,
accountEmail: self.accountEmail,
accountOrganization: self.accountOrganization,
loginMethod: primary == nil ? "Amp" : "Amp Free")
let ampUsage: AmpUsageDetails? = if self.individualCredits != nil || !self.workspaceBalances.isEmpty {
AmpUsageDetails(
individualCredits: self.individualCredits,
workspaceBalances: self.workspaceBalances)
} else {
nil
}
return UsageSnapshot(
primary: primary,
secondary: nil,
tertiary: nil,
ampUsage: ampUsage,
providerCost: nil,
updatedAt: self.updatedAt,
identity: identity)
}
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,37 @@
import Foundation
extension AntigravityStatusProbe {
static func preferredLocalSnapshot(
_ snapshots: [AntigravityStatusSnapshot],
matchingAccountEmail expectedAccountEmail: String?) -> AntigravityStatusSnapshot?
{
var candidates = snapshots
if let expected = self.normalizedAccountEmail(expectedAccountEmail) {
let matches = snapshots.filter {
guard let found = self.normalizedAccountEmail($0.accountEmail) else { return false }
return found.caseInsensitiveCompare(expected) == .orderedSame
}
if !matches.isEmpty {
candidates = matches
}
}
var bestSnapshot: AntigravityStatusSnapshot?
var bestScore = Int.min
for snapshot in candidates {
let score = self.localSnapshotScore(snapshot)
if score > bestScore {
bestSnapshot = snapshot
bestScore = score
}
}
return bestSnapshot
}
private static func normalizedAccountEmail(_ email: String?) -> String? {
guard let trimmed = email?.trimmingCharacters(in: .whitespacesAndNewlines), !trimmed.isEmpty else {
return nil
}
return trimmed
}
}
@@ -0,0 +1,109 @@
import Foundation
#if canImport(FoundationNetworking)
import FoundationNetworking
#endif
enum LocalhostTrustPolicy {
static func shouldAcceptServerTrust(
host: String,
authenticationMethod: String,
hasServerTrust: Bool) -> Bool
{
#if !os(Linux)
guard authenticationMethod == NSURLAuthenticationMethodServerTrust else { return false }
#endif
let normalizedHost = host.lowercased()
guard normalizedHost == "127.0.0.1" || normalizedHost == "localhost" else { return false }
return hasServerTrust
}
}
final class LocalhostSessionDelegate: NSObject {
func data(for request: URLRequest, session: URLSession) async throws -> (Data, URLResponse) {
let state = LocalhostSessionTaskState()
return try await withTaskCancellationHandler {
try await withCheckedThrowingContinuation { continuation in
let task = session.dataTask(with: request) { data, response, error in
if let error {
continuation.resume(throwing: error)
return
}
guard let data, let response else {
continuation.resume(throwing: AntigravityStatusProbeError.apiError("Invalid response"))
return
}
continuation.resume(returning: (data, response))
}
state.setTask(task)
task.resume()
}
} onCancel: {
state.cancel()
}
}
private func challengeResult(_ challenge: URLAuthenticationChallenge) -> (
disposition: URLSession.AuthChallengeDisposition,
credential: URLCredential?)
{
#if os(Linux)
return (.performDefaultHandling, nil)
#else
let protectionSpace = challenge.protectionSpace
let trust = protectionSpace.serverTrust
guard LocalhostTrustPolicy.shouldAcceptServerTrust(
host: protectionSpace.host,
authenticationMethod: protectionSpace.authenticationMethod,
hasServerTrust: trust != nil),
let trust
else {
return (.performDefaultHandling, nil)
}
return (.useCredential, URLCredential(trust: trust))
#endif
}
}
extension LocalhostSessionDelegate: URLSessionDelegate {
func urlSession(
_ session: URLSession,
didReceive challenge: URLAuthenticationChallenge) async -> (URLSession.AuthChallengeDisposition, URLCredential?)
{
self.challengeResult(challenge)
}
}
extension LocalhostSessionDelegate: URLSessionTaskDelegate {
func urlSession(
_ session: URLSession,
task: URLSessionTask,
didReceive challenge: URLAuthenticationChallenge) async -> (URLSession.AuthChallengeDisposition, URLCredential?)
{
self.challengeResult(challenge)
}
}
private final class LocalhostSessionTaskState: @unchecked Sendable {
private let lock = NSLock()
private var task: URLSessionDataTask?
private var isCancelled = false
func setTask(_ task: URLSessionDataTask) {
self.lock.lock()
self.task = task
let shouldCancel = self.isCancelled
self.lock.unlock()
if shouldCancel {
task.cancel()
}
}
func cancel() {
self.lock.lock()
self.isCancelled = true
let task = self.task
self.lock.unlock()
task?.cancel()
}
}
@@ -0,0 +1,516 @@
import Foundation
public struct AntigravityOAuthCredentials: Codable, Sendable, Equatable {
public var accessToken: String?
public var refreshToken: String?
public var expiryDateMilliseconds: Double?
public var idToken: String?
public var email: String?
public var projectID: String?
public var clientID: String?
public var clientSecret: String?
public init(
accessToken: String?,
refreshToken: String?,
expiryDate: Date?,
idToken: String? = nil,
email: String? = nil,
projectID: String? = nil,
clientID: String? = nil,
clientSecret: String? = nil)
{
self.accessToken = accessToken
self.refreshToken = refreshToken
self.expiryDateMilliseconds = expiryDate.map { $0.timeIntervalSince1970 * 1000 }
self.idToken = idToken
self.email = email
self.projectID = projectID
self.clientID = clientID
self.clientSecret = clientSecret
}
public var expiryDate: Date? {
guard let expiryDateMilliseconds else { return nil }
return Date(timeIntervalSince1970: expiryDateMilliseconds / 1000)
}
/// Email of the Google account these credentials authenticate, preferring the
/// signed `id_token` claim (what the remote OAuth fetcher reports) and falling
/// back to the stored `email` field. Used to verify that an ambient local/CLI
/// Antigravity snapshot belongs to the account the user explicitly selected.
public var resolvedAccountEmail: String? {
Self.email(fromIDToken: self.idToken) ?? self.email?.trimmedNonEmptyEmail
}
static func email(fromIDToken idToken: String?) -> String? {
guard let idToken else { return nil }
let parts = idToken.components(separatedBy: ".")
guard parts.count >= 2 else { return nil }
var payload = parts[1]
.replacingOccurrences(of: "-", with: "+")
.replacingOccurrences(of: "_", with: "/")
let remainder = payload.count % 4
if remainder > 0 {
payload += String(repeating: "=", count: 4 - remainder)
}
guard let data = Data(base64Encoded: payload, options: .ignoreUnknownCharacters),
let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any]
else {
return nil
}
return (json["email"] as? String)?.trimmedNonEmptyEmail
}
public init(from decoder: Decoder) throws {
let container = try decoder.container(keyedBy: CodingKeys.self)
self.accessToken =
try container.decodeIfPresent(String.self, forKey: .accessTokenSnake)
?? container.decodeIfPresent(String.self, forKey: .accessTokenCamel)
self.refreshToken =
try container.decodeIfPresent(String.self, forKey: .refreshTokenSnake)
?? container.decodeIfPresent(String.self, forKey: .refreshTokenCamel)
self.idToken =
try container.decodeIfPresent(String.self, forKey: .idTokenSnake)
?? container.decodeIfPresent(String.self, forKey: .idTokenCamel)
self.email = try container.decodeIfPresent(String.self, forKey: .email)
self.projectID =
try container.decodeIfPresent(String.self, forKey: .projectIDSnake)
?? container.decodeIfPresent(String.self, forKey: .projectIDCamel)
self.clientID =
try container.decodeIfPresent(String.self, forKey: .clientIDSnake)
?? container.decodeIfPresent(String.self, forKey: .clientIDCamel)
self.clientSecret =
try container.decodeIfPresent(String.self, forKey: .clientSecretSnake)
?? container.decodeIfPresent(String.self, forKey: .clientSecretCamel)
if let expiryDateMilliseconds = try container.decodeIfPresent(Double.self, forKey: .expiryDateSnake)
?? container.decodeIfPresent(Double.self, forKey: .expiresAtCamel)
{
self.expiryDateMilliseconds = expiryDateMilliseconds
} else if let expiryDateMilliseconds = try container.decodeIfPresent(Int.self, forKey: .expiryDateSnake)
?? container.decodeIfPresent(Int.self, forKey: .expiresAtCamel)
{
self.expiryDateMilliseconds = Double(expiryDateMilliseconds)
} else {
self.expiryDateMilliseconds = nil
}
}
public func encode(to encoder: Encoder) throws {
var container = encoder.container(keyedBy: CodingKeys.self)
try container.encodeIfPresent(self.accessToken, forKey: .accessTokenSnake)
try container.encodeIfPresent(self.refreshToken, forKey: .refreshTokenSnake)
try container.encodeIfPresent(self.expiryDateMilliseconds, forKey: .expiryDateSnake)
try container.encodeIfPresent(self.idToken, forKey: .idTokenSnake)
try container.encodeIfPresent(self.email, forKey: .email)
try container.encodeIfPresent(self.projectID, forKey: .projectIDSnake)
try container.encodeIfPresent(self.clientID, forKey: .clientIDSnake)
try container.encodeIfPresent(self.clientSecret, forKey: .clientSecretSnake)
}
enum CodingKeys: String, CodingKey {
case accessTokenSnake = "access_token"
case accessTokenCamel = "accessToken"
case refreshTokenSnake = "refresh_token"
case refreshTokenCamel = "refreshToken"
case expiryDateSnake = "expiry_date"
case expiresAtCamel = "expiresAt"
case idTokenSnake = "id_token"
case idTokenCamel = "idToken"
case email
case projectIDSnake = "project_id"
case projectIDCamel = "projectId"
case clientIDSnake = "client_id"
case clientIDCamel = "clientId"
case clientSecretSnake = "client_secret"
case clientSecretCamel = "clientSecret"
}
}
public struct AntigravityOAuthClient: Sendable, Equatable {
public let clientID: String
public let clientSecret: String
public init(clientID: String, clientSecret: String) {
self.clientID = clientID
self.clientSecret = clientSecret
}
}
public enum AntigravityOAuthConfig {
public static var configuredClientID: String? {
let value = ProcessInfo.processInfo.environment["ANTIGRAVITY_OAUTH_CLIENT_ID"]
return value?.trimmingCharacters(in: .whitespacesAndNewlines).nilIfEmpty
}
public static var configuredClientSecret: String? {
let value = ProcessInfo.processInfo.environment["ANTIGRAVITY_OAUTH_CLIENT_SECRET"]
return value?.trimmingCharacters(in: .whitespacesAndNewlines).nilIfEmpty
}
public static let authURL = URL(string: "https://accounts.google.com/o/oauth2/v2/auth")!
public static let tokenURL = URL(string: "https://oauth2.googleapis.com/token")!
public static let userInfoURL = URL(string: "https://www.googleapis.com/oauth2/v2/userinfo")!
public static let scopes = [
"https://www.googleapis.com/auth/cloud-platform",
"https://www.googleapis.com/auth/userinfo.email",
]
public static let missingCredentialsMessage =
"""
Antigravity OAuth client is not configured. Install Antigravity.app or set \
ANTIGRAVITY_OAUTH_CLIENT_ID and ANTIGRAVITY_OAUTH_CLIENT_SECRET before logging in.
"""
public static func resolvedClient() -> AntigravityOAuthClient? {
if let client = environmentClient() {
return client
}
return Self.discoverClientFromInstalledApp()
}
private static func environmentClient() -> AntigravityOAuthClient? {
guard let clientID = configuredClientID,
let clientSecret = configuredClientSecret
else {
return nil
}
return AntigravityOAuthClient(clientID: clientID, clientSecret: clientSecret)
}
static func discoverClientFromInstalledApp(
applicationRoots: [URL]? = nil,
fileManager: FileManager = .default) -> AntigravityOAuthClient?
{
for url in self.candidateOAuthClientArtifactURLs(
applicationRoots: applicationRoots,
fileManager: fileManager)
where fileManager.fileExists(atPath: url.path)
{
guard let data = try? Data(contentsOf: url),
let client = Self.parseClient(fromInstalledArtifactData: data)
else {
continue
}
return client
}
return nil
}
static func candidateOAuthClientArtifactURLs(
applicationRoots: [URL]? = nil,
fileManager: FileManager = .default) -> [URL]
{
let roots = [
URL(fileURLWithPath: "/Applications", isDirectory: true),
fileManager.homeDirectoryForCurrentUser.appendingPathComponent("Applications", isDirectory: true),
]
let applicationRoots = applicationRoots ?? roots
let appBundleURLs = self.candidateAntigravityAppBundleURLs(
applicationRoots: applicationRoots,
fileManager: fileManager)
let relativePaths = [
"Contents/Resources/app/extensions/antigravity/bin/language_server_macos_arm",
"Contents/Resources/app/extensions/antigravity/bin/language_server_macos_x64",
"Contents/Resources/app/extensions/antigravity/bin/language_server_macos",
"Contents/Resources/app/out/main.js",
"Contents/Resources/bin/language_server",
"Contents/Resources/bin/language_server_macos",
]
return appBundleURLs.flatMap { bundleURL in
relativePaths.map { bundleURL.appendingPathComponent($0) }
}
}
private static func candidateAntigravityAppBundleURLs(
applicationRoots: [URL],
fileManager: FileManager) -> [URL]
{
var urls: [URL] = []
for root in applicationRoots {
urls.append(root.appendingPathComponent("Antigravity.app", isDirectory: true))
let appURLs = (try? fileManager.contentsOfDirectory(
at: root,
includingPropertiesForKeys: [.isDirectoryKey],
options: [.skipsHiddenFiles])) ?? []
for appURL in appURLs where appURL.pathExtension == "app" {
guard self.isAntigravityAppBundle(appURL) else { continue }
urls.append(appURL)
}
}
var seen = Set<String>()
return urls.filter { url in
let key = url.standardizedFileURL.path
guard !seen.contains(key) else { return false }
seen.insert(key)
return true
}
}
private static func isAntigravityAppBundle(_ url: URL) -> Bool {
switch Bundle(url: url)?.bundleIdentifier {
case "com.google.antigravity", "com.google.antigravity-ide":
true
default:
false
}
}
static func parseClient(fromInstalledArtifactData data: Data) -> AntigravityOAuthClient? {
if let content = String(data: data, encoding: .utf8),
let client = parseClient(fromInstalledArtifactText: content)
{
return client
}
let clientIDs = Self.clientIDs(in: data)
let clientSecrets = Self.clientSecrets(in: data)
guard let client = Self.preferredBinaryClient(
clientIDs: clientIDs,
clientSecrets: clientSecrets)
else {
return nil
}
return client
}
static func parseClient(fromInstalledArtifactText content: String) -> AntigravityOAuthClient? {
let marker = "vs/platform/cloudCode/common/oauthClient.js"
let searchStart = content.range(of: marker)?.lowerBound ?? content.startIndex
let searchEnd = content.index(searchStart, offsetBy: 4000, limitedBy: content.endIndex) ?? content.endIndex
let haystack = String(content[searchStart..<searchEnd])
guard let clientID = Self.firstMatch(
pattern: #"[0-9]+-[A-Za-z0-9_-]+\.apps\.googleusercontent\.com"#,
in: haystack),
let clientSecret = Self.firstMatch(
pattern: #"GOCSPX-[A-Za-z0-9_-]{28}"#,
in: haystack)
else {
return nil
}
return AntigravityOAuthClient(clientID: clientID, clientSecret: clientSecret)
}
private static func clientIDs(in data: Data) -> [String] {
let suffix = Data(".apps.googleusercontent.com".utf8)
var searchRange = data.startIndex..<data.endIndex
var values: [String] = []
while let range = data.range(of: suffix, options: [], in: searchRange) {
var start = range.lowerBound
while start > data.startIndex {
let previous = data.index(before: start)
guard Self.isOAuthClientIDPrefixByte(data[previous]) else { break }
start = previous
}
let candidateData = Data(data[start..<range.upperBound])
if let candidate = String(data: candidateData, encoding: .ascii),
let clientID = Self.firstMatch(
pattern: #"[0-9]+-[A-Za-z0-9_-]+\.apps\.googleusercontent\.com"#,
in: candidate)
{
values.append(clientID)
}
searchRange = range.upperBound..<data.endIndex
}
return self.unique(values)
}
private static func clientSecrets(in data: Data) -> [String] {
let prefix = Data("GOCSPX-".utf8)
let secretLength = 35
var searchRange = data.startIndex..<data.endIndex
var values: [String] = []
while let range = data.range(of: prefix, options: [], in: searchRange) {
let end = range.lowerBound + secretLength
if end <= data.endIndex {
let candidateData = Data(data[range.lowerBound..<end])
if candidateData.dropFirst(prefix.count).allSatisfy(Self.isOAuthClientSecretByte),
let candidate = String(data: candidateData, encoding: .ascii)
{
values.append(candidate)
}
}
searchRange = range.upperBound..<data.endIndex
}
return self.unique(values)
}
private static func preferredBinaryClient(
clientIDs: [String],
clientSecrets: [String]) -> AntigravityOAuthClient?
{
guard !clientIDs.isEmpty,
!clientSecrets.isEmpty
else {
return nil
}
if clientSecrets.count == 1, clientIDs.count > 1 {
return AntigravityOAuthClient(clientID: clientIDs[clientIDs.count - 1], clientSecret: clientSecrets[0])
}
let clientSecret: String = if clientSecrets.count == clientIDs.count, clientSecrets.count > 1 {
// Antigravity 2's language_server binary stores the secret table before the client id table.
clientSecrets[clientSecrets.count - 1]
} else {
clientSecrets[0]
}
return AntigravityOAuthClient(clientID: clientIDs[0], clientSecret: clientSecret)
}
private static func isOAuthClientIDPrefixByte(_ byte: UInt8) -> Bool {
(byte >= 48 && byte <= 57)
|| (byte >= 65 && byte <= 90)
|| (byte >= 97 && byte <= 122)
|| byte == 45
|| byte == 95
}
private static func isOAuthClientSecretByte(_ byte: UInt8) -> Bool {
self.isOAuthClientIDPrefixByte(byte)
}
private static func firstMatch(pattern: String, in text: String) -> String? {
guard let regex = try? NSRegularExpression(pattern: pattern) else { return nil }
let range = NSRange(text.startIndex..<text.endIndex, in: text)
guard let match = regex.firstMatch(in: text, range: range),
let swiftRange = Range(match.range, in: text)
else {
return nil
}
return String(text[swiftRange])
}
private static func unique(_ values: [String]) -> [String] {
var seen = Set<String>()
return values.filter { value in
guard !seen.contains(value) else { return false }
seen.insert(value)
return true
}
}
}
public struct AntigravityOAuthCredentialsStore: @unchecked Sendable {
public static let environmentCredentialsKey = "ANTIGRAVITY_OAUTH_CREDENTIALS_JSON"
private static let fileLock = NSLock()
public let fileURL: URL
private let fileManager: FileManager
public init(fileURL: URL = Self.defaultURL(), fileManager: FileManager = .default) {
self.fileURL = fileURL
self.fileManager = fileManager
}
public func load() throws -> AntigravityOAuthCredentials? {
try Self.fileLock.withLock {
try self.loadUnlocked()
}
}
private func loadUnlocked() throws -> AntigravityOAuthCredentials? {
guard self.fileManager.fileExists(atPath: self.fileURL.path) else { return nil }
let data = try Data(contentsOf: self.fileURL)
return try JSONDecoder().decode(AntigravityOAuthCredentials.self, from: data)
}
public func save(_ credentials: AntigravityOAuthCredentials) throws {
try Self.fileLock.withLock {
let data = try JSONEncoder.antigravityCredentials.encode(credentials)
let directory = self.fileURL.deletingLastPathComponent()
if !self.fileManager.fileExists(atPath: directory.path) {
try self.fileManager.createDirectory(at: directory, withIntermediateDirectories: true)
}
try data.write(to: self.fileURL, options: [.atomic])
try self.applySecurePermissionsIfNeeded()
}
}
public func deleteIfPresent() throws {
try Self.fileLock.withLock {
try self.deleteIfPresentUnlocked()
}
}
public func deleteIfPresent(
matching predicate: (AntigravityOAuthCredentials) -> Bool) throws
{
try Self.fileLock.withLock {
guard let credentials = try self.loadUnlocked(), predicate(credentials) else { return }
try self.deleteIfPresentUnlocked()
}
}
private func deleteIfPresentUnlocked() throws {
guard self.fileManager.fileExists(atPath: self.fileURL.path) else { return }
try self.fileManager.removeItem(at: self.fileURL)
}
public static func defaultDirectoryURL(home: URL = FileManager.default.homeDirectoryForCurrentUser) -> URL {
home
.appendingPathComponent(".codexbar", isDirectory: true)
.appendingPathComponent("antigravity", isDirectory: true)
}
public static func defaultURL(home: URL = FileManager.default.homeDirectoryForCurrentUser) -> URL {
self.defaultDirectoryURL(home: home)
.appendingPathComponent("oauth_creds.json")
}
public static func tokenAccountValue(for credentials: AntigravityOAuthCredentials) throws -> String {
let data = try JSONEncoder.antigravityCredentials.encode(credentials)
guard let value = String(data: data, encoding: .utf8) else {
throw CocoaError(.coderInvalidValue)
}
return value
}
public static func credentials(fromTokenAccountValue value: String) -> AntigravityOAuthCredentials? {
let trimmed = value.trimmingCharacters(in: .whitespacesAndNewlines)
guard !trimmed.isEmpty, let data = trimmed.data(using: .utf8) else { return nil }
return try? JSONDecoder().decode(AntigravityOAuthCredentials.self, from: data)
}
private func applySecurePermissionsIfNeeded() throws {
#if os(macOS) || os(Linux)
try self.fileManager.setAttributes([
.posixPermissions: NSNumber(value: Int16(0o600)),
], ofItemAtPath: self.fileURL.path)
#endif
}
}
extension JSONEncoder {
fileprivate static let antigravityCredentials: JSONEncoder = {
let encoder = JSONEncoder()
encoder.outputFormatting = [.prettyPrinted, .sortedKeys]
return encoder
}()
}
extension String {
fileprivate var nilIfEmpty: String? {
self.isEmpty ? nil : self
}
fileprivate var trimmedNonEmptyEmail: String? {
let trimmed = self.trimmingCharacters(in: .whitespacesAndNewlines)
return trimmed.isEmpty ? nil : trimmed
}
}
@@ -0,0 +1,673 @@
import Foundation
public enum AntigravityProviderDescriptor {
public static let descriptor: ProviderDescriptor = Self.makeDescriptor()
static func makeDescriptor() -> ProviderDescriptor {
ProviderDescriptor(
id: .antigravity,
metadata: ProviderMetadata(
id: .antigravity,
displayName: "Antigravity",
sessionLabel: "Gemini Models",
weeklyLabel: "Claude and GPT",
opusLabel: nil,
supportsOpus: false,
supportsCredits: false,
creditsHint: "",
toggleTitle: "Show Antigravity usage (experimental)",
cliName: "antigravity",
defaultEnabled: false,
isPrimaryProvider: false,
usesAccountFallback: false,
dashboardURL: nil,
statusPageURL: nil,
statusLinkURL: "https://www.google.com/appsstatus/dashboard/products/npdyhgECDJ6tB66MxXyo/history",
statusWorkspaceProductID: "npdyhgECDJ6tB66MxXyo"),
branding: ProviderBranding(
iconStyle: .antigravity,
iconResourceName: "ProviderIcon-antigravity",
color: ProviderColor(red: 96 / 255, green: 186 / 255, blue: 126 / 255)),
tokenCost: ProviderTokenCostConfig(
supportsTokenCost: false,
noDataMessage: { "Antigravity cost summary is not supported." }),
fetchPlan: ProviderFetchPlan(
sourceModes: [.auto, .cli, .oauth],
pipeline: ProviderFetchPipeline(resolveStrategies: self.resolveStrategies)),
cli: ProviderCLIConfig(
name: "antigravity",
versionDetector: nil))
}
private static func resolveStrategies(context: ProviderFetchContext) async -> [any ProviderFetchStrategy] {
let app = AntigravityStatusFetchStrategy(source: .app)
let cli = AntigravityCLIHTTPSFetchStrategy()
let ide = AntigravityStatusFetchStrategy(source: .ide)
let oauth = AntigravityOAuthFetchStrategy()
switch context.sourceMode {
case .cli:
return [app, cli, ide]
case .oauth:
return [oauth]
case .auto:
if context.selectedTokenAccountID != nil ||
context.env[AntigravityOAuthCredentialsStore.environmentCredentialsKey] != nil ||
self.hasSharedOAuthCredentials(context: context)
{
return [app, cli, ide, oauth]
}
return [app, cli, ide]
case .web, .api:
return []
}
}
private static func hasSharedOAuthCredentials(context: ProviderFetchContext) -> Bool {
let homeURL = context.env["HOME"]
.flatMap { $0.isEmpty ? nil : URL(fileURLWithPath: $0, isDirectory: true) }
?? FileManager.default.homeDirectoryForCurrentUser
let fileURL = AntigravityOAuthCredentialsStore.defaultURL(home: homeURL)
return FileManager.default.fileExists(atPath: fileURL.path)
}
}
struct AntigravityStatusFetchStrategy: ProviderFetchStrategy {
enum Source {
case app
case ide
var id: String {
switch self {
case .app: "antigravity.app-local"
case .ide: "antigravity.ide-local"
}
}
var processScope: AntigravityStatusProbe.ProcessScope {
switch self {
case .app: .appOnly
case .ide: .ideOnly
}
}
var sourceLabel: String {
switch self {
case .app: "app"
case .ide: "ide"
}
}
}
let source: Source
var id: String {
self.source.id
}
let kind: ProviderFetchKind = .localProbe
init(source: Source = .app) {
self.source = source
}
func isAvailable(_: ProviderFetchContext) async -> Bool {
true
}
func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult {
let probe = AntigravityStatusProbe(processScope: self.source.processScope)
let selectedAccountEmail: String? = if context.sourceMode == .auto, context.selectedTokenAccountID != nil {
AntigravitySelectedAccountGuard.selectedAccountEmail(context: context)
} else {
nil
}
let snap = try await probe.fetch(matchingAccountEmail: selectedAccountEmail)
let usage = try snap.toUsageSnapshot()
try AntigravitySelectedAccountGuard.validate(usage, context: context)
return self.makeResult(
usage: usage,
sourceLabel: self.source.sourceLabel)
}
func shouldFallback(on _: Error, context: ProviderFetchContext) -> Bool {
context.sourceMode == .auto || context.sourceMode == .cli
}
}
/// When the Antigravity 2.0 app is closed or unavailable, this strategy spawns
/// or reuses ``agy`` and talks to the localhost server embedded in that CLI
/// process. ``agy`` is an interactive REPL, not a query command, so CodexBar
/// never scrapes TUI output here; it only keeps the process alive long enough
/// for the server to answer quota endpoints.
struct AntigravityCLIHTTPSFetchStrategy: ProviderFetchStrategy {
static let sourceLabel = "cli"
let id: String = "antigravity.cli-https"
let kind: ProviderFetchKind = .cli
private static let log = CodexBarLog.logger(LogCategories.antigravity)
struct SnapshotWaitDependencies {
let pollIntervalNanoseconds: UInt64
let listeningPorts: @Sendable (Int, TimeInterval) async throws -> [Int]
let drainOutput: @Sendable () async -> Data
let fetchSnapshot: @Sendable ([Int]) async throws -> AntigravityStatusSnapshot
let now: @Sendable () -> Date
init(
pollIntervalNanoseconds: UInt64,
listeningPorts: @escaping @Sendable (Int, TimeInterval) async throws -> [Int],
drainOutput: @escaping @Sendable () async -> Data,
fetchSnapshot: @escaping @Sendable ([Int]) async throws -> AntigravityStatusSnapshot,
now: @escaping @Sendable () -> Date = Date.init)
{
self.pollIntervalNanoseconds = pollIntervalNanoseconds
self.listeningPorts = listeningPorts
self.drainOutput = drainOutput
self.fetchSnapshot = fetchSnapshot
self.now = now
}
}
/// Seams for discovering and reusing an already-running ``agy`` CLI language
/// server, so a fresh spawn (and its multi-second ``GetUserStatus`` warm-up)
/// can be skipped when a warm server is already present.
struct WarmAgyDependencies {
let processInfos: @Sendable (TimeInterval) async throws -> [AntigravityStatusProbe.ProcessInfoResult]
let listeningPorts: @Sendable (Int, TimeInterval) async throws -> [Int]
let fetchSnapshot: @Sendable ([Int], TimeInterval) async throws -> AntigravityStatusSnapshot
let processOwnerUserID: @Sendable (Int) -> UInt32?
let currentUserID: @Sendable () -> UInt32
/// The pid of an ``agy`` that CodexBar itself spawned and manages through
/// ``AntigravityCLISession`` (if any). Such a process must NOT be reused
/// through the warm path: doing so bypasses `beginProbe`/`finishProbe`, so
/// the idle timer is never cancelled/extended and `stopIfIdle` could tear
/// the managed session down mid-poll. Externally owned `agy` (an IDE, a
/// long-lived `agy`, or another CodexBar host) has no such accounting.
let ownedPID: @Sendable () async -> Int?
let now: @Sendable () -> Date
init(
processInfos: @escaping @Sendable (TimeInterval) async throws
-> [AntigravityStatusProbe.ProcessInfoResult],
listeningPorts: @escaping @Sendable (Int, TimeInterval) async throws -> [Int],
fetchSnapshot: @escaping @Sendable ([Int], TimeInterval) async throws -> AntigravityStatusSnapshot,
processOwnerUserID: @escaping @Sendable (Int) -> UInt32? = { _ in 0 },
currentUserID: @escaping @Sendable () -> UInt32 = { 0 },
ownedPID: @escaping @Sendable () async -> Int? = { nil },
now: @escaping @Sendable () -> Date = Date.init)
{
self.processInfos = processInfos
self.listeningPorts = listeningPorts
self.fetchSnapshot = fetchSnapshot
self.processOwnerUserID = processOwnerUserID
self.currentUserID = currentUserID
self.ownedPID = ownedPID
self.now = now
}
}
/// Discover an already-running, authenticated ``agy`` CLI language server and
/// reuse its listening ports instead of spawning a fresh process.
///
/// One-shot CLI invocations otherwise spawn a brand-new ``agy`` on every
/// call; a fresh server binds its port quickly but ``GetUserStatus`` returns
/// transient initialization failures for a few seconds, so the readiness
/// deadline is occasionally missed. When a warm CLI server is already up, we
/// can talk to it immediately it needs no CSRF token (``cliHTTPS``).
///
/// Returns the snapshot from the first warm server that answers with
/// parseable usage for the requested account, or `nil` when none is found or
/// none answers in which case the caller falls back to the existing spawn
/// path unchanged.
static func tryWarmAgyFetch(
timeout: TimeInterval,
expectedBinaryPath: String? = nil,
expectedAccountEmail: String? = nil,
dependencies: WarmAgyDependencies) async throws -> AntigravityStatusSnapshot?
{
try Task.checkCancellation()
let deadline = dependencies.now().addingTimeInterval(timeout)
guard let discoveryTimeout = Self.remainingWarmProbeTime(deadline: deadline, now: dependencies.now) else {
return nil
}
let processInfos: [AntigravityStatusProbe.ProcessInfoResult]
do {
processInfos = try await dependencies.processInfos(discoveryTimeout)
} catch let error as CancellationError {
throw error
} catch {
return nil
}
try Task.checkCancellation()
let ownedPID = await dependencies.ownedPID()
try Task.checkCancellation()
let currentUserID = dependencies.currentUserID()
// Only the CLI's language server needs no CSRF token; the IDE/app servers
// require one and must not be reused through this token-less fast path.
// Also exclude any `agy` CodexBar itself spawned and manages: reusing it
// here would bypass session lifecycle accounting (see `ownedPID`).
let cliProcesses = processInfos.filter { info in
info.pid != ownedPID &&
dependencies.processOwnerUserID(info.pid) == currentUserID &&
AntigravityStatusProbe.antigravityProcessKind(info.commandLine) == .cli
}
guard !cliProcesses.isEmpty else { return nil }
for info in cliProcesses {
if let expectedBinaryPath {
guard Self.commandLine(info.commandLine, matchesBinaryPath: expectedBinaryPath)
else {
continue
}
}
guard let portTimeout = Self.remainingWarmProbeTime(deadline: deadline, now: dependencies.now) else {
return nil
}
let ports: [Int]
do {
ports = try await dependencies.listeningPorts(info.pid, portTimeout)
} catch let error as CancellationError {
throw error
} catch {
continue
}
try Task.checkCancellation()
guard !ports.isEmpty else { continue }
guard let fetchTimeout = Self.remainingWarmProbeTime(deadline: deadline, now: dependencies.now) else {
return nil
}
let snapshot: AntigravityStatusSnapshot
do {
snapshot = try await dependencies.fetchSnapshot(ports, fetchTimeout)
} catch let error as CancellationError {
throw error
} catch {
continue
}
try Task.checkCancellation()
guard (try? snapshot.toUsageSnapshot()) != nil,
AntigravitySelectedAccountGuard.matches(
snapshotAccountEmail: snapshot.accountEmail,
expectedAccountEmail: expectedAccountEmail)
else {
continue
}
Self.log.debug("Antigravity CLI HTTPS reusing warm agy", metadata: [
"pid": "\(info.pid)",
"ports": ports.map(String.init).joined(separator: ","),
])
return snapshot
}
try Task.checkCancellation()
return nil
}
private static func remainingWarmProbeTime(
deadline: Date,
now: @Sendable () -> Date) -> TimeInterval?
{
let remaining = deadline.timeIntervalSince(now())
return remaining > 0 ? remaining : nil
}
private static func commandLine(_ commandLine: String, matchesBinaryPath binaryPath: String) -> Bool {
let candidates = [
URL(fileURLWithPath: binaryPath).standardizedFileURL.path,
URL(fileURLWithPath: binaryPath).resolvingSymlinksInPath().standardizedFileURL.path,
]
return candidates.contains { candidate in
commandLine == candidate || commandLine.hasPrefix("\(candidate) ")
}
}
/// Production wiring for ``tryWarmAgyFetch``: list processes via `ps`, find
/// listening ports via `lsof`, and probe the token-less CLI HTTPS endpoint.
static func liveWarmAgyDependencies() -> WarmAgyDependencies {
WarmAgyDependencies(
processInfos: { timeout in
// A missing-CSRF/notRunning throw means no reusable server; the
// caller maps any throw to "no warm agy" and spawns instead.
try await AntigravityStatusProbe.detectProcessInfos(
timeout: timeout,
scope: .ideAndCLI)
},
listeningPorts: { pid, timeout in
try await AntigravityStatusProbe.listeningPorts(pid: pid, timeout: timeout)
},
fetchSnapshot: { ports, timeout in
let deadline = Date().addingTimeInterval(timeout)
return try await AntigravityStatusProbe(timeout: timeout)
.fetchFromPorts(ports, deadline: deadline)
},
processOwnerUserID: { pid in
AntigravityProcessIdentityProvider().ownerUserID(for: pid_t(pid))
},
currentUserID: {
AntigravityProcessIdentityProvider.currentUserID
},
ownedPID: {
// The pid of the `agy` CodexBar manages through the shared
// session, so the warm scan never reuses our own process.
await AntigravityCLISession.shared.pid.map(Int.init)
})
}
func isAvailable(_ context: ProviderFetchContext) async -> Bool {
BinaryLocator.resolveAntigravityBinary(env: context.env) != nil
}
func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult {
guard let binary = BinaryLocator.resolveAntigravityBinary(env: context.env) else {
throw AntigravityStatusProbeError.notRunning
}
let expectedAccountEmail: String? = if context.sourceMode == .auto,
context.selectedTokenAccountID != nil
{
AntigravitySelectedAccountGuard.selectedAccountEmail(context: context)
} else {
nil
}
let result = try await self.fetchUsingWarmSession(
binary: binary,
idleWindow: context.persistentCLISessionIdleWindow,
resetAfterFetch: Self.shouldResetSessionAfterFetch(context),
expectedAccountEmail: expectedAccountEmail)
try AntigravitySelectedAccountGuard.validate(result.usage, context: context)
return result
}
private func fetchUsingWarmSession(
binary: String,
idleWindow: TimeInterval?,
resetAfterFetch: Bool,
expectedAccountEmail: String?) async throws -> ProviderFetchResult
{
try await self.fetchUsingWarmSession(
binary: binary,
idleWindow: idleWindow,
resetAfterFetch: resetAfterFetch,
expectedAccountEmail: expectedAccountEmail,
warmDependencies: Self.liveWarmAgyDependencies(),
spawnFetch: { binary, idleWindow, resetAfterFetch in
try await self.fetchBySpawning(
binary: binary,
idleWindow: idleWindow,
resetAfterFetch: resetAfterFetch)
})
}
/// Testable core of the CLI fetch: try the warm-reuse fast path first, then
/// fall back to spawning. The `spawnFetch` seam lets tests assert the spawn
/// path is skipped when a warm server is reused.
func fetchUsingWarmSession(
binary: String,
idleWindow: TimeInterval?,
resetAfterFetch: Bool,
expectedAccountEmail: String? = nil,
warmDependencies: WarmAgyDependencies,
spawnFetch: @Sendable (String, TimeInterval?, Bool) async throws -> ProviderFetchResult)
async throws -> ProviderFetchResult
{
// Fast path: reuse an already-running, authenticated `agy` CLI server if
// one is present, avoiding a fresh spawn and its multi-second warm-up.
// When none is found (or none answers), fall through to the spawn path.
//
// The warm path deliberately does NOT touch `AntigravityCLISession`
// (`beginProbe`/`finishProbe`): a discovered `agy` is owned by another
// process (an IDE, a long-lived `agy`, or another CodexBar host), so
// CodexBar must not manage its lifecycle, idle timeout, or
// `resetAfterFetch` teardown. Those apply only to processes CodexBar
// itself spawns on the fallback path below.
// Long-lived hosts already keep a managed session warm. Restrict external
// process reuse to one-shot CLI calls so app/server lifecycle accounting
// stays entirely inside AntigravityCLISession.
if resetAfterFetch, let warmSnapshot = try await Self.tryWarmAgyFetch(
timeout: 2.0,
expectedBinaryPath: binary,
expectedAccountEmail: expectedAccountEmail,
dependencies: warmDependencies)
{
// `tryWarmAgyFetch` only returns a snapshot whose `toUsageSnapshot()`
// already succeeded, so this conversion must not silently fail.
let warmUsage = try warmSnapshot.toUsageSnapshot()
return self.makeResult(
usage: warmUsage,
sourceLabel: Self.sourceLabel)
}
try Task.checkCancellation()
return try await spawnFetch(binary, idleWindow, resetAfterFetch)
}
/// Spawn (or reuse CodexBar's own warm) `agy` session and wait for the CLI
/// HTTPS endpoint to report ready. This is the original behavior, unchanged.
private func fetchBySpawning(
binary: String,
idleWindow: TimeInterval?,
resetAfterFetch: Bool) async throws -> ProviderFetchResult
{
let session = AntigravityCLISession.shared
let pid = try await session.beginProbe(binary: binary, idleWindow: idleWindow)
let deadline = Date().addingTimeInterval(5.0)
let snap: AntigravityStatusSnapshot
let usage: UsageSnapshot
do {
snap = try await Self.waitForSnapshot(
pid: pid,
deadline: deadline,
dependencies: SnapshotWaitDependencies(
pollIntervalNanoseconds: 200_000_000,
listeningPorts: { pid, timeout in
try await AntigravityStatusProbe.listeningPorts(pid: pid, timeout: timeout)
},
drainOutput: {
await session.drainOutput()
},
fetchSnapshot: { ports in
let timeout = min(2.0, max(0.2, deadline.timeIntervalSinceNow))
return try await AntigravityStatusProbe(timeout: timeout)
.fetchFromPorts(ports, deadline: deadline)
}))
usage = try snap.toUsageSnapshot()
await session.finishProbe(success: true, resetAfterFetch: resetAfterFetch)
} catch {
let authenticationRequired = (error as? AntigravityStatusProbeError) == .authenticationRequired
await session.finishProbe(
success: false,
resetAfterFetch: resetAfterFetch || authenticationRequired,
forceTerminate: authenticationRequired)
throw error
}
return self.makeResult(
usage: usage,
sourceLabel: Self.sourceLabel)
}
static func shouldResetSessionAfterFetch(_ context: ProviderFetchContext) -> Bool {
// Long-lived hosts (the app, `codexbar serve`) keep the warm `agy`
// session between fetches; only one-shot CLI invocations reset it.
context.runtime == .cli && !context.persistsCLISessions
}
/// Waits for real API readiness, not just socket readiness. Fresh ``agy``
/// processes bind ports quickly, but ``GetUserStatus`` can return transient
/// initialization failures for a few seconds after the port appears.
static func waitForSnapshot(
pid: pid_t,
deadline: Date,
dependencies: SnapshotWaitDependencies) async throws -> AntigravityStatusSnapshot
{
var lastFetchError: Error?
while dependencies.now() < deadline {
try await Self.checkAuthenticationPrompt(dependencies)
let remaining = deadline.timeIntervalSince(dependencies.now())
let portProbeTimeout = min(2.0, max(0.2, remaining))
let ports: [Int]
do {
ports = try await dependencies.listeningPorts(Int(pid), portProbeTimeout)
} catch {
guard Self.isNoListeningPortsError(error) else {
try await Self.checkAuthenticationPrompt(dependencies)
throw error
}
ports = []
}
if !ports.isEmpty {
var readySnapshot: AntigravityStatusSnapshot?
do {
let snapshot = try await dependencies.fetchSnapshot(ports)
_ = try snapshot.toUsageSnapshot()
readySnapshot = snapshot
} catch {
try await Self.checkAuthenticationPrompt(dependencies)
lastFetchError = error
Self.log.debug("Antigravity CLI HTTPS endpoint not ready", metadata: [
"pid": "\(pid)",
"ports": ports.map(String.init).joined(separator: ","),
"error": error.localizedDescription,
])
}
if let readySnapshot {
try await Self.checkAuthenticationPrompt(dependencies)
return readySnapshot
}
}
let remainingNanoseconds = UInt64(
max(0, deadline.timeIntervalSince(dependencies.now())) * 1_000_000_000)
guard remainingNanoseconds > 0 else { break }
let sleepNanoseconds = min(dependencies.pollIntervalNanoseconds, remainingNanoseconds)
if sleepNanoseconds > 0 {
try await Task.sleep(nanoseconds: sleepNanoseconds)
}
}
try await Self.checkAuthenticationPrompt(dependencies)
if let lastFetchError {
throw lastFetchError
}
Self.log.warning("Antigravity CLI HTTPS: no ports found for pid \(pid)")
throw AntigravityStatusProbeError.portDetectionFailed(
"Antigravity CLI started but no listening ports found")
}
static func containsAuthenticationPrompt(_ output: Data) -> Bool {
AntigravityCLIAuthenticationPrompt.contains(output)
}
private static func checkAuthenticationPrompt(_ dependencies: SnapshotWaitDependencies) async throws {
let terminalOutput = await dependencies.drainOutput()
if Self.containsAuthenticationPrompt(terminalOutput) {
throw AntigravityStatusProbeError.authenticationRequired
}
}
private static func isNoListeningPortsError(_ error: Error) -> Bool {
if case let AntigravityStatusProbeError.portDetectionFailed(message) = error {
return message == "no listening ports found"
}
if case let SubprocessRunnerError.nonZeroExit(code, stderr) = error {
return code == 1 && stderr.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty
}
return false
}
func shouldFallback(on _: Error, context: ProviderFetchContext) -> Bool {
context.sourceMode == .auto || context.sourceMode == .cli
}
}
struct AntigravityOAuthFetchStrategy: ProviderFetchStrategy {
let id: String = "antigravity.oauth"
let kind: ProviderFetchKind = .oauth
func isAvailable(_: ProviderFetchContext) async -> Bool {
true
}
static func usageSnapshot(
from snapshot: AntigravityStatusSnapshot,
updatedAt: Date = Date()) throws -> UsageSnapshot
{
if snapshot.modelQuotas.isEmpty {
return UsageSnapshot(
primary: nil,
secondary: nil,
tertiary: nil,
updatedAt: updatedAt,
identity: ProviderIdentitySnapshot(
providerID: .antigravity,
accountEmail: snapshot.accountEmail,
accountOrganization: nil,
loginMethod: snapshot.accountPlan))
}
return try snapshot.toUsageSnapshot()
}
func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult {
let fetcher = AntigravityRemoteUsageFetcher(
environment: context.env,
credentialsUpdateHandler: { credentials in
guard let accountID = context.selectedTokenAccountID,
let updater = context.tokenAccountTokenUpdater
else {
return
}
let token = try AntigravityOAuthCredentialsStore.tokenAccountValue(for: credentials)
await updater(.antigravity, accountID, token)
})
let snapshot = try await fetcher.fetch()
let usage = try Self.usageSnapshot(from: snapshot)
return self.makeResult(
usage: usage,
sourceLabel: "oauth")
}
func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool {
false
}
}
/// Guards ambient Antigravity snapshots against the explicitly selected account.
///
/// The local desktop probe and the ``agy`` CLI HTTPS server report whichever
/// Antigravity account is signed into the local session. When the user has
/// selected a specific saved Google account, an ambient probe can return a
/// *different* account's quota. Only the OAuth strategy is account-scoped (it
/// fetches with the selected account's injected credentials), so in ``auto``
/// mode we reject a snapshot whose identity does not match the selected account
/// and let the pipeline fall through to OAuth. Explicit ``cli``/``oauth`` source
/// modes stay authoritative and are never second-guessed here.
enum AntigravitySelectedAccountGuard {
static func matches(snapshotAccountEmail: String?, expectedAccountEmail: String?) -> Bool {
guard let expected = self.normalizedEmail(expectedAccountEmail) else { return true }
guard let found = self.normalizedEmail(snapshotAccountEmail) else { return false }
return found.caseInsensitiveCompare(expected) == .orderedSame
}
static func validate(_ usage: UsageSnapshot, context: ProviderFetchContext) throws {
guard context.sourceMode == .auto, context.selectedTokenAccountID != nil else { return }
let expected = self.selectedAccountEmail(context: context)
let found = self.normalizedEmail(usage.identity?.accountEmail)
guard let expected, let found, found.caseInsensitiveCompare(expected) == .orderedSame else {
throw AntigravityStatusProbeError.accountMismatch(expected: expected, found: found)
}
}
/// Email of the selected token account, read from the same injected
/// credentials the OAuth strategy would use (`ANTIGRAVITY_OAUTH_CREDENTIALS_JSON`).
static func selectedAccountEmail(context: ProviderFetchContext) -> String? {
guard let value = context.env[AntigravityOAuthCredentialsStore.environmentCredentialsKey],
let credentials = AntigravityOAuthCredentialsStore.credentials(fromTokenAccountValue: value)
else {
return nil
}
return credentials.resolvedAccountEmail
}
private static func normalizedEmail(_ email: String?) -> String? {
guard let trimmed = email?.trimmingCharacters(in: .whitespacesAndNewlines), !trimmed.isEmpty else {
return nil
}
return trimmed
}
}
@@ -0,0 +1,173 @@
import Foundation
// swiftformat:disable:next redundantSendable
struct AntigravityQuotaSummary: Sendable, Equatable {
let description: String?
let groups: [AntigravityQuotaSummaryGroup]
}
// swiftformat:disable:next redundantSendable
struct AntigravityQuotaSummaryGroup: Sendable, Equatable {
let displayName: String
let description: String?
let buckets: [AntigravityQuotaSummaryBucket]
}
// swiftformat:disable:next redundantSendable
struct AntigravityQuotaSummaryBucket: Sendable, Equatable {
let bucketId: String
let displayName: String
let remainingFraction: Double?
let resetTime: Date?
let resetDescription: String?
let disabled: Bool
init(
bucketId: String,
displayName: String,
remainingFraction: Double?,
resetTime: Date? = nil,
resetDescription: String?,
disabled: Bool)
{
self.bucketId = bucketId
self.displayName = displayName
self.remainingFraction = remainingFraction
self.resetTime = resetTime
self.resetDescription = resetDescription
self.disabled = disabled
}
}
extension AntigravityStatusProbe {
static func parseQuotaSummaryResponse(_ data: Data) throws -> AntigravityStatusSnapshot {
let decoder = JSONDecoder()
let response = try decoder.decode(QuotaSummaryResponse.self, from: data)
if let invalid = Self.invalidCode(response.code) {
throw AntigravityStatusProbeError.apiError(invalid)
}
let payload = response.response ?? response.summary ?? response.rootPayload
guard let payload else {
throw AntigravityStatusProbeError.parseFailed("Missing quota summary")
}
let groups = payload.groups.compactMap(self.quotaSummaryGroup(from:))
guard !groups.isEmpty else {
throw AntigravityStatusProbeError.parseFailed("Missing quota groups")
}
return AntigravityStatusSnapshot(
quotaSummary: AntigravityQuotaSummary(
description: payload.description,
groups: groups),
accountEmail: nil,
accountPlan: nil,
source: .local)
}
private static func quotaSummaryGroup(from payload: QuotaSummaryGroupPayload) -> AntigravityQuotaSummaryGroup? {
let displayName = payload.displayName?.trimmingCharacters(in: .whitespacesAndNewlines)
let buckets = (payload.buckets ?? []).compactMap(self.quotaSummaryBucket(from:))
guard !buckets.isEmpty else { return nil }
return AntigravityQuotaSummaryGroup(
displayName: self.nonEmpty(displayName) ?? "Quota",
description: payload.description,
buckets: buckets)
}
private static func quotaSummaryBucket(from payload: QuotaSummaryBucketPayload) -> AntigravityQuotaSummaryBucket? {
let bucketId = payload.bucketId?.trimmingCharacters(in: .whitespacesAndNewlines)
let displayName = payload.displayName?.trimmingCharacters(in: .whitespacesAndNewlines)
guard let resolvedBucketId = bucketId, !resolvedBucketId.isEmpty else { return nil }
let resetTime = payload.resetTime.flatMap { Self.parseDate($0) }
return AntigravityQuotaSummaryBucket(
bucketId: resolvedBucketId,
displayName: self.nonEmpty(displayName) ?? resolvedBucketId,
remainingFraction: payload.resolvedRemainingFraction,
resetTime: resetTime,
resetDescription: payload.description,
disabled: payload.disabled ?? false)
}
private static func nonEmpty(_ value: String?) -> String? {
guard let value, !value.isEmpty else { return nil }
return value
}
}
private struct QuotaSummaryResponse: Decodable {
let code: CodeValue?
let message: String?
let response: QuotaSummaryPayload?
let summary: QuotaSummaryPayload?
let description: String?
let groups: [QuotaSummaryGroupPayload]?
var rootPayload: QuotaSummaryPayload? {
guard let groups else { return nil }
return QuotaSummaryPayload(description: self.description, groups: groups)
}
}
private struct QuotaSummaryPayload: Decodable {
let description: String?
let groups: [QuotaSummaryGroupPayload]
init(description: String?, groups: [QuotaSummaryGroupPayload]) {
self.description = description
self.groups = groups
}
private enum CodingKeys: String, CodingKey {
case description
case groups
}
init(from decoder: Decoder) throws {
let container = try decoder.container(keyedBy: CodingKeys.self)
self.description = try container.decodeIfPresent(String.self, forKey: .description)
self.groups = try container.decodeIfPresent([QuotaSummaryGroupPayload].self, forKey: .groups) ?? []
}
}
private struct QuotaSummaryGroupPayload: Decodable {
let displayName: String?
let description: String?
let buckets: [QuotaSummaryBucketPayload]?
}
private struct QuotaSummaryBucketPayload: Decodable {
let bucketId: String?
let displayName: String?
let description: String?
let disabled: Bool?
let remainingFraction: Double?
let remaining: QuotaSummaryRemainingPayload?
let resetTime: String?
var resolvedRemainingFraction: Double? {
self.remainingFraction ?? self.remaining?.remainingFraction
}
}
private struct QuotaSummaryRemainingPayload: Decodable {
let remainingFraction: Double?
private enum CodingKeys: String, CodingKey {
case remainingFraction
case oneofCase = "case"
case value
}
init(from decoder: Decoder) throws {
let container = try decoder.container(keyedBy: CodingKeys.self)
if let remainingFraction = try container.decodeIfPresent(Double.self, forKey: .remainingFraction) {
self.remainingFraction = remainingFraction
return
}
let oneofCase = try container.decodeIfPresent(String.self, forKey: .oneofCase)
if oneofCase == "remainingFraction" {
self.remainingFraction = try container.decodeIfPresent(Double.self, forKey: .value)
} else {
self.remainingFraction = nil
}
}
}
@@ -0,0 +1,777 @@
import Foundation
#if canImport(FoundationNetworking)
import FoundationNetworking
#endif
public enum AntigravityRemoteFetchError: LocalizedError, Sendable, Equatable {
case notLoggedIn
case permissionDenied(String)
case apiError(String)
case parseFailed(String)
public var errorDescription: String? {
switch self {
case .notLoggedIn:
"Antigravity Google auth not found. Use Antigravity login to authenticate."
case let .permissionDenied(message):
"Antigravity remote API permission denied: \(message)"
case let .apiError(message):
"Antigravity remote API error: \(message)"
case let .parseFailed(message):
"Could not parse Antigravity remote usage: \(message)"
}
}
}
public struct AntigravityRemoteUsageFetcher: Sendable {
public var timeout: TimeInterval = 10.0
public var homeDirectory: String
public var environment: [String: String]
public var dataLoader: @Sendable (URLRequest) async throws -> (Data, URLResponse)
public var oauthClientResolver: @Sendable () -> AntigravityOAuthClient?
public var credentialsUpdateHandler: @Sendable (AntigravityOAuthCredentials) async throws -> Void
private static let log = CodexBarLog.logger(LogCategories.antigravity)
private static let userAgent = "antigravity"
private static let baseURL = "https://cloudcode-pa.googleapis.com"
private static let loadCodeAssistEndpoint = "\(baseURL)/v1internal:loadCodeAssist"
private static let onboardUserEndpoint = "\(baseURL)/v1internal:onboardUser"
private static let fetchAvailableModelsEndpoint = "\(baseURL)/v1internal:fetchAvailableModels"
private static let retrieveUserQuotaEndpoint = "\(baseURL)/v1internal:retrieveUserQuota"
private static let refreshSafetyWindow: TimeInterval = 60
private struct FetchContext {
let timeout: TimeInterval
let store: AntigravityOAuthCredentialsStore?
let dataLoader: @Sendable (URLRequest) async throws -> (Data, URLResponse)
let oauthClientResolver: @Sendable () -> AntigravityOAuthClient?
let credentialsUpdateHandler: @Sendable (AntigravityOAuthCredentials) async throws -> Void
func persistCredentials(_ credentials: AntigravityOAuthCredentials) async throws {
if let store {
try store.save(credentials)
} else {
try await self.credentialsUpdateHandler(credentials)
}
}
}
public init(
timeout: TimeInterval = 10.0,
homeDirectory: String = NSHomeDirectory(),
environment: [String: String] = ProcessInfo.processInfo.environment,
dataLoader: @escaping @Sendable (URLRequest) async throws -> (Data, URLResponse) = { request in
try await ProviderHTTPClient.shared.data(for: request)
},
oauthClientResolver: @escaping @Sendable () -> AntigravityOAuthClient? = {
AntigravityOAuthConfig.resolvedClient()
},
credentialsUpdateHandler: @escaping @Sendable (AntigravityOAuthCredentials) async throws -> Void = { _ in })
{
self.timeout = timeout
self.homeDirectory = homeDirectory
self.environment = environment
self.dataLoader = dataLoader
self.oauthClientResolver = oauthClientResolver
self.credentialsUpdateHandler = credentialsUpdateHandler
}
public func fetch() async throws -> AntigravityStatusSnapshot {
let source = try Self.resolveCredentialSource(homeDirectory: self.homeDirectory, environment: self.environment)
guard let credentials = source.credentials else {
throw AntigravityRemoteFetchError.notLoggedIn
}
return try await self.fetchSnapshot(
using: credentials,
store: source.store)
}
private func fetchSnapshot(
using initialCredentials: AntigravityOAuthCredentials,
store: AntigravityOAuthCredentialsStore?) async throws
-> AntigravityStatusSnapshot
{
guard let storedAccessToken = initialCredentials.accessToken?.trimmedNonEmpty else {
throw AntigravityRemoteFetchError.notLoggedIn
}
var credentials = initialCredentials
var accessToken = storedAccessToken
let context = FetchContext(
timeout: self.timeout,
store: store,
dataLoader: self.dataLoader,
oauthClientResolver: self.oauthClientResolver,
credentialsUpdateHandler: self.credentialsUpdateHandler)
if Self.shouldRefresh(expiryDate: credentials.expiryDate, now: Date()) {
guard let refreshToken = credentials.refreshToken?.trimmedNonEmpty else {
throw AntigravityRemoteFetchError.notLoggedIn
}
let refreshed = try await Self.refreshAccessToken(
credentials: credentials,
refreshToken: refreshToken,
context: context)
accessToken = refreshed.accessToken
credentials = refreshed.credentials
if let store {
credentials = try store.load() ?? credentials
}
credentials.accessToken = credentials.accessToken?.trimmedNonEmpty ?? accessToken
}
let claims = Self.extractClaims(from: credentials)
let codeAssist = try await Self.loadCodeAssist(
accessToken: accessToken,
timeout: self.timeout,
dataLoader: self.dataLoader)
let projectId = try await Self.resolveProjectID(
accessToken: accessToken,
storedProjectID: credentials.projectID?.trimmedNonEmpty,
initialResponse: codeAssist,
context: context)
if let projectId, credentials.projectID?.trimmedNonEmpty != projectId {
credentials.projectID = projectId
do {
try await context.persistCredentials(credentials)
} catch {
Self.log.warning("Could not persist Antigravity project ID: \(error.localizedDescription)")
}
}
let models = try await Self.fetchModelQuotas(
accessToken: accessToken,
projectId: projectId,
timeout: self.timeout,
dataLoader: self.dataLoader)
return AntigravityStatusSnapshot(
modelQuotas: models,
accountEmail: claims.email,
accountPlan: Self.resolvePlan(response: codeAssist, claims: claims),
source: .remote)
}
private static func shouldRefresh(expiryDate: Date?, now: Date) -> Bool {
guard let expiryDate else { return false }
return expiryDate.timeIntervalSince(now) <= Self.refreshSafetyWindow
}
private static func loadCodeAssist(
accessToken: String,
timeout: TimeInterval,
dataLoader: @escaping @Sendable (URLRequest) async throws -> (Data, URLResponse)) async throws
-> CodeAssistResponse
{
let body = [
"metadata": [
"ideType": "ANTIGRAVITY",
"platform": "PLATFORM_UNSPECIFIED",
"pluginType": "GEMINI",
],
]
return try await Self.sendRequest(
endpoint: Self.loadCodeAssistEndpoint,
accessToken: accessToken,
body: body,
timeout: timeout,
dataLoader: dataLoader)
}
private static func fetchAvailableModels(
accessToken: String,
projectId: String?,
timeout: TimeInterval,
dataLoader: @escaping @Sendable (URLRequest) async throws -> (Data, URLResponse)) async throws
-> FetchAvailableModelsResponse
{
let body: [String: Any] = if let projectId = projectId?.trimmedNonEmpty {
["project": projectId]
} else {
[:]
}
return try await Self.sendRequest(
endpoint: Self.fetchAvailableModelsEndpoint,
accessToken: accessToken,
body: body,
timeout: timeout,
dataLoader: dataLoader)
}
private static func fetchModelQuotas(
accessToken: String,
projectId: String?,
timeout: TimeInterval,
dataLoader: @escaping @Sendable (URLRequest) async throws -> (Data, URLResponse)) async throws
-> [AntigravityModelQuota]
{
do {
let response = try await Self.fetchAvailableModels(
accessToken: accessToken,
projectId: projectId,
timeout: timeout,
dataLoader: dataLoader)
let modelQuotas = try Self.parseModelQuotas(response)
if Self.shouldVerifyFullRemoteQuotas(modelQuotas) {
let quotaBuckets = try await Self.fetchQuotaBucketsIfPermitted(
accessToken: accessToken,
projectId: projectId,
timeout: timeout,
dataLoader: dataLoader)
guard let quotaBuckets, Self.hasQuotaFractionData(quotaBuckets) else {
return []
}
return Self.mergeVerifiedQuotas(modelQuotas: modelQuotas, verifiedQuotas: quotaBuckets)
}
return modelQuotas
} catch let error as AntigravityRemoteFetchError {
guard case .permissionDenied = error else {
throw error
}
Self.log.info("Falling back to retrieveUserQuota for Antigravity remote usage")
return try await Self.fetchQuotaBucketsIfPermitted(
accessToken: accessToken,
projectId: projectId,
timeout: timeout,
dataLoader: dataLoader) ?? []
}
}
private static func mergeVerifiedQuotas(
modelQuotas: [AntigravityModelQuota],
verifiedQuotas: [AntigravityModelQuota]) -> [AntigravityModelQuota]
{
var verifiedByModelID = Dictionary(
uniqueKeysWithValues: verifiedQuotas.map { (Self.quotaKey($0), $0) })
var merged = modelQuotas.compactMap { modelQuota -> AntigravityModelQuota? in
guard let verifiedQuota = verifiedByModelID.removeValue(forKey: Self.quotaKey(modelQuota)) else {
return nil
}
let resetTime = verifiedQuota.resetTime ?? modelQuota.resetTime
return AntigravityModelQuota(
label: modelQuota.label,
modelId: modelQuota.modelId,
remainingFraction: verifiedQuota.remainingFraction ?? modelQuota.remainingFraction,
resetTime: resetTime,
resetDescription: resetTime.map { UsageFormatter.resetDescription(from: $0) })
}
let unmatchedVerifiedQuotas = verifiedByModelID.values
.filter { $0.remainingFraction != nil }
.sorted { lhs, rhs in
lhs.modelId.localizedCaseInsensitiveCompare(rhs.modelId) == .orderedAscending
}
merged.append(contentsOf: unmatchedVerifiedQuotas)
return merged
}
private static func quotaKey(_ quota: AntigravityModelQuota) -> String {
quota.modelId.trimmingCharacters(in: .whitespacesAndNewlines).lowercased()
}
private static func shouldVerifyFullRemoteQuotas(_ quotas: [AntigravityModelQuota]) -> Bool {
guard !quotas.isEmpty else { return false }
return quotas.allSatisfy { quota in
guard let remaining = quota.remainingFraction else { return false }
return remaining >= 0.999
}
}
private static func hasQuotaFractionData(_ quotas: [AntigravityModelQuota]) -> Bool {
quotas.contains { quota in
quota.remainingFraction != nil
}
}
private static func fetchQuotaBucketsIfPermitted(
accessToken: String,
projectId: String?,
timeout: TimeInterval,
dataLoader: @escaping @Sendable (URLRequest) async throws -> (Data, URLResponse)) async throws
-> [AntigravityModelQuota]?
{
do {
let response = try await Self.retrieveUserQuota(
accessToken: accessToken,
projectId: projectId,
timeout: timeout,
dataLoader: dataLoader)
return try Self.parseQuotaBuckets(response)
} catch let quotaError as AntigravityRemoteFetchError {
guard case .permissionDenied = quotaError else {
throw quotaError
}
Self.log.info("Antigravity remote quota endpoint is not permitted for this account")
return nil
}
}
private static func retrieveUserQuota(
accessToken: String,
projectId: String?,
timeout: TimeInterval,
dataLoader: @escaping @Sendable (URLRequest) async throws -> (Data, URLResponse)) async throws
-> RetrieveUserQuotaResponse
{
let body: [String: Any] = if let projectId = projectId?.trimmedNonEmpty {
["project": projectId]
} else {
[:]
}
return try await Self.sendRequest(
endpoint: Self.retrieveUserQuotaEndpoint,
accessToken: accessToken,
body: body,
timeout: timeout,
dataLoader: dataLoader)
}
private static func resolveProjectID(
accessToken: String,
storedProjectID: String?,
initialResponse: CodeAssistResponse,
context: FetchContext) async throws
-> String?
{
if let storedProjectID {
return storedProjectID
}
if let projectID = initialResponse.projectID {
return projectID
}
guard let tierID = Self.pickOnboardTier(from: initialResponse) else {
return nil
}
let onboardBody: [String: Any] = [
"tierId": tierID,
"metadata": [
"ideType": "ANTIGRAVITY",
"platform": "PLATFORM_UNSPECIFIED",
"pluginType": "GEMINI",
],
]
do {
let onboardResponse: OnboardResponse = try await Self.sendRequest(
endpoint: Self.onboardUserEndpoint,
accessToken: accessToken,
body: onboardBody,
timeout: context.timeout,
dataLoader: context.dataLoader)
if let projectID = onboardResponse.projectID {
return projectID
}
} catch {
Self.log.warning("Antigravity onboarding request failed", metadata: [
"error": "\(error.localizedDescription)",
])
}
for _ in 0..<5 {
try? await Task.sleep(for: .milliseconds(2000))
let refreshed = try await Self.loadCodeAssist(
accessToken: accessToken,
timeout: context.timeout,
dataLoader: context.dataLoader)
if let projectID = refreshed.projectID {
return projectID
}
}
return nil
}
private static func sendRequest<Response: Decodable>(
endpoint: String,
accessToken: String,
body: [String: Any],
timeout: TimeInterval,
dataLoader: @escaping @Sendable (URLRequest) async throws -> (Data, URLResponse)) async throws
-> Response
{
guard let url = URL(string: endpoint) else {
throw AntigravityRemoteFetchError.apiError("Invalid endpoint URL")
}
var request = URLRequest(url: url)
request.httpMethod = "POST"
request.timeoutInterval = timeout
request.setValue("Bearer \(accessToken)", forHTTPHeaderField: "Authorization")
request.setValue("application/json", forHTTPHeaderField: "Content-Type")
request.setValue(Self.userAgent, forHTTPHeaderField: "User-Agent")
request.httpBody = try JSONSerialization.data(withJSONObject: body)
let httpResponse = try await ProviderHTTPTransportHandler(dataLoader).response(for: request)
switch httpResponse.statusCode {
case 200:
break
case 401:
throw AntigravityRemoteFetchError.notLoggedIn
case 403:
let message = String(data: httpResponse.data, encoding: .utf8)?.trimmedNonEmpty ?? "HTTP 403"
throw AntigravityRemoteFetchError.permissionDenied(message)
default:
let message = String(data: httpResponse.data, encoding: .utf8)?.trimmedNonEmpty
?? "HTTP \(httpResponse.statusCode)"
throw AntigravityRemoteFetchError.apiError("HTTP \(httpResponse.statusCode): \(message)")
}
do {
return try JSONDecoder().decode(Response.self, from: httpResponse.data)
} catch {
throw AntigravityRemoteFetchError.parseFailed(error.localizedDescription)
}
}
private static func parseModelQuotas(_ response: FetchAvailableModelsResponse) throws -> [AntigravityModelQuota] {
let models = response.models ?? [:]
return models.compactMap { modelID, model in
guard let quotaInfo = model.quotaInfo else { return nil }
let resetTime = quotaInfo.resetTime.flatMap(Self.parseResetTime(_:))
let label = model.displayName?.trimmedNonEmpty
?? model.label?.trimmedNonEmpty
?? modelID
return AntigravityModelQuota(
label: label,
modelId: modelID,
remainingFraction: quotaInfo.remainingFraction,
resetTime: resetTime,
resetDescription: resetTime.map { UsageFormatter.resetDescription(from: $0) })
}
}
private static func parseQuotaBuckets(_ response: RetrieveUserQuotaResponse) throws -> [AntigravityModelQuota] {
guard let buckets = response.buckets else {
throw AntigravityRemoteFetchError.parseFailed("No quota buckets in response")
}
guard !buckets.isEmpty else { return [] }
var modelQuotaMap: [String: (fraction: Double?, resetTime: String?)] = [:]
for bucket in buckets {
guard let modelID = bucket.modelId?.trimmedNonEmpty else { continue }
let next = (bucket.remainingFraction, bucket.resetTime)
if let existing = modelQuotaMap[modelID] {
let existingValue = existing.fraction ?? Double.greatestFiniteMagnitude
let nextValue = next.0 ?? Double.greatestFiniteMagnitude
if nextValue < existingValue {
modelQuotaMap[modelID] = next
}
} else {
modelQuotaMap[modelID] = next
}
}
return modelQuotaMap.keys.sorted().compactMap { modelID in
guard let info = modelQuotaMap[modelID] else { return nil }
let resetTime = info.resetTime.flatMap(Self.parseResetTime(_:))
return AntigravityModelQuota(
label: modelID,
modelId: modelID,
remainingFraction: info.fraction,
resetTime: resetTime,
resetDescription: resetTime.map { UsageFormatter.resetDescription(from: $0) })
}
}
private static func resolvePlan(response: CodeAssistResponse, claims: TokenClaims) -> String? {
if let planType = response.planInfo?.planType?.trimmedNonEmpty {
return planType
}
switch (response.currentTier?.id?.trimmedNonEmpty, claims.hostedDomain) {
case ("standard-tier", _):
return "Paid"
case ("free-tier", .some):
return "Workspace"
case ("free-tier", .none):
return "Free"
case ("legacy-tier", _):
return "Legacy"
default:
return response.currentTier?.name?.trimmedNonEmpty
}
}
private static func pickOnboardTier(from response: CodeAssistResponse) -> String? {
if let defaultTier = response.allowedTiers?
.first(where: { $0.isDefault == true && $0.id?.trimmedNonEmpty != nil })?.id?.trimmedNonEmpty
{
return defaultTier
}
if let firstTier = response.allowedTiers?
.first(where: { $0.id?.trimmedNonEmpty != nil })?.id?.trimmedNonEmpty
{
return firstTier
}
if let paidTier = response.paidTier?.id?.trimmedNonEmpty {
return paidTier
}
if let currentTier = response.currentTier?.id?.trimmedNonEmpty {
return currentTier
}
return nil
}
private static func parseResetTime(_ value: String) -> Date? {
let formatter = ISO8601DateFormatter()
formatter.formatOptions = [.withInternetDateTime, .withFractionalSeconds]
if let date = formatter.date(from: value) {
return date
}
formatter.formatOptions = [.withInternetDateTime]
return formatter.date(from: value)
}
private static func credentialsStore(homeDirectory: String) -> AntigravityOAuthCredentialsStore {
let homeURL = URL(fileURLWithPath: homeDirectory, isDirectory: true)
return AntigravityOAuthCredentialsStore(fileURL: AntigravityOAuthCredentialsStore.defaultURL(home: homeURL))
}
private static func resolveCredentialSource(
homeDirectory: String,
environment: [String: String]) throws -> (
credentials: AntigravityOAuthCredentials?,
store: AntigravityOAuthCredentialsStore?)
{
let primaryStore = Self.credentialsStore(homeDirectory: homeDirectory)
if let tokenValue = environment[AntigravityOAuthCredentialsStore.environmentCredentialsKey] {
guard let credentials = AntigravityOAuthCredentialsStore.credentials(fromTokenAccountValue: tokenValue)
else {
throw AntigravityRemoteFetchError.parseFailed("Could not decode selected account credentials.")
}
return (credentials, nil)
}
return try (primaryStore.load(), primaryStore)
}
private struct RefreshResult {
let accessToken: String
let credentials: AntigravityOAuthCredentials
}
private static func refreshAccessToken(
credentials: AntigravityOAuthCredentials,
refreshToken: String,
context: FetchContext) async throws
-> RefreshResult
{
let oauthClient = try Self.refreshOAuthClient(
from: credentials,
oauthClientResolver: context.oauthClientResolver)
var request = URLRequest(url: AntigravityOAuthConfig.tokenURL)
request.httpMethod = "POST"
request.timeoutInterval = context.timeout
request.setValue("application/x-www-form-urlencoded", forHTTPHeaderField: "Content-Type")
request.httpBody = Self.formBody([
"client_id": oauthClient.clientID,
"client_secret": oauthClient.clientSecret,
"refresh_token": refreshToken,
"grant_type": "refresh_token",
])
let httpResponse = try await ProviderHTTPTransportHandler(context.dataLoader).response(for: request)
guard httpResponse.statusCode == 200 else {
throw AntigravityRemoteFetchError.notLoggedIn
}
guard let json = try JSONSerialization.jsonObject(with: httpResponse.data) as? [String: Any],
let accessToken = json["access_token"] as? String
else {
throw AntigravityRemoteFetchError.parseFailed("Could not parse refresh response")
}
let updatedCredentials = Self.updatedCredentials(credentials, refreshResponse: json)
try await context.persistCredentials(updatedCredentials)
return RefreshResult(accessToken: accessToken, credentials: updatedCredentials)
}
private static func refreshOAuthClient(
from credentials: AntigravityOAuthCredentials,
oauthClientResolver: @escaping @Sendable () -> AntigravityOAuthClient?) throws
-> AntigravityOAuthClient
{
if let clientID = credentials.clientID?.trimmedNonEmpty,
let clientSecret = credentials.clientSecret?.trimmedNonEmpty
{
return AntigravityOAuthClient(clientID: clientID, clientSecret: clientSecret)
}
guard let client = oauthClientResolver() else {
throw AntigravityRemoteFetchError.apiError(AntigravityOAuthConfig.missingCredentialsMessage)
}
return client
}
private static func updatedCredentials(
_ credentials: AntigravityOAuthCredentials,
refreshResponse: [String: Any]) -> AntigravityOAuthCredentials
{
var credentials = credentials
if let accessToken = refreshResponse["access_token"] as? String {
credentials.accessToken = accessToken
}
if let expiresIn = refreshResponse["expires_in"] as? Double {
credentials.expiryDateMilliseconds = (Date().timeIntervalSince1970 + expiresIn) * 1000
}
if let expiresIn = refreshResponse["expires_in"] as? Int {
credentials.expiryDateMilliseconds = (Date().timeIntervalSince1970 + Double(expiresIn)) * 1000
}
if let idToken = refreshResponse["id_token"] as? String {
credentials.idToken = idToken
}
return credentials
}
private static func formBody(_ values: [String: String]) -> Data? {
var components = URLComponents()
components.queryItems = values.map { key, value in
URLQueryItem(name: key, value: value)
}
return components.query?.data(using: .utf8)
}
private struct TokenClaims {
let email: String?
let hostedDomain: String?
}
private static func extractClaims(from credentials: AntigravityOAuthCredentials) -> TokenClaims {
let tokenClaims = Self.extractClaimsFromToken(credentials.idToken)
return TokenClaims(
email: tokenClaims.email ?? credentials.email?.trimmedNonEmpty,
hostedDomain: tokenClaims.hostedDomain)
}
private static func extractClaimsFromToken(_ idToken: String?) -> TokenClaims {
guard let idToken else {
return TokenClaims(email: nil, hostedDomain: nil)
}
let parts = idToken.components(separatedBy: ".")
guard parts.count >= 2 else {
return TokenClaims(email: nil, hostedDomain: nil)
}
var payload = parts[1]
.replacingOccurrences(of: "-", with: "+")
.replacingOccurrences(of: "_", with: "/")
let remainder = payload.count % 4
if remainder > 0 {
payload += String(repeating: "=", count: 4 - remainder)
}
guard let data = Data(base64Encoded: payload, options: .ignoreUnknownCharacters),
let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any]
else {
return TokenClaims(email: nil, hostedDomain: nil)
}
return TokenClaims(
email: (json["email"] as? String)?.trimmedNonEmpty,
hostedDomain: (json["hd"] as? String)?.trimmedNonEmpty)
}
}
extension String {
fileprivate var trimmedNonEmpty: String? {
let trimmed = self.trimmingCharacters(in: .whitespacesAndNewlines)
return trimmed.isEmpty ? nil : trimmed
}
}
private struct ProjectReference: Decodable {
let value: String?
init(from decoder: Decoder) throws {
let single = try decoder.singleValueContainer()
if let stringValue = try? single.decode(String.self) {
self.value = stringValue
return
}
let keyed = try decoder.container(keyedBy: CodingKeys.self)
self.value = try keyed.decodeIfPresent(String.self, forKey: .id)
?? keyed.decodeIfPresent(String.self, forKey: .projectID)
}
private enum CodingKeys: String, CodingKey {
case id
case projectID = "projectId"
}
}
private struct CodeAssistResponse: Decodable {
let planInfo: CodeAssistPlanInfo?
let currentTier: TierInfo?
let paidTier: TierInfo?
let allowedTiers: [AllowedTier]?
let cloudaicompanionProject: ProjectReference?
var projectID: String? {
self.cloudaicompanionProject?.value?.trimmingCharacters(in: .whitespacesAndNewlines).nilIfEmpty
}
}
private struct CodeAssistPlanInfo: Decodable {
let planType: String?
}
private struct TierInfo: Decodable {
let id: String?
let name: String?
}
private struct AllowedTier: Decodable {
let id: String?
let isDefault: Bool?
}
private struct OnboardResponse: Decodable {
let response: OnboardInnerResponse?
var projectID: String? {
self.response?.cloudaicompanionProject?.value?.trimmingCharacters(in: .whitespacesAndNewlines).nilIfEmpty
}
}
private struct OnboardInnerResponse: Decodable {
let cloudaicompanionProject: ProjectReference?
}
private struct FetchAvailableModelsResponse: Decodable {
let models: [String: AntigravityRemoteModel]?
}
private struct RetrieveUserQuotaResponse: Decodable {
let buckets: [RetrieveUserQuotaBucket]?
}
private struct RetrieveUserQuotaBucket: Decodable {
let modelId: String?
let remainingFraction: Double?
let resetTime: String?
}
private struct AntigravityRemoteModel: Decodable {
let displayName: String?
let label: String?
let quotaInfo: AntigravityRemoteQuotaInfo?
}
private struct AntigravityRemoteQuotaInfo: Decodable {
let remainingFraction: Double?
let resetTime: String?
}
extension String? {
fileprivate var trimmedNonEmpty: String? {
self?.trimmingCharacters(in: .whitespacesAndNewlines).nilIfEmpty
}
}
extension String {
fileprivate var nilIfEmpty: String? {
self.isEmpty ? nil : self
}
}
@@ -0,0 +1,128 @@
import Foundation
extension AntigravityStatusProbe {
private static let processProbeLog = CodexBarLog.logger(LogCategories.antigravity)
private enum ProcessSnapshotFetchFailure {
case antigravity(AntigravityStatusProbeError)
case url(URLError)
case cancellation
case other(String)
init(_ error: Error) {
if let error = error as? AntigravityStatusProbeError {
self = .antigravity(error)
} else if let error = error as? URLError, error.code == .cancelled {
self = .cancellation
} else if let error = error as? URLError {
self = .url(error)
} else if error is CancellationError {
self = .cancellation
} else {
self = .other(error.localizedDescription)
}
}
var error: Error {
switch self {
case let .antigravity(error):
error
case let .url(error):
error
case .cancellation:
CancellationError()
case let .other(message):
AntigravityStatusProbeError.apiError(message)
}
}
}
private enum ProcessSnapshotFetchOutcome {
case success(index: Int, snapshot: AntigravityStatusSnapshot)
case failure(index: Int, pid: Int, failure: ProcessSnapshotFetchFailure)
}
static func fetchProcessSnapshots(
processInfos: [ProcessInfoResult],
fetch: @escaping @Sendable (ProcessInfoResult) async throws -> AntigravityStatusSnapshot)
async throws -> (snapshots: [AntigravityStatusSnapshot], lastError: Error?)
{
let outcomes = await withTaskGroup(of: ProcessSnapshotFetchOutcome.self) { group in
for (index, processInfo) in processInfos.enumerated() {
group.addTask {
do {
return try await .success(index: index, snapshot: fetch(processInfo))
} catch {
return .failure(index: index, pid: processInfo.pid, failure: .init(error))
}
}
}
var ordered = [ProcessSnapshotFetchOutcome?](repeating: nil, count: processInfos.count)
for await outcome in group {
switch outcome {
case let .success(index, _), let .failure(index, _, _):
ordered[index] = outcome
}
}
return ordered.compactMap(\.self)
}
var snapshots: [AntigravityStatusSnapshot] = []
var lastError: Error?
for outcome in outcomes {
switch outcome {
case let .success(_, snapshot):
snapshots.append(snapshot)
case let .failure(_, pid, failure):
if case .cancellation = failure {
throw CancellationError()
}
let error = failure.error
lastError = error
Self.processProbeLog.debug("Antigravity local process probe failed", metadata: [
"pid": "\(pid)",
"error": error.localizedDescription,
])
}
}
return (snapshots, lastError)
}
static func fetch(
processInfo: ProcessInfoResult,
timeout: TimeInterval,
deadline: Date) async throws -> AntigravityStatusSnapshot
{
guard let portTimeout = timeoutForNextAttempt(timeout: timeout, deadline: deadline) else {
throw AntigravityStatusProbeError.timedOut
}
let ports = try await Self.listeningPorts(pid: processInfo.pid, timeout: portTimeout)
let endpoint = try await Self.resolveWorkingEndpoint(
candidateEndpoints: Self.connectionCandidates(
listeningPorts: ports,
languageServerCSRFToken: processInfo.csrfToken,
extensionServerPort: processInfo.extensionPort,
extensionServerCSRFToken: processInfo.extensionServerCSRFToken),
timeout: timeout,
deadline: deadline)
let context = RequestContext(
endpoints: Self.requestEndpoints(
resolvedEndpoint: endpoint,
listeningPorts: ports,
languageServerCSRFToken: processInfo.csrfToken,
extensionServerPort: processInfo.extensionPort,
extensionServerCSRFToken: processInfo.extensionServerCSRFToken),
timeout: timeout,
deadline: deadline)
return try await Self.fetchSnapshot(context: context)
}
static func timeoutForNextAttempt(timeout: TimeInterval, deadline: Date?) -> TimeInterval? {
guard let deadline else { return timeout }
let remaining = deadline.timeIntervalSinceNow
guard remaining > 0 else { return nil }
return min(timeout, remaining)
}
}
@@ -0,0 +1,23 @@
extension AntigravityStatusProbe {
static func cliEndpoints(ports: [Int]) -> [AntigravityConnectionEndpoint] {
ports.flatMap { port in
self.localProbeSchemes.map { scheme in
AntigravityConnectionEndpoint(
scheme: scheme,
port: port,
csrfToken: "",
source: .cliHTTPS)
}
}
}
static var localProbeSchemes: [String] {
#if os(Linux)
// FoundationNetworking cannot trust Antigravity's self-signed TLS cert.
// Requests remain pinned to 127.0.0.1 in makeRequest.
["https", "http"]
#else
["https"]
#endif
}
}
@@ -0,0 +1,137 @@
import Foundation
/// Parses Linux `/proc/<pid>/net/tcp{,6}` output to recover the listening ports
/// owned by a process. The parsing is platform-independent for focused tests.
enum ProcNetTCPListeningPortParser {
/// The `st` column value for a socket in the LISTEN state.
private static let listenState = "0A"
/// Extracts the socket inode from a `/proc/<pid>/fd` symlink destination such
/// as `socket:[12345]`. Returns nil for non-socket descriptors.
static func socketInode(fromLink destination: String) -> String? {
let prefix = "socket:["
guard destination.hasPrefix(prefix), destination.hasSuffix("]") else { return nil }
let inode = destination.dropFirst(prefix.count).dropLast()
return inode.isEmpty ? nil : String(inode)
}
/// Returns the local ports of LISTEN sockets whose inode is in `socketInodes`.
///
/// `content` is the raw text of a process-scoped `tcp` or `tcp6` table. Each
/// row encodes the local endpoint as `ADDRESS:PORT` (for example,
/// `0100007F:1F90` uses port 8080) and the owning socket inode in column ten.
static func listeningPorts(_ content: String, socketInodes: Set<String>) -> Set<Int> {
var ports: Set<Int> = []
for line in content.split(separator: "\n") {
let columns = line.split(separator: " ", omittingEmptySubsequences: true)
// Columns: sl local_address rem_address st ... uid timeout inode
guard columns.count > 9,
columns[3] == self.listenState,
socketInodes.contains(String(columns[9]))
else { continue }
let localAddress = columns[1]
guard let separator = localAddress.lastIndex(of: ":"),
let port = Int(localAddress[localAddress.index(after: separator)...], radix: 16),
(0...Int(UInt16.max)).contains(port)
else { continue }
ports.insert(port)
}
return ports
}
}
extension AntigravityStatusProbe {
/// Resolves the TCP ports the process `pid` is listening on. Uses `lsof` when
/// present (the common denominator across macOS and Linux) and falls back to
/// the kernel's `/proc` interface on Linux hosts without `lsof`.
static func listeningPorts(pid: Int, timeout: TimeInterval) async throws -> [Int] {
let lsof = ["/usr/sbin/lsof", "/usr/bin/lsof"].first(where: {
FileManager.default.isExecutableFile(atPath: $0)
})
if let lsof {
return try await Self.lsofListeningPorts(lsof: lsof, pid: pid, timeout: timeout)
}
#if os(Linux)
// `lsof` is frequently absent on minimal Linux hosts. Fall back to the
// kernel's /proc interface, mirroring the /proc/<pid>/cwd fallback that
// LocalAgentSessionScanner.cwdByPID already uses when lsof is missing.
let ports = Self.procListeningPorts(pid: pid)
if ports.isEmpty {
throw AntigravityStatusProbeError.portDetectionFailed("no listening ports found")
}
return ports
#else
throw AntigravityStatusProbeError.portDetectionFailed("lsof not available")
#endif
}
private static func lsofListeningPorts(
lsof: String,
pid: Int,
timeout: TimeInterval) async throws -> [Int]
{
let env = ProcessInfo.processInfo.environment
let result: SubprocessResult
do {
result = try await SubprocessRunner.run(
binary: lsof,
arguments: ["-nP", "-iTCP", "-sTCP:LISTEN", "-a", "-p", String(pid)],
environment: env,
timeout: timeout,
label: "antigravity-lsof")
} catch let SubprocessRunnerError.nonZeroExit(code, stderr)
where code == 1 && stderr.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty
{
throw AntigravityStatusProbeError.portDetectionFailed("no listening ports found")
}
let ports = Self.parseListeningPorts(result.stdout)
if ports.isEmpty {
throw AntigravityStatusProbeError.portDetectionFailed("no listening ports found")
}
return ports
}
private static func parseListeningPorts(_ output: String) -> [Int] {
guard let regex = try? NSRegularExpression(pattern: #":(\d+)\s+\(LISTEN\)"#) else { return [] }
let range = NSRange(output.startIndex..<output.endIndex, in: output)
var ports: Set<Int> = []
regex.enumerateMatches(in: output, options: [], range: range) { match, _, _ in
guard let match,
let range = Range(match.range(at: 1), in: output),
let value = Int(output[range]) else { return }
ports.insert(value)
}
return ports.sorted()
}
/// Recovers the listening ports owned by `pid` by matching its open socket
/// inodes against the TCP tables from the same process/network namespace.
static func procListeningPorts(pid: Int, procRoot: String = "/proc") -> [Int] {
let processRoot = "\(procRoot)/\(pid)"
let inodes = Self.socketInodes(processRoot: processRoot)
guard !inodes.isEmpty else { return [] }
var ports: Set<Int> = []
for path in ["\(processRoot)/net/tcp", "\(processRoot)/net/tcp6"] {
guard let content = try? String(contentsOfFile: path, encoding: .utf8) else { continue }
ports.formUnion(ProcNetTCPListeningPortParser.listeningPorts(content, socketInodes: inodes))
}
return ports.sorted()
}
/// Collects the socket inodes referenced by the process's open descriptors.
private static func socketInodes(processRoot: String) -> Set<String> {
let fdDirectory = "\(processRoot)/fd"
guard let entries = try? FileManager.default.contentsOfDirectory(atPath: fdDirectory) else { return [] }
var inodes: Set<String> = []
for entry in entries {
guard let destination = try? FileManager.default.destinationOfSymbolicLink(
atPath: "\(fdDirectory)/\(entry)"),
let inode = ProcNetTCPListeningPortParser.socketInode(fromLink: destination)
else { continue }
inodes.insert(inode)
}
return inodes
}
}
@@ -0,0 +1,111 @@
import Foundation
struct UserStatusResponse: Decodable {
let code: CodeValue?
let message: String?
let userStatus: UserStatus?
}
struct CommandModelConfigResponse: Decodable {
let code: CodeValue?
let message: String?
let clientModelConfigs: [ModelConfig]?
}
struct UserStatus: Decodable {
let email: String?
let planStatus: PlanStatus?
let cascadeModelConfigData: ModelConfigData?
let userTier: UserTier?
}
struct UserTier: Decodable {
let id: String?
let name: String?
let description: String?
var preferredName: String? {
guard let value = self.name?.trimmingCharacters(in: .whitespacesAndNewlines) else { return nil }
return value.isEmpty ? nil : value
}
}
struct PlanStatus: Decodable {
let planInfo: PlanInfo?
}
struct PlanInfo: Decodable {
let planName: String?
let planDisplayName: String?
let displayName: String?
let productName: String?
let planShortName: String?
var preferredName: String? {
let candidates = [
self.planDisplayName,
self.displayName,
self.productName,
self.planName,
self.planShortName,
]
for candidate in candidates {
guard let value = candidate?.trimmingCharacters(in: .whitespacesAndNewlines) else { continue }
if !value.isEmpty { return value }
}
return nil
}
}
struct ModelConfigData: Decodable {
let clientModelConfigs: [ModelConfig]?
}
struct ModelConfig: Decodable {
let label: String
let modelOrAlias: ModelAlias
let quotaInfo: QuotaInfo?
}
struct ModelAlias: Decodable {
let model: String
}
struct QuotaInfo: Decodable {
let remainingFraction: Double?
let resetTime: String?
}
enum CodeValue: Decodable {
case int(Int)
case string(String)
var isOK: Bool {
switch self {
case let .int(value):
value == 0
case let .string(value):
value.lowercased() == "ok" || value.lowercased() == "success" || value == "0"
}
}
var rawValue: String {
switch self {
case let .int(value): "\(value)"
case let .string(value): value
}
}
init(from decoder: Decoder) throws {
let container = try decoder.singleValueContainer()
if let value = try? container.decode(Int.self) {
self = .int(value)
return
}
if let value = try? container.decode(String.self) {
self = .string(value)
return
}
throw DecodingError.dataCorruptedError(in: container, debugDescription: "Unsupported code type")
}
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,30 @@
import Foundation
public enum AntigravityUsageDataSource: String, CaseIterable, Identifiable, Sendable {
case auto
case oauth
case cli
public var id: String {
self.rawValue
}
public var displayName: String {
switch self {
case .auto: "Auto"
case .oauth: "Google OAuth"
case .cli: "Local API / agy CLI"
}
}
public var sourceLabel: String {
switch self {
case .auto:
"auto"
case .oauth:
"oauth"
case .cli:
"cli"
}
}
}
@@ -0,0 +1,191 @@
import Foundation
#if os(macOS)
/// Fetches Augment usage via `auggie account status` CLI command
public struct AuggieCLIProbe: Sendable {
private static let log = CodexBarLog.logger(LogCategories.auggieCLI)
public init() {}
public func fetch() async throws -> AugmentStatusSnapshot {
let output = try await self.runAuggieAccountStatus()
return try self.parse(output)
}
/// Timeout for the `auggie account status` command.
private static let commandTimeout: TimeInterval = 15
private func runAuggieAccountStatus() async throws -> String {
let env = ProcessInfo.processInfo.environment
let loginPATH = LoginShellPathCache.shared.current
let executable = BinaryLocator.resolveAuggieBinary(env: env, loginPATH: loginPATH) ?? "auggie"
var pathEnv = env
pathEnv["PATH"] = PathBuilder.effectivePATH(
purposes: [.tty, .nodeTooling],
env: env,
loginPATH: loginPATH)
let result = try await SubprocessRunner.run(
binary: executable,
arguments: ["account", "status"],
environment: pathEnv,
timeout: Self.commandTimeout,
label: "auggie-account-status")
let output = result.stdout
let errorOutput = result.stderr
guard !output.isEmpty else {
if !errorOutput.isEmpty {
Self.log.error("Auggie stderr: \(errorOutput)")
}
throw AuggieCLIError.noOutput
}
// Check for auth errors
if output.contains("Authentication failed") || output.contains("auggie login") {
throw AuggieCLIError.notAuthenticated
}
return output
}
func parse(_ output: String) throws -> AugmentStatusSnapshot {
// Legacy output:
// Max Plan 450,000 credits / month
// 11,657 remaining · 953,170 / 964,827 credits used
// 2 days remaining in this billing cycle (ends 1/8/2026)
//
// Current output (2026+):
// 319,054 credits remaining Max Plan
// 450,000 credits / month
// 9 days remaining in this billing cycle (ends 6/9/2026)
var maxCredits: Int?
var remaining: Int?
var used: Int?
var total: Int?
var billingCycleEnd: Date?
for line in output.split(separator: "\n") {
let trimmed = line.trimmingCharacters(in: .whitespaces)
if trimmed.contains("credits / month") {
if let match = trimmed.range(of: #"([\d,]+)\s+credits\s*/\s*month"#, options: .regularExpression) {
let numberStr = String(trimmed[match]).replacingOccurrences(of: ",", with: "")
.replacingOccurrences(of: " credits", with: "")
.replacingOccurrences(of: " / month", with: "")
maxCredits = Int(numberStr)
total = total ?? Int(numberStr)
}
} else if trimmed.contains("Max Plan"), trimmed.contains("credits"), !trimmed.contains("remaining") {
if let match = trimmed.range(of: #"([\d,]+)\s+credits"#, options: .regularExpression) {
let numberStr = String(trimmed[match]).replacingOccurrences(of: ",", with: "")
.replacingOccurrences(of: " credits", with: "")
maxCredits = Int(numberStr)
}
}
if trimmed.contains("credits remaining"), !trimmed.contains("billing cycle") {
if let match = trimmed.range(of: #"([\d,]+)\s+credits\s+remaining"#, options: .regularExpression) {
let numberStr = String(trimmed[match]).replacingOccurrences(of: ",", with: "")
.replacingOccurrences(of: " credits", with: "")
.replacingOccurrences(of: " remaining", with: "")
remaining = Int(numberStr)
}
}
// Parse "11,657 remaining · 953,170 / 964,827 credits used"
if trimmed.contains("remaining"), trimmed.contains("credits used") {
if let remMatch = trimmed.range(of: #"([\d,]+)\s+remaining"#, options: .regularExpression) {
let numStr = String(trimmed[remMatch])
.replacingOccurrences(of: ",", with: "")
.replacingOccurrences(of: " remaining", with: "")
remaining = Int(numStr)
}
if let usedMatch = trimmed.range(
of: #"([\d,]+)\s*/\s*([\d,]+)\s+credits used"#,
options: .regularExpression)
{
let parts = String(trimmed[usedMatch])
.replacingOccurrences(of: " credits used", with: "")
.split(separator: "/")
if parts.count == 2 {
used = Int(parts[0].replacingOccurrences(of: ",", with: "")
.trimmingCharacters(in: .whitespaces))
total = Int(parts[1].replacingOccurrences(of: ",", with: "")
.trimmingCharacters(in: .whitespaces))
}
}
}
if trimmed.contains("billing cycle"), trimmed.contains("ends") {
if let dateMatch = trimmed.range(of: #"ends\s+([\d/]+)"#, options: .regularExpression) {
let dateStr = String(trimmed[dateMatch])
.replacingOccurrences(of: "ends", with: "")
.trimmingCharacters(in: .whitespaces)
let formatter = DateFormatter()
formatter.dateFormat = "M/d/yyyy"
formatter.locale = Locale(identifier: "en_US_POSIX")
formatter.timeZone = TimeZone.current
billingCycleEnd = formatter.date(from: dateStr)
}
}
}
guard let finalRemaining = remaining else {
Self.log.error("Failed to parse auggie output: \(output)")
throw AuggieCLIError.parseError("Could not extract credits from output")
}
let finalTotal = total ?? maxCredits
guard let finalTotal else {
Self.log.error("Failed to parse auggie output: \(output)")
throw AuggieCLIError.parseError("Could not extract credits from output")
}
let finalUsed = used ?? max(0, finalTotal - finalRemaining)
return AugmentStatusSnapshot(
creditsRemaining: Double(finalRemaining),
creditsUsed: Double(finalUsed),
creditsLimit: Double(finalTotal),
billingCycleEnd: billingCycleEnd,
accountEmail: nil,
accountPlan: maxCredits.map { "\($0.formatted()) credits/month" },
rawJSON: nil)
}
}
#else
public struct AuggieCLIProbe: Sendable {
public init() {}
public func fetch() async throws -> AugmentStatusSnapshot {
throw AugmentStatusProbeError.notSupported
}
}
#endif
public enum AuggieCLIError: LocalizedError {
case noOutput
case notAuthenticated
case parseError(String)
public var errorDescription: String? {
switch self {
case .noOutput:
"Auggie CLI returned no output"
case .notAuthenticated:
"Not authenticated. Run 'auggie login' to authenticate."
case let .parseError(msg):
"Failed to parse auggie output: \(msg)"
}
}
}

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