macos-ai-menubar-usage/AIMenubarUsage/Providers/CursorProvider.swift

import Foundation

enum CursorAPIError: Error, LocalizedError {
    case notAuthenticated
    case invalidResponse
    case httpError(statusCode: Int)

    var errorDescription: String? {
        switch self {
        case .notAuthenticated:
            "Not authenticated. Please log in to Cursor."
        case .invalidResponse:
            "Could not parse usage data from Cursor."
        case .httpError(let statusCode):
            "Cursor API returned HTTP \(statusCode)."
        }
    }
}

struct CursorUsageBreakdown: Decodable, Sendable {
    let included: Int
    let bonus: Int
    let total: Int
}

struct CursorPlanUsage: Decodable, Sendable {
    let enabled: Bool
    let used: Int
    let limit: Int
    let remaining: Int
    let breakdown: CursorUsageBreakdown
    let autoPercentUsed: Double
    let apiPercentUsed: Double
    let totalPercentUsed: Double
}

struct CursorOnDemandUsage: Decodable, Sendable {
    let enabled: Bool
    let used: Int
    let limit: Int?
    let remaining: Int?
}

struct CursorIndividualUsage: Decodable, Sendable {
    let plan: CursorPlanUsage
    let onDemand: CursorOnDemandUsage
}

struct CursorUsageSummary: Decodable, Sendable {
    let billingCycleStart: String
    let billingCycleEnd: String
    let membershipType: String
    let limitType: String
    let isUnlimited: Bool
    let individualUsage: CursorIndividualUsage
}

struct CursorTeamMember: Decodable, Sendable {
    let email: String?
    let name: String?
}

struct CursorTeamMembersResponse: Decodable, Sendable {
    let teamMembers: [CursorTeamMember]?
}

struct CursorSpendEntry: Decodable, Sendable {
    let spendCents: Int?
    let fastPremiumRequests: Int?
}

struct CursorSpendResponse: Decodable, Sendable {
    let teamMemberSpend: [CursorSpendEntry]?
}

enum CursorAPI {
    private static let usageSummaryURL = URL(string: "https://cursor.com/api/usage-summary")!
    private static let spendURL = URL(string: "https://api.cursor.com/teams/spend")!
    private static let membersURL = URL(string: "https://api.cursor.com/teams/members")!

    static func fetchUsageSummary(cookieValue: String) async throws -> CursorUsageSummary {
        var request = URLRequest(url: usageSummaryURL)
        request.httpMethod = "GET"
        request.setValue("WorkosCursorSessionToken=\(cookieValue)", forHTTPHeaderField: "Cookie")

        let (data, response) = try await URLSession.shared.data(for: request)
        try validateHTTPResponse(response)

        do {
            return try JSONDecoder().decode(CursorUsageSummary.self, from: data)
        } catch {
            throw CursorAPIError.invalidResponse
        }
    }

    static func fetchTeamSpend(apiKey: String) async throws -> CursorSpendResponse {
        var request = URLRequest(url: spendURL)
        request.httpMethod = "POST"
        request.setValue("application/json", forHTTPHeaderField: "Content-Type")
        request.httpBody = Data("{}".utf8)
        request.setValue(
            "Basic \(Data("\(apiKey):".utf8).base64EncodedString())",
            forHTTPHeaderField: "Authorization"
        )

        let (data, response) = try await URLSession.shared.data(for: request)
        try validateAdminHTTPResponse(response)

        do {
            return try JSONDecoder().decode(CursorSpendResponse.self, from: data)
        } catch {
            throw CursorAPIError.invalidResponse
        }
    }

    static func fetchTeamMembers(apiKey: String) async throws -> CursorTeamMembersResponse {
        var request = URLRequest(url: membersURL)
        request.httpMethod = "GET"
        request.setValue(
            "Basic \(Data("\(apiKey):".utf8).base64EncodedString())",
            forHTTPHeaderField: "Authorization"
        )

        let (data, response) = try await URLSession.shared.data(for: request)
        try validateAdminHTTPResponse(response)

        do {
            return try JSONDecoder().decode(CursorTeamMembersResponse.self, from: data)
        } catch {
            throw CursorAPIError.invalidResponse
        }
    }

