125 lines · 4.8 KB · Swift
RawDownload
import Foundation public struct CodexStatusSnapshot: Sendable {    public let credits: Double?    public let fiveHourPercentLeft: Int?    public let weeklyPercentLeft: Int?    public let fiveHourResetDescription: String?    public let weeklyResetDescription: String?    public let rawText: String     public init(        credits: Double?,        fiveHourPercentLeft: Int?,        weeklyPercentLeft: Int?,        fiveHourResetDescription: String?,        weeklyResetDescription: String?,        rawText: String)    {        self.credits = credits        self.fiveHourPercentLeft = fiveHourPercentLeft        self.weeklyPercentLeft = weeklyPercentLeft        self.fiveHourResetDescription = fiveHourResetDescription        self.weeklyResetDescription = weeklyResetDescription        self.rawText = rawText    }} public enum CodexStatusProbeError: LocalizedError, Sendable {    case codexNotInstalled    case parseFailed(String)    case timedOut    case updateRequired(String)     public var errorDescription: String? {        switch self {        case .codexNotInstalled:            "Codex CLI missing. Install via `npm i -g @openai/codex` (or bun install) and restart."        case .parseFailed:            "Could not parse Codex status; will retry shortly."        case .timedOut:            "Codex status probe timed out."        case let .updateRequired(msg):            "Codex CLI update needed: \(msg)"        }    }} /// Runs `codex` inside a PTY, sends `/status`, captures text, and parses credits/limits.public struct CodexStatusProbe {    public var codexBinary: String = "codex"    public var timeout: TimeInterval = 18.0     public init() {}     public init(codexBinary: String = "codex", timeout: TimeInterval = 18.0) {        self.codexBinary = codexBinary        self.timeout = timeout    }     public func fetch() async throws -> CodexStatusSnapshot {        guard TTYCommandRunner.which(self.codexBinary) != nil else { throw CodexStatusProbeError.codexNotInstalled }        do {            return try self.runAndParse(rows: 60, cols: 200, timeout: self.timeout)        } catch let error as CodexStatusProbeError {            // Codex sometimes returns an incomplete screen on the first try; retry once with a longer window.            switch error {            case .parseFailed, .timedOut:                return try self.runAndParse(rows: 70, cols: 220, timeout: max(self.timeout, 24.0))            default:                throw error            }        }    }     // MARK: - Parsing     public static func parse(text: String) throws -> CodexStatusSnapshot {        let clean = TextParsing.stripANSICodes(text)        guard !clean.isEmpty else { throw CodexStatusProbeError.timedOut }        if clean.localizedCaseInsensitiveContains("data not available yet") {            throw CodexStatusProbeError.parseFailed("data not available yet")        }        if self.containsUpdatePrompt(clean) {            throw CodexStatusProbeError.updateRequired(                "Run `bun install -g @openai/codex` to continue (update prompt blocking /status).")        }        let credits = TextParsing.firstNumber(pattern: #"Credits:\s*([0-9][0-9.,]*)"#, text: clean)        // Pull reset info from the same lines that contain the percentages.        let fiveLine = TextParsing.firstLine(matching: #"5h limit[^\n]*"#, text: clean)        let weekLine = TextParsing.firstLine(matching: #"Weekly limit[^\n]*"#, text: clean)        let fivePct = fiveLine.flatMap(TextParsing.percentLeft(fromLine:))        let weekPct = weekLine.flatMap(TextParsing.percentLeft(fromLine:))        let fiveReset = fiveLine.flatMap(TextParsing.resetString(fromLine:))        let weekReset = weekLine.flatMap(TextParsing.resetString(fromLine:))        if credits == nil, fivePct == nil, weekPct == nil {            throw CodexStatusProbeError.parseFailed(clean.prefix(400).description)        }        return CodexStatusSnapshot(            credits: credits,            fiveHourPercentLeft: fivePct,            weeklyPercentLeft: weekPct,            fiveHourResetDescription: fiveReset,            weeklyResetDescription: weekReset,            rawText: clean)    }     private func runAndParse(rows: UInt16, cols: UInt16, timeout: TimeInterval) throws -> CodexStatusSnapshot {        let runner = TTYCommandRunner()        let script = "/status\n"        let result = try runner.run(            binary: self.codexBinary,            send: script,            options: .init(                rows: rows,                cols: cols,                timeout: timeout,                extraArgs: ["-s", "read-only", "-a", "untrusted"]))        return try Self.parse(text: result.text)    }     private static func containsUpdatePrompt(_ text: String) -> Bool {        let lower = text.lowercased()        return lower.contains("update available") && lower.contains("codex")    }}