From a0fc7bd848391aca569089a4eca16670532854e6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lingxiao=20Li=20=EF=BC=88=E6=9D=8E=E5=87=8C=E9=9C=84?= =?UTF-8?q?=EF=BC=89?= Date: Tue, 14 Jul 2026 18:14:38 +0800 Subject: [PATCH 1/3] fix(kiro): Add direct API fetch strategy for enterprise/IdC account usage Fixes Kiro usage fetching for enterprise/IdC accounts where kiro-cli fails because it passes profileArn to the legacy REST API, causing a 400 error. The fix adds a new KiroAPIFetchStrategy that: - Reads credentials directly from kiro-cli's SQLite database - Makes direct HTTP calls to AWS Q getUsageLimits API - Uses fixed KiroIDE-0.9.2 version (doesn't require profileArn) - Falls back to CLI strategy if API fails This is the same approach used by kiro.rs (v0.6.11) which correctly handles enterprise accounts by NOT passing profileArn to usage APIs. Reference: https://github.com/ZyphrZero/kiro.rs/issues/XX Root cause: Enterprise credentials have a profileArn that gets resolved via ListAvailableProfiles. When passed to getUsageLimits, it causes 'Improperly formed request' 400 error. --- .../Providers/Kiro/KiroAPIFetchStrategy.swift | 408 ++++++++++++++++++ .../Kiro/KiroProviderDescriptor.swift | 6 +- 2 files changed, 412 insertions(+), 2 deletions(-) create mode 100644 Sources/CodexBarCore/Providers/Kiro/KiroAPIFetchStrategy.swift diff --git a/Sources/CodexBarCore/Providers/Kiro/KiroAPIFetchStrategy.swift b/Sources/CodexBarCore/Providers/Kiro/KiroAPIFetchStrategy.swift new file mode 100644 index 0000000000..1d67b93522 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Kiro/KiroAPIFetchStrategy.swift @@ -0,0 +1,408 @@ +import Foundation +import SQLite3 +#if canImport(Darwin) +import Darwin +#endif + +/// Direct API fetch strategy for Kiro usage that bypasses `kiro-cli` and calls +/// the AWS Q `getUsageLimits` REST endpoint directly. +/// +/// This fixes enterprise/IdC account usage fetching where `kiro-cli` fails because +/// it incorrectly passes `profileArn` to the legacy REST API, causing a 400 error. +/// The Kiro 0.9.2 version of these APIs does NOT accept `profileArn`. +/// +/// Reference: https://github.com/ZyphrZero/kiro.rs (v0.6.11 fix) +struct KiroAPIFetchStrategy: ProviderFetchStrategy { + let id: String = "kiro.api" + let kind: ProviderFetchKind = .apiToken + + /// Fixed Kiro IDE version for usage REST APIs. + /// This version does not require `profileArn`, which is the key fix for enterprise accounts. + private static let usageAPIKiroVersion = "0.9.2" + + /// AWS Q REST API regions (only these two serve getUsageLimits) + private static let apiRegions = ["us-east-1", "eu-central-1"] + + func isAvailable(_ context: ProviderFetchContext) async -> Bool { + Self.resolveCredentials() != nil + } + + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + guard let credentials = Self.resolveCredentials() else { + throw KiroAPIError.credentialsNotFound + } + + let snapshot = try await Self.fetchUsage(credentials: credentials) + return self.makeResult( + usage: snapshot.toUsageSnapshot(), + sourceLabel: "api") + } + + func shouldFallback(on error: Error, context: ProviderFetchContext) -> Bool { + // Fall back to CLI if API fails + true + } + + // MARK: - Credential Resolution + + /// Read credentials from kiro-cli's SQLite database + private static func resolveCredentials() -> KiroCLICredentials? { + let dbPath = NSString(string: "~/Library/Application Support/kiro-cli/data.sqlite3") + .expandingTildeInPath + + guard FileManager.default.fileExists(atPath: dbPath) else { + return nil + } + + var db: OpaquePointer? + guard sqlite3_open_v2(dbPath, &db, SQLITE_OPEN_READONLY, nil) == SQLITE_OK else { + sqlite3_close(db) + return nil + } + defer { sqlite3_close(db) } + + let query = "SELECT value FROM auth_kv WHERE key = 'kirocli:odic:token'" + var stmt: OpaquePointer? + guard sqlite3_prepare_v2(db, query, -1, &stmt, nil) == SQLITE_OK else { + sqlite3_finalize(stmt) + return nil + } + defer { sqlite3_finalize(stmt) } + + guard sqlite3_step(stmt) == SQLITE_ROW, + let blob = sqlite3_column_text(stmt, 0) + else { + return nil + } + + let jsonString = String(cString: blob) + guard let data = jsonString.data(using: .utf8), + let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any], + let accessToken = json["access_token"] as? String, + !accessToken.isEmpty + else { + return nil + } + + let region = json["region"] as? String + let startURL = json["start_url"] as? String + + return KiroCLICredentials( + accessToken: accessToken, + region: region, + startURL: startURL) + } + + // MARK: - API Fetch + + /// Fetch usage directly from AWS Q API (no profileArn, fixed Kiro version) + private static func fetchUsage(credentials: KiroCLICredentials) async throws -> KiroUsageSnapshot { + let regions = regionCandidates(for: credentials.region) + + var lastError: Error? + for region in regions { + do { + return try await fetchUsageFromRegion( + region: region, + accessToken: credentials.accessToken, + startURL: credentials.startURL) + } catch let error as URLError where error.code == .badServerResponse { + lastError = error + continue + } catch { + throw error + } + } + + throw lastError ?? KiroAPIError.allRegionsFailed + } + + private static func fetchUsageFromRegion( + region: String, + accessToken: String, + startURL: String?) async throws -> KiroUsageSnapshot + { + let host = "q.\(region).amazonaws.com" + let urlString = "https://\(host)/getUsageLimits?origin=AI_EDITOR&resourceType=AGENTIC_REQUEST&isEmailRequired=true" + + guard let url = URL(string: urlString) else { + throw KiroAPIError.invalidURL + } + + var request = URLRequest(url: url) + request.httpMethod = "GET" + + // Machine ID (simplified - use a stable identifier) + let machineId = generateMachineId() + + // User-Agent matching Kiro 0.9.2 (the version that doesn't require profileArn) + let userAgent = "aws-sdk-js/1.0.0 ua/2.1 os/macos lang/js md/nodejs#20.0.0 api/codewhispererruntime#1.0.0 m/N,E KiroIDE-\(usageAPIKiroVersion)-\(machineId)" + let amzUserAgent = "aws-sdk-js/1.0.0 KiroIDE-\(usageAPIKiroVersion)-\(machineId)" + + request.setValue(amzUserAgent, forHTTPHeaderField: "x-amz-user-agent") + request.setValue(userAgent, forHTTPHeaderField: "user-agent") + request.setValue(host, forHTTPHeaderField: "host") + request.setValue(UUID().uuidString, forHTTPHeaderField: "amz-sdk-invocation-id") + request.setValue("attempt=1; max=1", forHTTPHeaderField: "amz-sdk-request") + request.setValue("Bearer \(accessToken)", forHTTPHeaderField: "Authorization") + request.setValue("close", forHTTPHeaderField: "Connection") + + // For enterprise/IdC accounts, we need to detect the auth method. + // If startURL is present, it's likely an enterprise account. + // However, we do NOT pass profileArn - that's the key fix. + if let startURL, !startURL.isEmpty { + // Enterprise accounts may need EXTERNAL_IDP token type + // But we don't have enough info here to determine the exact type + // The API should work without it for usage endpoints + } + + let (data, response) = try await URLSession.shared.data(for: request) + + guard let httpResponse = response as? HTTPURLResponse else { + throw KiroAPIError.invalidResponse + } + + switch httpResponse.statusCode { + case 200: + let decoder = JSONDecoder() + decoder.dateDecodingStrategy = .secondsSince1970 + let usageResponse = try decoder.decode(KiroUsageLimitsResponse.self, from: data) + return try usageResponse.toSnapshot() + case 403: + // Try next region + throw URLError(.badServerResponse) + case 401: + throw KiroAPIError.authenticationFailed + case 429: + throw KiroAPIError.rateLimited + default: + let body = String(data: data, encoding: .utf8) ?? "" + throw KiroAPIError.httpError(statusCode: httpResponse.statusCode, body: body) + } + } + + /// Region candidate ordering: prefer us-east-1, fallback to eu-central-1 + private static func regionCandidates(for ssoRegion: String?) -> [String] { + guard let ssoRegion else { return apiRegions } + if ssoRegion == "eu-central-1" || ssoRegion.hasPrefix("eu-") { + return ["eu-central-1", "us-east-1"] + } + return apiRegions + } + + /// Generate a stable machine ID from the device's hardware UUID + private static func generateMachineId() -> String { + #if canImport(Darwin) + let platformExpert = IOServiceGetMatchingService( + kIOMainPortDefault, + IOServiceMatching("IOPlatformExpertDevice")) + defer { IOObjectRelease(platformExpert) } + + guard platformExpert != 0, + let uuidCF = IORegistryEntryCreateCFProperty( + platformExpert, + kIOPlatformUUIDKey as CFString, + kCFAllocatorDefault, 0) + else { + return UUID().uuidString.replacingOccurrences(of: "-", with: "") + } + + let uuid = uuidCF.takeRetainedValue() as? String ?? UUID().uuidString + return uuid.replacingOccurrences(of: "-", with: "") + #else + return UUID().uuidString.replacingOccurrences(of: "-", with: "") + #endif + } +} + +// MARK: - Supporting Types + +private struct KiroCLICredentials { + let accessToken: String + let region: String? + let startURL: String? +} + +enum KiroAPIError: Error, LocalizedError { + case credentialsNotFound + case invalidURL + case invalidResponse + case authenticationFailed + case rateLimited + case allRegionsFailed + case httpError(statusCode: Int, body: String) + + var errorDescription: String? { + switch self { + case .credentialsNotFound: + return "Kiro credentials not found. Please run 'kiro-cli login' first." + case .invalidURL: + return "Invalid API URL" + case .invalidResponse: + return "Invalid response from Kiro API" + case .authenticationFailed: + return "Kiro authentication failed. Please run 'kiro-cli login' again." + case .rateLimited: + return "Kiro API rate limited. Please try again later." + case .allRegionsFailed: + return "All Kiro API regions failed" + case .httpError(let statusCode, let body): + return "Kiro API error: \(statusCode) \(body)" + } + } +} + +// MARK: - API Response Models + +struct KiroUsageLimitsResponse: Decodable { + let nextDateReset: Double? + let subscriptionInfo: KiroSubscriptionInfo? + let usageBreakdownList: [KiroUsageBreakdown]? + let overageConfiguration: KiroOverageConfiguration? + let userInfo: KiroUserInfo? + + enum CodingKeys: String, CodingKey { + case nextDateReset + case subscriptionInfo + case usageBreakdownList + case overageConfiguration + case userInfo + } + + func toSnapshot() throws -> KiroUsageSnapshot { + let planName = subscriptionInfo?.subscriptionTitle ?? "KIRO FREE" + let email = userInfo?.email + + guard let breakdown = usageBreakdownList?.first else { + // No usage data - might be a managed/enterprise plan + return KiroUsageSnapshot( + planName: planName, + displayPlanName: KiroStatusProbe.displayPlanName(planName), + accountEmail: email, + authMethod: nil, + creditsUsed: 0, + creditsTotal: 0, + creditsPercent: 0, + bonusCreditsUsed: nil, + bonusCreditsTotal: nil, + bonusExpiryDays: nil, + overagesStatus: overageConfiguration?.overageStatus, + overageCreditsUsed: nil, + estimatedOverageCostUSD: nil, + manageURL: nil, + contextUsage: nil, + resetsAt: nextDateReset.map { Date(timeIntervalSince1970: $0) }, + updatedAt: Date()) + } + + let creditsUsed = breakdown.currentUsageWithPrecision ?? Double(breakdown.currentUsage ?? 0) + let creditsTotal = breakdown.usageLimitWithPrecision ?? Double(breakdown.usageLimit ?? 0) + let creditsPercent = creditsTotal > 0 ? (creditsUsed / creditsTotal) * 100.0 : 0 + + // Calculate bonus credits + var bonusUsed: Double? + var bonusTotal: Double? + var bonusExpiry: Int? + + let activeBonuses = (breakdown.bonuses ?? []).filter { $0.status == "ACTIVE" } + if !activeBonuses.isEmpty { + bonusUsed = activeBonuses.reduce(0) { $0 + ($1.currentUsage ?? 0) } + bonusTotal = activeBonuses.reduce(0) { $0 + ($1.usageLimit ?? 0) } + } + + // Overages + let overageEnabled = overageConfiguration?.overageEnabled ?? false + let overagesStatus = overageConfiguration?.overageStatus + + return KiroUsageSnapshot( + planName: planName, + displayPlanName: KiroStatusProbe.displayPlanName(planName), + accountEmail: email, + authMethod: nil, + creditsUsed: creditsUsed, + creditsTotal: creditsTotal, + creditsPercent: creditsPercent, + bonusCreditsUsed: bonusUsed, + bonusCreditsTotal: bonusTotal, + bonusExpiryDays: bonusExpiry, + overagesStatus: overagesStatus, + overageCreditsUsed: nil, + estimatedOverageCostUSD: nil, + manageURL: nil, + contextUsage: nil, + resetsAt: nextDateReset.map { Date(timeIntervalSince1970: $0) }, + updatedAt: Date()) + } +} + +struct KiroSubscriptionInfo: Decodable { + let subscriptionTitle: String? + let overageCapability: String? +} + +struct KiroUsageBreakdown: Decodable { + let currentUsage: Int? + let currentUsageWithPrecision: Double? + let usageLimit: Int? + let usageLimitWithPrecision: Double? + let bonuses: [KiroBonus]? + let freeTrialInfo: KiroFreeTrialInfo? + + enum CodingKeys: String, CodingKey { + case currentUsage + case currentUsageWithPrecision + case usageLimit + case usageLimitWithPrecision + case bonuses + case freeTrialInfo + } + + init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + currentUsage = try container.decodeIfPresent(Int.self, forKey: .currentUsage) + currentUsageWithPrecision = try container.decodeIfPresent(Double.self, forKey: .currentUsageWithPrecision) + usageLimit = try container.decodeIfPresent(Int.self, forKey: .usageLimit) + usageLimitWithPrecision = try container.decodeIfPresent(Double.self, forKey: .usageLimitWithPrecision) + bonuses = try container.decodeIfPresent([KiroBonus].self, forKey: .bonuses) + freeTrialInfo = try container.decodeIfPresent(KiroFreeTrialInfo.self, forKey: .freeTrialInfo) + } +} + +struct KiroBonus: Decodable { + let currentUsage: Double? + let usageLimit: Double? + let status: String? + + enum CodingKeys: String, CodingKey { + case currentUsage + case usageLimit + case status + } +} + +struct KiroFreeTrialInfo: Decodable { + let currentUsage: Int? + let currentUsageWithPrecision: Double? + let usageLimit: Int? + let usageLimitWithPrecision: Double? + let freeTrialExpiry: Double? + let freeTrialStatus: String? + + enum CodingKeys: String, CodingKey { + case currentUsage + case currentUsageWithPrecision + case usageLimit + case usageLimitWithPrecision + case freeTrialExpiry + case freeTrialStatus + } +} + +struct KiroOverageConfiguration: Decodable { + let overageEnabled: Bool? + let overageStatus: String? +} + +struct KiroUserInfo: Decodable { + let email: String? +} diff --git a/Sources/CodexBarCore/Providers/Kiro/KiroProviderDescriptor.swift b/Sources/CodexBarCore/Providers/Kiro/KiroProviderDescriptor.swift index 804b346c08..9a47bd4efb 100644 --- a/Sources/CodexBarCore/Providers/Kiro/KiroProviderDescriptor.swift +++ b/Sources/CodexBarCore/Providers/Kiro/KiroProviderDescriptor.swift @@ -31,8 +31,10 @@ public enum KiroProviderDescriptor { supportsTokenCost: false, noDataMessage: { "Kiro cost summary is not supported." }), fetchPlan: ProviderFetchPlan( - sourceModes: [.auto, .cli], - pipeline: ProviderFetchPipeline(resolveStrategies: { _ in [KiroCLIFetchStrategy()] })), + sourceModes: [.auto, .api, .cli], + pipeline: ProviderFetchPipeline(resolveStrategies: { _ in + [KiroAPIFetchStrategy(), KiroCLIFetchStrategy()] + })), cli: ProviderCLIConfig( name: "kiro", aliases: ["kiro-cli"], From 8b8684e98ac394e2fdbc882a90f2c0278d19b6f3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lingxiao=20Li=20=EF=BC=88=E6=9D=8E=E5=87=8C=E9=9C=84?= =?UTF-8?q?=EF=BC=89?= Date: Wed, 15 Jul 2026 08:39:06 +0800 Subject: [PATCH 2/3] fix: respect source mode and support all credential keys - Add source mode handling: auto (API+CLI), api (API only), cli (CLI only) - Support multiple credential keys: kirocli:odic:token, kirocli:oidc:token, kirocli:social:token - shouldFallback only in auto mode, not when explicit source selected --- .../Providers/Kiro/KiroAPIFetchStrategy.swift | 27 ++++++++++++++++--- .../Kiro/KiroProviderDescriptor.swift | 17 +++++++++--- 2 files changed, 37 insertions(+), 7 deletions(-) diff --git a/Sources/CodexBarCore/Providers/Kiro/KiroAPIFetchStrategy.swift b/Sources/CodexBarCore/Providers/Kiro/KiroAPIFetchStrategy.swift index 1d67b93522..1a6dafcb4c 100644 --- a/Sources/CodexBarCore/Providers/Kiro/KiroAPIFetchStrategy.swift +++ b/Sources/CodexBarCore/Providers/Kiro/KiroAPIFetchStrategy.swift @@ -39,12 +39,19 @@ struct KiroAPIFetchStrategy: ProviderFetchStrategy { } func shouldFallback(on error: Error, context: ProviderFetchContext) -> Bool { - // Fall back to CLI if API fails - true + // Only fall back to CLI in auto mode + context.sourceMode == .auto } // MARK: - Credential Resolution + /// Credential keys to try, in order of preference + private static let credentialKeys = [ + "kirocli:odic:token", // Current (typo in kiro-cli, kept for compat) + "kirocli:oidc:token", // Correct spelling, newer versions + "kirocli:social:token", // Social login (GitHub/Google) + ] + /// Read credentials from kiro-cli's SQLite database private static func resolveCredentials() -> KiroCLICredentials? { let dbPath = NSString(string: "~/Library/Application Support/kiro-cli/data.sqlite3") @@ -61,7 +68,19 @@ struct KiroAPIFetchStrategy: ProviderFetchStrategy { } defer { sqlite3_close(db) } - let query = "SELECT value FROM auth_kv WHERE key = 'kirocli:odic:token'" + // Try each credential key until we find one with a valid token + for key in credentialKeys { + if let creds = queryCredentials(db: db, key: key) { + return creds + } + } + + return nil + } + + private static func queryCredentials(db: OpaquePointer?, key: String) -> KiroCLICredentials? { + guard let db else { return nil } + let query = "SELECT value FROM auth_kv WHERE key = ?" var stmt: OpaquePointer? guard sqlite3_prepare_v2(db, query, -1, &stmt, nil) == SQLITE_OK else { sqlite3_finalize(stmt) @@ -69,6 +88,8 @@ struct KiroAPIFetchStrategy: ProviderFetchStrategy { } defer { sqlite3_finalize(stmt) } + sqlite3_bind_text(stmt, 1, key, -1, unsafeBitCast(-1, to: sqlite3_destructor_type.self)) + guard sqlite3_step(stmt) == SQLITE_ROW, let blob = sqlite3_column_text(stmt, 0) else { diff --git a/Sources/CodexBarCore/Providers/Kiro/KiroProviderDescriptor.swift b/Sources/CodexBarCore/Providers/Kiro/KiroProviderDescriptor.swift index 9a47bd4efb..a9651cc9f8 100644 --- a/Sources/CodexBarCore/Providers/Kiro/KiroProviderDescriptor.swift +++ b/Sources/CodexBarCore/Providers/Kiro/KiroProviderDescriptor.swift @@ -32,8 +32,17 @@ public enum KiroProviderDescriptor { noDataMessage: { "Kiro cost summary is not supported." }), fetchPlan: ProviderFetchPlan( sourceModes: [.auto, .api, .cli], - pipeline: ProviderFetchPipeline(resolveStrategies: { _ in - [KiroAPIFetchStrategy(), KiroCLIFetchStrategy()] + pipeline: ProviderFetchPipeline(resolveStrategies: { context in + switch context.sourceMode { + case .auto: + [KiroAPIFetchStrategy(), KiroCLIFetchStrategy()] + case .api: + [KiroAPIFetchStrategy()] + case .cli: + [KiroCLIFetchStrategy()] + case .web, .oauth: + [] + } })), cli: ProviderCLIConfig( name: "kiro", @@ -58,7 +67,7 @@ struct KiroCLIFetchStrategy: ProviderFetchStrategy { sourceLabel: "cli") } - func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool { - false + func shouldFallback(on _: Error, context: ProviderFetchContext) -> Bool { + context.sourceMode == .auto } } From 27dc2457d6804c006142ceb31c38011bdb11fb06 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lingxiao=20Li=20=EF=BC=88=E6=9D=8E=E5=87=8C=E9=9C=84?= =?UTF-8?q?=EF=BC=89?= Date: Fri, 17 Jul 2026 10:43:53 +0800 Subject: [PATCH 3/3] Address Kiro API review feedback for enterprise usage. Refactor the direct fetch path into injectable credential and HTTP seams, add token refresh and EXTERNAL_IDP handling, narrow auto mode to enterprise credentials, and cover the new behavior with focused tests. --- .../Providers/Kiro/KiroAPIFetchStrategy.swift | 406 +----------- .../Providers/Kiro/KiroAPIUsageFetcher.swift | 616 ++++++++++++++++++ .../Kiro/KiroCLICredentialStore.swift | 224 +++++++ .../KiroAPIUsageFetcherTests.swift | 384 +++++++++++ 4 files changed, 1230 insertions(+), 400 deletions(-) create mode 100644 Sources/CodexBarCore/Providers/Kiro/KiroAPIUsageFetcher.swift create mode 100644 Sources/CodexBarCore/Providers/Kiro/KiroCLICredentialStore.swift create mode 100644 Tests/CodexBarTests/KiroAPIUsageFetcherTests.swift diff --git a/Sources/CodexBarCore/Providers/Kiro/KiroAPIFetchStrategy.swift b/Sources/CodexBarCore/Providers/Kiro/KiroAPIFetchStrategy.swift index 1a6dafcb4c..621a798e46 100644 --- a/Sources/CodexBarCore/Providers/Kiro/KiroAPIFetchStrategy.swift +++ b/Sources/CodexBarCore/Providers/Kiro/KiroAPIFetchStrategy.swift @@ -1,8 +1,4 @@ import Foundation -import SQLite3 -#if canImport(Darwin) -import Darwin -#endif /// Direct API fetch strategy for Kiro usage that bypasses `kiro-cli` and calls /// the AWS Q `getUsageLimits` REST endpoint directly. @@ -16,414 +12,24 @@ struct KiroAPIFetchStrategy: ProviderFetchStrategy { let id: String = "kiro.api" let kind: ProviderFetchKind = .apiToken - /// Fixed Kiro IDE version for usage REST APIs. - /// This version does not require `profileArn`, which is the key fix for enterprise accounts. - private static let usageAPIKiroVersion = "0.9.2" + private let fetcher: KiroAPIUsageFetcher - /// AWS Q REST API regions (only these two serve getUsageLimits) - private static let apiRegions = ["us-east-1", "eu-central-1"] + init(fetcher: KiroAPIUsageFetcher = KiroAPIUsageFetcher()) { + self.fetcher = fetcher + } func isAvailable(_ context: ProviderFetchContext) async -> Bool { - Self.resolveCredentials() != nil + self.fetcher.hasCredentials(allowSocial: context.sourceMode == .api) } func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { - guard let credentials = Self.resolveCredentials() else { - throw KiroAPIError.credentialsNotFound - } - - let snapshot = try await Self.fetchUsage(credentials: credentials) + let snapshot = try await self.fetcher.fetchUsage(allowSocial: context.sourceMode == .api) return self.makeResult( usage: snapshot.toUsageSnapshot(), sourceLabel: "api") } func shouldFallback(on error: Error, context: ProviderFetchContext) -> Bool { - // Only fall back to CLI in auto mode context.sourceMode == .auto } - - // MARK: - Credential Resolution - - /// Credential keys to try, in order of preference - private static let credentialKeys = [ - "kirocli:odic:token", // Current (typo in kiro-cli, kept for compat) - "kirocli:oidc:token", // Correct spelling, newer versions - "kirocli:social:token", // Social login (GitHub/Google) - ] - - /// Read credentials from kiro-cli's SQLite database - private static func resolveCredentials() -> KiroCLICredentials? { - let dbPath = NSString(string: "~/Library/Application Support/kiro-cli/data.sqlite3") - .expandingTildeInPath - - guard FileManager.default.fileExists(atPath: dbPath) else { - return nil - } - - var db: OpaquePointer? - guard sqlite3_open_v2(dbPath, &db, SQLITE_OPEN_READONLY, nil) == SQLITE_OK else { - sqlite3_close(db) - return nil - } - defer { sqlite3_close(db) } - - // Try each credential key until we find one with a valid token - for key in credentialKeys { - if let creds = queryCredentials(db: db, key: key) { - return creds - } - } - - return nil - } - - private static func queryCredentials(db: OpaquePointer?, key: String) -> KiroCLICredentials? { - guard let db else { return nil } - let query = "SELECT value FROM auth_kv WHERE key = ?" - var stmt: OpaquePointer? - guard sqlite3_prepare_v2(db, query, -1, &stmt, nil) == SQLITE_OK else { - sqlite3_finalize(stmt) - return nil - } - defer { sqlite3_finalize(stmt) } - - sqlite3_bind_text(stmt, 1, key, -1, unsafeBitCast(-1, to: sqlite3_destructor_type.self)) - - guard sqlite3_step(stmt) == SQLITE_ROW, - let blob = sqlite3_column_text(stmt, 0) - else { - return nil - } - - let jsonString = String(cString: blob) - guard let data = jsonString.data(using: .utf8), - let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any], - let accessToken = json["access_token"] as? String, - !accessToken.isEmpty - else { - return nil - } - - let region = json["region"] as? String - let startURL = json["start_url"] as? String - - return KiroCLICredentials( - accessToken: accessToken, - region: region, - startURL: startURL) - } - - // MARK: - API Fetch - - /// Fetch usage directly from AWS Q API (no profileArn, fixed Kiro version) - private static func fetchUsage(credentials: KiroCLICredentials) async throws -> KiroUsageSnapshot { - let regions = regionCandidates(for: credentials.region) - - var lastError: Error? - for region in regions { - do { - return try await fetchUsageFromRegion( - region: region, - accessToken: credentials.accessToken, - startURL: credentials.startURL) - } catch let error as URLError where error.code == .badServerResponse { - lastError = error - continue - } catch { - throw error - } - } - - throw lastError ?? KiroAPIError.allRegionsFailed - } - - private static func fetchUsageFromRegion( - region: String, - accessToken: String, - startURL: String?) async throws -> KiroUsageSnapshot - { - let host = "q.\(region).amazonaws.com" - let urlString = "https://\(host)/getUsageLimits?origin=AI_EDITOR&resourceType=AGENTIC_REQUEST&isEmailRequired=true" - - guard let url = URL(string: urlString) else { - throw KiroAPIError.invalidURL - } - - var request = URLRequest(url: url) - request.httpMethod = "GET" - - // Machine ID (simplified - use a stable identifier) - let machineId = generateMachineId() - - // User-Agent matching Kiro 0.9.2 (the version that doesn't require profileArn) - let userAgent = "aws-sdk-js/1.0.0 ua/2.1 os/macos lang/js md/nodejs#20.0.0 api/codewhispererruntime#1.0.0 m/N,E KiroIDE-\(usageAPIKiroVersion)-\(machineId)" - let amzUserAgent = "aws-sdk-js/1.0.0 KiroIDE-\(usageAPIKiroVersion)-\(machineId)" - - request.setValue(amzUserAgent, forHTTPHeaderField: "x-amz-user-agent") - request.setValue(userAgent, forHTTPHeaderField: "user-agent") - request.setValue(host, forHTTPHeaderField: "host") - request.setValue(UUID().uuidString, forHTTPHeaderField: "amz-sdk-invocation-id") - request.setValue("attempt=1; max=1", forHTTPHeaderField: "amz-sdk-request") - request.setValue("Bearer \(accessToken)", forHTTPHeaderField: "Authorization") - request.setValue("close", forHTTPHeaderField: "Connection") - - // For enterprise/IdC accounts, we need to detect the auth method. - // If startURL is present, it's likely an enterprise account. - // However, we do NOT pass profileArn - that's the key fix. - if let startURL, !startURL.isEmpty { - // Enterprise accounts may need EXTERNAL_IDP token type - // But we don't have enough info here to determine the exact type - // The API should work without it for usage endpoints - } - - let (data, response) = try await URLSession.shared.data(for: request) - - guard let httpResponse = response as? HTTPURLResponse else { - throw KiroAPIError.invalidResponse - } - - switch httpResponse.statusCode { - case 200: - let decoder = JSONDecoder() - decoder.dateDecodingStrategy = .secondsSince1970 - let usageResponse = try decoder.decode(KiroUsageLimitsResponse.self, from: data) - return try usageResponse.toSnapshot() - case 403: - // Try next region - throw URLError(.badServerResponse) - case 401: - throw KiroAPIError.authenticationFailed - case 429: - throw KiroAPIError.rateLimited - default: - let body = String(data: data, encoding: .utf8) ?? "" - throw KiroAPIError.httpError(statusCode: httpResponse.statusCode, body: body) - } - } - - /// Region candidate ordering: prefer us-east-1, fallback to eu-central-1 - private static func regionCandidates(for ssoRegion: String?) -> [String] { - guard let ssoRegion else { return apiRegions } - if ssoRegion == "eu-central-1" || ssoRegion.hasPrefix("eu-") { - return ["eu-central-1", "us-east-1"] - } - return apiRegions - } - - /// Generate a stable machine ID from the device's hardware UUID - private static func generateMachineId() -> String { - #if canImport(Darwin) - let platformExpert = IOServiceGetMatchingService( - kIOMainPortDefault, - IOServiceMatching("IOPlatformExpertDevice")) - defer { IOObjectRelease(platformExpert) } - - guard platformExpert != 0, - let uuidCF = IORegistryEntryCreateCFProperty( - platformExpert, - kIOPlatformUUIDKey as CFString, - kCFAllocatorDefault, 0) - else { - return UUID().uuidString.replacingOccurrences(of: "-", with: "") - } - - let uuid = uuidCF.takeRetainedValue() as? String ?? UUID().uuidString - return uuid.replacingOccurrences(of: "-", with: "") - #else - return UUID().uuidString.replacingOccurrences(of: "-", with: "") - #endif - } -} - -// MARK: - Supporting Types - -private struct KiroCLICredentials { - let accessToken: String - let region: String? - let startURL: String? -} - -enum KiroAPIError: Error, LocalizedError { - case credentialsNotFound - case invalidURL - case invalidResponse - case authenticationFailed - case rateLimited - case allRegionsFailed - case httpError(statusCode: Int, body: String) - - var errorDescription: String? { - switch self { - case .credentialsNotFound: - return "Kiro credentials not found. Please run 'kiro-cli login' first." - case .invalidURL: - return "Invalid API URL" - case .invalidResponse: - return "Invalid response from Kiro API" - case .authenticationFailed: - return "Kiro authentication failed. Please run 'kiro-cli login' again." - case .rateLimited: - return "Kiro API rate limited. Please try again later." - case .allRegionsFailed: - return "All Kiro API regions failed" - case .httpError(let statusCode, let body): - return "Kiro API error: \(statusCode) \(body)" - } - } -} - -// MARK: - API Response Models - -struct KiroUsageLimitsResponse: Decodable { - let nextDateReset: Double? - let subscriptionInfo: KiroSubscriptionInfo? - let usageBreakdownList: [KiroUsageBreakdown]? - let overageConfiguration: KiroOverageConfiguration? - let userInfo: KiroUserInfo? - - enum CodingKeys: String, CodingKey { - case nextDateReset - case subscriptionInfo - case usageBreakdownList - case overageConfiguration - case userInfo - } - - func toSnapshot() throws -> KiroUsageSnapshot { - let planName = subscriptionInfo?.subscriptionTitle ?? "KIRO FREE" - let email = userInfo?.email - - guard let breakdown = usageBreakdownList?.first else { - // No usage data - might be a managed/enterprise plan - return KiroUsageSnapshot( - planName: planName, - displayPlanName: KiroStatusProbe.displayPlanName(planName), - accountEmail: email, - authMethod: nil, - creditsUsed: 0, - creditsTotal: 0, - creditsPercent: 0, - bonusCreditsUsed: nil, - bonusCreditsTotal: nil, - bonusExpiryDays: nil, - overagesStatus: overageConfiguration?.overageStatus, - overageCreditsUsed: nil, - estimatedOverageCostUSD: nil, - manageURL: nil, - contextUsage: nil, - resetsAt: nextDateReset.map { Date(timeIntervalSince1970: $0) }, - updatedAt: Date()) - } - - let creditsUsed = breakdown.currentUsageWithPrecision ?? Double(breakdown.currentUsage ?? 0) - let creditsTotal = breakdown.usageLimitWithPrecision ?? Double(breakdown.usageLimit ?? 0) - let creditsPercent = creditsTotal > 0 ? (creditsUsed / creditsTotal) * 100.0 : 0 - - // Calculate bonus credits - var bonusUsed: Double? - var bonusTotal: Double? - var bonusExpiry: Int? - - let activeBonuses = (breakdown.bonuses ?? []).filter { $0.status == "ACTIVE" } - if !activeBonuses.isEmpty { - bonusUsed = activeBonuses.reduce(0) { $0 + ($1.currentUsage ?? 0) } - bonusTotal = activeBonuses.reduce(0) { $0 + ($1.usageLimit ?? 0) } - } - - // Overages - let overageEnabled = overageConfiguration?.overageEnabled ?? false - let overagesStatus = overageConfiguration?.overageStatus - - return KiroUsageSnapshot( - planName: planName, - displayPlanName: KiroStatusProbe.displayPlanName(planName), - accountEmail: email, - authMethod: nil, - creditsUsed: creditsUsed, - creditsTotal: creditsTotal, - creditsPercent: creditsPercent, - bonusCreditsUsed: bonusUsed, - bonusCreditsTotal: bonusTotal, - bonusExpiryDays: bonusExpiry, - overagesStatus: overagesStatus, - overageCreditsUsed: nil, - estimatedOverageCostUSD: nil, - manageURL: nil, - contextUsage: nil, - resetsAt: nextDateReset.map { Date(timeIntervalSince1970: $0) }, - updatedAt: Date()) - } -} - -struct KiroSubscriptionInfo: Decodable { - let subscriptionTitle: String? - let overageCapability: String? -} - -struct KiroUsageBreakdown: Decodable { - let currentUsage: Int? - let currentUsageWithPrecision: Double? - let usageLimit: Int? - let usageLimitWithPrecision: Double? - let bonuses: [KiroBonus]? - let freeTrialInfo: KiroFreeTrialInfo? - - enum CodingKeys: String, CodingKey { - case currentUsage - case currentUsageWithPrecision - case usageLimit - case usageLimitWithPrecision - case bonuses - case freeTrialInfo - } - - init(from decoder: Decoder) throws { - let container = try decoder.container(keyedBy: CodingKeys.self) - currentUsage = try container.decodeIfPresent(Int.self, forKey: .currentUsage) - currentUsageWithPrecision = try container.decodeIfPresent(Double.self, forKey: .currentUsageWithPrecision) - usageLimit = try container.decodeIfPresent(Int.self, forKey: .usageLimit) - usageLimitWithPrecision = try container.decodeIfPresent(Double.self, forKey: .usageLimitWithPrecision) - bonuses = try container.decodeIfPresent([KiroBonus].self, forKey: .bonuses) - freeTrialInfo = try container.decodeIfPresent(KiroFreeTrialInfo.self, forKey: .freeTrialInfo) - } -} - -struct KiroBonus: Decodable { - let currentUsage: Double? - let usageLimit: Double? - let status: String? - - enum CodingKeys: String, CodingKey { - case currentUsage - case usageLimit - case status - } -} - -struct KiroFreeTrialInfo: Decodable { - let currentUsage: Int? - let currentUsageWithPrecision: Double? - let usageLimit: Int? - let usageLimitWithPrecision: Double? - let freeTrialExpiry: Double? - let freeTrialStatus: String? - - enum CodingKeys: String, CodingKey { - case currentUsage - case currentUsageWithPrecision - case usageLimit - case usageLimitWithPrecision - case freeTrialExpiry - case freeTrialStatus - } -} - -struct KiroOverageConfiguration: Decodable { - let overageEnabled: Bool? - let overageStatus: String? -} - -struct KiroUserInfo: Decodable { - let email: String? } diff --git a/Sources/CodexBarCore/Providers/Kiro/KiroAPIUsageFetcher.swift b/Sources/CodexBarCore/Providers/Kiro/KiroAPIUsageFetcher.swift new file mode 100644 index 0000000000..6e706d6cc7 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Kiro/KiroAPIUsageFetcher.swift @@ -0,0 +1,616 @@ +import Foundation +#if canImport(CryptoKit) +import CryptoKit +#else +import Crypto +#endif + +struct KiroAPIUsageFetcher: Sendable { + private static let usageAPIKiroVersion = "0.9.2" + private static let apiRegions = ["us-east-1", "eu-central-1"] + private static let idcAmzUserAgent = + "aws-sdk-js/3.738.0 ua/2.1 os/other lang/js md/browser#unknown_unknown api/sso-oidc#3.738.0 m/E KiroIDE" + private static let maxErrorBodyLength = 240 + + let credentialStore: KiroCLICredentialStore + let transport: any ProviderHTTPTransport + + init( + credentialStore: KiroCLICredentialStore = KiroCLICredentialStore(), + transport: any ProviderHTTPTransport = ProviderHTTPClient.shared) + { + self.credentialStore = credentialStore + self.transport = transport + } + + func hasCredentials(allowSocial: Bool) -> Bool { + self.credentialStore.loadCredentials(allowSocial: allowSocial) != nil + } + + func fetchUsage(allowSocial: Bool) async throws -> KiroUsageSnapshot { + guard var credentials = self.credentialStore.loadCredentials(allowSocial: allowSocial) else { + throw KiroAPIError.credentialsNotFound + } + + if credentials.needsRefresh() { + credentials = try await self.refreshCredentials(credentials) + } + + return try await self.fetchUsage(credentials: credentials) + } + + // MARK: - Token refresh + + private func refreshCredentials(_ credentials: KiroCLICredentials) async throws -> KiroCLICredentials { + guard let refreshToken = credentials.refreshToken, !refreshToken.isEmpty else { + throw KiroAPIError.authenticationFailed + } + guard refreshToken.count >= 100, !refreshToken.contains("...") else { + throw KiroAPIError.refreshTokenUnavailable + } + + if credentials.isExternalIDP { + return try await self.refreshExternalIDPToken(credentials, refreshToken: refreshToken) + } + + let authMethod = credentials.canonicalAuthMethod + if authMethod == "idc" { + return try await self.refreshIDCToken(credentials, refreshToken: refreshToken) + } + return try await self.refreshSocialToken(credentials, refreshToken: refreshToken) + } + + private func refreshSocialToken( + _ credentials: KiroCLICredentials, + refreshToken: String) async throws -> KiroCLICredentials + { + let region = credentials.effectiveAuthRegion + let host = "prod.\(region).auth.desktop.kiro.dev" + guard let url = URL(string: "https://\(host)/refreshToken") else { + throw KiroAPIError.invalidURL + } + + let machineID = Self.machineID(for: credentials) + var request = URLRequest(url: url) + request.httpMethod = "POST" + request.setValue("application/json", forHTTPHeaderField: "Content-Type") + request.setValue("application/json, text/plain, */*", forHTTPHeaderField: "Accept") + request.setValue( + "KiroIDE-\(Self.usageAPIKiroVersion)-\(machineID)", + forHTTPHeaderField: "User-Agent") + request.setValue(host, forHTTPHeaderField: "host") + request.httpBody = try JSONSerialization.data(withJSONObject: ["refreshToken": refreshToken]) + + let response = try await self.transport.response(for: request) + try Self.validateRefreshResponse(response) + + let json = try Self.decodeJSON(response.data) + guard let accessToken = Self.string(json, keys: "access_token", "accessToken"), !accessToken.isEmpty else { + throw KiroAPIError.authenticationFailed + } + + return credentials.refreshed( + accessToken: accessToken, + refreshToken: Self.string(json, keys: "refresh_token", "refreshToken") ?? refreshToken, + expiresAt: Self.expiresAt(from: json)) + } + + private func refreshIDCToken( + _ credentials: KiroCLICredentials, + refreshToken: String) async throws -> KiroCLICredentials + { + guard let clientID = credentials.clientID, !clientID.isEmpty else { + throw KiroAPIError.authenticationFailed + } + guard let clientSecret = credentials.clientSecret, !clientSecret.isEmpty else { + throw KiroAPIError.authenticationFailed + } + + let region = credentials.effectiveAuthRegion + let host = "oidc.\(region).amazonaws.com" + guard let url = URL(string: "https://\(host)/token") else { + throw KiroAPIError.invalidURL + } + + var request = URLRequest(url: url) + request.httpMethod = "POST" + request.setValue("application/json", forHTTPHeaderField: "Content-Type") + request.setValue(host, forHTTPHeaderField: "Host") + request.setValue(Self.idcAmzUserAgent, forHTTPHeaderField: "x-amz-user-agent") + request.setValue("node", forHTTPHeaderField: "User-Agent") + request.httpBody = try JSONSerialization.data(withJSONObject: [ + "clientId": clientID, + "clientSecret": clientSecret, + "refreshToken": refreshToken, + "grantType": "refresh_token", + ]) + + let response = try await self.transport.response(for: request) + try Self.validateRefreshResponse(response) + + let json = try Self.decodeJSON(response.data) + guard let accessToken = Self.string(json, keys: "access_token", "accessToken"), !accessToken.isEmpty else { + throw KiroAPIError.authenticationFailed + } + + return credentials.refreshed( + accessToken: accessToken, + refreshToken: Self.string(json, keys: "refresh_token", "refreshToken") ?? refreshToken, + expiresAt: Self.expiresAt(from: json)) + } + + private func refreshExternalIDPToken( + _ credentials: KiroCLICredentials, + refreshToken: String) async throws -> KiroCLICredentials + { + guard let clientID = credentials.clientID, !clientID.isEmpty else { + throw KiroAPIError.authenticationFailed + } + guard let tokenEndpoint = credentials.tokenEndpoint, !tokenEndpoint.isEmpty else { + throw KiroAPIError.authenticationFailed + } + try KiroCLICredentials.validateExternalIDPEndpoint(tokenEndpoint) + + guard let url = URL(string: tokenEndpoint) else { + throw KiroAPIError.invalidURL + } + + var bodyComponents = [ + URLQueryItem(name: "client_id", value: clientID), + URLQueryItem(name: "grant_type", value: "refresh_token"), + URLQueryItem(name: "refresh_token", value: refreshToken), + ] + if let scopes = credentials.scopes, !scopes.isEmpty { + bodyComponents.append(URLQueryItem(name: "scope", value: scopes)) + } + + var request = URLRequest(url: url) + request.httpMethod = "POST" + request.setValue("application/json", forHTTPHeaderField: "Accept") + request.setValue("application/x-www-form-urlencoded", forHTTPHeaderField: "Content-Type") + request.httpBody = bodyComponents + .percentEncoded() + .data(using: .utf8) + + let response = try await self.transport.response(for: request) + try Self.validateRefreshResponse(response) + + let json = try Self.decodeJSON(response.data) + guard let accessToken = Self.string(json, keys: "access_token", "accessToken"), !accessToken.isEmpty else { + throw KiroAPIError.authenticationFailed + } + + return credentials.refreshed( + accessToken: accessToken, + refreshToken: Self.string(json, keys: "refresh_token", "refreshToken") ?? refreshToken, + expiresAt: Self.expiresAt(from: json)) + } + + // MARK: - Usage fetch + + private func fetchUsage(credentials: KiroCLICredentials) async throws -> KiroUsageSnapshot { + let regions = Self.regionCandidates(for: credentials.region) + var lastError: Error? + + for region in regions { + do { + return try await self.fetchUsageFromRegion(region: region, credentials: credentials) + } catch let error as URLError where error.code == .badServerResponse { + lastError = error + continue + } catch { + throw error + } + } + + throw lastError ?? KiroAPIError.allRegionsFailed + } + + private func fetchUsageFromRegion( + region: String, + credentials: KiroCLICredentials) async throws -> KiroUsageSnapshot + { + let host = "q.\(region).amazonaws.com" + let urlString = + "https://\(host)/getUsageLimits?origin=AI_EDITOR&resourceType=AGENTIC_REQUEST&isEmailRequired=true" + guard let url = URL(string: urlString) else { + throw KiroAPIError.invalidURL + } + + let machineID = Self.machineID(for: credentials) + var request = URLRequest(url: url) + request.httpMethod = "GET" + request.timeoutInterval = 30 + + let userAgent = [ + "aws-sdk-js/1.0.0 ua/2.1 os/macos lang/js md/nodejs#20.0.0", + "api/codewhispererruntime#1.0.0 m/N,E KiroIDE-\(Self.usageAPIKiroVersion)-\(machineID)", + ].joined(separator: " ") + let amzUserAgent = "aws-sdk-js/1.0.0 KiroIDE-\(Self.usageAPIKiroVersion)-\(machineID)" + + request.setValue(amzUserAgent, forHTTPHeaderField: "x-amz-user-agent") + request.setValue(userAgent, forHTTPHeaderField: "user-agent") + request.setValue(host, forHTTPHeaderField: "host") + request.setValue(UUID().uuidString, forHTTPHeaderField: "amz-sdk-invocation-id") + request.setValue("attempt=1; max=1", forHTTPHeaderField: "amz-sdk-request") + request.setValue("Bearer \(credentials.accessToken)", forHTTPHeaderField: "Authorization") + request.setValue("close", forHTTPHeaderField: "Connection") + if let tokenType = credentials.tokenTypeHeader { + request.setValue(tokenType, forHTTPHeaderField: "tokentype") + } + + let response = try await self.transport.response( + for: request, + retryPolicy: .transientIdempotent) + + switch response.statusCode { + case 200: + let usageResponse = try JSONDecoder().decode(KiroUsageLimitsResponse.self, from: response.data) + return try usageResponse.toSnapshot() + case 403: + throw URLError(.badServerResponse) + case 401: + throw KiroAPIError.authenticationFailed + case 429: + throw KiroAPIError.rateLimited + default: + throw KiroAPIError.httpError( + statusCode: response.statusCode, + summary: Self.sanitizedResponseBodySummary(response.data)) + } + } + + // MARK: - Helpers + + private static func regionCandidates(for ssoRegion: String?) -> [String] { + guard let ssoRegion else { return self.apiRegions } + if ssoRegion == "eu-central-1" || ssoRegion.hasPrefix("eu-") { + return ["eu-central-1", "us-east-1"] + } + return self.apiRegions + } + + static func machineID(for credentials: KiroCLICredentials) -> String { + if let configured = credentials.machineID, + let normalized = normalizeMachineID(configured) + { + return normalized + } + if let refreshToken = credentials.refreshToken, !refreshToken.isEmpty { + return Self.sha256Hex("KotlinNativeAPI/\(refreshToken)") + } + return Self.sha256Hex("KiroFallback/\(credentials.storageKey)") + } + + static func normalizeMachineID(_ raw: String) -> String? { + let trimmed = raw.trimmingCharacters(in: .whitespacesAndNewlines) + if trimmed.count == 64, trimmed.allSatisfy(\.isHexDigit) { + return trimmed.lowercased() + } + let withoutDashes = trimmed.filter { $0 != "-" } + if withoutDashes.count == 32, withoutDashes.allSatisfy(\.isHexDigit) { + return String(repeating: withoutDashes.lowercased(), count: 2) + } + return nil + } + + private static func sha256Hex(_ input: String) -> String { + let digest = SHA256.hash(data: Data(input.utf8)) + return digest.map { String(format: "%02x", $0) }.joined() + } + + private static func validateRefreshResponse(_ response: ProviderHTTPResponse) throws { + switch response.statusCode { + case 200...299: + return + case 400 where String(data: response.data, encoding: .utf8)?.contains("invalid_grant") == true: + throw KiroAPIError.refreshTokenExpired + case 401: + throw KiroAPIError.authenticationFailed + case 429: + throw KiroAPIError.rateLimited + default: + throw KiroAPIError.httpError( + statusCode: response.statusCode, + summary: self.sanitizedResponseBodySummary(response.data)) + } + } + + private static func decodeJSON(_ data: Data) throws -> [String: Any] { + guard let json = try JSONSerialization.jsonObject(with: data) as? [String: Any] else { + throw KiroAPIError.invalidResponse + } + return json + } + + private static func string(_ json: [String: Any], keys: String...) -> String? { + for key in keys { + guard let value = json[key] as? String else { continue } + let trimmed = value.trimmingCharacters(in: .whitespacesAndNewlines) + if !trimmed.isEmpty { return trimmed } + } + return nil + } + + private static func expiresAt(from json: [String: Any]) -> Date? { + if let expiresIn = json["expires_in"] as? Int ?? json["expiresIn"] as? Int { + return Date().addingTimeInterval(TimeInterval(expiresIn)) + } + if let expiresIn = json["expires_in"] as? Double ?? json["expiresIn"] as? Double { + return Date().addingTimeInterval(expiresIn) + } + if let expiresAt = self.string(json, keys: "expires_at", "expiresAt") { + return Self.parseISO8601Date(expiresAt) + } + return nil + } + + static func sanitizedResponseBodySummary(_ data: Data) -> String { + guard !data.isEmpty else { return "empty body" } + guard let rawBody = String(data: data, encoding: .utf8) else { + return "non-text body (\(data.count) bytes)" + } + + let body = Self.redactSensitiveBodyContent(rawBody) + .replacingOccurrences(of: #"\s+"#, with: " ", options: .regularExpression) + .trimmingCharacters(in: .whitespacesAndNewlines) + guard !body.isEmpty else { return "non-text body (\(data.count) bytes)" } + guard body.count > Self.maxErrorBodyLength else { return body } + let index = body.index(body.startIndex, offsetBy: Self.maxErrorBodyLength) + return "\(body[.. String { + let replacements: [(String, String)] = [ + (#"(?i)(bearer\s+)[A-Za-z0-9._\-]+"#, "$1[REDACTED]"), + ( + #"(?i)(\"(?:api_?key|authorization|token|access_token|refresh_token)\"\s*:\s*\")([^\"]+)(\")"#, + "$1[REDACTED]$3"), + ( + #"(?i)((?:api_?key|authorization|token|access_token|refresh_token)\s*[=:]\s*)([^,\s]+)"#, + "$1[REDACTED]"), + ] + return replacements.reduce(text) { partial, replacement in + partial.replacingOccurrences( + of: replacement.0, + with: replacement.1, + options: .regularExpression) + } + } + + #if DEBUG + static func _machineIDForTesting(_ credentials: KiroCLICredentials) -> String { + self.machineID(for: credentials) + } + + static func _sanitizedResponseBodySummaryForTesting(_ body: String) -> String { + self.sanitizedResponseBodySummary(Data(body.utf8)) + } + #endif +} + +extension KiroCLICredentials { + fileprivate func refreshed(accessToken: String, refreshToken: String, expiresAt: Date?) -> KiroCLICredentials { + KiroCLICredentials( + storageKey: self.storageKey, + accessToken: accessToken, + refreshToken: refreshToken, + expiresAt: expiresAt, + region: self.region, + authRegion: self.authRegion, + startURL: self.startURL, + tokenEndpoint: self.tokenEndpoint, + scopes: self.scopes, + clientID: self.clientID, + clientSecret: self.clientSecret, + authMethod: self.authMethod, + provider: self.provider, + machineID: self.machineID) + } +} + +extension [URLQueryItem] { + fileprivate func percentEncoded() -> String { + var components = URLComponents() + components.queryItems = self + return components.percentEncodedQuery ?? "" + } +} + +extension Character { + fileprivate var isHexDigit: Bool { + ("0"..."9").contains(self) || ("a"..."f").contains(self) || ("A"..."F").contains(self) + } +} + +enum KiroAPIError: Error, LocalizedError, Equatable { + case credentialsNotFound + case invalidURL + case invalidResponse + case authenticationFailed + case refreshTokenUnavailable + case refreshTokenExpired + case rateLimited + case allRegionsFailed + case invalidExternalIDPEndpoint + case httpError(statusCode: Int, summary: String) + + var errorDescription: String? { + switch self { + case .credentialsNotFound: + "Kiro credentials not found. Please run 'kiro-cli login' first." + case .invalidURL: + "Invalid Kiro API URL" + case .invalidResponse: + "Invalid response from Kiro API" + case .authenticationFailed: + "Kiro authentication failed. Please run 'kiro-cli login' again." + case .refreshTokenUnavailable: + "Kiro refresh token is unavailable. Please run 'kiro-cli login' again." + case .refreshTokenExpired: + "Kiro session expired. Please run 'kiro-cli login' again." + case .rateLimited: + "Kiro API rate limited. Please try again later." + case .allRegionsFailed: + "All Kiro API regions failed" + case .invalidExternalIDPEndpoint: + "Kiro external IdP endpoint is not allowed" + case let .httpError(statusCode, summary): + "Kiro API error (\(statusCode)): \(summary)" + } + } +} + +// MARK: - API Response Models + +struct KiroUsageLimitsResponse: Decodable { + let nextDateReset: KiroFlexibleResetDate? + let subscriptionInfo: KiroSubscriptionInfo? + let usageBreakdownList: [KiroUsageBreakdown]? + let overageConfiguration: KiroOverageConfiguration? + let userInfo: KiroUserInfo? + + func toSnapshot() throws -> KiroUsageSnapshot { + let planName = self.subscriptionInfo?.subscriptionTitle ?? "KIRO FREE" + let email = self.userInfo?.email + + guard let breakdown = self.usageBreakdownList?.first else { + return KiroUsageSnapshot( + planName: planName, + displayPlanName: KiroStatusProbe.displayPlanName(planName), + accountEmail: email, + authMethod: nil, + creditsUsed: 0, + creditsTotal: 0, + creditsPercent: 0, + bonusCreditsUsed: nil, + bonusCreditsTotal: nil, + bonusExpiryDays: nil, + overagesStatus: self.overageConfiguration?.overageStatus, + overageCreditsUsed: nil, + estimatedOverageCostUSD: nil, + manageURL: nil, + contextUsage: nil, + resetsAt: self.nextDateReset?.date, + updatedAt: Date()) + } + + let creditsUsed = breakdown.currentUsageWithPrecision ?? Double(breakdown.currentUsage ?? 0) + let creditsTotal = breakdown.usageLimitWithPrecision ?? Double(breakdown.usageLimit ?? 0) + let creditsPercent = creditsTotal > 0 ? (creditsUsed / creditsTotal) * 100.0 : 0 + + var bonusUsed: Double? + var bonusTotal: Double? + let activeBonuses = (breakdown.bonuses ?? []).filter { $0.status == "ACTIVE" } + if !activeBonuses.isEmpty { + bonusUsed = activeBonuses.reduce(0) { $0 + ($1.currentUsage ?? 0) } + bonusTotal = activeBonuses.reduce(0) { $0 + ($1.usageLimit ?? 0) } + } + + return KiroUsageSnapshot( + planName: planName, + displayPlanName: KiroStatusProbe.displayPlanName(planName), + accountEmail: email, + authMethod: nil, + creditsUsed: creditsUsed, + creditsTotal: creditsTotal, + creditsPercent: creditsPercent, + bonusCreditsUsed: bonusUsed, + bonusCreditsTotal: bonusTotal, + bonusExpiryDays: nil, + overagesStatus: self.overageConfiguration?.overageStatus, + overageCreditsUsed: nil, + estimatedOverageCostUSD: nil, + manageURL: nil, + contextUsage: nil, + resetsAt: self.nextDateReset?.date, + updatedAt: Date()) + } +} + +struct KiroFlexibleResetDate: Decodable { + let date: Date? + + init(from decoder: Decoder) throws { + let container = try decoder.singleValueContainer() + if container.decodeNil() { + self.date = nil + return + } + if let timestamp = try? container.decode(Double.self) { + self.date = Date(timeIntervalSince1970: timestamp) + return + } + if let stringValue = try? container.decode(String.self) { + self.date = KiroAPIUsageFetcher.parseISO8601Date(stringValue) + return + } + self.date = nil + } +} + +struct KiroSubscriptionInfo: Decodable { + let subscriptionTitle: String? + let overageCapability: String? +} + +struct KiroUsageBreakdown: Decodable { + let currentUsage: Int? + let currentUsageWithPrecision: Double? + let usageLimit: Int? + let usageLimitWithPrecision: Double? + let bonuses: [KiroBonus]? + let freeTrialInfo: KiroFreeTrialInfo? + + enum CodingKeys: String, CodingKey { + case currentUsage + case currentUsageWithPrecision + case usageLimit + case usageLimitWithPrecision + case bonuses + case freeTrialInfo + } + + init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + self.currentUsage = try container.decodeIfPresent(Int.self, forKey: .currentUsage) + self.currentUsageWithPrecision = try container.decodeIfPresent(Double.self, forKey: .currentUsageWithPrecision) + self.usageLimit = try container.decodeIfPresent(Int.self, forKey: .usageLimit) + self.usageLimitWithPrecision = try container.decodeIfPresent(Double.self, forKey: .usageLimitWithPrecision) + self.bonuses = try container.decodeIfPresent([KiroBonus].self, forKey: .bonuses) + self.freeTrialInfo = try container.decodeIfPresent(KiroFreeTrialInfo.self, forKey: .freeTrialInfo) + } +} + +struct KiroBonus: Decodable { + let currentUsage: Double? + let usageLimit: Double? + let status: String? +} + +struct KiroFreeTrialInfo: Decodable { + let currentUsage: Int? + let currentUsageWithPrecision: Double? + let usageLimit: Int? + let usageLimitWithPrecision: Double? + let freeTrialExpiry: Double? + let freeTrialStatus: String? +} + +struct KiroOverageConfiguration: Decodable { + let overageEnabled: Bool? + let overageStatus: String? +} + +struct KiroUserInfo: Decodable { + let email: String? +} + +extension KiroAPIUsageFetcher { + static func parseISO8601Date(_ value: String) -> Date? { + let formatter = ISO8601DateFormatter() + if let date = formatter.date(from: value) { return date } + formatter.formatOptions = [.withInternetDateTime, .withFractionalSeconds] + return formatter.date(from: value) + } +} diff --git a/Sources/CodexBarCore/Providers/Kiro/KiroCLICredentialStore.swift b/Sources/CodexBarCore/Providers/Kiro/KiroCLICredentialStore.swift new file mode 100644 index 0000000000..cc8e132d0f --- /dev/null +++ b/Sources/CodexBarCore/Providers/Kiro/KiroCLICredentialStore.swift @@ -0,0 +1,224 @@ +import Foundation +import SQLite3 + +/// Reads Kiro CLI OAuth material from the local SQLite auth store. +/// +/// ponytail: injectable `databaseURL` keeps filesystem access out of fetch strategies and +/// mirrors `OpenCodeGoLocalUsageReader`. +struct KiroCLICredentialStore: Sendable { + static let enterpriseStorageKeySuffixes = ["oidc:token", "odic:token"] + static let socialStorageKey = "kirocli:social:token" + + private static let storageKeys = [ + "kirocli:odic:token", + "kirocli:oidc:token", + socialStorageKey, + ] + + let databaseURL: URL + + init(homeDirectory: URL = FileManager.default.homeDirectoryForCurrentUser) { + self.databaseURL = homeDirectory + .appendingPathComponent("Library/Application Support/kiro-cli", isDirectory: true) + .appendingPathComponent("data.sqlite3", isDirectory: false) + } + + init(databaseURL: URL) { + self.databaseURL = databaseURL + } + + func loadCredentials(allowSocial: Bool) -> KiroCLICredentials? { + guard FileManager.default.fileExists(atPath: self.databaseURL.path) else { + return nil + } + + var db: OpaquePointer? + guard sqlite3_open_v2(self.databaseURL.path, &db, SQLITE_OPEN_READONLY, nil) == SQLITE_OK else { + sqlite3_close(db) + return nil + } + defer { sqlite3_close(db) } + + for key in Self.storageKeys { + let isSocial = key == Self.socialStorageKey + if isSocial, !allowSocial { continue } + if let credentials = Self.queryCredentials(db: db, storageKey: key) { + if !allowSocial, !credentials.isEnterpriseAffected { continue } + return credentials + } + } + return nil + } + + private static func queryCredentials(db: OpaquePointer?, storageKey: String) -> KiroCLICredentials? { + guard let db else { return nil } + let query = "SELECT value FROM auth_kv WHERE key = ?" + var stmt: OpaquePointer? + guard sqlite3_prepare_v2(db, query, -1, &stmt, nil) == SQLITE_OK else { + sqlite3_finalize(stmt) + return nil + } + defer { sqlite3_finalize(stmt) } + + sqlite3_bind_text(stmt, 1, storageKey, -1, unsafeBitCast(-1, to: sqlite3_destructor_type.self)) + + guard sqlite3_step(stmt) == SQLITE_ROW, + let blob = sqlite3_column_text(stmt, 0) + else { + return nil + } + + let jsonString = String(cString: blob) + guard let data = jsonString.data(using: .utf8), + let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] + else { + return nil + } + + guard let accessToken = Self.string(json, keys: "access_token", "accessToken"), + !accessToken.isEmpty + else { + return nil + } + + return KiroCLICredentials( + storageKey: storageKey, + accessToken: accessToken, + refreshToken: Self.string(json, keys: "refresh_token", "refreshToken"), + expiresAt: Self.parseExpires(json), + region: Self.string(json, keys: "region"), + authRegion: Self.string(json, keys: "auth_region", "authRegion"), + startURL: Self.string(json, keys: "start_url", "startUrl"), + tokenEndpoint: Self.string(json, keys: "token_endpoint", "tokenEndpoint"), + scopes: Self.string(json, keys: "scopes", "scope"), + clientID: Self.string(json, keys: "client_id", "clientId"), + clientSecret: Self.string(json, keys: "client_secret", "clientSecret"), + authMethod: Self.string(json, keys: "auth_method", "authMethod"), + provider: Self.string(json, keys: "provider"), + machineID: Self.string(json, keys: "machine_id", "machineId")) + } + + private static func string(_ json: [String: Any], keys: String...) -> String? { + for key in keys { + guard let value = json[key] as? String else { continue } + let trimmed = value.trimmingCharacters(in: .whitespacesAndNewlines) + if !trimmed.isEmpty { return trimmed } + } + return nil + } + + private static func parseExpires(_ json: [String: Any]) -> Date? { + if let expiresAt = self.string(json, keys: "expires_at", "expiresAt") { + return parseISO8601Date(expiresAt) + } + if let expiresIn = json["expires_in"] as? Int ?? json["expiresIn"] as? Int { + return Date().addingTimeInterval(TimeInterval(expiresIn)) + } + if let expiresIn = json["expires_in"] as? Double ?? json["expiresIn"] as? Double { + return Date().addingTimeInterval(expiresIn) + } + return nil + } +} + +struct KiroCLICredentials: Sendable, Equatable { + let storageKey: String + let accessToken: String + let refreshToken: String? + let expiresAt: Date? + let region: String? + let authRegion: String? + let startURL: String? + let tokenEndpoint: String? + let scopes: String? + let clientID: String? + let clientSecret: String? + let authMethod: String? + let provider: String? + let machineID: String? + + var canonicalAuthMethod: String { + Self.canonicalizeAuthMethod(self.authMethod, tokenEndpoint: self.tokenEndpoint) + } + + var isExternalIDP: Bool { + self.canonicalAuthMethod == "external_idp" + } + + var isEnterpriseAffected: Bool { + if KiroCLICredentialStore.enterpriseStorageKeySuffixes.contains(where: { self.storageKey.hasSuffix($0) }) { + return true + } + let method = self.canonicalAuthMethod + if method == "idc" || method == "external_idp" { return true } + if let startURL = self.startURL, !startURL.isEmpty { return true } + if let tokenEndpoint = self.tokenEndpoint, !tokenEndpoint.isEmpty { return true } + return false + } + + var tokenTypeHeader: String? { + if self.canonicalAuthMethod == "api_key" { return "API_KEY" } + if self.isExternalIDP { return "EXTERNAL_IDP" } + return nil + } + + var effectiveAuthRegion: String { + self.authRegion ?? self.region ?? "us-east-1" + } + + func needsRefresh(now: Date = Date(), buffer: TimeInterval = 5 * 60) -> Bool { + guard let expiresAt else { return true } + return expiresAt.timeIntervalSince(now) <= buffer + } + + static func canonicalizeAuthMethod(_ raw: String?, tokenEndpoint: String?) -> String { + let trimmed = raw?.trimmingCharacters(in: .whitespacesAndNewlines) ?? "" + let normalized: String = switch trimmed.lowercased() { + case "builder-id", "iam": "idc" + case "api_key", "apikey": "api_key" + case "external_idp", "azuread", "azure", "entra", "entra-id", "microsoft", "m365", "office365", "external": + "external_idp" + case "": "" + default: trimmed.lowercased() + } + + if normalized == "external_idp" { return "external_idp" } + if let tokenEndpoint, !tokenEndpoint.isEmpty { return "external_idp" } + return normalized.isEmpty ? "social" : normalized + } + + static func validateExternalIDPEndpoint(_ rawURL: String) throws { + guard let url = URL(string: rawURL.trimmingCharacters(in: .whitespacesAndNewlines)), + let host = url.host?.lowercased(), + !host.isEmpty + else { + throw KiroAPIError.invalidExternalIDPEndpoint + } + + guard url.scheme?.lowercased() == "https" else { + throw KiroAPIError.invalidExternalIDPEndpoint + } + + if host.contains(":") || host.split(separator: ".").allSatisfy({ $0.allSatisfy(\.isNumber) }) { + throw KiroAPIError.invalidExternalIDPEndpoint + } + + let allowedSuffixes = [ + ".microsoftonline.com", + ".microsoftonline.us", + ".microsoftonline.cn", + ] + guard allowedSuffixes.contains(where: { host.hasSuffix($0) }) else { + throw KiroAPIError.invalidExternalIDPEndpoint + } + } +} + +extension KiroCLICredentialStore { + fileprivate static func parseISO8601Date(_ value: String) -> Date? { + let formatter = ISO8601DateFormatter() + if let date = formatter.date(from: value) { return date } + formatter.formatOptions = [.withInternetDateTime, .withFractionalSeconds] + return formatter.date(from: value) + } +} diff --git a/Tests/CodexBarTests/KiroAPIUsageFetcherTests.swift b/Tests/CodexBarTests/KiroAPIUsageFetcherTests.swift new file mode 100644 index 0000000000..9f82e5d656 --- /dev/null +++ b/Tests/CodexBarTests/KiroAPIUsageFetcherTests.swift @@ -0,0 +1,384 @@ +import Foundation +import SQLite3 +import Testing +@testable import CodexBarCore + +@Suite(.serialized) +struct KiroAPIUsageFetcherTests { + private struct StubClaudeFetcher: ClaudeUsageFetching { + func loadLatestUsage(model _: String) async throws -> ClaudeUsageSnapshot { + throw ClaudeUsageError.parseFailed("stub") + } + + func debugRawProbe(model _: String) async -> String { + "stub" + } + + func detectVersion() -> String? { + nil + } + } + + private func makeFetchContext(sourceMode: ProviderSourceMode) -> ProviderFetchContext { + let env: [String: String] = [:] + return ProviderFetchContext( + runtime: .cli, + sourceMode: sourceMode, + includeCredits: false, + webTimeout: 30, + webDebugDumpHTML: false, + verbose: false, + env: env, + settings: ProviderSettingsSnapshot.make(), + fetcher: UsageFetcher(environment: env), + claudeFetcher: StubClaudeFetcher(), + browserDetection: BrowserDetection(cacheTTL: 0)) + } + + @Test + func `auto mode ignores social credentials`() throws { + let env = try self.makeCredentialEnvironment( + key: KiroCLICredentialStore.socialStorageKey, + json: Self.socialCredentialJSON) + defer { try? FileManager.default.removeItem(at: env.root) } + + let store = KiroCLICredentialStore(databaseURL: env.databaseURL) + #expect(store.loadCredentials(allowSocial: false) == nil) + #expect(store.loadCredentials(allowSocial: true) != nil) + } + + @Test + func `credential keys prefer enterprise oidc before social`() throws { + let env = try self.makeCredentialEnvironment( + entries: [ + (KiroCLICredentialStore.socialStorageKey, Self.socialCredentialJSON), + ("kirocli:oidc:token", Self.idcCredentialJSON), + ]) + defer { try? FileManager.default.removeItem(at: env.root) } + + let store = KiroCLICredentialStore(databaseURL: env.databaseURL) + let credentials = store.loadCredentials(allowSocial: false) + #expect(credentials?.storageKey == "kirocli:oidc:token") + #expect(credentials?.isEnterpriseAffected == true) + } + + @Test + func `odic typo key remains supported`() throws { + let env = try self.makeCredentialEnvironment( + key: "kirocli:odic:token", + json: Self.idcCredentialJSON) + defer { try? FileManager.default.removeItem(at: env.root) } + + let credentials = KiroCLICredentialStore(databaseURL: env.databaseURL) + .loadCredentials(allowSocial: false) + #expect(credentials?.storageKey == "kirocli:odic:token") + } + + @Test + func `external idp auth method sends tokentype header`() async throws { + let env = try self.makeCredentialEnvironment( + key: "kirocli:oidc:token", + json: Self.externalIDPCredentialJSON) + defer { try? FileManager.default.removeItem(at: env.root) } + + let transport = ProviderHTTPTransportStub { request in + #expect(request.value(forHTTPHeaderField: "tokentype") == "EXTERNAL_IDP") + #expect(request.url?.absoluteString.contains("profileArn") == false) + return try Self.httpResponse(data: Self.usageFixture(reset: 1_700_000_000), statusCode: 200) + } + + let fetcher = KiroAPIUsageFetcher( + credentialStore: KiroCLICredentialStore(databaseURL: env.databaseURL), + transport: transport) + let snapshot = try await fetcher.fetchUsage(allowSocial: false) + #expect(snapshot.planName == "KIRO POWER") + #expect(snapshot.creditsUsed == 12.5) + } + + @Test + func `idc credentials do not send tokentype header`() async throws { + let env = try self.makeCredentialEnvironment( + key: "kirocli:oidc:token", + json: Self.idcCredentialJSON) + defer { try? FileManager.default.removeItem(at: env.root) } + + let transport = ProviderHTTPTransportStub { request in + #expect(request.value(forHTTPHeaderField: "tokentype") == nil) + return try Self.httpResponse(data: Self.usageFixture(reset: "2026-07-01T00:00:00.000Z"), statusCode: 200) + } + + let fetcher = KiroAPIUsageFetcher( + credentialStore: KiroCLICredentialStore(databaseURL: env.databaseURL), + transport: transport) + let snapshot = try await fetcher.fetchUsage(allowSocial: false) + #expect(snapshot.resetsAt != nil) + } + + @Test + func `region fallback tries second endpoint after 403`() async throws { + let env = try self.makeCredentialEnvironment( + key: "kirocli:oidc:token", + json: Self.idcCredentialJSON) + defer { try? FileManager.default.removeItem(at: env.root) } + + let transport = ProviderHTTPTransportStub { request in + guard let host = request.url?.host else { + return try Self.httpResponse(data: Data(), statusCode: 500) + } + if host == "q.us-east-1.amazonaws.com" { + return try Self.httpResponse(data: Data("forbidden".utf8), statusCode: 403) + } + return try Self.httpResponse(data: Self.usageFixture(reset: 1_700_000_000), statusCode: 200) + } + + let fetcher = KiroAPIUsageFetcher( + credentialStore: KiroCLICredentialStore(databaseURL: env.databaseURL), + transport: transport) + let snapshot = try await fetcher.fetchUsage(allowSocial: false) + #expect(snapshot.creditsTotal == 50) + } + + @Test + func `expired idc credentials refresh before usage fetch`() async throws { + let env = try self.makeCredentialEnvironment( + key: "kirocli:oidc:token", + json: Self.expiredIDCCredentialJSON) + defer { try? FileManager.default.removeItem(at: env.root) } + + let transport = ProviderHTTPTransportStub { request in + if request.url?.host?.hasPrefix("oidc.") == true { + let body = """ + {"accessToken":"fresh-token","refreshToken":"refresh-abc","expiresIn":3600} + """ + return try Self.httpResponse(data: Data(body.utf8), statusCode: 200) + } + #expect(request.value(forHTTPHeaderField: "Authorization") == "Bearer fresh-token") + return try Self.httpResponse(data: Self.usageFixture(reset: 1_700_000_000), statusCode: 200) + } + + let fetcher = KiroAPIUsageFetcher( + credentialStore: KiroCLICredentialStore(databaseURL: env.databaseURL), + transport: transport) + _ = try await fetcher.fetchUsage(allowSocial: false) + } + + @Test + func `refresh failure surfaces authentication error`() async throws { + let env = try self.makeCredentialEnvironment( + key: "kirocli:oidc:token", + json: Self.expiredIDCCredentialJSON) + defer { try? FileManager.default.removeItem(at: env.root) } + + let transport = ProviderHTTPTransportStub { request in + if request.url?.host?.hasPrefix("oidc.") == true { + return try Self.httpResponse(data: Data("{\"error\":\"invalid_grant\"}".utf8), statusCode: 400) + } + return try Self.httpResponse(data: Data(), statusCode: 500) + } + + let fetcher = KiroAPIUsageFetcher( + credentialStore: KiroCLICredentialStore(databaseURL: env.databaseURL), + transport: transport) + + await #expect(throws: KiroAPIError.refreshTokenExpired) { + _ = try await fetcher.fetchUsage(allowSocial: false) + } + } + + @Test + func `machine id derives from refresh token instead of hardware uuid`() { + let credentials = KiroCLICredentials( + storageKey: "kirocli:oidc:token", + accessToken: "token", + refreshToken: "refresh-token-value", + expiresAt: nil, + region: "us-east-1", + authRegion: nil, + startURL: nil, + tokenEndpoint: nil, + scopes: nil, + clientID: nil, + clientSecret: nil, + authMethod: "idc", + provider: nil, + machineID: nil) + + let machineID = KiroAPIUsageFetcher._machineIDForTesting(credentials) + #expect(machineID.count == 64) + #expect(machineID == KiroAPIUsageFetcher._machineIDForTesting(credentials)) + } + + @Test + func `http errors redact sensitive response bodies`() { + let summary = KiroAPIUsageFetcher._sanitizedResponseBodySummaryForTesting( + #"{"access_token":"secret-token","message":"bad request"}"#) + #expect(summary.contains("secret-token") == false) + #expect(summary.contains("[REDACTED]")) + } + + @Test + func `provider descriptor respects explicit source modes`() async { + let descriptor = KiroProviderDescriptor.descriptor + let autoContext = self.makeFetchContext(sourceMode: .auto) + + let apiStrategies = await descriptor.fetchPlan.pipeline.resolveStrategies( + self.makeFetchContext(sourceMode: .api)) + #expect(apiStrategies.count == 1) + #expect(apiStrategies[0].id == "kiro.api") + + let cliStrategies = await descriptor.fetchPlan.pipeline.resolveStrategies( + self.makeFetchContext(sourceMode: .cli)) + #expect(cliStrategies.count == 1) + #expect(cliStrategies[0].id == "kiro.cli") + + let autoStrategies = await descriptor.fetchPlan.pipeline.resolveStrategies(autoContext) + #expect(autoStrategies.count == 2) + #expect(autoStrategies[0].id == "kiro.api") + #expect(autoStrategies[1].id == "kiro.cli") + } + + @Test + func `api strategy only falls back in auto mode`() { + let strategy = KiroAPIFetchStrategy() + let autoContext = self.makeFetchContext(sourceMode: .auto) + let apiContext = self.makeFetchContext(sourceMode: .api) + + #expect(strategy.shouldFallback(on: KiroAPIError.authenticationFailed, context: autoContext)) + #expect(!strategy.shouldFallback(on: KiroAPIError.authenticationFailed, context: apiContext)) + } + + // MARK: - Fixtures + + private static let socialCredentialJSON = """ + {"access_token":"social-token","refresh_token":"\(String(repeating: "s", count: 120))","auth_method":"social"} + """ + + private static let idcCredentialJSON = """ + { + "access_token": "idc-token", + "refresh_token": "\(String(repeating: "r", count: 120))", + "auth_method": "idc", + "region": "us-east-1", + "client_id": "client-id", + "client_secret": "client-secret", + "expires_at": "2099-01-01T00:00:00Z" + } + """ + + private static let expiredIDCCredentialJSON = """ + { + "access_token": "stale-token", + "refresh_token": "\(String(repeating: "r", count: 120))", + "auth_method": "idc", + "region": "us-east-1", + "client_id": "client-id", + "client_secret": "client-secret", + "expires_at": "2020-01-01T00:00:00Z" + } + """ + + private static let externalIDPCredentialJSON = """ + { + "access_token": "external-token", + "refresh_token": "\(String(repeating: "e", count: 120))", + "auth_method": "external_idp", + "client_id": "client-id", + "token_endpoint": "https://login.microsoftonline.com/tenant/oauth2/v2.0/token", + "scopes": "openid offline_access", + "expires_at": "2099-01-01T00:00:00Z" + } + """ + + private static func usageFixture(reset: Double) -> Data { + self.usageFixture(resetJSON: String(reset)) + } + + private static func usageFixture(reset: String) -> Data { + self.usageFixture(resetJSON: "\"\(reset)\"") + } + + private static func usageFixture(resetJSON: String) -> Data { + let body = """ + { + "nextDateReset": \(resetJSON), + "subscriptionInfo": { "subscriptionTitle": "KIRO POWER" }, + "usageBreakdownList": [{ + "currentUsage": 12, + "currentUsageWithPrecision": 12.5, + "usageLimit": 50, + "usageLimitWithPrecision": 50.0 + }] + } + """ + return Data(body.utf8) + } + + private static func httpResponse(data: Data, statusCode: Int) throws -> (Data, URLResponse) { + let response = HTTPURLResponse( + url: URL(string: "https://q.us-east-1.amazonaws.com/getUsageLimits")!, + statusCode: statusCode, + httpVersion: "HTTP/1.1", + headerFields: nil)! + return (data, response) + } + + private struct CredentialEnvironment { + let root: URL + let databaseURL: URL + } + + private func makeCredentialEnvironment( + key: String, + json: String) throws -> CredentialEnvironment + { + try self.makeCredentialEnvironment(entries: [(key, json)]) + } + + private func makeCredentialEnvironment(entries: [(String, String)]) throws -> CredentialEnvironment { + let root = FileManager.default.temporaryDirectory + .appendingPathComponent("KiroAPIUsageFetcherTests-\(UUID().uuidString)", isDirectory: true) + try FileManager.default.createDirectory(at: root, withIntermediateDirectories: true) + let databaseURL = root.appendingPathComponent("data.sqlite3") + + var db: OpaquePointer? + guard sqlite3_open(databaseURL.path, &db) == SQLITE_OK else { + throw CredentialTestError.sqliteOpen + } + defer { sqlite3_close(db) } + + let createSQL = "CREATE TABLE auth_kv (key TEXT PRIMARY KEY, value TEXT NOT NULL);" + guard sqlite3_exec(db, createSQL, nil, nil, nil) == SQLITE_OK else { + throw CredentialTestError.sqliteExec + } + + for (key, json) in entries { + var stmt: OpaquePointer? + guard sqlite3_prepare_v2( + db, + "INSERT INTO auth_kv(key, value) VALUES(?, ?);", + -1, + &stmt, + nil) == SQLITE_OK + else { + throw CredentialTestError.sqlitePrepare + } + defer { sqlite3_finalize(stmt) } + + let transient = unsafeBitCast(-1, to: sqlite3_destructor_type.self) + sqlite3_bind_text(stmt, 1, key, -1, transient) + sqlite3_bind_text(stmt, 2, json, -1, transient) + guard sqlite3_step(stmt) == SQLITE_DONE else { + throw CredentialTestError.sqliteStep + } + } + + return CredentialEnvironment(root: root, databaseURL: databaseURL) + } + + private enum CredentialTestError: Error { + case sqliteOpen + case sqliteExec + case sqlitePrepare + case sqliteStep + } +}