    private static func validateHTTPResponse(_ response: URLResponse) throws {
        guard let httpResponse = response as? HTTPURLResponse else {
            throw CursorAPIError.invalidResponse
        }
        switch httpResponse.statusCode {
        case 200:
            return
        case 401:
            throw CursorAPIError.notAuthenticated
        default:
            throw CursorAPIError.httpError(statusCode: httpResponse.statusCode)
        }
    }

    private static func validateAdminHTTPResponse(_ response: URLResponse) throws {
        guard let httpResponse = response as? HTTPURLResponse else {
            throw CursorAPIError.invalidResponse
        }
        switch httpResponse.statusCode {
        case 200:
            return
        case 401, 403:
            throw CursorAPIError.notAuthenticated
        default:
            throw CursorAPIError.httpError(statusCode: httpResponse.statusCode)
        }
    }

    static func snapshot(from summary: CursorUsageSummary, fetchedAt: Date, isStale: Bool) -> ProviderUsageSnapshot {
        let isoFormatter = ISO8601DateFormatter()
        isoFormatter.formatOptions = [.withInternetDateTime, .withFractionalSeconds]

        var values: [UsageMetric: MetricValue] = [
            .cursorTotalPercentUsed: .percent(summary.individualUsage.plan.totalPercentUsed),
            .cursorApiPercentUsed: .percent(summary.individualUsage.plan.apiPercentUsed),
            .cursorAutoPercentUsed: .percent(summary.individualUsage.plan.autoPercentUsed),
            .cursorPlanUsed: .integer(summary.individualUsage.plan.used),
            .cursorPlanLimit: .integer(summary.individualUsage.plan.limit),
            .cursorPlanRemaining: .integer(summary.individualUsage.plan.remaining),
            .cursorOnDemandUsed: .currency(Double(summary.individualUsage.onDemand.used) / 100.0),
            .cursorMembershipType: .text(summary.membershipType),
        ]

        if let start = isoFormatter.date(from: summary.billingCycleStart)
            ?? ISO8601DateFormatter().date(from: summary.billingCycleStart) {
            values[.cursorBillingCycleStart] = .date(start)
        }
        if let end = isoFormatter.date(from: summary.billingCycleEnd)
            ?? ISO8601DateFormatter().date(from: summary.billingCycleEnd) {
            values[.cursorBillingCycleEnd] = .date(end)
        }

        return ProviderUsageSnapshot(
            providerID: .cursor,
            fetchedAt: fetchedAt,
            values: values,
            isStale: isStale
        )
    }

    static func adminSnapshot(
        spend: CursorSpendResponse,
        members: CursorTeamMembersResponse,
        fetchedAt: Date,
        isStale: Bool
    ) -> ProviderUsageSnapshot {
        let totalSpend = spend.teamMemberSpend?.reduce(0) { $0 + ($1.spendCents ?? 0) } ?? 0
        let memberCount = members.teamMembers?.count ?? 0

        return ProviderUsageSnapshot(
            providerID: .cursor,
            fetchedAt: fetchedAt,
            values: [
                .cursorAdminSpendCents: .currency(Double(totalSpend) / 100.0),
                .cursorAdminMemberCount: .integer(memberCount),
            ],
            isStale: isStale
        )
    }
}

struct CursorProvider: UsageProvider {
    let id: ProviderID = .cursor
    let displayName = "Cursor"

