import CUsage
import Foundation
/// One quota window, matching the contract in packages/ai/docs/usage-feeds.md.
public struct UsageLimit: Codable, Sendable, Equatable {
/// Stable machine key within the provider; `label` is display text.
public var id: String
public var label: String
public var usedPercent: Double?
public var windowLabel: String?
public var windowSeconds: Int64?
public var resetsAt: Int64?
/// The model family or feature this limit is restricted to; nil counts all usage.
public var scope: String?
public var status: String
public var detail: String?
}
/// What a provider's feed reports for one account.
public struct UsageReport: Codable, Sendable, Equatable {
public var provider: String?
public var providerName: String?
public var accountId: String?
public var email: String?
public var planType: String?
public var limits: [UsageLimit]
public var balance: String?
public var notes: [String]?
public var error: String?
}
/// A provider with a feed, as the core describes it.
public struct UsageProviderInfo: Codable, Sendable, Equatable {
public var provider: String
public var providerName: String
public var credentialKinds: [String]
}
/// The request the core wants sent. The host does the sending.
///
/// Headers arrive as ordered pairs across the C ABI, because HTTP allows a name to repeat and a
/// dictionary cannot. `headers` is the convenience view for the common case.
public struct UsageRequest: Codable, Sendable, Equatable {
public var kind: String = "http"
// HTTP request fields
public var method: String = "GET"
public var url: String = ""
public var headerPairs: [[String]]?
public var body: String?
// Command request fields
public var command: String?
public var args: [String]?
public var stdin: String?
public var timeoutMs: Int?
public var headers: [String: String] {
Dictionary(
(headerPairs ?? []).compactMap { pair in pair.count == 2 ? (pair[0], pair[1]) : nil },
uniquingKeysWith: { _, last in last }
)
}
enum CodingKeys: String, CodingKey {
case kind, method, url, body, command, args, stdin, timeoutMs
case headerPairs = "headers"
}
public init(
kind: String = "http",
method: String = "GET",
url: String = "",
headerPairs: [[String]]? = nil,
body: String? = nil,
command: String? = nil,
args: [String]? = nil,
stdin: String? = nil,
timeoutMs: Int? = nil
) {
self.kind = kind
self.method = method
self.url = url
self.headerPairs = headerPairs
self.body = body
self.command = command
self.args = args
self.stdin = stdin
self.timeoutMs = timeoutMs
}
public init(from decoder: Decoder) throws {
let container = try decoder.container(keyedBy: CodingKeys.self)
self.kind = try container.decodeIfPresent(String.self, forKey: .kind) ?? "http"
self.method = try container.decodeIfPresent(String.self, forKey: .method) ?? "GET"
self.url = try container.decodeIfPresent(String.self, forKey: .url) ?? ""
self.headerPairs = try container.decodeIfPresent([[String]].self, forKey: .headerPairs)
self.body = try container.decodeIfPresent(String.self, forKey: .body)
self.command = try container.decodeIfPresent(String.self, forKey: .command)
self.args = try container.decodeIfPresent([String].self, forKey: .args)
self.stdin = try container.decodeIfPresent(String.self, forKey: .stdin)
self.timeoutMs = try container.decodeIfPresent(Int.self, forKey: .timeoutMs)
}
}
/// One response handed back to the core.
public struct UsageResponse: Codable, Sendable, Equatable {
public var status: Int
public var body: String
public init(status: Int, body: String) {
self.status = status
self.body = body
}
}
/// Either the next requests to send, or the finished report.
public enum UsageStep: Sendable {
case requests([UsageRequest])
case report(UsageReport)
/// A provider that could not report says so in its report rather than by throwing.
public var reportValue: UsageReport? {
if case .report(let report) = self { return report }
return nil
}
}
/// Opaque cross-round state managed by the host.
public final class UsageState: @unchecked Sendable {
let pointer: OpaquePointer?
public init() {
self.pointer = usage_state_new()
}
deinit {
if let pointer {
usage_state_free(pointer)
}
}
}
public enum UsageCoreError: Error, CustomStringConvertible {
case missingLibraryResult
case invalidJSON(String)
case providerRefused(String)
case incompatibleAbiVersion(expected: UInt32, actual: UInt32)
public var description: String {
switch self {
case .missingLibraryResult: return "the usage core returned nothing"
case .invalidJSON(let text): return "unexpected JSON from the usage core: \(text)"
case .providerRefused(let message): return message
case .incompatibleAbiVersion(let expected, let actual):
return "incompatible usage core ABI version: expected \(expected), got \(actual)"
}
}
}
public enum UsageCore {
public static let expectedAbiVersion: UInt32 = 4
public static var abiVersion: UInt32 { usage_abi_version() }
public static func checkAbiVersion() throws {
let version = abiVersion
guard version == expectedAbiVersion else {
throw UsageCoreError.incompatibleAbiVersion(expected: expectedAbiVersion, actual: version)
}
}
public static func providers() throws -> [UsageProviderInfo] {
try checkAbiVersion()
let json = try consume(usage_providers())
return try JSONDecoder().decode([UsageProviderInfo].self, from: Data(json.utf8))
}
/// One step of the host loop: hand the responses you already have (none on the first call) and
/// get back either the requests to send or the report.
public static func advance(
state: UsageState? = nil,
provider: String,
credentialJSON: String,
responses: [UsageResponse],
nowMs: Int64
) throws -> UsageStep {
try checkAbiVersion()
let responsesJSON = String(data: try JSONEncoder().encode(responses), encoding: .utf8) ?? "[]"
let raw = try consume(
provider.withCString { providerPtr in
credentialJSON.withCString { credentialPtr in
responsesJSON.withCString { responsesPtr in
usage_advance(
state?.pointer,
providerPtr, strlen(providerPtr),
credentialPtr, strlen(credentialPtr),
responsesPtr, strlen(responsesPtr),
nowMs
)
}
}
}
)
guard let data = raw.data(using: .utf8), let object = try JSONSerialization.jsonObject(with: data) as? [String: Any] else {
throw UsageCoreError.invalidJSON(raw)
}
if object["requests"] != nil {
return .requests(try JSONDecoder().decode(UsageRequestsEnvelope.self, from: data).requests)
}
if let error = object["error"] as? String {
throw UsageCoreError.providerRefused(error)
}
return .report(try JSONDecoder().decode(UsageReportEnvelope.self, from: data).report)
}
/// Account id and email from an OAuth access token, without a request.
public static func credentialsFromAccessToken(_ jwt: String) throws -> [String: String] {
let json = try consume(
jwt.withCString { pointer in
usage_credentials_from_access_token(pointer, strlen(pointer))
}
)
guard let data = json.data(using: .utf8) else { throw UsageCoreError.invalidJSON(json) }
return (try? JSONDecoder().decode([String: String].self, from: data)) ?? [:]
}
private struct UsageRequestsEnvelope: Codable { var requests: [UsageRequest] }
private struct UsageReportEnvelope: Codable { var report: UsageReport }
/// Every call hands back a string the caller owns. NULL means allocation failed, which for a
/// quota read is a failure like any other, so it becomes an error instead of a crash.
private static func consume(_ pointer: UnsafeMutablePointer<CChar>?) throws -> String {
guard let pointer else { throw UsageCoreError.missingLibraryResult }
defer { usage_free(pointer) }
return String(cString: pointer)
}
}