import Foundation
/// Drives a feed with URLSession and command execution: ask the core for requests, send them,
/// hand the responses back, repeat until it reports. The loop is the same one the TypeScript host
/// runs, so a provider with two dependent calls (teams, then billing) works here with no extra machinery.
public enum UsageDriver {
public typealias CommandRunner = @Sendable (UsageRequest) async throws -> UsageResponse
public static func fetch(
provider: String,
credentialJSON: String,
nowMs: Int64 = Int64(Date().timeIntervalSince1970 * 1000),
session: URLSession = URLSession(configuration: .ephemeral),
runCommand: CommandRunner? = nil,
maxRounds: Int = 4
) async throws -> UsageReport {
let state = UsageState()
var responses: [UsageResponse] = []
for _ in 0..<maxRounds {
switch try UsageCore.advance(state: state, provider: provider, credentialJSON: credentialJSON, responses: responses, nowMs: nowMs) {
case .report(let report):
return report
case .requests(let requests):
responses = []
for request in requests {
if request.kind == "command" {
if let runCommand {
do {
responses.append(try await runCommand(request))
} catch {
responses.append(UsageResponse(status: 126, body: error.localizedDescription))
}
} else {
// The one answer a feed is allowed to read as "this host will not run things".
responses.append(UsageResponse(status: 127, body: "this host does not run commands (\(request.command ?? ""))"))
}
} else {
responses.append(try await send(request, session: session))
}
}
}
}
return UsageReport(
provider: provider,
providerName: nil,
accountId: nil,
email: nil,
planType: nil,
limits: [],
balance: nil,
notes: nil,
error: "usage feed for \(provider) did not finish in \(maxRounds) rounds"
)
}
private final class PipeBuffer: @unchecked Sendable {
private let lock = NSLock()
private let maxBytes: Int
private var data = Data()
init(maxBytes: Int) {
self.maxBytes = maxBytes
}
func append(_ chunk: Data) {
if chunk.isEmpty { return }
lock.withLock {
if data.count < maxBytes {
data.append(chunk.prefix(maxBytes - data.count))
}
}
}
func snapshot() -> Data {
lock.withLock { data }
}
}
/// Run a command request using Process with the environment allowlist, timeout, and stdout cap.
@Sendable public static func runCommand(_ request: UsageRequest) async throws -> UsageResponse {
guard let cmd = request.command, !cmd.isEmpty else {
return UsageResponse(status: 127, body: "no command specified")
}
let process = Process()
let stdoutPipe = Pipe()
let stderrPipe = Pipe()
let stdinPipe = Pipe()
// Environment allowlist: PATH, HOME, LANG, LC_ALL, LC_CTYPE, TZ and proxy variables
let allowedKeys: Set<String> = [
"PATH", "HOME", "LANG", "LC_ALL", "LC_CTYPE", "TZ",
"HTTP_PROXY", "HTTPS_PROXY", "ALL_PROXY", "NO_PROXY",
"http_proxy", "https_proxy", "all_proxy", "no_proxy",
]
let systemEnv = ProcessInfo.processInfo.environment
var filteredEnv: [String: String] = [:]
for key in allowedKeys {
if let val = systemEnv[key] {
filteredEnv[key] = val
}
}
process.environment = filteredEnv
let args = request.args ?? []
if cmd.contains("/") {
process.executableURL = URL(fileURLWithPath: cmd)
process.arguments = args
} else {
let pathVar = filteredEnv["PATH"] ?? "/usr/bin:/bin:/usr/local/bin"
var foundURL: URL? = nil
for dir in pathVar.split(separator: ":") {
let candidate = URL(fileURLWithPath: String(dir)).appendingPathComponent(cmd)
if FileManager.default.isExecutableFile(atPath: candidate.path) {
foundURL = candidate
break
}
}
if let foundURL {
process.executableURL = foundURL
process.arguments = args
} else {
// Not falling back to `/usr/bin/env`: it would exit 127 with the same code the protocol
// reserves for "this host does not run commands", and the feed would read a missing
// binary as a refusal. 126 says what actually happened.
return UsageResponse(status: 126, body: "\(cmd) is not on PATH")
}
}
process.standardOutput = stdoutPipe
process.standardError = stderrPipe
if request.stdin != nil {
process.standardInput = stdinPipe
}
let timeoutSec: TimeInterval = Double(request.timeoutMs ?? 15_000) / 1000.0
let maxBytes = 64 * 1024
let stdoutBuffer = PipeBuffer(maxBytes: maxBytes)
let stderrBuffer = PipeBuffer(maxBytes: maxBytes)
stdoutPipe.fileHandleForReading.readabilityHandler = { handle in
stdoutBuffer.append(handle.availableData)
}
stderrPipe.fileHandleForReading.readabilityHandler = { handle in
stderrBuffer.append(handle.availableData)
}
do {
try process.run()
} catch {
stdoutPipe.fileHandleForReading.readabilityHandler = nil
stderrPipe.fileHandleForReading.readabilityHandler = nil
// 126 is "found it, could not run it", which is not the same answer as 127 ("this host does
// not run commands"): a feed that treats 127 as "fall back to another path" must not be told
// that when the real problem is a missing binary.
return UsageResponse(status: 126, body: error.localizedDescription)
}
if let stdin = request.stdin, let data = stdin.data(using: .utf8) {
try? stdinPipe.fileHandleForWriting.write(contentsOf: data)
try? stdinPipe.fileHandleForWriting.close()
}
let deadline = Date().addingTimeInterval(timeoutSec)
while process.isRunning && Date() < deadline {
try? await Task.sleep(nanoseconds: 20_000_000)
}
if process.isRunning {
process.terminate()
var termWait = 0
while process.isRunning && termWait < 10 {
try? await Task.sleep(nanoseconds: 20_000_000)
termWait += 1
}
}
stdoutPipe.fileHandleForReading.readabilityHandler = nil
stderrPipe.fileHandleForReading.readabilityHandler = nil
stdoutBuffer.append(stdoutPipe.fileHandleForReading.availableData)
stderrBuffer.append(stderrPipe.fileHandleForReading.availableData)
var body = String(data: stdoutBuffer.snapshot(), encoding: .utf8) ?? ""
let exitCode = Int(process.terminationStatus)
// A command that wrote nothing to stdout but failed still has something to say, and saying it
// on stderr is how a CLI reports a missing argument or a bad credential. Without this the feed
// only sees an empty body and cannot tell the user why nothing came back.
if exitCode != 0, body.isEmpty, let errorText = String(data: stderrBuffer.snapshot(), encoding: .utf8), !errorText.isEmpty {
body = errorText
}
return UsageResponse(status: exitCode, body: body)
}
/// The core asks for one credential kind; a real app resolves it from its own store (or from
/// prime-agent's auth.json). The shape is provider-defined and documented in usage-feeds.md.
public static func credentialJSON(kind: String, value: String, accountId: String? = nil, email: String? = nil, extra: [String: String] = [:]) -> String {
var object: [String: Any] = ["kind": kind]
switch kind {
case "oauth": object["access"] = value
case "apiKey": object["apiKey"] = value
case "session": object["session"] = (try? JSONSerialization.jsonObject(with: Data(value.utf8))) ?? [:]
default: object[kind] = value
}
if let accountId { object["accountId"] = accountId }
if let email { object["email"] = email }
for (key, val) in extra { object[key] = val }
let data = try? JSONSerialization.data(withJSONObject: object)
return data.flatMap { String(data: $0, encoding: .utf8) } ?? "{}"
}
private static func send(_ request: UsageRequest, session: URLSession) async throws -> UsageResponse {
guard let url = URL(string: request.url) else {
return UsageResponse(status: 0, body: "invalid url: \(request.url)")
}
var urlRequest = URLRequest(url: url)
urlRequest.httpMethod = request.method
for (name, value) in request.headers {
urlRequest.setValue(value, forHTTPHeaderField: name)
}
if let body = request.body {
urlRequest.httpBody = Data(body.utf8)
}
do {
let (data, response) = try await session.data(for: urlRequest)
let status = (response as? HTTPURLResponse)?.statusCode ?? 0
return UsageResponse(status: status, body: String(data: data, encoding: .utf8) ?? "")
} catch {
// A transport failure is data for the feed, not an exception: the report carries it.
return UsageResponse(status: 0, body: "transport error: \(error.localizedDescription)")
}
}
}