    func resolveCredentials(authMode: AuthMode) throws -> AuthCredentials {
        switch authMode {
        case .auto:
            let session = try CursorAuthService.loadSessionCredentials()
            return .cursorSession(cookieValue: session.cookieValue)
        case .manualSession:
            let account = KeychainService.keychainAccount(provider: .cursor, authMode: .manualSession)
            guard let token = KeychainService.load(account: account), !token.isEmpty else {
                throw CursorAPIError.notAuthenticated
            }
            return .cursorSession(cookieValue: CursorAuthService.normalizeSessionCookie(token))
        case .adminAPI:
            let account = KeychainService.keychainAccount(provider: .cursor, authMode: .adminAPI)
            guard let apiKey = KeychainService.load(account: account), !apiKey.isEmpty else {
                throw CursorAPIError.notAuthenticated
            }
            return .cursorAdminAPI(apiKey: apiKey)
        }
    }

    func fetchUsage(credentials: AuthCredentials, cache: UsageCache) async throws -> ProviderUsageSnapshot {
        switch credentials {
        case .cursorSession(let cookieValue):
            return try await fetchSessionUsage(cookieValue: cookieValue, cache: cache)
        case .cursorAdminAPI(let apiKey):
            return try await fetchAdminUsage(apiKey: apiKey, cache: cache)
        default:
            throw CursorAPIError.invalidResponse
        }
    }

    private func fetchSessionUsage(cookieValue: String, cache: UsageCache) async throws -> ProviderUsageSnapshot {
        if let cached = await cache.load(provider: .cursor, maxAge: 60),
           let summary = try? JSONDecoder().decode(CursorUsageSummary.self, from: cached.payload) {
            return CursorAPI.snapshot(from: summary, fetchedAt: cached.fetchedAt, isStale: false)
        }

        do {
            let summary = try await CursorAPI.fetchUsageSummary(cookieValue: cookieValue)
            if let data = try? JSONEncoder().encode(summary) {
                await cache.save(provider: .cursor, payload: data)
            }
            return CursorAPI.snapshot(from: summary, fetchedAt: Date(), isStale: false)
        } catch {
            if let stale = await cache.loadStale(provider: .cursor),
               let summary = try? JSONDecoder().decode(CursorUsageSummary.self, from: stale.payload) {
                return CursorAPI.snapshot(from: summary, fetchedAt: stale.fetchedAt, isStale: true)
            }
            throw error
        }
    }

    private func fetchAdminUsage(apiKey: String, cache: UsageCache) async throws -> ProviderUsageSnapshot {
        struct AdminPayload: Codable {
            let spend: CursorSpendResponse
            let members: CursorTeamMembersResponse
        }

        if let cached = await cache.load(provider: .cursor, maxAge: 300),
           let payload = try? JSONDecoder().decode(AdminPayload.self, from: cached.payload) {
            return CursorAPI.adminSnapshot(
                spend: payload.spend,
                members: payload.members,
                fetchedAt: cached.fetchedAt,
                isStale: false
            )
        }

        do {
            async let spend = CursorAPI.fetchTeamSpend(apiKey: apiKey)
            async let members = CursorAPI.fetchTeamMembers(apiKey: apiKey)
            let result = try await (spend: spend, members: members)

            let payload = AdminPayload(spend: result.spend, members: result.members)
            if let data = try? JSONEncoder().encode(payload) {
                await cache.save(provider: .cursor, payload: data)
            }

            return CursorAPI.adminSnapshot(
                spend: result.spend,
                members: result.members,
                fetchedAt: Date(),
                isStale: false
            )
        } catch {
            if let stale = await cache.loadStale(provider: .cursor),
               let payload = try? JSONDecoder().decode(AdminPayload.self, from: stale.payload) {
                return CursorAPI.adminSnapshot(
                    spend: payload.spend,
                    members: payload.members,
                    fetchedAt: stale.fetchedAt,
                    isStale: true
                )
            }
            throw error
        }
    }
}

extension CursorUsageSummary: Encodable {}
extension CursorSpendResponse: Encodable {}
extension CursorTeamMembersResponse: Encodable {}
extension CursorUsageBreakdown: Encodable {}
extension CursorPlanUsage: Encodable {}
extension CursorOnDemandUsage: Encodable {}
extension CursorIndividualUsage: Encodable {}
extension CursorTeamMember: Encodable {}
extension CursorSpendEntry: Encodable {}