906 lines · 33.7 KB · Swift
RawDownload
import Foundation#if canImport(Darwin)import Darwin#elseif canImport(Glibc)import Glibc#elseif canImport(Musl)import Musl#endif public struct KiroUsageSnapshot: Sendable {    public let planName: String    public let displayPlanName: String    public let accountEmail: String?    public let authMethod: String?    public let creditsUsed: Double    public let creditsTotal: Double    public let creditsPercent: Double    public let bonusCreditsUsed: Double?    public let bonusCreditsTotal: Double?    public let bonusExpiryDays: Int?    public let overagesStatus: String?    public let overageCreditsUsed: Double?    public let estimatedOverageCostUSD: Double?    public let manageURL: String?    public let contextUsage: KiroContextUsageSnapshot?    public let resetsAt: Date?    public let updatedAt: Date     public init(        planName: String,        displayPlanName: String? = nil,        accountEmail: String? = nil,        authMethod: String? = nil,        creditsUsed: Double,        creditsTotal: Double,        creditsPercent: Double,        bonusCreditsUsed: Double?,        bonusCreditsTotal: Double?,        bonusExpiryDays: Int?,        overagesStatus: String? = nil,        overageCreditsUsed: Double? = nil,        estimatedOverageCostUSD: Double? = nil,        manageURL: String? = nil,        contextUsage: KiroContextUsageSnapshot? = nil,        resetsAt: Date?,        updatedAt: Date)    {        self.planName = planName        self.displayPlanName = displayPlanName ?? KiroStatusProbe.displayPlanName(planName)        self.accountEmail = accountEmail        self.authMethod = authMethod        self.creditsUsed = creditsUsed        self.creditsTotal = creditsTotal        self.creditsPercent = creditsPercent        self.bonusCreditsUsed = bonusCreditsUsed        self.bonusCreditsTotal = bonusCreditsTotal        self.bonusExpiryDays = bonusExpiryDays        self.overagesStatus = overagesStatus        self.overageCreditsUsed = overageCreditsUsed        self.estimatedOverageCostUSD = estimatedOverageCostUSD        self.manageURL = manageURL        self.contextUsage = contextUsage        self.resetsAt = resetsAt        self.updatedAt = updatedAt    }     public func toUsageSnapshot() -> UsageSnapshot {        let primary = RateWindow(            usedPercent: self.creditsPercent,            windowMinutes: nil,            resetsAt: self.resetsAt,            resetDescription: nil)         var secondary: RateWindow?        if let bonusUsed = self.bonusCreditsUsed,           let bonusTotal = self.bonusCreditsTotal,           bonusTotal > 0        {            let bonusPercent = (bonusUsed / bonusTotal) * 100.0            var expiryDate: Date?            if let days = self.bonusExpiryDays {                expiryDate = Calendar.current.date(byAdding: .day, value: days, to: Date())            }            secondary = RateWindow(                usedPercent: bonusPercent,                windowMinutes: nil,                resetsAt: expiryDate,                resetDescription: self.bonusExpiryDays.map { "expires in \($0)d" })        }         let identity = ProviderIdentitySnapshot(            providerID: .kiro,            accountEmail: self.accountEmail,            accountOrganization: nil,            loginMethod: self.authMethod)         let kiroUsage = KiroUsageDetails(            planName: self.planName,            displayPlanName: self.displayPlanName,            creditsUsed: self.creditsUsed,            creditsTotal: self.creditsTotal,            creditsRemaining: self.creditsRemaining,            bonusCreditsUsed: self.bonusCreditsUsed,            bonusCreditsTotal: self.bonusCreditsTotal,            bonusCreditsRemaining: self.bonusCreditsRemaining,            bonusExpiryDays: self.bonusExpiryDays,            overagesStatus: self.overagesStatus,            overageCreditsUsed: self.overageCreditsUsed,            estimatedOverageCostUSD: self.estimatedOverageCostUSD,            manageURL: self.manageURL,            contextUsage: self.contextUsage)         return UsageSnapshot(            primary: primary,            secondary: secondary,            tertiary: nil,            kiroUsage: kiroUsage,            providerCost: nil,            zaiUsage: nil,            updatedAt: self.updatedAt,            identity: identity)    }     public var creditsRemaining: Double {        max(0, self.creditsTotal - self.creditsUsed)    }     public var bonusCreditsRemaining: Double? {        guard let bonusCreditsUsed, let bonusCreditsTotal else { return nil }        return max(0, bonusCreditsTotal - bonusCreditsUsed)    }} public struct KiroContextUsageSnapshot: Codable, Equatable, Sendable {    public let totalPercentUsed: Double    public let contextFilesPercent: Double?    public let toolsPercent: Double?    public let kiroResponsesPercent: Double?    public let promptsPercent: Double?     public init(        totalPercentUsed: Double,        contextFilesPercent: Double?,        toolsPercent: Double?,        kiroResponsesPercent: Double?,        promptsPercent: Double?)    {        self.totalPercentUsed = totalPercentUsed        self.contextFilesPercent = contextFilesPercent        self.toolsPercent = toolsPercent        self.kiroResponsesPercent = kiroResponsesPercent        self.promptsPercent = promptsPercent    }} public struct KiroUsageDetails: Codable, Equatable, Sendable {    public let planName: String    public let displayPlanName: String    public let creditsUsed: Double    public let creditsTotal: Double    public let creditsRemaining: Double    public let bonusCreditsUsed: Double?    public let bonusCreditsTotal: Double?    public let bonusCreditsRemaining: Double?    public let bonusExpiryDays: Int?    public let overagesStatus: String?    public let overageCreditsUsed: Double?    public let estimatedOverageCostUSD: Double?    public let manageURL: String?    public let contextUsage: KiroContextUsageSnapshot?     public init(        planName: String,        displayPlanName: String,        creditsUsed: Double,        creditsTotal: Double,        creditsRemaining: Double,        bonusCreditsUsed: Double?,        bonusCreditsTotal: Double?,        bonusCreditsRemaining: Double?,        bonusExpiryDays: Int?,        overagesStatus: String?,        overageCreditsUsed: Double?,        estimatedOverageCostUSD: Double?,        manageURL: String?,        contextUsage: KiroContextUsageSnapshot?)    {        self.planName = planName        self.displayPlanName = displayPlanName        self.creditsUsed = creditsUsed        self.creditsTotal = creditsTotal        self.creditsRemaining = creditsRemaining        self.bonusCreditsUsed = bonusCreditsUsed        self.bonusCreditsTotal = bonusCreditsTotal        self.bonusCreditsRemaining = bonusCreditsRemaining        self.bonusExpiryDays = bonusExpiryDays        self.overagesStatus = overagesStatus        self.overageCreditsUsed = overageCreditsUsed        self.estimatedOverageCostUSD = estimatedOverageCostUSD        self.manageURL = manageURL        self.contextUsage = contextUsage    }} public enum KiroStatusProbeError: LocalizedError, Sendable {    case cliNotFound    case notLoggedIn    case cliFailed(String)    case parseError(String)    case timeout     public var errorDescription: String? {        switch self {        case .cliNotFound:            "kiro-cli not found. Install it from https://kiro.dev"        case .notLoggedIn:            "Not logged in to Kiro. Run 'kiro-cli login' first."        case let .cliFailed(message):            message        case let .parseError(msg):            "Failed to parse Kiro usage: \(msg)"        case .timeout:            "Kiro CLI timed out."        }    }} public struct KiroStatusProbe: Sendable {    private let cliBinaryResolver: @Sendable () -> String?    private let accountProbeTimeout: TimeInterval     public init() {        self.cliBinaryResolver = { TTYCommandRunner.which("kiro-cli") }        self.accountProbeTimeout = 3.0    }     init(cliBinaryResolver: @escaping @Sendable () -> String?, accountProbeTimeout: TimeInterval = 3.0) {        self.cliBinaryResolver = cliBinaryResolver        self.accountProbeTimeout = accountProbeTimeout    }     private static let logger = CodexBarLog.logger(LogCategories.kiro)     public static func detectVersion() -> String? {        guard let path = TTYCommandRunner.which("kiro-cli"),              let output = ProviderVersionDetector.run(                  path: path,                  args: ["--version"],                  mergeStandardError: true)        else {            self.logger.debug("kiro-cli version detection failed")            return nil        }        // Output is like "kiro-cli 1.23.1"        if output.hasPrefix("kiro-cli ") {            return String(output.dropFirst("kiro-cli ".count))        }        return output    }     public func fetch() async throws -> KiroUsageSnapshot {        let accountTask = Task { await self.fetchAccountStatus() }         let output: String        do {            output = try await self.runUsageCommand()        } catch is CancellationError {            accountTask.cancel()            _ = await accountTask.value            throw CancellationError()        } catch {            if try await self.awaitAccountStatus(accountTask) == .notLoggedIn {                throw KiroStatusProbeError.notLoggedIn            }            throw error        }         var contextUsage: KiroContextUsageSnapshot?        do {            contextUsage = try await self.fetchContextUsage()        } catch is CancellationError {            accountTask.cancel()            _ = await accountTask.value            throw CancellationError()        } catch {            Self.logger.debug("Kiro context usage probe failed: \(error.localizedDescription)")        }         let accountStatus = try await self.awaitAccountStatus(accountTask)        let accountInfo = accountStatus.account        do {            return try self.parse(                output: output,                accountEmail: accountInfo?.email,                authMethod: accountInfo?.authMethod,                contextUsage: contextUsage)        } catch KiroStatusProbeError.parseError where accountStatus == .notLoggedIn {            throw KiroStatusProbeError.notLoggedIn        }    }     struct KiroCLIResult {        let stdout: String        let stderr: String        let terminationStatus: Int32        let terminatedForIdle: Bool    }     struct KiroAccountInfo: Equatable {        let authMethod: String?        let email: String?    }     private enum KiroAccountProbeStatus: Equatable {        case account(KiroAccountInfo)        case notLoggedIn        case unavailable         var account: KiroAccountInfo? {            guard case let .account(info) = self else { return nil }            return info        }    }     private func fetchAccountStatus() async -> KiroAccountProbeStatus {        do {            return try await .account(self.ensureLoggedIn())        } catch KiroStatusProbeError.notLoggedIn {            return .notLoggedIn        } catch {            Self.logger.debug("Kiro account probe failed: \(error.localizedDescription)")            return .unavailable        }    }     private func awaitAccountStatus(        _ task: Task<KiroAccountProbeStatus, Never>) async throws -> KiroAccountProbeStatus    {        let status = await withTaskCancellationHandler {            await task.value        } onCancel: {            task.cancel()        }        try Task.checkCancellation()        return status    }     private func ensureLoggedIn() async throws -> KiroAccountInfo {        let result = try await self.runCommand(arguments: ["whoami"], timeout: self.accountProbeTimeout)        return try self.validateWhoAmIOutput(            stdout: result.stdout,            stderr: result.stderr,            terminationStatus: result.terminationStatus)    }     func validateWhoAmIOutput(stdout: String, stderr: String, terminationStatus: Int32) throws -> KiroAccountInfo {        let trimmedStdout = stdout.trimmingCharacters(in: .whitespacesAndNewlines)        let trimmedStderr = stderr.trimmingCharacters(in: .whitespacesAndNewlines)        let combined = trimmedStderr.isEmpty ? trimmedStdout : trimmedStderr        let lowered = combined.lowercased()         if lowered.contains("not logged in") || lowered.contains("login required") {            throw KiroStatusProbeError.notLoggedIn        }         if terminationStatus != 0 {            let message = combined.isEmpty                ? "Kiro CLI failed with status \(terminationStatus)."                : combined            throw KiroStatusProbeError.cliFailed(message)        }         if combined.isEmpty {            throw KiroStatusProbeError.cliFailed("Kiro CLI whoami returned no output.")        }         return self.parseWhoAmIOutput(combined)    }     func parseWhoAmIOutput(_ output: String) -> KiroAccountInfo {        let stripped = Self.stripANSI(output)        var authMethod: String?        var email: String?        for rawLine in stripped.components(separatedBy: .newlines) {            let line = rawLine.trimmingCharacters(in: .whitespacesAndNewlines)            guard !line.isEmpty else { continue }            if line.localizedCaseInsensitiveContains("logged in with") {                authMethod = line.replacingOccurrences(                    of: #"(?i)^\s*logged in with\s+"#,                    with: "",                    options: [.regularExpression])                    .trimmingCharacters(in: .whitespacesAndNewlines)            } else if line.localizedCaseInsensitiveContains("email:") {                email = line.replacingOccurrences(                    of: #"(?i)^\s*email:\s*"#,                    with: "",                    options: [.regularExpression])                    .trimmingCharacters(in: .whitespacesAndNewlines)            } else if email == nil,                      !line.contains(" "),                      line.contains("@")            {                email = line            }        }        return KiroAccountInfo(            authMethod: authMethod?.nilIfEmpty,            email: email?.nilIfEmpty)    }     private func runUsageCommand() async throws -> String {        let result = try await self.runCommand(            arguments: ["chat", "--no-interactive", "/usage"],            timeout: 20.0,            idleTimeout: 10.0)        let trimmedStdout = result.stdout.trimmingCharacters(in: .whitespacesAndNewlines)        let trimmedStderr = result.stderr.trimmingCharacters(in: .whitespacesAndNewlines)        let combinedOutput = trimmedStderr.isEmpty ? trimmedStdout : trimmedStderr        let combinedStripped = Self.stripANSI(combinedOutput).lowercased()         if combinedStripped.contains("not logged in")            || combinedStripped.contains("login required")            || combinedStripped.contains("failed to initialize auth portal")            || combinedStripped.contains("kiro-cli login")            || combinedStripped.contains("oauth error")        {            throw KiroStatusProbeError.notLoggedIn        }         if result.terminatedForIdle, !Self.isUsageOutputComplete(combinedOutput) {            throw KiroStatusProbeError.timeout        }         if !trimmedStdout.isEmpty {            return result.stdout        }         if !trimmedStderr.isEmpty {            return result.stderr        }         if result.terminationStatus != 0 {            let message = combinedOutput.isEmpty                ? "Kiro CLI failed with status \(result.terminationStatus)."                : combinedOutput            throw KiroStatusProbeError.cliFailed(message)        }         return result.stdout    }     private func fetchContextUsage() async throws -> KiroContextUsageSnapshot? {        let result = try await self.runCommand(            arguments: ["chat", "--no-interactive", "/context"],            timeout: 8.0,            idleTimeout: 3.0)        let output = result.stdout.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty            ? result.stderr            : result.stdout        return self.parseContextUsage(output: output)    }     func runCommand(        arguments: [String],        timeout: TimeInterval,        idleTimeout: TimeInterval = 5.0) async throws -> KiroCLIResult    {        guard let binary = self.cliBinaryResolver() else {            throw KiroStatusProbeError.cliNotFound        }         let stdoutPipe = Pipe()        let stderrPipe = Pipe()         var env = ProcessInfo.processInfo.environment        env["TERM"] = "xterm-256color"         final class ActivityState: @unchecked Sendable {            private let lock = NSLock()            private var _lastActivityAt = Date()            private var _hasReceivedOutput = false             var lastActivityAt: Date {                self.lock.withLock { self._lastActivityAt }            }             var hasReceivedOutput: Bool {                self.lock.withLock { self._hasReceivedOutput }            }             func markActivity() {                self.lock.withLock {                    self._lastActivityAt = Date()                    self._hasReceivedOutput = true                }            }        }         let state = ActivityState()        let stdoutCapture = ProcessPipeCapture(pipe: stdoutPipe, onData: { state.markActivity() })

Lines 1–500 of 906.