chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,327 @@
|
||||
import Foundation
|
||||
|
||||
public struct VertexAIOAuthCredentials: Sendable {
|
||||
public let accessToken: String
|
||||
public let refreshToken: String
|
||||
public let clientId: String
|
||||
public let clientSecret: String
|
||||
public let projectId: String?
|
||||
public let email: String?
|
||||
public let expiryDate: Date?
|
||||
|
||||
public init(
|
||||
accessToken: String,
|
||||
refreshToken: String,
|
||||
clientId: String,
|
||||
clientSecret: String,
|
||||
projectId: String?,
|
||||
email: String?,
|
||||
expiryDate: Date?)
|
||||
{
|
||||
self.accessToken = accessToken
|
||||
self.refreshToken = refreshToken
|
||||
self.clientId = clientId
|
||||
self.clientSecret = clientSecret
|
||||
self.projectId = projectId
|
||||
self.email = email
|
||||
self.expiryDate = expiryDate
|
||||
}
|
||||
|
||||
public var needsRefresh: Bool {
|
||||
guard let expiryDate else { return true }
|
||||
// Refresh 5 minutes before expiry
|
||||
return Date().addingTimeInterval(300) > expiryDate
|
||||
}
|
||||
}
|
||||
|
||||
public enum VertexAIOAuthCredentialsError: LocalizedError, Sendable {
|
||||
case notFound
|
||||
case decodeFailed(String)
|
||||
case missingTokens
|
||||
case missingClientCredentials
|
||||
|
||||
public var errorDescription: String? {
|
||||
switch self {
|
||||
case .notFound:
|
||||
"gcloud credentials not found. Run `gcloud auth application-default login` to authenticate."
|
||||
case let .decodeFailed(message):
|
||||
"Failed to decode gcloud credentials: \(message)"
|
||||
case .missingTokens:
|
||||
"gcloud credentials exist but contain no tokens."
|
||||
case .missingClientCredentials:
|
||||
"gcloud credentials missing client ID or secret."
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
public enum VertexAIOAuthCredentialsStore {
|
||||
#if DEBUG
|
||||
@TaskLocal static var gcloudAccessTokenOverrideForTesting: (@Sendable ([String: String]) async throws -> String)?
|
||||
#endif
|
||||
|
||||
private struct ServiceAccountMetadata {
|
||||
let email: String
|
||||
let projectId: String?
|
||||
}
|
||||
|
||||
private static func credentialsFilePath(
|
||||
environment: [String: String] = ProcessInfo.processInfo.environment) -> URL
|
||||
{
|
||||
if let path = environment["GOOGLE_APPLICATION_CREDENTIALS"]?.trimmingCharacters(
|
||||
in: .whitespacesAndNewlines),
|
||||
!path.isEmpty
|
||||
{
|
||||
return URL(fileURLWithPath: path)
|
||||
}
|
||||
|
||||
let home = FileManager.default.homeDirectoryForCurrentUser
|
||||
// gcloud application default credentials location
|
||||
if let configDir = environment["CLOUDSDK_CONFIG"]?.trimmingCharacters(
|
||||
in: .whitespacesAndNewlines),
|
||||
!configDir.isEmpty
|
||||
{
|
||||
return URL(fileURLWithPath: configDir)
|
||||
.appendingPathComponent("application_default_credentials.json")
|
||||
}
|
||||
return home
|
||||
.appendingPathComponent(".config")
|
||||
.appendingPathComponent("gcloud")
|
||||
.appendingPathComponent("application_default_credentials.json")
|
||||
}
|
||||
|
||||
private static func projectFilePath(
|
||||
environment: [String: String] = ProcessInfo.processInfo.environment) -> URL
|
||||
{
|
||||
let home = FileManager.default.homeDirectoryForCurrentUser
|
||||
if let configDir = environment["CLOUDSDK_CONFIG"]?.trimmingCharacters(
|
||||
in: .whitespacesAndNewlines),
|
||||
!configDir.isEmpty
|
||||
{
|
||||
return URL(fileURLWithPath: configDir)
|
||||
.appendingPathComponent("configurations")
|
||||
.appendingPathComponent("config_default")
|
||||
}
|
||||
return home
|
||||
.appendingPathComponent(".config")
|
||||
.appendingPathComponent("gcloud")
|
||||
.appendingPathComponent("configurations")
|
||||
.appendingPathComponent("config_default")
|
||||
}
|
||||
|
||||
public static func hasCredentials(
|
||||
environment: [String: String] = ProcessInfo.processInfo.environment) -> Bool
|
||||
{
|
||||
let url = self.credentialsFilePath(environment: environment)
|
||||
guard FileManager.default.fileExists(atPath: url.path),
|
||||
let data = try? Data(contentsOf: url),
|
||||
let json = try? self.parseJSONObject(data: data)
|
||||
else {
|
||||
return false
|
||||
}
|
||||
|
||||
if self.parseServiceAccountMetadata(json: json) != nil {
|
||||
return true
|
||||
}
|
||||
|
||||
return (try? self.parseUserCredentials(json: json, environment: environment)) != nil
|
||||
}
|
||||
|
||||
public static func load(
|
||||
environment: [String: String] = ProcessInfo.processInfo.environment) throws -> VertexAIOAuthCredentials
|
||||
{
|
||||
let url = self.credentialsFilePath(environment: environment)
|
||||
guard FileManager.default.fileExists(atPath: url.path) else {
|
||||
throw VertexAIOAuthCredentialsError.notFound
|
||||
}
|
||||
|
||||
let data = try Data(contentsOf: url)
|
||||
return try self.parse(data: data, environment: environment)
|
||||
}
|
||||
|
||||
public static func loadForFetch(
|
||||
environment: [String: String] = ProcessInfo.processInfo.environment) async throws -> VertexAIOAuthCredentials
|
||||
{
|
||||
let url = self.credentialsFilePath(environment: environment)
|
||||
guard FileManager.default.fileExists(atPath: url.path) else {
|
||||
throw VertexAIOAuthCredentialsError.notFound
|
||||
}
|
||||
|
||||
let data = try Data(contentsOf: url)
|
||||
let json = try self.parseJSONObject(data: data)
|
||||
if let serviceAccount = self.parseServiceAccountMetadata(json: json) {
|
||||
let token = try await self.printAccessToken(environment: environment)
|
||||
return VertexAIOAuthCredentials(
|
||||
accessToken: token,
|
||||
refreshToken: "",
|
||||
clientId: "",
|
||||
clientSecret: "",
|
||||
projectId: serviceAccount.projectId ?? self.loadProjectId(environment: environment),
|
||||
email: serviceAccount.email,
|
||||
expiryDate: Date().addingTimeInterval(50 * 60))
|
||||
}
|
||||
|
||||
return try self.parseUserCredentials(json: json, environment: environment)
|
||||
}
|
||||
|
||||
public static func parse(data: Data) throws -> VertexAIOAuthCredentials {
|
||||
try self.parse(data: data, environment: ProcessInfo.processInfo.environment)
|
||||
}
|
||||
|
||||
public static func parse(
|
||||
data: Data,
|
||||
environment: [String: String]) throws -> VertexAIOAuthCredentials
|
||||
{
|
||||
let json = try self.parseJSONObject(data: data)
|
||||
return try self.parseUserCredentials(json: json, environment: environment)
|
||||
}
|
||||
|
||||
private static func parseJSONObject(data: Data) throws -> [String: Any] {
|
||||
guard let json = try JSONSerialization.jsonObject(with: data) as? [String: Any] else {
|
||||
throw VertexAIOAuthCredentialsError.decodeFailed("Invalid JSON")
|
||||
}
|
||||
return json
|
||||
}
|
||||
|
||||
private static func parseUserCredentials(
|
||||
json: [String: Any],
|
||||
environment: [String: String]) throws -> VertexAIOAuthCredentials
|
||||
{
|
||||
// Check for service account credentials
|
||||
if self.parseServiceAccountMetadata(json: json) != nil {
|
||||
throw VertexAIOAuthCredentialsError.decodeFailed(
|
||||
"Service account credentials require `gcloud auth application-default print-access-token`.")
|
||||
}
|
||||
|
||||
// User credentials from gcloud auth application-default login
|
||||
guard let clientId = json["client_id"] as? String,
|
||||
let clientSecret = json["client_secret"] as? String
|
||||
else {
|
||||
throw VertexAIOAuthCredentialsError.missingClientCredentials
|
||||
}
|
||||
|
||||
guard let refreshToken = json["refresh_token"] as? String, !refreshToken.isEmpty else {
|
||||
throw VertexAIOAuthCredentialsError.missingTokens
|
||||
}
|
||||
|
||||
// Access token may not be present in the file; we'll need to refresh
|
||||
let accessToken = json["access_token"] as? String ?? ""
|
||||
|
||||
// Try to get project ID from gcloud config
|
||||
let projectId = Self.loadProjectId(environment: environment)
|
||||
|
||||
// Try to extract email from ID token if present
|
||||
let email = Self.extractEmailFromIdToken(json["id_token"] as? String)
|
||||
|
||||
// Parse expiry if present
|
||||
var expiryDate: Date?
|
||||
if let expiryStr = json["token_expiry"] as? String {
|
||||
let formatter = ISO8601DateFormatter()
|
||||
expiryDate = formatter.date(from: expiryStr)
|
||||
}
|
||||
|
||||
return VertexAIOAuthCredentials(
|
||||
accessToken: accessToken,
|
||||
refreshToken: refreshToken,
|
||||
clientId: clientId,
|
||||
clientSecret: clientSecret,
|
||||
projectId: projectId,
|
||||
email: email,
|
||||
expiryDate: expiryDate)
|
||||
}
|
||||
|
||||
public static func save(_ credentials: VertexAIOAuthCredentials) throws {
|
||||
// We don't modify gcloud's credentials file; just cache the access token in memory
|
||||
// The refresh happens on each app launch if needed
|
||||
}
|
||||
|
||||
private static func parseServiceAccountMetadata(json: [String: Any]) -> ServiceAccountMetadata? {
|
||||
guard let email = (json["client_email"] as? String)?.trimmingCharacters(in: .whitespacesAndNewlines),
|
||||
!email.isEmpty,
|
||||
let privateKey = (json["private_key"] as? String)?.trimmingCharacters(in: .whitespacesAndNewlines),
|
||||
!privateKey.isEmpty
|
||||
else {
|
||||
return nil
|
||||
}
|
||||
|
||||
let projectId = (json["project_id"] as? String)?
|
||||
.trimmingCharacters(in: .whitespacesAndNewlines)
|
||||
return ServiceAccountMetadata(
|
||||
email: email,
|
||||
projectId: projectId?.isEmpty == false ? projectId : nil)
|
||||
}
|
||||
|
||||
private static func printAccessToken(environment: [String: String]) async throws -> String {
|
||||
#if DEBUG
|
||||
if let override = self.gcloudAccessTokenOverrideForTesting {
|
||||
let token = try await override(environment)
|
||||
return try self.cleanAccessToken(token)
|
||||
}
|
||||
#endif
|
||||
|
||||
let env = TTYCommandRunner.enrichedEnvironment(baseEnv: environment)
|
||||
let result = try await SubprocessRunner.run(
|
||||
binary: "/usr/bin/env",
|
||||
arguments: ["gcloud", "auth", "application-default", "print-access-token"],
|
||||
environment: env,
|
||||
timeout: 20,
|
||||
label: "vertexai-gcloud-adc-token")
|
||||
return try self.cleanAccessToken(result.stdout)
|
||||
}
|
||||
|
||||
private static func cleanAccessToken(_ token: String) throws -> String {
|
||||
let trimmed = token.trimmingCharacters(in: .whitespacesAndNewlines)
|
||||
guard !trimmed.isEmpty else {
|
||||
throw VertexAIOAuthCredentialsError.missingTokens
|
||||
}
|
||||
return trimmed
|
||||
}
|
||||
|
||||
private static func loadProjectId(environment: [String: String]) -> String? {
|
||||
let configPath = self.projectFilePath(environment: environment)
|
||||
guard let content = try? String(contentsOf: configPath, encoding: .utf8) else {
|
||||
return environment["GOOGLE_CLOUD_PROJECT"]
|
||||
?? environment["GCLOUD_PROJECT"]
|
||||
?? environment["CLOUDSDK_CORE_PROJECT"]
|
||||
}
|
||||
|
||||
// Parse INI-style config for project
|
||||
for line in content.components(separatedBy: .newlines) {
|
||||
let trimmed = line.trimmingCharacters(in: .whitespaces)
|
||||
if trimmed.hasPrefix("project") {
|
||||
let parts = trimmed.components(separatedBy: "=")
|
||||
if parts.count >= 2 {
|
||||
return parts[1].trimmingCharacters(in: .whitespaces)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Try environment variable
|
||||
return environment["GOOGLE_CLOUD_PROJECT"]
|
||||
?? environment["GCLOUD_PROJECT"]
|
||||
?? environment["CLOUDSDK_CORE_PROJECT"]
|
||||
}
|
||||
|
||||
private static func extractEmailFromIdToken(_ token: String?) -> String? {
|
||||
guard let token, !token.isEmpty else { return nil }
|
||||
|
||||
let parts = token.components(separatedBy: ".")
|
||||
guard parts.count >= 2 else { return nil }
|
||||
|
||||
var payload = parts[1]
|
||||
.replacingOccurrences(of: "-", with: "+")
|
||||
.replacingOccurrences(of: "_", with: "/")
|
||||
|
||||
let remainder = payload.count % 4
|
||||
if remainder > 0 {
|
||||
payload += String(repeating: "=", count: 4 - remainder)
|
||||
}
|
||||
|
||||
guard let data = Data(base64Encoded: payload, options: .ignoreUnknownCharacters),
|
||||
let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any]
|
||||
else {
|
||||
return nil
|
||||
}
|
||||
|
||||
return json["email"] as? String
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,128 @@
|
||||
import Foundation
|
||||
#if canImport(FoundationNetworking)
|
||||
import FoundationNetworking
|
||||
#endif
|
||||
|
||||
public enum VertexAITokenRefresher {
|
||||
private static let tokenEndpoint = URL(string: "https://oauth2.googleapis.com/token")!
|
||||
|
||||
public enum RefreshError: LocalizedError, Sendable {
|
||||
case expired
|
||||
case revoked
|
||||
case networkError(Error)
|
||||
case invalidResponse(String)
|
||||
|
||||
public var errorDescription: String? {
|
||||
switch self {
|
||||
case .expired:
|
||||
"Refresh token expired. Run `gcloud auth application-default login` again."
|
||||
case .revoked:
|
||||
"Refresh token was revoked. Run `gcloud auth application-default login` again."
|
||||
case let .networkError(error):
|
||||
"Network error during token refresh: \(error.localizedDescription)"
|
||||
case let .invalidResponse(message):
|
||||
"Invalid refresh response: \(message)"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
public static func refresh(_ credentials: VertexAIOAuthCredentials) async throws -> VertexAIOAuthCredentials {
|
||||
guard !credentials.refreshToken.isEmpty else {
|
||||
throw RefreshError.invalidResponse("No refresh token available")
|
||||
}
|
||||
|
||||
var request = URLRequest(url: Self.tokenEndpoint)
|
||||
request.httpMethod = "POST"
|
||||
request.timeoutInterval = 30
|
||||
request.setValue("application/x-www-form-urlencoded", forHTTPHeaderField: "Content-Type")
|
||||
|
||||
let bodyParams = [
|
||||
"client_id": credentials.clientId,
|
||||
"client_secret": credentials.clientSecret,
|
||||
"refresh_token": credentials.refreshToken,
|
||||
"grant_type": "refresh_token",
|
||||
]
|
||||
|
||||
let bodyString = bodyParams
|
||||
.map { "\($0.key)=\($0.value.addingPercentEncoding(withAllowedCharacters: .urlQueryAllowed) ?? $0.value)" }
|
||||
.joined(separator: "&")
|
||||
request.httpBody = bodyString.data(using: .utf8)
|
||||
|
||||
do {
|
||||
let response = try await ProviderHTTPClient.shared.response(for: request)
|
||||
let data = response.data
|
||||
|
||||
if response.statusCode == 400 || response.statusCode == 401 {
|
||||
if let errorCode = Self.extractErrorCode(from: data) {
|
||||
switch errorCode.lowercased() {
|
||||
case "invalid_grant":
|
||||
throw RefreshError.expired
|
||||
case "unauthorized_client":
|
||||
throw RefreshError.revoked
|
||||
default:
|
||||
throw RefreshError.invalidResponse("Error: \(errorCode)")
|
||||
}
|
||||
}
|
||||
throw RefreshError.expired
|
||||
}
|
||||
|
||||
guard response.statusCode == 200 else {
|
||||
throw RefreshError.invalidResponse("Status \(response.statusCode)")
|
||||
}
|
||||
|
||||
guard let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] else {
|
||||
throw RefreshError.invalidResponse("Invalid JSON")
|
||||
}
|
||||
|
||||
let newAccessToken = json["access_token"] as? String ?? credentials.accessToken
|
||||
let expiresIn = json["expires_in"] as? Double ?? 3600
|
||||
let newExpiryDate = Date().addingTimeInterval(expiresIn)
|
||||
|
||||
// Extract email from new ID token if present
|
||||
let idToken = json["id_token"] as? String
|
||||
let email = Self.extractEmailFromIdToken(idToken) ?? credentials.email
|
||||
|
||||
return VertexAIOAuthCredentials(
|
||||
accessToken: newAccessToken,
|
||||
refreshToken: credentials.refreshToken,
|
||||
clientId: credentials.clientId,
|
||||
clientSecret: credentials.clientSecret,
|
||||
projectId: credentials.projectId,
|
||||
email: email,
|
||||
expiryDate: newExpiryDate)
|
||||
} catch let error as RefreshError {
|
||||
throw error
|
||||
} catch {
|
||||
throw RefreshError.networkError(error)
|
||||
}
|
||||
}
|
||||
|
||||
private static func extractErrorCode(from data: Data) -> String? {
|
||||
guard let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] else { return nil }
|
||||
return json["error"] as? String
|
||||
}
|
||||
|
||||
private static func extractEmailFromIdToken(_ token: String?) -> String? {
|
||||
guard let token, !token.isEmpty else { return nil }
|
||||
|
||||
let parts = token.components(separatedBy: ".")
|
||||
guard parts.count >= 2 else { return nil }
|
||||
|
||||
var payload = parts[1]
|
||||
.replacingOccurrences(of: "-", with: "+")
|
||||
.replacingOccurrences(of: "_", with: "/")
|
||||
|
||||
let remainder = payload.count % 4
|
||||
if remainder > 0 {
|
||||
payload += String(repeating: "=", count: 4 - remainder)
|
||||
}
|
||||
|
||||
guard let data = Data(base64Encoded: payload, options: .ignoreUnknownCharacters),
|
||||
let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any]
|
||||
else {
|
||||
return nil
|
||||
}
|
||||
|
||||
return json["email"] as? String
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,292 @@
|
||||
import Foundation
|
||||
#if canImport(FoundationNetworking)
|
||||
import FoundationNetworking
|
||||
#endif
|
||||
|
||||
public enum VertexAIFetchError: LocalizedError, Sendable {
|
||||
case unauthorized
|
||||
case forbidden
|
||||
case noProject
|
||||
case networkError(Error)
|
||||
case invalidResponse(String)
|
||||
case noData
|
||||
|
||||
public var errorDescription: String? {
|
||||
switch self {
|
||||
case .unauthorized:
|
||||
"Vertex AI request unauthorized. Run `gcloud auth application-default login`."
|
||||
case .forbidden:
|
||||
"Access forbidden. Check your IAM permissions for Cloud Monitoring."
|
||||
case .noProject:
|
||||
"No Google Cloud project configured. Run `gcloud config set project PROJECT_ID`."
|
||||
case let .networkError(error):
|
||||
"Vertex AI network error: \(error.localizedDescription)"
|
||||
case let .invalidResponse(message):
|
||||
"Vertex AI response was invalid: \(message)"
|
||||
case .noData:
|
||||
"No Vertex AI usage data found for the current project."
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
public struct VertexAIUsageResponse: Sendable {
|
||||
public let requestsUsedPercent: Double
|
||||
public let tokensUsedPercent: Double?
|
||||
public let resetsAt: Date?
|
||||
public let resetDescription: String?
|
||||
public let rawData: String?
|
||||
|
||||
public init(
|
||||
requestsUsedPercent: Double,
|
||||
tokensUsedPercent: Double?,
|
||||
resetsAt: Date?,
|
||||
resetDescription: String?,
|
||||
rawData: String?)
|
||||
{
|
||||
self.requestsUsedPercent = requestsUsedPercent
|
||||
self.tokensUsedPercent = tokensUsedPercent
|
||||
self.resetsAt = resetsAt
|
||||
self.resetDescription = resetDescription
|
||||
self.rawData = rawData
|
||||
}
|
||||
}
|
||||
|
||||
public enum VertexAIUsageFetcher {
|
||||
private static let log = CodexBarLog.logger(LogCategories.vertexAIFetcher)
|
||||
|
||||
// Cloud Monitoring API endpoint for time series
|
||||
private static let monitoringEndpoint = "https://monitoring.googleapis.com/v3/projects"
|
||||
private static let usageWindowSeconds: TimeInterval = 24 * 60 * 60
|
||||
|
||||
public static func fetchUsage(
|
||||
accessToken: String,
|
||||
projectId: String?) async throws -> VertexAIUsageResponse
|
||||
{
|
||||
guard let projectId, !projectId.isEmpty else {
|
||||
throw VertexAIFetchError.noProject
|
||||
}
|
||||
|
||||
return try await Self.fetchQuotaUsage(
|
||||
accessToken: accessToken,
|
||||
projectId: projectId)
|
||||
}
|
||||
|
||||
private static func fetchQuotaUsage(
|
||||
accessToken: String,
|
||||
projectId: String) async throws -> VertexAIUsageResponse
|
||||
{
|
||||
let usageFilter = """
|
||||
metric.type="serviceruntime.googleapis.com/quota/allocation/usage" \
|
||||
AND resource.type="consumer_quota" \
|
||||
AND resource.label.service="aiplatform.googleapis.com"
|
||||
"""
|
||||
let limitFilter = """
|
||||
metric.type="serviceruntime.googleapis.com/quota/limit" \
|
||||
AND resource.type="consumer_quota" \
|
||||
AND resource.label.service="aiplatform.googleapis.com"
|
||||
"""
|
||||
|
||||
let usageSeries = try await Self.fetchTimeSeries(
|
||||
accessToken: accessToken,
|
||||
projectId: projectId,
|
||||
filter: usageFilter)
|
||||
let limitSeries = try await Self.fetchTimeSeries(
|
||||
accessToken: accessToken,
|
||||
projectId: projectId,
|
||||
filter: limitFilter)
|
||||
|
||||
let usageByKey = Self.aggregate(series: usageSeries)
|
||||
let limitByKey = Self.aggregate(series: limitSeries)
|
||||
|
||||
guard !usageByKey.isEmpty, !limitByKey.isEmpty else {
|
||||
throw VertexAIFetchError.noData
|
||||
}
|
||||
|
||||
var maxPercent: Double?
|
||||
var matchedCount = 0
|
||||
var matchedKeys: Set<QuotaKey> = []
|
||||
for (key, limit) in limitByKey {
|
||||
guard limit > 0, let usage = usageByKey[key] else { continue }
|
||||
matchedKeys.insert(key)
|
||||
matchedCount += 1
|
||||
let percent = (usage / limit) * 100.0
|
||||
maxPercent = max(maxPercent ?? percent, percent)
|
||||
}
|
||||
|
||||
guard let usedPercent = maxPercent, matchedCount > 0 else {
|
||||
throw VertexAIFetchError.noData
|
||||
}
|
||||
|
||||
let unmatchedUsage = Set(usageByKey.keys).subtracting(matchedKeys).count
|
||||
let unmatchedLimit = Set(limitByKey.keys).subtracting(matchedKeys).count
|
||||
Self.log.debug("Quota series preview", metadata: [
|
||||
"usageKeys": Self.previewKeys(usageByKey),
|
||||
"limitKeys": Self.previewKeys(limitByKey),
|
||||
])
|
||||
Self.log.info("Parsed quota", metadata: [
|
||||
"usedPercent": "\(usedPercent)",
|
||||
"usageSeries": "\(usageByKey.count)",
|
||||
"limitSeries": "\(limitByKey.count)",
|
||||
"matchedSeries": "\(matchedCount)",
|
||||
"unmatchedUsage": "\(unmatchedUsage)",
|
||||
"unmatchedLimit": "\(unmatchedLimit)",
|
||||
])
|
||||
|
||||
return VertexAIUsageResponse(
|
||||
requestsUsedPercent: usedPercent,
|
||||
tokensUsedPercent: nil,
|
||||
resetsAt: nil,
|
||||
resetDescription: nil,
|
||||
rawData: nil)
|
||||
}
|
||||
|
||||
private struct MonitoringTimeSeriesResponse: Decodable {
|
||||
let timeSeries: [MonitoringTimeSeries]?
|
||||
let nextPageToken: String?
|
||||
}
|
||||
|
||||
private struct MonitoringTimeSeries: Decodable {
|
||||
let metric: MonitoringMetric
|
||||
let resource: MonitoringResource
|
||||
let points: [MonitoringPoint]
|
||||
}
|
||||
|
||||
private struct MonitoringMetric: Decodable {
|
||||
let type: String?
|
||||
let labels: [String: String]?
|
||||
}
|
||||
|
||||
private struct MonitoringResource: Decodable {
|
||||
let type: String?
|
||||
let labels: [String: String]?
|
||||
}
|
||||
|
||||
private struct MonitoringPoint: Decodable {
|
||||
let value: MonitoringValue
|
||||
}
|
||||
|
||||
private struct MonitoringValue: Decodable {
|
||||
let doubleValue: Double?
|
||||
let int64Value: String?
|
||||
}
|
||||
|
||||
private struct QuotaKey: Hashable {
|
||||
let quotaMetric: String
|
||||
let limitName: String
|
||||
let location: String
|
||||
}
|
||||
|
||||
private static func fetchTimeSeries(
|
||||
accessToken: String,
|
||||
projectId: String,
|
||||
filter: String) async throws -> [MonitoringTimeSeries]
|
||||
{
|
||||
let now = Date()
|
||||
let start = now.addingTimeInterval(-Self.usageWindowSeconds)
|
||||
let formatter = ISO8601DateFormatter()
|
||||
var pageToken: String?
|
||||
var allSeries: [MonitoringTimeSeries] = []
|
||||
|
||||
repeat {
|
||||
guard var components = URLComponents(
|
||||
string: "\(Self.monitoringEndpoint)/\(projectId)/timeSeries")
|
||||
else {
|
||||
throw VertexAIFetchError.invalidResponse("Invalid Monitoring URL")
|
||||
}
|
||||
|
||||
var queryItems = [
|
||||
URLQueryItem(name: "filter", value: filter),
|
||||
URLQueryItem(name: "interval.startTime", value: formatter.string(from: start)),
|
||||
URLQueryItem(name: "interval.endTime", value: formatter.string(from: now)),
|
||||
URLQueryItem(name: "aggregation.alignmentPeriod", value: "3600s"),
|
||||
URLQueryItem(name: "aggregation.perSeriesAligner", value: "ALIGN_MAX"),
|
||||
URLQueryItem(name: "view", value: "FULL"),
|
||||
]
|
||||
if let pageToken {
|
||||
queryItems.append(URLQueryItem(name: "pageToken", value: pageToken))
|
||||
}
|
||||
components.queryItems = queryItems
|
||||
|
||||
guard let url = components.url else {
|
||||
throw VertexAIFetchError.invalidResponse("Invalid Monitoring URL")
|
||||
}
|
||||
|
||||
var request = URLRequest(url: url)
|
||||
request.setValue("Bearer \(accessToken)", forHTTPHeaderField: "Authorization")
|
||||
request.timeoutInterval = 30
|
||||
|
||||
let response: ProviderHTTPResponse
|
||||
|
||||
do {
|
||||
response = try await ProviderHTTPClient.shared.response(for: request)
|
||||
} catch {
|
||||
throw VertexAIFetchError.networkError(error)
|
||||
}
|
||||
|
||||
switch response.statusCode {
|
||||
case 401:
|
||||
throw VertexAIFetchError.unauthorized
|
||||
case 403:
|
||||
throw VertexAIFetchError.forbidden
|
||||
case 200:
|
||||
break
|
||||
default:
|
||||
let body = String(data: response.data, encoding: .utf8) ?? ""
|
||||
throw VertexAIFetchError.invalidResponse("HTTP \(response.statusCode): \(body)")
|
||||
}
|
||||
|
||||
let decoded = try JSONDecoder().decode(MonitoringTimeSeriesResponse.self, from: response.data)
|
||||
if let series = decoded.timeSeries {
|
||||
allSeries.append(contentsOf: series)
|
||||
}
|
||||
pageToken = decoded.nextPageToken?.isEmpty == false ? decoded.nextPageToken : nil
|
||||
} while pageToken != nil
|
||||
|
||||
return allSeries
|
||||
}
|
||||
|
||||
private static func aggregate(series: [MonitoringTimeSeries]) -> [QuotaKey: Double] {
|
||||
var buckets: [QuotaKey: Double] = [:]
|
||||
|
||||
for entry in series {
|
||||
guard let key = Self.quotaKey(from: entry),
|
||||
let value = Self.maxPointValue(from: entry.points)
|
||||
else {
|
||||
continue
|
||||
}
|
||||
buckets[key] = max(buckets[key] ?? 0, value)
|
||||
}
|
||||
|
||||
return buckets
|
||||
}
|
||||
|
||||
private static func quotaKey(from series: MonitoringTimeSeries) -> QuotaKey? {
|
||||
let metricLabels = series.metric.labels ?? [:]
|
||||
let resourceLabels = series.resource.labels ?? [:]
|
||||
let quotaMetric = metricLabels["quota_metric"]
|
||||
?? resourceLabels["quota_id"]
|
||||
guard let quotaMetric, !quotaMetric.isEmpty else { return nil }
|
||||
let limitName = metricLabels["limit_name"] ?? ""
|
||||
let location = resourceLabels["location"] ?? "global"
|
||||
return QuotaKey(quotaMetric: quotaMetric, limitName: limitName, location: location)
|
||||
}
|
||||
|
||||
private static func maxPointValue(from points: [MonitoringPoint]) -> Double? {
|
||||
points.compactMap(self.pointValue).max()
|
||||
}
|
||||
|
||||
private static func pointValue(from point: MonitoringPoint) -> Double? {
|
||||
if let doubleValue = point.value.doubleValue { return doubleValue }
|
||||
if let int64Value = point.value.int64Value { return Double(int64Value) }
|
||||
return nil
|
||||
}
|
||||
|
||||
private static func previewKeys(_ map: [QuotaKey: Double], maxCount: Int = 3) -> String {
|
||||
guard !map.isEmpty else { return "none" }
|
||||
let keys = map.keys.prefix(maxCount).map { key in
|
||||
"\(key.quotaMetric)|\(key.limitName)|\(key.location)"
|
||||
}
|
||||
let suffix = map.count > maxCount ? " +\(map.count - maxCount)" : ""
|
||||
return keys.joined(separator: ", ") + suffix
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,110 @@
|
||||
import Foundation
|
||||
|
||||
public enum VertexAIProviderDescriptor {
|
||||
public static let descriptor: ProviderDescriptor = Self.makeDescriptor()
|
||||
|
||||
static func makeDescriptor() -> ProviderDescriptor {
|
||||
ProviderDescriptor(
|
||||
id: .vertexai,
|
||||
metadata: ProviderMetadata(
|
||||
id: .vertexai,
|
||||
displayName: "Vertex AI",
|
||||
sessionLabel: "Requests",
|
||||
weeklyLabel: "Tokens",
|
||||
opusLabel: nil,
|
||||
supportsOpus: false,
|
||||
supportsCredits: false,
|
||||
creditsHint: "",
|
||||
toggleTitle: "Show Vertex AI usage",
|
||||
cliName: "vertexai",
|
||||
defaultEnabled: false,
|
||||
isPrimaryProvider: false,
|
||||
usesAccountFallback: false,
|
||||
dashboardURL: "https://console.cloud.google.com/vertex-ai",
|
||||
statusPageURL: nil,
|
||||
statusLinkURL: "https://status.cloud.google.com"),
|
||||
branding: ProviderBranding(
|
||||
iconStyle: .vertexai,
|
||||
iconResourceName: "ProviderIcon-vertexai",
|
||||
color: ProviderColor(red: 66 / 255, green: 133 / 255, blue: 244 / 255)),
|
||||
tokenCost: ProviderTokenCostConfig(
|
||||
supportsTokenCost: true,
|
||||
noDataMessage: { "No Vertex AI cost data found in Claude logs. Ensure entries include Vertex metadata."
|
||||
}),
|
||||
fetchPlan: ProviderFetchPlan(
|
||||
sourceModes: [.auto, .oauth],
|
||||
pipeline: ProviderFetchPipeline(resolveStrategies: { _ in [VertexAIOAuthFetchStrategy()] })),
|
||||
cli: ProviderCLIConfig(
|
||||
name: "vertexai",
|
||||
versionDetector: nil))
|
||||
}
|
||||
}
|
||||
|
||||
struct VertexAIOAuthFetchStrategy: ProviderFetchStrategy {
|
||||
let id: String = "vertexai.oauth"
|
||||
let kind: ProviderFetchKind = .oauth
|
||||
|
||||
func isAvailable(_ context: ProviderFetchContext) async -> Bool {
|
||||
VertexAIOAuthCredentialsStore.hasCredentials(environment: context.env)
|
||||
}
|
||||
|
||||
func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult {
|
||||
var credentials = try await VertexAIOAuthCredentialsStore.loadForFetch(environment: context.env)
|
||||
|
||||
// Refresh token if expired
|
||||
if credentials.needsRefresh {
|
||||
credentials = try await VertexAITokenRefresher.refresh(credentials)
|
||||
try VertexAIOAuthCredentialsStore.save(credentials)
|
||||
}
|
||||
|
||||
// Fetch quota usage from Cloud Monitoring. If no data is found (e.g., no recent
|
||||
// Vertex AI requests), return an empty snapshot so token costs can still display.
|
||||
let usage: VertexAIUsageResponse?
|
||||
do {
|
||||
usage = try await VertexAIUsageFetcher.fetchUsage(
|
||||
accessToken: credentials.accessToken,
|
||||
projectId: credentials.projectId)
|
||||
} catch VertexAIFetchError.noData {
|
||||
// No quota data is fine - token costs from local logs can still be shown.
|
||||
usage = nil
|
||||
}
|
||||
|
||||
return self.makeResult(
|
||||
usage: Self.mapUsage(usage, credentials: credentials),
|
||||
sourceLabel: "oauth")
|
||||
}
|
||||
|
||||
func shouldFallback(on error: Error, context _: ProviderFetchContext) -> Bool {
|
||||
if error is VertexAIOAuthCredentialsError { return true }
|
||||
if let fetchError = error as? VertexAIFetchError {
|
||||
switch fetchError {
|
||||
case .unauthorized, .forbidden:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
private static func mapUsage(
|
||||
_ response: VertexAIUsageResponse?,
|
||||
credentials: VertexAIOAuthCredentials) -> UsageSnapshot
|
||||
{
|
||||
// Token cost is fetched separately via CostUsageScanner from local Claude logs.
|
||||
// Quota usage from Cloud Monitoring is optional - we still show token costs if unavailable.
|
||||
|
||||
let identity = ProviderIdentitySnapshot(
|
||||
providerID: .vertexai,
|
||||
accountEmail: credentials.email,
|
||||
accountOrganization: credentials.projectId,
|
||||
loginMethod: "gcloud")
|
||||
|
||||
return UsageSnapshot(
|
||||
primary: nil,
|
||||
secondary: nil,
|
||||
providerCost: nil,
|
||||
updatedAt: Date(),
|
||||
identity: identity)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user