@sriinnu/runic 2.0.0 → 2.11.1

package/Sources/RunicCore/UsageFetcher.swift

@@ -1,694 +0,0 @@
-import Foundation
-#if canImport(Darwin)
-import Darwin
-#else
-import Glibc
-#endif
-
-public struct RateWindow: Codable, Equatable, Sendable {
-    public let usedPercent: Double
-    public let windowMinutes: Int?
-    public let resetsAt: Date?
-    /// Optional textual reset description (used by Claude CLI UI scrape).
-    public let resetDescription: String?
-    /// Optional label for this quota window (for example, a model name).
-    public let label: String?
-
-    public init(
-        usedPercent: Double,
-        windowMinutes: Int?,
-        resetsAt: Date?,
-        resetDescription: String?,
-        label: String? = nil)
-    {
-        self.usedPercent = usedPercent
-        self.windowMinutes = windowMinutes
-        self.resetsAt = resetsAt
-        self.resetDescription = resetDescription
-        self.label = label
-    }
-
-    public var remainingPercent: Double {
-        max(0, 100 - self.usedPercent)
-    }
-}
-
-public struct ProviderIdentitySnapshot: Codable, Sendable {
-    public let providerID: UsageProvider?
-    public let accountEmail: String?
-    public let accountOrganization: String?
-    public let loginMethod: String?
-
-    public init(
-        providerID: UsageProvider?,
-        accountEmail: String?,
-        accountOrganization: String?,
-        loginMethod: String?)
-    {
-        self.providerID = providerID
-        self.accountEmail = accountEmail
-        self.accountOrganization = accountOrganization
-        self.loginMethod = loginMethod
-    }
-
-    public func scoped(to provider: UsageProvider) -> ProviderIdentitySnapshot {
-        if self.providerID == provider { return self }
-        return ProviderIdentitySnapshot(
-            providerID: provider,
-            accountEmail: self.accountEmail,
-            accountOrganization: self.accountOrganization,
-            loginMethod: self.loginMethod)
-    }
-}
-
-public struct UsageSnapshot: Codable, Sendable {
-    public let primary: RateWindow
-    public let secondary: RateWindow?
-    public let tertiary: RateWindow?
-    public let providerCost: ProviderCostSnapshot?
-    public let zaiUsage: ZaiUsageSnapshot?
-    public let updatedAt: Date
-    public let identity: ProviderIdentitySnapshot?
-
-    private enum CodingKeys: String, CodingKey {
-        case primary
-        case secondary
-        case tertiary
-        case providerCost
-        case updatedAt
-        case identity
-        case accountEmail
-        case accountOrganization
-        case loginMethod
-    }
-
-    public init(
-        primary: RateWindow,
-        secondary: RateWindow?,
-        tertiary: RateWindow? = nil,
-        providerCost: ProviderCostSnapshot? = nil,
-        zaiUsage: ZaiUsageSnapshot? = nil,
-        updatedAt: Date,
-        identity: ProviderIdentitySnapshot? = nil)
-    {
-        self.primary = primary
-        self.secondary = secondary
-        self.tertiary = tertiary
-        self.providerCost = providerCost
-        self.zaiUsage = zaiUsage
-        self.updatedAt = updatedAt
-        self.identity = identity
-    }
-
-    public init(from decoder: Decoder) throws {
-        let container = try decoder.container(keyedBy: CodingKeys.self)
-        self.primary = try container.decode(RateWindow.self, forKey: .primary)
-        self.secondary = try container.decodeIfPresent(RateWindow.self, forKey: .secondary)
-        self.tertiary = try container.decodeIfPresent(RateWindow.self, forKey: .tertiary)
-        self.providerCost = try container.decodeIfPresent(ProviderCostSnapshot.self, forKey: .providerCost)
-        self.zaiUsage = nil // Not persisted, fetched fresh each time
-        self.updatedAt = try container.decode(Date.self, forKey: .updatedAt)
-        if let identity = try container.decodeIfPresent(ProviderIdentitySnapshot.self, forKey: .identity) {
-            self.identity = identity
-        } else {
-            let email = try container.decodeIfPresent(String.self, forKey: .accountEmail)
-            let organization = try container.decodeIfPresent(String.self, forKey: .accountOrganization)
-            let loginMethod = try container.decodeIfPresent(String.self, forKey: .loginMethod)
-            if email != nil || organization != nil || loginMethod != nil {
-                self.identity = ProviderIdentitySnapshot(
-                    providerID: nil,
-                    accountEmail: email,
-                    accountOrganization: organization,
-                    loginMethod: loginMethod)
-            } else {
-                self.identity = nil
-            }
-        }
-    }
-
-    public func encode(to encoder: Encoder) throws {
-        var container = encoder.container(keyedBy: CodingKeys.self)
-        try container.encode(self.primary, forKey: .primary)
-        if let secondary = self.secondary {
-            try container.encode(secondary, forKey: .secondary)
-        } else {
-            try container.encodeNil(forKey: .secondary)
-        }
-        try container.encodeIfPresent(self.tertiary, forKey: .tertiary)
-        try container.encodeIfPresent(self.providerCost, forKey: .providerCost)
-        try container.encode(self.updatedAt, forKey: .updatedAt)
-        try container.encodeIfPresent(self.identity, forKey: .identity)
-        try container.encodeIfPresent(self.identity?.accountEmail, forKey: .accountEmail)
-        try container.encodeIfPresent(self.identity?.accountOrganization, forKey: .accountOrganization)
-        try container.encodeIfPresent(self.identity?.loginMethod, forKey: .loginMethod)
-    }
-
-    public func identity(for provider: UsageProvider) -> ProviderIdentitySnapshot? {
-        guard let identity, identity.providerID == provider else { return nil }
-        return identity
-    }
-
-    public func accountEmail(for provider: UsageProvider) -> String? {
-        self.identity(for: provider)?.accountEmail
-    }
-
-    public func accountOrganization(for provider: UsageProvider) -> String? {
-        self.identity(for: provider)?.accountOrganization
-    }
-
-    public func loginMethod(for provider: UsageProvider) -> String? {
-        self.identity(for: provider)?.loginMethod
-    }
-
-    public func scoped(to provider: UsageProvider) -> UsageSnapshot {
-        guard let identity else { return self }
-        let scopedIdentity = identity.scoped(to: provider)
-        if scopedIdentity.providerID == identity.providerID { return self }
-        return UsageSnapshot(
-            primary: self.primary,
-            secondary: self.secondary,
-            tertiary: self.tertiary,
-            providerCost: self.providerCost,
-            zaiUsage: self.zaiUsage,
-            updatedAt: self.updatedAt,
-            identity: scopedIdentity)
-    }
-}
-
-public struct AccountInfo: Equatable, Sendable {
-    public let email: String?
-    public let plan: String?
-
-    public init(email: String?, plan: String?) {
-        self.email = email
-        self.plan = plan
-    }
-}
-
-public enum UsageError: LocalizedError, Sendable {
-    case noSessions
-    case noRateLimitsFound
-    case decodeFailed
-
-    public var errorDescription: String? {
-        switch self {
-        case .noSessions:
-            "No Codex sessions found yet. Run at least one Codex prompt first."
-        case .noRateLimitsFound:
-            "Found sessions, but no rate limit events yet."
-        case .decodeFailed:
-            "Could not parse Codex session log."
-        }
-    }
-}
-
-// MARK: - Codex RPC client (local process)
-
-private struct RPCAccountResponse: Decodable {
-    let account: RPCAccountDetails?
-    let requiresOpenaiAuth: Bool?
-}
-
-private enum RPCAccountDetails: Decodable {
-    case apiKey
-    case chatgpt(email: String, planType: String)
-
-    enum CodingKeys: String, CodingKey {
-        case type
-        case email
-        case planType
-    }
-
-    init(from decoder: Decoder) throws {
-        let container = try decoder.container(keyedBy: CodingKeys.self)
-        let type = try container.decode(String.self, forKey: .type)
-        switch type.lowercased() {
-        case "apikey":
-            self = .apiKey
-        case "chatgpt":
-            let email = try container.decodeIfPresent(String.self, forKey: .email) ?? "unknown"
-            let plan = try container.decodeIfPresent(String.self, forKey: .planType) ?? "unknown"
-            self = .chatgpt(email: email, planType: plan)
-        default:
-            throw DecodingError.dataCorruptedError(
-                forKey: .type,
-                in: container,
-                debugDescription: "Unknown account type \(type)")
-        }
-    }
-}
-
-private struct RPCRateLimitsResponse: Decodable, Encodable {
-    let rateLimits: RPCRateLimitSnapshot
-}
-
-private struct RPCRateLimitSnapshot: Decodable, Encodable {
-    let primary: RPCRateLimitWindow?
-    let secondary: RPCRateLimitWindow?
-    let credits: RPCCreditsSnapshot?
-}
-
-private struct RPCRateLimitWindow: Decodable, Encodable {
-    let usedPercent: Double
-    let windowDurationMins: Int?
-    let resetsAt: Int?
-}
-
-private struct RPCCreditsSnapshot: Decodable, Encodable {
-    let hasCredits: Bool
-    let unlimited: Bool
-    let balance: String?
-}
-
-private enum RPCWireError: Error, CustomStringConvertible {
-    case startFailed(String)
-    case requestFailed(String)
-    case malformed(String)
-
-    var description: String {
-        switch self {
-        case let .startFailed(message):
-            "Failed to start codex app-server: \(message)"
-        case let .requestFailed(message):
-            "RPC request failed: \(message)"
-        case let .malformed(message):
-            "Malformed response: \(message)"
-        }
-    }
-}
-
-/// RPC helper used on background tasks; safe because we confine it to the owning task.
-private final class CodexRPCClient: @unchecked Sendable {
-    private let process = Process()
-    private let stdinPipe = Pipe()
-    private let stdoutPipe = Pipe()
-    private let stderrPipe = Pipe()
-    private let stdoutLineStream: AsyncStream<Data>
-    private let stdoutLineContinuation: AsyncStream<Data>.Continuation
-    private var nextID = 1
-    private var processGroup: pid_t?
-
-    private final class LineBuffer: @unchecked Sendable {
-        private let lock = NSLock()
-        private var buffer = Data()
-
-        func appendAndDrainLines(_ data: Data) -> [Data] {
-            self.lock.lock()
-            defer { self.lock.unlock() }
-
-            self.buffer.append(data)
-            var out: [Data] = []
-            while let newline = self.buffer.firstIndex(of: 0x0A) {
-                let lineData = Data(self.buffer[..<newline])
-                self.buffer.removeSubrange(...newline)
-                if !lineData.isEmpty {
-                    out.append(lineData)
-                }
-            }
-            return out
-        }
-    }
-
-    private static func debugWriteStderr(_ message: String) {
-        #if !os(Linux)
-        fputs(message, stderr)
-        #endif
-    }
-
-    init(
-        executable: String = "codex",
-        arguments: [String] = ["-s", "read-only", "-a", "untrusted", "app-server"]) throws
-    {
-        var stdoutContinuation: AsyncStream<Data>.Continuation!
-        self.stdoutLineStream = AsyncStream<Data> { continuation in
-            stdoutContinuation = continuation
-        }
-        self.stdoutLineContinuation = stdoutContinuation
-
-        let resolvedExec = BinaryLocator.resolveCodexBinary()
-            ?? TTYCommandRunner.which(executable)
-
-        guard let resolvedExec else {
-            throw RPCWireError.startFailed(
-                "Codex CLI not found. Install with `npm i -g @openai/codex` (or bun) then relaunch Runic.")
-        }
-        var env = ProcessInfo.processInfo.environment
-        env["PATH"] = PathBuilder.effectivePATH(
-            purposes: [.rpc, .nodeTooling],
-            env: env)
-
-        self.process.environment = env
-        self.process.executableURL = URL(fileURLWithPath: "/usr/bin/env")
-        self.process.arguments = [resolvedExec] + arguments
-        self.process.standardInput = self.stdinPipe
-        self.process.standardOutput = self.stdoutPipe
-        self.process.standardError = self.stderrPipe
-
-        do {
-            try self.process.run()
-        } catch {
-            throw RPCWireError.startFailed(error.localizedDescription)
-        }
-
-        let pid = self.process.processIdentifier
-        if setpgid(pid, pid) == 0 {
-            self.processGroup = pid
-        }
-
-        let stdoutHandle = self.stdoutPipe.fileHandleForReading
-        let stdoutLineContinuation = self.stdoutLineContinuation
-        let stdoutBuffer = LineBuffer()
-        stdoutHandle.readabilityHandler = { handle in
-            let data = handle.availableData
-            if data.isEmpty {
-                handle.readabilityHandler = nil
-                stdoutLineContinuation.finish()
-                return
-            }
-
-            let lines = stdoutBuffer.appendAndDrainLines(data)
-
-            for lineData in lines {
-                stdoutLineContinuation.yield(lineData)
-            }
-        }
-
-        let stderrHandle = self.stderrPipe.fileHandleForReading
-        stderrHandle.readabilityHandler = { handle in
-            let data = handle.availableData
-            // When the child closes stderr, availableData returns empty and will keep re-firing; clear the handler
-            // to avoid a busy read loop on the file-descriptor monitoring queue.
-            if data.isEmpty {
-                handle.readabilityHandler = nil
-                return
-            }
-            guard let text = String(data: data, encoding: .utf8), !text.isEmpty else { return }
-            for line in text.split(whereSeparator: \.isNewline) {
-                Self.debugWriteStderr("[codex stderr] \(line)\n")
-            }
-        }
-    }
-
-    func initialize(clientName: String, clientVersion: String) async throws {
-        _ = try await self.request(
-            method: "initialize",
-            params: ["clientInfo": ["name": clientName, "version": clientVersion]])
-        try self.sendNotification(method: "initialized")
-    }
-
-    func fetchAccount() async throws -> RPCAccountResponse {
-        let message = try await self.request(method: "account/read")
-        return try self.decodeResult(from: message)
-    }
-
-    func fetchRateLimits() async throws -> RPCRateLimitsResponse {
-        let message = try await self.request(method: "account/rateLimits/read")
-        return try self.decodeResult(from: message)
-    }
-
-    func shutdown() {
-        self.stdoutPipe.fileHandleForReading.readabilityHandler = nil
-        self.stderrPipe.fileHandleForReading.readabilityHandler = nil
-        if self.process.isRunning {
-            self.process.terminate()
-        }
-        if let pgid = self.processGroup {
-            kill(-pgid, SIGTERM)
-        }
-        let waitDeadline = Date().addingTimeInterval(2.0)
-        while self.process.isRunning, Date() < waitDeadline {
-            usleep(100_000)
-        }
-        if self.process.isRunning {
-            if let pgid = self.processGroup {
-                kill(-pgid, SIGKILL)
-            }
-            kill(self.process.processIdentifier, SIGKILL)
-        }
-        try? self.stdinPipe.fileHandleForWriting.close()
-        try? self.stdoutPipe.fileHandleForReading.close()
-        try? self.stderrPipe.fileHandleForReading.close()
-    }
-
-    // MARK: - JSON-RPC helpers
-
-    private func request(method: String, params: [String: Any]? = nil) async throws -> [String: Any] {
-        let id = self.nextID
-        self.nextID += 1
-        try self.sendRequest(id: id, method: method, params: params)
-
-        while true {
-            let message = try await self.readNextMessage()
-
-            if message["id"] == nil, let methodName = message["method"] as? String {
-                Self.debugWriteStderr("[codex notify] \(methodName)\n")
-                continue
-            }
-
-            guard let messageID = self.jsonID(message["id"]), messageID == id else { continue }
-
-            if let error = message["error"] as? [String: Any], let messageText = error["message"] as? String {
-                throw RPCWireError.requestFailed(messageText)
-            }
-
-            return message
-        }
-    }
-
-    private func sendNotification(method: String, params: [String: Any]? = nil) throws {
-        let paramsValue: Any = params ?? [:]
-        try self.sendPayload(["method": method, "params": paramsValue])
-    }
-
-    private func sendRequest(id: Int, method: String, params: [String: Any]?) throws {
-        let paramsValue: Any = params ?? [:]
-        let payload: [String: Any] = ["id": id, "method": method, "params": paramsValue]
-        try self.sendPayload(payload)
-    }
-
-    private func sendPayload(_ payload: [String: Any]) throws {
-        let data = try JSONSerialization.data(withJSONObject: payload)
-        self.stdinPipe.fileHandleForWriting.write(data)
-        self.stdinPipe.fileHandleForWriting.write(Data([0x0A]))
-    }
-
-    private func readNextMessage() async throws -> [String: Any] {
-        for await lineData in self.stdoutLineStream {
-            if lineData.isEmpty { continue }
-            if let json = try? JSONSerialization.jsonObject(with: lineData) as? [String: Any] {
-                return json
-            }
-        }
-        throw RPCWireError.malformed("codex app-server closed stdout")
-    }
-
-    private func decodeResult<T: Decodable>(from message: [String: Any]) throws -> T {
-        guard let result = message["result"] else {
-            throw RPCWireError.malformed("missing result field")
-        }
-        let data = try JSONSerialization.data(withJSONObject: result)
-        let decoder = JSONDecoder()
-        return try decoder.decode(T.self, from: data)
-    }
-
-    private func jsonID(_ value: Any?) -> Int? {
-        switch value {
-        case let int as Int:
-            int
-        case let number as NSNumber:
-            number.intValue
-        default:
-            nil
-        }
-    }
-}
-
-// MARK: - Public fetcher used by the app
-
-public struct UsageFetcher: Sendable {
-    private let environment: [String: String]
-
-    public init(environment: [String: String] = ProcessInfo.processInfo.environment) {
-        self.environment = environment
-        LoginShellPathCache.shared.captureOnce()
-    }
-
-    public func loadLatestUsage() async throws -> UsageSnapshot {
-        try await self.withFallback(primary: self.loadRPCUsage, secondary: self.loadTTYUsage)
-    }
-
-    private func loadRPCUsage() async throws -> UsageSnapshot {
-        let rpc = try CodexRPCClient()
-        defer { rpc.shutdown() }
-
-        try await rpc.initialize(clientName: "runic", clientVersion: "0.5.4")
-        // The app-server answers on a single stdout stream, so keep requests
-        // serialized to avoid starving one reader when multiple awaiters race
-        // for the same pipe.
-        let limits = try await rpc.fetchRateLimits().rateLimits
-        let account = try? await rpc.fetchAccount()
-
-        guard let primary = Self.makeWindow(from: limits.primary),
-              let secondary = Self.makeWindow(from: limits.secondary)
-        else {
-            throw UsageError.noRateLimitsFound
-        }
-
-        let identity = ProviderIdentitySnapshot(
-            providerID: .codex,
-            accountEmail: account?.account.flatMap { details in
-                if case let .chatgpt(email, _) = details { email } else { nil }
-            },
-            accountOrganization: nil,
-            loginMethod: account?.account.flatMap { details in
-                if case let .chatgpt(_, plan) = details { plan } else { nil }
-            })
-        return UsageSnapshot(
-            primary: primary,
-            secondary: secondary,
-            tertiary: nil,
-            updatedAt: Date(),
-            identity: identity)
-    }
-
-    private func loadTTYUsage() async throws -> UsageSnapshot {
-        let status = try await CodexStatusProbe().fetch()
-        guard let fiveLeft = status.fiveHourPercentLeft, let weekLeft = status.weeklyPercentLeft else {
-            throw UsageError.noRateLimitsFound
-        }
-
-        let primary = RateWindow(
-            usedPercent: max(0, 100 - Double(fiveLeft)),
-            windowMinutes: 300,
-            resetsAt: nil,
-            resetDescription: status.fiveHourResetDescription)
-        let secondary = RateWindow(
-            usedPercent: max(0, 100 - Double(weekLeft)),
-            windowMinutes: 10080,
-            resetsAt: nil,
-            resetDescription: status.weeklyResetDescription)
-
-        return UsageSnapshot(
-            primary: primary,
-            secondary: secondary,
-            tertiary: nil,
-            updatedAt: Date(),
-            identity: nil)
-    }
-
-    public func loadLatestCredits() async throws -> CreditsSnapshot {
-        try await self.withFallback(primary: self.loadRPCCredits, secondary: self.loadTTYCredits)
-    }
-
-    private func loadRPCCredits() async throws -> CreditsSnapshot {
-        let rpc = try CodexRPCClient()
-        defer { rpc.shutdown() }
-        try await rpc.initialize(clientName: "runic", clientVersion: "0.5.4")
-        let limits = try await rpc.fetchRateLimits().rateLimits
-        guard let credits = limits.credits else { throw UsageError.noRateLimitsFound }
-        let remaining = Self.parseCredits(credits.balance)
-        return CreditsSnapshot(remaining: remaining, events: [], updatedAt: Date())
-    }
-
-    private func loadTTYCredits() async throws -> CreditsSnapshot {
-        let status = try await CodexStatusProbe().fetch()
-        guard let credits = status.credits else { throw UsageError.noRateLimitsFound }
-        return CreditsSnapshot(remaining: credits, events: [], updatedAt: Date())
-    }
-
-    private func withFallback<T>(
-        primary: @escaping () async throws -> T,
-        secondary: @escaping () async throws -> T) async throws -> T
-    {
-        do {
-            return try await primary()
-        } catch let primaryError {
-            do {
-                return try await secondary()
-            } catch {
-                // Preserve the original failure so callers see the primary path error.
-                throw primaryError
-            }
-        }
-    }
-
-    public func debugRawRateLimits() async -> String {
-        do {
-            let rpc = try CodexRPCClient()
-            defer { rpc.shutdown() }
-            try await rpc.initialize(clientName: "runic", clientVersion: "0.5.4")
-            let limits = try await rpc.fetchRateLimits()
-            let data = try JSONEncoder().encode(limits)
-            return String(data: data, encoding: .utf8) ?? "<unprintable>"
-        } catch {
-            return "Codex RPC probe failed: \(error)"
-        }
-    }
-
-    public func loadAccountInfo() -> AccountInfo {
-        // Keep using auth.json for quick startup (non-blocking, no RPC spin-up required).
-        let authURL = URL(fileURLWithPath: self.environment["CODEX_HOME"] ?? "\(NSHomeDirectory())/.codex")
-            .appendingPathComponent("auth.json")
-        guard let data = try? Data(contentsOf: authURL),
-              let auth = try? JSONDecoder().decode(AuthFile.self, from: data),
-              let idToken = auth.tokens?.idToken
-        else {
-            return AccountInfo(email: nil, plan: nil)
-        }
-
-        guard let payload = UsageFetcher.parseJWT(idToken) else {
-            return AccountInfo(email: nil, plan: nil)
-        }
-
-        let authDict = payload["https://api.openai.com/auth"] as? [String: Any]
-        let profileDict = payload["https://api.openai.com/profile"] as? [String: Any]
-
-        let plan = (authDict?["chatgpt_plan_type"] as? String)
-            ?? (payload["chatgpt_plan_type"] as? String)
-
-        let email = (payload["email"] as? String)
-            ?? (profileDict?["email"] as? String)
-
-        return AccountInfo(email: email, plan: plan)
-    }
-
-    // MARK: - Helpers
-
-    private static func makeWindow(from rpc: RPCRateLimitWindow?) -> RateWindow? {
-        guard let rpc else { return nil }
-        let resetsAtDate = rpc.resetsAt.map { Date(timeIntervalSince1970: TimeInterval($0)) }
-        let resetDescription = resetsAtDate.map { UsageFormatter.resetDescription(from: $0) }
-        return RateWindow(
-            usedPercent: rpc.usedPercent,
-            windowMinutes: rpc.windowDurationMins,
-            resetsAt: resetsAtDate,
-            resetDescription: resetDescription)
-    }
-
-    private static func parseCredits(_ balance: String?) -> Double {
-        guard let balance, let val = Double(balance) else { return 0 }
-        return val
-    }
-
-    public static func parseJWT(_ token: String) -> [String: Any]? {
-        let parts = token.split(separator: ".")
-        guard parts.count >= 2 else { return nil }
-        let payloadPart = parts[1]
-
-        var padded = String(payloadPart)
-            .replacingOccurrences(of: "-", with: "+")
-            .replacingOccurrences(of: "_", with: "/")
-        while padded.count % 4 != 0 {
-            padded.append("=")
-        }
-        guard let data = Data(base64Encoded: padded) else { return nil }
-        guard let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] else { return nil }
-        return json
-    }
-}
-
-/// Minimal auth.json struct preserved from previous implementation
-private struct AuthFile: Decodable {
-    struct Tokens: Decodable { let idToken: String? }
-    let tokens: Tokens?
-}