243 lines · 9.4 KB · Swift
RawDownload
import Foundation public struct ProviderTokenCostConfig: Sendable {    public let supportsTokenCost: Bool    public let noDataMessage: @Sendable () -> String     public init(supportsTokenCost: Bool, noDataMessage: @escaping @Sendable () -> String) {        self.supportsTokenCost = supportsTokenCost        self.noDataMessage = noDataMessage    }} public enum ProviderPaceWindowRule: Sendable {    case unsupported    case resetDatePresent    case windowDurationPresent    case windowDuration(minutes: Int)    case custom(@Sendable (_ window: RateWindow, _ now: Date) -> Bool)     public func matches(window: RateWindow, now: Date) -> Bool {        switch self {        case .unsupported:            false        case .resetDatePresent:            window.resetsAt != nil        case .windowDurationPresent:            window.windowMinutes != nil        case let .windowDuration(minutes):            window.windowMinutes == minutes        case let .custom(predicate):            predicate(window, now)        }    }} public enum ProviderPaceDurationRule: Sendable {    case unsupported    case windowDurationMissing    case windowDuration(minutes: Int)     public func matches(window: RateWindow) -> Bool {        switch self {        case .unsupported:            false        case .windowDurationMissing:            window.windowMinutes == nil        case let .windowDuration(minutes):            window.windowMinutes == minutes        }    }} public struct ProviderPaceCapability: Sendable {    public static let monthlyWindowSentinelMinutes = 30 * 24 * 60    public static let unsupported = ProviderPaceCapability()    public static let calendarMonthResetWindow = ProviderPaceCapability(        resetWindowPace: .windowDuration(minutes: ProviderPaceCapability.monthlyWindowSentinelMinutes),        inferredMonthlyDuration: .windowDuration(minutes: ProviderPaceCapability.monthlyWindowSentinelMinutes))     public let resetWindowPace: ProviderPaceWindowRule    public let inferredMonthlyDuration: ProviderPaceDurationRule     public init(        resetWindowPace: ProviderPaceWindowRule = .unsupported,        inferredMonthlyDuration: ProviderPaceDurationRule = .unsupported)    {        self.resetWindowPace = resetWindowPace        self.inferredMonthlyDuration = inferredMonthlyDuration    }     public func supportsResetWindowPace(window: RateWindow, now: Date) -> Bool {        self.resetWindowPace.matches(window: window, now: now)    }     public func usesInferredMonthlyDuration(window: RateWindow) -> Bool {        self.inferredMonthlyDuration.matches(window: window)    }} public struct ProviderDescriptor: Sendable {    public let id: UsageProvider    public let metadata: ProviderMetadata    public let branding: ProviderBranding    public let tokenCost: ProviderTokenCostConfig    public let pace: ProviderPaceCapability    public let fetchPlan: ProviderFetchPlan    public let cli: ProviderCLIConfig     public init(        id: UsageProvider,        metadata: ProviderMetadata,        branding: ProviderBranding,        tokenCost: ProviderTokenCostConfig,        pace: ProviderPaceCapability = .unsupported,        fetchPlan: ProviderFetchPlan,        cli: ProviderCLIConfig)    {        self.id = id        self.metadata = metadata        self.branding = branding        self.tokenCost = tokenCost        self.pace = pace        self.fetchPlan = fetchPlan        self.cli = cli    }     public func fetchOutcome(context: ProviderFetchContext) async -> ProviderFetchOutcome {        await self.fetchPlan.fetchOutcome(context: context, provider: self.id)    }     public func fetch(context: ProviderFetchContext) async throws -> ProviderFetchResult {        let outcome = await self.fetchOutcome(context: context)        return try outcome.result.get()    }} public enum ProviderDescriptorRegistry {    private final class Store: @unchecked Sendable {        var ordered: [ProviderDescriptor] = []        var byID: [UsageProvider: ProviderDescriptor] = [:]    }     private static let lock = NSLock()    private static let store = Store()    private static let descriptorsByID: [UsageProvider: ProviderDescriptor] = [        .codex: CodexProviderDescriptor.descriptor,        .openai: OpenAIAPIProviderDescriptor.descriptor,        .azureopenai: AzureOpenAIProviderDescriptor.descriptor,        .claude: ClaudeProviderDescriptor.descriptor,        .clinepass: ClinePassProviderDescriptor.descriptor,        .cursor: CursorProviderDescriptor.descriptor,        .opencode: OpenCodeProviderDescriptor.descriptor,        .opencodego: OpenCodeGoProviderDescriptor.descriptor,        .alibaba: AlibabaCodingPlanProviderDescriptor.descriptor,        .alibabatokenplan: AlibabaTokenPlanProviderDescriptor.descriptor,        .factory: FactoryProviderDescriptor.descriptor,        .gemini: GeminiProviderDescriptor.descriptor,        .antigravity: AntigravityProviderDescriptor.descriptor,        .copilot: CopilotProviderDescriptor.descriptor,        .devin: DevinProviderDescriptor.descriptor,        .zai: ZaiProviderDescriptor.descriptor,        .minimax: MiniMaxProviderDescriptor.descriptor,        .manus: ManusProviderDescriptor.descriptor,        .kimi: KimiProviderDescriptor.descriptor,        .kilo: KiloProviderDescriptor.descriptor,        .kiro: KiroProviderDescriptor.descriptor,        .vertexai: VertexAIProviderDescriptor.descriptor,        .augment: AugmentProviderDescriptor.descriptor,        .jetbrains: JetBrainsProviderDescriptor.descriptor,        .moonshot: MoonshotProviderDescriptor.descriptor,        .amp: AmpProviderDescriptor.descriptor,        .t3chat: T3ChatProviderDescriptor.descriptor,        .ollama: OllamaProviderDescriptor.descriptor,        .synthetic: SyntheticProviderDescriptor.descriptor,        .openrouter: OpenRouterProviderDescriptor.descriptor,        .elevenlabs: ElevenLabsProviderDescriptor.descriptor,        .warp: WarpProviderDescriptor.descriptor,        .windsurf: WindsurfProviderDescriptor.descriptor,        .zed: ZedProviderDescriptor.descriptor,        .perplexity: PerplexityProviderDescriptor.descriptor,        .mimo: MiMoProviderDescriptor.descriptor,        .doubao: DoubaoProviderDescriptor.descriptor,        .sakana: SakanaProviderDescriptor.descriptor,        .abacus: AbacusProviderDescriptor.descriptor,        .mistral: MistralProviderDescriptor.descriptor,        .deepseek: DeepSeekProviderDescriptor.descriptor,        .deepinfra: DeepInfraProviderDescriptor.descriptor,        .codebuff: CodebuffProviderDescriptor.descriptor,        .crof: CrofProviderDescriptor.descriptor,        .venice: VeniceProviderDescriptor.descriptor,        .commandcode: CommandCodeProviderDescriptor.descriptor,        .qoder: QoderProviderDescriptor.descriptor,        .stepfun: StepFunProviderDescriptor.descriptor,        .bedrock: BedrockProviderDescriptor.descriptor,        .grok: GrokProviderDescriptor.descriptor,        .groq: GroqProviderDescriptor.descriptor,        .llmproxy: LLMProxyProviderDescriptor.descriptor,        .litellm: LiteLLMProviderDescriptor.descriptor,        .deepgram: DeepgramProviderDescriptor.descriptor,        .poe: PoeProviderDescriptor.descriptor,        .chutes: ChutesProviderDescriptor.descriptor,        .neuralwatt: NeuralWattProviderDescriptor.descriptor,        .clawrouter: ClawRouterProviderDescriptor.descriptor,        .longcat: LongCatProviderDescriptor.descriptor,        .sub2api: Sub2APIProviderDescriptor.descriptor,        .wayfinder: WayfinderProviderDescriptor.descriptor,        .zenmux: ZenMuxProviderDescriptor.descriptor,        .aiand: AiAndProviderDescriptor.descriptor,    ]    private static let bootstrap: Void = {        for provider in UsageProvider.allCases {            guard let descriptor = descriptorsByID[provider] else {                preconditionFailure("Missing ProviderDescriptor for \(provider.rawValue)")            }            _ = ProviderDescriptorRegistry.register(descriptor)        }    }()     private static func ensureBootstrapped() {        _ = self.bootstrap    }     @discardableResult    public static func register(_ descriptor: ProviderDescriptor) -> ProviderDescriptor {        self.lock.lock()        defer { self.lock.unlock() }        if self.store.byID[descriptor.id] == nil {            self.store.ordered.append(descriptor)        }        self.store.byID[descriptor.id] = descriptor        return descriptor    }     public static var all: [ProviderDescriptor] {        self.ensureBootstrapped()        self.lock.lock()        defer { self.lock.unlock() }        return self.store.ordered    }     public static var metadata: [UsageProvider: ProviderMetadata] {        Dictionary(uniqueKeysWithValues: self.all.map { ($0.id, $0.metadata) })    }     public static func descriptor(for id: UsageProvider) -> ProviderDescriptor {        self.ensureBootstrapped()        if let found = self.store.byID[id] { return found }        if let found = self.all.first(where: { $0.id == id }) { return found }        fatalError("Missing ProviderDescriptor for \(id.rawValue)")    }     public static var cliNameMap: [String: UsageProvider] {        self.ensureBootstrapped()        var map: [String: UsageProvider] = [:]        for descriptor in self.all {            map[descriptor.cli.name] = descriptor.id            for alias in descriptor.cli.aliases {                map[alias] = descriptor.id            }        }        return map    }}