chore: import upstream snapshot with attribution
This commit is contained in:
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
+23
@@ -0,0 +1,23 @@
|
||||
extension AntigravityStatusProbe {
|
||||
static func cliEndpoints(ports: [Int]) -> [AntigravityConnectionEndpoint] {
|
||||
ports.flatMap { port in
|
||||
self.localProbeSchemes.map { scheme in
|
||||
AntigravityConnectionEndpoint(
|
||||
scheme: scheme,
|
||||
port: port,
|
||||
csrfToken: "",
|
||||
source: .cliHTTPS)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
static var localProbeSchemes: [String] {
|
||||
#if os(Linux)
|
||||
// FoundationNetworking cannot trust Antigravity's self-signed TLS cert.
|
||||
// Requests remain pinned to 127.0.0.1 in makeRequest.
|
||||
["https", "http"]
|
||||
#else
|
||||
["https"]
|
||||
#endif
|
||||
}
|
||||
}
|
||||
+137
@@ -0,0 +1,137 @@
|
||||
import Foundation
|
||||
|
||||
/// Parses Linux `/proc/<pid>/net/tcp{,6}` output to recover the listening ports
|
||||
/// owned by a process. The parsing is platform-independent for focused tests.
|
||||
enum ProcNetTCPListeningPortParser {
|
||||
/// The `st` column value for a socket in the LISTEN state.
|
||||
private static let listenState = "0A"
|
||||
|
||||
/// Extracts the socket inode from a `/proc/<pid>/fd` symlink destination such
|
||||
/// as `socket:[12345]`. Returns nil for non-socket descriptors.
|
||||
static func socketInode(fromLink destination: String) -> String? {
|
||||
let prefix = "socket:["
|
||||
guard destination.hasPrefix(prefix), destination.hasSuffix("]") else { return nil }
|
||||
let inode = destination.dropFirst(prefix.count).dropLast()
|
||||
return inode.isEmpty ? nil : String(inode)
|
||||
}
|
||||
|
||||
/// Returns the local ports of LISTEN sockets whose inode is in `socketInodes`.
|
||||
///
|
||||
/// `content` is the raw text of a process-scoped `tcp` or `tcp6` table. Each
|
||||
/// row encodes the local endpoint as `ADDRESS:PORT` (for example,
|
||||
/// `0100007F:1F90` uses port 8080) and the owning socket inode in column ten.
|
||||
static func listeningPorts(_ content: String, socketInodes: Set<String>) -> Set<Int> {
|
||||
var ports: Set<Int> = []
|
||||
for line in content.split(separator: "\n") {
|
||||
let columns = line.split(separator: " ", omittingEmptySubsequences: true)
|
||||
// Columns: sl local_address rem_address st ... uid timeout inode
|
||||
guard columns.count > 9,
|
||||
columns[3] == self.listenState,
|
||||
socketInodes.contains(String(columns[9]))
|
||||
else { continue }
|
||||
let localAddress = columns[1]
|
||||
guard let separator = localAddress.lastIndex(of: ":"),
|
||||
let port = Int(localAddress[localAddress.index(after: separator)...], radix: 16),
|
||||
(0...Int(UInt16.max)).contains(port)
|
||||
else { continue }
|
||||
ports.insert(port)
|
||||
}
|
||||
return ports
|
||||
}
|
||||
}
|
||||
|
||||
extension AntigravityStatusProbe {
|
||||
/// Resolves the TCP ports the process `pid` is listening on. Uses `lsof` when
|
||||
/// present (the common denominator across macOS and Linux) and falls back to
|
||||
/// the kernel's `/proc` interface on Linux hosts without `lsof`.
|
||||
static func listeningPorts(pid: Int, timeout: TimeInterval) async throws -> [Int] {
|
||||
let lsof = ["/usr/sbin/lsof", "/usr/bin/lsof"].first(where: {
|
||||
FileManager.default.isExecutableFile(atPath: $0)
|
||||
})
|
||||
|
||||
if let lsof {
|
||||
return try await Self.lsofListeningPorts(lsof: lsof, pid: pid, timeout: timeout)
|
||||
}
|
||||
|
||||
#if os(Linux)
|
||||
// `lsof` is frequently absent on minimal Linux hosts. Fall back to the
|
||||
// kernel's /proc interface, mirroring the /proc/<pid>/cwd fallback that
|
||||
// LocalAgentSessionScanner.cwdByPID already uses when lsof is missing.
|
||||
let ports = Self.procListeningPorts(pid: pid)
|
||||
if ports.isEmpty {
|
||||
throw AntigravityStatusProbeError.portDetectionFailed("no listening ports found")
|
||||
}
|
||||
return ports
|
||||
#else
|
||||
throw AntigravityStatusProbeError.portDetectionFailed("lsof not available")
|
||||
#endif
|
||||
}
|
||||
|
||||
private static func lsofListeningPorts(
|
||||
lsof: String,
|
||||
pid: Int,
|
||||
timeout: TimeInterval) async throws -> [Int]
|
||||
{
|
||||
let env = ProcessInfo.processInfo.environment
|
||||
let result: SubprocessResult
|
||||
do {
|
||||
result = try await SubprocessRunner.run(
|
||||
binary: lsof,
|
||||
arguments: ["-nP", "-iTCP", "-sTCP:LISTEN", "-a", "-p", String(pid)],
|
||||
environment: env,
|
||||
timeout: timeout,
|
||||
label: "antigravity-lsof")
|
||||
} catch let SubprocessRunnerError.nonZeroExit(code, stderr)
|
||||
where code == 1 && stderr.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty
|
||||
{
|
||||
throw AntigravityStatusProbeError.portDetectionFailed("no listening ports found")
|
||||
}
|
||||
let ports = Self.parseListeningPorts(result.stdout)
|
||||
if ports.isEmpty {
|
||||
throw AntigravityStatusProbeError.portDetectionFailed("no listening ports found")
|
||||
}
|
||||
return ports
|
||||
}
|
||||
|
||||
private static func parseListeningPorts(_ output: String) -> [Int] {
|
||||
guard let regex = try? NSRegularExpression(pattern: #":(\d+)\s+\(LISTEN\)"#) else { return [] }
|
||||
let range = NSRange(output.startIndex..<output.endIndex, in: output)
|
||||
var ports: Set<Int> = []
|
||||
regex.enumerateMatches(in: output, options: [], range: range) { match, _, _ in
|
||||
guard let match,
|
||||
let range = Range(match.range(at: 1), in: output),
|
||||
let value = Int(output[range]) else { return }
|
||||
ports.insert(value)
|
||||
}
|
||||
return ports.sorted()
|
||||
}
|
||||
|
||||
/// Recovers the listening ports owned by `pid` by matching its open socket
|
||||
/// inodes against the TCP tables from the same process/network namespace.
|
||||
static func procListeningPorts(pid: Int, procRoot: String = "/proc") -> [Int] {
|
||||
let processRoot = "\(procRoot)/\(pid)"
|
||||
let inodes = Self.socketInodes(processRoot: processRoot)
|
||||
guard !inodes.isEmpty else { return [] }
|
||||
var ports: Set<Int> = []
|
||||
for path in ["\(processRoot)/net/tcp", "\(processRoot)/net/tcp6"] {
|
||||
guard let content = try? String(contentsOfFile: path, encoding: .utf8) else { continue }
|
||||
ports.formUnion(ProcNetTCPListeningPortParser.listeningPorts(content, socketInodes: inodes))
|
||||
}
|
||||
return ports.sorted()
|
||||
}
|
||||
|
||||
/// Collects the socket inodes referenced by the process's open descriptors.
|
||||
private static func socketInodes(processRoot: String) -> Set<String> {
|
||||
let fdDirectory = "\(processRoot)/fd"
|
||||
guard let entries = try? FileManager.default.contentsOfDirectory(atPath: fdDirectory) else { return [] }
|
||||
var inodes: Set<String> = []
|
||||
for entry in entries {
|
||||
guard let destination = try? FileManager.default.destinationOfSymbolicLink(
|
||||
atPath: "\(fdDirectory)/\(entry)"),
|
||||
let inode = ProcNetTCPListeningPortParser.socketInode(fromLink: destination)
|
||||
else { continue }
|
||||
inodes.insert(inode)
|
||||
}
|
||||
return inodes
|
||||
}
|
||||
}
|
||||
+111
@@ -0,0 +1,111 @@
|
||||
import Foundation
|
||||
|
||||
struct UserStatusResponse: Decodable {
|
||||
let code: CodeValue?
|
||||
let message: String?
|
||||
let userStatus: UserStatus?
|
||||
}
|
||||
|
||||
struct CommandModelConfigResponse: Decodable {
|
||||
let code: CodeValue?
|
||||
let message: String?
|
||||
let clientModelConfigs: [ModelConfig]?
|
||||
}
|
||||
|
||||
struct UserStatus: Decodable {
|
||||
let email: String?
|
||||
let planStatus: PlanStatus?
|
||||
let cascadeModelConfigData: ModelConfigData?
|
||||
let userTier: UserTier?
|
||||
}
|
||||
|
||||
struct UserTier: Decodable {
|
||||
let id: String?
|
||||
let name: String?
|
||||
let description: String?
|
||||
|
||||
var preferredName: String? {
|
||||
guard let value = self.name?.trimmingCharacters(in: .whitespacesAndNewlines) else { return nil }
|
||||
return value.isEmpty ? nil : value
|
||||
}
|
||||
}
|
||||
|
||||
struct PlanStatus: Decodable {
|
||||
let planInfo: PlanInfo?
|
||||
}
|
||||
|
||||
struct PlanInfo: Decodable {
|
||||
let planName: String?
|
||||
let planDisplayName: String?
|
||||
let displayName: String?
|
||||
let productName: String?
|
||||
let planShortName: String?
|
||||
|
||||
var preferredName: String? {
|
||||
let candidates = [
|
||||
self.planDisplayName,
|
||||
self.displayName,
|
||||
self.productName,
|
||||
self.planName,
|
||||
self.planShortName,
|
||||
]
|
||||
for candidate in candidates {
|
||||
guard let value = candidate?.trimmingCharacters(in: .whitespacesAndNewlines) else { continue }
|
||||
if !value.isEmpty { return value }
|
||||
}
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
struct ModelConfigData: Decodable {
|
||||
let clientModelConfigs: [ModelConfig]?
|
||||
}
|
||||
|
||||
struct ModelConfig: Decodable {
|
||||
let label: String
|
||||
let modelOrAlias: ModelAlias
|
||||
let quotaInfo: QuotaInfo?
|
||||
}
|
||||
|
||||
struct ModelAlias: Decodable {
|
||||
let model: String
|
||||
}
|
||||
|
||||
struct QuotaInfo: Decodable {
|
||||
let remainingFraction: Double?
|
||||
let resetTime: String?
|
||||
}
|
||||
|
||||
enum CodeValue: Decodable {
|
||||
case int(Int)
|
||||
case string(String)
|
||||
|
||||
var isOK: Bool {
|
||||
switch self {
|
||||
case let .int(value):
|
||||
value == 0
|
||||
case let .string(value):
|
||||
value.lowercased() == "ok" || value.lowercased() == "success" || value == "0"
|
||||
}
|
||||
}
|
||||
|
||||
var rawValue: String {
|
||||
switch self {
|
||||
case let .int(value): "\(value)"
|
||||
case let .string(value): value
|
||||
}
|
||||
}
|
||||
|
||||
init(from decoder: Decoder) throws {
|
||||
let container = try decoder.singleValueContainer()
|
||||
if let value = try? container.decode(Int.self) {
|
||||
self = .int(value)
|
||||
return
|
||||
}
|
||||
if let value = try? container.decode(String.self) {
|
||||
self = .string(value)
|
||||
return
|
||||
}
|
||||
throw DecodingError.dataCorruptedError(in: container, debugDescription: "Unsupported code type")
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,30 @@
|
||||
import Foundation
|
||||
|
||||
public enum AntigravityUsageDataSource: String, CaseIterable, Identifiable, Sendable {
|
||||
case auto
|
||||
case oauth
|
||||
case cli
|
||||
|
||||
public var id: String {
|
||||
self.rawValue
|
||||
}
|
||||
|
||||
public var displayName: String {
|
||||
switch self {
|
||||
case .auto: "Auto"
|
||||
case .oauth: "Google OAuth"
|
||||
case .cli: "Local API / agy CLI"
|
||||
}
|
||||
}
|
||||
|
||||
public var sourceLabel: String {
|
||||
switch self {
|
||||
case .auto:
|
||||
"auto"
|
||||
case .oauth:
|
||||
"oauth"
|
||||
case .cli:
|
||||
"cli"
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,191 @@
|
||||
import Foundation
|
||||
|
||||
#if os(macOS)
|
||||
|
||||
/// Fetches Augment usage via `auggie account status` CLI command
|
||||
public struct AuggieCLIProbe: Sendable {
|
||||
private static let log = CodexBarLog.logger(LogCategories.auggieCLI)
|
||||
|
||||
public init() {}
|
||||
|
||||
public func fetch() async throws -> AugmentStatusSnapshot {
|
||||
let output = try await self.runAuggieAccountStatus()
|
||||
return try self.parse(output)
|
||||
}
|
||||
|
||||
/// Timeout for the `auggie account status` command.
|
||||
private static let commandTimeout: TimeInterval = 15
|
||||
|
||||
private func runAuggieAccountStatus() async throws -> String {
|
||||
let env = ProcessInfo.processInfo.environment
|
||||
let loginPATH = LoginShellPathCache.shared.current
|
||||
let executable = BinaryLocator.resolveAuggieBinary(env: env, loginPATH: loginPATH) ?? "auggie"
|
||||
|
||||
var pathEnv = env
|
||||
pathEnv["PATH"] = PathBuilder.effectivePATH(
|
||||
purposes: [.tty, .nodeTooling],
|
||||
env: env,
|
||||
loginPATH: loginPATH)
|
||||
|
||||
let result = try await SubprocessRunner.run(
|
||||
binary: executable,
|
||||
arguments: ["account", "status"],
|
||||
environment: pathEnv,
|
||||
timeout: Self.commandTimeout,
|
||||
label: "auggie-account-status")
|
||||
|
||||
let output = result.stdout
|
||||
let errorOutput = result.stderr
|
||||
|
||||
guard !output.isEmpty else {
|
||||
if !errorOutput.isEmpty {
|
||||
Self.log.error("Auggie stderr: \(errorOutput)")
|
||||
}
|
||||
throw AuggieCLIError.noOutput
|
||||
}
|
||||
|
||||
// Check for auth errors
|
||||
if output.contains("Authentication failed") || output.contains("auggie login") {
|
||||
throw AuggieCLIError.notAuthenticated
|
||||
}
|
||||
|
||||
return output
|
||||
}
|
||||
|
||||
func parse(_ output: String) throws -> AugmentStatusSnapshot {
|
||||
// Legacy output:
|
||||
// Max Plan 450,000 credits / month
|
||||
// 11,657 remaining · 953,170 / 964,827 credits used
|
||||
// 2 days remaining in this billing cycle (ends 1/8/2026)
|
||||
//
|
||||
// Current output (2026+):
|
||||
// 319,054 credits remaining Max Plan
|
||||
// 450,000 credits / month
|
||||
// 9 days remaining in this billing cycle (ends 6/9/2026)
|
||||
|
||||
var maxCredits: Int?
|
||||
var remaining: Int?
|
||||
var used: Int?
|
||||
var total: Int?
|
||||
var billingCycleEnd: Date?
|
||||
|
||||
for line in output.split(separator: "\n") {
|
||||
let trimmed = line.trimmingCharacters(in: .whitespaces)
|
||||
|
||||
if trimmed.contains("credits / month") {
|
||||
if let match = trimmed.range(of: #"([\d,]+)\s+credits\s*/\s*month"#, options: .regularExpression) {
|
||||
let numberStr = String(trimmed[match]).replacingOccurrences(of: ",", with: "")
|
||||
.replacingOccurrences(of: " credits", with: "")
|
||||
.replacingOccurrences(of: " / month", with: "")
|
||||
maxCredits = Int(numberStr)
|
||||
total = total ?? Int(numberStr)
|
||||
}
|
||||
} else if trimmed.contains("Max Plan"), trimmed.contains("credits"), !trimmed.contains("remaining") {
|
||||
if let match = trimmed.range(of: #"([\d,]+)\s+credits"#, options: .regularExpression) {
|
||||
let numberStr = String(trimmed[match]).replacingOccurrences(of: ",", with: "")
|
||||
.replacingOccurrences(of: " credits", with: "")
|
||||
maxCredits = Int(numberStr)
|
||||
}
|
||||
}
|
||||
|
||||
if trimmed.contains("credits remaining"), !trimmed.contains("billing cycle") {
|
||||
if let match = trimmed.range(of: #"([\d,]+)\s+credits\s+remaining"#, options: .regularExpression) {
|
||||
let numberStr = String(trimmed[match]).replacingOccurrences(of: ",", with: "")
|
||||
.replacingOccurrences(of: " credits", with: "")
|
||||
.replacingOccurrences(of: " remaining", with: "")
|
||||
remaining = Int(numberStr)
|
||||
}
|
||||
}
|
||||
|
||||
// Parse "11,657 remaining · 953,170 / 964,827 credits used"
|
||||
if trimmed.contains("remaining"), trimmed.contains("credits used") {
|
||||
if let remMatch = trimmed.range(of: #"([\d,]+)\s+remaining"#, options: .regularExpression) {
|
||||
let numStr = String(trimmed[remMatch])
|
||||
.replacingOccurrences(of: ",", with: "")
|
||||
.replacingOccurrences(of: " remaining", with: "")
|
||||
remaining = Int(numStr)
|
||||
}
|
||||
|
||||
if let usedMatch = trimmed.range(
|
||||
of: #"([\d,]+)\s*/\s*([\d,]+)\s+credits used"#,
|
||||
options: .regularExpression)
|
||||
{
|
||||
let parts = String(trimmed[usedMatch])
|
||||
.replacingOccurrences(of: " credits used", with: "")
|
||||
.split(separator: "/")
|
||||
if parts.count == 2 {
|
||||
used = Int(parts[0].replacingOccurrences(of: ",", with: "")
|
||||
.trimmingCharacters(in: .whitespaces))
|
||||
total = Int(parts[1].replacingOccurrences(of: ",", with: "")
|
||||
.trimmingCharacters(in: .whitespaces))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if trimmed.contains("billing cycle"), trimmed.contains("ends") {
|
||||
if let dateMatch = trimmed.range(of: #"ends\s+([\d/]+)"#, options: .regularExpression) {
|
||||
let dateStr = String(trimmed[dateMatch])
|
||||
.replacingOccurrences(of: "ends", with: "")
|
||||
.trimmingCharacters(in: .whitespaces)
|
||||
|
||||
let formatter = DateFormatter()
|
||||
formatter.dateFormat = "M/d/yyyy"
|
||||
formatter.locale = Locale(identifier: "en_US_POSIX")
|
||||
formatter.timeZone = TimeZone.current
|
||||
billingCycleEnd = formatter.date(from: dateStr)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
guard let finalRemaining = remaining else {
|
||||
Self.log.error("Failed to parse auggie output: \(output)")
|
||||
throw AuggieCLIError.parseError("Could not extract credits from output")
|
||||
}
|
||||
|
||||
let finalTotal = total ?? maxCredits
|
||||
guard let finalTotal else {
|
||||
Self.log.error("Failed to parse auggie output: \(output)")
|
||||
throw AuggieCLIError.parseError("Could not extract credits from output")
|
||||
}
|
||||
|
||||
let finalUsed = used ?? max(0, finalTotal - finalRemaining)
|
||||
|
||||
return AugmentStatusSnapshot(
|
||||
creditsRemaining: Double(finalRemaining),
|
||||
creditsUsed: Double(finalUsed),
|
||||
creditsLimit: Double(finalTotal),
|
||||
billingCycleEnd: billingCycleEnd,
|
||||
accountEmail: nil,
|
||||
accountPlan: maxCredits.map { "\($0.formatted()) credits/month" },
|
||||
rawJSON: nil)
|
||||
}
|
||||
}
|
||||
|
||||
#else
|
||||
|
||||
public struct AuggieCLIProbe: Sendable {
|
||||
public init() {}
|
||||
|
||||
public func fetch() async throws -> AugmentStatusSnapshot {
|
||||
throw AugmentStatusProbeError.notSupported
|
||||
}
|
||||
}
|
||||
|
||||
#endif
|
||||
|
||||
public enum AuggieCLIError: LocalizedError {
|
||||
case noOutput
|
||||
case notAuthenticated
|
||||
case parseError(String)
|
||||
|
||||
public var errorDescription: String? {
|
||||
switch self {
|
||||
case .noOutput:
|
||||
"Auggie CLI returned no output"
|
||||
case .notAuthenticated:
|
||||
"Not authenticated. Run 'auggie login' to authenticate."
|
||||
case let .parseError(msg):
|
||||
"Failed to parse auggie output: \(msg)"
|
||||
}
|
||||
}
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user