diff --git a/Sources/CodexBar/CostHistoryChartMenuView.swift b/Sources/CodexBar/CostHistoryChartMenuView.swift index 4d2e99b525..219eff00c6 100644 --- a/Sources/CodexBar/CostHistoryChartMenuView.swift +++ b/Sources/CodexBar/CostHistoryChartMenuView.swift @@ -2,6 +2,15 @@ import Charts import CodexBarCore import SwiftUI +private func claudeCodeModelProviderText(_ provider: CostUsageAttribution.ModelProvider) -> String { + switch provider { + case .openAI: "OpenAI model via Claude Code" + case .anthropic: "Anthropic model via Claude Code" + case .google: "Google model via Claude Code" + case .unknown: "Unknown model provider via Claude Code" + } +} + @MainActor struct CostHistoryChartMenuView: View { typealias DailyEntry = CostUsageDailyReport.Entry @@ -577,7 +586,7 @@ struct CostHistoryChartMenuView: View { } private static func hasModeSubtitle(_ item: CostUsageDailyReport.ModelBreakdown) -> Bool { - item.standardCostUSD != nil || item.priorityCostUSD != nil + item.attribution != nil || item.standardCostUSD != nil || item.priorityCostUSD != nil } private static func detailRowsViewportHeight(rowCount: Int, rowHeight: CGFloat) -> CGFloat { @@ -801,26 +810,6 @@ struct CostHistoryChartMenuView: View { } } - static func orderedBreakdownItems( - _ breakdown: [CostUsageDailyReport.ModelBreakdown]) -> [CostUsageDailyReport.ModelBreakdown] - { - breakdown.sorted { lhs, rhs in - let lCost = lhs.costUSD ?? -1 - let rCost = rhs.costUSD ?? -1 - if lCost != rCost { - return lCost > rCost - } - - let lTokens = lhs.totalTokens ?? -1 - let rTokens = rhs.totalTokens ?? -1 - if lTokens != rTokens { - return lTokens > rTokens - } - - return lhs.modelName > rhs.modelName - } - } - static func detailViewportRowCount(itemCount: Int) -> Int { min(max(itemCount, 0), self.maxVisibleDetailLines) } @@ -843,6 +832,19 @@ struct CostHistoryChartMenuView: View { private func modelBreakdownModeSubtitle(_ item: CostUsageDailyReport.ModelBreakdown) -> String? { var parts: [String] = [] + if let attribution = item.attribution { + switch attribution.route { + case .cliProxyAPI: + let route = if let upstream = attribution.upstream { + "\(upstream.displayName) · CLIProxyAPI via Claude Code" + } else { + "CLIProxyAPI via Claude Code" + } + parts.append(route) + case .unknown: + parts.append(claudeCodeModelProviderText(attribution.modelProvider)) + } + } if let standardCost = item.standardCostUSD { var standardPart = "Std \(self.costString(standardCost))" if let standardTokens = item.standardTokens { @@ -880,6 +882,32 @@ struct CostHistoryChartMenuView: View { } extension CostHistoryChartMenuView { + static func orderedBreakdownItems( + _ breakdown: [CostUsageDailyReport.ModelBreakdown]) -> [CostUsageDailyReport.ModelBreakdown] + { + breakdown.sorted { lhs, rhs in + let lCost = lhs.costUSD ?? -1 + let rCost = rhs.costUSD ?? -1 + if lCost != rCost { + return lCost > rCost + } + + let lTokens = lhs.totalTokens ?? -1 + let rTokens = rhs.totalTokens ?? -1 + if lTokens != rTokens { + return lTokens > rTokens + } + + if lhs.modelName != rhs.modelName { + return lhs.modelName > rhs.modelName + } + + let lhsAttribution = lhs.attribution?.deterministicSortKey ?? "" + let rhsAttribution = rhs.attribution?.deterministicSortKey ?? "" + return lhsAttribution > rhsAttribution + } + } + struct RenderFingerprint: Equatable { let currencyCode: String let historyDays: Int @@ -901,6 +929,7 @@ extension CostHistoryChartMenuView { struct VisibleModelBreakdownFingerprint: Equatable { let modelName: String + let attribution: CostUsageAttribution? let costBitPattern: UInt64? let totalTokens: Int? let standardCostBitPattern: UInt64? @@ -980,6 +1009,7 @@ extension CostHistoryChartMenuView { models: session.modelBreakdowns.map { item in VisibleModelBreakdownFingerprint( modelName: item.modelName, + attribution: item.attribution, costBitPattern: item.costUSD.map(\.bitPattern), totalTokens: item.totalTokens, standardCostBitPattern: item.standardCostUSD.map(\.bitPattern), @@ -999,6 +1029,7 @@ extension CostHistoryChartMenuView { modelBreakdowns: self.orderedBreakdownItems(entry.modelBreakdowns ?? []).map { item in VisibleModelBreakdownFingerprint( modelName: item.modelName, + attribution: item.attribution, costBitPattern: item.costUSD.map(\.bitPattern), totalTokens: item.totalTokens, standardCostBitPattern: item.standardCostUSD.map(\.bitPattern), diff --git a/Sources/CodexBar/PreferencesSpendDashboardPane.swift b/Sources/CodexBar/PreferencesSpendDashboardPane.swift index 4dc454e745..31badae386 100644 --- a/Sources/CodexBar/PreferencesSpendDashboardPane.swift +++ b/Sources/CodexBar/PreferencesSpendDashboardPane.swift @@ -27,6 +27,49 @@ func spendDashboardCoverageText(covered: Int, requested: Int) -> String { "\(L("Coverage")): \(codexBarLocalizedInteger(covered)) / \(codexBarLocalizedInteger(requested))" } +func spendDashboardShouldUseAmbientCodexSubscription( + rowID: String, + codexRowCount: Int) -> Bool +{ + rowID != SpendDashboardSource.codexProxySourceID && codexRowCount == 1 +} + +func spendDashboardCodexAccountRowCount( + _ rows: [SpendDashboardModel.ProviderRow]) -> Int +{ + rows.count { + $0.provider == .codex && $0.id != SpendDashboardSource.codexProxySourceID + } +} + +func spendDashboardSubscriptionCount( + _ rows: [SpendDashboardModel.ProviderRow]) -> Int +{ + rows.count { $0.id != SpendDashboardSource.codexProxySourceID } +} + +func spendDashboardModelSourceText( + providerName: String, + attribution: CostUsageAttribution?) -> String +{ + guard let attribution else { return providerName } + switch attribution.route { + case .cliProxyAPI: + if let upstream = attribution.upstream { + return "\(upstream.displayName) · CLIProxyAPI via Claude Code" + } + return "CLIProxyAPI via Claude Code" + case .unknown: + let modelProvider = switch attribution.modelProvider { + case .openAI: "OpenAI model" + case .anthropic: "Anthropic model" + case .google: "Google model" + case .unknown: "Unknown model provider" + } + return "\(providerName) · \(modelProvider) via Claude Code" + } +} + func codexCostCatchUpProgressText(_ activity: CodexCostCatchUpActivity) -> String { if activity.totalBytes > 0 { let processed = ByteCountFormatter.string( @@ -60,11 +103,31 @@ func spendDashboardModelHistoryPresentation( return group.modelHistoryCompleteness == .incomplete ? .partial : .complete } +func spendDashboardCLIProxyAPIConfigurationPresentation( + loadResult: KeychainCacheStore.LoadResult, + currentBaseURL: String, + hasSavedConfiguration: Bool) -> (baseURL: String, hasSavedConfiguration: Bool) +{ + switch loadResult { + case let .found(configuration): + (configuration.baseURL, true) + case .missing, .invalid: + (currentBaseURL, false) + case .temporarilyUnavailable: + (currentBaseURL, hasSavedConfiguration) + } +} + @MainActor struct SpendDashboardPane: View { @Bindable var settings: SettingsStore @Bindable var store: UsageStore @State private var controller: SpendDashboardController + @State private var cliProxyAPIBaseURL = CLIProxyAPIConnectionSettings.defaultBaseURL + @State private var cliProxyAPIManagementKey = "" + @State private var cliProxyAPIHasSavedConfiguration = false + @State private var cliProxyAPIStatus: String? + @State private var cliProxyAPIIsSaving = false @State private var isVisible = false init(settings: SettingsStore, store: UsageStore) { @@ -83,6 +146,7 @@ struct SpendDashboardPane: View { self.header self.codexCostCatchUpPanel self.content + self.cliProxyAPISetup self.provenance self.shareAction } @@ -93,6 +157,7 @@ struct SpendDashboardPane: View { self.isVisible = true self.controller.refreshDateWindow() self.controller.update(configuration: self.configuration) + self.loadCLIProxyAPIConfiguration() if !self.controller.isRefreshing { self.synchronizeCodexCostCatchUp() } @@ -333,6 +398,159 @@ struct SpendDashboardPane: View { } } + private var cliProxyAPISetup: some View { + SpendDashboardPanel { + VStack(alignment: .leading, spacing: 12) { + HStack { + Label("CLIProxyAPI attribution", systemImage: "point.3.connected.trianglepath.dotted") + .font(.headline) + Spacer() + if self.cliProxyAPIHasSavedConfiguration { + Label("Configured", systemImage: "checkmark.circle.fill") + .font(.caption) + .foregroundStyle(.secondary) + } + } + + Text( + "Connect CLIProxyAPI’s management usage queue to identify the exact upstream provider " + + "and whether it used OAuth or an API key. Enter the same plaintext " + + "remote-management secret key configured in CLIProxyAPI.") + .font(.caption) + .foregroundStyle(.secondary) + + Grid(alignment: .leading, horizontalSpacing: 12, verticalSpacing: 10) { + GridRow { + Text("Local server URL") + .foregroundStyle(.secondary) + TextField( + CLIProxyAPIConnectionSettings.defaultBaseURL, + text: self.$cliProxyAPIBaseURL) + .textFieldStyle(.roundedBorder) + } + GridRow { + Text("Management key") + .foregroundStyle(.secondary) + SecureField( + self.cliProxyAPIHasSavedConfiguration ? "Saved — enter to replace" : "Required", + text: self.$cliProxyAPIManagementKey) + .textFieldStyle(.roundedBorder) + } + } + + HStack(spacing: 8) { + Button(self.cliProxyAPIHasSavedConfiguration ? "Save & test" : "Connect & test") { + Task { + await self.saveAndTestCLIProxyAPIConfiguration() + } + } + .disabled(self.cliProxyAPIIsSaving) + + if self.cliProxyAPIHasSavedConfiguration { + Button("Remove", role: .destructive) { + Task { + await self.removeCLIProxyAPIConfiguration() + } + } + .disabled(self.cliProxyAPIIsSaving) + } + + if self.cliProxyAPIIsSaving { + ProgressView() + .controlSize(.small) + } + if let cliProxyAPIStatus { + Text(cliProxyAPIStatus) + .font(.caption) + .foregroundStyle(.secondary) + } + Spacer() + } + + Text( + "Reading the queue removes the returned records from CLIProxyAPI. CodexBar stores a " + + "sanitized local copy for cost attribution and does not retain source, account, " + + "API-key, response-header, or failure-body fields. Records are retained for up to " + + "366 days; Remove and Clear cost cache delete both retained and pending records.") + .font(.caption2) + .foregroundStyle(.tertiary) + } + } + } + + private func loadCLIProxyAPIConfiguration() { + let presentation = spendDashboardCLIProxyAPIConfigurationPresentation( + loadResult: CLIProxyAPIConnectionSettingsStore.loadResult(), + currentBaseURL: self.cliProxyAPIBaseURL, + hasSavedConfiguration: self.cliProxyAPIHasSavedConfiguration) + self.cliProxyAPIBaseURL = presentation.baseURL + self.cliProxyAPIHasSavedConfiguration = presentation.hasSavedConfiguration + } + + private func saveAndTestCLIProxyAPIConfiguration() async { + self.cliProxyAPIIsSaving = true + defer { self.cliProxyAPIIsSaving = false } + + let existingKey = CLIProxyAPIConnectionSettingsStore.load()?.managementKey ?? "" + let enteredKey = self.cliProxyAPIManagementKey.trimmingCharacters(in: .whitespacesAndNewlines) + let configuration = CLIProxyAPIConnectionSettings( + baseURL: self.cliProxyAPIBaseURL, + managementKey: enteredKey.isEmpty ? existingKey : enteredKey) + guard configuration.isConfigured else { + self.cliProxyAPIStatus = "Enter a loopback URL and management key." + return + } + let collectorTask = self.store.stopCLIProxyAPIUsageCollector() + await collectorTask?.value + let saved = await Task.detached(priority: .utility) { + CLIProxyAPIConnectionSettingsStore.save(configuration) + }.value + guard saved else { + self.store.startCLIProxyAPIUsageCollector() + self.cliProxyAPIStatus = "Could not save the management key." + return + } + + self.cliProxyAPIManagementKey = "" + self.cliProxyAPIHasSavedConfiguration = true + switch await self.store.collectCLIProxyAPIUsageNow() { + case .disabled: + self.cliProxyAPIStatus = "Enable Track costs to test." + case .notConfigured: + self.cliProxyAPIStatus = "Configuration was not available." + case let .collected(count): + self.cliProxyAPIStatus = count == 0 + ? "Connected. No queued records." + : "Connected. Imported \(count) records." + self.controller.refresh() + case let .failed(message): + self.cliProxyAPIStatus = "Saved, but test failed: \(message)" + } + await self.store.refreshCLIProxyAPICostAttribution() + self.store.startCLIProxyAPIUsageCollector() + } + + private func removeCLIProxyAPIConfiguration() async { + self.cliProxyAPIIsSaving = true + defer { self.cliProxyAPIIsSaving = false } + + switch await self.store.removeCLIProxyAPIConfiguration() { + case .removed: + self.cliProxyAPIManagementKey = "" + self.cliProxyAPIHasSavedConfiguration = false + self.cliProxyAPIStatus = "Configuration and local telemetry removed." + self.controller.refresh() + case .configurationRemovalFailed: + self.store.startCLIProxyAPIUsageCollector() + self.cliProxyAPIStatus = "Could not remove the saved configuration. Local telemetry was preserved." + case .telemetryCleanupFailed: + self.cliProxyAPIManagementKey = "" + self.cliProxyAPIHasSavedConfiguration = false + self.cliProxyAPIStatus = "Configuration removed, but some local telemetry could not be deleted." + self.controller.refresh() + } + } + private var shareAction: some View { HStack { Spacer() @@ -354,9 +572,8 @@ struct SpendDashboardPane: View { private var subscriptionNames: [String: ShareStatsSubscriptionName] { var names: [String: ShareStatsSubscriptionName] = [:] - let codexRowCount = self.controller.model.groups - .flatMap(\.providers) - .count { $0.provider == .codex } + let codexRowCount = spendDashboardCodexAccountRowCount( + self.controller.model.groups.flatMap(\.providers)) for group in self.controller.model.groups { for row in group.providers { let snapshots: [UsageSnapshot?] = if row.provider == .codex, @@ -366,7 +583,11 @@ struct SpendDashboardPane: View { self.store.codexAccountSnapshots.first { row.id == "codex:\($0.id)" }?.snapshot, - codexRowCount == 1 ? self.store.snapshot(for: .codex) : nil, + spendDashboardShouldUseAmbientCodexSubscription( + rowID: row.id, + codexRowCount: codexRowCount) + ? self.store.snapshot(for: .codex) + : nil, ] } else { [self.store.snapshot(for: row.provider.instanceID)] @@ -439,7 +660,8 @@ private struct SpendCurrencySection: View { value: self.group.totalTokens.map(UsageFormatter.tokenCountString) ?? "—") SpendSummaryValue( title: L("Subscriptions"), - value: codexBarLocalizedInteger(self.group.providers.count)) + value: codexBarLocalizedInteger( + spendDashboardSubscriptionCount(self.group.providers))) Spacer() } } @@ -542,7 +764,11 @@ private struct SpendModelPanel: View { SpendProviderIcon(provider: row.provider) VStack(alignment: .leading, spacing: 2) { Text(row.modelName).lineLimit(1) - Text(row.providerName).font(.caption).foregroundStyle(.secondary) + Text(spendDashboardModelSourceText( + providerName: row.providerName, + attribution: row.attribution)) + .font(.caption) + .foregroundStyle(.secondary) } Spacer() Text(row.totalCost.map { diff --git a/Sources/CodexBar/ShareStatsPayload.swift b/Sources/CodexBar/ShareStatsPayload.swift index ed8a447361..b468f650c4 100644 --- a/Sources/CodexBar/ShareStatsPayload.swift +++ b/Sources/CodexBar/ShareStatsPayload.swift @@ -321,8 +321,9 @@ enum ShareStatsBuilder { subscriptionNames: [String: ShareStatsSubscriptionName] = [:]) -> ShareStatsPayload? { let providers = model.groups.flatMap { group in - group.providers.map { row in - ShareStatsProviderPayload( + group.providers.compactMap { row -> ShareStatsProviderPayload? in + guard row.id != SpendDashboardSource.codexProxySourceID else { return nil } + return ShareStatsProviderPayload( provider: row.provider, providerName: row.displayName, subscriptionName: subscriptionNames[row.id]?.displayName, @@ -338,11 +339,12 @@ enum ShareStatsBuilder { group.models.compactMap { row -> ShareStatsModelPayload? in let estimatedCost = self.finiteCost(row.totalCost) guard let modelName = ShareStatsSanitizer.modelName(row.modelName), + let sharedProvider = self.sharedModelProvider(for: row), row.totalTokens != nil else { return nil } return ShareStatsModelPayload( - provider: row.provider, - providerName: row.providerName, + provider: sharedProvider.provider, + providerName: sharedProvider.name, modelName: modelName, currencyCode: group.currencyCode, totalTokens: row.totalTokens, @@ -393,6 +395,24 @@ enum ShareStatsBuilder { return payload.hasShareableData ? payload : nil } + private static func sharedModelProvider( + for row: SpendDashboardModel.ModelRow) -> (provider: UsageProvider, name: String)? + { + guard row.attribution?.route == .cliProxyAPI else { + return (row.provider, row.providerName) + } + guard let upstream = row.attribution?.upstream else { return nil } + let normalizedProvider = upstream.provider.trimmingCharacters(in: .whitespacesAndNewlines).lowercased() + let provider: UsageProvider? = switch normalizedProvider { + case "anthropic": .claude + case "aistudio", "gemini-interactions", "google": .gemini + case "vertex": .vertexai + default: UsageProvider(rawValue: normalizedProvider) + } + guard let provider else { return nil } + return (provider, ProviderDescriptorRegistry.descriptor(for: provider).metadata.displayName) + } + private static func finiteCost(_ value: Double?) -> Double? { guard let value, value.isFinite, value >= 0 else { return nil } return value diff --git a/Sources/CodexBar/SpendDashboardController.swift b/Sources/CodexBar/SpendDashboardController.swift index e204fb5cb6..a4762fc666 100644 --- a/Sources/CodexBar/SpendDashboardController.swift +++ b/Sources/CodexBar/SpendDashboardController.swift @@ -117,14 +117,31 @@ struct CodexSpendSnapshotLoadContext: Sendable { let includePiSessions: Bool } +struct CodexProxySpendSnapshotLoadContext: Sendable { + let now: Date + let force: Bool + let historyDays: Int + let refreshPricingInBackground: Bool +} + enum SpendDashboardSource { typealias CodexSnapshotLoader = @Sendable (CodexSpendSnapshotLoadContext) async throws -> CostUsageTokenSnapshot + typealias CodexProxySnapshotLoader = @Sendable (CodexProxySpendSnapshotLoadContext) async throws + -> CostUsageTokenSnapshot typealias CachedCodexSnapshotLoader = @Sendable (CodexSpendSnapshotLoadContext) async -> CostUsageTokenSnapshot? typealias CodexCacheRootResolver = @Sendable (CodexSpendScanRequest) -> URL static let scanDays = 30 + static let codexProxySourceID = "codex:cliproxyapi" + private static let isRunningTests: Bool = { + let environment = ProcessInfo.processInfo.environment + return environment["XCTestConfigurationFilePath"] != nil + || environment["TESTING_LIBRARY_VERSION"] != nil + || environment["SWIFT_TESTING"] != nil + || NSClassFromString("XCTestCase") != nil + }() @MainActor static func configuration(settings: SettingsStore, store: UsageStore) -> SpendDashboardConfiguration { @@ -146,12 +163,17 @@ enum SpendDashboardSource { providers: [UsageProvider], codexRequests: [CodexSpendScanRequest]) -> SpendDashboardConfiguration { - SpendDashboardConfiguration( + var codexDisplayNames = self.codexDisplayNamesByID(codexRequests) + if providers.contains(.codex) { + let providerName = store.metadata(for: .codex).displayName + codexDisplayNames[Self.codexProxySourceID] = "\(providerName) · CLIProxyAPI" + } + return SpendDashboardConfiguration( costUsageEnabled: settings.costUsageEnabled, preferredCurrencyCode: settings.preferredCurrencyCode, providerIDs: providers.map(\.rawValue), codexAccountIdentities: codexRequests.map { "\($0.id)|\($0.cacheIdentity)" }, - codexAccountDisplayNames: self.codexDisplayNamesByID(codexRequests), + codexAccountDisplayNames: codexDisplayNames, sourceOwnershipFingerprints: self.sourceOwnershipFingerprints( providers: providers, settings: settings, @@ -253,9 +275,17 @@ enum SpendDashboardSource { } static func load(_ request: SpendDashboardLoadRequest) async -> SpendDashboardLoadResult { - await self.load(request, codexSnapshotLoader: { context in - try await self.loadCodexSnapshot(context) - }) + let codexProxySnapshotLoader: CodexProxySnapshotLoader? = if self.isRunningTests { + nil + } else { + self.loadCodexProxySnapshot + } + return await self.load( + request, + codexSnapshotLoader: { context in + try await self.loadCodexSnapshot(context) + }, + codexProxySnapshotLoader: codexProxySnapshotLoader) } static func loadCached(_ request: SpendDashboardLoadRequest) async -> SpendDashboardLoadResult { @@ -325,6 +355,17 @@ enum SpendDashboardSource { static func load( _ request: SpendDashboardLoadRequest, codexSnapshotLoader: CodexSnapshotLoader) async -> SpendDashboardLoadResult + { + await self.load( + request, + codexSnapshotLoader: codexSnapshotLoader, + codexProxySnapshotLoader: nil) + } + + static func load( + _ request: SpendDashboardLoadRequest, + codexSnapshotLoader: CodexSnapshotLoader, + codexProxySnapshotLoader: CodexProxySnapshotLoader?) async -> SpendDashboardLoadResult { var inputs = request.capturedInputs var failedSourceIDs = request.unavailableSourceIDs @@ -368,6 +409,37 @@ enum SpendDashboardSource { failedSourceIDs.insert(sourceID) } } + if self.shouldLoadCodexProxy(providerIDs: request.configuration.providerIDs), + let codexProxySnapshotLoader + { + do { + let snapshot = try await codexProxySnapshotLoader(CodexProxySpendSnapshotLoadContext( + now: request.now, + force: request.force, + historyDays: Self.scanDays, + refreshPricingInBackground: false)) + try Task.checkCancellation() + if !snapshot.daily.isEmpty { + let providerName = ProviderDescriptorRegistry.descriptor(for: .codex).metadata.displayName + inputs.append(SpendDashboardModel.ProviderInput( + id: Self.codexProxySourceID, + provider: .codex, + displayName: "\(providerName) · CLIProxyAPI", + modelProviderName: providerName, + snapshot: snapshot)) + } + } catch is CancellationError { + failedSourceIDs.formUnion(request.codexRequests.map { "codex:\($0.id)" }) + failedSourceIDs.insert(Self.codexProxySourceID) + invalidatedSourceIDs.insert(Self.codexProxySourceID) + return SpendDashboardLoadResult( + inputs: [], + failedSourceIDs: failedSourceIDs, + invalidatedSourceIDs: invalidatedSourceIDs) + } catch { + failedSourceIDs.insert(Self.codexProxySourceID) + } + } let lateInvalidatedSourceIDs = Set(request.codexRequests.compactMap { account in self.codexAuthFingerprintMatches(account) ? nil @@ -393,7 +465,18 @@ enum SpendDashboardSource { codexHomePath: context.account.homePath, historyDays: context.historyDays, refreshPricingInBackground: context.refreshPricingInBackground, - includePiSessions: context.includePiSessions) + includePiSessions: context.includePiSessions, + includeClaudeProxyUsage: false) + } + + private static func loadCodexProxySnapshot( + _ context: CodexProxySpendSnapshotLoadContext) async throws -> CostUsageTokenSnapshot + { + try await CostUsageFetcher().loadCodexProxyTokenSnapshot( + now: context.now, + forceRefresh: context.force, + historyDays: context.historyDays, + refreshPricingInBackground: context.refreshPricingInBackground) } @MainActor @@ -403,6 +486,11 @@ enum SpendDashboardSource { } } + static func shouldLoadCodexProxy(providerIDs: [String]) -> Bool { + providerIDs.contains(UsageProvider.codex.rawValue) + || providerIDs.contains(UsageProvider.claude.rawValue) + } + @MainActor static func codexRequests(settings: SettingsStore, store: UsageStore) -> [CodexSpendScanRequest] { let accounts = settings.codexVisibleAccountProjection.visibleAccounts @@ -468,6 +556,16 @@ enum SpendDashboardSource { encoder.append(entry.modelBreakdowns?.count) for breakdown in entry.modelBreakdowns ?? [] { encoder.append(breakdown.modelName) + encoder.append(breakdown.attribution?.client.rawValue ?? "") + encoder.append(breakdown.attribution?.route.rawValue ?? "") + encoder.append(breakdown.attribution?.modelProvider.rawValue ?? "") + encoder.append(breakdown.attribution?.upstream?.provider ?? "") + encoder.append(breakdown.attribution?.upstream?.authType.rawValue ?? "") + encoder.append(breakdown.attribution?.upstream?.model ?? "") + encoder.append(breakdown.attribution?.upstream?.executorType ?? "") + for evidence in breakdown.attribution?.evidence ?? [] { + encoder.append(evidence.rawValue) + } encoder.append(breakdown.totalTokens) encoder.append(breakdown.requestCount) encoder.append(breakdown.costUSD) @@ -1021,7 +1119,12 @@ final class SpendDashboardController { let forceFailed = outcome.result.failedSourceIDs let invalidated = outcome.result.invalidatedSourceIDs let barrierFailed = capture.unavailableSourceIDs - let forcedCodexIDs = Set(outcome.request.codexRequests.map { "codex:\($0.id)" }) + var forcedCodexIDs = Set(outcome.request.codexRequests.map { "codex:\($0.id)" }) + if SpendDashboardSource.shouldLoadCodexProxy( + providerIDs: outcome.request.configuration.providerIDs) + { + forcedCodexIDs.insert(SpendDashboardSource.codexProxySourceID) + } let confirmedNonemptyInputs = outcome.confirmedNonemptyInputs let confirmedNonemptyIDs = Set(confirmedNonemptyInputs.map(\.id)) var inputs = capture.capturedInputs.filter { diff --git a/Sources/CodexBar/SpendDashboardModel.swift b/Sources/CodexBar/SpendDashboardModel.swift index 96d3b5db9a..ade5c859b6 100644 --- a/Sources/CodexBar/SpendDashboardModel.swift +++ b/Sources/CodexBar/SpendDashboardModel.swift @@ -35,15 +35,43 @@ struct SpendDashboardModel: Equatable, Sendable { } struct ModelRow: Identifiable, Equatable, Sendable { + struct ID: Hashable, Sendable { + let provider: UsageProvider + let modelName: String + let attribution: CostUsageAttribution? + } + let rank: Int let provider: UsageProvider let providerName: String let modelName: String let totalTokens: Int? let totalCost: Double? + let attribution: CostUsageAttribution? - var id: String { - "\(self.provider.rawValue):\(self.modelName)" + init( + rank: Int, + provider: UsageProvider, + providerName: String, + modelName: String, + totalTokens: Int?, + totalCost: Double?, + attribution: CostUsageAttribution? = nil) + { + self.rank = rank + self.provider = provider + self.providerName = providerName + self.modelName = modelName + self.totalTokens = totalTokens + self.totalCost = totalCost + self.attribution = attribution + } + + var id: ID { + ID( + provider: self.provider, + modelName: self.modelName, + attribution: self.attribution) } } @@ -146,6 +174,7 @@ struct SpendDashboardModel: Equatable, Sendable { private struct ModelKey: Hashable { let provider: UsageProvider let modelName: String + let attribution: CostUsageAttribution? } private struct ModelAccumulator { @@ -311,7 +340,10 @@ struct SpendDashboardModel: Equatable, Sendable { for breakdown in breakdowns { let name = breakdown.modelName.trimmingCharacters(in: .whitespacesAndNewlines) guard !name.isEmpty else { continue } - let key = ModelKey(provider: input.provider, modelName: name) + let key = ModelKey( + provider: input.provider, + modelName: name, + attribution: breakdown.attribution) var aggregate = aggregates[key] ?? ModelAccumulator( providerName: input.modelProviderName, tokens: 0, @@ -350,7 +382,8 @@ struct SpendDashboardModel: Equatable, Sendable { providerName: value.providerName, modelName: key.modelName, totalTokens: value.sawTokens && !value.invalidTokens && !value.overflowedTokens ? value.tokens : nil, - totalCost: value.sawCost && !value.invalidCost && !value.overflowedCost ? value.cost : nil) + totalCost: value.sawCost && !value.invalidCost && !value.overflowedCost ? value.cost : nil, + attribution: key.attribution) } .sorted { lhs, rhs in switch (lhs.totalCost, rhs.totalCost) { @@ -361,7 +394,12 @@ struct SpendDashboardModel: Equatable, Sendable { if lhs.providerName != rhs.providerName { return lhs.providerName < rhs.providerName } - return lhs.modelName < rhs.modelName + if lhs.modelName != rhs.modelName { + return lhs.modelName < rhs.modelName + } + let lhsAttribution = lhs.attribution?.deterministicSortKey ?? "" + let rhsAttribution = rhs.attribution?.deterministicSortKey ?? "" + return lhsAttribution < rhsAttribution } } .enumerated() @@ -372,7 +410,8 @@ struct SpendDashboardModel: Equatable, Sendable { providerName: row.providerName, modelName: row.modelName, totalTokens: row.totalTokens, - totalCost: row.totalCost) + totalCost: row.totalCost, + attribution: row.attribution) } return ModelSummary(rows: rows, completeness: completeness) } diff --git a/Sources/CodexBar/UsageStore+CLIProxyAPI.swift b/Sources/CodexBar/UsageStore+CLIProxyAPI.swift new file mode 100644 index 0000000000..10c0cba321 --- /dev/null +++ b/Sources/CodexBar/UsageStore+CLIProxyAPI.swift @@ -0,0 +1,181 @@ +import CodexBarCore +import Foundation + +struct CLIProxyAPIUsageCollectorState: Equatable { + enum ConfigurationAvailability: Equatable { + case unknown + case available + case unavailable + } + + var configurationAvailability: ConfigurationAvailability = .unknown + var configurationGeneration: String? +} + +@MainActor +extension UsageStore { + private static let cliProxyAPIUsageCollectionInterval: Duration = .seconds(30) + private static let cliProxyAPIPendingPruneInterval: Duration = .seconds(24 * 60 * 60) + + func startCLIProxyAPIUsageCollector(initialConfigurationGeneration: String? = nil) { + self.stopCLIProxyAPIUsageCollector() + self.cliProxyAPICleanupRetryTask?.cancel() + self.cliProxyAPICleanupRetryTask = nil + let pendingPruneInterval = Self.cliProxyAPIPendingPruneInterval + let initialConfigurationGeneration = initialConfigurationGeneration ?? + CostUsageCacheLocations.cliProxyAPIConfigurationGeneration() + self.cliProxyAPIUsageCollectorTask = Task.detached(priority: .utility) { [weak self] in + var nextPendingPruneAt: ContinuousClock.Instant? + var collectorState = CLIProxyAPIUsageCollectorState( + configurationGeneration: initialConfigurationGeneration) + while !Task.isCancelled { + let now = ContinuousClock.now + if nextPendingPruneAt.map({ now >= $0 }) ?? true { + _ = CLIProxyAPIUsageCollector.pruneExpiredUsage() + nextPendingPruneAt = now.advanced(by: pendingPruneInterval) + } + guard let result = await self?.collectCLIProxyAPIUsageNow() else { return } + collectorState = await self?.handleCLIProxyAPIUsageCollectionResult( + result, + collectorState: collectorState) ?? collectorState + do { + try await Task.sleep(for: Self.cliProxyAPIUsageCollectionInterval) + } catch { + return + } + } + } + } + + @discardableResult + func stopCLIProxyAPIUsageCollector() -> Task? { + let task = self.cliProxyAPIUsageCollectorTask + task?.cancel() + self.cliProxyAPIUsageCollectorTask = nil + return task + } + + @discardableResult + func removeCLIProxyAPIConfiguration( + remove: (() async -> CLIProxyAPIConfigurationRemovalResult)? = nil, + scheduleCleanupRetry: (() -> Void)? = nil) async + -> CLIProxyAPIConfigurationRemovalResult + { + let collectorTask = self.stopCLIProxyAPIUsageCollector() + self.cliProxyAPICleanupRetryTask?.cancel() + self.cliProxyAPICleanupRetryTask = nil + await collectorTask?.value + let result = if let remove { + await remove() + } else { + await Task.detached(priority: .utility) { + CLIProxyAPIConnectionSettingsStore.removeAndPurgeTelemetry() + }.value + } + if result != .configurationRemovalFailed { + self.invalidateCLIProxyAPICostAttribution() + } + if result == .telemetryCleanupFailed { + if let scheduleCleanupRetry { + scheduleCleanupRetry() + } else { + self.startCLIProxyAPICleanupRetry() + } + } + return result + } + + @discardableResult + func startCLIProxyAPICleanupRetry( + retryInterval: Duration = .seconds(30), + maintenance: @escaping @Sendable () -> Bool = { + CLIProxyAPIUsageCollector.pruneExpiredUsage() + }) -> Task + { + self.cliProxyAPICleanupRetryTask?.cancel() + let task = Task.detached(priority: .utility) { + while !Task.isCancelled { + if maintenance() { return } + do { + try await Task.sleep(for: retryInterval) + } catch { + return + } + } + } + self.cliProxyAPICleanupRetryTask = task + return task + } + + func collectCLIProxyAPIUsageNow( + collector: (@Sendable () async -> CLIProxyAPIUsageCollectionResult)? = nil) async + -> CLIProxyAPIUsageCollectionResult + { + guard self.settings.costUsageEnabled else { return .disabled } + if let collector { + return await collector() + } + return await CLIProxyAPIUsageCollector.collect(shouldContinue: { [weak self] in + await self?.settings.costUsageEnabled == true + }) + } + + func handleCLIProxyAPIUsageCollectionResult( + _ result: CLIProxyAPIUsageCollectionResult, + collectorState: CLIProxyAPIUsageCollectorState, + isExplicitlyDisconnected: () -> Bool = { + CostUsageCacheLocations.isCLIProxyAPIExplicitlyDisconnected() + }, + configurationGeneration: () -> String? = { + CostUsageCacheLocations.cliProxyAPIConfigurationGeneration() + }, + refresh: ((UsageProvider, Bool) async -> Void)? = nil) async -> CLIProxyAPIUsageCollectorState + { + var collectorState = collectorState + switch result { + case .notConfigured: + if collectorState.configurationAvailability == .available || + (collectorState.configurationAvailability == .unknown && isExplicitlyDisconnected()) + { + self.invalidateCLIProxyAPICostAttribution(widgetReason: "cliproxyapi-disconnected") + } + collectorState.configurationAvailability = .unavailable + collectorState.configurationGeneration = configurationGeneration() + case .collected: + let currentGeneration = configurationGeneration() + if collectorState.configurationAvailability == .unavailable || + (collectorState.configurationGeneration != nil && + collectorState.configurationGeneration != currentGeneration) + { + await self.refreshCLIProxyAPICostAttribution(refresh: refresh) + } + collectorState.configurationAvailability = .available + collectorState.configurationGeneration = currentGeneration + case .failed: + let currentGeneration = configurationGeneration() + if collectorState.configurationGeneration != nil, + collectorState.configurationGeneration != currentGeneration + { + self.invalidateCLIProxyAPICostAttribution(widgetReason: "cliproxyapi-configuration-changed") + collectorState.configurationAvailability = .unavailable + collectorState.configurationGeneration = currentGeneration + } + case .disabled: + break + } + return collectorState + } + + func refreshCLIProxyAPICostAttribution( + refresh: ((UsageProvider, Bool) async -> Void)? = nil) async + { + self.invalidateCLIProxyAPICostAttribution(widgetReason: "cliproxyapi-reconnected") + for provider in [UsageProvider.claude, .codex] { + if let refresh { + await refresh(provider, true) + } else { + await self.refreshTokenUsageNow(for: provider, force: true) + } + } + } +} diff --git a/Sources/CodexBar/UsageStore+TokenCost.swift b/Sources/CodexBar/UsageStore+TokenCost.swift index c1505d58c9..1e71d64dcd 100644 --- a/Sources/CodexBar/UsageStore+TokenCost.swift +++ b/Sources/CodexBar/UsageStore+TokenCost.swift @@ -11,6 +11,12 @@ struct CurrentProviderConfigTokenPublication: Sendable, Equatable { let publicationRevision: UInt64 } +struct TokenRefreshPublicationGuard { + let provider: UsageStore.ProviderPublicationRevision + let tokenSnapshot: UInt64 + let providerConfig: UInt64 +} + struct TokenSnapshotPublication: Sendable, Equatable { let snapshot: CostUsageTokenSnapshot? let publicationRevision: UInt64 @@ -116,6 +122,13 @@ extension UsageStore { self.tokenSnapshotPublicationRevisions[provider.instanceID] ?? 0 } + func tokenRefreshPublicationGuard(for provider: UsageProvider) -> TokenRefreshPublicationGuard { + TokenRefreshPublicationGuard( + provider: self.providerPublicationRevision(for: provider), + tokenSnapshot: self.tokenSnapshotPublicationRevision(for: provider), + providerConfig: self.settings.providerConfigRevision(for: provider)) + } + func publishTokenSnapshot(_ snapshot: CostUsageTokenSnapshot, for provider: UsageProvider) { self.tokenSnapshots[provider.instanceID] = snapshot self.publishTokenSnapshotState(snapshot, for: provider) @@ -126,6 +139,20 @@ extension UsageStore { self.publishTokenSnapshotState(nil, for: provider) } + func invalidateCLIProxyAPICostAttribution(widgetReason: String = "cliproxyapi-removed") { + self.cancelCodexCostCatchUp() + self.cancelSpendDashboardCodexCostCatchUp() + self.spendDashboardCodexCostCatchUpRevision &+= 1 + for provider in [UsageProvider.codex, .claude] { + self.publishConfirmedEmptyTokenSnapshot(for: provider) + self.tokenErrors[provider.instanceID] = nil + self.tokenFailureGates[provider.instanceID]?.reset() + self.lastTokenFetchAt.removeValue(forKey: provider.instanceID) + self.lastTokenFetchScope.removeValue(forKey: provider.instanceID) + } + self.persistWidgetSnapshot(reason: widgetReason) + } + private func publishTokenSnapshotState(_ snapshot: CostUsageTokenSnapshot?, for provider: UsageProvider) { self.tokenSnapshotPublicationRevisions[provider.instanceID, default: 0] &+= 1 self.tokenSnapshotPublications[provider.instanceID] = TokenSnapshotPublication( @@ -150,6 +177,9 @@ extension UsageStore { } func clearTokenSnapshots() { + for provider in UsageProvider.allCases { + self.tokenSnapshotPublicationRevisions[provider.instanceID, default: 0] &+= 1 + } self.tokenSnapshots.removeAll() self.tokenSnapshotPublications.removeAll() } @@ -210,6 +240,7 @@ extension UsageStore { let costUsageSettingsRevision = self.settings.costUsageSettingsRevision let tokenSnapshotScopeSignature = self.tokenSnapshotScopeSignature(for: .codex) let tokenSnapshotPublicationRevision = self.tokenSnapshotPublicationRevision(for: .codex) + let cliProxyAPIConfigurationGeneration = self.costUsageFetcher.cliProxyAPIConfigurationGeneration() return Task { @MainActor [weak self] in guard let self else { return } guard self.tokenSnapshotPublicationForCurrentProviderConfig(for: .codex) == nil else { return } @@ -244,6 +275,7 @@ extension UsageStore { self.settings.costUsageHistoryDays == historyDays, self.tokenSnapshotScopeSignature(for: .codex) == tokenSnapshotScopeSignature, self.tokenSnapshotPublicationRevision(for: .codex) == tokenSnapshotPublicationRevision, + self.costUsageFetcher.cliProxyAPIConfigurationGeneration() == cliProxyAPIConfigurationGeneration, self.tokenSnapshotPublicationForCurrentProviderConfig(for: .codex) == nil else { return @@ -365,14 +397,14 @@ extension UsageStore { func tokenRefreshPublicationIsCurrent( provider: UsageProvider, - publicationRevision: ProviderPublicationRevision, - providerConfigRevision: UInt64, + publicationGuard: TokenRefreshPublicationGuard, historyDays: Int, costScopeSignature: String, fetchedCredentialScopeFingerprint: String? = nil) -> Bool { - guard self.providerPublicationRevisionIsCurrent(publicationRevision, for: provider), - self.settings.providerConfigRevision(for: provider) == providerConfigRevision, + guard self.providerPublicationRevisionIsCurrent(publicationGuard.provider, for: provider), + self.tokenSnapshotPublicationRevision(for: provider) == publicationGuard.tokenSnapshot, + self.settings.providerConfigRevision(for: provider) == publicationGuard.providerConfig, self.settings.costUsageEnabled, self.isEnabled(provider), self.settings.costUsageHistoryDays == historyDays @@ -445,41 +477,63 @@ extension UsageStore { nonisolated static func costUsageCacheDirectory( fileManager: FileManager = .default) -> URL { - let root = fileManager.urls(for: .cachesDirectory, in: .userDomainMask).first! - return root - .appendingPathComponent("CodexBar", isDirectory: true) - .appendingPathComponent("cost-usage", isDirectory: true) - } - - func clearCostUsageCache() async -> String? { - let errorMessage: String? = await Task.detached(priority: .utility) { - let fm = FileManager.default - let cacheDirs = [ - Self.costUsageCacheDirectory(fileManager: fm), - ] - - for cacheDir in cacheDirs { - do { - try fm.removeItem(at: cacheDir) - } catch let error as NSError { - if error.domain == NSCocoaErrorDomain, error.code == NSFileNoSuchFileError { - continue - } - return error.localizedDescription - } + CostUsageCacheLocations.directories(fileManager: fileManager)[0] + } + + func clearCostUsageCache( + clearDirectories: (@Sendable () async -> (cleared: Int, errorMessage: String?))? = nil, + fileManager: FileManager = .default) async -> String? + { + guard !self.costUsageCacheClearInProgress else { return nil } + self.costUsageCacheClearInProgress = true + defer { self.costUsageCacheClearInProgress = false } + + let collectorTask = self.stopCLIProxyAPIUsageCollector() + await collectorTask?.value + defer { + if collectorTask != nil { + self.startCLIProxyAPIUsageCollector() } - return nil - }.value + } - guard errorMessage == nil else { return errorMessage } + await self.drainTokenRefreshesForCostCacheClear() + self.cancelCodexCostCatchUp() + self.cancelSpendDashboardCodexCostCatchUp() + + let cacheDirectories = CostUsageCacheLocations.directories(fileManager: fileManager) + let cliProxyAPIStateRoot = cacheDirectories[1].deletingLastPathComponent() + let clearResult: (didClear: Bool, errorMessage: String?) + if let clearDirectories { + let result = await clearDirectories() + clearResult = (result.errorMessage == nil || result.cleared > 0, result.errorMessage) + } else { + let result = await Task.detached(priority: .utility) { + CostUsageCacheLocations.clearAllCostUsageCaches( + in: cacheDirectories, + stateRoot: cliProxyAPIStateRoot) + }.value + clearResult = (result.errorDescription == nil || result.cleared > 0, result.errorDescription) + } + + guard clearResult.didClear else { return clearResult.errorMessage } self.clearTokenSnapshots() + self.spendDashboardCodexCostCatchUpRevision &+= 1 self.tokenErrors.removeAll() self.lastTokenFetchAt.removeAll() self.lastTokenFetchScope.removeAll() self.tokenFailureGates[.codex]?.reset() self.tokenFailureGates[.claude]?.reset() - return nil + return clearResult.errorMessage + } + + /// Fast failures may retry on the next scheduled pass instead of waiting out the fetch + /// TTL; timed-out scans keep the TTL so a slow corpus cannot thrash back-to-back rescans. + nonisolated static func tokenFetchFailureAllowsEarlyRetry(_ error: Error) -> Bool { + if case CostUsageError.timedOut = error { + return false + } + return true } nonisolated static func tokenCostNoDataMessage(for provider: UsageProvider) -> String { diff --git a/Sources/CodexBar/UsageStore+TokenRefreshSequence.swift b/Sources/CodexBar/UsageStore+TokenRefreshSequence.swift index 8b465d26cb..815ad6c0bf 100644 --- a/Sources/CodexBar/UsageStore+TokenRefreshSequence.swift +++ b/Sources/CodexBar/UsageStore+TokenRefreshSequence.swift @@ -24,7 +24,10 @@ extension UsageStore { } func scheduleTokenRefresh() { - guard self.tokenRefreshSequenceTask == nil, !self.hasForcedRefreshEnrichmentInFlight else { return } + guard !self.costUsageCacheClearInProgress, + self.tokenRefreshSequenceTask == nil, + !self.hasForcedRefreshEnrichmentInFlight + else { return } if self.startPendingTokenRefreshRetryIfPossible() { return } @@ -37,6 +40,7 @@ extension UsageStore { } func refreshTokenUsageNow(for provider: UsageProvider, force: Bool) async { + guard !self.costUsageCacheClearInProgress else { return } if force, self.tokenRefreshSequenceTask != nil, let activeProvider = self.tokenRefreshSequenceProvider, @@ -55,6 +59,17 @@ extension UsageStore { await self.awaitTokenRefreshSequence(task) } + func drainTokenRefreshesForCostCacheClear() async { + self.tokenRefreshRetryProviders.removeAll() + let sequenceTask = self.tokenRefreshSequenceTask + sequenceTask?.cancel() + await sequenceTask?.value + while !self.tokenRefreshInFlight.isEmpty { + await Task.yield() + } + self.tokenRefreshRetryProviders.removeAll() + } + private func serializedTokenRefreshTask( force: Bool, scope: TokenRefreshSequenceScope) async -> Task? diff --git a/Sources/CodexBar/UsageStore.swift b/Sources/CodexBar/UsageStore.swift index 9fd80274de..758df565ed 100644 --- a/Sources/CodexBar/UsageStore.swift +++ b/Sources/CodexBar/UsageStore.swift @@ -184,6 +184,7 @@ final class UsageStore { var tokenSnapshotPublicationRevisions: [ProviderInstanceID: UInt64] = [:] var tokenErrors: [ProviderInstanceID: String] = [:] var tokenRefreshInFlight: Set = [] + @ObservationIgnored var costUsageCacheClearInProgress = false var codexCostCatchUpActivity: CodexCostCatchUpActivity? var spendDashboardCodexCostCatchUpActivity: CodexCostCatchUpActivity? var spendDashboardCodexCostCatchUpRevision: UInt64 = 0 @@ -412,6 +413,8 @@ final class UsageStore { /// Background load task; cleared on deinit and on the cancel test seam. @ObservationIgnored var planUtilizationHistoryLoadTask: Task? + @ObservationIgnored var cliProxyAPIUsageCollectorTask: Task? + @ObservationIgnored var cliProxyAPICleanupRetryTask: Task? /// Set once after the load completes. Gates mutation paths and sync menu /// accessors so they cannot race the decode or write empty history back to disk. @ObservationIgnored var planUtilizationHistoryLoaded: Bool = false @@ -520,6 +523,7 @@ final class UsageStore { effectivePATH: PathBuilder.effectivePATH(purposes: [.rpc, .tty, .nodeTooling]), loginShellPATH: LoginShellPathCache.shared.current?.joined(separator: ":")) guard self.startupBehavior.automaticallyStartsBackgroundWork else { return } + let cliProxyAPIConfigurationGeneration = self.costUsageFetcher.cliProxyAPIConfigurationGeneration() self.hydrateCachedTokenSnapshots() self.detectVersions() self.updateProviderRuntimes() @@ -537,6 +541,8 @@ final class UsageStore { Task { await self.refresh(enrichmentMode: .automatic) } self.startTimer() self.startTokenTimer() + self.startCLIProxyAPIUsageCollector( + initialConfigurationGeneration: cliProxyAPIConfigurationGeneration) } var iconStyle: IconStyle { @@ -927,6 +933,8 @@ final class UsageStore { self.codexPlanHistoryBackfillTask?.cancel() self.resetBoundaryRefreshTask?.cancel() self.planUtilizationHistoryLoadTask?.cancel() + self.cliProxyAPIUsageCollectorTask?.cancel() + self.cliProxyAPICleanupRetryTask?.cancel() } enum SessionQuotaWindowSource: String { @@ -1476,8 +1484,7 @@ extension UsageStore { } let costScope = self.tokenCostScope(for: provider) let costScopeSignature = self.tokenSnapshotScopeSignature(for: provider) - let publicationRevision = self.providerPublicationRevision(for: provider) - let providerConfigRevision = self.settings.providerConfigRevision(for: provider) + let publicationGuard = self.tokenRefreshPublicationGuard(for: provider) if !force, self.tokenRefreshCanReuseCurrentSnapshot( provider: provider, now: now, @@ -1521,8 +1528,7 @@ extension UsageStore { snapshot: snapshot) guard self.tokenRefreshPublicationIsCurrent( provider: provider, - publicationRevision: publicationRevision, - providerConfigRevision: providerConfigRevision, + publicationGuard: publicationGuard, historyDays: historyDays, costScopeSignature: costScopeSignature, fetchedCredentialScopeFingerprint: snapshot.credentialScopeFingerprint) @@ -1555,8 +1561,7 @@ extension UsageStore { } catch { guard self.tokenRefreshPublicationIsCurrent( provider: provider, - publicationRevision: publicationRevision, - providerConfigRevision: providerConfigRevision, + publicationGuard: publicationGuard, historyDays: historyDays, costScopeSignature: costScopeSignature) else { @@ -1641,15 +1646,6 @@ extension UsageStore { self.lastTokenFetchAt.removeValue(forKey: provider.instanceID) self.lastTokenFetchScope.removeValue(forKey: provider.instanceID) } - - /// Fast failures may retry on the next scheduled pass instead of waiting out the fetch - /// TTL; timed-out scans keep the TTL so a slow corpus cannot thrash back-to-back rescans. - nonisolated static func tokenFetchFailureAllowsEarlyRetry(_ error: Error) -> Bool { - if case CostUsageError.timedOut = error { - return false - } - return true - } } extension UsageStore { diff --git a/Sources/CodexBarCLI/CLICacheCommand.swift b/Sources/CodexBarCLI/CLICacheCommand.swift index 78d63e96ec..e06574c306 100644 --- a/Sources/CodexBarCLI/CLICacheCommand.swift +++ b/Sources/CodexBarCLI/CLICacheCommand.swift @@ -53,19 +53,12 @@ extension CodexBarCLI { } if clearCost { - let fm = FileManager.default - let cacheDir = Self.costUsageCacheDirectory(fileManager: fm) - var cleared = 0 - var costError: String? - if fm.fileExists(atPath: cacheDir.path) { - do { - try fm.removeItem(at: cacheDir) - cleared = 1 - } catch { - costError = error.localizedDescription - } - } - results.append(CacheClearResult(cache: "cost", provider: nil, cleared: cleared, error: costError)) + let result = CostUsageCacheLocations.clearAllCostUsageCaches() + results.append(CacheClearResult( + cache: "cost", + provider: nil, + cleared: result.cleared, + error: result.errorDescription)) } switch output.format { @@ -144,9 +137,6 @@ private struct CacheClearResult: Encodable { extension CodexBarCLI { /// Mirrors the cost usage cache directory used by the app (UsageStore.costUsageCacheDirectory). static func costUsageCacheDirectory(fileManager: FileManager = .default) -> URL { - let root = fileManager.urls(for: .cachesDirectory, in: .userDomainMask).first! - return root - .appendingPathComponent("CodexBar", isDirectory: true) - .appendingPathComponent("cost-usage", isDirectory: true) + CostUsageCacheLocations.directories(fileManager: fileManager)[0] } } diff --git a/Sources/CodexBarCore/CLIProxyAPIAttributionResolver.swift b/Sources/CodexBarCore/CLIProxyAPIAttributionResolver.swift new file mode 100644 index 0000000000..23fdc9e526 --- /dev/null +++ b/Sources/CodexBarCore/CLIProxyAPIAttributionResolver.swift @@ -0,0 +1,850 @@ +import Foundation + +struct CLIProxyAPIAttributionResolver: Sendable { + struct Observation: Sendable, Equatable { + let sourceID: String? + let sessionID: String + let model: String + let timestamp: Date? + + init(sourceID: String? = nil, sessionID: String, model: String, timestamp: Date?) { + self.sourceID = sourceID + self.sessionID = sessionID + self.model = model + self.timestamp = timestamp + } + } + + struct AuthProvider: Sendable, Equatable, Hashable { + let provider: String + let authType: CostUsageAttribution.Upstream.AuthType + } + + struct TokenSignature: Sendable, Equatable { + let input: Int + let cacheRead: Int + let cacheCreate: Int + let output: Int + } + + struct Request: Sendable { + let model: String + let modelProvider: CostUsageAttribution.ModelProvider + let sessionID: String? + let timestampUnixMs: Int64? + let tokens: TokenSignature? + } + + private struct ObservationKey: Hashable { + let sourceID: String? + let sessionID: String + let canonicalModel: String + let timestamp: Date? + } + + private struct UsageRecordKey: Hashable { + let sourceID: Int + } + + private struct IndexedUsageRecord { + let sourceID: Int + let record: CLIProxyAPIUsageRecord + } + + private struct UsageRecordMatch { + let key: UsageRecordKey + let record: CLIProxyAPIUsageRecord + } + + private static let requestBodyMarker = "=== REQUEST BODY ===" + private static let responseMarkers = ["=== API RESPONSE ===", "=== RESPONSE ==="] + private static let maxLogPrefixBytes = 2 * 1024 * 1024 + private static let maximumRouteMatchDistance: TimeInterval = 60 * 60 + private static let maximumTelemetryMatchDistance: TimeInterval = 5 + private static let observationCache = ObservationCache() + + private let observationsBySessionID: [String: [Observation]] + private let observationsByCanonicalModel: [String: [Observation]] + private let usageRecordsByCanonicalModel: [String: [IndexedUsageRecord]] + private let authProviders: [AuthProvider] + private let codexOAuthModelRoutes: [String: String] + private let hasConfiguredOpenAIAPIUpstream: Bool + + init( + observations: [Observation], + usageRecords: [CLIProxyAPIUsageRecord] = [], + authProviders: [AuthProvider] = [], + codexOAuthModelAliases: [String: String] = [:], + hasConfiguredOpenAIAPIUpstream: Bool = false) + { + self.observationsBySessionID = Dictionary(grouping: observations, by: \.sessionID) + self.observationsByCanonicalModel = Dictionary( + grouping: observations, + by: { Self.canonicalModel($0.model) }) + self.usageRecordsByCanonicalModel = Self.indexUsageRecords(usageRecords) + self.authProviders = authProviders + self.codexOAuthModelRoutes = codexOAuthModelAliases.reduce(into: [:]) { result, entry in + let upstreamModel = entry.value.trimmingCharacters(in: .whitespacesAndNewlines) + guard !upstreamModel.isEmpty else { return } + result[Self.canonicalModel(entry.key)] = upstreamModel + result[Self.canonicalModel(upstreamModel)] = upstreamModel + } + self.hasConfiguredOpenAIAPIUpstream = hasConfiguredOpenAIAPIUpstream + } + + static func load( + home: URL, + cacheRoot: URL? = nil, + fileManager: FileManager = .default, + forceReload: Bool = false, + usageRecords: [CLIProxyAPIUsageRecord]? = nil, + checkCancellation: (() throws -> Void)? = nil) throws -> Self + { + let observations = try self.loadObservations( + logDirectory: home.appendingPathComponent("logs", isDirectory: true), + fileManager: fileManager, + forceReload: forceReload, + checkCancellation: checkCancellation) + return Self( + observations: observations, + usageRecords: usageRecords ?? CLIProxyAPIUsageCacheIO.load(cacheRoot: cacheRoot), + authProviders: self.loadAuthProviders(home: home, fileManager: fileManager), + codexOAuthModelAliases: self.loadCodexOAuthModelAliases(home: home, fileManager: fileManager), + hasConfiguredOpenAIAPIUpstream: self.hasConfiguredOpenAIAPIUpstream( + home: home, + fileManager: fileManager)) + } + + func attribution( + model: String, + modelProvider: CostUsageAttribution.ModelProvider, + sessionID: String?, + timestampUnixMs: Int64?, + tokens: TokenSignature?) -> CostUsageAttribution + { + let request = Request( + model: model, + modelProvider: modelProvider, + sessionID: sessionID, + timestampUnixMs: timestampUnixMs, + tokens: tokens) + let routeObservation = self.matchingObservation( + model: model, + sessionID: sessionID, + timestampUnixMs: timestampUnixMs) + let telemetryObservation = self.matchingObservation( + model: model, + sessionID: sessionID, + timestampUnixMs: timestampUnixMs) + let usageRecord = telemetryObservation.flatMap { + self.matchingUsageRecord( + observation: $0, + model: model, + tokens: tokens) + } + return self.attribution( + request: request, + routeObservation: routeObservation, + usageRecord: usageRecord) + } + + func attributions(for requests: [Request]) -> [CostUsageAttribution] { + var prepared = requests.map { request in + let observationCandidates = self.matchingObservations( + model: request.model, + sessionID: request.sessionID, + timestampUnixMs: request.timestampUnixMs) + let observationMatches = observationCandidates.compactMap { observation in + self.closestUsageRecordMatch( + observation: observation, + model: request.model, + tokens: request.tokens) + } + let matchKeys = Set(observationMatches.map(\.key)) + let usageRecordMatch = observationMatches.count == observationCandidates.count && matchKeys.count == 1 + ? observationMatches.first + : nil + let observation = observationCandidates.count == 1 ? observationCandidates[0] : nil + return ( + request: request, + routeObservation: observation, + telemetryObservation: observation, + usageRecordMatch: usageRecordMatch, + observationCandidates: observationCandidates) + } + var claimedObservationKeys = Set(prepared.compactMap(\.routeObservation).map(Self.observationKey)) + for index in prepared.indices where prepared[index].routeObservation == nil + && prepared[index].usageRecordMatch != nil + { + let candidates = prepared[index].observationCandidates.sorted { + ($0.sourceID ?? "") < ($1.sourceID ?? "") + } + guard let observation = candidates.first(where: { + !claimedObservationKeys.contains(Self.observationKey($0)) + }) else { continue } + prepared[index].routeObservation = observation + prepared[index].telemetryObservation = observation + claimedObservationKeys.insert(Self.observationKey(observation)) + } + var routeCandidatesByObservation: + [ObservationKey: [(index: Int, request: Request, observation: Observation)]] = [:] + for (index, item) in prepared.enumerated() { + guard let observation = item.routeObservation else { continue } + routeCandidatesByObservation[Self.observationKey(observation), default: []].append( + (index: index, request: item.request, observation: observation)) + } + let matchCounts = Dictionary( + grouping: prepared.compactMap(\.usageRecordMatch?.key), + by: { $0 }).mapValues(\.count) + var routeOwnerByObservation: [ObservationKey: Int] = [:] + for (key, candidates) in routeCandidatesByObservation { + if candidates.count == 1, + let candidate = candidates.first, + prepared[candidate.index].usageRecordMatch != nil + || Self.isCloseRouteMatch(candidate) + { + routeOwnerByObservation[key] = candidate.index + } else if candidates.count > 1, + let observationTimestamp = candidates.first?.observation.timestamp, + let candidate = Self.uniqueClosest( + candidates, + target: observationTimestamp, + timestamp: { Self.timestamp(for: $0.request) }), + prepared[candidate.index].usageRecordMatch.map({ matchCounts[$0.key] == 1 }) == true + || Self.isCloseRouteMatch(candidate) + { + routeOwnerByObservation[key] = candidate.index + } + } + let representedObservations = Set(prepared.compactMap { item -> ObservationKey? in + guard item.usageRecordMatch != nil, + let observation = item.telemetryObservation + else { return nil } + return Self.observationKey(observation) + }) + + return prepared.enumerated().map { index, item in + let routeObservation = item.routeObservation.flatMap { observation in + routeOwnerByObservation[Self.observationKey(observation)] == index + ? observation + : nil + } + let usageRecord = item.usageRecordMatch.flatMap { match -> CLIProxyAPIUsageRecord? in + guard matchCounts[match.key] == 1, + self.allPlausibleObservationsRepresented( + for: match.record, + model: item.request.model, + representedObservations: representedObservations) + else { return nil } + return match.record + } + return self.attribution( + request: item.request, + routeObservation: routeObservation, + usageRecord: usageRecord) + } + } + + private static func isCloseRouteMatch( + _ candidate: (index: Int, request: Request, observation: Observation)) -> Bool + { + guard let requestTimestamp = timestamp(for: candidate.request), + let observationTimestamp = candidate.observation.timestamp + else { return false } + return abs(requestTimestamp.timeIntervalSince(observationTimestamp)) <= Self.maximumTelemetryMatchDistance + } + + func hasMatchingObservation(for request: Request) -> Bool { + self.matchingObservation( + model: request.model, + sessionID: request.sessionID, + timestampUnixMs: request.timestampUnixMs) != nil + } + + private func attribution( + request: Request, + routeObservation: Observation?, + usageRecord: CLIProxyAPIUsageRecord?) -> CostUsageAttribution + { + let canonicalModel = Self.canonicalModel(request.model) + let configuredCodexModel = self.codexOAuthModelRoutes[canonicalModel] + let routeObserved = routeObservation != nil || usageRecord != nil + let resolvedModelProvider: CostUsageAttribution.ModelProvider = + request.modelProvider == .unknown && configuredCodexModel != nil && routeObserved + ? .openAI + : request.modelProvider + let inventoryUpstream = usageRecord == nil + ? self.authInventoryUpstream( + model: request.model, + modelProvider: resolvedModelProvider, + routeObserved: routeObservation != nil, + configuredCodexModel: configuredCodexModel) + : nil + let routeConfirmed = routeObservation != nil || usageRecord != nil || inventoryUpstream != nil + var evidence: Set = [.modelProvider] + if routeObservation != nil { + evidence.insert(.cliProxyRequestLog) + } + if configuredCodexModel != nil, inventoryUpstream != nil { + evidence.insert(.cliProxyModelAlias) + } + if usageRecord != nil { + evidence.insert(.cliProxyUsageTelemetry) + } + if inventoryUpstream != nil { + evidence.insert(.cliProxyAuthInventory) + } + + return CostUsageAttribution( + client: .claudeCode, + route: routeConfirmed ? .cliProxyAPI : .unknown, + modelProvider: resolvedModelProvider, + upstream: usageRecord.map(Self.upstream) ?? inventoryUpstream, + evidence: evidence.sorted { $0.rawValue < $1.rawValue }) + } + + private func matchingObservation( + model: String, + sessionID: String?, + timestampUnixMs: Int64?) -> Observation? + { + let candidates = self.matchingObservations( + model: model, + sessionID: sessionID, + timestampUnixMs: timestampUnixMs) + return candidates.count == 1 ? candidates[0] : nil + } + + private func matchingObservations( + model: String, + sessionID: String?, + timestampUnixMs: Int64?) -> [Observation] + { + guard let sessionID = sessionID?.trimmingCharacters(in: .whitespacesAndNewlines), + !sessionID.isEmpty, + let observations = self.observationsBySessionID[sessionID] + else { return [] } + + let canonicalModel = Self.canonicalModel(model) + let matchingModels = observations.filter { Self.canonicalModel($0.model) == canonicalModel } + guard !matchingModels.isEmpty else { return [] } + guard let timestampUnixMs else { + return matchingModels.count == 1 ? matchingModels : [] + } + let timestamp = Date(timeIntervalSince1970: Double(timestampUnixMs) / 1000) + let ranked = matchingModels.compactMap { observation -> (Observation, TimeInterval)? in + guard let observationTimestamp = observation.timestamp else { return nil } + let distance = abs(observationTimestamp.timeIntervalSince(timestamp)) + return distance <= Self.maximumRouteMatchDistance ? (observation, distance) : nil + } + guard let closestDistance = ranked.map(\.1).min() else { return [] } + return ranked.filter { $0.1 == closestDistance }.map(\.0) + } + + private static func timestamp(for request: Request) -> Date? { + request.timestampUnixMs.map { + Date(timeIntervalSince1970: Double($0) / 1000) + } + } + + private func matchingUsageRecord( + observation: Observation, + model: String, + tokens: TokenSignature?) -> CLIProxyAPIUsageRecord? + { + let candidates = self.usageRecordMatches( + observation: observation, + model: model, + tokens: tokens) + guard candidates.count == 1, let candidate = candidates.first else { return nil } + return self.plausibleObservations(for: candidate.record, model: model).count == 1 + ? candidate.record + : nil + } + + private func closestUsageRecordMatch( + observation: Observation, + model: String, + tokens: TokenSignature?) -> UsageRecordMatch? + { + guard let observationTimestamp = observation.timestamp else { return nil } + let candidates = self.usageRecordMatches( + observation: observation, + model: model, + tokens: tokens) + return Self.uniqueClosest( + candidates, + target: observationTimestamp, + timestamp: { $0.record.timestamp }) + } + + private func usageRecordMatches( + observation: Observation, + model: String, + tokens: TokenSignature?) -> [UsageRecordMatch] + { + guard let observationTimestamp = observation.timestamp else { return [] } + let canonicalModel = Self.canonicalModel(model) + guard let records = self.usageRecordsByCanonicalModel[canonicalModel] else { return [] } + let earliest = observationTimestamp.addingTimeInterval(-Self.maximumTelemetryMatchDistance) + let latest = observationTimestamp.addingTimeInterval(Self.maximumTelemetryMatchDistance) + let startIndex = Self.firstRecordIndex(atOrAfter: earliest, in: records) + var candidates: [UsageRecordMatch] = [] + for index in startIndex.. [Observation] + { + let canonicalModel = Self.canonicalModel(model) + return self.observationsByCanonicalModel[canonicalModel]?.filter { + guard let timestamp = $0.timestamp else { return false } + return abs(timestamp.timeIntervalSince(record.timestamp)) <= Self.maximumTelemetryMatchDistance + } ?? [] + } + + private func allPlausibleObservationsRepresented( + for record: CLIProxyAPIUsageRecord, + model: String, + representedObservations: Set) -> Bool + { + let plausibleKeys = self.plausibleObservations(for: record, model: model).map(Self.observationKey) + return !plausibleKeys.isEmpty + && Set(plausibleKeys).count == plausibleKeys.count + && plausibleKeys.allSatisfy(representedObservations.contains) + } + + private static func observationKey(_ observation: Observation) -> ObservationKey { + ObservationKey( + sourceID: observation.sourceID, + sessionID: observation.sessionID, + canonicalModel: self.canonicalModel(observation.model), + timestamp: observation.timestamp) + } + + private static func indexUsageRecords( + _ records: [CLIProxyAPIUsageRecord]) -> [String: [IndexedUsageRecord]] + { + var recordsByModel: [String: [IndexedUsageRecord]] = [:] + for (sourceID, record) in records.enumerated() where !record.failed + && record.generate + && self.isClaudeMessagesGenerationEndpoint(record.endpoint) + { + let models = Set([self.canonicalModel(record.alias), self.canonicalModel(record.model)]) + for model in models where !model.isEmpty { + recordsByModel[model, default: []].append(IndexedUsageRecord( + sourceID: sourceID, + record: record)) + } + } + return recordsByModel.mapValues { records in + records.sorted { $0.record.timestamp < $1.record.timestamp } + } + } + + private static func firstRecordIndex( + atOrAfter timestamp: Date, + in records: [IndexedUsageRecord]) -> Int + { + var lowerBound = 0 + var upperBound = records.count + while lowerBound < upperBound { + let midpoint = lowerBound + (upperBound - lowerBound) / 2 + if records[midpoint].record.timestamp < timestamp { + lowerBound = midpoint + 1 + } else { + upperBound = midpoint + } + } + return lowerBound + } + + private func authInventoryUpstream( + model: String, + modelProvider: CostUsageAttribution.ModelProvider, + routeObserved: Bool, + configuredCodexModel: String?) -> CostUsageAttribution.Upstream? + { + let providers = Array(Set(self.authProviders)) + let codexProviders = providers.filter { + $0.provider.caseInsensitiveCompare("codex") == .orderedSame + } + if let configuredCodexModel { + guard routeObserved, + codexProviders.count == 1, + let provider = codexProviders.first + else { return nil } + return CostUsageAttribution.Upstream( + provider: provider.provider, + authType: provider.authType, + model: configuredCodexModel) + } + + guard routeObserved, + modelProvider == .openAI, + !self.observationsBySessionID.isEmpty, + !self.hasConfiguredOpenAIAPIUpstream, + providers.count == 1, + let provider = providers.first, + provider.provider.caseInsensitiveCompare("codex") == .orderedSame + else { return nil } + return CostUsageAttribution.Upstream( + provider: provider.provider, + authType: provider.authType, + model: model.trimmingCharacters(in: .whitespacesAndNewlines)) + } + + private static func tokensMatch( + _ tokens: TokenSignature, + _ telemetry: CLIProxyAPIUsageRecord.Tokens) -> Bool + { + guard telemetry.output == tokens.output else { return false } + if telemetry.cacheRead != 0 || telemetry.cacheCreation != 0 { + return telemetry.input == tokens.input + && telemetry.cacheRead == tokens.cacheRead + && telemetry.cacheCreation == tokens.cacheCreate + } + let claudeInputTotal = tokens.input + tokens.cacheRead + tokens.cacheCreate + return telemetry.input == tokens.input + || telemetry.input == claudeInputTotal + || telemetry.input + telemetry.cached == claudeInputTotal + } + + private static func uniqueClosest( + _ candidates: [T], + target: Date, + timestamp: (T) -> Date?) -> T? + { + let ranked = candidates.compactMap { candidate -> (candidate: T, distance: TimeInterval)? in + guard let date = timestamp(candidate) else { return nil } + return (candidate, abs(date.timeIntervalSince(target))) + } + .sorted { $0.distance < $1.distance } + guard let first = ranked.first else { + let undated = candidates.filter { timestamp($0) == nil } + return undated.count == 1 ? undated[0] : nil + } + guard ranked.count == 1 || ranked[1].distance > first.distance else { return nil } + return first.candidate + } + + private static func upstream( + _ record: CLIProxyAPIUsageRecord) -> CostUsageAttribution.Upstream + { + let authType: CostUsageAttribution.Upstream.AuthType = switch record.authType + .trimmingCharacters(in: .whitespacesAndNewlines) + .lowercased() + { + case "oauth": .oauth + case "api_key", "api-key", "apikey": .apiKey + default: .unknown + } + return CostUsageAttribution.Upstream( + provider: record.provider.trimmingCharacters(in: .whitespacesAndNewlines), + authType: authType, + model: record.model.trimmingCharacters(in: .whitespacesAndNewlines), + executorType: record.executorType?.trimmingCharacters(in: .whitespacesAndNewlines)) + } + + private struct AuthFile: Decodable { + let type: String? + let disabled: Bool? + } + + private final class ObservationCache: @unchecked Sendable { + private struct Metadata: Equatable { + let modificationDate: Date? + let size: Int? + } + + private struct Entry { + let metadata: Metadata + let observation: Observation? + } + + private let lock = NSLock() + private var entriesByDirectory: [String: [String: Entry]] = [:] + + func load( + logDirectory: URL, + fileManager: FileManager, + forceReload: Bool, + checkCancellation: (() throws -> Void)?, + parse: (URL) -> Observation?) throws -> [Observation] + { + try self.lock.withLock { + let directoryKey = logDirectory.standardizedFileURL.path + let resourceKeys: Set = [ + .contentModificationDateKey, + .fileSizeKey, + .isRegularFileKey, + ] + guard let urls = try? fileManager.contentsOfDirectory( + at: logDirectory, + includingPropertiesForKeys: Array(resourceKeys), + options: [.skipsHiddenFiles]) + else { + self.entriesByDirectory.removeValue(forKey: directoryKey) + return [] + } + + let cachedEntries = forceReload ? [:] : self.entriesByDirectory[directoryKey] ?? [:] + var currentEntries: [String: Entry] = [:] + var observations: [Observation] = [] + for url in urls where url.pathExtension.lowercased() == "log" { + try checkCancellation?() + guard let values = try? url.resourceValues(forKeys: resourceKeys), + values.isRegularFile == true + else { continue } + + let path = url.standardizedFileURL.path + let metadata = Metadata( + modificationDate: values.contentModificationDate, + size: values.fileSize) + let entry = if let cached = cachedEntries[path], + cached.metadata == metadata + { + cached + } else { + Entry(metadata: metadata, observation: parse(url)) + } + currentEntries[path] = entry + if let observation = entry.observation { + observations.append(observation) + } + } + + if currentEntries.isEmpty { + self.entriesByDirectory.removeValue(forKey: directoryKey) + } else { + self.entriesByDirectory[directoryKey] = currentEntries + } + return observations + } + } + } + + private static func loadAuthProviders( + home: URL, + fileManager: FileManager) -> [AuthProvider] + { + guard let urls = try? fileManager.contentsOfDirectory( + at: home, + includingPropertiesForKeys: [.isRegularFileKey], + options: [.skipsHiddenFiles]) + else { return [] } + + let decoder = JSONDecoder() + let providers = urls.compactMap { url -> AuthProvider? in + guard url.pathExtension.lowercased() == "json", + let values = try? url.resourceValues(forKeys: [.isRegularFileKey]), + values.isRegularFile == true, + let data = try? Data(contentsOf: url), + let auth = try? decoder.decode(AuthFile.self, from: data), + auth.disabled != true, + let rawType = auth.type?.trimmingCharacters(in: .whitespacesAndNewlines), + !rawType.isEmpty + else { return nil } + let isCodex = rawType.caseInsensitiveCompare("codex") == .orderedSame + return AuthProvider( + provider: isCodex ? "codex" : rawType.lowercased(), + authType: isCodex ? .oauth : .unknown) + } + return Array(Set(providers)) + } + + static func parseCodexOAuthModelAliases(_ text: String) -> [String: String] { + var aliases: [String: String] = [:] + var rootIndent: Int? + var codexIndent: Int? + var currentName: String? + + for rawLine in text.split(whereSeparator: \.isNewline) { + let line = String(rawLine) + let trimmed = line.trimmingCharacters(in: .whitespaces) + guard !trimmed.isEmpty, !trimmed.hasPrefix("#") else { continue } + let indent = line.prefix { $0 == " " }.count + + if rootIndent == nil { + if trimmed == "oauth-model-alias:" { + rootIndent = indent + } + continue + } + + guard let rootIndent else { continue } + if indent <= rootIndent { + break + } + if codexIndent == nil { + if trimmed == "codex:" { + codexIndent = indent + } + continue + } + + guard let codexIndent else { continue } + if indent <= codexIndent { + break + } + if trimmed.hasPrefix("- name:") { + currentName = self.simpleYAMLScalar(String(trimmed.dropFirst("- name:".count))) + } else if trimmed.hasPrefix("alias:"), let currentName { + let alias = self.simpleYAMLScalar(String(trimmed.dropFirst("alias:".count))) + if !alias.isEmpty, !currentName.isEmpty { + aliases[alias] = currentName + } + } + } + return aliases + } + + static func hasCodexOAuthModelAliasRoute( + home: URL, + fileManager: FileManager = .default) -> Bool + { + guard !self.loadCodexOAuthModelAliases(home: home, fileManager: fileManager).isEmpty else { + return false + } + return self.loadAuthProviders(home: home, fileManager: fileManager).contains { + $0.provider.caseInsensitiveCompare("codex") == .orderedSame + } + } + + private static func loadCodexOAuthModelAliases( + home: URL, + fileManager: FileManager) -> [String: String] + { + let url = home.appendingPathComponent("config.yaml", isDirectory: false) + guard fileManager.fileExists(atPath: url.path), + let text = try? String(contentsOf: url, encoding: .utf8) + else { return [:] } + return self.parseCodexOAuthModelAliases(text) + } + + private static func simpleYAMLScalar(_ raw: String) -> String { + let trimmed = raw.trimmingCharacters(in: .whitespaces) + guard let first = trimmed.first else { return "" } + if first == "\"" || first == "'" { + let remainder = trimmed.dropFirst() + guard let end = remainder.firstIndex(of: first) else { return String(remainder) } + return String(remainder[.. Bool + { + let url = home.appendingPathComponent("config.yaml", isDirectory: false) + guard fileManager.fileExists(atPath: url.path), + let text = try? String(contentsOf: url, encoding: .utf8) + else { return false } + let conflictingKeys = ["codex-api-key", "openai-compatibility"] + return text.split(whereSeparator: \.isNewline).contains { line in + guard line.first?.isWhitespace != true else { return false } + let trimmed = line.trimmingCharacters(in: .whitespacesAndNewlines) + guard !trimmed.hasPrefix("#") else { return false } + return conflictingKeys.contains { trimmed.hasPrefix("\($0):") } + } + } + + private static func loadObservations( + logDirectory: URL, + fileManager: FileManager, + forceReload: Bool, + checkCancellation: (() throws -> Void)?) throws -> [Observation] + { + try self.observationCache.load( + logDirectory: logDirectory, + fileManager: fileManager, + forceReload: forceReload, + checkCancellation: checkCancellation) + { url in + self.parseObservation(url: url) + } + } + + private static func parseObservation(url: URL) -> Observation? { + guard let handle = try? FileHandle(forReadingFrom: url) else { return nil } + defer { try? handle.close() } + guard let data = try? handle.read(upToCount: self.maxLogPrefixBytes), + let text = String(data: data, encoding: .utf8), + let bodyMarkerRange = text.range(of: self.requestBodyMarker) + else { return nil } + + let info = String(text[.. Bool { + let candidate = value.split(whereSeparator: \.isWhitespace).last.map(String.init) ?? value + guard let components = URLComponents(string: candidate) else { return false } + return components.path == "/v1/messages" + } + + private static func field(_ name: String, in text: String) -> String? { + let prefix = "\(name):" + for line in text.split(whereSeparator: \.isNewline) { + guard line.hasPrefix(prefix) else { continue } + return line.dropFirst(prefix.count).trimmingCharacters(in: .whitespacesAndNewlines) + } + return nil + } + + private static func topLevelJSONStringValue(forKey key: String, in text: String) -> String? { + guard let data = text.data(using: .utf8), + let object = try? JSONSerialization.jsonObject(with: data) as? [String: Any], + let value = object[key] as? String + else { + let escapedKey = NSRegularExpression.escapedPattern(for: key) + guard let regex = try? NSRegularExpression( + pattern: "\"\(escapedKey)\"\\s*:\\s*\"([^\"\\\\]*(?:\\\\.[^\"\\\\]*)*)\""), + let match = regex.firstMatch( + in: text, + range: NSRange(text.startIndex..., in: text)), + let valueRange = Range(match.range(at: 1), in: text) + else { return nil } + return String(text[valueRange]) + } + return value + } + + private static func canonicalModel(_ raw: String) -> String { + let codexNormalized = CostUsagePricing.normalizeCodexModel(raw) + return CostUsagePricing.normalizeClaudeModel(codexNormalized) + .trimmingCharacters(in: .whitespacesAndNewlines) + .lowercased() + } +} diff --git a/Sources/CodexBarCore/CLIProxyAPIUsageTelemetry.swift b/Sources/CodexBarCore/CLIProxyAPIUsageTelemetry.swift new file mode 100644 index 0000000000..498dbd5a98 --- /dev/null +++ b/Sources/CodexBarCore/CLIProxyAPIUsageTelemetry.swift @@ -0,0 +1,1304 @@ +import Foundation +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif + +struct CLIProxyAPIUsageRecord: Codable, Equatable, Sendable { + struct Tokens: Codable, Equatable, Sendable { + let input: Int + let output: Int + let reasoning: Int + let cached: Int + let cacheRead: Int + let cacheCreation: Int + let total: Int + + private enum CodingKeys: String, CodingKey { + case input = "input_tokens" + case output = "output_tokens" + case reasoning = "reasoning_tokens" + case cached = "cached_tokens" + case cacheRead = "cache_read_tokens" + case cacheCreation = "cache_creation_tokens" + case total = "total_tokens" + } + + init( + input: Int, + output: Int, + reasoning: Int = 0, + cached: Int = 0, + cacheRead: Int = 0, + cacheCreation: Int = 0, + total: Int) + { + self.input = input + self.output = output + self.reasoning = reasoning + self.cached = cached + self.cacheRead = cacheRead + self.cacheCreation = cacheCreation + self.total = total + } + + init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + self.input = try container.decodeIfPresent(Int.self, forKey: .input) ?? 0 + self.output = try container.decodeIfPresent(Int.self, forKey: .output) ?? 0 + self.reasoning = try container.decodeIfPresent(Int.self, forKey: .reasoning) ?? 0 + self.cached = try container.decodeIfPresent(Int.self, forKey: .cached) ?? 0 + self.cacheRead = try container.decodeIfPresent(Int.self, forKey: .cacheRead) ?? 0 + self.cacheCreation = try container.decodeIfPresent(Int.self, forKey: .cacheCreation) ?? 0 + self.total = try container.decodeIfPresent(Int.self, forKey: .total) ?? 0 + } + } + + let timestamp: Date + let provider: String + let executorType: String? + let model: String + let alias: String + let endpoint: String + let authType: String + let requestID: String + let localOccurrenceID: String? + let failed: Bool + let generate: Bool + let tokens: Tokens + + private enum CodingKeys: String, CodingKey { + case timestamp + case provider + case executorType = "executor_type" + case model + case alias + case endpoint + case authType = "auth_type" + case requestID = "request_id" + case localOccurrenceID = "codexbar_occurrence_id" + case failed + case generate + case tokens + } + + init( + timestamp: Date, + provider: String, + executorType: String? = nil, + model: String, + alias: String, + endpoint: String, + authType: String, + requestID: String, + localOccurrenceID: String? = nil, + failed: Bool = false, + generate: Bool = true, + tokens: Tokens) + { + self.timestamp = timestamp + self.provider = provider + self.executorType = executorType + self.model = model + self.alias = alias + self.endpoint = endpoint + self.authType = authType + self.requestID = requestID + self.localOccurrenceID = localOccurrenceID + self.failed = failed + self.generate = generate + self.tokens = tokens + } + + init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + self.timestamp = try container.decode(Date.self, forKey: .timestamp) + self.provider = try container.decode(String.self, forKey: .provider) + self.executorType = try container.decodeIfPresent(String.self, forKey: .executorType) + self.model = try container.decode(String.self, forKey: .model) + self.alias = try container.decodeIfPresent(String.self, forKey: .alias) ?? self.model + self.endpoint = try container.decodeIfPresent(String.self, forKey: .endpoint) ?? "" + self.authType = try container.decodeIfPresent(String.self, forKey: .authType) ?? "" + self.requestID = try container.decodeIfPresent(String.self, forKey: .requestID) ?? "" + self.localOccurrenceID = try container.decodeIfPresent(String.self, forKey: .localOccurrenceID) + self.failed = try container.decodeIfPresent(Bool.self, forKey: .failed) ?? false + self.generate = try container.decodeIfPresent(Bool.self, forKey: .generate) ?? true + self.tokens = try container.decodeIfPresent(Tokens.self, forKey: .tokens) + ?? Tokens(input: 0, output: 0, total: 0) + } + + func assigningNewLocalOccurrenceID() -> Self { + guard self.requestID.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty else { return self } + return Self( + timestamp: self.timestamp, + provider: self.provider, + executorType: self.executorType, + model: self.model, + alias: self.alias, + endpoint: self.endpoint, + authType: self.authType, + requestID: self.requestID, + localOccurrenceID: UUID().uuidString.lowercased(), + failed: self.failed, + generate: self.generate, + tokens: self.tokens) + } +} + +enum CLIProxyAPIUsageCacheIO { + private struct Cache: Codable, Equatable { + var version: Int = 1 + var records: [CLIProxyAPIUsageRecord] = [] + } + + private enum CacheReadResult { + case missing + case valid(Cache) + case invalid + } + + private static let cacheLock = NSLock() + private static let maximumRecordAge: TimeInterval = 366 * 24 * 60 * 60 + + static func withExclusiveAccess(_ body: () throws -> T) rethrows -> T { + try self.cacheLock.withLock(body) + } + + static func pruneUnserialized( + cacheRoot: URL?, + now: Date) -> Bool + { + let legacyCacheRoot = cacheRoot == nil ? self.defaultLegacyCacheRoot() : nil + return self.withExclusiveAccess { + guard let currentCache = self.loadCache( + cacheRoot: cacheRoot, + legacyCacheRoot: legacyCacheRoot) + else { return false } + let cutoff = now.addingTimeInterval(-self.maximumRecordAge) + let retainedCache = Cache(records: currentCache.records.filter { $0.timestamp >= cutoff }) + return retainedCache == currentCache || self.save(retainedCache, cacheRoot: cacheRoot) + } + } + + static func load( + cacheRoot: URL? = nil, + now: Date = Date()) -> [CLIProxyAPIUsageRecord] + { + let legacyCacheRoot = cacheRoot == nil ? self.defaultLegacyCacheRoot() : nil + return self.load( + cacheRoot: cacheRoot, + legacyCacheRoot: legacyCacheRoot, + now: now) + } + + /// Reads and prunes the cache while the caller already owns the CLIProxyAPI interprocess lock. + static func loadAssumingInterprocessLockHeld( + cacheRoot: URL?, + now: Date = Date()) -> [CLIProxyAPIUsageRecord] + { + let legacyCacheRoot = cacheRoot == nil ? self.defaultLegacyCacheRoot() : nil + let cutoff = now.addingTimeInterval(-self.maximumRecordAge) + return self.withExclusiveAccess { + guard let currentCache = self.loadCache( + cacheRoot: cacheRoot, + legacyCacheRoot: legacyCacheRoot) + else { return [] } + let retainedRecords = currentCache.records.filter { $0.timestamp >= cutoff } + let retainedCache = Cache(records: retainedRecords) + if retainedCache != currentCache { + _ = self.save(retainedCache, cacheRoot: cacheRoot) + } + return retainedRecords + } + } + + static func load( + cacheRoot: URL?, + legacyCacheRoot: URL?, + now: Date = Date()) -> [CLIProxyAPIUsageRecord] + { + let cutoff = now.addingTimeInterval(-self.maximumRecordAge) + if self.hasLegacyCacheToMigrate( + cacheRoot: cacheRoot, + legacyCacheRoot: legacyCacheRoot) + { + do { + return try CostUsageCacheLocations.withCLIProxyAPIInterprocessLock( + stateRoot: cacheRoot?.deletingLastPathComponent()) + { + self.withExclusiveAccess { + guard let currentCache = self.loadCache( + cacheRoot: cacheRoot, + legacyCacheRoot: legacyCacheRoot) + else { return [] } + let retainedRecords = currentCache.records.filter { $0.timestamp >= cutoff } + let retainedCache = Cache(records: retainedRecords) + if retainedCache != currentCache { + _ = self.save(retainedCache, cacheRoot: cacheRoot) + } + return retainedRecords + } + } + } catch { + return self.withExclusiveAccess { + self.loadCache( + cacheRoot: cacheRoot, + legacyCacheRoot: nil)?.records.filter { $0.timestamp >= cutoff } ?? [] + } + } + } + + let initialSnapshot: (records: [CLIProxyAPIUsageRecord], needsPruning: Bool) = self.withExclusiveAccess { + guard let existingCache = self.loadCache( + cacheRoot: cacheRoot, + legacyCacheRoot: nil) + else { return ([], false) } + let retainedRecords = existingCache.records.filter { $0.timestamp >= cutoff } + return (retainedRecords, retainedRecords.count != existingCache.records.count) + } + guard initialSnapshot.needsPruning else { return initialSnapshot.records } + + do { + return try CostUsageCacheLocations.withCLIProxyAPIInterprocessLock( + stateRoot: cacheRoot?.deletingLastPathComponent()) + { + self.withExclusiveAccess { + guard let currentCache = self.loadCache( + cacheRoot: cacheRoot, + legacyCacheRoot: nil) + else { return [] } + let retainedRecords = currentCache.records.filter { $0.timestamp >= cutoff } + let retainedCache = Cache(records: retainedRecords) + if retainedCache != currentCache { + _ = self.save(retainedCache, cacheRoot: cacheRoot) + } + return retainedRecords + } + } + } catch { + return initialSnapshot.records + } + } + + @discardableResult + static func merge( + _ records: [CLIProxyAPIUsageRecord], + cacheRoot: URL? = nil, + now: Date = Date()) -> Int? + { + let legacyCacheRoot = cacheRoot == nil ? self.defaultLegacyCacheRoot() : nil + return self.merge( + records, + cacheRoot: cacheRoot, + legacyCacheRoot: legacyCacheRoot, + now: now) + } + + @discardableResult + static func merge( + _ records: [CLIProxyAPIUsageRecord], + cacheRoot: URL?, + legacyCacheRoot: URL?, + now: Date = Date()) -> Int? + { + self.withExclusiveAccess { + let cutoff = now.addingTimeInterval(-self.maximumRecordAge) + guard let existingCache = self.loadCache( + cacheRoot: cacheRoot, + legacyCacheRoot: legacyCacheRoot) + else { return nil } + var byKey = self.recordsByKey(existingCache.records.filter { $0.timestamp >= cutoff }) + let priorCount = byKey.count + for (key, record) in self.recordsByKey(records.filter { $0.timestamp >= cutoff }) { + byKey[key] = record + } + let cache = Cache(records: byKey.values.sorted { $0.timestamp < $1.timestamp }) + if cache == existingCache { + return 0 + } + guard self.save(cache, cacheRoot: cacheRoot) else { return nil } + return max(0, byKey.count - priorCount) + } + } + + static func cacheFileURL(cacheRoot: URL? = nil) -> URL { + let root = cacheRoot ?? FileManager.default.urls( + for: .applicationSupportDirectory, + in: .userDomainMask).first! + .appendingPathComponent("CodexBar", isDirectory: true) + return root + .appendingPathComponent("cost-usage", isDirectory: true) + .appendingPathComponent(CostUsageCacheLocations.cliProxyAPIUsageFileName, isDirectory: false) + } + + static func legacyCacheFileURL(cacheRoot: URL? = nil) -> URL { + let root = cacheRoot ?? self.defaultLegacyCacheRoot() + return root + .appendingPathComponent("cost-usage", isDirectory: true) + .appendingPathComponent(CostUsageCacheLocations.cliProxyAPIUsageFileName, isDirectory: false) + } + + private static let decoder: JSONDecoder = { + let decoder = JSONDecoder() + decoder.dateDecodingStrategy = .custom { decoder in + let container = try decoder.singleValueContainer() + let value = try container.decode(String.self) + guard let date = CostUsageDateParser.parse(value) else { + throw DecodingError.dataCorruptedError( + in: container, + debugDescription: "Invalid CLIProxyAPI usage timestamp.") + } + return date + } + return decoder + }() + + private static let encoder: JSONEncoder = { + let encoder = JSONEncoder() + encoder.dateEncodingStrategy = .custom { date, encoder in + let formatter = ISO8601DateFormatter() + formatter.formatOptions = [.withInternetDateTime, .withFractionalSeconds] + var container = encoder.singleValueContainer() + try container.encode(formatter.string(from: date)) + } + encoder.outputFormatting = [.sortedKeys] + return encoder + }() + + private static func recordKey(_ record: CLIProxyAPIUsageRecord) -> String { + let requestID = record.requestID.trimmingCharacters(in: .whitespacesAndNewlines) + if !requestID.isEmpty { + return "request:\(requestID)" + } + let localOccurrenceID = record.localOccurrenceID?.trimmingCharacters(in: .whitespacesAndNewlines) ?? "" + if !localOccurrenceID.isEmpty { + return "occurrence:\(localOccurrenceID)" + } + let timestamp = Int64(record.timestamp.timeIntervalSince1970 * 1000) + return [ + "fallback", + String(timestamp), + record.provider, + record.model, + record.alias, + record.endpoint, + record.authType, + String(record.tokens.input), + String(record.tokens.cacheRead), + String(record.tokens.cacheCreation), + String(record.tokens.output), + ].joined(separator: ":") + } + + private static func recordsByKey( + _ records: [CLIProxyAPIUsageRecord]) -> [String: CLIProxyAPIUsageRecord] + { + var fallbackOccurrences: [String: Int] = [:] + var recordsByKey: [String: CLIProxyAPIUsageRecord] = [:] + for record in records { + let baseKey = self.recordKey(record) + let requestID = record.requestID.trimmingCharacters(in: .whitespacesAndNewlines) + let localOccurrenceID = record.localOccurrenceID? + .trimmingCharacters(in: .whitespacesAndNewlines) ?? "" + if !requestID.isEmpty || !localOccurrenceID.isEmpty { + recordsByKey[baseKey] = record + continue + } + let occurrence = fallbackOccurrences[baseKey, default: 0] + fallbackOccurrences[baseKey] = occurrence + 1 + recordsByKey["\(baseKey):occurrence:\(occurrence)"] = record + } + return recordsByKey + } + + private static func defaultLegacyCacheRoot() -> URL { + FileManager.default.urls(for: .cachesDirectory, in: .userDomainMask).first! + .appendingPathComponent("CodexBar", isDirectory: true) + } + + private static func hasLegacyCacheToMigrate( + cacheRoot: URL?, + legacyCacheRoot: URL?, + fileManager: FileManager = .default) -> Bool + { + guard let legacyCacheRoot else { return false } + let durableURL = self.cacheFileURL(cacheRoot: cacheRoot) + let legacyURL = self.legacyCacheFileURL(cacheRoot: legacyCacheRoot) + return legacyURL.standardizedFileURL != durableURL.standardizedFileURL + && fileManager.fileExists(atPath: legacyURL.path) + } + + private static func loadCache(cacheRoot: URL?, legacyCacheRoot: URL?) -> Cache? { + let durableURL = self.cacheFileURL(cacheRoot: cacheRoot) + let durableCache: Cache? + switch self.readCache(at: durableURL) { + case .missing: + durableCache = nil + case let .valid(cache): + durableCache = cache + case .invalid: + return nil + } + guard let legacyCacheRoot else { return durableCache ?? Cache() } + + let legacyURL = self.legacyCacheFileURL(cacheRoot: legacyCacheRoot) + guard legacyURL.standardizedFileURL != durableURL.standardizedFileURL else { + return durableCache ?? Cache() + } + let legacyCache: Cache + switch self.readCache(at: legacyURL) { + case let .valid(cache): + legacyCache = cache + case .missing, .invalid: + return durableCache ?? Cache() + } + + let migratedCache = self.mergedCaches(legacy: legacyCache, durable: durableCache) + if self.save(migratedCache, cacheRoot: cacheRoot) { + try? FileManager.default.removeItem(at: legacyURL) + } + return migratedCache + } + + private static func readCache(at url: URL, fileManager: FileManager = .default) -> CacheReadResult { + guard fileManager.fileExists(atPath: url.path) else { return .missing } + guard let data = try? Data(contentsOf: url), + let cache = try? self.decoder.decode(Cache.self, from: data), + cache.version == 1 + else { return .invalid } + return .valid(cache) + } + + private static func mergedCaches(legacy: Cache, durable: Cache?) -> Cache { + var byKey = self.recordsByKey(legacy.records) + for (key, record) in self.recordsByKey(durable?.records ?? []) { + byKey[key] = record + } + return Cache(records: byKey.values.sorted { $0.timestamp < $1.timestamp }) + } + + private static func save(_ cache: Cache, cacheRoot: URL?) -> Bool { + let url = self.cacheFileURL(cacheRoot: cacheRoot) + let directory = url.deletingLastPathComponent() + do { + try FileManager.default.createDirectory(at: directory, withIntermediateDirectories: true) + let data = try self.encoder.encode(cache) + try data.write(to: url, options: [.atomic]) + return true + } catch { + return false + } + } +} + +enum CLIProxyAPIUsagePendingIO { + private struct PendingBatch: Codable { + var version: Int = 1 + var records: [CLIProxyAPIUsageRecord] = [] + } + + private static let maximumRecordAge: TimeInterval = 366 * 24 * 60 * 60 + + static func load( + pendingRoot: URL? = nil, + now: Date = Date()) -> [CLIProxyAPIUsageRecord]? + { + let url = self.pendingFileURL(pendingRoot: pendingRoot) + guard FileManager.default.fileExists(atPath: url.path) else { return [] } + guard let data = try? Data(contentsOf: url), + let pendingBatch = try? self.decoder.decode(PendingBatch.self, from: data), + pendingBatch.version == 1 + else { return nil } + let cutoff = now.addingTimeInterval(-self.maximumRecordAge) + let retainedRecords = pendingBatch.records.filter { $0.timestamp >= cutoff } + if retainedRecords.count != pendingBatch.records.count { + guard retainedRecords.isEmpty + ? self.clear(pendingRoot: pendingRoot) + : self.save(retainedRecords, pendingRoot: pendingRoot) + else { return nil } + } + return retainedRecords + } + + static func save(_ records: [CLIProxyAPIUsageRecord], pendingRoot: URL? = nil) -> Bool { + let url = self.pendingFileURL(pendingRoot: pendingRoot) + do { + try FileManager.default.createDirectory( + at: url.deletingLastPathComponent(), + withIntermediateDirectories: true) + let data = try self.encoder.encode(PendingBatch(records: records)) + try data.write(to: url, options: [.atomic]) + return true + } catch { + return false + } + } + + static func clear(pendingRoot: URL? = nil) -> Bool { + let url = self.pendingFileURL(pendingRoot: pendingRoot) + guard FileManager.default.fileExists(atPath: url.path) else { return true } + do { + try FileManager.default.removeItem(at: url) + return true + } catch { + return false + } + } + + static func pendingFileURL(pendingRoot: URL? = nil) -> URL { + let root = pendingRoot ?? FileManager.default.urls( + for: .applicationSupportDirectory, + in: .userDomainMask).first! + .appendingPathComponent("CodexBar", isDirectory: true) + return root + .appendingPathComponent("cost-usage", isDirectory: true) + .appendingPathComponent(CostUsageCacheLocations.cliProxyAPIPendingFileName, isDirectory: false) + } + + private static let decoder: JSONDecoder = { + let decoder = JSONDecoder() + decoder.dateDecodingStrategy = .custom { decoder in + let container = try decoder.singleValueContainer() + let value = try container.decode(String.self) + guard let date = CostUsageDateParser.parse(value) else { + throw DecodingError.dataCorruptedError( + in: container, + debugDescription: "Invalid pending CLIProxyAPI usage timestamp.") + } + return date + } + return decoder + }() + + private static let encoder: JSONEncoder = { + let encoder = JSONEncoder() + encoder.dateEncodingStrategy = .custom { date, encoder in + let formatter = ISO8601DateFormatter() + formatter.formatOptions = [.withInternetDateTime, .withFractionalSeconds] + var container = encoder.singleValueContainer() + try container.encode(formatter.string(from: date)) + } + encoder.outputFormatting = [.sortedKeys] + return encoder + }() +} + +public struct CLIProxyAPIConnectionSettings: Codable, Equatable, Sendable { + public static let defaultBaseURL = "http://127.0.0.1:8317" + + public let baseURL: String + public let managementKey: String + + public init(baseURL: String = Self.defaultBaseURL, managementKey: String) { + self.baseURL = baseURL.trimmingCharacters(in: .whitespacesAndNewlines) + self.managementKey = managementKey.trimmingCharacters(in: .whitespacesAndNewlines) + } + + public var isConfigured: Bool { + !self.managementKey.isEmpty && self.resolvedBaseURL != nil + } + + var resolvedBaseURL: URL? { + let value = self.baseURL.isEmpty ? Self.defaultBaseURL : self.baseURL + guard let url = URL(string: value), + let scheme = url.scheme?.lowercased(), + scheme == "http" || scheme == "https", + let host = url.host?.lowercased(), + ["127.0.0.1", "::1", "localhost"].contains(host) + else { return nil } + return url + } +} + +public enum CLIProxyAPIConnectionSettingsStore { + enum StoredSettingsSnapshot: Sendable { + case found(CLIProxyAPIConnectionSettings) + case missing + case unavailable + } + + enum ArtifactDisposition: Equatable, Sendable { + case preserve + case purge + } + + struct SerializedSaveOperations: Sendable { + let isDisconnected: @Sendable () -> Bool + let loadStored: @Sendable () -> StoredSettingsSnapshot + let store: @Sendable (CLIProxyAPIConnectionSettings) -> Bool + let setDisconnectedState: @Sendable (Bool) -> Bool + let restore: @Sendable (StoredSettingsSnapshot) -> Bool + } + + struct SerializedRemovalOperations: Sendable { + let isDisconnected: @Sendable () -> Bool + let loadStored: @Sendable () -> StoredSettingsSnapshot + let clearConfiguration: @Sendable () -> Bool + let setDisconnectedState: @Sendable (Bool) -> Bool + let restore: @Sendable (StoredSettingsSnapshot) -> Bool + } + + struct SerializedRemovalSnapshot: Sendable { + let wasDisconnected: Bool + let storedSettings: StoredSettingsSnapshot + } + + private static let key = KeychainCacheStore.Key( + category: "integration", + identifier: "cliproxyapi-management") + + public static func load() -> CLIProxyAPIConnectionSettings? { + guard case let .found(settings) = self.loadResult() else { return nil } + return settings + } + + public static func loadResult() -> KeychainCacheStore.LoadResult { + self.loadResult( + isDisconnected: { CostUsageCacheLocations.isCLIProxyAPIExplicitlyDisconnected() }, + loadStored: { KeychainCacheStore.load(key: self.key, as: CLIProxyAPIConnectionSettings.self) }) + } + + static func loadResult( + isDisconnected: () -> Bool, + loadStored: () -> KeychainCacheStore.LoadResult) + -> KeychainCacheStore.LoadResult + { + guard !isDisconnected() else { return .missing } + return loadStored() + } + + static func load( + isDisconnected: () -> Bool, + loadStored: () -> CLIProxyAPIConnectionSettings?) -> CLIProxyAPIConnectionSettings? + { + guard !isDisconnected() else { return nil } + return loadStored() + } + + @discardableResult + public static func save(_ settings: CLIProxyAPIConnectionSettings) -> Bool { + let fileManager = FileManager.default + let directories = CostUsageCacheLocations.directories(fileManager: fileManager) + return self.saveSerialized( + settings, + artifactDirectories: directories, + stateRoot: directories[1].deletingLastPathComponent(), + fileManager: fileManager, + operations: SerializedSaveOperations( + isDisconnected: { CostUsageCacheLocations.isCLIProxyAPIExplicitlyDisconnected() }, + loadStored: { + self.storedSettingsSnapshot(from: KeychainCacheStore.load( + key: self.key, + as: CLIProxyAPIConnectionSettings.self)) + }, + store: { KeychainCacheStore.storeResult(key: self.key, entry: $0) }, + setDisconnectedState: { CostUsageCacheLocations.setCLIProxyAPIExplicitlyDisconnected($0) }, + restore: { storedSettings in + switch storedSettings { + case let .found(previousSettings): + KeychainCacheStore.storeResult(key: self.key, entry: previousSettings) + case .missing: + KeychainCacheStore.clear(key: self.key) + case .unavailable: + false + } + })) + } + + static func storedSettingsSnapshot( + from result: KeychainCacheStore.LoadResult) -> StoredSettingsSnapshot + { + switch result { + case let .found(settings): .found(settings) + case .missing: .missing + case .temporarilyUnavailable, .invalid: .unavailable + } + } + + static func artifactDisposition( + _ settings: CLIProxyAPIConnectionSettings, + isDisconnected: Bool, + storedSettings: StoredSettingsSnapshot) -> ArtifactDisposition? + { + switch storedSettings { + case let .found(currentSettings): + if !isDisconnected, currentSettings == settings { + return .preserve + } + return .purge + case .missing: + return .purge + case .unavailable: + return nil + } + } + + static func saveSerialized( + _ settings: CLIProxyAPIConnectionSettings, + artifactDirectories: [URL] = [], + stateRoot: URL?, + fileManager: FileManager, + operations: SerializedSaveOperations) -> Bool + { + guard settings.isConfigured else { return false } + do { + return try CostUsageCacheLocations.withCLIProxyAPIInterprocessLock( + stateRoot: stateRoot, + fileManager: fileManager) + { + let wasDisconnected = operations.isDisconnected() + let storedSettings = operations.loadStored() + guard let artifactDisposition = self.artifactDisposition( + settings, + isDisconnected: wasDisconnected, + storedSettings: storedSettings) + else { return false } + guard let generationUpdate = CostUsageCacheLocations + .prepareCLIProxyAPIConfigurationGenerationUpdate( + stateRoot: stateRoot, + fileManager: fileManager) + else { return false } + defer { + CostUsageCacheLocations.discardCLIProxyAPIConfigurationGenerationUpdate( + generationUpdate, + fileManager: fileManager) + } + let artifactsUpdate: CostUsageCacheLocations.CLIProxyAPIArtifactsUpdate? + switch artifactDisposition { + case .preserve: + artifactsUpdate = nil + case .purge: + guard let update = CostUsageCacheLocations.prepareCLIProxyAPIArtifactsUpdate( + in: artifactDirectories, + stateRoot: stateRoot, + expectedGeneration: generationUpdate.generation, + fileManager: fileManager, + disconnectedStateAfterCommit: false, + disconnectedStateAfterRollback: wasDisconnected, + prepareState: { operations.setDisconnectedState(true) }) + else { + _ = operations.setDisconnectedState(wasDisconnected) + return false + } + artifactsUpdate = update + } + + func rollback() { + if let artifactsUpdate { + guard CostUsageCacheLocations.markCLIProxyAPIArtifactsUpdateForRollback( + artifactsUpdate, + fileManager: fileManager) + else { return } + } + guard operations.restore(storedSettings) else { return } + if let artifactsUpdate { + guard CostUsageCacheLocations.markCLIProxyAPIArtifactsRollbackCredentialsRestored( + artifactsUpdate, + fileManager: fileManager) + else { return } + } + guard operations.setDisconnectedState(wasDisconnected) else { return } + if let artifactsUpdate { + _ = CostUsageCacheLocations.restoreCLIProxyAPIArtifactsUpdate( + artifactsUpdate, + fileManager: fileManager) + } + } + + guard CostUsageCacheLocations.commitCLIProxyAPIConfigurationGenerationUpdate( + generationUpdate, + fileManager: fileManager) + else { + rollback() + return false + } + // Publish the telemetry invalidation before replacing credentials. If the process exits + // during the Keychain write, recovery will discard the staged artifacts instead of exposing + // telemetry collected under the previous credentials with the replacement configuration. + guard operations.store(settings) else { + rollback() + return false + } + guard operations.setDisconnectedState(false) else { + rollback() + return false + } + if let artifactsUpdate { + CostUsageCacheLocations.discardCLIProxyAPIArtifactsUpdate( + artifactsUpdate, + fileManager: fileManager) + } + return true + } + } catch { + return false + } + } + + @discardableResult + public static func clear() -> Bool { + do { + return try CostUsageCacheLocations.withCLIProxyAPIInterprocessLock(stateRoot: nil) { + self.clearUnserialized() + } + } catch { + return false + } + } + + public static func removeAndPurgeTelemetry() -> CLIProxyAPIConfigurationRemovalResult { + let fileManager = FileManager.default + let directories = CostUsageCacheLocations.directories(fileManager: fileManager) + return self.removeAndPurgeTelemetry( + in: directories, + stateRoot: directories[1].deletingLastPathComponent(), + fileManager: fileManager, + operations: SerializedRemovalOperations( + isDisconnected: { CostUsageCacheLocations.isCLIProxyAPIExplicitlyDisconnected() }, + loadStored: { + self.storedSettingsSnapshot(from: KeychainCacheStore.load( + key: self.key, + as: CLIProxyAPIConnectionSettings.self)) + }, + clearConfiguration: { self.clearUnserialized() }, + setDisconnectedState: { CostUsageCacheLocations.setCLIProxyAPIExplicitlyDisconnected($0) }, + restore: { storedSettings in + switch storedSettings { + case let .found(previousSettings): + KeychainCacheStore.storeResult(key: self.key, entry: previousSettings) + case .missing: + KeychainCacheStore.clear(key: self.key) + case .unavailable: + false + } + })) + } + + static func removeAndPurgeTelemetry( + in directories: [URL], + stateRoot: URL?, + fileManager: FileManager, + operations: SerializedRemovalOperations) -> CLIProxyAPIConfigurationRemovalResult + { + do { + return try CostUsageCacheLocations.withCLIProxyAPIInterprocessLock( + stateRoot: stateRoot, + fileManager: fileManager) + { + let wasDisconnected = operations.isDisconnected() + let storedSettings = operations.loadStored() + if case .unavailable = storedSettings { return .configurationRemovalFailed } + return self.removeAndPurgeTelemetryUnserialized( + in: directories, + stateRoot: stateRoot, + fileManager: fileManager, + snapshot: .init( + wasDisconnected: wasDisconnected, + storedSettings: storedSettings), + operations: operations) + } + } catch { + return .configurationRemovalFailed + } + } + + private static func removeAndPurgeTelemetryUnserialized( + in directories: [URL], + stateRoot: URL?, + fileManager: FileManager, + snapshot: SerializedRemovalSnapshot, + operations: SerializedRemovalOperations) -> CLIProxyAPIConfigurationRemovalResult + { + guard let generationUpdate = CostUsageCacheLocations + .prepareCLIProxyAPIConfigurationGenerationUpdate( + stateRoot: stateRoot, + fileManager: fileManager) + else { return .configurationRemovalFailed } + defer { + CostUsageCacheLocations.discardCLIProxyAPIConfigurationGenerationUpdate( + generationUpdate, + fileManager: fileManager) + } + guard let artifactsUpdate = CostUsageCacheLocations.prepareCLIProxyAPIArtifactsUpdate( + in: directories, + stateRoot: stateRoot, + expectedGeneration: generationUpdate.generation, + fileManager: fileManager) + else { return .configurationRemovalFailed } + + func rollback() { + guard CostUsageCacheLocations.markCLIProxyAPIArtifactsUpdateForRollback( + artifactsUpdate, + fileManager: fileManager), + operations.restore(snapshot.storedSettings), + CostUsageCacheLocations.markCLIProxyAPIArtifactsRollbackCredentialsRestored( + artifactsUpdate, + fileManager: fileManager), + operations.setDisconnectedState(snapshot.wasDisconnected) + else { return } + _ = CostUsageCacheLocations.restoreCLIProxyAPIArtifactsUpdate( + artifactsUpdate, + fileManager: fileManager) + } + + guard CostUsageCacheLocations.commitCLIProxyAPIConfigurationGenerationUpdate( + generationUpdate, + fileManager: fileManager) + else { + rollback() + return .configurationRemovalFailed + } + // Publish the telemetry invalidation before deleting credentials. Crash recovery can then only + // finalize the purge; it must never restore data after the integration has been removed. + guard operations.clearConfiguration() else { + rollback() + return .configurationRemovalFailed + } + return CostUsageCacheLocations.discardCLIProxyAPIArtifactsUpdate( + artifactsUpdate, + fileManager: fileManager) ? .removed : .telemetryCleanupFailed + } + + private static func clearUnserialized() -> Bool { + self.clearUnserialized( + isDisconnected: { CostUsageCacheLocations.isCLIProxyAPIExplicitlyDisconnected() }, + setDisconnectedState: { CostUsageCacheLocations.setCLIProxyAPIExplicitlyDisconnected($0) }, + clearConfiguration: { KeychainCacheStore.clearResult(key: self.key) }) + } + + static func clearUnserialized( + isDisconnected: () -> Bool, + setDisconnectedState: (Bool) -> Bool, + clearConfiguration: () -> KeychainCacheStore.ClearResult) -> Bool + { + let wasDisconnected = isDisconnected() + guard wasDisconnected || setDisconnectedState(true) else { + return false + } + + switch clearConfiguration() { + case .removed, .missing: + return true + case .failed: + if !wasDisconnected { + _ = setDisconnectedState(false) + } + return false + } + } +} + +public enum CLIProxyAPIConfigurationRemovalResult: Equatable, Sendable { + case removed + case configurationRemovalFailed + case telemetryCleanupFailed +} + +public enum CLIProxyAPIUsageCollectionResult: Equatable, Sendable { + case disabled + case notConfigured + case collected(Int) + case failed(String) +} + +private actor CLIProxyAPIUsageCollectionGate { + private var isLocked = false + private var waiters: [CheckedContinuation] = [] + + func perform(_ operation: @Sendable () async -> T) async -> T { + await self.acquire() + let result = await operation() + self.release() + return result + } + + private func acquire() async { + if !self.isLocked { + self.isLocked = true + return + } + await withCheckedContinuation { continuation in + self.waiters.append(continuation) + } + } + + private func release() { + guard !self.waiters.isEmpty else { + self.isLocked = false + return + } + self.waiters.removeFirst().resume() + } +} + +public enum CLIProxyAPIUsageCollector { + private static let maximumBatches = 10 + private static let batchSize = 100 + private static let collectionGate = CLIProxyAPIUsageCollectionGate() + + @discardableResult + public static func pruneExpiredUsage(now: Date = Date()) -> Bool { + self.pruneExpiredUsage( + cacheRoot: nil, + pendingRoot: nil, + stateRoot: nil, + now: now) + } + + @discardableResult + static func pruneExpiredUsage( + cacheRoot: URL?, + pendingRoot: URL?, + stateRoot: URL?, + now: Date, + fileManager: FileManager = .default) -> Bool + { + do { + return try CostUsageCacheLocations.withCLIProxyAPIInterprocessLock( + stateRoot: stateRoot, + fileManager: fileManager) + { + let pendingPruned = CLIProxyAPIUsagePendingIO.load(pendingRoot: pendingRoot, now: now) != nil + let durablePruned = CLIProxyAPIUsageCacheIO.pruneUnserialized(cacheRoot: cacheRoot, now: now) + return pendingPruned && durablePruned + } + } catch { + return false + } + } + + public static func collect( + cacheRoot: URL? = nil, + shouldContinue: @escaping @Sendable () async -> Bool = { true }) async + -> CLIProxyAPIUsageCollectionResult + { + await self.collect( + cacheRoot: cacheRoot, + settingsResult: CLIProxyAPIConnectionSettingsStore.loadResult(), + shouldContinue: shouldContinue) + } + + static func collect( + cacheRoot: URL? = nil, + settingsResult: KeychainCacheStore.LoadResult, + shouldContinue: @escaping @Sendable () async -> Bool = { true }) async + -> CLIProxyAPIUsageCollectionResult + { + switch settingsResult { + case let .found(settings): + await self.collect( + cacheRoot: cacheRoot, + settings: settings, + shouldContinue: shouldContinue) + case .temporarilyUnavailable: + .failed("CLIProxyAPI configuration is temporarily unavailable.") + case .missing, .invalid: + .notConfigured + } + } + + public static func collect( + cacheRoot: URL? = nil, + settings: CLIProxyAPIConnectionSettings?, + shouldContinue: @escaping @Sendable () async -> Bool = { true }) async + -> CLIProxyAPIUsageCollectionResult + { + guard let settings, settings.isConfigured else { return .notConfigured } + return await self.collect( + cacheRoot: cacheRoot, + configurationIsCurrent: { + guard !CostUsageCacheLocations.isCLIProxyAPIExplicitlyDisconnected(stateRoot: cacheRoot) + else { return false } + return switch CLIProxyAPIConnectionSettingsStore.loadResult() { + case let .found(currentSettings): currentSettings == settings + case .temporarilyUnavailable: true + case .missing, .invalid: false + } + }, + shouldContinue: shouldContinue, + client: CLIProxyAPIUsageQueueClient(settings: settings)) + } + + static func collect( + cacheRoot: URL? = nil, + pendingRoot: URL? = nil, + configurationIsCurrent: @escaping @Sendable () -> Bool = { true }, + shouldContinue: @escaping @Sendable () async -> Bool = { true }, + client: CLIProxyAPIUsageQueueClient) async -> CLIProxyAPIUsageCollectionResult + { + guard configurationIsCurrent() else { return .notConfigured } + return await self.collectionGate.perform { + do { + return try await CostUsageCacheLocations.withCLIProxyAPIInterprocessLock( + stateRoot: cacheRoot?.deletingLastPathComponent()) + { + guard configurationIsCurrent() else { return .notConfigured } + return await self.collectUnserialized( + cacheRoot: cacheRoot, + pendingRoot: pendingRoot, + configurationIsCurrent: configurationIsCurrent, + shouldContinue: shouldContinue, + client: client) + } + } catch { + return .failed("Could not lock CLIProxyAPI usage telemetry: \(error.localizedDescription)") + } + } + } + + private static func collectUnserialized( + cacheRoot: URL?, + pendingRoot: URL?, + configurationIsCurrent: @escaping @Sendable () -> Bool, + shouldContinue: @escaping @Sendable () async -> Bool, + client: CLIProxyAPIUsageQueueClient) async -> CLIProxyAPIUsageCollectionResult + { + do { + var added = 0 + let effectivePendingRoot = pendingRoot ?? cacheRoot + guard let pendingRecords = CLIProxyAPIUsagePendingIO.load(pendingRoot: effectivePendingRoot) else { + return .failed("Could not load pending CLIProxyAPI usage telemetry.") + } + if !pendingRecords.isEmpty { + guard let pendingAdded = CLIProxyAPIUsageCacheIO.merge( + pendingRecords, + cacheRoot: cacheRoot) + else { + return .failed("Could not save CLIProxyAPI usage telemetry.") + } + added += pendingAdded + guard CLIProxyAPIUsagePendingIO.clear(pendingRoot: effectivePendingRoot) else { + return .failed("Could not clear pending CLIProxyAPI usage telemetry.") + } + } + + for _ in 0.. (Data, URLResponse) + + struct PoppedBatch: Sendable { + let records: [CLIProxyAPIUsageRecord] + let receivedCount: Int + } + + private struct LossyRecord: Decodable { + let value: CLIProxyAPIUsageRecord? + + init(from decoder: Decoder) { + self.value = try? CLIProxyAPIUsageRecord(from: decoder) + } + } + + private static let log = CodexBarLog.logger(LogCategories.tokenCost) + + enum ClientError: LocalizedError { + case invalidBaseURL + case invalidResponse + case httpError(Int) + + var errorDescription: String? { + switch self { + case .invalidBaseURL: "Invalid CLIProxyAPI URL." + case .invalidResponse: "CLIProxyAPI returned an invalid response." + case let .httpError(status): "CLIProxyAPI returned HTTP \(status)." + } + } + } + + let settings: CLIProxyAPIConnectionSettings + let dataLoader: DataLoader + + init( + settings: CLIProxyAPIConnectionSettings, + dataLoader: @escaping DataLoader = Self.liveDataLoader) + { + self.settings = settings + self.dataLoader = dataLoader + } + + func pop(count: Int) async throws -> PoppedBatch { + guard let baseURL = self.settings.resolvedBaseURL, + var components = URLComponents( + url: baseURL.appendingPathComponent("v0/management/usage-queue"), + resolvingAgainstBaseURL: false) + else { throw ClientError.invalidBaseURL } + components.queryItems = [URLQueryItem(name: "count", value: String(max(1, count)))] + guard let url = components.url else { throw ClientError.invalidBaseURL } + + var request = URLRequest(url: url) + request.timeoutInterval = 5 + request.setValue("Bearer \(self.settings.managementKey)", forHTTPHeaderField: "Authorization") + request.setValue("application/json", forHTTPHeaderField: "Accept") + + let (data, response) = try await self.dataLoader(request) + guard let httpResponse = response as? HTTPURLResponse else { + throw ClientError.invalidResponse + } + guard (200..<300).contains(httpResponse.statusCode) else { + throw ClientError.httpError(httpResponse.statusCode) + } + let decoder = JSONDecoder() + decoder.dateDecodingStrategy = .custom { decoder in + let container = try decoder.singleValueContainer() + let value = try container.decode(String.self) + guard let date = CostUsageDateParser.parse(value) else { + throw DecodingError.dataCorruptedError( + in: container, + debugDescription: "Invalid CLIProxyAPI usage timestamp.") + } + return date + } + let decoded = try decoder.decode([LossyRecord].self, from: data) + let records = decoded.compactMap(\.value) + let malformedCount = decoded.count - records.count + if malformedCount > 0 { + Self.log.warning( + "Ignored malformed CLIProxyAPI usage records", + metadata: ["count": String(malformedCount)]) + } + return PoppedBatch(records: records, receivedCount: decoded.count) + } + + private static func liveDataLoader(_ request: URLRequest) async throws -> (Data, URLResponse) { + let configuration = URLSessionConfiguration.ephemeral + configuration.timeoutIntervalForRequest = 5 + configuration.timeoutIntervalForResource = 10 + return try await URLSession(configuration: configuration).data(for: request) + } +} diff --git a/Sources/CodexBarCore/CostUsageCacheLocations.swift b/Sources/CodexBarCore/CostUsageCacheLocations.swift new file mode 100644 index 0000000000..875b9cf81d --- /dev/null +++ b/Sources/CodexBarCore/CostUsageCacheLocations.swift @@ -0,0 +1,621 @@ +import Foundation +#if canImport(Darwin) +import Darwin +#elseif canImport(Glibc) +import Glibc +#elseif canImport(Musl) +import Musl +#endif + +public struct CostUsageCacheClearResult: Equatable, Sendable { + public let cleared: Int + public let errorDescription: String? +} + +public enum CostUsageCacheLocations { + struct CLIProxyAPIArtifactsUpdate: Sendable { + struct Move: Sendable { + let originalURL: URL + let stagedURL: URL + } + + let moves: [Move] + let manifestURL: URL? + + init(moves: [Move], manifestURL: URL? = nil) { + self.moves = moves + self.manifestURL = manifestURL + } + } + + struct CLIProxyAPIConfigurationGenerationUpdate { + let stagedURL: URL + let destinationURL: URL + let generation: String + } + + private struct CLIProxyAPIArtifactsTransactionManifest: Codable { + struct Move: Codable { + let originalPath: String + let stagedPath: String + } + + let expectedGeneration: String + let moves: [Move] + let disconnectedStateAfterCommit: Bool? + let disconnectedStateAfterRollback: Bool? + let forceRollback: Bool? + let rollbackCredentialsRestored: Bool? + } + + static let cliProxyAPIUsageFileName = "cliproxyapi-usage-v1.json" + static let cliProxyAPIPendingFileName = "cliproxyapi-pending-v1.json" + private static let cliProxyAPIDisconnectedFileName = "cliproxyapi-disconnected-v1" + private static let cliProxyAPIConfigurationGenerationFileName = "cliproxyapi-configuration-generation-v1" + private static let cliProxyAPIArtifactsTransactionFileName = "cliproxyapi-artifacts-transaction-v1.json" + + public static func directories(fileManager: FileManager = .default) -> [URL] { + let cacheRoot = fileManager.urls(for: .cachesDirectory, in: .userDomainMask).first! + let applicationSupportRoot = fileManager.urls( + for: .applicationSupportDirectory, + in: .userDomainMask).first! + return [cacheRoot, applicationSupportRoot].map { root in + root + .appendingPathComponent("CodexBar", isDirectory: true) + .appendingPathComponent("cost-usage", isDirectory: true) + } + } + + public static func clearAllCostUsageCaches( + fileManager: FileManager = .default) -> CostUsageCacheClearResult + { + self.clearAllCostUsageCaches( + in: self.directories(fileManager: fileManager), + stateRoot: nil, + fileManager: fileManager) + } + + public static func clearAllCostUsageCaches( + in directories: [URL], + stateRoot: URL?, + fileManager: FileManager = .default) -> CostUsageCacheClearResult + { + do { + return try self.withCLIProxyAPIInterprocessLock( + stateRoot: stateRoot, + fileManager: fileManager) + { + guard self.advanceCLIProxyAPIConfigurationGeneration( + stateRoot: stateRoot, + fileManager: fileManager) + else { + return CostUsageCacheClearResult( + cleared: 0, + errorDescription: CocoaError(.fileWriteUnknown).localizedDescription) + } + var cleared = 0 + for directory in directories where fileManager.fileExists(atPath: directory.path) { + do { + try fileManager.removeItem(at: directory) + cleared += 1 + } catch { + return CostUsageCacheClearResult( + cleared: cleared, + errorDescription: error.localizedDescription) + } + } + return CostUsageCacheClearResult(cleared: cleared, errorDescription: nil) + } + } catch { + return CostUsageCacheClearResult(cleared: 0, errorDescription: error.localizedDescription) + } + } + + static func withCLIProxyAPIInterprocessLock( + stateRoot: URL?, + fileManager: FileManager = .default, + operation: () throws -> T) throws -> T + { + let descriptor = try self.acquireCLIProxyAPILock(stateRoot: stateRoot, fileManager: fileManager) + defer { self.releaseCLIProxyAPILock(descriptor) } + guard self.recoverCLIProxyAPIArtifactsTransaction(stateRoot: stateRoot, fileManager: fileManager) else { + throw CocoaError(.fileReadUnknown) + } + return try operation() + } + + static func withCLIProxyAPIInterprocessLock( + stateRoot: URL?, + fileManager: FileManager = .default, + operation: () async throws -> T) async throws -> T + { + let descriptor = try self.acquireCLIProxyAPILock(stateRoot: stateRoot, fileManager: fileManager) + defer { self.releaseCLIProxyAPILock(descriptor) } + guard self.recoverCLIProxyAPIArtifactsTransaction(stateRoot: stateRoot, fileManager: fileManager) else { + throw CocoaError(.fileReadUnknown) + } + return try await operation() + } + + @discardableResult + public static func clearCLIProxyAPIArtifacts(fileManager: FileManager = .default) -> Bool { + let directories = self.directories(fileManager: fileManager) + return self.clearCLIProxyAPIArtifacts( + in: directories, + stateRoot: directories[1].deletingLastPathComponent(), + fileManager: fileManager) + } + + @discardableResult + static func clearCLIProxyAPIArtifacts( + in directories: [URL], + stateRoot: URL?, + fileManager: FileManager = .default) -> Bool + { + do { + return try self.withCLIProxyAPIInterprocessLock( + stateRoot: stateRoot, + fileManager: fileManager) + { + guard self.advanceCLIProxyAPIConfigurationGeneration( + stateRoot: stateRoot, + fileManager: fileManager) + else { return false } + return self.clearCLIProxyAPIArtifactsUnserialized( + in: directories, + fileManager: fileManager) + } + } catch { + return false + } + } + + static func clearCLIProxyAPIArtifactsUnserialized( + in directories: [URL], + fileManager: FileManager) -> Bool + { + var succeeded = true + for url in self.cliProxyAPIArtifactURLs(in: directories) { + guard fileManager.fileExists(atPath: url.path) else { continue } + do { + try fileManager.removeItem(at: url) + } catch { + succeeded = false + } + } + return succeeded + } + + static func prepareCLIProxyAPIArtifactsUpdate( + in directories: [URL], + stateRoot: URL?, + expectedGeneration: String, + fileManager: FileManager, + disconnectedStateAfterCommit: Bool? = nil, + disconnectedStateAfterRollback: Bool? = nil, + prepareState: () -> Bool = { true }) -> CLIProxyAPIArtifactsUpdate? + { + let identifier = UUID().uuidString + let moves = self.cliProxyAPIArtifactURLs(in: directories) + .filter { fileManager.fileExists(atPath: $0.path) } + .map { originalURL in + CLIProxyAPIArtifactsUpdate.Move( + originalURL: originalURL, + stagedURL: originalURL + .deletingLastPathComponent() + .appendingPathComponent( + ".\(originalURL.lastPathComponent).\(identifier).replacement-backup", + isDirectory: false)) + } + let manifestURL = self.cliProxyAPIArtifactsTransactionURL( + stateRoot: stateRoot, + fileManager: fileManager) + let manifest = CLIProxyAPIArtifactsTransactionManifest( + expectedGeneration: expectedGeneration, + moves: moves.map { + .init(originalPath: $0.originalURL.path, stagedPath: $0.stagedURL.path) + }, + disconnectedStateAfterCommit: disconnectedStateAfterCommit, + disconnectedStateAfterRollback: disconnectedStateAfterRollback, + forceRollback: nil, + rollbackCredentialsRestored: nil) + do { + try fileManager.createDirectory( + at: manifestURL.deletingLastPathComponent(), + withIntermediateDirectories: true) + try JSONEncoder().encode(manifest).write(to: manifestURL, options: [.atomic]) + } catch { + return nil + } + + let update = CLIProxyAPIArtifactsUpdate(moves: moves, manifestURL: manifestURL) + guard prepareState() else { + _ = self.removeCLIProxyAPIArtifactsManifest(manifestURL, fileManager: fileManager) + return nil + } + for move in moves { + do { + try fileManager.moveItem(at: move.originalURL, to: move.stagedURL) + } catch { + _ = self.restoreCLIProxyAPIArtifactsUpdate(update, fileManager: fileManager) + return nil + } + } + return update + } + + static func prepareCLIProxyAPIArtifactsUpdate( + in directories: [URL], + fileExists: (URL) -> Bool, + moveItem: (URL, URL) throws -> Void) -> CLIProxyAPIArtifactsUpdate? + { + let identifier = UUID().uuidString + var moves: [CLIProxyAPIArtifactsUpdate.Move] = [] + for originalURL in self.cliProxyAPIArtifactURLs(in: directories) where fileExists(originalURL) { + let stagedURL = originalURL + .deletingLastPathComponent() + .appendingPathComponent( + ".\(originalURL.lastPathComponent).\(identifier).replacement-backup", + isDirectory: false) + do { + try moveItem(originalURL, stagedURL) + moves.append(.init(originalURL: originalURL, stagedURL: stagedURL)) + } catch { + _ = self.restoreCLIProxyAPIArtifactsUpdate( + .init(moves: moves), + fileExists: fileExists, + moveItem: moveItem) + return nil + } + } + return CLIProxyAPIArtifactsUpdate(moves: moves) + } + + @discardableResult + static func restoreCLIProxyAPIArtifactsUpdate( + _ update: CLIProxyAPIArtifactsUpdate, + fileManager: FileManager) -> Bool + { + let restored = self.restoreCLIProxyAPIArtifactsUpdate( + update, + fileExists: { fileManager.fileExists(atPath: $0.path) }, + moveItem: { try fileManager.moveItem(at: $0, to: $1) }) + guard restored else { return false } + return self.removeCLIProxyAPIArtifactsManifest(update.manifestURL, fileManager: fileManager) + } + + @discardableResult + static func markCLIProxyAPIArtifactsUpdateForRollback( + _ update: CLIProxyAPIArtifactsUpdate, + fileManager: FileManager) -> Bool + { + guard let manifestURL = update.manifestURL else { return true } + guard let data = try? Data(contentsOf: manifestURL), + let manifest = try? JSONDecoder().decode(CLIProxyAPIArtifactsTransactionManifest.self, from: data) + else { return false } + guard manifest.forceRollback != true else { return true } + let rollbackManifest = CLIProxyAPIArtifactsTransactionManifest( + expectedGeneration: manifest.expectedGeneration, + moves: manifest.moves, + disconnectedStateAfterCommit: manifest.disconnectedStateAfterCommit, + disconnectedStateAfterRollback: manifest.disconnectedStateAfterRollback, + forceRollback: true, + rollbackCredentialsRestored: manifest.rollbackCredentialsRestored) + do { + try JSONEncoder().encode(rollbackManifest).write(to: manifestURL, options: [.atomic]) + return true + } catch { + return false + } + } + + @discardableResult + static func markCLIProxyAPIArtifactsRollbackCredentialsRestored( + _ update: CLIProxyAPIArtifactsUpdate, + fileManager: FileManager) -> Bool + { + guard let manifestURL = update.manifestURL else { return true } + guard let data = try? Data(contentsOf: manifestURL), + let manifest = try? JSONDecoder().decode(CLIProxyAPIArtifactsTransactionManifest.self, from: data) + else { return false } + guard manifest.forceRollback == true else { return false } + guard manifest.rollbackCredentialsRestored != true else { return true } + let restoredManifest = CLIProxyAPIArtifactsTransactionManifest( + expectedGeneration: manifest.expectedGeneration, + moves: manifest.moves, + disconnectedStateAfterCommit: manifest.disconnectedStateAfterCommit, + disconnectedStateAfterRollback: manifest.disconnectedStateAfterRollback, + forceRollback: true, + rollbackCredentialsRestored: true) + do { + try JSONEncoder().encode(restoredManifest).write(to: manifestURL, options: [.atomic]) + return true + } catch { + return false + } + } + + @discardableResult + static func restoreCLIProxyAPIArtifactsUpdate( + _ update: CLIProxyAPIArtifactsUpdate, + fileExists: (URL) -> Bool, + moveItem: (URL, URL) throws -> Void) -> Bool + { + var succeeded = true + for move in update.moves.reversed() where fileExists(move.stagedURL) { + guard !fileExists(move.originalURL) else { + succeeded = false + continue + } + do { + try moveItem(move.stagedURL, move.originalURL) + } catch { + succeeded = false + } + } + return succeeded + } + + @discardableResult + static func discardCLIProxyAPIArtifactsUpdate( + _ update: CLIProxyAPIArtifactsUpdate, + fileManager: FileManager) -> Bool + { + var succeeded = true + for move in update.moves { + guard fileManager.fileExists(atPath: move.stagedURL.path) else { continue } + do { + try fileManager.removeItem(at: move.stagedURL) + } catch { + succeeded = false + } + } + guard succeeded else { return false } + return self.removeCLIProxyAPIArtifactsManifest(update.manifestURL, fileManager: fileManager) + } + + @discardableResult + static func recoverCLIProxyAPIArtifactsTransaction( + stateRoot: URL?, + fileManager: FileManager = .default) -> Bool + { + let manifestURL = self.cliProxyAPIArtifactsTransactionURL( + stateRoot: stateRoot, + fileManager: fileManager) + guard fileManager.fileExists(atPath: manifestURL.path) else { return true } + guard let data = try? Data(contentsOf: manifestURL), + let manifest = try? JSONDecoder().decode(CLIProxyAPIArtifactsTransactionManifest.self, from: data) + else { return false } + let update = CLIProxyAPIArtifactsUpdate( + moves: manifest.moves.map { + .init( + originalURL: URL(fileURLWithPath: $0.originalPath), + stagedURL: URL(fileURLWithPath: $0.stagedPath)) + }, + manifestURL: manifestURL) + if manifest.forceRollback == true, manifest.rollbackCredentialsRestored != true { + guard self.setCLIProxyAPIExplicitlyDisconnected( + true, + stateRoot: stateRoot, + fileManager: fileManager) + else { return false } + return self.discardCLIProxyAPIArtifactsUpdate(update, fileManager: fileManager) + } + let didCommit = manifest.forceRollback != true && + self.cliProxyAPIConfigurationGeneration(stateRoot: stateRoot, fileManager: fileManager) == + manifest.expectedGeneration + let disconnectedState = didCommit + ? manifest.disconnectedStateAfterCommit + : manifest.disconnectedStateAfterRollback + if let disconnectedState, + !self.setCLIProxyAPIExplicitlyDisconnected( + disconnectedState, + stateRoot: stateRoot, + fileManager: fileManager) + { + return false + } + if didCommit { + return self.discardCLIProxyAPIArtifactsUpdate(update, fileManager: fileManager) + } + return self.restoreCLIProxyAPIArtifactsUpdate(update, fileManager: fileManager) + } + + private static func removeCLIProxyAPIArtifactsManifest( + _ manifestURL: URL?, + fileManager: FileManager) -> Bool + { + guard let manifestURL, fileManager.fileExists(atPath: manifestURL.path) else { return true } + do { + try fileManager.removeItem(at: manifestURL) + return true + } catch { + return false + } + } + + private static func cliProxyAPIArtifactURLs(in directories: [URL]) -> [URL] { + var seenPaths: Set = [] + return directories.flatMap { directory in + [ + directory.appendingPathComponent(self.cliProxyAPIUsageFileName, isDirectory: false), + directory.appendingPathComponent(self.cliProxyAPIPendingFileName, isDirectory: false), + CostUsageCacheIO.cacheFileURL( + provider: .claude, + cacheRoot: directory.deletingLastPathComponent()), + ] + }.filter { seenPaths.insert($0.path).inserted } + } + + public static func isCLIProxyAPIExplicitlyDisconnected( + stateRoot: URL? = nil, + fileManager: FileManager = .default) -> Bool + { + fileManager.fileExists(atPath: self.cliProxyAPIDisconnectedURL( + stateRoot: stateRoot, + fileManager: fileManager).path) + } + + public static func cliProxyAPIConfigurationGeneration( + stateRoot: URL? = nil, + fileManager: FileManager = .default) -> String? + { + let url = self.cliProxyAPIConfigurationGenerationURL( + stateRoot: stateRoot, + fileManager: fileManager) + guard let data = try? Data(contentsOf: url), + let generation = String(data: data, encoding: .utf8), + !generation.isEmpty + else { return nil } + return generation + } + + static func prepareCLIProxyAPIConfigurationGenerationUpdate( + stateRoot: URL? = nil, + fileManager: FileManager = .default) -> CLIProxyAPIConfigurationGenerationUpdate? + { + let destinationURL = self.cliProxyAPIConfigurationGenerationURL( + stateRoot: stateRoot, + fileManager: fileManager) + let stagedURL = destinationURL + .deletingLastPathComponent() + .appendingPathComponent(".cliproxyapi-generation-\(UUID().uuidString).tmp", isDirectory: false) + let generation = UUID().uuidString + do { + try fileManager.createDirectory( + at: destinationURL.deletingLastPathComponent(), + withIntermediateDirectories: true) + try Data(generation.utf8).write(to: stagedURL, options: [.atomic]) + return CLIProxyAPIConfigurationGenerationUpdate( + stagedURL: stagedURL, + destinationURL: destinationURL, + generation: generation) + } catch { + try? fileManager.removeItem(at: stagedURL) + return nil + } + } + + static func commitCLIProxyAPIConfigurationGenerationUpdate( + _ update: CLIProxyAPIConfigurationGenerationUpdate, + fileManager: FileManager = .default) -> Bool + { + rename(update.stagedURL.path, update.destinationURL.path) == 0 + } + + static func discardCLIProxyAPIConfigurationGenerationUpdate( + _ update: CLIProxyAPIConfigurationGenerationUpdate, + fileManager: FileManager = .default) + { + try? fileManager.removeItem(at: update.stagedURL) + } + + private static func advanceCLIProxyAPIConfigurationGeneration( + stateRoot: URL?, + fileManager: FileManager) -> Bool + { + guard let update = self.prepareCLIProxyAPIConfigurationGenerationUpdate( + stateRoot: stateRoot, + fileManager: fileManager) + else { return false } + guard self.commitCLIProxyAPIConfigurationGenerationUpdate(update, fileManager: fileManager) else { + self.discardCLIProxyAPIConfigurationGenerationUpdate(update, fileManager: fileManager) + return false + } + return true + } + + @discardableResult + static func setCLIProxyAPIExplicitlyDisconnected( + _ disconnected: Bool, + stateRoot: URL? = nil, + fileManager: FileManager = .default) -> Bool + { + let url = self.cliProxyAPIDisconnectedURL( + stateRoot: stateRoot, + fileManager: fileManager) + if disconnected { + guard !fileManager.fileExists(atPath: url.path) else { return true } + do { + try fileManager.createDirectory( + at: url.deletingLastPathComponent(), + withIntermediateDirectories: true) + try Data().write(to: url, options: [.atomic]) + return true + } catch { + return false + } + } + + guard fileManager.fileExists(atPath: url.path) else { return true } + do { + try fileManager.removeItem(at: url) + return true + } catch { + return false + } + } + + private static func cliProxyAPIDisconnectedURL( + stateRoot: URL?, + fileManager: FileManager) -> URL + { + let root = stateRoot ?? fileManager.urls( + for: .applicationSupportDirectory, + in: .userDomainMask).first! + .appendingPathComponent("CodexBar", isDirectory: true) + return root.appendingPathComponent(self.cliProxyAPIDisconnectedFileName, isDirectory: false) + } + + private static func cliProxyAPIConfigurationGenerationURL( + stateRoot: URL?, + fileManager: FileManager) -> URL + { + let root = stateRoot ?? fileManager.urls( + for: .applicationSupportDirectory, + in: .userDomainMask).first! + .appendingPathComponent("CodexBar", isDirectory: true) + return root.appendingPathComponent( + self.cliProxyAPIConfigurationGenerationFileName, + isDirectory: false) + } + + private static func cliProxyAPIArtifactsTransactionURL( + stateRoot: URL?, + fileManager: FileManager) -> URL + { + let root = stateRoot ?? fileManager.urls( + for: .applicationSupportDirectory, + in: .userDomainMask).first! + .appendingPathComponent("CodexBar", isDirectory: true) + return root.appendingPathComponent(self.cliProxyAPIArtifactsTransactionFileName, isDirectory: false) + } + + private static func acquireCLIProxyAPILock( + stateRoot: URL?, + fileManager: FileManager) throws -> Int32 + { + let root = stateRoot ?? fileManager.urls( + for: .applicationSupportDirectory, + in: .userDomainMask).first! + .appendingPathComponent("CodexBar", isDirectory: true) + try fileManager.createDirectory(at: root, withIntermediateDirectories: true) + let lockURL = root.appendingPathComponent("cliproxyapi-collection.lock", isDirectory: false) + let descriptor = open(lockURL.path, O_CREAT | O_RDWR | O_CLOEXEC, S_IRUSR | S_IWUSR) + guard descriptor >= 0 else { + throw POSIXError(POSIXErrorCode(rawValue: errno) ?? .EIO) + } + while flock(descriptor, LOCK_EX) != 0 { + guard errno == EINTR else { + close(descriptor) + throw POSIXError(POSIXErrorCode(rawValue: errno) ?? .EIO) + } + } + return descriptor + } + + private static func releaseCLIProxyAPILock(_ descriptor: Int32) { + _ = flock(descriptor, LOCK_UN) + close(descriptor) + } +} diff --git a/Sources/CodexBarCore/CostUsageFetcher.swift b/Sources/CodexBarCore/CostUsageFetcher.swift index 0662e18f7d..7c4ce6a4a1 100644 --- a/Sources/CodexBarCore/CostUsageFetcher.swift +++ b/Sources/CodexBarCore/CostUsageFetcher.swift @@ -1,5 +1,7 @@ import Foundation +// swiftlint:disable file_length + public enum CostUsageError: LocalizedError, Sendable { case unsupportedProvider(UsageProvider) case timedOut(seconds: Int) @@ -67,7 +69,7 @@ public struct CostUsageFetcher: Sendable { private let scannerOptions: CostUsageScanner.Options? public init(cacheRoot: URL? = nil) { - self.scannerOptions = cacheRoot.map { CostUsageScanner.Options(cacheRoot: $0) } + self.scannerOptions = Self.defaultScannerOptions(cacheRoot: cacheRoot) } init(scannerOptions: CostUsageScanner.Options) { @@ -165,7 +167,8 @@ public struct CostUsageFetcher: Sendable { cursorCookieHeaderOverride: String? = nil, allowPricingRefresh: Bool = true, refreshPricingInBackground: Bool = true, - includePiSessions: Bool = true) async throws -> CostUsageTokenSnapshot + includePiSessions: Bool = true, + includeClaudeProxyUsage: Bool = true) async throws -> CostUsageTokenSnapshot { try await Self.loadTokenSnapshot( provider: provider, @@ -179,6 +182,7 @@ public struct CostUsageFetcher: Sendable { allowPricingRefresh: allowPricingRefresh, refreshPricingInBackground: refreshPricingInBackground, includePiSessions: includePiSessions, + includeClaudeProxyUsage: includeClaudeProxyUsage, bypassScannerDebounce: false, scannerOptions: self.scannerOptionsOverride()) } @@ -195,6 +199,7 @@ public struct CostUsageFetcher: Sendable { allowPricingRefresh: Bool = true, refreshPricingInBackground: Bool = true, includePiSessions: Bool = true, + includeClaudeProxyUsage: Bool = true, bypassScannerDebounce: Bool) async throws -> CostUsageTokenSnapshot { try await Self.loadTokenSnapshot( @@ -209,10 +214,46 @@ public struct CostUsageFetcher: Sendable { allowPricingRefresh: allowPricingRefresh, refreshPricingInBackground: refreshPricingInBackground, includePiSessions: includePiSessions, + includeClaudeProxyUsage: includeClaudeProxyUsage, bypassScannerDebounce: bypassScannerDebounce, scannerOptions: self.scannerOptionsOverride()) } + package func loadCodexProxyTokenSnapshot( + now: Date = Date(), + forceRefresh: Bool = false, + historyDays: Int = 30, + allowPricingRefresh: Bool = true, + refreshPricingInBackground: Bool = true) async throws -> CostUsageTokenSnapshot + { + try await self.loadCodexProxyTokenSnapshot( + now: now, + forceRefresh: forceRefresh, + historyDays: historyDays, + allowPricingRefresh: allowPricingRefresh, + refreshPricingInBackground: refreshPricingInBackground, + modelsDevClient: ModelsDevClient()) + } + + func loadCodexProxyTokenSnapshot( + now: Date, + forceRefresh: Bool, + historyDays: Int = 30, + allowPricingRefresh: Bool = true, + refreshPricingInBackground: Bool, + modelsDevClient: ModelsDevClient) async throws -> CostUsageTokenSnapshot + { + try await Self.loadCodexProxyTokenSnapshot(CodexProxyTokenSnapshotOptions( + now: now, + forceRefresh: forceRefresh, + historyDays: historyDays, + allowPricingRefresh: allowPricingRefresh, + refreshPricingInBackground: refreshPricingInBackground, + scannerOptions: self.scannerOptionsOverride(), + modelsDevClient: modelsDevClient, + retryUnknownPricing: true)) + } + @available(*, deprecated, message: "Codex token-cost scans are uncapped; this limit is ignored.") public func loadTokenSnapshot( provider: UsageProvider, @@ -242,6 +283,14 @@ public struct CostUsageFetcher: Sendable { self.scannerOptions } + package func cliProxyAPIConfigurationGeneration() -> String? { + let options = Self.resolvedScannerOptions( + self.scannerOptionsOverride(), + provider: .codex, + codexHomePath: nil) + return CostUsageCacheLocations.cliProxyAPIConfigurationGeneration(stateRoot: options.cacheRoot) + } + package func codexScanCatchUpStatus( codexHomePath: String? = nil) async -> CodexScanCatchUpStatus { @@ -333,12 +382,15 @@ public struct CostUsageFetcher: Sendable { staleSnapshotUpdatedAt: pending ? cache.codexPreviousReport?.updatedAt : nil) } - private static func resolvedScannerOptions( + static func resolvedScannerOptions( _ override: CostUsageScanner.Options?, provider: UsageProvider, codexHomePath: String?) -> CostUsageScanner.Options { var options = override ?? CostUsageScanner.Options() + if override == nil { + options.cliProxyAPIHome = Self.defaultCLIProxyAPIHome() + } if provider == .codex, let codexHomePath = codexHomePath?.trimmingCharacters(in: .whitespacesAndNewlines), !codexHomePath.isEmpty @@ -349,6 +401,23 @@ public struct CostUsageFetcher: Sendable { return options } + static func defaultScannerOptions( + cacheRoot: URL?, + homeDirectory: URL = FileManager.default.homeDirectoryForCurrentUser) -> CostUsageScanner.Options? + { + cacheRoot.map { + CostUsageScanner.Options( + cacheRoot: $0, + cliProxyAPIHome: Self.defaultCLIProxyAPIHome(homeDirectory: homeDirectory)) + } + } + + private static func defaultCLIProxyAPIHome( + homeDirectory: URL = FileManager.default.homeDirectoryForCurrentUser) -> URL + { + homeDirectory.appendingPathComponent(".cli-proxy-api", isDirectory: true) + } + static func loadTokenSnapshot( provider: UsageProvider, environment: [String: String] = ProcessInfo.processInfo.environment, @@ -361,6 +430,7 @@ public struct CostUsageFetcher: Sendable { allowPricingRefresh: Bool = true, refreshPricingInBackground: Bool = true, includePiSessions: Bool = true, + includeClaudeProxyUsage: Bool = true, bypassScannerDebounce: Bool = false, scannerOptions overrideScannerOptions: CostUsageScanner.Options? = nil, piScannerOptions overridePiScannerOptions: PiSessionCostScanner @@ -388,10 +458,14 @@ public struct CostUsageFetcher: Sendable { overrideScannerOptions, provider: provider, codexHomePath: codexHomePath) + let cliProxyAPIAttributionEnabled = Self.isCLIProxyAPIAttributionEnabled(options: options) + if !cliProxyAPIAttributionEnabled { + options.cliProxyAPIHome = nil + } // Rolling window is inclusive, so a 30-day display starts 29 days before `now`. let since = options.calendar.date(byAdding: .day, value: -(clampedHistoryDays - 1), to: now) ?? now let scopedCodexHomePath = codexHomePath?.trimmingCharacters(in: .whitespacesAndNewlines) - let shouldMergePiUsage = provider != .codex || scopedCodexHomePath?.isEmpty != false + let shouldMergeGlobalCodexUsage = provider != .codex || scopedCodexHomePath?.isEmpty != false await Self.refreshPricingIfAllowed( options: PricingRefreshOptions( provider: provider, @@ -408,6 +482,9 @@ public struct CostUsageFetcher: Sendable { allowVertexClaudeFallback: allowVertexClaudeFallback, forceRefresh: forceRefresh, bypassScannerDebounce: bypassScannerDebounce) + if provider == .claude { + options.claudeAttributionFilter = cliProxyAPIAttributionEnabled ? .excludeCodexBackend : .all + } var resolvedPiOptions = overridePiScannerOptions ?? PiSessionCostScanner.Options() if resolvedPiOptions.cacheRoot == nil { resolvedPiOptions.cacheRoot = options.cacheRoot @@ -422,7 +499,8 @@ public struct CostUsageFetcher: Sendable { let localScanOptions = LocalTokenScanOptions( allowVertexClaudeFallback: allowVertexClaudeFallback, includePiSessions: includePiSessions, - shouldMergePiUsage: shouldMergePiUsage, + includeClaudeProxyUsage: includeClaudeProxyUsage, + shouldMergeGlobalCodexUsage: shouldMergeGlobalCodexUsage, scanOptions: scanOptions, piOptions: piOptions) let scanResult = try await Self.loadLocalTokenScanResult( @@ -453,6 +531,7 @@ public struct CostUsageFetcher: Sendable { allowPricingRefresh: allowPricingRefresh, refreshPricingInBackground: false, includePiSessions: includePiSessions, + includeClaudeProxyUsage: includeClaudeProxyUsage, scannerOptions: options, piScannerOptions: piOptions, modelsDevClient: modelsDevClient, @@ -479,7 +558,8 @@ public struct CostUsageFetcher: Sendable { private struct LocalTokenScanOptions: Sendable { let allowVertexClaudeFallback: Bool let includePiSessions: Bool - let shouldMergePiUsage: Bool + let includeClaudeProxyUsage: Bool + let shouldMergeGlobalCodexUsage: Bool let scanOptions: CostUsageScanner.Options let piOptions: PiSessionCostScanner.Options } @@ -523,6 +603,7 @@ public struct CostUsageFetcher: Sendable { var projects: [CostUsageProjectBreakdown] = [] var sessions: [CostUsageSessionBreakdown] = [] var piDaily: CostUsageDailyReport? + var claudeProxyDaily: CostUsageDailyReport? var staleSnapshotUpdatedAt: Date? if provider == .codex { let roots = CostUsageScanner.codexSessionsRoots(options: options.scanOptions) @@ -531,6 +612,15 @@ public struct CostUsageFetcher: Sendable { scopedTo: roots) let range = CostUsageScanner.CostUsageDayRange( since: since, until: now, calendar: options.scanOptions.calendar) + let supplemental = try Self.loadCodexSupplementalScan( + options: options.scanOptions, + range: range, + now: now, + includeClaudeProxy: options.includeClaudeProxyUsage + && options.shouldMergeGlobalCodexUsage, + checkCancellation: checkCancellation) + claudeProxyDaily = supplemental.claudeProxyDaily + daily = claudeProxyDaily.map { daily.merged(with: $0) } ?? daily if let previous = CostUsageScanner.codexPreviousReport( cache: cache, range: range, @@ -538,19 +628,12 @@ public struct CostUsageFetcher: Sendable { { staleSnapshotUpdatedAt = previous.updatedAt } else { - projects = CostUsageScanner.buildCodexProjectBreakdownsFromCache( - cache: cache, - range: range, - modelsDevCacheRoot: options.scanOptions.cacheRoot) - sessions = CostUsageScanner.buildCodexSessionBreakdownsFromCache( - cache: cache, - range: range, - modelsDevCacheRoot: options.scanOptions.cacheRoot, - sessionRoots: roots) + projects = supplemental.projects + sessions = supplemental.sessions } } if options.includePiSessions, - provider == .claude || (provider == .codex && options.shouldMergePiUsage) + provider == .claude || (provider == .codex && options.shouldMergeGlobalCodexUsage) { let piReport = try PiSessionCostScanner.loadDailyReportCancellable( provider: provider, @@ -566,11 +649,13 @@ public struct CostUsageFetcher: Sendable { daily = CostUsageDailyReport.merged([daily, piReport]) } if provider == .codex { - projects = Self.mergedProjectBreakdowns( - projects + [piDaily.flatMap(Self.unknownProjectBreakdown(from:))].compactMap(\.self)) - if piDaily?.data.isEmpty == false { - sessions = [] - } + let finalized = Self.finalizeCodexSupplementalScan( + projects: projects, + sessions: sessions, + claudeProxyDaily: claudeProxyDaily, + piDaily: piDaily) + projects = finalized.projects + sessions = finalized.sessions } return LocalTokenScanResult( daily: daily, @@ -608,8 +693,7 @@ public struct CostUsageFetcher: Sendable { } private struct UnknownPricingRefreshRequest: Sendable { - let providerID: String - let modelIDs: Set + let modelIDsByProviderID: [String: Set] let now: Date let cacheRoot: URL? let client: ModelsDevClient @@ -623,48 +707,107 @@ public struct CostUsageFetcher: Sendable { client: ModelsDevClient) -> UnknownPricingRefreshRequest? { guard provider == .codex || provider == .claude else { return nil } - let unknownModelIDs = Set(daily.data.flatMap { entry in - entry.modelBreakdowns?.compactMap { breakdown -> String? in - guard breakdown.costUSD == nil else { return nil } + var modelIDsByProviderID: [String: Set] = [:] + for entry in daily.data { + for breakdown in entry.modelBreakdowns ?? [] { + guard breakdown.costUSD == nil else { continue } + let upstreamModel = breakdown.attribution?.route == .cliProxyAPI + ? breakdown.attribution?.upstream?.model?.trimmingCharacters(in: .whitespacesAndNewlines) + : nil + let pricingModel = upstreamModel.flatMap { $0.isEmpty ? nil : $0 } ?? breakdown.modelName if provider == .codex, - CostUsagePricing.isCodexUnattributedModel(breakdown.modelName) + CostUsagePricing.isCodexUnattributedModel(pricingModel) { - return nil + continue + } + let providerIDs = Self.modelsDevProviderIDs( + pricingModel: pricingModel, + attribution: breakdown.attribution, + fallbackProvider: provider, + cacheRoot: cacheRoot) + for providerID in providerIDs { + modelIDsByProviderID[providerID, default: []].insert(pricingModel) } - return breakdown.modelName - } ?? [] - }) - guard !unknownModelIDs.isEmpty else { return nil } + } + } + guard !modelIDsByProviderID.isEmpty else { return nil } return UnknownPricingRefreshRequest( - providerID: provider == .codex ? "openai" : "anthropic", - modelIDs: unknownModelIDs, + modelIDsByProviderID: modelIDsByProviderID, now: now, cacheRoot: cacheRoot, client: client) } + private static func modelsDevProviderIDs( + pricingModel: String, + attribution: CostUsageAttribution?, + fallbackProvider: UsageProvider, + cacheRoot: URL?) -> Set + { + let knownProviderID: String? = switch CostUsagePricing.modelProvider( + for: pricingModel, + modelsDevCacheRoot: cacheRoot) + { + case .openAI: "openai" + case .anthropic: "anthropic" + case .google: "google" + case .unknown: nil + } + if let knownProviderID { + return [knownProviderID] + } + + if attribution?.route == .cliProxyAPI { + let upstreamProvider = attribution?.upstream?.provider + .trimmingCharacters(in: .whitespacesAndNewlines) + .lowercased() + switch upstreamProvider { + case "codex", "openai": return ["openai"] + case "claude", "anthropic": return ["anthropic"] + case "aistudio", "gemini", "gemini-interactions", "google", "vertex": return ["google"] + default: break + } + return switch attribution?.upstream?.executorType?.lowercased() { + case "codexexecutor": ["openai"] + case "claudeexecutor": ["anthropic"] + case "geminiexecutor": ["google"] + case "openaicompatexecutor": ["openai", "anthropic", "google"] + default: [fallbackProvider == .codex ? "openai" : "anthropic"] + } + } + return [fallbackProvider == .codex ? "openai" : "anthropic"] + } + private static func refreshUnknownPricingIfNeeded( _ request: UnknownPricingRefreshRequest, inBackground: Bool) async -> Bool { if inBackground { Task.detached(priority: .utility) { - _ = await ModelsDevPricingPipeline.refreshForUnknownModelsIfNeeded( - providerID: request.providerID, - modelIDs: request.modelIDs, - now: request.now, - cacheRoot: request.cacheRoot, - client: request.client) + for providerID in request.modelIDsByProviderID.keys.sorted() { + _ = await ModelsDevPricingPipeline.refreshForUnknownModelsIfNeeded( + providerID: providerID, + modelIDs: request.modelIDsByProviderID[providerID] ?? [], + now: request.now, + cacheRoot: request.cacheRoot, + client: request.client) + } } return false } - return await ModelsDevPricingPipeline.refreshForUnknownModelsIfNeeded( - providerID: request.providerID, - modelIDs: request.modelIDs, - now: request.now, - cacheRoot: request.cacheRoot, - client: request.client) == .pricingAvailable + + var pricingAvailable = false + for providerID in request.modelIDsByProviderID.keys.sorted() { + let result = await ModelsDevPricingPipeline.refreshForUnknownModelsIfNeeded( + providerID: providerID, + modelIDs: request.modelIDsByProviderID[providerID] ?? [], + now: request.now, + cacheRoot: request.cacheRoot, + client: request.client) + pricingAvailable = result == .pricingAvailable || pricingAvailable + } + return pricingAvailable } static func loadCachedCodexTokenSnapshot( @@ -686,6 +829,7 @@ public struct CostUsageFetcher: Sendable { scannerOptions: overrideScannerOptions)?.snapshot } + // swiftlint:disable:next function_body_length static func loadCachedCodexTokenSnapshotResult( now: Date = Date(), codexHomePath: String? = nil, @@ -701,9 +845,11 @@ public struct CostUsageFetcher: Sendable { return nil } + typealias CachedResult = CachedCodexTokenSnapshotResult? + // Decoding the persisted scan cache parses multi-megabyte JSON; keep it off the // cooperative pool alongside the scans themselves. - let cachedSnapshot: CachedCodexTokenSnapshotResult?? = try? await CostUsageScanExecutor.run { _ in + let cachedSnapshot: CachedResult? = try? await CostUsageScanExecutor.run { checkCancellation in let clampedHistoryDays = max(1, min(365, historyDays)) let options = Self.resolvedScannerOptions( overrideScannerOptions, @@ -735,6 +881,7 @@ public struct CostUsageFetcher: Sendable { var nativeScanAt: Date? var scanTimes: [Date] = [] var piMerged = false + var claudeProxyMerged = false var staleSnapshotUpdatedAt: Date? if let previous = CostUsageScanner.codexPreviousReport( @@ -798,6 +945,24 @@ public struct CostUsageFetcher: Sendable { } } + if shouldMergePiUsage, + let proxy = try Self.loadCachedCodexProxyReport( + options: options, + range: range, + now: now, + checkCancellation: checkCancellation) + { + reports.append(proxy.report) + claudeProxyMerged = true + if let scanAt = proxy.scanAt { + scanTimes.append(scanAt) + } + if includeProjectAndSessionBreakdowns, let project = proxy.project { + projects.append(project) + sessions = [] + } + } + if includePiSessions, shouldMergePiUsage, let piResult = PiSessionCostScanner.loadCachedDailyReportResult( @@ -813,10 +978,12 @@ public struct CostUsageFetcher: Sendable { if let piLastScanAt = piResult.lastScanAt { scanTimes.append(piLastScanAt) } - if let piProject = Self.unknownProjectBreakdown(from: piResult.report) { + if includeProjectAndSessionBreakdowns, + let piProject = Self.unknownProjectBreakdown(from: piResult.report) + { projects.append(piProject) } - if !piResult.report.data.isEmpty { + if includeProjectAndSessionBreakdowns, !piResult.report.data.isEmpty { sessions = [] } } @@ -835,12 +1002,61 @@ public struct CostUsageFetcher: Sendable { projects: Self.mergedProjectBreakdowns(projects), sessions: sessions, updatedAt: scanTimes.min()), - lastRefreshAt: piMerged || staleSnapshotUpdatedAt != nil ? nil : nativeScanAt, + lastRefreshAt: piMerged || claudeProxyMerged || staleSnapshotUpdatedAt != nil ? nil : nativeScanAt, staleSnapshotUpdatedAt: staleSnapshotUpdatedAt) } return cachedSnapshot.flatMap(\.self) } + private static func loadCachedCodexProxyReport( + options: CostUsageScanner.Options, + range: CostUsageScanner.CostUsageDayRange, + now: Date, + checkCancellation: @escaping CostUsageScanner.CancellationCheck) throws -> ( + report: CostUsageDailyReport, + scanAt: Date?, + project: CostUsageProjectBreakdown?)? + { + guard self.isCLIProxyAPIAttributionEnabled(options: options) else { return nil } + let claudeCache = CostUsageCacheIO.load( + provider: .claude, + cacheRoot: options.cacheRoot, + calendar: range.calendar) + guard !claudeCache.days.isEmpty, + !CostUsageScanner.requestedWindowExpandsCache(range: range, cache: claudeCache) + else { return nil } + + let attributionResolver: CLIProxyAPIAttributionResolver? = if let home = options.cliProxyAPIHome { + try CLIProxyAPIAttributionResolver.load( + home: home, + cacheRoot: options.cacheRoot, + forceReload: options.forceRescan, + checkCancellation: checkCancellation) + } else { + nil + } + let report = CostUsageScanner.buildClaudeReportFromCache( + cache: claudeCache, + range: range, + attributionFilter: .codexBackendOnly, + attributionResolver: attributionResolver, + modelsDevCatalog: CostUsagePricing.modelsDevCatalog( + now: now, + cacheRoot: options.cacheRoot), + modelsDevCacheRoot: options.cacheRoot) + guard !report.data.isEmpty else { return nil } + + let scanAt = claudeCache.lastScanUnixMs > 0 + ? Date(timeIntervalSince1970: TimeInterval(claudeCache.lastScanUnixMs) / 1000) + : nil + return ( + report, + scanAt, + self.unknownProjectBreakdown( + from: report, + name: "Claude Code via CLIProxyAPI")) + } + /// Providers whose token-cost snapshot `loadTokenSnapshot` can produce. Cursor is /// macOS-only because it reuses the macOS Cursor session resolution. static func supportsTokenSnapshot(_ provider: UsageProvider) -> Bool { @@ -1033,6 +1249,214 @@ public struct CostUsageFetcher: Sendable { sessions: sessions, updatedAt: updatedAt ?? now) } +} + +extension CostUsageFetcher { + private struct CodexProxyTokenSnapshotOptions { + let now: Date + let forceRefresh: Bool + let historyDays: Int + let allowPricingRefresh: Bool + let refreshPricingInBackground: Bool + let scannerOptions: CostUsageScanner.Options? + let modelsDevClient: ModelsDevClient + let retryUnknownPricing: Bool + } + + private struct CodexSupplementalScan { + let projects: [CostUsageProjectBreakdown] + let sessions: [CostUsageSessionBreakdown] + let claudeProxyDaily: CostUsageDailyReport? + } + + private static func loadCodexProxyTokenSnapshot( + _ request: CodexProxyTokenSnapshotOptions) async throws -> CostUsageTokenSnapshot + { + let clampedHistoryDays = max(1, min(365, request.historyDays)) + let since = Calendar.current.date( + byAdding: .day, + value: -(clampedHistoryDays - 1), + to: request.now) ?? request.now + var options = Self.resolvedScannerOptions( + request.scannerOptions, + provider: .codex, + codexHomePath: nil) + await Self.refreshPricingIfAllowed( + options: PricingRefreshOptions( + provider: .codex, + isAllowed: request.allowPricingRefresh, + retryUnknown: request.retryUnknownPricing, + inBackground: request.refreshPricingInBackground), + now: request.now, + cacheRoot: options.cacheRoot, + client: request.modelsDevClient) + if request.forceRefresh { + options.refreshMinIntervalSeconds = 0 + } + + let scanOptions = options + let proxyDaily = try await CostUsageScanExecutor.run { checkCancellation in + let range = CostUsageScanner.CostUsageDayRange(since: since, until: request.now) + let supplemental = try Self.loadCodexSupplementalScan( + options: scanOptions, + range: range, + now: request.now, + includeClaudeProxy: true, + checkCancellation: checkCancellation) + return supplemental.claudeProxyDaily + } + let daily = proxyDaily ?? CostUsageDailyReport(data: [], summary: nil) + let projects = proxyDaily.flatMap { + Self.unknownProjectBreakdown( + from: $0, + name: "Claude Code via CLIProxyAPI") + }.map { [$0] } ?? [] + if request.allowPricingRefresh, + request.retryUnknownPricing, + let refreshRequest = Self.unknownPricingRefreshRequest( + provider: .codex, + daily: daily, + now: request.now, + cacheRoot: options.cacheRoot, + client: request.modelsDevClient), + await Self.refreshUnknownPricingIfNeeded( + refreshRequest, + inBackground: request.refreshPricingInBackground) + { + return try await Self.loadCodexProxyTokenSnapshot(CodexProxyTokenSnapshotOptions( + now: request.now, + forceRefresh: request.forceRefresh, + historyDays: request.historyDays, + allowPricingRefresh: request.allowPricingRefresh, + refreshPricingInBackground: false, + scannerOptions: options, + modelsDevClient: request.modelsDevClient, + retryUnknownPricing: false)) + } + return Self.tokenSnapshot( + from: daily, + now: request.now, + historyDays: clampedHistoryDays, + projects: projects) + } + + private static func loadCodexSupplementalScan( + options: CostUsageScanner.Options, + range: CostUsageScanner.CostUsageDayRange, + now: Date, + includeClaudeProxy: Bool, + checkCancellation: @escaping CostUsageScanner.CancellationCheck) throws -> CodexSupplementalScan + { + let roots = CostUsageScanner.codexSessionsRoots(options: options) + let cache = CostUsageScanner.codexCache( + CostUsageCacheIO.load(provider: .codex, cacheRoot: options.cacheRoot), + scopedTo: roots) + let projects = CostUsageScanner.buildCodexProjectBreakdownsFromCache( + cache: cache, + range: range, + modelsDevCacheRoot: options.cacheRoot) + let sessions = CostUsageScanner.buildCodexSessionBreakdownsFromCache( + cache: cache, + range: range, + modelsDevCacheRoot: options.cacheRoot, + sessionRoots: roots) + + guard includeClaudeProxy, + self.hasCodexProxyEvidence(options: options) + else { + return CodexSupplementalScan( + projects: projects, + sessions: sessions, + claudeProxyDaily: nil) + } + var proxyOptions = options + proxyOptions.claudeLogProviderFilter = .excludeVertexAI + proxyOptions.claudeAttributionFilter = .codexBackendOnly + let proxyDaily = try CostUsageScanner.loadClaudeDaily( + provider: .claude, + range: range, + now: now, + options: proxyOptions, + checkCancellation: checkCancellation) + return CodexSupplementalScan( + projects: projects, + sessions: sessions, + claudeProxyDaily: proxyDaily.data.isEmpty ? nil : proxyDaily) + } + + static func hasCodexProxyEvidence( + options: CostUsageScanner.Options, + fileManager: FileManager = .default) -> Bool + { + guard self.isCLIProxyAPIAttributionEnabled(options: options, fileManager: fileManager) else { return false } + + if CLIProxyAPIUsageCacheIO.load(cacheRoot: options.cacheRoot).contains(where: { + $0.provider.caseInsensitiveCompare("codex") == .orderedSame + }) { + return true + } + + let cachedClaude = CostUsageCacheIO.load(provider: .claude, cacheRoot: options.cacheRoot) + if cachedClaude.files.values.contains(where: { usage in + usage.claudeRows?.contains { + $0.attribution?.route == .cliProxyAPI + && $0.attribution?.upstream?.isCodex == true + } == true + }) { + return true + } + + guard let home = options.cliProxyAPIHome else { return false } + if CLIProxyAPIAttributionResolver.hasCodexOAuthModelAliasRoute( + home: home, + fileManager: fileManager) + { + return true + } + let logDirectory = home.appendingPathComponent("logs", isDirectory: true) + guard let urls = try? fileManager.contentsOfDirectory( + at: logDirectory, + includingPropertiesForKeys: [.isRegularFileKey], + options: [.skipsHiddenFiles]) + else { return false } + return urls.contains { url in + guard url.pathExtension.caseInsensitiveCompare("log") == .orderedSame, + let values = try? url.resourceValues(forKeys: [.isRegularFileKey]) + else { return false } + return values.isRegularFile == true + } + } + + private static func isCLIProxyAPIAttributionEnabled( + options: CostUsageScanner.Options, + fileManager: FileManager = .default) -> Bool + { + !CostUsageCacheLocations.isCLIProxyAPIExplicitlyDisconnected( + stateRoot: options.cacheRoot, + fileManager: fileManager) + } + + private static func finalizeCodexSupplementalScan( + projects: [CostUsageProjectBreakdown], + sessions: [CostUsageSessionBreakdown], + claudeProxyDaily: CostUsageDailyReport?, + piDaily: CostUsageDailyReport?) -> ( + projects: [CostUsageProjectBreakdown], + sessions: [CostUsageSessionBreakdown]) + { + let mergedProjects = Self.mergedProjectBreakdowns( + projects + [ + claudeProxyDaily.flatMap { + Self.unknownProjectBreakdown( + from: $0, + name: "Claude Code via CLIProxyAPI") + }, + piDaily.flatMap { Self.unknownProjectBreakdown(from: $0) }, + ].compactMap(\.self)) + let hasUnattributedSessions = piDaily?.data.isEmpty == false + || claudeProxyDaily?.data.isEmpty == false + return (mergedProjects, hasUnattributedSessions ? [] : sessions) + } package static func resolvedCodexScanDurationPerRefresh( provider: UsageProvider, @@ -1071,10 +1495,13 @@ public struct CostUsageFetcher: Sendable { configuredDuration: options.maxCodexScanDurationPerRefresh) } - private static func unknownProjectBreakdown(from daily: CostUsageDailyReport) -> CostUsageProjectBreakdown? { + private static func unknownProjectBreakdown( + from daily: CostUsageDailyReport, + name: String = CostUsageProjectBreakdown.unknownProjectName) -> CostUsageProjectBreakdown? + { guard !daily.data.isEmpty else { return nil } return CostUsageProjectBreakdown( - name: CostUsageProjectBreakdown.unknownProjectName, + name: name, path: nil, totalTokens: daily.summary?.totalTokens, totalCostUSD: daily.summary?.totalCostUSD, @@ -1082,7 +1509,7 @@ public struct CostUsageFetcher: Sendable { modelBreakdowns: self.projectModelBreakdowns(from: daily.data), sources: [ CostUsageProjectSourceBreakdown( - name: CostUsageProjectBreakdown.unknownProjectName, + name: name, path: nil, totalTokens: daily.summary?.totalTokens, totalCostUSD: daily.summary?.totalCostUSD, @@ -1094,12 +1521,12 @@ public struct CostUsageFetcher: Sendable { private static func mergedProjectBreakdowns( _ projects: [CostUsageProjectBreakdown]) -> [CostUsageProjectBreakdown] { - var dailyByPath: [String: [CostUsageDailyReport]] = [:] - var namesByPath: [String: String] = [:] - var sourceDailyByProjectPath: [String: [String: [CostUsageDailyReport]]] = [:] - var sourceNamesByProjectPath: [String: [String: String]] = [:] + var dailyByPath: [ProjectMergeKey: [CostUsageDailyReport]] = [:] + var namesByPath: [ProjectMergeKey: String] = [:] + var sourceDailyByProjectPath: [ProjectMergeKey: [ProjectMergeKey: [CostUsageDailyReport]]] = [:] + var sourceNamesByProjectPath: [ProjectMergeKey: [ProjectMergeKey: String]] = [:] for project in projects { - let key = project.path ?? "" + let key = ProjectMergeKey(name: project.name, path: project.path) namesByPath[key] = project.name dailyByPath[key, default: []].append(CostUsageDailyReport(data: project.daily, summary: nil)) let sources = project.sources.isEmpty @@ -1114,7 +1541,7 @@ public struct CostUsageFetcher: Sendable { ] : project.sources for source in sources { - let sourceKey = source.path ?? "" + let sourceKey = ProjectMergeKey(name: source.name, path: source.path) sourceNamesByProjectPath[key, default: [:]][sourceKey] = source.name sourceDailyByProjectPath[key, default: [:]][sourceKey, default: []] .append(CostUsageDailyReport(data: source.daily, summary: nil)) @@ -1124,7 +1551,7 @@ public struct CostUsageFetcher: Sendable { let merged = CostUsageDailyReport.merged(reports) return CostUsageProjectBreakdown( name: namesByPath[key] ?? CostUsageProjectBreakdown.unknownProjectName, - path: key.isEmpty ? nil : key, + path: key.path, totalTokens: merged.summary?.totalTokens, totalCostUSD: merged.summary?.totalCostUSD, daily: merged.data, @@ -1148,15 +1575,25 @@ public struct CostUsageFetcher: Sendable { } } + private struct ProjectMergeKey: Hashable { + let path: String? + let syntheticName: String? + + init(name: String, path: String?) { + self.path = path + self.syntheticName = path == nil ? name : nil + } + } + private static func mergedProjectSources( - sourceDailyByPath: [String: [CostUsageDailyReport]], - sourceNamesByPath: [String: String]) -> [CostUsageProjectSourceBreakdown] + sourceDailyByPath: [ProjectMergeKey: [CostUsageDailyReport]], + sourceNamesByPath: [ProjectMergeKey: String]) -> [CostUsageProjectSourceBreakdown] { sourceDailyByPath.map { key, reports in let merged = CostUsageDailyReport.merged(reports) return CostUsageProjectSourceBreakdown( name: sourceNamesByPath[key] ?? CostUsageProjectBreakdown.unknownProjectName, - path: key.isEmpty ? nil : key, + path: key.path, totalTokens: merged.summary?.totalTokens, totalCostUSD: merged.summary?.totalCostUSD, daily: merged.data, @@ -1194,28 +1631,42 @@ public struct CostUsageFetcher: Sendable { } } - func build(modelName: String) -> CostUsageDailyReport.ModelBreakdown { + func build( + modelName: String, + attribution: CostUsageAttribution?) -> CostUsageDailyReport.ModelBreakdown + { CostUsageDailyReport.ModelBreakdown( modelName: modelName, costUSD: self.sawCost ? self.costUSD : nil, - totalTokens: self.sawTotalTokens ? self.totalTokens : nil) + totalTokens: self.sawTotalTokens ? self.totalTokens : nil, + attribution: attribution) } } - private static func projectModelBreakdowns( + private struct ProjectBreakdownKey: Hashable { + let modelName: String + let attribution: CostUsageAttribution? + } + + static func projectModelBreakdowns( from entries: [CostUsageDailyReport.Entry]) -> [CostUsageDailyReport.ModelBreakdown]? { - var accumulators: [String: ProjectBreakdownAccumulator] = [:] + var accumulators: [ProjectBreakdownKey: ProjectBreakdownAccumulator] = [:] for entry in entries { for breakdown in entry.modelBreakdowns ?? [] { - var accumulator = accumulators[breakdown.modelName] ?? ProjectBreakdownAccumulator() + let key = ProjectBreakdownKey( + modelName: breakdown.modelName, + attribution: breakdown.attribution) + var accumulator = accumulators[key] ?? ProjectBreakdownAccumulator() accumulator.add(breakdown) - accumulators[breakdown.modelName] = accumulator + accumulators[key] = accumulator } } guard !accumulators.isEmpty else { return nil } - return accumulators.map { modelName, accumulator in - accumulator.build(modelName: modelName) + return accumulators.map { key, accumulator in + accumulator.build( + modelName: key.modelName, + attribution: key.attribution) } .sorted { lhs, rhs in let lhsCost = lhs.costUSD ?? -1 @@ -1228,10 +1679,17 @@ public struct CostUsageFetcher: Sendable { if lhsTokens != rhsTokens { return lhsTokens > rhsTokens } - return lhs.modelName > rhs.modelName + if lhs.modelName != rhs.modelName { + return lhs.modelName > rhs.modelName + } + let lhsAttribution = lhs.attribution?.deterministicSortKey ?? "" + let rhsAttribution = rhs.attribution?.deterministicSortKey ?? "" + return lhsAttribution > rhsAttribution } } +} +extension CostUsageFetcher { static func selectCurrentSession(from sessions: [CostUsageSessionReport.Entry]) -> CostUsageSessionReport.Entry? { diff --git a/Sources/CodexBarCore/CostUsageModels.swift b/Sources/CodexBarCore/CostUsageModels.swift index 88f2760bbb..e3e8162801 100644 --- a/Sources/CodexBarCore/CostUsageModels.swift +++ b/Sources/CodexBarCore/CostUsageModels.swift @@ -272,6 +272,132 @@ public struct CostUsageProjectSourceBreakdown: Sendable, Equatable { } } +public struct CostUsageAttribution: Sendable, Codable, Equatable, Hashable { + public enum Client: String, Sendable, Codable, Hashable { + case claudeCode + } + + public enum Route: String, Sendable, Codable, Hashable { + case unknown + case cliProxyAPI + } + + public enum ModelProvider: String, Sendable, Codable, Hashable { + case openAI + case anthropic + case google + case unknown + } + + public struct Upstream: Sendable, Codable, Equatable, Hashable { + public enum AuthType: String, Sendable, Codable, Hashable { + case oauth + case apiKey + case unknown + } + + public let provider: String + public let authType: AuthType + public let model: String? + public let executorType: String? + + public init( + provider: String, + authType: AuthType, + model: String? = nil, + executorType: String? = nil) + { + self.provider = provider + self.authType = authType + self.model = model + self.executorType = executorType + } + + public var isCodex: Bool { + self.provider.trimmingCharacters(in: .whitespacesAndNewlines) + .caseInsensitiveCompare("codex") == .orderedSame + } + + public var providerDisplayName: String { + switch self.provider.trimmingCharacters(in: .whitespacesAndNewlines).lowercased() { + case "codex": "Codex" + case "claude": "Claude" + case "gemini": "Gemini" + case "gemini-interactions": "Gemini Interactions" + case "aistudio": "AI Studio" + case "vertex": "Vertex AI" + case "antigravity": "Antigravity" + case "xai": "xAI" + case "kimi": "Kimi" + case "openrouter": "OpenRouter" + case let provider where provider.isEmpty: "Unknown" + case let provider: provider + } + } + + public var authDisplayName: String? { + switch self.authType { + case .oauth: "OAuth" + case .apiKey: "API key" + case .unknown: nil + } + } + + public var displayName: String { + guard let authDisplayName else { return self.providerDisplayName } + return "\(self.providerDisplayName) \(authDisplayName)" + } + } + + public enum Evidence: String, Sendable, Codable, Hashable { + case cliProxyAuthInventory + case cliProxyModelAlias + case modelProvider + case cliProxyRequestLog + case cliProxyUsageTelemetry + } + + public let client: Client + public let route: Route + public let modelProvider: ModelProvider + public let upstream: Upstream? + public let evidence: [Evidence] + + public init( + client: Client, + route: Route, + modelProvider: ModelProvider = .unknown, + upstream: Upstream? = nil, + evidence: [Evidence] = []) + { + self.client = client + self.route = route + self.modelProvider = modelProvider + self.upstream = upstream + self.evidence = evidence + } + + package var deterministicSortKey: String { + let fields = [ + self.client.rawValue, + self.route.rawValue, + self.modelProvider.rawValue, + self.upstream == nil ? "0" : "1", + self.upstream?.provider ?? "", + self.upstream?.authType.rawValue ?? "", + self.upstream?.model == nil ? "0" : "1", + self.upstream?.model ?? "", + self.upstream?.executorType == nil ? "0" : "1", + self.upstream?.executorType ?? "", + String(self.evidence.count), + ] + self.evidence.map(\.rawValue) + + return fields + .map { "\($0.utf8.count):\($0)" } + .joined() + } +} + public struct CostUsageDailyReport: Sendable, Decodable { public struct ModelBreakdown: Sendable, Decodable, Equatable { public let modelName: String @@ -282,6 +408,7 @@ public struct CostUsageDailyReport: Sendable, Decodable { public let priorityCostUSD: Double? public let standardTokens: Int? public let priorityTokens: Int? + public let attribution: CostUsageAttribution? private enum CodingKeys: String, CodingKey { case modelName @@ -294,6 +421,7 @@ public struct CostUsageDailyReport: Sendable, Decodable { case priorityCostUSD case standardTokens case priorityTokens + case attribution } public init(from decoder: Decoder) throws { @@ -310,6 +438,7 @@ public struct CostUsageDailyReport: Sendable, Decodable { self.priorityCostUSD = try container.decodeIfPresent(Double.self, forKey: .priorityCostUSD) self.standardTokens = try container.decodeIfPresent(Int.self, forKey: .standardTokens) self.priorityTokens = try container.decodeIfPresent(Int.self, forKey: .priorityTokens) + self.attribution = try container.decodeIfPresent(CostUsageAttribution.self, forKey: .attribution) } public init( @@ -320,7 +449,8 @@ public struct CostUsageDailyReport: Sendable, Decodable { standardCostUSD: Double? = nil, priorityCostUSD: Double? = nil, standardTokens: Int? = nil, - priorityTokens: Int? = nil) + priorityTokens: Int? = nil, + attribution: CostUsageAttribution? = nil) { self.modelName = modelName self.costUSD = costUSD @@ -330,6 +460,7 @@ public struct CostUsageDailyReport: Sendable, Decodable { self.priorityCostUSD = priorityCostUSD self.standardTokens = standardTokens self.priorityTokens = priorityTokens + self.attribution = attribution } } @@ -527,6 +658,11 @@ public struct CostUsageDailyReport: Sendable, Decodable { } extension CostUsageDailyReport { + private struct BreakdownKey: Hashable { + let modelName: String + let attribution: CostUsageAttribution? + } + private struct BreakdownAccumulator { var totalTokens: Int = 0 var sawTotalTokens = false @@ -568,15 +704,16 @@ extension CostUsageDailyReport { } } - func build(modelName: String) -> ModelBreakdown { + func build(key: BreakdownKey) -> ModelBreakdown { ModelBreakdown( - modelName: modelName, + modelName: key.modelName, costUSD: self.sawCost ? self.costUSD : nil, totalTokens: self.sawTotalTokens ? self.totalTokens : nil, standardCostUSD: self.sawStandardCost ? self.standardCostUSD : nil, priorityCostUSD: self.sawPriorityCost ? self.priorityCostUSD : nil, standardTokens: self.sawStandardTokens ? self.standardTokens : nil, - priorityTokens: self.sawPriorityTokens ? self.priorityTokens : nil) + priorityTokens: self.sawPriorityTokens ? self.priorityTokens : nil, + attribution: key.attribution) } } @@ -595,7 +732,7 @@ extension CostUsageDailyReport { var costUSD: Double = 0 var sawCost = false var modelsUsed: Set = [] - var breakdowns: [String: BreakdownAccumulator] = [:] + var breakdowns: [BreakdownKey: BreakdownAccumulator] = [:] mutating func add(_ entry: Entry) { let entryDerivedTotalTokens = (entry.inputTokens ?? 0) @@ -633,9 +770,12 @@ extension CostUsageDailyReport { } if let modelBreakdowns = entry.modelBreakdowns { for breakdown in modelBreakdowns { - var accumulator = self.breakdowns[breakdown.modelName] ?? BreakdownAccumulator() + let key = BreakdownKey( + modelName: breakdown.modelName, + attribution: breakdown.attribution) + var accumulator = self.breakdowns[key] ?? BreakdownAccumulator() accumulator.add(breakdown) - self.breakdowns[breakdown.modelName] = accumulator + self.breakdowns[key] = accumulator self.modelsUsed.insert(breakdown.modelName) } } @@ -657,8 +797,8 @@ extension CostUsageDailyReport { guard !self.breakdowns.isEmpty else { return nil } return CostUsageDailyReport.sortedModelBreakdowns( self.breakdowns - .map { modelName, accumulator in - accumulator.build(modelName: modelName) + .map { key, accumulator in + accumulator.build(key: key) }) }() let modelsUsed = self.modelsUsed.isEmpty ? nil : self.modelsUsed.sorted() @@ -767,7 +907,12 @@ extension CostUsageDailyReport { return lhsTokens > rhsTokens } - return lhs.modelName > rhs.modelName + if lhs.modelName != rhs.modelName { + return lhs.modelName > rhs.modelName + } + let lhsAttribution = lhs.attribution?.deterministicSortKey ?? "" + let rhsAttribution = rhs.attribution?.deterministicSortKey ?? "" + return lhsAttribution > rhsAttribution } } } diff --git a/Sources/CodexBarCore/Generated/CodexParserHash.generated.swift b/Sources/CodexBarCore/Generated/CodexParserHash.generated.swift index 7401a1244a..cca23d26a2 100644 --- a/Sources/CodexBarCore/Generated/CodexParserHash.generated.swift +++ b/Sources/CodexBarCore/Generated/CodexParserHash.generated.swift @@ -1,5 +1,5 @@ // Generated by Scripts/regenerate-codex-parser-hash.sh. Do not edit by hand. enum CodexParserHash { - static let value = "843ca061c36bbea1" + static let value = "31c724c0918cdfe4" } diff --git a/Sources/CodexBarCore/Vendored/CostUsage/CostUsagePricing+ModelProvider.swift b/Sources/CodexBarCore/Vendored/CostUsage/CostUsagePricing+ModelProvider.swift new file mode 100644 index 0000000000..687eba810f --- /dev/null +++ b/Sources/CodexBarCore/Vendored/CostUsage/CostUsagePricing+ModelProvider.swift @@ -0,0 +1,42 @@ +import Foundation + +extension CostUsagePricing { + static func modelProvider( + for model: String, + modelsDevCatalog: ModelsDevCatalog? = nil, + modelsDevCacheRoot: URL? = nil) -> CostUsageAttribution.ModelProvider + { + if self.isOpenAIModel( + model, + modelsDevCatalog: modelsDevCatalog, + modelsDevCacheRoot: modelsDevCacheRoot) + { + return .openAI + } + + if self.claudeCostUSD( + model: model, + inputTokens: 0, + cacheReadInputTokens: 0, + cacheCreationInputTokens: 0, + outputTokens: 0, + modelsDevCatalog: modelsDevCatalog, + modelsDevCacheRoot: modelsDevCacheRoot) != nil + { + return .anthropic + } + + let trimmed = model.trimmingCharacters(in: .whitespacesAndNewlines) + if trimmed.lowercased().hasPrefix("gemini-") + || modelsDevCatalog?.pricing(providerID: "google", modelID: trimmed) != nil + || ModelsDevPricingPipeline.lookup( + providerID: "google", + modelID: trimmed, + cacheRoot: modelsDevCacheRoot) != nil + { + return .google + } + + return .unknown + } +} diff --git a/Sources/CodexBarCore/Vendored/CostUsage/CostUsagePricing.swift b/Sources/CodexBarCore/Vendored/CostUsage/CostUsagePricing.swift index 1bd0c5e624..32171187f5 100644 --- a/Sources/CodexBarCore/Vendored/CostUsage/CostUsagePricing.swift +++ b/Sources/CodexBarCore/Vendored/CostUsage/CostUsagePricing.swift @@ -451,6 +451,22 @@ enum CostUsagePricing { self.normalizeCodexModel(raw) == self.codexUnattributedModel } + static func isOpenAIModel( + _ model: String, + modelsDevCatalog: ModelsDevCatalog? = nil, + modelsDevCacheRoot: URL? = nil) -> Bool + { + let normalized = self.normalizeCodexModel(model) + if normalized != self.codexUnattributedModel, self.codex[normalized] != nil { + return true + } + return self.modelsDevLookup( + providerID: self.codexModelsDevProviderID, + model: model, + catalog: modelsDevCatalog, + cacheRoot: modelsDevCacheRoot) != nil + } + static func codexDisplayLabel(model: String) -> String? { let key = self.normalizeCodexModel(model) return self.codex[key]?.displayLabel @@ -552,6 +568,31 @@ enum CostUsagePricing { outputTokens: outputTokens) } + static func claudeProxyCodexCostUSD( + model: String, + inputTokens: Int, + cacheReadInputTokens: Int, + cacheCreationInputTokens: Int, + outputTokens: Int, + modelsDevCatalog: ModelsDevCatalog? = nil, + modelsDevCacheRoot: URL? = nil) -> Double? + { + let uncachedInput = max(0, inputTokens) + let cachedInput = max(0, cacheReadInputTokens) + let cacheWriteInput = max(0, cacheCreationInputTokens) + let totalInput = [uncachedInput, cachedInput, cacheWriteInput].reduce(0) { total, component in + total > Int.max - component ? Int.max : total + component + } + return self.codexCostUSD( + model: model, + inputTokens: totalInput, + cachedInputTokens: cachedInput, + outputTokens: outputTokens, + cacheWriteInputTokens: cacheWriteInput, + modelsDevCatalog: modelsDevCatalog, + modelsDevCacheRoot: modelsDevCacheRoot) + } + static func codexPriorityCostUSD( model: String, inputTokens: Int, @@ -781,3 +822,30 @@ enum CostUsagePricing { cacheRoot: cacheRoot) } } + +extension CostUsagePricing { + static func claudeProxyGoogleCostUSD( + model: String, + inputTokens: Int, + cacheReadInputTokens: Int, + cacheCreationInputTokens: Int, + outputTokens: Int, + modelsDevCatalog: ModelsDevCatalog? = nil, + modelsDevCacheRoot: URL? = nil) -> Double? + { + guard let lookup = self.modelsDevLookup( + providerID: "google", + model: model, + catalog: modelsDevCatalog, + cacheRoot: modelsDevCacheRoot) + else { return nil } + return self.claudeCostUSD( + pricing: lookup.pricing, + tokens: ClaudeCostTokens( + input: inputTokens, + cacheRead: cacheReadInputTokens, + cacheCreation: cacheCreationInputTokens, + cacheCreation1h: 0, + output: outputTokens)) + } +} diff --git a/Sources/CodexBarCore/Vendored/CostUsage/CostUsageScanner+CacheHelpers.swift b/Sources/CodexBarCore/Vendored/CostUsage/CostUsageScanner+CacheHelpers.swift index 3820a109f8..dea3ff58f5 100644 --- a/Sources/CodexBarCore/Vendored/CostUsage/CostUsageScanner+CacheHelpers.swift +++ b/Sources/CodexBarCore/Vendored/CostUsage/CostUsageScanner+CacheHelpers.swift @@ -1702,7 +1702,12 @@ extension CostUsageScanner { return lhsTokens > rhsTokens } - return lhs.modelName > rhs.modelName + if lhs.modelName != rhs.modelName { + return lhs.modelName > rhs.modelName + } + let lhsAttribution = lhs.attribution?.deterministicSortKey ?? "" + let rhsAttribution = rhs.attribution?.deterministicSortKey ?? "" + return lhsAttribution > rhsAttribution } } diff --git a/Sources/CodexBarCore/Vendored/CostUsage/CostUsageScanner+Claude.swift b/Sources/CodexBarCore/Vendored/CostUsage/CostUsageScanner+Claude.swift index e3f4ff8ceb..2c14416dd0 100644 --- a/Sources/CodexBarCore/Vendored/CostUsage/CostUsageScanner+Claude.swift +++ b/Sources/CodexBarCore/Vendored/CostUsage/CostUsageScanner+Claude.swift @@ -13,9 +13,18 @@ extension CostUsageScanner { let costPriced: Bool } + private struct ClaudeRawTokens { + let input: Int + let cacheRead: Int + let cacheCreate: Int + let cacheCreate1h: Int + let output: Int + } + private struct ClaudeDayModelKey: Hashable { let day: String let model: String + let attribution: CostUsageAttribution? } private struct ClaudeRepricedCost { @@ -24,6 +33,100 @@ extension CostUsageScanner { var unresolved = false } + private struct ClaudeModelResolution { + let normalizedModel: String + let cost: Double? + let attribution: CostUsageAttribution? + } + + private struct ClaudeModelResolutionContext { + let pricingDate: Date + let sessionID: String? + let timestampUnixMs: Int64? + let attributionResolver: CLIProxyAPIAttributionResolver? + let modelsDevCatalog: ModelsDevCatalog? + let modelsDevCacheRoot: URL? + } + + private static func resolveClaudeModel( + model: String, + tokens: ClaudeRawTokens, + context: ClaudeModelResolutionContext) -> ClaudeModelResolution + { + let modelProvider = CostUsagePricing.modelProvider( + for: model, + modelsDevCatalog: context.modelsDevCatalog, + modelsDevCacheRoot: context.modelsDevCacheRoot) + let resolvedAttribution = context.attributionResolver?.attribution( + model: model, + modelProvider: modelProvider, + sessionID: context.sessionID, + timestampUnixMs: context.timestampUnixMs, + tokens: .init( + input: tokens.input, + cacheRead: tokens.cacheRead, + cacheCreate: tokens.cacheCreate, + output: tokens.output)) + ?? CostUsageAttribution( + client: .claudeCode, + route: .unknown, + modelProvider: modelProvider, + evidence: [.modelProvider]) + let attribution = resolvedAttribution.route == .cliProxyAPI || modelProvider != .anthropic + ? resolvedAttribution + : nil + let upstreamModel = resolvedAttribution.route == .cliProxyAPI + ? resolvedAttribution.upstream?.model?.trimmingCharacters(in: .whitespacesAndNewlines) + : nil + let pricingModel = upstreamModel.flatMap { $0.isEmpty ? nil : $0 } ?? model + let pricingProvider = CostUsagePricing.modelProvider( + for: pricingModel, + modelsDevCatalog: context.modelsDevCatalog, + modelsDevCacheRoot: context.modelsDevCacheRoot) + let cost: Double? = if pricingProvider == .openAI { + CostUsagePricing.claudeProxyCodexCostUSD( + model: pricingModel, + inputTokens: tokens.input, + cacheReadInputTokens: tokens.cacheRead, + cacheCreationInputTokens: tokens.cacheCreate, + outputTokens: tokens.output, + modelsDevCatalog: context.modelsDevCatalog, + modelsDevCacheRoot: context.modelsDevCacheRoot) + } else if pricingProvider == .anthropic { + CostUsagePricing.claudeCostUSD( + model: pricingModel, + inputTokens: tokens.input, + cacheReadInputTokens: tokens.cacheRead, + cacheCreationInputTokens: tokens.cacheCreate, + cacheCreationInputTokens1h: tokens.cacheCreate1h, + outputTokens: tokens.output, + pricingDate: context.pricingDate, + modelsDevCatalog: context.modelsDevCatalog, + modelsDevCacheRoot: context.modelsDevCacheRoot) + } else if pricingProvider == .google { + CostUsagePricing.claudeProxyGoogleCostUSD( + model: pricingModel, + inputTokens: tokens.input, + cacheReadInputTokens: tokens.cacheRead, + cacheCreationInputTokens: tokens.cacheCreate, + outputTokens: tokens.output, + modelsDevCatalog: context.modelsDevCatalog, + modelsDevCacheRoot: context.modelsDevCacheRoot) + } else { nil } + let normalizedModel = switch modelProvider { + case .openAI: + CostUsagePricing.normalizeCodexModel(model) + case .anthropic: + CostUsagePricing.normalizeClaudeModel(model) + case .google, .unknown: + model.trimmingCharacters(in: .whitespacesAndNewlines) + } + return ClaudeModelResolution( + normalizedModel: normalizedModel, + cost: cost, + attribution: attribution) + } + static func defaultClaudeProjectsRoots( options: Options, environment: [String: String] = ProcessInfo.processInfo.environment, @@ -81,6 +184,7 @@ extension CostUsageScanner { range: CostUsageDayRange, providerFilter: ClaudeLogProviderFilter, startOffset: Int64 = 0, + attributionResolver: CLIProxyAPIAttributionResolver? = nil, modelsDevCatalog: ModelsDevCatalog? = nil, modelsDevCacheRoot: URL? = nil) -> ClaudeParseResult { @@ -90,6 +194,7 @@ extension CostUsageScanner { range: range, providerFilter: providerFilter, startOffset: startOffset, + attributionResolver: attributionResolver, modelsDevCatalog: modelsDevCatalog, modelsDevCacheRoot: modelsDevCacheRoot, checkCancellation: nil)) ?? ClaudeParseResult(days: [:], rows: [], parsedBytes: startOffset) @@ -100,6 +205,7 @@ extension CostUsageScanner { range: CostUsageDayRange, providerFilter: ClaudeLogProviderFilter, startOffset: Int64 = 0, + attributionResolver: CLIProxyAPIAttributionResolver? = nil, modelsDevCatalog: ModelsDevCatalog? = nil, modelsDevCacheRoot: URL? = nil, checkCancellation: CancellationCheck? = nil) throws -> ClaudeParseResult @@ -183,17 +289,28 @@ extension CostUsageScanner { let output = max(0, toInt(usage["output_tokens"])) if input == 0, cacheCreate == 0, cacheRead == 0, output == 0 { return } - let cost = CostUsagePricing.claudeCostUSD( + let rawTokens = ClaudeRawTokens( + input: input, + cacheRead: cacheRead, + cacheCreate: cacheCreate, + cacheCreate1h: cacheCreate1h, + output: output) + let sessionId = obj["sessionId"] as? String + ?? obj["session_id"] as? String + ?? (obj["metadata"] as? [String: Any])?["sessionId"] as? String + ?? (message["metadata"] as? [String: Any])?["sessionId"] as? String + let timestampUnixMs = Int64((timestamp.timeIntervalSince1970 * 1000).rounded()) + let modelResolution = Self.resolveClaudeModel( model: model, - inputTokens: input, - cacheReadInputTokens: cacheRead, - cacheCreationInputTokens: cacheCreate, - cacheCreationInputTokens1h: cacheCreate1h, - outputTokens: output, - pricingDate: timestamp, - modelsDevCatalog: modelsDevCatalog, - modelsDevCacheRoot: modelsDevCacheRoot) - let costNanos = cost.map { Int(($0 * costScale).rounded()) } ?? 0 + tokens: rawTokens, + context: ClaudeModelResolutionContext( + pricingDate: timestamp, + sessionID: sessionId, + timestampUnixMs: timestampUnixMs, + attributionResolver: attributionResolver, + modelsDevCatalog: modelsDevCatalog, + modelsDevCacheRoot: modelsDevCacheRoot)) + let costNanos = modelResolution.cost.map { Int(($0 * costScale).rounded()) } ?? 0 let tokens = ClaudeTokens( input: input, cacheRead: cacheRead, @@ -201,7 +318,7 @@ extension CostUsageScanner { cacheCreate1h: cacheCreate1h, output: output, costNanos: costNanos, - costPriced: cost != nil) + costPriced: modelResolution.cost != nil) guard CostUsageDayRange.isInRange( dayKey: dayKey, @@ -211,18 +328,13 @@ extension CostUsageScanner { let messageId = message["id"] as? String let requestId = obj["requestId"] as? String - let sessionId = obj["sessionId"] as? String - ?? obj["session_id"] as? String - ?? (obj["metadata"] as? [String: Any])?["sessionId"] as? String - ?? (message["metadata"] as? [String: Any])?["sessionId"] as? String - let normalizedModel = CostUsagePricing.normalizeClaudeModel(model) let row = ClaudeUsageRow( dayKey: dayKey, - model: normalizedModel, + model: modelResolution.normalizedModel, sessionId: sessionId, messageId: messageId, requestId: requestId, - timestampUnixMs: Int64((timestamp.timeIntervalSince1970 * 1000).rounded()), + timestampUnixMs: timestampUnixMs, isSidechain: toBool(obj["isSidechain"]), pathRole: pathRole, input: tokens.input, @@ -231,7 +343,8 @@ extension CostUsageScanner { cacheCreate1h: tokens.cacheCreate1h, output: tokens.output, costNanos: tokens.costNanos, - costPriced: tokens.costPriced) + costPriced: tokens.costPriced, + attribution: modelResolution.attribution) // Streaming chunks share message.id + requestId inside a file. // Keep overwriting so the final cumulative chunk wins. @@ -350,6 +463,119 @@ extension CostUsageScanner { return rows } + private static func claudeAttributionReconciliationRows( + cache: CostUsageCache) -> [(key: ClaudeAttributionReconciliationKey, row: ClaudeUsageRow)] + { + var rows: [(key: ClaudeAttributionReconciliationKey, row: ClaudeUsageRow)] = [] + var winners: [String: (path: String, row: ClaudeUsageRow)] = [:] + + for path in cache.files.keys.sorted() { + guard let fileRows = cache.files[path]?.claudeRows else { continue } + for (index, row) in fileRows.enumerated() { + guard let canonicalKey = Self.claudeCanonicalRowKey(row) else { + rows.append((key: .unkeyed(path: path, index: index), row: row)) + continue + } + let candidate = (path: path, row: row) + if let existing = winners[canonicalKey] { + if Self.claudeRowWins(lhs: candidate, rhs: existing) { + winners[canonicalKey] = candidate + } + } else { + winners[canonicalKey] = candidate + } + } + } + + rows.append(contentsOf: winners.keys.sorted().compactMap { key in + winners[key].map { (key: .canonical(key), row: $0.row) } + }) + return rows + } + + private static func reconcileClaudeAttributions( + cache: inout CostUsageCache, + attributionResolver: CLIProxyAPIAttributionResolver, + modelsDevCatalog: ModelsDevCatalog?, + modelsDevCacheRoot: URL?) + { + let items = Self.claudeAttributionReconciliationRows(cache: cache).map { item in + let row = item.row + let modelProvider = if let cachedProvider = row.attribution?.modelProvider, + cachedProvider != .unknown + { + cachedProvider + } else { + CostUsagePricing.modelProvider( + for: row.model, + modelsDevCatalog: modelsDevCatalog, + modelsDevCacheRoot: modelsDevCacheRoot) + } + return ClaudeAttributionReconciliationItem( + key: item.key, + request: CLIProxyAPIAttributionResolver.Request( + model: row.model, + modelProvider: modelProvider, + sessionID: row.sessionId, + timestampUnixMs: row.timestampUnixMs, + tokens: .init( + input: row.input, + cacheRead: row.cacheRead, + cacheCreate: row.cacheCreate, + output: row.output)), + modelProvider: modelProvider, + cachedAttribution: row.attribution) + } + let requests = items.map(\.request) + let liveAttributions = attributionResolver.attributions(for: requests) + var replacementKeys: Set = [] + var replacements: [ClaudeAttributionReconciliationKey: CostUsageAttribution] = [:] + for (index, item) in items.enumerated() + where attributionResolver.hasMatchingObservation(for: item.request) + { + replacementKeys.insert(item.key) + let liveAttribution = liveAttributions[index] + let replacement: CostUsageAttribution? = if liveAttribution.route == .cliProxyAPI { + Self.preferredCLIProxyAPIAttribution( + live: liveAttribution, + cached: item.cachedAttribution) + } else if item.modelProvider != .anthropic { + liveAttribution + } else { + nil + } + replacements[item.key] = replacement + } + guard !replacementKeys.isEmpty else { return } + + for path in cache.files.keys { + guard var file = cache.files[path], let rows = file.claudeRows else { continue } + file.claudeRows = rows.enumerated().map { index, row in + let key = Self.claudeCanonicalRowKey(row).map(ClaudeAttributionReconciliationKey.canonical) + ?? .unkeyed(path: path, index: index) + guard replacementKeys.contains(key) else { return row } + return ClaudeUsageRow( + dayKey: row.dayKey, + model: row.model, + sessionId: row.sessionId, + messageId: row.messageId, + requestId: row.requestId, + timestampUnixMs: row.timestampUnixMs, + isSidechain: row.isSidechain, + pathRole: row.pathRole, + input: row.input, + cacheRead: row.cacheRead, + cacheCreate: row.cacheCreate, + cacheCreate1h: row.cacheCreate1h, + output: row.output, + costNanos: row.costNanos, + costPriced: row.costPriced, + attribution: replacements[key]) + } + cache.files[path] = file + } + } + private static func rebuildClaudeDays(cache: inout CostUsageCache) { var days: [String: [String: [Int]]] = [:] @@ -513,6 +739,7 @@ extension CostUsageScanner { let range: CostUsageDayRange let providerFilter: ClaudeLogProviderFilter let forceFullScan: Bool + let attributionResolver: CLIProxyAPIAttributionResolver? let modelsDevCatalog: ModelsDevCatalog? let modelsDevCacheRoot: URL? let checkCancellation: CancellationCheck? @@ -522,6 +749,7 @@ extension CostUsageScanner { range: CostUsageDayRange, providerFilter: ClaudeLogProviderFilter, forceFullScan: Bool, + attributionResolver: CLIProxyAPIAttributionResolver?, modelsDevCatalog: ModelsDevCatalog?, modelsDevCacheRoot: URL?, checkCancellation: CancellationCheck?) @@ -531,6 +759,7 @@ extension CostUsageScanner { self.range = range self.providerFilter = providerFilter self.forceFullScan = forceFullScan + self.attributionResolver = attributionResolver self.modelsDevCatalog = modelsDevCatalog self.modelsDevCacheRoot = modelsDevCacheRoot self.checkCancellation = checkCancellation @@ -565,6 +794,7 @@ extension CostUsageScanner { range: state.range, providerFilter: state.providerFilter, startOffset: startOffset, + attributionResolver: state.attributionResolver, modelsDevCatalog: state.modelsDevCatalog, modelsDevCacheRoot: state.modelsDevCacheRoot, checkCancellation: state.checkCancellation) @@ -582,6 +812,7 @@ extension CostUsageScanner { fileURL: url, range: state.range, providerFilter: state.providerFilter, + attributionResolver: state.attributionResolver, modelsDevCatalog: state.modelsDevCatalog, modelsDevCacheRoot: state.modelsDevCacheRoot, checkCancellation: state.checkCancellation) @@ -654,6 +885,41 @@ extension CostUsageScanner { // Root mtime caching removed — see comment above. } + private struct ClaudeCLIProxyAPIAttributionState { + let configurationGeneration: String? + let resolver: CLIProxyAPIAttributionResolver? + } + + private static func captureClaudeCLIProxyAPIAttributionState( + options: Options, + checkCancellation: CancellationCheck?) throws -> ClaudeCLIProxyAPIAttributionState + { + try CostUsageCacheLocations.withCLIProxyAPIInterprocessLock( + stateRoot: options.cacheRoot) + { + let configurationGeneration = CostUsageCacheLocations.cliProxyAPIConfigurationGeneration( + stateRoot: options.cacheRoot) + let attributionResolver: CLIProxyAPIAttributionResolver? + if !CostUsageCacheLocations.isCLIProxyAPIExplicitlyDisconnected(stateRoot: options.cacheRoot), + let home = options.cliProxyAPIHome + { + let usageRecords = CLIProxyAPIUsageCacheIO.loadAssumingInterprocessLockHeld( + cacheRoot: options.cacheRoot) + attributionResolver = try CLIProxyAPIAttributionResolver.load( + home: home, + cacheRoot: options.cacheRoot, + forceReload: options.forceRescan, + usageRecords: usageRecords, + checkCancellation: checkCancellation) + } else { + attributionResolver = nil + } + return ClaudeCLIProxyAPIAttributionState( + configurationGeneration: configurationGeneration, + resolver: attributionResolver) + } + } + static func loadClaudeDaily( provider: UsageProvider, range: CostUsageDayRange, @@ -661,6 +927,11 @@ extension CostUsageScanner { options: Options, checkCancellation: CancellationCheck?) throws -> CostUsageDailyReport { + let cliProxyAPIAttributionState = try self.captureClaudeCLIProxyAPIAttributionState( + options: options, + checkCancellation: checkCancellation) + let cliProxyAPIConfigurationGeneration = cliProxyAPIAttributionState.configurationGeneration + let attributionResolver = cliProxyAPIAttributionState.resolver var cache = CostUsageCacheIO.load( provider: provider, cacheRoot: options.cacheRoot, @@ -669,8 +940,12 @@ extension CostUsageScanner { let refreshMs = Int64(max(0, options.refreshMinIntervalSeconds) * 1000) let windowExpanded = Self.requestedWindowExpandsCache(range: range, cache: cache) + let requiresRowBackfill = cache.files.values.contains { + $0.claudeRows == nil && !$0.days.isEmpty + } let shouldRefresh = options.forceRescan || windowExpanded + || requiresRowBackfill || refreshMs == 0 || cache.lastScanUnixMs == 0 || nowMs - cache.lastScanUnixMs > refreshMs @@ -689,7 +964,10 @@ extension CostUsageScanner { cache: cache, range: range, providerFilter: providerFilter, - forceFullScan: options.forceRescan || windowExpanded, + forceFullScan: options.forceRescan + || windowExpanded + || requiresRowBackfill, + attributionResolver: attributionResolver, modelsDevCatalog: modelsDevCatalog, modelsDevCacheRoot: options.cacheRoot, checkCancellation: checkCancellation) @@ -710,30 +988,329 @@ extension CostUsageScanner { cache.files.removeValue(forKey: key) } + if let attributionResolver { + Self.reconcileClaudeAttributions( + cache: &cache, + attributionResolver: attributionResolver, + modelsDevCatalog: modelsDevCatalog, + modelsDevCacheRoot: options.cacheRoot) + } Self.rebuildClaudeDays(cache: &cache) Self.pruneDays(cache: &cache, sinceKey: range.scanSinceKey, untilKey: range.scanUntilKey) cache.scanSinceKey = range.scanSinceKey cache.scanUntilKey = range.scanUntilKey cache.lastScanUnixMs = nowMs try checkCancellation?() - CostUsageCacheIO.save( - provider: provider, - cache: cache, - cacheRoot: options.cacheRoot, - calendar: range.calendar) + try CostUsageCacheLocations.withCLIProxyAPIInterprocessLock( + stateRoot: options.cacheRoot) + { + guard CostUsageCacheLocations.cliProxyAPIConfigurationGeneration( + stateRoot: options.cacheRoot) == cliProxyAPIConfigurationGeneration + else { throw CancellationError() } + if CostUsageCacheLocations.isCLIProxyAPIExplicitlyDisconnected( + stateRoot: options.cacheRoot) + { + for path in cache.files.keys { + guard var file = cache.files[path], let rows = file.claudeRows else { continue } + file.claudeRows = rows.map { row in + guard row.attribution?.route == .cliProxyAPI else { return row } + return ClaudeUsageRow( + dayKey: row.dayKey, + model: row.model, + sessionId: row.sessionId, + messageId: row.messageId, + requestId: row.requestId, + timestampUnixMs: row.timestampUnixMs, + isSidechain: row.isSidechain, + pathRole: row.pathRole, + input: row.input, + cacheRead: row.cacheRead, + cacheCreate: row.cacheCreate, + cacheCreate1h: row.cacheCreate1h, + output: row.output, + costNanos: row.costNanos, + costPriced: row.costPriced, + attribution: nil) + } + cache.files[path] = file + } + } + CostUsageCacheIO.save( + provider: provider, + cache: cache, + cacheRoot: options.cacheRoot, + calendar: range.calendar) + } } let modelsDevCatalog = CostUsagePricing.modelsDevCatalog(now: now, cacheRoot: options.cacheRoot) - return Self.buildClaudeReportFromCache( - cache: cache, - range: range, + return try CostUsageCacheLocations.withCLIProxyAPIInterprocessLock( + stateRoot: options.cacheRoot) + { + guard CostUsageCacheLocations.cliProxyAPIConfigurationGeneration( + stateRoot: options.cacheRoot) == cliProxyAPIConfigurationGeneration + else { throw CancellationError() } + + let reportAttributionEnabled = !CostUsageCacheLocations.isCLIProxyAPIExplicitlyDisconnected( + stateRoot: options.cacheRoot) + let reportAttributionFilter: ClaudeAttributionFilter = if reportAttributionEnabled { + options.claudeAttributionFilter + } else { + switch options.claudeAttributionFilter { + case .all, .excludeCodexBackend: .all + case .codexBackendOnly: .codexBackendOnly + } + } + let report = Self.buildClaudeReportFromCache( + cache: cache, + range: range, + attributionFilter: reportAttributionFilter, + attributionResolver: reportAttributionEnabled ? attributionResolver : nil, + allowCachedCLIProxyAPIAttribution: reportAttributionEnabled, + modelsDevCatalog: modelsDevCatalog, + modelsDevCacheRoot: options.cacheRoot) + try checkCancellation?() + guard CostUsageCacheLocations.cliProxyAPIConfigurationGeneration( + stateRoot: options.cacheRoot) == cliProxyAPIConfigurationGeneration + else { throw CancellationError() } + return report + } + } + + private struct ClaudeReportAggregation { + var dayModels: [String: [ClaudeDayModelKey: [Int]]] = [:] + var repricedCosts: [ClaudeDayModelKey: ClaudeRepricedCost] = [:] + } + + private enum ClaudeAttributionReconciliationKey: Hashable { + case canonical(String) + case unkeyed(path: String, index: Int) + } + + private struct ClaudeAttributionReconciliationItem { + let key: ClaudeAttributionReconciliationKey + let request: CLIProxyAPIAttributionResolver.Request + let modelProvider: CostUsageAttribution.ModelProvider + let cachedAttribution: CostUsageAttribution? + } + + private struct ClaudeAttributionAggregationContext { + let filter: ClaudeAttributionFilter + let resolver: CLIProxyAPIAttributionResolver? + let allowCachedCLIProxyAPIAttribution: Bool + } + + private static func aggregateClaudeRows( + cache: CostUsageCache, + attributionContext: ClaudeAttributionAggregationContext, + modelsDevCatalog: ModelsDevCatalog?, + modelsDevCacheRoot: URL?) -> ClaudeReportAggregation + { + var result = ClaudeReportAggregation() + let rowsWithProviders = Self.reconciledClaudeRows(cache: cache).map { row in + let modelProvider = if let cachedProvider = row.attribution?.modelProvider, + cachedProvider != .unknown + { + cachedProvider + } else { + CostUsagePricing.modelProvider( + for: row.model, + modelsDevCatalog: modelsDevCatalog, + modelsDevCacheRoot: modelsDevCacheRoot) + } + return (row: row, modelProvider: modelProvider) + } + let requests = rowsWithProviders.map { item in + CLIProxyAPIAttributionResolver.Request( + model: item.row.model, + modelProvider: item.modelProvider, + sessionID: item.row.sessionId, + timestampUnixMs: item.row.timestampUnixMs, + tokens: .init( + input: item.row.input, + cacheRead: item.row.cacheRead, + cacheCreate: item.row.cacheCreate, + output: item.row.output)) + } + let liveAttributions: [CostUsageAttribution?] = if let attributionResolver = attributionContext.resolver { + attributionResolver.attributions(for: requests).map(Optional.some) + } else { + Array(repeating: nil, count: rowsWithProviders.count) + } + + for (index, item) in rowsWithProviders.enumerated() { + let row = item.row + let modelProvider = item.modelProvider + let request = requests[index] + let liveAttribution = liveAttributions[index] + let cachedAttribution: CostUsageAttribution? = + if !attributionContext.allowCachedCLIProxyAPIAttribution, + row.attribution?.route == .cliProxyAPI { + nil + } else { + row.attribution + } + let attribution: CostUsageAttribution? = if let liveAttribution, + liveAttribution.route == .cliProxyAPI + { + Self.preferredCLIProxyAPIAttribution( + live: liveAttribution, + cached: cachedAttribution) + } else if attributionContext.allowCachedCLIProxyAPIAttribution, + row.attribution?.route == .cliProxyAPI, + attributionContext.resolver?.hasMatchingObservation(for: request) != true + { + row.attribution + } else if modelProvider != .anthropic { + liveAttribution ?? cachedAttribution + } else { + nil + } + let isCodexBackend = attribution?.route == .cliProxyAPI + && attribution?.upstream?.isCodex == true + let isUnresolvedAttribution = if attribution?.route == .cliProxyAPI { + attribution?.upstream == nil + } else { + modelProvider != .anthropic && modelProvider != .unknown + } + let includeRow = switch attributionContext.filter { + case .all: true + case .codexBackendOnly: isCodexBackend + case .excludeCodexBackend: !isCodexBackend && !isUnresolvedAttribution + } + guard includeRow else { continue } + + var models = result.dayModels[row.dayKey] ?? [:] + let key = ClaudeDayModelKey( + day: row.dayKey, + model: row.model, + attribution: attribution) + var packed = models[key] ?? [0, 0, 0, 0, 0, 0] + packed[0] += row.input + packed[1] += row.cacheRead + packed[2] += row.cacheCreate + packed[3] += row.output + packed[5] += 1 + models[key] = packed + result.dayModels[row.dayKey] = models + + var cost = result.repricedCosts[key] ?? ClaudeRepricedCost() + cost.sampleCount += 1 + let wasPriced = row.costPriced ?? (row.costNanos > 0) + let upstreamModel = attribution?.route == .cliProxyAPI + ? attribution?.upstream?.model?.trimmingCharacters(in: .whitespacesAndNewlines) + : nil + let pricingModel = upstreamModel.flatMap { $0.isEmpty ? nil : $0 } ?? row.model + let cachedUpstreamModel = row.attribution?.route == .cliProxyAPI + ? row.attribution?.upstream?.model?.trimmingCharacters(in: .whitespacesAndNewlines) + : nil + let cachedPricingModel = cachedUpstreamModel.flatMap { $0.isEmpty ? nil : $0 } ?? row.model + let pricingProvider = CostUsagePricing.modelProvider( + for: pricingModel, + modelsDevCatalog: modelsDevCatalog, + modelsDevCacheRoot: modelsDevCacheRoot) + let currentCost = Self.currentClaudeRowCost( + row, + pricingModel: pricingModel, + pricingProvider: pricingProvider, + modelsDevCatalog: modelsDevCatalog, + modelsDevCacheRoot: modelsDevCacheRoot) + let resolvedCost = Self.resolvedClaudeRowCost( + wasPriced: wasPriced, + cachedCostNanos: row.costNanos, + cachedPricingModel: cachedPricingModel, + pricingModel: pricingModel, + currentCost: currentCost) + if let resolvedCost { + cost.total += resolvedCost + } else { + cost.unresolved = true + } + result.repricedCosts[key] = cost + } + return result + } + + static func preferredCLIProxyAPIAttribution( + live: CostUsageAttribution, + cached: CostUsageAttribution?) -> CostUsageAttribution + { + guard live.route == .cliProxyAPI, + live.upstream == nil, + let cached, + cached.route == .cliProxyAPI, + cached.upstream != nil, + cached.evidence.contains(.cliProxyUsageTelemetry) + else { return live } + return cached + } + + static func resolvedClaudeRowCost( + wasPriced: Bool, + cachedCostNanos: Int, + cachedPricingModel: String, + pricingModel: String, + currentCost: Double?) -> Double? + { + let pricingModelUnchanged = cachedPricingModel.caseInsensitiveCompare(pricingModel) == .orderedSame + if wasPriced, cachedCostNanos == 0, pricingModelUnchanged { + return 0 + } + if let currentCost { + return currentCost + } + guard wasPriced, pricingModelUnchanged else { return nil } + return Double(cachedCostNanos) / Self.costScale + } + + private static func currentClaudeRowCost( + _ row: ClaudeUsageRow, + pricingModel: String, + pricingProvider: CostUsageAttribution.ModelProvider, + modelsDevCatalog: ModelsDevCatalog?, + modelsDevCacheRoot: URL?) -> Double? + { + if pricingProvider == .openAI { + return CostUsagePricing.claudeProxyCodexCostUSD( + model: pricingModel, + inputTokens: row.input, + cacheReadInputTokens: row.cacheRead, + cacheCreationInputTokens: row.cacheCreate, + outputTokens: row.output, + modelsDevCatalog: modelsDevCatalog, + modelsDevCacheRoot: modelsDevCacheRoot) + } + if pricingProvider == .google { + return CostUsagePricing.claudeProxyGoogleCostUSD( + model: pricingModel, + inputTokens: row.input, + cacheReadInputTokens: row.cacheRead, + cacheCreationInputTokens: row.cacheCreate, + outputTokens: row.output, + modelsDevCatalog: modelsDevCatalog, + modelsDevCacheRoot: modelsDevCacheRoot) + } + guard pricingProvider == .anthropic else { return nil } + return CostUsagePricing.claudeCostUSD( + model: pricingModel, + inputTokens: row.input, + cacheReadInputTokens: row.cacheRead, + cacheCreationInputTokens: row.cacheCreate, + cacheCreationInputTokens1h: row.cacheCreate1h ?? 0, + outputTokens: row.output, + pricingDate: row.timestampUnixMs.map { + Date(timeIntervalSince1970: Double($0) / 1000) + }, modelsDevCatalog: modelsDevCatalog, - modelsDevCacheRoot: options.cacheRoot) + modelsDevCacheRoot: modelsDevCacheRoot) } - private static func buildClaudeReportFromCache( + static func buildClaudeReportFromCache( cache: CostUsageCache, range: CostUsageDayRange, + attributionFilter: ClaudeAttributionFilter = .all, + attributionResolver: CLIProxyAPIAttributionResolver? = nil, + allowCachedCLIProxyAPIAttribution: Bool = true, modelsDevCatalog: ModelsDevCatalog? = nil, modelsDevCacheRoot: URL? = nil) -> CostUsageDailyReport { @@ -745,50 +1322,27 @@ extension CostUsageScanner { var totalTokens = 0 var totalCost: Double = 0 var costSeen = false - let costScale = 1_000_000_000.0 - var repricedCosts: [ClaudeDayModelKey: ClaudeRepricedCost] = [:] - - for row in Self.reconciledClaudeRows(cache: cache) { - let key = ClaudeDayModelKey(day: row.dayKey, model: row.model) - var aggregate = repricedCosts[key] ?? ClaudeRepricedCost() - aggregate.sampleCount += 1 - let isPriced = row.costPriced ?? (row.costNanos > 0) - let currentPricingCost = CostUsagePricing.claudeCostUSD( - model: row.model, - inputTokens: row.input, - cacheReadInputTokens: row.cacheRead, - cacheCreationInputTokens: row.cacheCreate, - cacheCreationInputTokens1h: row.cacheCreate1h ?? 0, - outputTokens: row.output, - pricingDate: row.timestampUnixMs.map { - Date(timeIntervalSince1970: Double($0) / 1000) - }, - modelsDevCatalog: modelsDevCatalog, - modelsDevCacheRoot: modelsDevCacheRoot) - let cost: Double? = if isPriced, row.costNanos == 0 { - 0 - } else if let currentPricingCost { - currentPricingCost - } else if isPriced { - Double(row.costNanos) / costScale - } else { - nil - } - if let cost { - aggregate.total += cost - } else { - aggregate.unresolved = true - } - repricedCosts[key] = aggregate - } + let aggregation = Self.aggregateClaudeRows( + cache: cache, + attributionContext: .init( + filter: attributionFilter, + resolver: attributionResolver, + allowCachedCLIProxyAPIAttribution: allowCachedCLIProxyAPIAttribution), + modelsDevCatalog: modelsDevCatalog, + modelsDevCacheRoot: modelsDevCacheRoot) + let dayModels = aggregation.dayModels + let repricedCosts = aggregation.repricedCosts - let dayKeys = cache.days.keys.sorted().filter { + let dayKeys = dayModels.keys.sorted().filter { CostUsageDayRange.isInRange(dayKey: $0, since: range.sinceKey, until: range.untilKey) } for day in dayKeys { - guard let models = cache.days[day] else { continue } - let modelNames = models.keys.sorted() + guard let models = dayModels[day] else { continue } + let modelKeys = models.keys.sorted { + if $0.model != $1.model { return $0.model < $1.model } + return ($0.attribution?.upstream?.provider ?? "") < ($1.attribution?.upstream?.provider ?? "") + } var dayInput = 0 var dayOutput = 0 @@ -799,8 +1353,9 @@ extension CostUsageScanner { var dayCost: Double = 0 var dayCostSeen = false - for model in modelNames { - let packed = models[model] ?? [0, 0, 0, 0] + for modelKey in modelKeys { + let model = modelKey.model + let packed = models[modelKey] ?? [0, 0, 0, 0] let input = packed[safe: 0] ?? 0 let cacheRead = packed[safe: 1] ?? 0 let cacheCreate = packed[safe: 2] ?? 0 @@ -814,7 +1369,7 @@ extension CostUsageScanner { dayCacheCreate += cacheCreate dayOutput += output - let repricedCost = repricedCosts[ClaudeDayModelKey(day: day, model: model)] + let repricedCost = repricedCosts[modelKey] let currentPricingCost: Double? = if let repricedCost, repricedCost.sampleCount == sampleCount, !repricedCost.unresolved @@ -828,7 +1383,8 @@ extension CostUsageScanner { CostUsageDailyReport.ModelBreakdown( modelName: model, costUSD: cost, - totalTokens: totalTokens)) + totalTokens: totalTokens, + attribution: modelKey.attribution)) if let cost { dayCost += cost dayCostSeen = true @@ -847,7 +1403,7 @@ extension CostUsageScanner { cacheCreationTokens: dayCacheCreate, totalTokens: dayTotal, costUSD: entryCost, - modelsUsed: modelNames, + modelsUsed: Array(Set(modelKeys.map(\.model))).sorted(), modelBreakdowns: sortedBreakdown)) totalInput += dayInput diff --git a/Sources/CodexBarCore/Vendored/CostUsage/CostUsageScanner.swift b/Sources/CodexBarCore/Vendored/CostUsage/CostUsageScanner.swift index 853f54dd88..6d7c6ba81d 100644 --- a/Sources/CodexBarCore/Vendored/CostUsage/CostUsageScanner.swift +++ b/Sources/CodexBarCore/Vendored/CostUsage/CostUsageScanner.swift @@ -43,6 +43,12 @@ enum CostUsageScanner { case excludeVertexAI } + enum ClaudeAttributionFilter { + case all + case codexBackendOnly + case excludeCodexBackend + } + struct Options { var codexSessionsRoot: URL? var claudeProjectsRoots: [URL]? @@ -51,6 +57,8 @@ enum CostUsageScanner { var calendar: Calendar var refreshMinIntervalSeconds: TimeInterval = 60 var claudeLogProviderFilter: ClaudeLogProviderFilter = .all + var claudeAttributionFilter: ClaudeAttributionFilter = .all + var cliProxyAPIHome: URL? /// Force a full rescan, ignoring per-file cache and incremental offsets. var forceRescan: Bool = false /// Maximum bounded slice read from one Codex rollout per refresh. Larger files @@ -72,6 +80,8 @@ enum CostUsageScanner { codexTraceDatabaseURL: URL? = nil, calendar: Calendar = .current, claudeLogProviderFilter: ClaudeLogProviderFilter = .all, + claudeAttributionFilter: ClaudeAttributionFilter = .all, + cliProxyAPIHome: URL? = nil, forceRescan: Bool = false, maxCodexSessionFileBytes: Int64 = 256 * 1024 * 1024, maxCodexScanBytesPerRefresh: Int64 = 512 * 1024 * 1024, @@ -84,6 +94,8 @@ enum CostUsageScanner { self.codexTraceDatabaseURL = codexTraceDatabaseURL self.calendar = calendar self.claudeLogProviderFilter = claudeLogProviderFilter + self.claudeAttributionFilter = claudeAttributionFilter + self.cliProxyAPIHome = cliProxyAPIHome self.forceRescan = forceRescan self.maxCodexSessionFileBytes = max(0, maxCodexSessionFileBytes) self.maxCodexScanBytesPerRefresh = max(0, maxCodexScanBytesPerRefresh) @@ -1672,6 +1684,7 @@ enum CostUsageScanner { let output: Int let costNanos: Int let costPriced: Bool? + let attribution: CostUsageAttribution? } static func loadDailyReport( diff --git a/Tests/CodexBarTests/CLICacheTests.swift b/Tests/CodexBarTests/CLICacheTests.swift index 21ca7d869a..c5d1b03cb4 100644 --- a/Tests/CodexBarTests/CLICacheTests.swift +++ b/Tests/CodexBarTests/CLICacheTests.swift @@ -1,6 +1,9 @@ import Commander +import Dispatch +import Foundation import Testing @testable import CodexBarCLI +@testable import CodexBarCore struct CLICacheTests { @Test @@ -29,4 +32,62 @@ struct CLICacheTests { #expect(help.contains("--provider with --cookies")) #expect(help.contains("codexbar cache clear --cookies --provider claude")) } + + @Test + func `cost clear waits for the collector interprocess lock`() async throws { + let root = FileManager.default.temporaryDirectory + .appendingPathComponent("codexbar-cli-cost-clear-lock-\(UUID().uuidString)", isDirectory: true) + let cacheDirectory = root.appendingPathComponent("cost-usage", isDirectory: true) + try FileManager.default.createDirectory(at: cacheDirectory, withIntermediateDirectories: true) + try Data("telemetry".utf8).write(to: cacheDirectory.appendingPathComponent("usage.json")) + defer { try? FileManager.default.removeItem(at: root) } + + let lockAcquired = DispatchSemaphore(value: 0) + let releaseLock = DispatchSemaphore(value: 0) + let clearStarted = DispatchSemaphore(value: 0) + let clearFinished = DispatchSemaphore(value: 0) + let collector = Task.detached { + try await CostUsageCacheLocations.withCLIProxyAPIInterprocessLock(stateRoot: root) { + lockAcquired.signal() + _ = await Self.waitForSignal(releaseLock, timeout: .distantFuture) + } + } + #expect(await Self.waitForSignal(lockAcquired, timeout: .now() + 1)) + + let clear = Task.detached { + clearStarted.signal() + let result = CostUsageCacheLocations.clearAllCostUsageCaches( + in: [cacheDirectory], + stateRoot: root, + fileManager: .default) + clearFinished.signal() + return result + } + #expect(await Self.waitForSignal(clearStarted, timeout: .now() + 1)) + #expect(!Self.waitForSignalSync(clearFinished, timeout: .now() + .milliseconds(50))) + + releaseLock.signal() + try await collector.value + let result = await clear.value + #expect(result == CostUsageCacheClearResult(cleared: 1, errorDescription: nil)) + #expect(!FileManager.default.fileExists(atPath: cacheDirectory.path)) + } + + private static func waitForSignal( + _ semaphore: DispatchSemaphore, + timeout: DispatchTime) async -> Bool + { + await withCheckedContinuation { continuation in + DispatchQueue.global().async { + continuation.resume(returning: semaphore.wait(timeout: timeout) == .success) + } + } + } + + private static func waitForSignalSync( + _ semaphore: DispatchSemaphore, + timeout: DispatchTime) -> Bool + { + semaphore.wait(timeout: timeout) == .success + } } diff --git a/Tests/CodexBarTests/CLIProxyAPIAliasRegressionTests.swift b/Tests/CodexBarTests/CLIProxyAPIAliasRegressionTests.swift new file mode 100644 index 0000000000..bb16795f38 --- /dev/null +++ b/Tests/CodexBarTests/CLIProxyAPIAliasRegressionTests.swift @@ -0,0 +1,174 @@ +import Foundation +import Testing +@testable import CodexBarCore + +struct CLIProxyAPIAliasRegressionTests { + @Test + func `codex oauth model and alias do not prove a proxy route`() throws { + let fileManager = FileManager.default + let root = fileManager.temporaryDirectory + .appendingPathComponent("cliproxy-alias-\(UUID().uuidString)", isDirectory: true) + try fileManager.createDirectory(at: root, withIntermediateDirectories: true) + defer { try? fileManager.removeItem(at: root) } + + try Self.writeCodexAliasConfiguration(to: root, fileManager: fileManager) + let resolver = try CLIProxyAPIAttributionResolver.load(home: root, fileManager: fileManager) + for model in ["gpt-5.5", "proxy-codex-alias"] { + let attribution = resolver.attribution( + model: model, + modelProvider: .unknown, + sessionID: nil, + timestampUnixMs: nil, + tokens: Self.tokens) + + #expect(attribution.route == .unknown) + #expect(attribution.modelProvider == .unknown) + #expect(attribution.upstream == nil) + #expect(!attribution.evidence.contains(.cliProxyAuthInventory)) + #expect(!attribution.evidence.contains(.cliProxyModelAlias)) + #expect(!attribution.evidence.contains(.cliProxyRequestLog)) + } + } + + @Test + func `codex oauth model and alias resolve after request route evidence`() { + let timestamp = Date(timeIntervalSince1970: 1_784_179_200) + for model in ["gpt-5.5", "proxy-codex-alias"] { + let resolver = CLIProxyAPIAttributionResolver( + observations: [ + .init(sessionID: "proxied-session", model: model, timestamp: timestamp), + ], + authProviders: [ + .init(provider: "codex", authType: .oauth), + ], + codexOAuthModelAliases: ["proxy-codex-alias": "gpt-5.5"]) + let attribution = resolver.attribution( + model: model, + modelProvider: .unknown, + sessionID: "proxied-session", + timestampUnixMs: Int64(timestamp.timeIntervalSince1970 * 1000), + tokens: Self.tokens) + + #expect(attribution.route == .cliProxyAPI) + #expect(attribution.modelProvider == .openAI) + #expect(attribution.upstream == .init(provider: "codex", authType: .oauth, model: "gpt-5.5")) + #expect(attribution.evidence.contains(.cliProxyAuthInventory)) + #expect(attribution.evidence.contains(.cliProxyModelAlias)) + #expect(attribution.evidence.contains(.cliProxyRequestLog)) + } + } + + @Test + func `codex oauth alias parser ignores comments and other providers`() { + let configuration = """ + # oauth-model-alias: + # codex: + # - name: "ignored" + oauth-model-alias: + codex: + - name: 'gpt-5.5' + alias: 'proxy-codex-alias' # local alias + vertex: + - name: "gemini-test" + alias: "unrelated-alias" + """ + + #expect(CLIProxyAPIAttributionResolver.parseCodexOAuthModelAliases(configuration) == [ + "proxy-codex-alias": "gpt-5.5", + ]) + } + + @Test + func `weaker live route evidence preserves cached telemetry upstream`() { + let cached = CostUsageAttribution( + client: .claudeCode, + route: .cliProxyAPI, + modelProvider: .openAI, + upstream: .init(provider: "codex", authType: .oauth, model: "gpt-5.5"), + evidence: [.cliProxyRequestLog, .cliProxyUsageTelemetry, .modelProvider]) + let live = CostUsageAttribution( + client: .claudeCode, + route: .cliProxyAPI, + modelProvider: .openAI, + evidence: [.cliProxyRequestLog, .modelProvider]) + + #expect(CostUsageScanner.preferredCLIProxyAPIAttribution(live: live, cached: cached) == cached) + } + + @Test + func `codex oauth alias keeps direct historical usage with Claude`() async throws { + let env = try CostUsageTestEnvironment() + defer { env.cleanup() } + + let day = try env.makeLocalNoon(year: 2026, month: 7, day: 24) + _ = try env.writeClaudeProjectFile( + relativePath: "aliased-proxy/session.jsonl", + contents: env.jsonl((0..<1).map { index in + [ + "type": "assistant", + "timestamp": env.isoString(for: day.addingTimeInterval(TimeInterval(index))), + "sessionId": "aliased-proxy-session", + "requestId": "aliased-request-\(index)", + "message": [ + "id": "aliased-message-\(index)", + "model": "proxy-codex-alias", + "usage": ["input_tokens": 100, "output_tokens": 5], + ], + ] + })) + let cliProxyHome = env.root.appendingPathComponent("cli-proxy-api", isDirectory: true) + try FileManager.default.createDirectory(at: cliProxyHome, withIntermediateDirectories: true) + try Self.writeCodexAliasConfiguration(to: cliProxyHome, fileManager: .default) + let options = CostUsageScanner.Options( + codexSessionsRoot: env.codexSessionsRoot, + claudeProjectsRoots: [env.claudeProjectsRoot], + cacheRoot: env.cacheRoot, + cliProxyAPIHome: cliProxyHome) + + let codex = try await CostUsageFetcher.loadTokenSnapshot( + provider: .codex, + now: day, + allowPricingRefresh: false, + includePiSessions: false, + scannerOptions: options) + let claude = try await CostUsageFetcher.loadTokenSnapshot( + provider: .claude, + now: day, + allowPricingRefresh: false, + includePiSessions: false, + scannerOptions: options) + + #expect(codex.daily.isEmpty) + #expect(claude.daily.first?.totalTokens == 105) + #expect(claude.daily.first?.modelBreakdowns?.count == 1) + #expect(claude.daily.first?.modelBreakdowns?.allSatisfy { + $0.attribution?.route != .cliProxyAPI + } == true) + } + + private static let tokens = CLIProxyAPIAttributionResolver.TokenSignature( + input: 10, + cacheRead: 30, + cacheCreate: 40, + output: 20) + + private static func writeCodexAliasConfiguration( + to root: URL, + fileManager: FileManager) throws + { + try fileManager.createDirectory(at: root, withIntermediateDirectories: true) + let configuration = """ + oauth-model-alias: + codex: + - name: "gpt-5.5" + alias: "proxy-codex-alias" + force-mapping: true + vertex: + - name: "gemini-test" + alias: "unrelated-alias" + """ + try Data(configuration.utf8).write(to: root.appendingPathComponent("config.yaml")) + try Data(#"{"type":"codex","disabled":false}"#.utf8) + .write(to: root.appendingPathComponent("codex-auth.json")) + } +} diff --git a/Tests/CodexBarTests/CLIProxyAPIAttributionBatchTests.swift b/Tests/CodexBarTests/CLIProxyAPIAttributionBatchTests.swift new file mode 100644 index 0000000000..2ce3243744 --- /dev/null +++ b/Tests/CodexBarTests/CLIProxyAPIAttributionBatchTests.swift @@ -0,0 +1,232 @@ +import Foundation +import Testing +@testable import CodexBarCore + +struct CLIProxyAPIAttributionBatchTests { + @Test + func `batch attribution preserves uniquely matched concurrent proxy requests`() { + let timestamp = Date(timeIntervalSince1970: 1_784_179_200) + let otherTokens = CLIProxyAPIAttributionResolver.TokenSignature( + input: 100, + cacheRead: 300, + cacheCreate: 400, + output: 200) + let resolver = CLIProxyAPIAttributionResolver( + observations: [ + .init(sessionID: "session-1", model: "gpt-5.5", timestamp: timestamp), + .init( + sessionID: "session-2", + model: "gpt-5.5", + timestamp: timestamp.addingTimeInterval(2)), + ], + usageRecords: [ + Self.record( + timestamp: timestamp.addingTimeInterval(1), + provider: "codex", + authType: "oauth"), + Self.record( + timestamp: timestamp.addingTimeInterval(3), + provider: "openrouter", + authType: "api_key", + tokens: otherTokens), + ]) + let requests = [ + CLIProxyAPIAttributionResolver.Request( + model: "gpt-5.5", + modelProvider: .openAI, + sessionID: "session-1", + timestampUnixMs: Int64(timestamp.timeIntervalSince1970 * 1000), + tokens: Self.tokens), + CLIProxyAPIAttributionResolver.Request( + model: "gpt-5.5", + modelProvider: .openAI, + sessionID: "session-2", + timestampUnixMs: Int64(timestamp.addingTimeInterval(2).timeIntervalSince1970 * 1000), + tokens: otherTokens), + ] + + let attributions = resolver.attributions(for: requests) + + #expect(attributions.map(\.upstream?.provider) == ["codex", "openrouter"]) + #expect(attributions.allSatisfy { $0.evidence.contains(.cliProxyUsageTelemetry) }) + } + + @Test + func `batch attribution uses unique timestamps for concurrent requests with equal tokens`() { + let timestamp = Date(timeIntervalSince1970: 1_784_179_200) + let resolver = CLIProxyAPIAttributionResolver( + observations: [ + .init(sessionID: "session-1", model: "gpt-5.5", timestamp: timestamp), + .init( + sessionID: "session-2", + model: "gpt-5.5", + timestamp: timestamp.addingTimeInterval(4)), + ], + usageRecords: [ + Self.record(timestamp: timestamp, provider: "codex", authType: "oauth"), + Self.record( + timestamp: timestamp.addingTimeInterval(4), + provider: "openrouter", + authType: "api_key"), + ]) + let attributions = resolver.attributions(for: [ + .init( + model: "gpt-5.5", + modelProvider: .openAI, + sessionID: "session-1", + timestampUnixMs: Int64(timestamp.timeIntervalSince1970 * 1000), + tokens: Self.tokens), + .init( + model: "gpt-5.5", + modelProvider: .openAI, + sessionID: "session-2", + timestampUnixMs: Int64(timestamp.addingTimeInterval(4).timeIntervalSince1970 * 1000), + tokens: Self.tokens), + ]) + + #expect(attributions.map(\.upstream?.provider) == ["codex", "openrouter"]) + } + + @Test + func `batch route evidence belongs only to the closest request in a resumed session`() { + let timestamp = Date(timeIntervalSince1970: 1_784_179_200) + let resolver = CLIProxyAPIAttributionResolver( + observations: [ + .init(sessionID: "resumed-session", model: "gpt-5.5", timestamp: timestamp), + ]) + let attributions = resolver.attributions(for: [ + .init( + model: "gpt-5.5", + modelProvider: .openAI, + sessionID: "resumed-session", + timestampUnixMs: Int64(timestamp.addingTimeInterval(1).timeIntervalSince1970 * 1000), + tokens: Self.tokens), + .init( + model: "gpt-5.5", + modelProvider: .openAI, + sessionID: "resumed-session", + timestampUnixMs: Int64(timestamp.addingTimeInterval(30).timeIntervalSince1970 * 1000), + tokens: Self.tokens), + ]) + + #expect(attributions.map(\.route) == [.cliProxyAPI, .unknown]) + #expect(attributions[0].evidence.contains(.cliProxyRequestLog)) + #expect(!attributions[1].evidence.contains(.cliProxyRequestLog)) + } + + @Test + func `orphaned route evidence does not claim a later sole request`() { + let timestamp = Date(timeIntervalSince1970: 1_784_179_200) + let resolver = CLIProxyAPIAttributionResolver( + observations: [ + .init(sessionID: "resumed-session", model: "gpt-5.5", timestamp: timestamp), + ]) + let attribution = resolver.attributions(for: [ + .init( + model: "gpt-5.5", + modelProvider: .openAI, + sessionID: "resumed-session", + timestampUnixMs: Int64(timestamp.addingTimeInterval(30 * 60).timeIntervalSince1970 * 1000), + tokens: Self.tokens), + ])[0] + + #expect(attribution.route == .unknown) + #expect(!attribution.evidence.contains(.cliProxyRequestLog)) + } + + @Test + func `orphaned route evidence does not claim the closest of multiple later requests`() { + let timestamp = Date(timeIntervalSince1970: 1_784_179_200) + let resolver = CLIProxyAPIAttributionResolver( + observations: [ + .init(sessionID: "resumed-session", model: "gpt-5.5", timestamp: timestamp), + ]) + let attributions = resolver.attributions(for: [ + .init( + model: "gpt-5.5", + modelProvider: .openAI, + sessionID: "resumed-session", + timestampUnixMs: Int64(timestamp.addingTimeInterval(30 * 60).timeIntervalSince1970 * 1000), + tokens: Self.tokens), + .init( + model: "gpt-5.5", + modelProvider: .openAI, + sessionID: "resumed-session", + timestampUnixMs: Int64(timestamp.addingTimeInterval(40 * 60).timeIntervalSince1970 * 1000), + tokens: Self.tokens), + ]) + + #expect(attributions.map(\.route) == [.unknown, .unknown]) + #expect(attributions.allSatisfy { !$0.evidence.contains(.cliProxyRequestLog) }) + } + + @Test + func `uniquely matched telemetry confirms both requests sharing one route observation`() { + let timestamp = Date(timeIntervalSince1970: 1_784_179_200) + let otherTokens = CLIProxyAPIAttributionResolver.TokenSignature( + input: 100, + cacheRead: 300, + cacheCreate: 400, + output: 200) + let resolver = CLIProxyAPIAttributionResolver( + observations: [ + .init(sessionID: "resumed-session", model: "gpt-5.5", timestamp: timestamp), + ], + usageRecords: [ + Self.record(timestamp: timestamp, provider: "codex", authType: "oauth"), + Self.record( + timestamp: timestamp.addingTimeInterval(1), + provider: "openrouter", + authType: "api_key", + tokens: otherTokens), + ]) + let attributions = resolver.attributions(for: [ + .init( + model: "gpt-5.5", + modelProvider: .openAI, + sessionID: "resumed-session", + timestampUnixMs: Int64(timestamp.timeIntervalSince1970 * 1000), + tokens: Self.tokens), + .init( + model: "gpt-5.5", + modelProvider: .openAI, + sessionID: "resumed-session", + timestampUnixMs: Int64(timestamp.addingTimeInterval(1).timeIntervalSince1970 * 1000), + tokens: otherTokens), + ]) + + #expect(attributions.map(\.route) == [.cliProxyAPI, .cliProxyAPI]) + #expect(attributions.map(\.upstream?.provider) == ["codex", "openrouter"]) + #expect(attributions.allSatisfy { $0.evidence.contains(.cliProxyUsageTelemetry) }) + #expect(attributions.count(where: { $0.evidence.contains(.cliProxyRequestLog) }) == 1) + } + + private static let tokens = CLIProxyAPIAttributionResolver.TokenSignature( + input: 10, + cacheRead: 30, + cacheCreate: 40, + output: 20) + + private static func record( + timestamp: Date, + provider: String, + authType: String, + tokens: CLIProxyAPIAttributionResolver.TokenSignature = Self.tokens) -> CLIProxyAPIUsageRecord + { + CLIProxyAPIUsageRecord( + timestamp: timestamp, + provider: provider, + executorType: provider == "codex" ? "CodexExecutor" : "OpenAICompatExecutor", + model: "gpt-5.5", + alias: "gpt-5.5", + endpoint: "POST /v1/messages", + authType: authType, + requestID: "request-\(provider)-\(timestamp.timeIntervalSince1970)", + tokens: .init( + input: tokens.input, + output: tokens.output, + cacheRead: tokens.cacheRead, + cacheCreation: tokens.cacheCreate, + total: tokens.input + tokens.output + tokens.cacheRead + tokens.cacheCreate)) + } +} diff --git a/Tests/CodexBarTests/CLIProxyAPIAttributionResolverTests.swift b/Tests/CodexBarTests/CLIProxyAPIAttributionResolverTests.swift new file mode 100644 index 0000000000..a82967a3c7 --- /dev/null +++ b/Tests/CodexBarTests/CLIProxyAPIAttributionResolverTests.swift @@ -0,0 +1,1028 @@ +import Foundation +import Testing +@testable import CodexBarCore + +struct CLIProxyAPIAttributionResolverTests { + @Test + func `request log confirms route without guessing upstream`() { + let timestamp = Date(timeIntervalSince1970: 1_784_179_200) + let resolver = CLIProxyAPIAttributionResolver( + observations: [ + .init(sessionID: "session-1", model: "gpt-5.5", timestamp: timestamp), + ]) + + let attribution = resolver.attribution( + model: "gpt-5.5", + modelProvider: .openAI, + sessionID: "session-1", + timestampUnixMs: Int64(timestamp.timeIntervalSince1970 * 1000), + tokens: Self.tokens) + + #expect(attribution.route == .cliProxyAPI) + #expect(attribution.upstream == nil) + #expect(attribution.evidence == [.cliProxyRequestLog, .modelProvider]) + } + + @Test + func `request log does not confirm a distant request in the same session`() { + let timestamp = Date(timeIntervalSince1970: 1_784_179_200) + let resolver = CLIProxyAPIAttributionResolver( + observations: [ + .init(sessionID: "session-1", model: "gpt-5.5", timestamp: timestamp), + ]) + + let attribution = resolver.attribution( + model: "gpt-5.5", + modelProvider: .openAI, + sessionID: "session-1", + timestampUnixMs: Int64(timestamp.addingTimeInterval(3 * 60 * 60).timeIntervalSince1970 * 1000), + tokens: Self.tokens) + + #expect(attribution.route == .unknown) + #expect(attribution.upstream == nil) + #expect(attribution.evidence == [.modelProvider]) + } + + @Test + func `codex auth inventory identifies upstream after this session route is proven`() { + let resolver = CLIProxyAPIAttributionResolver( + observations: [ + .init(sessionID: "logged-session", model: "gpt-5.5", timestamp: nil), + ], + authProviders: [ + .init(provider: "codex", authType: .oauth), + ]) + + let attribution = resolver.attribution( + model: "gpt-5.5", + modelProvider: .openAI, + sessionID: "logged-session", + timestampUnixMs: nil, + tokens: Self.tokens) + + #expect(attribution.route == .cliProxyAPI) + #expect(attribution.upstream == .init( + provider: "codex", + authType: .oauth, + model: "gpt-5.5")) + #expect(attribution.evidence == [ + .cliProxyAuthInventory, + .cliProxyRequestLog, + .modelProvider, + ]) + } + + @Test + func `codex auth inventory does not transfer route proof between sessions`() { + let resolver = CLIProxyAPIAttributionResolver( + observations: [ + .init(sessionID: "logged-session", model: "gpt-5.5", timestamp: nil), + ], + authProviders: [ + .init(provider: "codex", authType: .oauth), + ]) + + let attribution = resolver.attribution( + model: "gpt-5.5", + modelProvider: .openAI, + sessionID: "unrelated-session", + timestampUnixMs: nil, + tokens: Self.tokens) + + #expect(attribution.route == .unknown) + #expect(attribution.upstream == nil) + #expect(attribution.evidence == [.modelProvider]) + } + + @Test + func `codex auth inventory stays ambiguous with another active provider`() throws { + let fileManager = FileManager.default + let root = fileManager.temporaryDirectory + .appendingPathComponent("cliproxy-auth-inventory-\(UUID().uuidString)", isDirectory: true) + let home = root.appendingPathComponent("home", isDirectory: true) + let logs = home.appendingPathComponent("logs", isDirectory: true) + try fileManager.createDirectory(at: logs, withIntermediateDirectories: true) + defer { try? fileManager.removeItem(at: root) } + + let timestamp = Date(timeIntervalSince1970: 1_784_179_200) + try Data(Self.requestLog(sessionID: "session-1", timestamp: timestamp).utf8) + .write(to: logs.appendingPathComponent("request.log")) + try Data(#"{"type":"codex"}"#.utf8) + .write(to: home.appendingPathComponent("codex.json")) + try Data(#"{"type":"openrouter"}"#.utf8) + .write(to: home.appendingPathComponent("openrouter.json")) + + let resolver = try CLIProxyAPIAttributionResolver.load(home: home, fileManager: fileManager) + let attribution = resolver.attribution( + model: "gpt-5.5", + modelProvider: .openAI, + sessionID: "session-1", + timestampUnixMs: Int64(timestamp.timeIntervalSince1970 * 1000), + tokens: Self.tokens) + + #expect(attribution.route == .cliProxyAPI) + #expect(attribution.upstream == nil) + #expect(attribution.evidence == [.cliProxyRequestLog, .modelProvider]) + } + + @Test + func `request telemetry identifies exact codex oauth upstream`() { + let timestamp = Date(timeIntervalSince1970: 1_784_179_200) + let resolver = CLIProxyAPIAttributionResolver( + observations: [ + .init(sessionID: "session-1", model: "gpt-5.5", timestamp: timestamp), + ], + usageRecords: [ + Self.record( + timestamp: timestamp.addingTimeInterval(1), + provider: "codex", + authType: "oauth"), + ]) + + let attribution = resolver.attribution( + model: "gpt-5.5", + modelProvider: .openAI, + sessionID: "session-1", + timestampUnixMs: Int64(timestamp.timeIntervalSince1970 * 1000), + tokens: Self.tokens) + + #expect(attribution.route == .cliProxyAPI) + #expect(attribution.upstream?.provider == "codex") + #expect(attribution.upstream?.authType == .oauth) + #expect(attribution.upstream?.model == "gpt-5.5") + #expect(attribution.evidence == [ + .cliProxyRequestLog, + .cliProxyUsageTelemetry, + .modelProvider, + ]) + } + + @Test + func `dated request log outranks an undated log for telemetry correlation`() { + let timestamp = Date(timeIntervalSince1970: 1_784_179_200) + let resolver = CLIProxyAPIAttributionResolver( + observations: [ + .init(sessionID: "session-1", model: "gpt-5.5", timestamp: nil), + .init(sessionID: "session-1", model: "gpt-5.5", timestamp: timestamp), + ], + usageRecords: [ + Self.record(timestamp: timestamp, provider: "openrouter", authType: "api_key"), + ], + authProviders: [ + .init(provider: "codex", authType: .oauth), + ]) + + let attribution = resolver.attribution( + model: "gpt-5.5", + modelProvider: .openAI, + sessionID: "session-1", + timestampUnixMs: Int64(timestamp.timeIntervalSince1970 * 1000), + tokens: Self.tokens) + + #expect(attribution.upstream?.provider == "openrouter") + #expect(attribution.upstream?.authType == .apiKey) + #expect(attribution.evidence.contains(.cliProxyUsageTelemetry)) + #expect(!attribution.evidence.contains(.cliProxyAuthInventory)) + } + + @Test + func `request telemetry preserves api key authentication type`() { + let timestamp = Date(timeIntervalSince1970: 1_784_179_200) + let resolver = CLIProxyAPIAttributionResolver( + observations: [ + .init(sessionID: "session-1", model: "gpt-5.5", timestamp: timestamp), + ], + usageRecords: [ + Self.record(timestamp: timestamp, provider: "openrouter", authType: "apikey"), + ]) + + let attribution = resolver.attribution( + model: "gpt-5.5", + modelProvider: .openAI, + sessionID: "session-1", + timestampUnixMs: Int64(timestamp.timeIntervalSince1970 * 1000), + tokens: Self.tokens) + + #expect(attribution.upstream?.provider == "openrouter") + #expect(attribution.upstream?.authType == .apiKey) + #expect(attribution.upstream?.displayName == "OpenRouter API key") + } + + @Test + func `ambiguous telemetry does not claim an upstream`() { + let timestamp = Date(timeIntervalSince1970: 1_784_179_200) + let resolver = CLIProxyAPIAttributionResolver( + observations: [ + .init(sessionID: "session-1", model: "gpt-5.5", timestamp: timestamp), + ], + usageRecords: [ + Self.record(timestamp: timestamp, provider: "codex", authType: "oauth"), + Self.record(timestamp: timestamp, provider: "openrouter", authType: "api_key"), + ]) + + let attribution = resolver.attribution( + model: "gpt-5.5", + modelProvider: .openAI, + sessionID: "session-1", + timestampUnixMs: Int64(timestamp.timeIntervalSince1970 * 1000), + tokens: Self.tokens) + + #expect(attribution.route == .cliProxyAPI) + #expect(attribution.upstream == nil) + #expect(!attribution.evidence.contains(.cliProxyUsageTelemetry)) + } + + @Test + func `telemetry plausible for two requests does not claim either upstream`() { + let timestamp = Date(timeIntervalSince1970: 1_784_179_200) + let resolver = CLIProxyAPIAttributionResolver( + observations: [ + .init(sessionID: "session-1", model: "gpt-5.5", timestamp: timestamp), + .init( + sessionID: "session-2", + model: "gpt-5.5", + timestamp: timestamp.addingTimeInterval(2)), + ], + usageRecords: [ + Self.record( + timestamp: timestamp.addingTimeInterval(1), + provider: "codex", + authType: "oauth"), + ]) + + let attributions = resolver.attributions(for: ["session-1", "session-2"].map { sessionID in + CLIProxyAPIAttributionResolver.Request( + model: "gpt-5.5", + modelProvider: .openAI, + sessionID: sessionID, + timestampUnixMs: Int64(timestamp.timeIntervalSince1970 * 1000), + tokens: Self.tokens) + }) + + for attribution in attributions { + #expect(attribution.route == .cliProxyAPI) + #expect(attribution.upstream == nil) + #expect(!attribution.evidence.contains(.cliProxyUsageTelemetry)) + } + } + + @Test + func `failed and token mismatched telemetry are ignored`() { + let timestamp = Date(timeIntervalSince1970: 1_784_179_200) + let resolver = CLIProxyAPIAttributionResolver( + observations: [ + .init(sessionID: "session-1", model: "gpt-5.5", timestamp: timestamp), + ], + usageRecords: [ + Self.record( + timestamp: timestamp, + provider: "codex", + authType: "oauth", + failed: true), + CLIProxyAPIUsageRecord( + timestamp: timestamp.addingTimeInterval(1), + provider: "codex", + model: "gpt-5.5", + alias: "gpt-5.5", + endpoint: "POST /v1/messages", + authType: "oauth", + requestID: "request-mismatch", + tokens: .init(input: 999, output: 999, total: 1998)), + ]) + + let attribution = resolver.attribution( + model: "gpt-5.5", + modelProvider: .openAI, + sessionID: "session-1", + timestampUnixMs: Int64(timestamp.timeIntervalSince1970 * 1000), + tokens: Self.tokens) + + #expect(attribution.route == .cliProxyAPI) + #expect(attribution.upstream == nil) + } + + @Test + func `telemetry index isolates the matching model and time window`() { + let timestamp = Date(timeIntervalSince1970: 1_784_179_200) + let unrelated = (0..<10000).map { index in + CLIProxyAPIUsageRecord( + timestamp: timestamp.addingTimeInterval(TimeInterval(index - 5000)), + provider: "openrouter", + model: "unrelated-model", + alias: "unrelated-model", + endpoint: "POST /v1/messages", + authType: "api_key", + requestID: "unrelated-\(index)", + tokens: .init(input: 10, output: 20, total: 30)) + } + let matching = Self.record( + timestamp: timestamp.addingTimeInterval(1), + provider: "codex", + authType: "oauth") + let resolver = CLIProxyAPIAttributionResolver( + observations: [ + .init(sessionID: "session-1", model: "gpt-5.5", timestamp: timestamp), + ], + usageRecords: unrelated + [matching]) + + let attribution = resolver.attribution( + model: "gpt-5.5", + modelProvider: .openAI, + sessionID: "session-1", + timestampUnixMs: Int64(timestamp.timeIntervalSince1970 * 1000), + tokens: Self.tokens) + + #expect(attribution.upstream?.provider == "codex") + #expect(attribution.upstream?.model == "gpt-5.5") + #expect(attribution.evidence.contains(.cliProxyUsageTelemetry)) + } + + @Test + func `model without correlated request does not claim cliproxyapi`() { + let resolver = CLIProxyAPIAttributionResolver( + observations: [ + .init(sessionID: "other-session", model: "gpt-5.5", timestamp: nil), + ], + usageRecords: [ + Self.record(timestamp: Date(), provider: "codex", authType: "oauth"), + ]) + + let attribution = resolver.attribution( + model: "gpt-5.5", + modelProvider: .openAI, + sessionID: "session-1", + timestampUnixMs: nil, + tokens: Self.tokens) + + #expect(attribution.route == .unknown) + #expect(attribution.upstream == nil) + #expect(attribution.evidence == [.modelProvider]) + } + + @Test + func `filesystem loader correlates sanitized cached telemetry`() throws { + let fileManager = FileManager.default + let root = fileManager.temporaryDirectory + .appendingPathComponent("cliproxy-attribution-\(UUID().uuidString)", isDirectory: true) + let home = root.appendingPathComponent("home", isDirectory: true) + let logs = home.appendingPathComponent("logs", isDirectory: true) + let cacheRoot = root.appendingPathComponent("cache", isDirectory: true) + try fileManager.createDirectory(at: logs, withIntermediateDirectories: true) + defer { try? fileManager.removeItem(at: root) } + + let requestLog = """ + === REQUEST INFO === + URL: /v1/messages + Method: POST + Timestamp: 2026-07-16T12:00:00Z + === HEADERS === + X-Claude-Code-Session-Id: session-1 + === REQUEST BODY === + {"model":"gpt-5.5"} + === RESPONSE === + Status: 200 + """ + try Data(requestLog.utf8).write(to: logs.appendingPathComponent("v1-messages.log")) + let timestamp = try #require(CostUsageDateParser.parse("2026-07-16T12:00:01Z")) + CLIProxyAPIUsageCacheIO.merge( + [Self.record(timestamp: timestamp, provider: "codex", authType: "oauth")], + cacheRoot: cacheRoot, + now: timestamp) + + let resolver = try CLIProxyAPIAttributionResolver.load( + home: home, + cacheRoot: cacheRoot, + fileManager: fileManager) + let attribution = resolver.attribution( + model: "gpt-5.5", + modelProvider: .openAI, + sessionID: "session-1", + timestampUnixMs: Int64(timestamp.timeIntervalSince1970 * 1000), + tokens: Self.tokens) + + #expect(attribution.route == .cliProxyAPI) + #expect(attribution.upstream?.isCodex == true) + #expect(attribution.evidence.contains(.cliProxyUsageTelemetry)) + } + + @Test + func `filesystem loader preserves observations beyond five hundred newer logs`() throws { + let fileManager = FileManager.default + let root = fileManager.temporaryDirectory + .appendingPathComponent("cliproxy-log-window-\(UUID().uuidString)", isDirectory: true) + let home = root.appendingPathComponent("home", isDirectory: true) + let logs = home.appendingPathComponent("logs", isDirectory: true) + try fileManager.createDirectory(at: logs, withIntermediateDirectories: true) + defer { try? fileManager.removeItem(at: root) } + let timestamp = try #require(CostUsageDateParser.parse("2026-01-01T12:00:00Z")) + let targetURL = logs.appendingPathComponent("target.log") + try Data(Self.requestLog( + sessionID: "target-session", + timestamp: timestamp).utf8).write(to: targetURL) + try fileManager.setAttributes([.modificationDate: timestamp], ofItemAtPath: targetURL.path) + for index in 0..<500 { + let newerTimestamp = timestamp.addingTimeInterval(TimeInterval(index + 1)) + let url = logs.appendingPathComponent("newer-\(index).log") + try Data(Self.requestLog( + sessionID: "newer-session-\(index)", + timestamp: newerTimestamp).utf8).write(to: url) + try fileManager.setAttributes( + [.modificationDate: newerTimestamp], + ofItemAtPath: url.path) + } + + let resolver = try CLIProxyAPIAttributionResolver.load(home: home, fileManager: fileManager) + let attribution = resolver.attribution( + model: "gpt-5.5", + modelProvider: .openAI, + sessionID: "target-session", + timestampUnixMs: Int64(timestamp.timeIntervalSince1970 * 1000), + tokens: Self.tokens) + + #expect(attribution.route == .cliProxyAPI) + #expect(attribution.evidence.contains(.cliProxyRequestLog)) + } + + @Test + func `filesystem loader checks cancellation before reading request logs`() throws { + let fileManager = FileManager.default + let root = fileManager.temporaryDirectory + .appendingPathComponent("cliproxy-log-cancellation-\(UUID().uuidString)", isDirectory: true) + let home = root.appendingPathComponent("home", isDirectory: true) + let logs = home.appendingPathComponent("logs", isDirectory: true) + try fileManager.createDirectory(at: logs, withIntermediateDirectories: true) + defer { try? fileManager.removeItem(at: root) } + try Data(Self.requestLog( + sessionID: "cancelled-session", + timestamp: Date()).utf8).write(to: logs.appendingPathComponent("request.log")) + + #expect(throws: CancellationError.self) { + try CLIProxyAPIAttributionResolver.load( + home: home, + fileManager: fileManager, + checkCancellation: { throw CancellationError() }) + } + } + + @Test + func `filesystem loader reuses unchanged logs and refreshes changed paths`() throws { + let fileManager = FileManager.default + let root = fileManager.temporaryDirectory + .appendingPathComponent("cliproxy-log-cache-\(UUID().uuidString)", isDirectory: true) + let home = root.appendingPathComponent("home", isDirectory: true) + let logs = home.appendingPathComponent("logs", isDirectory: true) + let logURL = logs.appendingPathComponent("request.log") + try fileManager.createDirectory(at: logs, withIntermediateDirectories: true) + defer { try? fileManager.removeItem(at: root) } + + let requestTimestamp = try #require(CostUsageDateParser.parse("2026-07-16T12:00:00Z")) + let pinnedModificationDate = Date(timeIntervalSince1970: 1_000_000) + let firstLog = Self.requestLog(sessionID: "session-one", timestamp: requestTimestamp) + let secondLog = Self.requestLog(sessionID: "session-two", timestamp: requestTimestamp) + let thirdLog = Self.requestLog(sessionID: "session-new", timestamp: requestTimestamp) + #expect(firstLog.utf8.count == secondLog.utf8.count) + #expect(secondLog.utf8.count == thirdLog.utf8.count) + + try Data(firstLog.utf8).write(to: logURL) + try fileManager.setAttributes( + [.modificationDate: pinnedModificationDate], + ofItemAtPath: logURL.path) + let firstResolver = try CLIProxyAPIAttributionResolver.load(home: home, fileManager: fileManager) + #expect(Self.route(for: "session-one", resolver: firstResolver) == .cliProxyAPI) + + try Data(secondLog.utf8).write(to: logURL) + try fileManager.setAttributes( + [.modificationDate: pinnedModificationDate], + ofItemAtPath: logURL.path) + let cachedResolver = try CLIProxyAPIAttributionResolver.load(home: home, fileManager: fileManager) + #expect(Self.route(for: "session-one", resolver: cachedResolver) == .cliProxyAPI) + #expect(Self.route(for: "session-two", resolver: cachedResolver) == .unknown) + + let forcedResolver = try CLIProxyAPIAttributionResolver.load( + home: home, + fileManager: fileManager, + forceReload: true) + #expect(Self.route(for: "session-one", resolver: forcedResolver) == .unknown) + #expect(Self.route(for: "session-two", resolver: forcedResolver) == .cliProxyAPI) + + try Data(firstLog.utf8).write(to: logURL) + try fileManager.setAttributes( + [.modificationDate: pinnedModificationDate.addingTimeInterval(1)], + ofItemAtPath: logURL.path) + let refreshedResolver = try CLIProxyAPIAttributionResolver.load(home: home, fileManager: fileManager) + #expect(Self.route(for: "session-one", resolver: refreshedResolver) == .cliProxyAPI) + #expect(Self.route(for: "session-two", resolver: refreshedResolver) == .unknown) + + try fileManager.removeItem(at: logURL) + _ = try CLIProxyAPIAttributionResolver.load(home: home, fileManager: fileManager) + try Data(thirdLog.utf8).write(to: logURL) + try fileManager.setAttributes( + [.modificationDate: pinnedModificationDate.addingTimeInterval(1)], + ofItemAtPath: logURL.path) + let recreatedResolver = try CLIProxyAPIAttributionResolver.load(home: home, fileManager: fileManager) + #expect(Self.route(for: "session-one", resolver: recreatedResolver) == .unknown) + #expect(Self.route(for: "session-new", resolver: recreatedResolver) == .cliProxyAPI) + } + + @Test + func `usage cache never persists source or api key fields`() throws { + let fileManager = FileManager.default + let cacheRoot = fileManager.temporaryDirectory + .appendingPathComponent("cliproxy-cache-\(UUID().uuidString)", isDirectory: true) + defer { try? fileManager.removeItem(at: cacheRoot) } + let payload = """ + [{ + "timestamp":"2026-07-16T12:00:00Z", + "source":"private@example.com", + "api_key":"secret-client-key", + "provider":"codex", + "executor_type":"CodexExecutor", + "model":"gpt-5.5", + "alias":"gpt-5.5", + "endpoint":"POST /v1/messages", + "auth_type":"oauth", + "request_id":"request-1", + "failed":false, + "generate":true, + "tokens":{ + "input_tokens":10, + "output_tokens":20, + "cache_read_tokens":30, + "cache_creation_tokens":40, + "total_tokens":100 + } + }] + """ + let decoder = JSONDecoder() + decoder.dateDecodingStrategy = .iso8601 + let records = try decoder.decode([CLIProxyAPIUsageRecord].self, from: Data(payload.utf8)) + let now = try #require(CostUsageDateParser.parse("2026-07-16T12:00:01Z")) + + #expect(CLIProxyAPIUsageCacheIO.merge(records, cacheRoot: cacheRoot, now: now) == 1) + let persisted = try String( + contentsOf: CLIProxyAPIUsageCacheIO.cacheFileURL(cacheRoot: cacheRoot), + encoding: .utf8) + #expect(!persisted.contains("private@example.com")) + #expect(!persisted.contains("secret-client-key")) + #expect(persisted.contains("\"provider\":\"codex\"")) + } + + @Test + func `usage queue client authenticates and decodes sanitized records`() async throws { + let responseBody = """ + [{ + "timestamp":"2026-07-16T12:00:00.123456789Z", + "source":"private@example.com", + "api_key":"secret-client-key", + "provider":"codex", + "executor_type":"CodexExecutor", + "model":"gpt-5.5", + "alias":"gpt-5.5", + "endpoint":"POST /v1/messages", + "auth_type":"oauth", + "request_id":"request-1", + "failed":false, + "generate":true, + "tokens":{"input_tokens":10,"output_tokens":20,"total_tokens":30} + }] + """ + let client = CLIProxyAPIUsageQueueClient( + settings: .init(managementKey: "management-secret"), + dataLoader: { request in + #expect(request.url?.absoluteString == + "http://127.0.0.1:8317/v0/management/usage-queue?count=100") + #expect(request.value(forHTTPHeaderField: "Authorization") == + "Bearer management-secret") + let url = try #require(request.url) + let response = try #require(HTTPURLResponse( + url: url, + statusCode: 200, + httpVersion: nil, + headerFields: nil)) + return (Data(responseBody.utf8), response) + }) + + let batch = try await client.pop(count: 100) + let records = batch.records + + #expect(records.count == 1) + #expect(batch.receivedCount == 1) + #expect(records[0].provider == "codex") + #expect(records[0].authType == "oauth") + #expect(records[0].tokens.total == 30) + } + + @Test + func `usage collector serializes queue pops and cache merges`() async throws { + let fileManager = FileManager.default + let cacheRoot = fileManager.temporaryDirectory + .appendingPathComponent("cliproxy-collector-\(UUID().uuidString)", isDirectory: true) + defer { try? fileManager.removeItem(at: cacheRoot) } + let probe = CLIProxyAPICollectionConcurrencyProbe() + let timestamp = try #require(CostUsageDateParser.parse("2026-07-16T12:00:00Z")) + + func client(requestID: String, seconds: TimeInterval) throws -> CLIProxyAPIUsageQueueClient { + let record = CLIProxyAPIUsageRecord( + timestamp: timestamp.addingTimeInterval(seconds), + provider: "codex", + model: "gpt-5.5", + alias: "gpt-5.5", + endpoint: "POST /v1/messages", + authType: "oauth", + requestID: requestID, + tokens: .init(input: 10, output: 20, total: 30)) + let encoder = JSONEncoder() + encoder.dateEncodingStrategy = .iso8601 + let data = try encoder.encode([record]) + return CLIProxyAPIUsageQueueClient( + settings: .init(managementKey: "management-secret"), + dataLoader: { request in + await probe.recordCall() + let url = try #require(request.url) + let response = try #require(HTTPURLResponse( + url: url, + statusCode: 200, + httpVersion: nil, + headerFields: nil)) + return (data, response) + }) + } + + let firstClient = try client(requestID: "request-1", seconds: 0) + let secondClient = try client(requestID: "request-2", seconds: 1) + let results = await withTaskGroup( + of: CLIProxyAPIUsageCollectionResult.self, + returning: [CLIProxyAPIUsageCollectionResult].self) + { group in + group.addTask { + await CLIProxyAPIUsageCollector.collect(cacheRoot: cacheRoot, client: firstClient) + } + group.addTask { + await CLIProxyAPIUsageCollector.collect(cacheRoot: cacheRoot, client: secondClient) + } + return await group.reduce(into: []) { $0.append($1) } + } + + #expect(results.allSatisfy { $0 == .collected(1) }) + #expect(await probe.maximumActiveCallCount() == 1) + #expect(Set(CLIProxyAPIUsageCacheIO.load(cacheRoot: cacheRoot).map(\.requestID)) == [ + "request-1", + "request-2", + ]) + } + + @Test + func `usage collector persists a full batch before a later pop fails`() async throws { + let fileManager = FileManager.default + let cacheRoot = fileManager.temporaryDirectory + .appendingPathComponent("cliproxy-partial-\(UUID().uuidString)", isDirectory: true) + defer { try? fileManager.removeItem(at: cacheRoot) } + let timestamp = try #require(CostUsageDateParser.parse("2026-07-16T12:00:00Z")) + let records = (0..<100).map { index in + CLIProxyAPIUsageRecord( + timestamp: timestamp.addingTimeInterval(TimeInterval(index)), + provider: "codex", + model: "gpt-5.5", + alias: "gpt-5.5", + endpoint: "POST /v1/messages", + authType: "oauth", + requestID: "request-\(index)", + tokens: .init(input: 10, output: 20, total: 30)) + } + let encoder = JSONEncoder() + encoder.dateEncodingStrategy = .iso8601 + let sequence = try CLIProxyAPIBatchSequence(firstPayload: encoder.encode(records)) + let client = CLIProxyAPIUsageQueueClient( + settings: .init(managementKey: "management-secret"), + dataLoader: { request in + try await sequence.load(request) + }) + + let result = await CLIProxyAPIUsageCollector.collect(cacheRoot: cacheRoot, client: client) + + #expect(result == .failed("The second queue pop failed.")) + #expect(CLIProxyAPIUsageCacheIO.load(cacheRoot: cacheRoot).count == 100) + } + + @Test + func `usage collector prunes expired cache records when the queue is empty`() async { + let fileManager = FileManager.default + let cacheRoot = fileManager.temporaryDirectory + .appendingPathComponent("cliproxy-prune-\(UUID().uuidString)", isDirectory: true) + defer { try? fileManager.removeItem(at: cacheRoot) } + let now = Date() + let expired = Self.record( + timestamp: now.addingTimeInterval(-367 * 24 * 60 * 60), + provider: "codex", + authType: "oauth") + let current = Self.record( + timestamp: now.addingTimeInterval(-24 * 60 * 60), + provider: "codex", + authType: "oauth") + #expect(CLIProxyAPIUsageCacheIO.merge( + [expired, current], + cacheRoot: cacheRoot, + now: expired.timestamp) == 2) + let client = CLIProxyAPIUsageQueueClient( + settings: .init(managementKey: "management-secret"), + dataLoader: { request in + let url = try #require(request.url) + let response = try #require(HTTPURLResponse( + url: url, + statusCode: 200, + httpVersion: nil, + headerFields: nil)) + return (Data("[]".utf8), response) + }) + + let result = await CLIProxyAPIUsageCollector.collect(cacheRoot: cacheRoot, client: client) + + #expect(result == .collected(0)) + #expect(CLIProxyAPIUsageCacheIO.load(cacheRoot: cacheRoot).map(\.requestID) == [current.requestID]) + } + + @Test + func `usage cache prunes expired records from disk during load`() { + let fileManager = FileManager.default + let cacheRoot = fileManager.temporaryDirectory + .appendingPathComponent("cliproxy-load-retention-\(UUID().uuidString)", isDirectory: true) + defer { try? fileManager.removeItem(at: cacheRoot) } + let now = Date() + let expired = Self.record( + timestamp: now.addingTimeInterval(-367 * 24 * 60 * 60), + provider: "codex", + authType: "oauth") + let current = Self.record( + timestamp: now.addingTimeInterval(-24 * 60 * 60), + provider: "codex", + authType: "oauth") + #expect(CLIProxyAPIUsageCacheIO.merge( + [expired, current], + cacheRoot: cacheRoot, + now: expired.timestamp) == 2) + + #expect(CLIProxyAPIUsageCacheIO.load( + cacheRoot: cacheRoot, + now: now).map(\.requestID) == [current.requestID]) + #expect(CLIProxyAPIUsageCacheIO.load( + cacheRoot: cacheRoot, + now: expired.timestamp).map(\.requestID) == [current.requestID]) + } + + @Test + func `empty usage poll does not rewrite an unchanged cache`() async throws { + let fileManager = FileManager.default + let cacheRoot = fileManager.temporaryDirectory + .appendingPathComponent("cliproxy-empty-poll-\(UUID().uuidString)", isDirectory: true) + defer { try? fileManager.removeItem(at: cacheRoot) } + let current = Self.record( + timestamp: Date(), + provider: "codex", + authType: "oauth") + #expect(CLIProxyAPIUsageCacheIO.merge([current], cacheRoot: cacheRoot) == 1) + let cacheURL = CLIProxyAPIUsageCacheIO.cacheFileURL(cacheRoot: cacheRoot) + let marker = Date(timeIntervalSince1970: 1_700_000_000) + try fileManager.setAttributes([.modificationDate: marker], ofItemAtPath: cacheURL.path) + let before = try #require( + fileManager.attributesOfItem(atPath: cacheURL.path)[.modificationDate] as? Date) + let client = CLIProxyAPIUsageQueueClient( + settings: .init(managementKey: "management-secret"), + dataLoader: { request in + let url = try #require(request.url) + let response = try #require(HTTPURLResponse( + url: url, + statusCode: 200, + httpVersion: nil, + headerFields: nil)) + return (Data("[]".utf8), response) + }) + + let result = await CLIProxyAPIUsageCollector.collect(cacheRoot: cacheRoot, client: client) + let after = try #require( + fileManager.attributesOfItem(atPath: cacheURL.path)[.modificationDate] as? Date) + + #expect(result == .collected(0)) + #expect(after == before) + } + + @Test + func `plain http management url is limited to loopback`() { + #expect(CLIProxyAPIConnectionSettings( + baseURL: "http://127.0.0.1:8317", + managementKey: "secret").isConfigured) + #expect(CLIProxyAPIConnectionSettings( + baseURL: "http://localhost:8317", + managementKey: "secret").isConfigured) + #expect(!CLIProxyAPIConnectionSettings( + baseURL: "http://192.168.1.10:8317", + managementKey: "secret").isConfigured) + #expect(!CLIProxyAPIConnectionSettings( + baseURL: "https://proxy.example.com", + managementKey: "secret").isConfigured) + #expect(CLIProxyAPIConnectionSettings( + baseURL: "https://[::1]:8317", + managementKey: "secret").isConfigured) + } + + private static let tokens = CLIProxyAPIAttributionResolver.TokenSignature( + input: 10, + cacheRead: 30, + cacheCreate: 40, + output: 20) + + private static func route( + for sessionID: String, + resolver: CLIProxyAPIAttributionResolver) -> CostUsageAttribution.Route + { + resolver.attribution( + model: "gpt-5.5", + modelProvider: .openAI, + sessionID: sessionID, + timestampUnixMs: nil, + tokens: self.tokens).route + } + + private static func requestLog(sessionID: String, timestamp: Date) -> String { + """ + === REQUEST INFO === + URL: /v1/messages + Method: POST + Timestamp: \(ISO8601DateFormatter().string(from: timestamp)) + === HEADERS === + X-Claude-Code-Session-Id: \(sessionID) + === REQUEST BODY === + {"model":"gpt-5.5"} + === RESPONSE === + Status: 200 + """ + } + + private static func record( + timestamp: Date, + provider: String, + authType: String, + failed: Bool = false) -> CLIProxyAPIUsageRecord + { + CLIProxyAPIUsageRecord( + timestamp: timestamp, + provider: provider, + executorType: provider == "codex" ? "CodexExecutor" : "OpenAICompatExecutor", + model: "gpt-5.5", + alias: "gpt-5.5", + endpoint: "POST /v1/messages", + authType: authType, + requestID: "request-\(provider)-\(authType)-\(timestamp.timeIntervalSince1970)", + failed: failed, + tokens: .init( + input: 10, + output: 20, + cacheRead: 30, + cacheCreation: 40, + total: 100)) + } +} + +struct CLIProxyAPIAttributionTimestampTests { + @Test + func `undated request log does not confirm a dated request`() { + let timestamp = Date(timeIntervalSince1970: 1_784_179_200) + let resolver = CLIProxyAPIAttributionResolver( + observations: [ + .init(sessionID: "session-1", model: "gpt-5.5", timestamp: nil), + ]) + + let attribution = resolver.attribution( + model: "gpt-5.5", + modelProvider: .openAI, + sessionID: "session-1", + timestampUnixMs: Int64(timestamp.timeIntervalSince1970 * 1000), + tokens: nil) + + #expect(attribution.route == .unknown) + #expect(attribution.upstream == nil) + #expect(attribution.evidence == [.modelProvider]) + } +} + +struct CLIProxyAPIAttributionEqualTimestampTests { + @Test + func `equal timestamp logs retain distinct observations for token matched telemetry`() { + let timestamp = Date(timeIntervalSince1970: 1_784_179_200) + let firstTokens = CLIProxyAPIAttributionResolver.TokenSignature( + input: 10, + cacheRead: 0, + cacheCreate: 0, + output: 20) + let secondTokens = CLIProxyAPIAttributionResolver.TokenSignature( + input: 30, + cacheRead: 0, + cacheCreate: 0, + output: 40) + let resolver = CLIProxyAPIAttributionResolver( + observations: [ + .init(sourceID: "first.log", sessionID: "session-1", model: "gpt-5.5", timestamp: timestamp), + .init(sourceID: "second.log", sessionID: "session-1", model: "gpt-5.5", timestamp: timestamp), + ], + usageRecords: [ + CLIProxyAPIUsageRecord( + timestamp: timestamp, + provider: "codex", + model: "gpt-5.5", + alias: "gpt-5.5", + endpoint: "POST /v1/messages", + authType: "oauth", + requestID: "request-first", + tokens: .init(input: 10, output: 20, total: 30)), + CLIProxyAPIUsageRecord( + timestamp: timestamp, + provider: "openrouter", + model: "gpt-5.5", + alias: "gpt-5.5", + endpoint: "POST /v1/messages", + authType: "api_key", + requestID: "request-second", + tokens: .init(input: 30, output: 40, total: 70)), + ]) + let attributions = resolver.attributions(for: [firstTokens, secondTokens].map { tokens in + CLIProxyAPIAttributionResolver.Request( + model: "gpt-5.5", + modelProvider: .openAI, + sessionID: "session-1", + timestampUnixMs: Int64(timestamp.timeIntervalSince1970 * 1000), + tokens: tokens) + }) + + #expect(attributions.map(\.route) == [.cliProxyAPI, .cliProxyAPI]) + #expect(attributions.map(\.upstream?.provider) == ["codex", "openrouter"]) + #expect(attributions.allSatisfy { $0.evidence.contains(.cliProxyUsageTelemetry) }) + } +} + +struct CLIProxyAPIAttributionEndpointTests { + @Test(arguments: [ + "/v1/messages", + "/v1/messages?beta=true", + "POST /v1/messages", + "https://localhost:8317/v1/messages?beta=true", + ]) + func `accepts Claude generation endpoint URL variants`(_ endpoint: String) { + #expect(CLIProxyAPIAttributionResolver.isClaudeMessagesGenerationEndpoint(endpoint)) + } + + @Test(arguments: [ + "/v1/messages/count_tokens", + "/v1/messages/batches", + "/v1/messageship", + "POST /v1/messages/count_tokens", + ]) + func `rejects non generation Claude endpoint variants`(_ endpoint: String) { + #expect(!CLIProxyAPIAttributionResolver.isClaudeMessagesGenerationEndpoint(endpoint)) + } +} + +private actor CLIProxyAPICollectionConcurrencyProbe { + private var activeCallCount = 0 + private var maximumActiveCount = 0 + + func recordCall() async { + self.activeCallCount += 1 + self.maximumActiveCount = max(self.maximumActiveCount, self.activeCallCount) + try? await Task.sleep(for: .milliseconds(100)) + self.activeCallCount -= 1 + } + + func maximumActiveCallCount() -> Int { + self.maximumActiveCount + } +} + +private actor CLIProxyAPIBatchSequence { + private enum SequenceError: LocalizedError { + case secondPopFailed + + var errorDescription: String? { + "The second queue pop failed." + } + } + + private let firstPayload: Data + private var callCount = 0 + + init(firstPayload: Data) { + self.firstPayload = firstPayload + } + + func load(_ request: URLRequest) throws -> (Data, URLResponse) { + self.callCount += 1 + guard self.callCount == 1 else { + throw SequenceError.secondPopFailed + } + let url = try #require(request.url) + let response = try #require(HTTPURLResponse( + url: url, + statusCode: 200, + httpVersion: nil, + headerFields: nil)) + return (self.firstPayload, response) + } +} diff --git a/Tests/CodexBarTests/CLIProxyAPIUsageCacheTests.swift b/Tests/CodexBarTests/CLIProxyAPIUsageCacheTests.swift new file mode 100644 index 0000000000..91ea41aaa0 --- /dev/null +++ b/Tests/CodexBarTests/CLIProxyAPIUsageCacheTests.swift @@ -0,0 +1,1229 @@ +import Foundation +import Testing +@testable import CodexBarCore + +struct CLIProxyAPIUsageCacheTests { + @Test + func `cost cache clear advances the durable generation for other processes`() throws { + let fileManager = FileManager.default + let root = fileManager.temporaryDirectory + .appendingPathComponent("cliproxy-clear-generation-\(UUID().uuidString)", isDirectory: true) + let costUsage = root.appendingPathComponent("cost-usage", isDirectory: true) + defer { try? fileManager.removeItem(at: root) } + + try fileManager.createDirectory(at: costUsage, withIntermediateDirectories: true) + try Data("cache".utf8).write(to: costUsage.appendingPathComponent("claude-v6.json")) + let initialUpdate = try #require(CostUsageCacheLocations.prepareCLIProxyAPIConfigurationGenerationUpdate( + stateRoot: root, + fileManager: fileManager)) + #expect(CostUsageCacheLocations.commitCLIProxyAPIConfigurationGenerationUpdate( + initialUpdate, + fileManager: fileManager)) + let initialGeneration = CostUsageCacheLocations.cliProxyAPIConfigurationGeneration( + stateRoot: root, + fileManager: fileManager) + + let result = CostUsageCacheLocations.clearAllCostUsageCaches( + in: [costUsage], + stateRoot: root, + fileManager: fileManager) + + #expect(result == CostUsageCacheClearResult(cleared: 1, errorDescription: nil)) + #expect(CostUsageCacheLocations.cliProxyAPIConfigurationGeneration( + stateRoot: root, + fileManager: fileManager) != initialGeneration) + } + + @Test + func `integration cleanup removes telemetry pending and derived Claude cache artifacts`() throws { + let fileManager = FileManager.default + let root = fileManager.temporaryDirectory + .appendingPathComponent("cliproxy-cleanup-\(UUID().uuidString)", isDirectory: true) + let legacy = root + .appendingPathComponent("legacy", isDirectory: true) + .appendingPathComponent("CodexBar", isDirectory: true) + .appendingPathComponent("cost-usage", isDirectory: true) + let durable = root + .appendingPathComponent("durable", isDirectory: true) + .appendingPathComponent("CodexBar", isDirectory: true) + .appendingPathComponent("cost-usage", isDirectory: true) + defer { try? fileManager.removeItem(at: root) } + + for directory in [legacy, durable] { + try fileManager.createDirectory(at: directory, withIntermediateDirectories: true) + try Data("usage".utf8).write(to: directory.appendingPathComponent( + CostUsageCacheLocations.cliProxyAPIUsageFileName)) + try Data("pending".utf8).write(to: directory.appendingPathComponent( + CostUsageCacheLocations.cliProxyAPIPendingFileName)) + } + for directory in [legacy, durable] { + let claudeCache = CostUsageCacheIO.cacheFileURL( + provider: .claude, + cacheRoot: directory.deletingLastPathComponent()) + try Data("derived attribution".utf8).write(to: claudeCache) + } + let unrelated = durable.appendingPathComponent("codex-v11.json") + try Data("keep".utf8).write(to: unrelated) + + let cleared = CostUsageCacheLocations.clearCLIProxyAPIArtifacts( + in: [legacy, durable], + stateRoot: root, + fileManager: fileManager) + + #expect(cleared) + for directory in [legacy, durable] { + #expect(!fileManager.fileExists(atPath: directory.appendingPathComponent( + CostUsageCacheLocations.cliProxyAPIUsageFileName).path)) + #expect(!fileManager.fileExists(atPath: directory.appendingPathComponent( + CostUsageCacheLocations.cliProxyAPIPendingFileName).path)) + #expect(!fileManager.fileExists(atPath: CostUsageCacheIO.cacheFileURL( + provider: .claude, + cacheRoot: directory.deletingLastPathComponent()).path)) + } + #expect(fileManager.fileExists(atPath: unrelated.path)) + } + + @Test + func `integration cleanup waits for the collector interprocess lock`() async throws { + let root = FileManager.default.temporaryDirectory + .appendingPathComponent("cliproxy-cleanup-lock-\(UUID().uuidString)", isDirectory: true) + let costUsage = root.appendingPathComponent("cost-usage", isDirectory: true) + try FileManager.default.createDirectory(at: costUsage, withIntermediateDirectories: true) + let usageFile = costUsage.appendingPathComponent(CostUsageCacheLocations.cliProxyAPIUsageFileName) + try Data("telemetry".utf8).write(to: usageFile) + defer { try? FileManager.default.removeItem(at: root) } + + let lockAcquired = DispatchSemaphore(value: 0) + let releaseLock = DispatchSemaphore(value: 0) + let clearStarted = DispatchSemaphore(value: 0) + let clearFinished = DispatchSemaphore(value: 0) + let collector = Task.detached { + try await CostUsageCacheLocations.withCLIProxyAPIInterprocessLock(stateRoot: root) { + lockAcquired.signal() + _ = await Self.waitForSignal(releaseLock, timeout: .distantFuture) + } + } + #expect(await Self.waitForSignal(lockAcquired, timeout: .now() + 1)) + + let clear = Task.detached { + clearStarted.signal() + let result = CostUsageCacheLocations.clearCLIProxyAPIArtifacts( + in: [costUsage], + stateRoot: root, + fileManager: .default) + clearFinished.signal() + return result + } + #expect(await Self.waitForSignal(clearStarted, timeout: .now() + 1)) + let finishedBeforeRelease = await Self.waitForSignal( + clearFinished, + timeout: .now() + .milliseconds(50)) + #expect(!finishedBeforeRelease) + + releaseLock.signal() + try await collector.value + #expect(await clear.value) + #expect(!FileManager.default.fileExists(atPath: usageFile.path)) + } + + @Test + func `explicit disconnect state survives artifact cleanup and can be cleared on reconnect`() throws { + let fileManager = FileManager.default + let root = fileManager.temporaryDirectory + .appendingPathComponent("cliproxy-disconnect-\(UUID().uuidString)", isDirectory: true) + let costUsage = root.appendingPathComponent("cost-usage", isDirectory: true) + defer { try? fileManager.removeItem(at: root) } + + #expect(CostUsageCacheLocations.setCLIProxyAPIExplicitlyDisconnected( + true, + stateRoot: root, + fileManager: fileManager)) + #expect(CostUsageCacheLocations.isCLIProxyAPIExplicitlyDisconnected( + stateRoot: root, + fileManager: fileManager)) + try fileManager.createDirectory(at: costUsage, withIntermediateDirectories: true) + #expect(CostUsageCacheLocations.clearCLIProxyAPIArtifacts( + in: [costUsage], + stateRoot: root, + fileManager: fileManager)) + #expect(CostUsageCacheLocations.isCLIProxyAPIExplicitlyDisconnected( + stateRoot: root, + fileManager: fileManager)) + + #expect(CostUsageCacheLocations.setCLIProxyAPIExplicitlyDisconnected( + false, + stateRoot: root, + fileManager: fileManager)) + #expect(!CostUsageCacheLocations.isCLIProxyAPIExplicitlyDisconnected( + stateRoot: root, + fileManager: fileManager)) + } + + @Test + func `explicit disconnect prevents saved connection settings from loading`() { + var didReadStoredSettings = false + let loaded = CLIProxyAPIConnectionSettingsStore.load( + isDisconnected: { true }, + loadStored: { + didReadStoredSettings = true + return CLIProxyAPIConnectionSettings(managementKey: "test-management-key") + }) + + #expect(loaded == nil) + #expect(!didReadStoredSettings) + } + + @Test + func `reconnect rolls back saved credentials when disconnect state cannot be cleared`() { + let root = FileManager.default.temporaryDirectory + .appendingPathComponent("cliproxy-disconnect-clear-failure-\(UUID().uuidString)", isDirectory: true) + defer { try? FileManager.default.removeItem(at: root) } + let didStore = LockIsolated(false) + let didRollback = LockIsolated(false) + let saved = CLIProxyAPIConnectionSettingsStore.saveSerialized( + CLIProxyAPIConnectionSettings(managementKey: "test-management-key"), + stateRoot: root, + fileManager: .default, + operations: .init( + isDisconnected: { true }, + loadStored: { .missing }, + store: { _ in + didStore.setValue(true) + return true + }, + setDisconnectedState: { disconnected in disconnected }, + restore: { _ in + didRollback.setValue(true) + return true + })) + + #expect(!saved) + #expect(didStore.value) + #expect(didRollback.value) + } + + @Test + func `reconnect publishes telemetry invalidation before storing credentials`() throws { + let fileManager = FileManager.default + let root = fileManager.temporaryDirectory + .appendingPathComponent("cliproxy-reconnect-stage-\(UUID().uuidString)", isDirectory: true) + let costUsage = root.appendingPathComponent("cost-usage", isDirectory: true) + let usageFile = costUsage.appendingPathComponent(CostUsageCacheLocations.cliProxyAPIUsageFileName) + try fileManager.createDirectory(at: costUsage, withIntermediateDirectories: true) + try Data("telemetry".utf8).write(to: usageFile) + defer { try? fileManager.removeItem(at: root) } + + let artifactWasStagedAtStore = LockIsolated(false) + let generationWasPublishedAtStore = LockIsolated(false) + let saved = CLIProxyAPIConnectionSettingsStore.saveSerialized( + CLIProxyAPIConnectionSettings(managementKey: "test-management-key"), + artifactDirectories: [costUsage], + stateRoot: root, + fileManager: fileManager, + operations: .init( + isDisconnected: { true }, + loadStored: { .missing }, + store: { _ in + artifactWasStagedAtStore.setValue(!FileManager.default.fileExists(atPath: usageFile.path)) + generationWasPublishedAtStore.setValue( + CostUsageCacheLocations.cliProxyAPIConfigurationGeneration( + stateRoot: root, + fileManager: .default) != nil) + return true + }, + setDisconnectedState: { _ in true }, + restore: { _ in true })) + + #expect(saved) + #expect(artifactWasStagedAtStore.value) + #expect(generationWasPublishedAtStore.value) + #expect(!fileManager.fileExists(atPath: usageFile.path)) + } + + @Test + func `active configuration replacement plans a telemetry purge`() { + let existing = CLIProxyAPIConnectionSettings( + baseURL: "http://127.0.0.1:8317", + managementKey: "old-management-key") + let replacement = CLIProxyAPIConnectionSettings( + baseURL: "http://127.0.0.1:8318", + managementKey: "new-management-key") + #expect(CLIProxyAPIConnectionSettingsStore.artifactDisposition( + existing, + isDisconnected: false, + storedSettings: .found(existing)) == .preserve) + + #expect(CLIProxyAPIConnectionSettingsStore.artifactDisposition( + replacement, + isDisconnected: false, + storedSettings: .found(existing)) == .purge) + } + + @Test + func `replacement keeps prior telemetry when credential storage fails`() throws { + let fileManager = FileManager.default + let root = fileManager.temporaryDirectory + .appendingPathComponent("cliproxy-replacement-store-failure-\(UUID().uuidString)", isDirectory: true) + let costUsage = root.appendingPathComponent("cost-usage", isDirectory: true) + let usageFile = costUsage.appendingPathComponent(CostUsageCacheLocations.cliProxyAPIUsageFileName) + try fileManager.createDirectory(at: costUsage, withIntermediateDirectories: true) + try Data("telemetry".utf8).write(to: usageFile) + defer { try? fileManager.removeItem(at: root) } + let existing = CLIProxyAPIConnectionSettings(managementKey: "old-management-key") + let replacement = CLIProxyAPIConnectionSettings(managementKey: "new-management-key") + let disconnected = LockIsolated(false) + let wasIsolatedAtStore = LockIsolated(false) + + let saved = CLIProxyAPIConnectionSettingsStore.saveSerialized( + replacement, + artifactDirectories: [costUsage], + stateRoot: root, + fileManager: fileManager, + operations: .init( + isDisconnected: { disconnected.value }, + loadStored: { .found(existing) }, + store: { _ in + wasIsolatedAtStore.setValue(disconnected.value) + return false + }, + setDisconnectedState: { + disconnected.setValue($0) + return true + }, + restore: { _ in true })) + + #expect(!saved) + #expect(wasIsolatedAtStore.value) + #expect(!disconnected.value) + #expect(fileManager.fileExists(atPath: usageFile.path)) + } + + @Test + func `failed replacement staging restores already moved telemetry`() throws { + let fileManager = FileManager.default + let root = fileManager.temporaryDirectory + .appendingPathComponent("cliproxy-replacement-stage-failure-\(UUID().uuidString)", isDirectory: true) + let firstDirectory = root.appendingPathComponent("first/cost-usage", isDirectory: true) + let secondDirectory = root.appendingPathComponent("second/cost-usage", isDirectory: true) + let firstURL = firstDirectory.appendingPathComponent(CostUsageCacheLocations.cliProxyAPIUsageFileName) + let secondURL = secondDirectory.appendingPathComponent(CostUsageCacheLocations.cliProxyAPIUsageFileName) + for url in [firstURL, secondURL] { + try fileManager.createDirectory(at: url.deletingLastPathComponent(), withIntermediateDirectories: true) + try Data("telemetry".utf8).write(to: url) + } + defer { try? fileManager.removeItem(at: root) } + + let update = CostUsageCacheLocations.prepareCLIProxyAPIArtifactsUpdate( + in: [firstDirectory, secondDirectory], + fileExists: { fileManager.fileExists(atPath: $0.path) }, + moveItem: { source, destination in + if source == secondURL { + throw CocoaError(.fileWriteUnknown) + } + try fileManager.moveItem(at: source, to: destination) + }) + + #expect(update == nil) + #expect(fileManager.fileExists(atPath: firstURL.path)) + #expect(fileManager.fileExists(atPath: secondURL.path)) + } + + @Test + func `save and removal advance the durable configuration generation`() throws { + let fileManager = FileManager.default + let root = fileManager.temporaryDirectory + .appendingPathComponent("cliproxy-generation-\(UUID().uuidString)", isDirectory: true) + defer { try? fileManager.removeItem(at: root) } + let settings = CLIProxyAPIConnectionSettings(managementKey: "test-management-key") + + #expect(CLIProxyAPIConnectionSettingsStore.saveSerialized( + settings, + stateRoot: root, + fileManager: fileManager, + operations: .init( + isDisconnected: { false }, + loadStored: { .missing }, + store: { _ in true }, + setDisconnectedState: { _ in true }, + restore: { _ in true }))) + let savedGeneration = try #require(CostUsageCacheLocations.cliProxyAPIConfigurationGeneration( + stateRoot: root, + fileManager: fileManager)) + + #expect(CLIProxyAPIConnectionSettingsStore.removeAndPurgeTelemetry( + in: [], + stateRoot: root, + fileManager: fileManager, + operations: .init( + isDisconnected: { false }, + loadStored: { .missing }, + clearConfiguration: { true }, + setDisconnectedState: { _ in true }, + restore: { _ in true })) == .removed) + let removedGeneration = try #require(CostUsageCacheLocations.cliProxyAPIConfigurationGeneration( + stateRoot: root, + fileManager: fileManager)) + + #expect(removedGeneration != savedGeneration) + } + + @Test + func `failed credential save keeps its published telemetry invalidation`() { + let fileManager = FileManager.default + let root = fileManager.temporaryDirectory + .appendingPathComponent("cliproxy-generation-failure-\(UUID().uuidString)", isDirectory: true) + defer { try? fileManager.removeItem(at: root) } + + #expect(!CLIProxyAPIConnectionSettingsStore.saveSerialized( + CLIProxyAPIConnectionSettings(managementKey: "test-management-key"), + stateRoot: root, + fileManager: fileManager, + operations: .init( + isDisconnected: { false }, + loadStored: { .missing }, + store: { _ in false }, + setDisconnectedState: { _ in true }, + restore: { _ in true }))) + #expect(CostUsageCacheLocations.cliProxyAPIConfigurationGeneration( + stateRoot: root, + fileManager: fileManager) != nil) + } + + @Test + func `generation publication failure restores replacement credentials state and telemetry`() throws { + let fileManager = FileManager.default + let root = fileManager.temporaryDirectory + .appendingPathComponent("cliproxy-generation-commit-failure-\(UUID().uuidString)", isDirectory: true) + let costUsage = root.appendingPathComponent("cost-usage", isDirectory: true) + let usageFile = costUsage.appendingPathComponent(CostUsageCacheLocations.cliProxyAPIUsageFileName) + let generationURL = root.appendingPathComponent( + "cliproxyapi-configuration-generation-v1", + isDirectory: false) + try fileManager.createDirectory(at: costUsage, withIntermediateDirectories: true) + try Data("telemetry".utf8).write(to: usageFile) + try fileManager.createDirectory(at: generationURL, withIntermediateDirectories: true) + defer { try? fileManager.removeItem(at: root) } + let existing = CLIProxyAPIConnectionSettings(managementKey: "old-management-key") + let replacement = CLIProxyAPIConnectionSettings(managementKey: "new-management-key") + let storedSettings = LockIsolated(existing) + let disconnected = LockIsolated(true) + let didStore = LockIsolated(false) + + let saved = CLIProxyAPIConnectionSettingsStore.saveSerialized( + replacement, + artifactDirectories: [costUsage], + stateRoot: root, + fileManager: fileManager, + operations: .init( + isDisconnected: { disconnected.value }, + loadStored: { .found(storedSettings.value) }, + store: { settings in + didStore.setValue(true) + storedSettings.setValue(settings) + return true + }, + setDisconnectedState: { value in + disconnected.setValue(value) + return true + }, + restore: { snapshot in + guard case let .found(settings) = snapshot else { return false } + storedSettings.setValue(settings) + return true + })) + + #expect(!saved) + #expect(!didStore.value) + #expect(storedSettings.value == existing) + #expect(disconnected.value) + #expect(fileManager.fileExists(atPath: usageFile.path)) + } + + @Test + func `generation publication failure rolls back configuration removal and telemetry`() throws { + let fileManager = FileManager.default + let root = fileManager.temporaryDirectory + .appendingPathComponent("cliproxy-removal-generation-failure-\(UUID().uuidString)", isDirectory: true) + let costUsage = root.appendingPathComponent("cost-usage", isDirectory: true) + let usageFile = costUsage.appendingPathComponent(CostUsageCacheLocations.cliProxyAPIUsageFileName) + let generationURL = root.appendingPathComponent( + "cliproxyapi-configuration-generation-v1", + isDirectory: false) + try fileManager.createDirectory(at: costUsage, withIntermediateDirectories: true) + try Data("telemetry".utf8).write(to: usageFile) + try fileManager.createDirectory(at: generationURL, withIntermediateDirectories: true) + defer { try? fileManager.removeItem(at: root) } + let existing = CLIProxyAPIConnectionSettings(managementKey: "old-management-key") + let storedSettings = LockIsolated(existing) + let disconnected = LockIsolated(false) + let didClear = LockIsolated(false) + + let result = CLIProxyAPIConnectionSettingsStore.removeAndPurgeTelemetry( + in: [costUsage], + stateRoot: root, + fileManager: fileManager, + operations: .init( + isDisconnected: { disconnected.value }, + loadStored: { .found(existing) }, + clearConfiguration: { + didClear.setValue(true) + return false + }, + setDisconnectedState: { value in + disconnected.setValue(value) + return true + }, + restore: { snapshot in + guard case let .found(settings) = snapshot else { return false } + storedSettings.setValue(settings) + return true + })) + + #expect(result == .configurationRemovalFailed) + #expect(!didClear.value) + #expect(storedSettings.value == existing) + #expect(!disconnected.value) + #expect(fileManager.fileExists(atPath: usageFile.path)) + } + + @Test + func `configuration removal waits for an in progress save transaction`() async throws { + let fileManager = FileManager.default + let root = fileManager.temporaryDirectory + .appendingPathComponent("cliproxy-settings-lock-\(UUID().uuidString)", isDirectory: true) + let costUsage = root.appendingPathComponent("cost-usage", isDirectory: true) + try fileManager.createDirectory(at: costUsage, withIntermediateDirectories: true) + let usageFile = costUsage.appendingPathComponent(CostUsageCacheLocations.cliProxyAPIUsageFileName) + try Data("telemetry".utf8).write(to: usageFile) + defer { try? fileManager.removeItem(at: root) } + + let saveEntered = DispatchSemaphore(value: 0) + let releaseSave = DispatchSemaphore(value: 0) + let removalEntered = DispatchSemaphore(value: 0) + let settings = CLIProxyAPIConnectionSettings(managementKey: "test-management-key") + let saveTask = Task.detached { + CLIProxyAPIConnectionSettingsStore.saveSerialized( + settings, + stateRoot: root, + fileManager: .default, + operations: .init( + isDisconnected: { false }, + loadStored: { .missing }, + store: { _ in + saveEntered.signal() + releaseSave.wait() + return true + }, + setDisconnectedState: { _ in true }, + restore: { _ in true })) + } + #expect(await Self.waitForSignal(saveEntered, timeout: .now() + 1)) + + let removalTask = Task.detached { + CLIProxyAPIConnectionSettingsStore.removeAndPurgeTelemetry( + in: [costUsage], + stateRoot: root, + fileManager: .default, + operations: .init( + isDisconnected: { false }, + loadStored: { .missing }, + clearConfiguration: { + removalEntered.signal() + return true + }, + setDisconnectedState: { _ in true }, + restore: { _ in true })) + } + let removalEnteredBeforeSaveFinished = await Self.waitForSignal( + removalEntered, + timeout: .now() + .milliseconds(50)) + #expect(!removalEnteredBeforeSaveFinished) + + releaseSave.signal() + #expect(await saveTask.value) + #expect(await removalTask.value == .removed) + #expect(!FileManager.default.fileExists(atPath: usageFile.path)) + } + + @Test + func `explicit disconnect prevents collection with persisted settings`() async { + let fileManager = FileManager.default + let root = fileManager.temporaryDirectory + .appendingPathComponent("cliproxy-disconnected-collection-\(UUID().uuidString)", isDirectory: true) + defer { try? fileManager.removeItem(at: root) } + #expect(CostUsageCacheLocations.setCLIProxyAPIExplicitlyDisconnected( + true, + stateRoot: root, + fileManager: fileManager)) + + let result = await CLIProxyAPIUsageCollector.collect( + cacheRoot: root, + settings: CLIProxyAPIConnectionSettings(managementKey: "secret")) + + #expect(result == .notConfigured) + } + + @Test + func `cost cache locations include durable telemetry storage`() throws { + let fileManager = FileManager.default + let directories = CostUsageCacheLocations.directories(fileManager: fileManager) + let cacheRoot = try #require(fileManager.urls( + for: .cachesDirectory, + in: .userDomainMask).first) + let applicationSupportRoot = try #require(fileManager.urls( + for: .applicationSupportDirectory, + in: .userDomainMask).first) + + #expect(directories == [cacheRoot, applicationSupportRoot].map { root in + root + .appendingPathComponent("CodexBar", isDirectory: true) + .appendingPathComponent("cost-usage", isDirectory: true) + }) + #expect(directories.contains( + CLIProxyAPIUsageCacheIO.cacheFileURL().deletingLastPathComponent())) + #expect(directories.contains( + CLIProxyAPIUsagePendingIO.pendingFileURL().deletingLastPathComponent())) + } + + @Test + func `default telemetry storage is durable application support`() throws { + let fileManager = FileManager.default + let durableURL = CLIProxyAPIUsageCacheIO.cacheFileURL() + let legacyURL = CLIProxyAPIUsageCacheIO.legacyCacheFileURL() + let durableRoot = try #require(fileManager.urls( + for: .applicationSupportDirectory, + in: .userDomainMask).first) + .appendingPathComponent("CodexBar", isDirectory: true) + let legacyRoot = try #require(fileManager.urls( + for: .cachesDirectory, + in: .userDomainMask).first) + .appendingPathComponent("CodexBar", isDirectory: true) + + #expect(durableURL.path.hasPrefix(durableRoot.path + "/")) + #expect(legacyURL.path.hasPrefix(legacyRoot.path + "/")) + #expect(durableURL != legacyURL) + } + + @Test + func `legacy purgeable telemetry migrates into durable storage`() throws { + let fileManager = FileManager.default + let root = fileManager.temporaryDirectory + .appendingPathComponent("cliproxy-usage-migration-\(UUID().uuidString)", isDirectory: true) + let durableRoot = root.appendingPathComponent("application-support", isDirectory: true) + let legacyRoot = root.appendingPathComponent("caches", isDirectory: true) + defer { try? fileManager.removeItem(at: root) } + let timestamp = try #require(CostUsageDateParser.parse("2026-07-16T12:00:00Z")) + let record = CLIProxyAPIUsageRecord( + timestamp: timestamp, + provider: "codex", + model: "gpt-5.4", + alias: "gpt-5.4", + endpoint: "POST /v1/messages", + authType: "oauth", + requestID: "request-1", + tokens: .init(input: 10, output: 20, total: 30)) + #expect(CLIProxyAPIUsageCacheIO.merge( + [record], + cacheRoot: legacyRoot, + now: timestamp) == 1) + let legacyURL = CLIProxyAPIUsageCacheIO.cacheFileURL(cacheRoot: legacyRoot) + #expect(fileManager.fileExists(atPath: legacyURL.path)) + + let migrated = CLIProxyAPIUsageCacheIO.load( + cacheRoot: durableRoot, + legacyCacheRoot: legacyRoot, + now: timestamp) + + #expect(migrated.map(\.requestID) == ["request-1"]) + #expect(fileManager.fileExists( + atPath: CLIProxyAPIUsageCacheIO.cacheFileURL(cacheRoot: durableRoot).path)) + #expect(!fileManager.fileExists(atPath: legacyURL.path)) + } + + @Test + func `legacy migration waits for the collector interprocess lock`() async throws { + let fileManager = FileManager.default + let root = fileManager.temporaryDirectory + .appendingPathComponent("cliproxy-usage-lock-\(UUID().uuidString)", isDirectory: true) + let durableRoot = root.appendingPathComponent("application-support", isDirectory: true) + let legacyRoot = root.appendingPathComponent("caches", isDirectory: true) + defer { try? fileManager.removeItem(at: root) } + let timestamp = try #require(CostUsageDateParser.parse("2026-07-16T12:00:00Z")) + let legacyRecord = Self.record(id: "legacy", timestamp: timestamp) + let collectedRecord = Self.record( + id: "collected", + timestamp: timestamp.addingTimeInterval(1)) + #expect(CLIProxyAPIUsageCacheIO.merge( + [legacyRecord], + cacheRoot: legacyRoot, + now: timestamp) == 1) + + let lockAcquired = DispatchSemaphore(value: 0) + let releaseLock = DispatchSemaphore(value: 0) + let loadStarted = DispatchSemaphore(value: 0) + let loadFinished = DispatchSemaphore(value: 0) + let lockHolder = Task.detached { + try CostUsageCacheLocations.withCLIProxyAPIInterprocessLock(stateRoot: root) { + lockAcquired.signal() + releaseLock.wait() + #expect(CLIProxyAPIUsageCacheIO.merge( + [collectedRecord], + cacheRoot: durableRoot, + legacyCacheRoot: nil, + now: timestamp) == 1) + } + } + #expect(await Self.waitForSignal(lockAcquired, timeout: .now() + 1)) + + let loadTask = Task.detached { + loadStarted.signal() + let records = CLIProxyAPIUsageCacheIO.load( + cacheRoot: durableRoot, + legacyCacheRoot: legacyRoot, + now: timestamp) + loadFinished.signal() + return records + } + #expect(await Self.waitForSignal(loadStarted, timeout: .now() + 1)) + let loadFinishedBeforeRelease = await Self.waitForSignal( + loadFinished, + timeout: .now() + .milliseconds(50)) + #expect(!loadFinishedBeforeRelease) + + releaseLock.signal() + try await lockHolder.value + #expect(await Set(loadTask.value.map(\.requestID)) == ["legacy", "collected"]) + let finalRecords = CLIProxyAPIUsageCacheIO.load( + cacheRoot: durableRoot, + legacyCacheRoot: legacyRoot, + now: timestamp) + #expect(Set(finalRecords.map(\.requestID)) == ["legacy", "collected"]) + } + + @Test + func `fallback record identity survives cache round trips within one second`() throws { + let fileManager = FileManager.default + let root = fileManager.temporaryDirectory + .appendingPathComponent("cliproxy-usage-fractional-\(UUID().uuidString)", isDirectory: true) + defer { try? fileManager.removeItem(at: root) } + let second = try #require(CostUsageDateParser.parse("2026-07-16T12:00:00Z")) + let records = [ + Self.record(id: "", timestamp: second.addingTimeInterval(0.1)), + Self.record(id: "", timestamp: second.addingTimeInterval(0.9)), + ] + + #expect(CLIProxyAPIUsageCacheIO.merge( + records, + cacheRoot: root, + now: second) == 2) + #expect(CLIProxyAPIUsageCacheIO.merge( + records, + cacheRoot: root, + now: second) == 0) + + let roundTripped = CLIProxyAPIUsageCacheIO.load( + cacheRoot: root, + now: second) + #expect(roundTripped.count == 2) + #expect( + roundTripped.map { Int64($0.timestamp.timeIntervalSince1970 * 1000) } + == records.map { Int64($0.timestamp.timeIntervalSince1970 * 1000) }) + } + + @Test + func `fallback record identity preserves identical occurrences across batches and replay`() throws { + let fileManager = FileManager.default + let root = fileManager.temporaryDirectory + .appendingPathComponent("cliproxy-usage-identical-\(UUID().uuidString)", isDirectory: true) + defer { try? fileManager.removeItem(at: root) } + let timestamp = try #require(CostUsageDateParser.parse("2026-07-16T12:00:00.123Z")) + let records = [ + Self.record(id: "", timestamp: timestamp).assigningNewLocalOccurrenceID(), + Self.record(id: "", timestamp: timestamp).assigningNewLocalOccurrenceID(), + ] + #expect(records[0].localOccurrenceID != records[1].localOccurrenceID) + + #expect(CLIProxyAPIUsageCacheIO.merge( + [records[0]], + cacheRoot: root, + now: timestamp) == 1) + #expect(CLIProxyAPIUsageCacheIO.merge( + [records[0]], + cacheRoot: root, + now: timestamp) == 0) + #expect(CLIProxyAPIUsageCacheIO.merge( + [records[1]], + cacheRoot: root, + now: timestamp) == 1) + #expect(CLIProxyAPIUsageCacheIO.merge( + [records[1]], + cacheRoot: root, + now: timestamp) == 0) + + let roundTripped = CLIProxyAPIUsageCacheIO.load( + cacheRoot: root, + now: timestamp) + #expect(roundTripped.count == 2) + #expect(Set(roundTripped.compactMap(\.localOccurrenceID)) == Set(records.compactMap(\.localOccurrenceID))) + } + + @Test + func `corrupt durable cache is preserved instead of overwritten`() throws { + let fileManager = FileManager.default + let root = fileManager.temporaryDirectory + .appendingPathComponent("cliproxy-corrupt-cache-\(UUID().uuidString)", isDirectory: true) + defer { try? fileManager.removeItem(at: root) } + let cacheURL = CLIProxyAPIUsageCacheIO.cacheFileURL(cacheRoot: root) + try fileManager.createDirectory( + at: cacheURL.deletingLastPathComponent(), + withIntermediateDirectories: true) + let corruptData = Data(#"{"version":2,"records":[]}"#.utf8) + try corruptData.write(to: cacheURL) + let timestamp = try #require(CostUsageDateParser.parse("2026-07-16T12:00:00Z")) + + let result = CLIProxyAPIUsageCacheIO.merge( + [Self.record(id: "new", timestamp: timestamp)], + cacheRoot: root, + now: timestamp) + + #expect(result == nil) + #expect(try Data(contentsOf: cacheURL) == corruptData) + } + + @Test + func `fallback record identity survives pending journal round trips within one second`() throws { + let fileManager = FileManager.default + let root = fileManager.temporaryDirectory + .appendingPathComponent("cliproxy-pending-fractional-\(UUID().uuidString)", isDirectory: true) + defer { try? fileManager.removeItem(at: root) } + let second = try #require(CostUsageDateParser.parse("2026-07-16T12:00:00Z")) + let records = [ + Self.record(id: "", timestamp: second.addingTimeInterval(0.1)), + Self.record(id: "", timestamp: second.addingTimeInterval(0.9)), + ] + + #expect(CLIProxyAPIUsagePendingIO.save(records, pendingRoot: root)) + let roundTripped = try #require(CLIProxyAPIUsagePendingIO.load(pendingRoot: root)) + + #expect(roundTripped.count == 2) + #expect( + roundTripped.map { Int64($0.timestamp.timeIntervalSince1970 * 1000) } + == records.map { Int64($0.timestamp.timeIntervalSince1970 * 1000) }) + } + + @Test + func `pending journal prunes expired records without collection`() throws { + let fileManager = FileManager.default + let root = fileManager.temporaryDirectory + .appendingPathComponent("cliproxy-pending-retention-\(UUID().uuidString)", isDirectory: true) + defer { try? fileManager.removeItem(at: root) } + let now = try #require(CostUsageDateParser.parse("2026-07-30T12:00:00Z")) + let records = [ + Self.record(id: "expired", timestamp: now.addingTimeInterval(-367 * 24 * 60 * 60)), + Self.record(id: "retained", timestamp: now.addingTimeInterval(-365 * 24 * 60 * 60)), + ] + + #expect(CLIProxyAPIUsagePendingIO.save(records, pendingRoot: root)) + #expect(CLIProxyAPIUsagePendingIO.load(pendingRoot: root, now: now)?.map(\.requestID) == ["retained"]) + #expect(CLIProxyAPIUsagePendingIO.load(pendingRoot: root, now: now)?.map(\.requestID) == ["retained"]) + } + + @Test + func `collector maintenance prunes pending and durable usage independently of collection`() throws { + let fileManager = FileManager.default + let root = fileManager.temporaryDirectory + .appendingPathComponent("cliproxy-pending-maintenance-\(UUID().uuidString)", isDirectory: true) + defer { try? fileManager.removeItem(at: root) } + let now = try #require(CostUsageDateParser.parse("2026-07-30T12:00:00Z")) + let records = [ + Self.record(id: "expired", timestamp: now.addingTimeInterval(-367 * 24 * 60 * 60)), + Self.record(id: "retained", timestamp: now.addingTimeInterval(-365 * 24 * 60 * 60)), + ] + + #expect(CLIProxyAPIUsagePendingIO.save(records, pendingRoot: root)) + #expect(CLIProxyAPIUsageCacheIO.merge( + records, + cacheRoot: root, + now: records[0].timestamp) == 2) + #expect(CLIProxyAPIUsageCollector.pruneExpiredUsage( + cacheRoot: root, + pendingRoot: root, + stateRoot: root, + now: now, + fileManager: fileManager)) + #expect(CLIProxyAPIUsagePendingIO.load(pendingRoot: root, now: now)?.map(\.requestID) == ["retained"]) + #expect(CLIProxyAPIUsageCacheIO.load( + cacheRoot: root, + now: now.addingTimeInterval(-365 * 24 * 60 * 60)).map(\.requestID) == ["retained"]) + } + + private static func record(id: String, timestamp: Date) -> CLIProxyAPIUsageRecord { + CLIProxyAPIUsageRecord( + timestamp: timestamp, + provider: "codex", + model: "gpt-5.4", + alias: "gpt-5.4", + endpoint: "POST /v1/messages", + authType: "oauth", + requestID: id, + tokens: .init(input: 10, output: 20, total: 30)) + } + + private static func waitForSignal( + _ semaphore: DispatchSemaphore, + timeout: DispatchTime) async -> Bool + { + await withCheckedContinuation { continuation in + DispatchQueue.global().async { + continuation.resume(returning: semaphore.wait(timeout: timeout) == .success) + } + } + } +} + +extension CLIProxyAPIUsageCacheTests { + @Test + func `failed save marker rollback retains recovery transaction`() throws { + let fileManager = FileManager.default + let root = fileManager.temporaryDirectory + .appendingPathComponent("cliproxy-save-marker-rollback-failure-\(UUID().uuidString)", isDirectory: true) + let costUsage = root.appendingPathComponent("cost-usage", isDirectory: true) + let usageFile = costUsage.appendingPathComponent(CostUsageCacheLocations.cliProxyAPIUsageFileName) + try fileManager.createDirectory(at: costUsage, withIntermediateDirectories: true) + try Data("telemetry".utf8).write(to: usageFile) + defer { try? fileManager.removeItem(at: root) } + let existing = CLIProxyAPIConnectionSettings(managementKey: "old-management-key") + let replacement = CLIProxyAPIConnectionSettings(managementKey: "new-management-key") + let storedSettings = LockIsolated(existing) + let disconnected = LockIsolated(false) + + let saved = CLIProxyAPIConnectionSettingsStore.saveSerialized( + replacement, + artifactDirectories: [costUsage], + stateRoot: root, + fileManager: fileManager, + operations: .init( + isDisconnected: { disconnected.value }, + loadStored: { .found(storedSettings.value) }, + store: { settings in + storedSettings.setValue(settings) + return true + }, + setDisconnectedState: { value in + guard value else { return false } + disconnected.setValue(true) + return true + }, + restore: { snapshot in + guard case let .found(settings) = snapshot else { return false } + storedSettings.setValue(settings) + return true + })) + + #expect(!saved) + #expect(storedSettings.value == existing) + #expect(disconnected.value) + #expect(!fileManager.fileExists(atPath: usageFile.path)) + #expect(fileManager.fileExists( + atPath: root.appendingPathComponent("cliproxyapi-artifacts-transaction-v1.json").path)) + + try CostUsageCacheLocations.withCLIProxyAPIInterprocessLock( + stateRoot: root, + fileManager: fileManager) {} + + #expect(fileManager.fileExists(atPath: usageFile.path)) + #expect(!fileManager.fileExists( + atPath: root.appendingPathComponent("cliproxyapi-artifacts-transaction-v1.json").path)) + } + + @Test + func `failed credential rollback keeps replacement telemetry isolated`() throws { + let fileManager = FileManager.default + let root = fileManager.temporaryDirectory + .appendingPathComponent("cliproxy-credential-rollback-failure-\(UUID().uuidString)", isDirectory: true) + let costUsage = root.appendingPathComponent("cost-usage", isDirectory: true) + let usageFile = costUsage.appendingPathComponent(CostUsageCacheLocations.cliProxyAPIUsageFileName) + try fileManager.createDirectory(at: costUsage, withIntermediateDirectories: true) + try Data("telemetry".utf8).write(to: usageFile) + defer { try? fileManager.removeItem(at: root) } + let existing = CLIProxyAPIConnectionSettings(managementKey: "old-management-key") + let replacement = CLIProxyAPIConnectionSettings(managementKey: "new-management-key") + let storedSettings = LockIsolated(existing) + let disconnected = LockIsolated(false) + let didAttemptRestore = LockIsolated(false) + + let saved = CLIProxyAPIConnectionSettingsStore.saveSerialized( + replacement, + artifactDirectories: [costUsage], + stateRoot: root, + fileManager: fileManager, + operations: .init( + isDisconnected: { disconnected.value }, + loadStored: { .found(storedSettings.value) }, + store: { settings in + storedSettings.setValue(settings) + return true + }, + setDisconnectedState: { value in + guard value else { return false } + disconnected.setValue(true) + return true + }, + restore: { _ in + didAttemptRestore.setValue(true) + return false + })) + + #expect(!saved) + #expect(didAttemptRestore.value) + #expect(storedSettings.value == replacement) + #expect(disconnected.value) + #expect(!fileManager.fileExists(atPath: usageFile.path)) + #expect(fileManager.fileExists( + atPath: root.appendingPathComponent("cliproxyapi-artifacts-transaction-v1.json").path)) + #expect(try fileManager.contentsOfDirectory(at: costUsage, includingPropertiesForKeys: nil) + .contains { $0.lastPathComponent.hasSuffix("replacement-backup") }) + + try CostUsageCacheLocations.withCLIProxyAPIInterprocessLock( + stateRoot: root, + fileManager: fileManager) {} + + #expect(storedSettings.value == replacement) + #expect(disconnected.value) + #expect(!fileManager.fileExists(atPath: usageFile.path)) + #expect(!fileManager.fileExists( + atPath: root.appendingPathComponent("cliproxyapi-artifacts-transaction-v1.json").path)) + #expect(try fileManager.contentsOfDirectory(at: costUsage, includingPropertiesForKeys: nil) + .allSatisfy { !$0.lastPathComponent.hasSuffix("replacement-backup") }) + } + + @Test + func `configuration removal accepts a credential removed after its snapshot`() throws { + let fileManager = FileManager.default + let root = fileManager.temporaryDirectory + .appendingPathComponent("cliproxy-concurrent-removal-\(UUID().uuidString)", isDirectory: true) + let costUsage = root.appendingPathComponent("cost-usage", isDirectory: true) + let usageFile = costUsage.appendingPathComponent(CostUsageCacheLocations.cliProxyAPIUsageFileName) + try fileManager.createDirectory(at: costUsage, withIntermediateDirectories: true) + try Data("telemetry".utf8).write(to: usageFile) + defer { try? fileManager.removeItem(at: root) } + let existing = CLIProxyAPIConnectionSettings(managementKey: "test-management-key") + let disconnected = LockIsolated(false) + + let result = CLIProxyAPIConnectionSettingsStore.removeAndPurgeTelemetry( + in: [costUsage], + stateRoot: root, + fileManager: fileManager, + operations: .init( + isDisconnected: { disconnected.value }, + loadStored: { .found(existing) }, + clearConfiguration: { + CLIProxyAPIConnectionSettingsStore.clearUnserialized( + isDisconnected: { disconnected.value }, + setDisconnectedState: { value in + disconnected.setValue(value) + return true + }, + clearConfiguration: { .missing }) + }, + setDisconnectedState: { value in + disconnected.setValue(value) + return true + }, + restore: { _ in true })) + + #expect(result == .removed) + #expect(disconnected.value) + #expect(!fileManager.fileExists(atPath: usageFile.path)) + } + + @Test + func `failed removal marker rollback retains recovery transaction`() throws { + let fileManager = FileManager.default + let root = fileManager.temporaryDirectory + .appendingPathComponent("cliproxy-removal-marker-rollback-failure-\(UUID().uuidString)", isDirectory: true) + let costUsage = root.appendingPathComponent("cost-usage", isDirectory: true) + let usageFile = costUsage.appendingPathComponent(CostUsageCacheLocations.cliProxyAPIUsageFileName) + try fileManager.createDirectory(at: costUsage, withIntermediateDirectories: true) + try Data("telemetry".utf8).write(to: usageFile) + defer { try? fileManager.removeItem(at: root) } + let existing = CLIProxyAPIConnectionSettings(managementKey: "test-management-key") + let storedSettings = LockIsolated(existing) + let disconnected = LockIsolated(false) + + let result = CLIProxyAPIConnectionSettingsStore.removeAndPurgeTelemetry( + in: [costUsage], + stateRoot: root, + fileManager: fileManager, + operations: .init( + isDisconnected: { disconnected.value }, + loadStored: { .found(existing) }, + clearConfiguration: { + storedSettings.setValue(nil) + disconnected.setValue(true) + return false + }, + setDisconnectedState: { value in + guard value else { return false } + disconnected.setValue(true) + return true + }, + restore: { snapshot in + guard case let .found(settings) = snapshot else { return false } + storedSettings.setValue(settings) + return true + })) + + #expect(result == .configurationRemovalFailed) + #expect(storedSettings.value == existing) + #expect(disconnected.value) + #expect(!fileManager.fileExists(atPath: usageFile.path)) + #expect(fileManager.fileExists( + atPath: root.appendingPathComponent("cliproxyapi-artifacts-transaction-v1.json").path)) + + try CostUsageCacheLocations.withCLIProxyAPIInterprocessLock( + stateRoot: root, + fileManager: fileManager) {} + + #expect(fileManager.fileExists(atPath: usageFile.path)) + #expect(!fileManager.fileExists( + atPath: root.appendingPathComponent("cliproxyapi-artifacts-transaction-v1.json").path)) + } +} + +struct CLIProxyAPITransactionRecoveryTests { + @Test + func `interrupted committed save clears its isolation marker on the next lock`() throws { + let fileManager = FileManager.default + let root = fileManager.temporaryDirectory + .appendingPathComponent("cliproxy-save-marker-finalize-\(UUID().uuidString)", isDirectory: true) + defer { try? fileManager.removeItem(at: root) } + + let generationUpdate = try #require(CostUsageCacheLocations + .prepareCLIProxyAPIConfigurationGenerationUpdate(stateRoot: root, fileManager: fileManager)) + let artifactsUpdate = try #require(CostUsageCacheLocations.prepareCLIProxyAPIArtifactsUpdate( + in: [], + stateRoot: root, + expectedGeneration: generationUpdate.generation, + fileManager: fileManager, + disconnectedStateAfterCommit: false, + disconnectedStateAfterRollback: false, + prepareState: { + CostUsageCacheLocations.setCLIProxyAPIExplicitlyDisconnected( + true, + stateRoot: root, + fileManager: fileManager) + })) + #expect(CostUsageCacheLocations.isCLIProxyAPIExplicitlyDisconnected( + stateRoot: root, + fileManager: fileManager)) + #expect(artifactsUpdate.manifestURL.map { fileManager.fileExists(atPath: $0.path) } == true) + #expect(CostUsageCacheLocations.commitCLIProxyAPIConfigurationGenerationUpdate( + generationUpdate, + fileManager: fileManager)) + + try CostUsageCacheLocations.withCLIProxyAPIInterprocessLock(stateRoot: root, fileManager: fileManager) {} + + #expect(!CostUsageCacheLocations.isCLIProxyAPIExplicitlyDisconnected( + stateRoot: root, + fileManager: fileManager)) + #expect(artifactsUpdate.manifestURL.map { !fileManager.fileExists(atPath: $0.path) } == true) + } + + @Test + func `interrupted uncommitted save restores its previous isolation marker on the next lock`() throws { + let fileManager = FileManager.default + let root = fileManager.temporaryDirectory + .appendingPathComponent("cliproxy-save-marker-rollback-\(UUID().uuidString)", isDirectory: true) + defer { try? fileManager.removeItem(at: root) } + + let generationUpdate = try #require(CostUsageCacheLocations + .prepareCLIProxyAPIConfigurationGenerationUpdate(stateRoot: root, fileManager: fileManager)) + let artifactsUpdate = try #require(CostUsageCacheLocations.prepareCLIProxyAPIArtifactsUpdate( + in: [], + stateRoot: root, + expectedGeneration: generationUpdate.generation, + fileManager: fileManager, + disconnectedStateAfterCommit: false, + disconnectedStateAfterRollback: false, + prepareState: { + CostUsageCacheLocations.setCLIProxyAPIExplicitlyDisconnected( + true, + stateRoot: root, + fileManager: fileManager) + })) + #expect(CostUsageCacheLocations.isCLIProxyAPIExplicitlyDisconnected( + stateRoot: root, + fileManager: fileManager)) + + try CostUsageCacheLocations.withCLIProxyAPIInterprocessLock(stateRoot: root, fileManager: fileManager) {} + + #expect(!CostUsageCacheLocations.isCLIProxyAPIExplicitlyDisconnected( + stateRoot: root, + fileManager: fileManager)) + #expect(artifactsUpdate.manifestURL.map { !fileManager.fileExists(atPath: $0.path) } == true) + CostUsageCacheLocations.discardCLIProxyAPIConfigurationGenerationUpdate( + generationUpdate, + fileManager: fileManager) + } + + @Test + func `interrupted uncommitted replacement restores staged artifacts on the next lock`() throws { + let fileManager = FileManager.default + let root = fileManager.temporaryDirectory + .appendingPathComponent("cliproxy-replacement-recovery-\(UUID().uuidString)", isDirectory: true) + let costUsage = root.appendingPathComponent("cost-usage", isDirectory: true) + let usageFile = costUsage.appendingPathComponent(CostUsageCacheLocations.cliProxyAPIUsageFileName) + try fileManager.createDirectory(at: costUsage, withIntermediateDirectories: true) + try Data("telemetry".utf8).write(to: usageFile) + defer { try? fileManager.removeItem(at: root) } + + let generationUpdate = try #require(CostUsageCacheLocations + .prepareCLIProxyAPIConfigurationGenerationUpdate(stateRoot: root, fileManager: fileManager)) + let artifactsUpdate = try #require(CostUsageCacheLocations.prepareCLIProxyAPIArtifactsUpdate( + in: [costUsage], + stateRoot: root, + expectedGeneration: generationUpdate.generation, + fileManager: fileManager)) + #expect(!fileManager.fileExists(atPath: usageFile.path)) + #expect(artifactsUpdate.manifestURL.map { fileManager.fileExists(atPath: $0.path) } == true) + + try CostUsageCacheLocations.withCLIProxyAPIInterprocessLock(stateRoot: root, fileManager: fileManager) {} + + #expect(fileManager.fileExists(atPath: usageFile.path)) + #expect(artifactsUpdate.moves.allSatisfy { !fileManager.fileExists(atPath: $0.stagedURL.path) }) + #expect(artifactsUpdate.manifestURL.map { !fileManager.fileExists(atPath: $0.path) } == true) + CostUsageCacheLocations.discardCLIProxyAPIConfigurationGenerationUpdate( + generationUpdate, + fileManager: fileManager) + } + + @Test + func `interrupted committed replacement discards staged artifacts on the next lock`() throws { + let fileManager = FileManager.default + let root = fileManager.temporaryDirectory + .appendingPathComponent("cliproxy-replacement-finalize-\(UUID().uuidString)", isDirectory: true) + let costUsage = root.appendingPathComponent("cost-usage", isDirectory: true) + let usageFile = costUsage.appendingPathComponent(CostUsageCacheLocations.cliProxyAPIUsageFileName) + try fileManager.createDirectory(at: costUsage, withIntermediateDirectories: true) + try Data("telemetry".utf8).write(to: usageFile) + defer { try? fileManager.removeItem(at: root) } + + let generationUpdate = try #require(CostUsageCacheLocations + .prepareCLIProxyAPIConfigurationGenerationUpdate(stateRoot: root, fileManager: fileManager)) + let artifactsUpdate = try #require(CostUsageCacheLocations.prepareCLIProxyAPIArtifactsUpdate( + in: [costUsage], + stateRoot: root, + expectedGeneration: generationUpdate.generation, + fileManager: fileManager)) + #expect(CostUsageCacheLocations.commitCLIProxyAPIConfigurationGenerationUpdate( + generationUpdate, + fileManager: fileManager)) + + try CostUsageCacheLocations.withCLIProxyAPIInterprocessLock(stateRoot: root, fileManager: fileManager) {} + + #expect(!fileManager.fileExists(atPath: usageFile.path)) + #expect(artifactsUpdate.moves.allSatisfy { !fileManager.fileExists(atPath: $0.stagedURL.path) }) + #expect(artifactsUpdate.manifestURL.map { !fileManager.fileExists(atPath: $0.path) } == true) + } +} diff --git a/Tests/CodexBarTests/CLIProxyAPIUsageCollectorTests.swift b/Tests/CodexBarTests/CLIProxyAPIUsageCollectorTests.swift new file mode 100644 index 0000000000..f5f921fc93 --- /dev/null +++ b/Tests/CodexBarTests/CLIProxyAPIUsageCollectorTests.swift @@ -0,0 +1,391 @@ +import Foundation +import Testing +@testable import CodexBarCore + +private actor CLIProxyAPICollectionContinuationProbe { + private(set) var popCount = 0 + private var continuationCheckCount = 0 + + func shouldContinue() -> Bool { + self.continuationCheckCount += 1 + return self.continuationCheckCount == 1 + } + + func recordPop() { + self.popCount += 1 + } +} + +struct CLIProxyAPIUsageCollectorTests { + @Test + func `queue client preserves valid records around a malformed entry`() async throws { + let timestamp = try #require(CostUsageDateParser.parse("2026-07-16T12:00:00Z")) + let records = ["request-before", "request-after"].map { requestID in + CLIProxyAPIUsageRecord( + timestamp: timestamp, + provider: "codex", + model: "gpt-5.5", + alias: "gpt-5.5", + endpoint: "POST /v1/messages", + authType: "oauth", + requestID: requestID, + tokens: .init(input: 10, output: 20, total: 30)) + } + let encoder = JSONEncoder() + encoder.dateEncodingStrategy = .iso8601 + let validObjects = try records.map { + try JSONSerialization.jsonObject(with: encoder.encode($0)) + } + let data = try JSONSerialization.data(withJSONObject: [ + validObjects[0], + ["timestamp": "not-a-date"], + validObjects[1], + ]) + let client = CLIProxyAPIUsageQueueClient( + settings: .init(managementKey: "management-secret"), + dataLoader: { request in + let url = try #require(request.url) + let response = try #require(HTTPURLResponse( + url: url, + statusCode: 200, + httpVersion: nil, + headerFields: nil)) + return (data, response) + }) + + let batch = try await client.pop(count: 100) + + #expect(batch.receivedCount == 3) + #expect(batch.records.map(\.requestID) == ["request-before", "request-after"]) + } + + @Test + func `persists an idless popped batch outside a failed cache for the next collection`() async throws { + let fileManager = FileManager.default + let cacheRoot = fileManager.temporaryDirectory + .appendingPathComponent("cliproxy-blocked-\(UUID().uuidString)", isDirectory: false) + let pendingRoot = fileManager.temporaryDirectory + .appendingPathComponent("cliproxy-pending-\(UUID().uuidString)", isDirectory: true) + try Data("not-a-directory".utf8).write(to: cacheRoot) + defer { + try? fileManager.removeItem(at: cacheRoot) + try? fileManager.removeItem(at: pendingRoot) + } + let timestamp = try #require(CostUsageDateParser.parse("2026-07-16T12:00:00Z")) + let record = CLIProxyAPIUsageRecord( + timestamp: timestamp, + provider: "codex", + model: "gpt-5.5", + alias: "gpt-5.5", + endpoint: "POST /v1/messages", + authType: "oauth", + requestID: "", + tokens: .init(input: 10, output: 20, total: 30)) + let encoder = JSONEncoder() + encoder.dateEncodingStrategy = .iso8601 + let data = try encoder.encode([record]) + let client = CLIProxyAPIUsageQueueClient( + settings: .init(managementKey: "management-secret"), + dataLoader: { request in + let url = try #require(request.url) + let response = try #require(HTTPURLResponse( + url: url, + statusCode: 200, + httpVersion: nil, + headerFields: nil)) + return (data, response) + }) + + let result = await CLIProxyAPIUsageCollector.collect( + cacheRoot: cacheRoot, + pendingRoot: pendingRoot, + client: client) + + #expect(result == .failed("Could not save CLIProxyAPI usage telemetry.")) + let pendingOccurrenceID = try #require( + CLIProxyAPIUsagePendingIO.load(pendingRoot: pendingRoot)?.first?.localOccurrenceID) + #expect(!pendingOccurrenceID.isEmpty) + try fileManager.removeItem(at: cacheRoot) + let retryClient = CLIProxyAPIUsageQueueClient( + settings: .init(managementKey: "management-secret"), + dataLoader: { request in + #expect(CLIProxyAPIUsageCacheIO.load(cacheRoot: cacheRoot).map(\.localOccurrenceID) == [ + pendingOccurrenceID, + ]) + let url = try #require(request.url) + let response = try #require(HTTPURLResponse( + url: url, + statusCode: 200, + httpVersion: nil, + headerFields: nil)) + return (Data("[]".utf8), response) + }) + + let retryResult = await CLIProxyAPIUsageCollector.collect( + cacheRoot: cacheRoot, + pendingRoot: pendingRoot, + client: retryClient) + + #expect(retryResult == .collected(1)) + #expect(CLIProxyAPIUsageCacheIO.load(cacheRoot: cacheRoot).map(\.localOccurrenceID) == [pendingOccurrenceID]) + #expect(!fileManager.fileExists( + atPath: CLIProxyAPIUsagePendingIO.pendingFileURL(pendingRoot: pendingRoot).path)) + } + + @Test + func `does not merge a popped batch when staging fails`() async throws { + let fileManager = FileManager.default + let cacheRoot = fileManager.temporaryDirectory + .appendingPathComponent("cliproxy-cache-\(UUID().uuidString)", isDirectory: true) + let pendingRoot = fileManager.temporaryDirectory + .appendingPathComponent("cliproxy-blocked-pending-\(UUID().uuidString)", isDirectory: false) + try Data("not-a-directory".utf8).write(to: pendingRoot) + defer { + try? fileManager.removeItem(at: cacheRoot) + try? fileManager.removeItem(at: pendingRoot) + } + let timestamp = try #require(CostUsageDateParser.parse("2026-07-16T12:00:00Z")) + let record = CLIProxyAPIUsageRecord( + timestamp: timestamp, + provider: "codex", + model: "gpt-5.5", + alias: "gpt-5.5", + endpoint: "POST /v1/messages", + authType: "oauth", + requestID: "request-1", + tokens: .init(input: 10, output: 20, total: 30)) + let encoder = JSONEncoder() + encoder.dateEncodingStrategy = .iso8601 + let data = try encoder.encode([record]) + let client = CLIProxyAPIUsageQueueClient( + settings: .init(managementKey: "management-secret"), + dataLoader: { request in + let url = try #require(request.url) + let response = try #require(HTTPURLResponse( + url: url, + statusCode: 200, + httpVersion: nil, + headerFields: nil)) + return (data, response) + }) + + let result = await CLIProxyAPIUsageCollector.collect( + cacheRoot: cacheRoot, + pendingRoot: pendingRoot, + client: client) + + #expect(result == .failed("Could not stage CLIProxyAPI usage telemetry.")) + #expect(CLIProxyAPIUsageCacheIO.load(cacheRoot: cacheRoot).isEmpty) + #expect(!fileManager.fileExists( + atPath: CLIProxyAPIUsageCacheIO.cacheFileURL(cacheRoot: cacheRoot).path)) + } + + @Test + func `collector rechecks opt out before every destructive pop`() async throws { + let fileManager = FileManager.default + let cacheRoot = fileManager.temporaryDirectory + .appendingPathComponent("cliproxy-opt-out-\(UUID().uuidString)", isDirectory: true) + defer { try? fileManager.removeItem(at: cacheRoot) } + let timestamp = try #require(CostUsageDateParser.parse("2026-07-16T12:00:00Z")) + let records = (0..<100).map { index in + CLIProxyAPIUsageRecord( + timestamp: timestamp.addingTimeInterval(TimeInterval(index)), + provider: "codex", + model: "gpt-5.5", + alias: "gpt-5.5", + endpoint: "POST /v1/messages", + authType: "oauth", + requestID: "request-\(index)", + tokens: .init(input: 10, output: 20, total: 30)) + } + let encoder = JSONEncoder() + encoder.dateEncodingStrategy = .iso8601 + let data = try encoder.encode(records) + let probe = CLIProxyAPICollectionContinuationProbe() + let client = CLIProxyAPIUsageQueueClient( + settings: .init(managementKey: "management-secret"), + dataLoader: { request in + await probe.recordPop() + let url = try #require(request.url) + let response = try #require(HTTPURLResponse( + url: url, + statusCode: 200, + httpVersion: nil, + headerFields: nil)) + return (data, response) + }) + + let result = await CLIProxyAPIUsageCollector.collect( + cacheRoot: cacheRoot, + shouldContinue: { + await probe.shouldContinue() + }, + client: client) + + #expect(result == .disabled) + #expect(await probe.popCount == 1) + #expect(CLIProxyAPIUsageCacheIO.load(cacheRoot: cacheRoot).count == 100) + } + + @Test + func `collector rechecks configuration after acquiring the interprocess lock`() async throws { + let root = FileManager.default.temporaryDirectory + .appendingPathComponent("cliproxy-queued-disconnect-\(UUID().uuidString)", isDirectory: true) + let cacheRoot = root.appendingPathComponent("cost-usage", isDirectory: true) + defer { try? FileManager.default.removeItem(at: root) } + let lockAcquired = DispatchSemaphore(value: 0) + let releaseLock = DispatchSemaphore(value: 0) + let configurationChecked = DispatchSemaphore(value: 0) + let popProbe = CLIProxyAPICollectionContinuationProbe() + let lockHolder = Task.detached { + try await CostUsageCacheLocations.withCLIProxyAPIInterprocessLock(stateRoot: root) { + lockAcquired.signal() + _ = await Self.waitForSignal(releaseLock, timeout: .distantFuture) + } + } + #expect(await Self.waitForSignal(lockAcquired, timeout: .now() + 1)) + let client = CLIProxyAPIUsageQueueClient( + settings: .init(managementKey: "management-secret"), + dataLoader: { request in + await popProbe.recordPop() + let url = try #require(request.url) + let response = try #require(HTTPURLResponse( + url: url, + statusCode: 200, + httpVersion: nil, + headerFields: nil)) + return (Data("[]".utf8), response) + }) + let collection = Task.detached { + await CLIProxyAPIUsageCollector.collect( + cacheRoot: cacheRoot, + configurationIsCurrent: { + configurationChecked.signal() + return !CostUsageCacheLocations.isCLIProxyAPIExplicitlyDisconnected(stateRoot: root) + }, + client: client) + } + #expect(await Self.waitForSignal(configurationChecked, timeout: .now() + 1)) + #expect(CostUsageCacheLocations.setCLIProxyAPIExplicitlyDisconnected(true, stateRoot: root)) + releaseLock.signal() + + try await lockHolder.value + #expect(await collection.value == .notConfigured) + #expect(await popProbe.popCount == 0) + } + + @Test + func `temporary credential unavailability remains retryable`() async { + let result = await CLIProxyAPIUsageCollector.collect( + settingsResult: .temporarilyUnavailable) + + #expect(result == .failed("CLIProxyAPI configuration is temporarily unavailable.")) + } + + @Test + func `collector rechecks configuration before every destructive pop`() async throws { + let fileManager = FileManager.default + let cacheRoot = fileManager.temporaryDirectory + .appendingPathComponent("cliproxy-mid-collection-disconnect-\(UUID().uuidString)", isDirectory: true) + defer { try? fileManager.removeItem(at: cacheRoot) } + let timestamp = try #require(CostUsageDateParser.parse("2026-07-16T12:00:00Z")) + let records = (0..<100).map { index in + CLIProxyAPIUsageRecord( + timestamp: timestamp.addingTimeInterval(TimeInterval(index)), + provider: "codex", + model: "gpt-5.5", + alias: "gpt-5.5", + endpoint: "POST /v1/messages", + authType: "oauth", + requestID: "request-\(index)", + tokens: .init(input: 10, output: 20, total: 30)) + } + let encoder = JSONEncoder() + encoder.dateEncodingStrategy = .iso8601 + let data = try encoder.encode(records) + let configurationIsCurrent = LockIsolated(true) + let popCount = LockIsolated(0) + let client = CLIProxyAPIUsageQueueClient( + settings: .init(managementKey: "management-secret"), + dataLoader: { request in + popCount.setValue(popCount.value + 1) + configurationIsCurrent.setValue(false) + let url = try #require(request.url) + let response = try #require(HTTPURLResponse( + url: url, + statusCode: 200, + httpVersion: nil, + headerFields: nil)) + return (data, response) + }) + + let result = await CLIProxyAPIUsageCollector.collect( + cacheRoot: cacheRoot, + configurationIsCurrent: { configurationIsCurrent.value }, + client: client) + + #expect(result == .notConfigured) + #expect(popCount.value == 1) + #expect(CLIProxyAPIUsageCacheIO.load(cacheRoot: cacheRoot).count == 100) + } + + @Test + func `collector stages an in flight destructive pop before honoring cancellation`() async throws { + let fileManager = FileManager.default + let cacheRoot = fileManager.temporaryDirectory + .appendingPathComponent("cliproxy-cancelled-pop-\(UUID().uuidString)", isDirectory: true) + defer { try? fileManager.removeItem(at: cacheRoot) } + let timestamp = try #require(CostUsageDateParser.parse("2026-07-16T12:00:00Z")) + let record = CLIProxyAPIUsageRecord( + timestamp: timestamp, + provider: "codex", + model: "gpt-5.5", + alias: "gpt-5.5", + endpoint: "POST /v1/messages", + authType: "oauth", + requestID: "request-in-flight", + tokens: .init(input: 10, output: 20, total: 30)) + let encoder = JSONEncoder() + encoder.dateEncodingStrategy = .iso8601 + let data = try encoder.encode([record]) + let popStarted = DispatchSemaphore(value: 0) + let releasePop = DispatchSemaphore(value: 0) + let client = CLIProxyAPIUsageQueueClient( + settings: .init(managementKey: "management-secret"), + dataLoader: { request in + popStarted.signal() + _ = await Self.waitForSignal(releasePop, timeout: .distantFuture) + try Task.checkCancellation() + let url = try #require(request.url) + let response = try #require(HTTPURLResponse( + url: url, + statusCode: 200, + httpVersion: nil, + headerFields: nil)) + return (data, response) + }) + let collection = Task { + await CLIProxyAPIUsageCollector.collect( + cacheRoot: cacheRoot, + client: client) + } + #expect(await Self.waitForSignal(popStarted, timeout: .now() + 1)) + + collection.cancel() + releasePop.signal() + + #expect(await collection.value == .collected(1)) + #expect(CLIProxyAPIUsageCacheIO.load(cacheRoot: cacheRoot).map(\.requestID) == ["request-in-flight"]) + } + + private static func waitForSignal( + _ semaphore: DispatchSemaphore, + timeout: DispatchTime) async -> Bool + { + await withCheckedContinuation { continuation in + DispatchQueue.global().async { + continuation.resume(returning: semaphore.wait(timeout: timeout) == .success) + } + } + } +} diff --git a/Tests/CodexBarTests/CLIProxyAPIUsageStoreTests.swift b/Tests/CodexBarTests/CLIProxyAPIUsageStoreTests.swift new file mode 100644 index 0000000000..a88ac3cc77 --- /dev/null +++ b/Tests/CodexBarTests/CLIProxyAPIUsageStoreTests.swift @@ -0,0 +1,651 @@ +import CodexBarCore +import Foundation +import Testing +@testable import CodexBar + +private actor CLIProxyAPIUsageCollectionRecorder { + private(set) var callCount = 0 + + func collect() -> CLIProxyAPIUsageCollectionResult { + self.callCount += 1 + return .collected(1) + } +} + +private actor CLIProxyAPIUsageCollectorCancellationRecorder { + private(set) var wasCancelled = false + + func recordCancellation() { + self.wasCancelled = true + } +} + +private final class CLIProxyAPICleanupRetryRecorder: @unchecked Sendable { + private let lock = NSLock() + private var attempts = 0 + + func attempt() -> Bool { + self.lock.lock() + defer { self.lock.unlock() } + self.attempts += 1 + return self.attempts >= 2 + } + + var count: Int { + self.lock.lock() + defer { self.lock.unlock() } + return self.attempts + } +} + +@MainActor +struct CLIProxyAPIUsageStoreTests { + @Test + func `cost tracking opt out prevents telemetry collection`() async { + let settings = testSettingsStore(suiteName: "CLIProxyAPIUsageStoreTests-\(UUID().uuidString)") + settings.costUsageEnabled = false + let root = FileManager.default.temporaryDirectory + .appendingPathComponent("cliproxy-usage-store-\(UUID().uuidString)", isDirectory: true) + defer { try? FileManager.default.removeItem(at: root) } + let environment = [ + "HOME": root.path, + "CODEX_HOME": root.appendingPathComponent(".codex", isDirectory: true).path, + ] + let store = UsageStore( + fetcher: UsageFetcher(environment: environment), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing, + environmentBase: environment) + let recorder = CLIProxyAPIUsageCollectionRecorder() + + let disabledResult = await store.collectCLIProxyAPIUsageNow { + await recorder.collect() + } + + #expect(disabledResult == .disabled) + #expect(await recorder.callCount == 0) + + settings.costUsageEnabled = true + let enabledResult = await store.collectCLIProxyAPIUsageNow { + await recorder.collect() + } + + #expect(enabledResult == .collected(1)) + #expect(await recorder.callCount == 1) + } + + @Test + func `removing the integration cancels and clears the active telemetry task`() async { + let settings = testSettingsStore(suiteName: "CLIProxyAPIUsageStoreTests-\(UUID().uuidString)") + let root = FileManager.default.temporaryDirectory + .appendingPathComponent("cliproxy-usage-store-\(UUID().uuidString)", isDirectory: true) + defer { try? FileManager.default.removeItem(at: root) } + let environment = [ + "HOME": root.path, + "CODEX_HOME": root.appendingPathComponent(".codex", isDirectory: true).path, + ] + let store = UsageStore( + fetcher: UsageFetcher(environment: environment), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing, + environmentBase: environment) + let recorder = CLIProxyAPIUsageCollectorCancellationRecorder() + let collectorFinished = LockIsolated(false) + let task = Task { + while !Task.isCancelled { + await Task.yield() + } + await recorder.recordCancellation() + let drainDelay = Task.detached { + try? await Task.sleep(for: .milliseconds(50)) + } + await drainDelay.value + collectorFinished.setValue(true) + } + store.cliProxyAPIUsageCollectorTask = task + var collectorFinishedBeforePurge = false + store.publishTokenSnapshot(Self.tokenSnapshot(), for: .codex) + store.publishTokenSnapshot(Self.tokenSnapshot(), for: .claude) + let codexPublicationRevision = store.tokenSnapshotPublicationRevision(for: .codex) + let claudePublicationRevision = store.tokenSnapshotPublicationRevision(for: .claude) + let claudePublicationGuard = store.tokenRefreshPublicationGuard(for: .claude) + let claudeScopeSignature = store.tokenSnapshotScopeSignature(for: .claude) + + let removed = await store.removeCLIProxyAPIConfiguration { + collectorFinishedBeforePurge = collectorFinished.value + return .removed + } + await task.value + + #expect(removed == .removed) + #expect(collectorFinishedBeforePurge) + #expect(store.cliProxyAPIUsageCollectorTask == nil) + #expect(store.tokenSnapshot(for: .codex) == nil) + #expect(store.tokenSnapshotPublicationRevision(for: .codex) == codexPublicationRevision + 1) + #expect(store.tokenSnapshot(for: .claude) == nil) + #expect(store.tokenSnapshotPublicationRevision(for: .claude) == claudePublicationRevision + 1) + #expect(!store.tokenRefreshPublicationIsCurrent( + provider: .claude, + publicationGuard: claudePublicationGuard, + historyDays: settings.costUsageHistoryDays, + costScopeSignature: claudeScopeSignature)) + #expect(await recorder.wasCancelled) + } + + @Test + func `removing the integration preserves telemetry when configuration removal fails`() async { + let settings = testSettingsStore(suiteName: "CLIProxyAPIUsageStoreTests-\(UUID().uuidString)") + let root = FileManager.default.temporaryDirectory + .appendingPathComponent("cliproxy-usage-store-\(UUID().uuidString)", isDirectory: true) + defer { try? FileManager.default.removeItem(at: root) } + let environment = [ + "HOME": root.path, + "CODEX_HOME": root.appendingPathComponent(".codex", isDirectory: true).path, + ] + let store = UsageStore( + fetcher: UsageFetcher(environment: environment), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing, + environmentBase: environment) + store.publishTokenSnapshot(Self.tokenSnapshot(), for: .codex) + let removed = await store.removeCLIProxyAPIConfiguration { + .configurationRemovalFailed + } + + #expect(removed == .configurationRemovalFailed) + #expect(store.tokenSnapshot(for: .codex) != nil) + } + + @Test + func `removing the integration reports telemetry cleanup failure after configuration removal`() async { + let settings = testSettingsStore(suiteName: "CLIProxyAPIUsageStoreTests-\(UUID().uuidString)") + let root = FileManager.default.temporaryDirectory + .appendingPathComponent("cliproxy-usage-store-\(UUID().uuidString)", isDirectory: true) + defer { try? FileManager.default.removeItem(at: root) } + let environment = [ + "HOME": root.path, + "CODEX_HOME": root.appendingPathComponent(".codex", isDirectory: true).path, + ] + let store = UsageStore( + fetcher: UsageFetcher(environment: environment), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing, + environmentBase: environment) + store.publishTokenSnapshot(Self.tokenSnapshot(), for: .codex) + var scheduledCleanupRetry = false + + let removed = await store.removeCLIProxyAPIConfiguration( + remove: { .telemetryCleanupFailed }, + scheduleCleanupRetry: { scheduledCleanupRetry = true }) + + #expect(removed == .telemetryCleanupFailed) + #expect(store.tokenSnapshot(for: .codex) == nil) + #expect(scheduledCleanupRetry) + } + + @Test + func `telemetry cleanup maintenance retries until transaction recovery succeeds`() async { + let settings = testSettingsStore(suiteName: "CLIProxyAPIUsageStoreTests-\(UUID().uuidString)") + let store = UsageStore( + fetcher: UsageFetcher(environment: [:]), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing, + environmentBase: [:]) + let recorder = CLIProxyAPICleanupRetryRecorder() + + let task = store.startCLIProxyAPICleanupRetry(retryInterval: .milliseconds(1)) { + recorder.attempt() + } + await task.value + + #expect(recorder.count == 2) + } + + @Test + func `reconnecting invalidates and force refreshes both proxy affected token snapshots`() async { + let settings = testSettingsStore(suiteName: "CLIProxyAPIUsageStoreTests-\(UUID().uuidString)") + let root = FileManager.default.temporaryDirectory + .appendingPathComponent("cliproxy-usage-store-\(UUID().uuidString)", isDirectory: true) + defer { try? FileManager.default.removeItem(at: root) } + let environment = [ + "HOME": root.path, + "CODEX_HOME": root.appendingPathComponent(".codex", isDirectory: true).path, + ] + let store = UsageStore( + fetcher: UsageFetcher(environment: environment), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing, + environmentBase: environment) + store.publishTokenSnapshot(Self.tokenSnapshot(), for: .codex) + store.publishTokenSnapshot(Self.tokenSnapshot(), for: .claude) + let codexPublicationRevision = store.tokenSnapshotPublicationRevision(for: .codex) + let claudePublicationRevision = store.tokenSnapshotPublicationRevision(for: .claude) + var refreshes: [(UsageProvider, Bool)] = [] + + await store.refreshCLIProxyAPICostAttribution { provider, force in + refreshes.append((provider, force)) + } + + #expect(store.tokenSnapshot(for: .codex) == nil) + #expect(store.tokenSnapshotPublicationRevision(for: .codex) == codexPublicationRevision + 1) + #expect(store.tokenSnapshot(for: .claude) == nil) + #expect(store.tokenSnapshotPublicationRevision(for: .claude) == claudePublicationRevision + 1) + #expect(refreshes.map(\.0) == [.claude, .codex]) + #expect(refreshes.map(\.1) == [true, true]) + } + + @Test + func `initial missing configuration preserves hydrated proxy snapshots`() async { + let settings = testSettingsStore(suiteName: "CLIProxyAPIUsageStoreTests-\(UUID().uuidString)") + settings.costUsageEnabled = true + let store = UsageStore( + fetcher: UsageFetcher(environment: [:]), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing, + environmentBase: [:]) + let snapshot = Self.tokenSnapshot() + store.publishTokenSnapshot(snapshot, for: .codex) + store.publishTokenSnapshot(snapshot, for: .claude) + let codexPublicationRevision = store.tokenSnapshotPublicationRevision(for: .codex) + let claudePublicationRevision = store.tokenSnapshotPublicationRevision(for: .claude) + + let collectorState = await store.handleCLIProxyAPIUsageCollectionResult( + .notConfigured, + collectorState: CLIProxyAPIUsageCollectorState(), + isExplicitlyDisconnected: { false }) + + #expect(collectorState.configurationAvailability == .unavailable) + #expect(store.tokenSnapshot(for: .codex) == snapshot) + #expect(store.tokenSnapshot(for: .claude) == snapshot) + #expect(store.tokenSnapshotPublicationRevision(for: .codex) == codexPublicationRevision) + #expect(store.tokenSnapshotPublicationRevision(for: .claude) == claudePublicationRevision) + } + + @Test + func `background disconnect invalidates once and remote reconnect refreshes snapshots`() async { + let settings = testSettingsStore(suiteName: "CLIProxyAPIUsageStoreTests-\(UUID().uuidString)") + settings.costUsageEnabled = true + let root = FileManager.default.temporaryDirectory + .appendingPathComponent("cliproxy-usage-store-\(UUID().uuidString)", isDirectory: true) + defer { try? FileManager.default.removeItem(at: root) } + let environment = [ + "HOME": root.path, + "CODEX_HOME": root.appendingPathComponent(".codex", isDirectory: true).path, + ] + let store = UsageStore( + fetcher: UsageFetcher(environment: environment), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing, + environmentBase: environment) + store.publishTokenSnapshot(Self.tokenSnapshot(), for: .codex) + store.publishTokenSnapshot(Self.tokenSnapshot(), for: .claude) + let codexPublicationRevision = store.tokenSnapshotPublicationRevision(for: .codex) + let claudePublicationRevision = store.tokenSnapshotPublicationRevision(for: .claude) + let dashboardRevision = store.spendDashboardCodexCostCatchUpRevision + let dashboardConfiguration = SpendDashboardSource.configuration(settings: settings, store: store) + + var collectorState = await store.handleCLIProxyAPIUsageCollectionResult( + .notConfigured, + collectorState: CLIProxyAPIUsageCollectorState(configurationAvailability: .available), + isExplicitlyDisconnected: { false }) + + #expect(collectorState.configurationAvailability == .unavailable) + #expect(store.tokenSnapshot(for: .codex) == nil) + #expect(store.tokenSnapshot(for: .claude) == nil) + #expect(store.tokenSnapshotPublicationRevision(for: .codex) == codexPublicationRevision + 1) + #expect(store.tokenSnapshotPublicationRevision(for: .claude) == claudePublicationRevision + 1) + #expect(store.spendDashboardCodexCostCatchUpRevision == dashboardRevision + 1) + #expect( + SpendDashboardSource.configuration(settings: settings, store: store).sourceRevisions != + dashboardConfiguration.sourceRevisions) + + collectorState = await store.handleCLIProxyAPIUsageCollectionResult( + .notConfigured, + collectorState: collectorState, + isExplicitlyDisconnected: { false }) + + #expect(store.tokenSnapshotPublicationRevision(for: .codex) == codexPublicationRevision + 1) + #expect(store.tokenSnapshotPublicationRevision(for: .claude) == claudePublicationRevision + 1) + #expect(store.spendDashboardCodexCostCatchUpRevision == dashboardRevision + 1) + + var refreshes: [(UsageProvider, Bool)] = [] + collectorState = await store.handleCLIProxyAPIUsageCollectionResult( + .failed("temporary failure"), + collectorState: collectorState, + isExplicitlyDisconnected: { false }) + #expect(collectorState.configurationAvailability == .unavailable) + collectorState = await store.handleCLIProxyAPIUsageCollectionResult( + .collected(0), + collectorState: collectorState, + isExplicitlyDisconnected: { false }, + refresh: { provider, force in + refreshes.append((provider, force)) + }) + #expect(collectorState.configurationAvailability == .available) + #expect(refreshes.map(\.0) == [.claude, .codex]) + #expect(refreshes.map(\.1) == [true, true]) + store.publishTokenSnapshot(Self.tokenSnapshot(), for: .codex) + collectorState = await store.handleCLIProxyAPIUsageCollectionResult( + .notConfigured, + collectorState: collectorState, + isExplicitlyDisconnected: { false }) + + #expect(collectorState.configurationAvailability == .unavailable) + #expect(store.tokenSnapshot(for: .codex) == nil) + #expect(store.tokenSnapshotPublicationRevision(for: .codex) == codexPublicationRevision + 4) + } + + @Test + func `configuration generation detects a reconnect missed between polls`() async { + let settings = testSettingsStore(suiteName: "CLIProxyAPIUsageStoreTests-\(UUID().uuidString)") + settings.costUsageEnabled = true + let store = UsageStore( + fetcher: UsageFetcher(environment: [:]), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing, + environmentBase: [:]) + let snapshot = Self.tokenSnapshot() + store.publishTokenSnapshot(snapshot, for: .codex) + store.publishTokenSnapshot(snapshot, for: .claude) + var refreshes: [(UsageProvider, Bool)] = [] + + let collectorState = await store.handleCLIProxyAPIUsageCollectionResult( + .collected(0), + collectorState: CLIProxyAPIUsageCollectorState( + configurationAvailability: .available, + configurationGeneration: "before-removal"), + isExplicitlyDisconnected: { false }, + configurationGeneration: { "after-reconnect" }, + refresh: { provider, force in + refreshes.append((provider, force)) + }) + + #expect(collectorState.configurationAvailability == .available) + #expect(collectorState.configurationGeneration == "after-reconnect") + #expect(store.tokenSnapshot(for: .codex) == nil) + #expect(store.tokenSnapshot(for: .claude) == nil) + #expect(refreshes.map(\.0) == [.claude, .codex]) + #expect(refreshes.map(\.1) == [true, true]) + } + + @Test + func `failed generation transition stays pending until collection succeeds`() async { + let settings = testSettingsStore(suiteName: "CLIProxyAPIUsageStoreTests-\(UUID().uuidString)") + settings.costUsageEnabled = true + let store = UsageStore( + fetcher: UsageFetcher(environment: [:]), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing, + environmentBase: [:]) + var refreshes: [(UsageProvider, Bool)] = [] + + var collectorState = await store.handleCLIProxyAPIUsageCollectionResult( + .failed("replacement unavailable"), + collectorState: CLIProxyAPIUsageCollectorState( + configurationAvailability: .available, + configurationGeneration: "old-generation"), + configurationGeneration: { "new-generation" }) + + #expect(collectorState.configurationAvailability == .unavailable) + #expect(collectorState.configurationGeneration == "new-generation") + + collectorState = await store.handleCLIProxyAPIUsageCollectionResult( + .collected(0), + collectorState: collectorState, + configurationGeneration: { "new-generation" }, + refresh: { provider, force in + refreshes.append((provider, force)) + }) + + #expect(collectorState.configurationAvailability == .available) + #expect(refreshes.map(\.0) == [.claude, .codex]) + #expect(refreshes.map(\.1) == [true, true]) + } + + @Test + func `first collection detects a generation change after startup hydration`() async { + let settings = testSettingsStore(suiteName: "CLIProxyAPIUsageStoreTests-\(UUID().uuidString)") + settings.costUsageEnabled = true + let store = UsageStore( + fetcher: UsageFetcher(environment: [:]), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing, + environmentBase: [:]) + let snapshot = Self.tokenSnapshot() + store.publishTokenSnapshot(snapshot, for: .codex) + store.publishTokenSnapshot(snapshot, for: .claude) + var refreshes: [(UsageProvider, Bool)] = [] + + let collectorState = await store.handleCLIProxyAPIUsageCollectionResult( + .collected(0), + collectorState: CLIProxyAPIUsageCollectorState( + configurationGeneration: "hydrated-generation"), + configurationGeneration: { "replacement-generation" }, + refresh: { provider, force in + refreshes.append((provider, force)) + }) + + #expect(collectorState.configurationAvailability == .available) + #expect(collectorState.configurationGeneration == "replacement-generation") + #expect(store.tokenSnapshot(for: .codex) == nil) + #expect(store.tokenSnapshot(for: .claude) == nil) + #expect(refreshes.map(\.0) == [.claude, .codex]) + #expect(refreshes.map(\.1) == [true, true]) + } + + @Test + func `first failed collection invalidates stale snapshots after configuration changes`() async { + let settings = testSettingsStore(suiteName: "CLIProxyAPIUsageStoreTests-\(UUID().uuidString)") + settings.costUsageEnabled = true + let store = UsageStore( + fetcher: UsageFetcher(environment: [:]), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing, + environmentBase: [:]) + let snapshot = Self.tokenSnapshot() + store.publishTokenSnapshot(snapshot, for: .codex) + store.publishTokenSnapshot(snapshot, for: .claude) + + var collectorState = await store.handleCLIProxyAPIUsageCollectionResult( + .failed("replacement endpoint unavailable"), + collectorState: CLIProxyAPIUsageCollectorState( + configurationGeneration: "before-replacement"), + configurationGeneration: { "after-replacement" }) + + #expect(collectorState.configurationAvailability == .unavailable) + #expect(collectorState.configurationGeneration == "after-replacement") + #expect(store.tokenSnapshot(for: .codex) == nil) + #expect(store.tokenSnapshot(for: .claude) == nil) + + store.publishTokenSnapshot(snapshot, for: .codex) + collectorState = await store.handleCLIProxyAPIUsageCollectionResult( + .failed("still unavailable"), + collectorState: collectorState, + configurationGeneration: { "after-replacement" }) + + #expect(store.tokenSnapshot(for: .codex) == snapshot) + } + + @Test + func `clearing cost cache drains the active proxy collector before deletion`() async { + let settings = testSettingsStore(suiteName: "CLIProxyAPIUsageStoreTests-\(UUID().uuidString)") + let root = FileManager.default.temporaryDirectory + .appendingPathComponent("cliproxy-usage-store-\(UUID().uuidString)", isDirectory: true) + defer { try? FileManager.default.removeItem(at: root) } + let environment = [ + "HOME": root.path, + "CODEX_HOME": root.appendingPathComponent(".codex", isDirectory: true).path, + ] + let store = UsageStore( + fetcher: UsageFetcher(environment: environment), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing, + environmentBase: environment) + let collectorFinished = LockIsolated(false) + let deletionStartedAfterDrain = LockIsolated(false) + store.cliProxyAPIUsageCollectorTask = Task { + while !Task.isCancelled { + await Task.yield() + } + let drainDelay = Task.detached { + try? await Task.sleep(for: .milliseconds(50)) + } + await drainDelay.value + collectorFinished.setValue(true) + } + + let error = await store.clearCostUsageCache(clearDirectories: { + deletionStartedAfterDrain.setValue(collectorFinished.value) + return (cleared: 0, errorMessage: nil) + }) + store.stopCLIProxyAPIUsageCollector() + + #expect(error == nil) + #expect(deletionStartedAfterDrain.value) + } + + @Test + func `clearing cost cache cancels and drains token scans before deletion`() async { + let settings = testSettingsStore(suiteName: "CLIProxyAPIUsageStoreTests-\(UUID().uuidString)") + settings.costUsageEnabled = true + let root = FileManager.default.temporaryDirectory + .appendingPathComponent("cliproxy-usage-store-\(UUID().uuidString)", isDirectory: true) + defer { try? FileManager.default.removeItem(at: root) } + let environment = [ + "HOME": root.path, + "CODEX_HOME": root.appendingPathComponent(".codex", isDirectory: true).path, + ] + let store = UsageStore( + fetcher: UsageFetcher(environment: environment), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing, + environmentBase: environment) + let refreshStarted = LockIsolated(false) + let refreshFinished = LockIsolated(false) + let deletionStartedAfterDrain = LockIsolated(false) + store.publishTokenSnapshot(Self.tokenSnapshot(), for: .codex) + let publicationRevision = store.tokenSnapshotPublicationRevision(for: .codex) + store._test_tokenUsageRefreshOverride = { _, _ in + refreshStarted.setValue(true) + while !Task.isCancelled { + await Task.yield() + } + refreshFinished.setValue(true) + } + defer { store._test_tokenUsageRefreshOverride = nil } + let refreshTask = Task { + await store.refreshTokenUsageNow(for: .codex, force: true) + } + while !refreshStarted.value { + await Task.yield() + } + + let error = await store.clearCostUsageCache(clearDirectories: { + deletionStartedAfterDrain.setValue(refreshFinished.value) + return (cleared: 0, errorMessage: nil) + }) + await refreshTask.value + + #expect(error == nil) + #expect(deletionStartedAfterDrain.value) + #expect(store.tokenSnapshot(for: .codex) == nil) + #expect(store.tokenSnapshotPublicationRevision(for: .codex) == publicationRevision + 1) + #expect(store.tokenRefreshInFlight.isEmpty) + #expect(store.tokenRefreshSequenceTask == nil) + } + + @Test + func `clearing cost cache uses the shared locked deletion path`() async throws { + let settings = testSettingsStore(suiteName: "CLIProxyAPIUsageStoreTests-\(UUID().uuidString)") + let root = FileManager.default.temporaryDirectory + .appendingPathComponent("cliproxy-usage-store-\(UUID().uuidString)", isDirectory: true) + let fileManager = CLIProxyAPITestFileManager(root: root) + let cacheDirectory = CostUsageCacheLocations.directories(fileManager: fileManager)[0] + try FileManager.default.createDirectory(at: cacheDirectory, withIntermediateDirectories: true) + let usageFile = cacheDirectory.appendingPathComponent("usage.json") + try Data("telemetry".utf8).write(to: usageFile) + defer { try? FileManager.default.removeItem(at: root) } + + let environment = [ + "HOME": root.path, + "CODEX_HOME": root.appendingPathComponent(".codex", isDirectory: true).path, + ] + let store = UsageStore( + fetcher: UsageFetcher(environment: environment), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing, + environmentBase: environment) + let error = await store.clearCostUsageCache(fileManager: fileManager) + + #expect(error == nil) + #expect(!FileManager.default.fileExists(atPath: cacheDirectory.path)) + } + + @Test + func `partial cost cache clear invalidates in memory snapshots`() async { + let settings = testSettingsStore(suiteName: "CLIProxyAPIUsageStoreTests-\(UUID().uuidString)") + let store = UsageStore( + fetcher: UsageFetcher(), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing, + environmentBase: [:]) + store.publishTokenSnapshot(Self.tokenSnapshot(), for: .codex) + + let error = await store.clearCostUsageCache(clearDirectories: { + (cleared: 1, errorMessage: "Could not remove every cache directory") + }) + + #expect(error != nil) + #expect(store.tokenSnapshot(for: .codex) == nil) + } + + private static func tokenSnapshot() -> CostUsageTokenSnapshot { + CostUsageTokenSnapshot( + sessionTokens: 10, + sessionCostUSD: 0.01, + last30DaysTokens: 10, + last30DaysCostUSD: 0.01, + currencyCode: "USD", + daily: [], + updatedAt: Date(timeIntervalSince1970: 1_784_203_200)) + } +} + +private final class CLIProxyAPITestFileManager: FileManager { + private let root: URL + + init(root: URL) { + self.root = root + super.init() + } + + override func urls( + for directory: FileManager.SearchPathDirectory, + in _: FileManager.SearchPathDomainMask) -> [URL] + { + switch directory { + case .cachesDirectory: + [self.root.appendingPathComponent("Caches", isDirectory: true)] + case .applicationSupportDirectory: + [self.root.appendingPathComponent("Application Support", isDirectory: true)] + default: + super.urls(for: directory, in: .userDomainMask) + } + } +} diff --git a/Tests/CodexBarTests/CostHistoryChartMenuViewTests.swift b/Tests/CodexBarTests/CostHistoryChartMenuViewTests.swift index 53f6d8d821..c4cb174caf 100644 --- a/Tests/CodexBarTests/CostHistoryChartMenuViewTests.swift +++ b/Tests/CodexBarTests/CostHistoryChartMenuViewTests.swift @@ -874,6 +874,46 @@ struct CostHistoryChartMenuViewTests { } extension CostHistoryChartMenuViewTests { + @Test + func `model breakdown uses complete attribution as its final ordering key`() { + let apiKeyAttribution = CostUsageAttribution( + client: .claudeCode, + route: .cliProxyAPI, + modelProvider: .openAI, + upstream: .init( + provider: "codex", + authType: .apiKey, + model: "gpt-5.4", + executorType: "codex"), + evidence: [.cliProxyRequestLog]) + let oauthAttribution = CostUsageAttribution( + client: .claudeCode, + route: .cliProxyAPI, + modelProvider: .openAI, + upstream: .init( + provider: "codex", + authType: .oauth, + model: "gpt-5.4", + executorType: "codex"), + evidence: [.cliProxyRequestLog]) + let breakdowns = [ + CostUsageDailyReport.ModelBreakdown( + modelName: "gpt-5.4", + costUSD: 0.1, + totalTokens: 10, + attribution: oauthAttribution), + CostUsageDailyReport.ModelBreakdown( + modelName: "gpt-5.4", + costUSD: 0.1, + totalTokens: 10, + attribution: apiKeyAttribution), + ] + + let sorted = CostHistoryChartMenuView.orderedBreakdownItems(breakdowns) + + #expect(sorted.map(\.attribution) == [apiKeyAttribution, oauthAttribution]) + } + @Test func `session labels distinguish concurrent uuid v7 identifiers`() { let first = CostHistoryChartMenuView.shortSessionID("019f6d91-970b-7e13-b08e-000000000001") diff --git a/Tests/CodexBarTests/CostUsageDailyReportMergeTests.swift b/Tests/CodexBarTests/CostUsageDailyReportMergeTests.swift index e7a42293c9..5ab569a8c9 100644 --- a/Tests/CodexBarTests/CostUsageDailyReportMergeTests.swift +++ b/Tests/CodexBarTests/CostUsageDailyReportMergeTests.swift @@ -3,6 +3,56 @@ import Testing @testable import CodexBarCore struct CostUsageDailyReportMergeTests { + @Test + func `merged report keeps native and claude code proxy model rows distinct`() throws { + let native = CostUsageDailyReport( + data: [ + CostUsageDailyReport.Entry( + date: "2026-07-24", + inputTokens: 100, + outputTokens: 10, + totalTokens: 110, + costUSD: 1, + modelsUsed: ["gpt-5.5"], + modelBreakdowns: [ + CostUsageDailyReport.ModelBreakdown( + modelName: "gpt-5.5", + costUSD: 1, + totalTokens: 110), + ]), + ], + summary: nil) + let proxyAttribution = CostUsageAttribution( + client: .claudeCode, + route: .cliProxyAPI, + modelProvider: .openAI, + upstream: .init(provider: "codex", authType: .oauth), + evidence: [.cliProxyRequestLog, .cliProxyUsageTelemetry, .modelProvider]) + let proxy = CostUsageDailyReport( + data: [ + CostUsageDailyReport.Entry( + date: "2026-07-24", + inputTokens: 50, + outputTokens: 5, + totalTokens: 55, + costUSD: 0.5, + modelsUsed: ["gpt-5.5"], + modelBreakdowns: [ + CostUsageDailyReport.ModelBreakdown( + modelName: "gpt-5.5", + costUSD: 0.5, + totalTokens: 55, + attribution: proxyAttribution), + ]), + ], + summary: nil) + + let breakdowns = try #require(native.merged(with: proxy).data.first?.modelBreakdowns) + #expect(breakdowns.count == 2) + #expect(breakdowns.first { $0.attribution == nil }?.totalTokens == 110) + #expect(breakdowns.first { $0.attribution == proxyAttribution }?.totalTokens == 55) + } + @Test func `merged report sums overlapping day totals and model breakdowns`() { let native = CostUsageDailyReport( @@ -129,6 +179,236 @@ struct CostUsageDailyReportMergeTests { #expect(abs((merged.summary?.totalCostUSD ?? 0) - 0.70) < 0.000001) } + @Test + func `merged report uses complete attribution to order otherwise equal model breakdowns`() throws { + let apiKeyAttribution = CostUsageAttribution( + client: .claudeCode, + route: .cliProxyAPI, + modelProvider: .openAI, + upstream: .init( + provider: "codex", + authType: .apiKey, + model: "gpt-5.4", + executorType: "codex"), + evidence: [.cliProxyRequestLog]) + let oauthAttribution = CostUsageAttribution( + client: .claudeCode, + route: .cliProxyAPI, + modelProvider: .openAI, + upstream: .init( + provider: "codex", + authType: .oauth, + model: "gpt-5.4", + executorType: "codex"), + evidence: [.cliProxyRequestLog]) + let report = CostUsageDailyReport( + data: [ + CostUsageDailyReport.Entry( + date: "2026-07-24", + inputTokens: nil, + outputTokens: nil, + totalTokens: 20, + costUSD: 0.2, + modelsUsed: ["gpt-5.4"], + modelBreakdowns: [ + CostUsageDailyReport.ModelBreakdown( + modelName: "gpt-5.4", + costUSD: 0.1, + totalTokens: 10, + attribution: apiKeyAttribution), + CostUsageDailyReport.ModelBreakdown( + modelName: "gpt-5.4", + costUSD: 0.1, + totalTokens: 10, + attribution: oauthAttribution), + ]), + ], + summary: nil) + + let breakdowns = try #require(CostUsageDailyReport.merged([report]).data.first?.modelBreakdowns) + #expect(breakdowns.map(\.attribution) == [apiKeyAttribution, oauthAttribution]) + } + + @Test + func `vendored cache sorter uses complete attribution for equal model breakdowns`() { + let apiKeyAttribution = CostUsageAttribution( + client: .claudeCode, + route: .cliProxyAPI, + modelProvider: .openAI, + upstream: .init( + provider: "codex", + authType: .apiKey, + model: "gpt-5.4", + executorType: "codex"), + evidence: [.cliProxyRequestLog]) + let oauthAttribution = CostUsageAttribution( + client: .claudeCode, + route: .cliProxyAPI, + modelProvider: .openAI, + upstream: .init( + provider: "codex", + authType: .oauth, + model: "gpt-5.4", + executorType: "codex"), + evidence: [.cliProxyRequestLog]) + let breakdowns = [ + CostUsageDailyReport.ModelBreakdown( + modelName: "gpt-5.4", + costUSD: 0.1, + totalTokens: 10, + attribution: oauthAttribution), + CostUsageDailyReport.ModelBreakdown( + modelName: "gpt-5.4", + costUSD: 0.1, + totalTokens: 10, + attribution: apiKeyAttribution), + ] + + let sorted = CostUsageScanner.sortedModelBreakdowns(breakdowns) + + #expect(sorted.map(\.attribution) == [apiKeyAttribution, oauthAttribution]) + } + + @Test + func `project breakdown sorter uses complete attribution for equal model breakdowns`() throws { + let apiKeyAttribution = CostUsageAttribution( + client: .claudeCode, + route: .cliProxyAPI, + modelProvider: .openAI, + upstream: .init( + provider: "codex", + authType: .apiKey, + model: "gpt-5.4", + executorType: "codex"), + evidence: [.cliProxyRequestLog]) + let oauthAttribution = CostUsageAttribution( + client: .claudeCode, + route: .cliProxyAPI, + modelProvider: .openAI, + upstream: .init( + provider: "codex", + authType: .oauth, + model: "gpt-5.4", + executorType: "codex"), + evidence: [.cliProxyRequestLog]) + let entry = CostUsageDailyReport.Entry( + date: "2026-07-16", + inputTokens: 20, + outputTokens: 0, + totalTokens: 20, + costUSD: 0.2, + modelsUsed: ["gpt-5.4"], + modelBreakdowns: [ + .init( + modelName: "gpt-5.4", + costUSD: 0.1, + totalTokens: 10, + attribution: oauthAttribution), + .init( + modelName: "gpt-5.4", + costUSD: 0.1, + totalTokens: 10, + attribution: apiKeyAttribution), + ]) + + let sorted = try #require(CostUsageFetcher.projectModelBreakdowns(from: [entry])) + + #expect(sorted.map(\.attribution) == [apiKeyAttribution, oauthAttribution]) + } + + @Test + func `attribution sort key includes every distinguishable field`() { + let attributions = [ + CostUsageAttribution( + client: .claudeCode, + route: .cliProxyAPI, + modelProvider: .openAI, + upstream: .init( + provider: "codex", + authType: .oauth, + model: "gpt-5.4", + executorType: "codex"), + evidence: [.cliProxyRequestLog, .modelProvider]), + CostUsageAttribution( + client: .claudeCode, + route: .unknown, + modelProvider: .openAI, + upstream: .init( + provider: "codex", + authType: .oauth, + model: "gpt-5.4", + executorType: "codex"), + evidence: [.cliProxyRequestLog, .modelProvider]), + CostUsageAttribution( + client: .claudeCode, + route: .cliProxyAPI, + modelProvider: .anthropic, + upstream: .init( + provider: "codex", + authType: .oauth, + model: "gpt-5.4", + executorType: "codex"), + evidence: [.cliProxyRequestLog, .modelProvider]), + CostUsageAttribution( + client: .claudeCode, + route: .cliProxyAPI, + modelProvider: .openAI, + evidence: [.cliProxyRequestLog, .modelProvider]), + CostUsageAttribution( + client: .claudeCode, + route: .cliProxyAPI, + modelProvider: .openAI, + upstream: .init( + provider: "openrouter", + authType: .oauth, + model: "gpt-5.4", + executorType: "codex"), + evidence: [.cliProxyRequestLog, .modelProvider]), + CostUsageAttribution( + client: .claudeCode, + route: .cliProxyAPI, + modelProvider: .openAI, + upstream: .init( + provider: "codex", + authType: .apiKey, + model: "gpt-5.4", + executorType: "codex"), + evidence: [.cliProxyRequestLog, .modelProvider]), + CostUsageAttribution( + client: .claudeCode, + route: .cliProxyAPI, + modelProvider: .openAI, + upstream: .init( + provider: "codex", + authType: .oauth, + model: "gpt-5.3", + executorType: "codex"), + evidence: [.cliProxyRequestLog, .modelProvider]), + CostUsageAttribution( + client: .claudeCode, + route: .cliProxyAPI, + modelProvider: .openAI, + upstream: .init( + provider: "codex", + authType: .oauth, + model: "gpt-5.4", + executorType: "openai"), + evidence: [.cliProxyRequestLog, .modelProvider]), + CostUsageAttribution( + client: .claudeCode, + route: .cliProxyAPI, + modelProvider: .openAI, + upstream: .init( + provider: "codex", + authType: .oauth, + model: "gpt-5.4", + executorType: "codex"), + evidence: [.modelProvider, .cliProxyRequestLog]), + ] + + #expect(Set(attributions.map(\.deterministicSortKey)).count == attributions.count) + } + @Test func `merged report includes derived totals when another same day entry has explicit total`() { let explicit = CostUsageDailyReport( diff --git a/Tests/CodexBarTests/CostUsageFetcherCLIProxyAttributionReconciliationTests.swift b/Tests/CodexBarTests/CostUsageFetcherCLIProxyAttributionReconciliationTests.swift new file mode 100644 index 0000000000..aeb511b67e --- /dev/null +++ b/Tests/CodexBarTests/CostUsageFetcherCLIProxyAttributionReconciliationTests.swift @@ -0,0 +1,85 @@ +import Foundation +import Testing +@testable import CodexBarCore + +@Suite(.serialized) +struct CostUsageFetcherCLIProxyAttributionReconciliationTests { + @Test + func `batch reconciliation includes unkeyed legacy Claude rows`() async throws { + let env = try CostUsageTestEnvironment() + defer { env.cleanup() } + + let day = try env.makeLocalNoon(year: 2026, month: 7, day: 24) + func assistant(seconds: TimeInterval) -> [String: Any] { + [ + "type": "assistant", + "timestamp": env.isoString(for: day.addingTimeInterval(seconds)), + "sessionId": "legacy-session", + "message": [ + "model": "gpt-5.5", + "usage": ["input_tokens": 100, "output_tokens": 5], + ], + ] + } + _ = try env.writeClaudeProjectFile( + relativePath: "legacy-proxy/session.jsonl", + contents: env.jsonl([ + assistant(seconds: 0), + assistant(seconds: 30), + ])) + + let cliProxyHome = env.root.appendingPathComponent("cli-proxy-api", isDirectory: true) + let cliProxyLogs = cliProxyHome.appendingPathComponent("logs", isDirectory: true) + try FileManager.default.createDirectory(at: cliProxyLogs, withIntermediateDirectories: true) + try Data(#"{"type":"codex"}"#.utf8) + .write(to: cliProxyHome.appendingPathComponent("codex-auth.json")) + let proxyLog = """ + === REQUEST INFO === + URL: /v1/messages + Timestamp: \(env.isoString(for: day)) + === HEADERS === + X-Claude-Code-Session-Id: legacy-session + === REQUEST BODY === + {"model":"gpt-5.5"} + === API RESPONSE === + """ + try Data(proxyLog.utf8).write(to: cliProxyLogs.appendingPathComponent("request.log")) + #expect(CLIProxyAPIUsageCacheIO.merge( + [ + CLIProxyAPIUsageRecord( + timestamp: day, + provider: "codex", + executorType: "CodexExecutor", + model: "gpt-5.5", + alias: "gpt-5.5", + endpoint: "/v1/messages", + authType: "oauth", + requestID: "legacy-proxy-request", + tokens: .init(input: 100, output: 5, total: 105)), + ], + cacheRoot: env.cacheRoot, + now: day) == 1) + let options = CostUsageScanner.Options( + codexSessionsRoot: env.codexSessionsRoot, + claudeProjectsRoots: [env.claudeProjectsRoot], + cacheRoot: env.cacheRoot, + cliProxyAPIHome: cliProxyHome) + + let codex = try await CostUsageFetcher.loadTokenSnapshot( + provider: .codex, + now: day, + allowPricingRefresh: false, + includePiSessions: false, + scannerOptions: options) + let claude = try await CostUsageFetcher.loadTokenSnapshot( + provider: .claude, + now: day, + allowPricingRefresh: false, + includePiSessions: false, + scannerOptions: options) + + #expect(codex.daily.first?.totalTokens == 105) + #expect(codex.daily.first?.modelBreakdowns?.first?.attribution?.route == .cliProxyAPI) + #expect(claude.daily.isEmpty) + } +} diff --git a/Tests/CodexBarTests/CostUsageFetcherCLIProxyConcurrencyTests.swift b/Tests/CodexBarTests/CostUsageFetcherCLIProxyConcurrencyTests.swift new file mode 100644 index 0000000000..e8ba386f7e --- /dev/null +++ b/Tests/CodexBarTests/CostUsageFetcherCLIProxyConcurrencyTests.swift @@ -0,0 +1,118 @@ +import Foundation +import Testing +@testable import CodexBarCore + +struct CostUsageFetcherCLIProxyConcurrencyTests { + @Test + func `concurrent proxy requests retain distinct token matched upstreams`() async throws { + let env = try CostUsageTestEnvironment() + defer { env.cleanup() } + + let day = try env.makeLocalNoon(year: 2026, month: 7, day: 24) + func assistant( + sessionID: String, + requestID: String, + seconds: TimeInterval, + input: Int, + output: Int) -> [String: Any] + { + [ + "type": "assistant", + "timestamp": env.isoString(for: day.addingTimeInterval(seconds)), + "sessionId": sessionID, + "requestId": requestID, + "message": [ + "id": "message-\(requestID)", + "model": "gpt-5.5", + "usage": ["input_tokens": input, "output_tokens": output], + ], + ] + } + _ = try env.writeClaudeProjectFile( + relativePath: "concurrent-proxy/session.jsonl", + contents: env.jsonl([ + assistant( + sessionID: "session-codex", + requestID: "codex", + seconds: 0, + input: 10, + output: 2), + assistant( + sessionID: "session-openrouter", + requestID: "openrouter", + seconds: 2, + input: 100, + output: 20), + ])) + + let cliProxyHome = env.root.appendingPathComponent("cli-proxy-api", isDirectory: true) + let logs = cliProxyHome.appendingPathComponent("logs", isDirectory: true) + try FileManager.default.createDirectory(at: logs, withIntermediateDirectories: true) + for (name, sessionID, seconds) in [ + ("codex", "session-codex", 0.0), + ("openrouter", "session-openrouter", 2.0), + ] { + let log = """ + === REQUEST INFO === + URL: /v1/messages + Timestamp: \(env.isoString(for: day.addingTimeInterval(seconds))) + === HEADERS === + X-Claude-Code-Session-Id: \(sessionID) + === REQUEST BODY === + {"model":"gpt-5.5"} + === API RESPONSE === + """ + try Data(log.utf8).write(to: logs.appendingPathComponent("\(name).log")) + } + CLIProxyAPIUsageCacheIO.merge( + [ + CLIProxyAPIUsageRecord( + timestamp: day, + provider: "codex", + executorType: "CodexExecutor", + model: "gpt-5.5", + alias: "gpt-5.5", + endpoint: "/v1/messages", + authType: "oauth", + requestID: "cliproxy-codex", + tokens: .init(input: 10, output: 2, total: 12)), + CLIProxyAPIUsageRecord( + timestamp: day.addingTimeInterval(2), + provider: "openrouter", + executorType: "OpenAICompatExecutor", + model: "gpt-5.5", + alias: "gpt-5.5", + endpoint: "/v1/messages", + authType: "api_key", + requestID: "cliproxy-openrouter", + tokens: .init(input: 100, output: 20, total: 120)), + ], + cacheRoot: env.cacheRoot, + now: day.addingTimeInterval(2)) + let options = CostUsageScanner.Options( + codexSessionsRoot: env.codexSessionsRoot, + claudeProjectsRoots: [env.claudeProjectsRoot], + cacheRoot: env.cacheRoot, + cliProxyAPIHome: cliProxyHome) + + let codex = try await CostUsageFetcher.loadTokenSnapshot( + provider: .codex, + now: day, + allowPricingRefresh: false, + includePiSessions: false, + scannerOptions: options) + let claude = try await CostUsageFetcher.loadTokenSnapshot( + provider: .claude, + now: day, + allowPricingRefresh: false, + includePiSessions: false, + scannerOptions: options) + + let codexBreakdown = try #require(codex.daily.first?.modelBreakdowns?.first) + let claudeBreakdown = try #require(claude.daily.first?.modelBreakdowns?.first) + #expect(codex.daily.first?.totalTokens == 12) + #expect(codexBreakdown.attribution?.upstream?.provider == "codex") + #expect(claude.daily.first?.totalTokens == 120) + #expect(claudeBreakdown.attribution?.upstream?.provider == "openrouter") + } +} diff --git a/Tests/CodexBarTests/CostUsageFetcherCacheSnapshotTests.swift b/Tests/CodexBarTests/CostUsageFetcherCacheSnapshotTests.swift index 685fea2a03..d19b8b40ce 100644 --- a/Tests/CodexBarTests/CostUsageFetcherCacheSnapshotTests.swift +++ b/Tests/CodexBarTests/CostUsageFetcherCacheSnapshotTests.swift @@ -3,6 +3,31 @@ import Testing @testable import CodexBarCore struct CostUsageFetcherCacheSnapshotTests { + @Test + func `cached codex token snapshot resolves the default cli proxy home`() { + let options = CostUsageFetcher.resolvedScannerOptions( + nil, + provider: .codex, + codexHomePath: nil) + + #expect(options.cliProxyAPIHome == FileManager.default.homeDirectoryForCurrentUser + .appendingPathComponent(".cli-proxy-api", isDirectory: true)) + } + + @Test + func `cache root scanner options retain the default cli proxy home`() throws { + let env = try CostUsageTestEnvironment() + defer { env.cleanup() } + + let options = try #require(CostUsageFetcher.defaultScannerOptions( + cacheRoot: env.cacheRoot, + homeDirectory: env.root)) + + #expect(options.cacheRoot == env.cacheRoot) + #expect(options.cliProxyAPIHome == env.root + .appendingPathComponent(".cli-proxy-api", isDirectory: true)) + } + @Test func `cached codex token snapshot loads from existing cache without rescanning`() async throws { let env = try CostUsageTestEnvironment() diff --git a/Tests/CodexBarTests/CostUsageFetcherCachedProxyDisconnectTests.swift b/Tests/CodexBarTests/CostUsageFetcherCachedProxyDisconnectTests.swift new file mode 100644 index 0000000000..b286d21b07 --- /dev/null +++ b/Tests/CodexBarTests/CostUsageFetcherCachedProxyDisconnectTests.swift @@ -0,0 +1,281 @@ +import Foundation +import Testing +@testable import CodexBarCore + +struct CostUsageFetcherCachedProxyDisconnectTests { + @Test + func `configuration replacement rejects an in flight Claude cache publication`() throws { + let env = try CostUsageTestEnvironment() + defer { env.cleanup() } + + let day = try env.makeLocalNoon(year: 2026, month: 7, day: 24) + _ = try env.writeClaudeProjectFile( + relativePath: "generation-race/session.jsonl", + contents: env.jsonl([[ + "type": "assistant", + "timestamp": env.isoString(for: day), + "sessionId": "session-generation-race", + "requestId": "request-generation-race", + "message": [ + "id": "message-generation-race", + "model": "claude-sonnet-4-6", + "usage": ["input_tokens": 100, "output_tokens": 5], + ], + ]])) + + let initialGeneration = try #require(CostUsageCacheLocations + .prepareCLIProxyAPIConfigurationGenerationUpdate( + stateRoot: env.cacheRoot, + fileManager: .default)) + #expect(CostUsageCacheLocations.commitCLIProxyAPIConfigurationGenerationUpdate(initialGeneration)) + + var options = CostUsageScanner.Options( + claudeProjectsRoots: [env.claudeProjectsRoot], + cacheRoot: env.cacheRoot) + options.forceRescan = true + var replacedConfiguration = false + + #expect(throws: CancellationError.self) { + _ = try CostUsageScanner.loadClaudeDaily( + provider: .claude, + range: CostUsageScanner.CostUsageDayRange(since: day, until: day), + now: day, + options: options, + checkCancellation: { + guard !replacedConfiguration else { return } + replacedConfiguration = true + let replacementGeneration = try #require(CostUsageCacheLocations + .prepareCLIProxyAPIConfigurationGenerationUpdate( + stateRoot: env.cacheRoot, + fileManager: .default)) + #expect(CostUsageCacheLocations.commitCLIProxyAPIConfigurationGenerationUpdate( + replacementGeneration)) + }) + } + #expect(replacedConfiguration) + #expect(CostUsageCacheIO.load( + provider: .claude, + cacheRoot: env.cacheRoot, + calendar: options.calendar).lastScanUnixMs == 0) + } + + @Test + func `configuration replacement rejects a report built from the previous generation`() throws { + let env = try CostUsageTestEnvironment() + defer { env.cleanup() } + + let day = try env.makeLocalNoon(year: 2026, month: 7, day: 24) + let calendar = Calendar(identifier: .gregorian) + var cache = CostUsageCache() + cache.lastScanUnixMs = Int64(day.timeIntervalSince1970 * 1000) + cache.scanSinceKey = "2026-07-24" + cache.scanUntilKey = "2026-07-24" + cache.days = ["2026-07-24": ["claude-sonnet-4-6": [100, 0, 0, 5, 0, 1]]] + CostUsageCacheIO.save( + provider: .claude, + cache: cache, + cacheRoot: env.cacheRoot, + calendar: calendar) + + let initialGeneration = try #require(CostUsageCacheLocations + .prepareCLIProxyAPIConfigurationGenerationUpdate( + stateRoot: env.cacheRoot, + fileManager: .default)) + #expect(CostUsageCacheLocations.commitCLIProxyAPIConfigurationGenerationUpdate(initialGeneration)) + + var options = CostUsageScanner.Options(cacheRoot: env.cacheRoot) + options.refreshMinIntervalSeconds = 3600 + var replacedConfiguration = false + + #expect(throws: CancellationError.self) { + _ = try CostUsageScanner.loadClaudeDaily( + provider: .claude, + range: CostUsageScanner.CostUsageDayRange(since: day, until: day), + now: day, + options: options, + checkCancellation: { + guard !replacedConfiguration else { return } + replacedConfiguration = true + let replacementGeneration = try #require(CostUsageCacheLocations + .prepareCLIProxyAPIConfigurationGenerationUpdate( + stateRoot: env.cacheRoot, + fileManager: .default)) + #expect(CostUsageCacheLocations.commitCLIProxyAPIConfigurationGenerationUpdate( + replacementGeneration)) + }) + } + #expect(replacedConfiguration) + } + + @Test + func `disconnect during proxy scan excludes the stale Codex report`() throws { + let env = try CostUsageTestEnvironment() + defer { env.cleanup() } + + let day = try env.makeLocalNoon(year: 2026, month: 7, day: 24) + _ = try env.writeClaudeProjectFile( + relativePath: "proxy-race/session.jsonl", + contents: env.jsonl([[ + "type": "assistant", + "timestamp": env.isoString(for: day), + "sessionId": "session-proxy-race", + "requestId": "request-proxy-race", + "message": [ + "id": "message-proxy-race", + "model": "claude-sonnet-4-6", + "usage": ["input_tokens": 100, "output_tokens": 5], + ], + ]])) + + let proxyHome = env.root.appendingPathComponent("cli-proxy-api", isDirectory: true) + let proxyLogs = proxyHome.appendingPathComponent("logs", isDirectory: true) + try FileManager.default.createDirectory(at: proxyLogs, withIntermediateDirectories: true) + try Data(#"{"type":"codex"}"#.utf8) + .write(to: proxyHome.appendingPathComponent("codex-auth.json")) + let proxyLog = """ + === REQUEST INFO === + URL: /v1/messages + Timestamp: \(env.isoString(for: day)) + === HEADERS === + X-Claude-Code-Session-Id: session-proxy-race + === REQUEST BODY === + {"model":"claude-sonnet-4-6"} + === API RESPONSE === + """ + try Data(proxyLog.utf8).write(to: proxyLogs.appendingPathComponent("request.log")) + #expect(CLIProxyAPIUsageCacheIO.merge( + [ + CLIProxyAPIUsageRecord( + timestamp: day, + provider: "codex", + executorType: "CodexExecutor", + model: "gpt-5.5", + alias: "claude-sonnet-4-6", + endpoint: "/v1/messages", + authType: "oauth", + requestID: "request-proxy-race", + tokens: .init(input: 100, output: 5, total: 105)), + ], + cacheRoot: env.cacheRoot, + now: day) == 1) + + var options = CostUsageScanner.Options( + claudeProjectsRoots: [env.claudeProjectsRoot], + cacheRoot: env.cacheRoot, + claudeAttributionFilter: .codexBackendOnly, + cliProxyAPIHome: proxyHome) + options.forceRescan = true + var didDisconnect = false + let report = try CostUsageScanner.loadClaudeDaily( + provider: .claude, + range: CostUsageScanner.CostUsageDayRange(since: day, until: day), + now: day, + options: options, + checkCancellation: { + guard !didDisconnect else { return } + didDisconnect = true + #expect(CostUsageCacheLocations.setCLIProxyAPIExplicitlyDisconnected( + true, + stateRoot: env.cacheRoot)) + }) + + #expect(didDisconnect) + #expect(report.data.isEmpty) + + options.forceRescan = false + options.claudeAttributionFilter = .excludeCodexBackend + let claudeReport = try CostUsageScanner.loadClaudeDaily( + provider: .claude, + range: CostUsageScanner.CostUsageDayRange(since: day, until: day), + now: day, + options: options, + checkCancellation: nil) + #expect(claudeReport.data.first?.totalTokens == 105) + #expect(claudeReport.data.first?.modelBreakdowns?.first?.attribution == nil) + } + + @Test(arguments: ["claude-sonnet-4-6", "gpt-5.5"]) + func `disconnect strips surviving cached proxy attribution`(model: String) async throws { + let env = try CostUsageTestEnvironment() + defer { env.cleanup() } + + let day = try env.makeLocalNoon(year: 2026, month: 7, day: 24) + _ = try env.writeClaudeProjectFile( + relativePath: "proxy/session.jsonl", + contents: env.jsonl([[ + "type": "assistant", + "timestamp": env.isoString(for: day), + "sessionId": "session-proxy", + "requestId": "request-proxy", + "message": [ + "id": "message-proxy", + "model": model, + "usage": ["input_tokens": 100, "output_tokens": 5], + ], + ]])) + + let proxyHome = env.root.appendingPathComponent("cli-proxy-api", isDirectory: true) + let proxyLogs = proxyHome.appendingPathComponent("logs", isDirectory: true) + try FileManager.default.createDirectory(at: proxyLogs, withIntermediateDirectories: true) + try Data(#"{"type":"codex"}"#.utf8) + .write(to: proxyHome.appendingPathComponent("codex-auth.json")) + let proxyLog = """ + === REQUEST INFO === + URL: /v1/messages + Timestamp: \(env.isoString(for: day)) + === HEADERS === + X-Claude-Code-Session-Id: session-proxy + === REQUEST BODY === + {"model":"\(model)"} + === API RESPONSE === + """ + try Data(proxyLog.utf8).write(to: proxyLogs.appendingPathComponent("request.log")) + CLIProxyAPIUsageCacheIO.merge( + [ + CLIProxyAPIUsageRecord( + timestamp: day, + provider: "codex", + executorType: "CodexExecutor", + model: "gpt-5.5", + alias: model, + endpoint: "/v1/messages", + authType: "oauth", + requestID: "proxy-request", + tokens: .init(input: 100, output: 5, total: 105)), + ], + cacheRoot: env.cacheRoot, + now: day) + + let options = CostUsageScanner.Options( + codexSessionsRoot: env.codexSessionsRoot, + claudeProjectsRoots: [env.claudeProjectsRoot], + cacheRoot: env.cacheRoot, + cliProxyAPIHome: proxyHome) + let attributedCodex = try await CostUsageFetcher.loadTokenSnapshot( + provider: .codex, + now: day, + allowPricingRefresh: false, + includePiSessions: false, + scannerOptions: options) + #expect(attributedCodex.daily.first?.totalTokens == 105) + #expect(attributedCodex.daily.first?.modelBreakdowns?.first?.attribution?.route == .cliProxyAPI) + + try FileManager.default.removeItem(at: proxyLogs) + try FileManager.default.removeItem(at: proxyHome.appendingPathComponent("codex-auth.json")) + #expect(CostUsageCacheLocations.setCLIProxyAPIExplicitlyDisconnected( + true, + stateRoot: env.cacheRoot)) + + let disconnectedClaude = try await CostUsageFetcher.loadTokenSnapshot( + provider: .claude, + now: day, + allowPricingRefresh: false, + includePiSessions: false, + scannerOptions: options) + #expect(disconnectedClaude.daily.first?.totalTokens == 105) + #expect(disconnectedClaude.daily.first?.modelBreakdowns?.first?.attribution == nil) + #expect(await CostUsageFetcher.loadCachedCodexTokenSnapshot( + now: day, + scannerOptions: options) == nil) + } +} diff --git a/Tests/CodexBarTests/CostUsageFetcherCachedProxyTimeZoneTests.swift b/Tests/CodexBarTests/CostUsageFetcherCachedProxyTimeZoneTests.swift new file mode 100644 index 0000000000..8762dff282 --- /dev/null +++ b/Tests/CodexBarTests/CostUsageFetcherCachedProxyTimeZoneTests.swift @@ -0,0 +1,89 @@ +import Foundation +import Testing +@testable import CodexBarCore + +struct CostUsageFetcherCachedProxyTimeZoneTests { + @Test + func `cached codex snapshot rejects claude proxy cache from another time zone`() async throws { + let env = try CostUsageTestEnvironment() + defer { env.cleanup() } + + let day = try env.makeLocalNoon(year: 2026, month: 7, day: 24) + _ = try env.writeClaudeProjectFile( + relativePath: "proxy/session.jsonl", + contents: env.jsonl([[ + "type": "assistant", + "timestamp": env.isoString(for: day), + "sessionId": "session-proxy", + "requestId": "request-proxy", + "message": [ + "id": "message-proxy", + "model": "claude-sonnet-4-6", + "usage": ["input_tokens": 100, "output_tokens": 5], + ], + ]])) + + let proxyHome = env.root.appendingPathComponent("cli-proxy-api", isDirectory: true) + let proxyLogs = proxyHome.appendingPathComponent("logs", isDirectory: true) + try FileManager.default.createDirectory(at: proxyLogs, withIntermediateDirectories: true) + try Data(#"{"type":"codex"}"#.utf8) + .write(to: proxyHome.appendingPathComponent("codex-auth.json")) + let proxyLog = """ + === REQUEST INFO === + URL: /v1/messages + Timestamp: \(env.isoString(for: day)) + === HEADERS === + X-Claude-Code-Session-Id: session-proxy + === REQUEST BODY === + {"model":"claude-sonnet-4-6"} + === API RESPONSE === + """ + try Data(proxyLog.utf8).write(to: proxyLogs.appendingPathComponent("request.log")) + CLIProxyAPIUsageCacheIO.merge( + [ + CLIProxyAPIUsageRecord( + timestamp: day, + provider: "codex", + executorType: "CodexExecutor", + model: "gpt-5.5", + alias: "claude-sonnet-4-6", + endpoint: "/v1/messages", + authType: "oauth", + requestID: "proxy-request", + tokens: .init(input: 100, output: 5, total: 105)), + ], + cacheRoot: env.cacheRoot, + now: day) + + let options = CostUsageScanner.Options( + codexSessionsRoot: env.codexSessionsRoot, + claudeProjectsRoots: [env.claudeProjectsRoot], + cacheRoot: env.cacheRoot, + cliProxyAPIHome: proxyHome) + _ = try await CostUsageFetcher.loadTokenSnapshot( + provider: .codex, + now: day, + allowPricingRefresh: false, + includePiSessions: false, + scannerOptions: options) + + let claudeCache = CostUsageCacheIO.load(provider: .claude, cacheRoot: env.cacheRoot) + #expect(!claudeCache.days.isEmpty) + var staleZoneCalendar = options.calendar + staleZoneCalendar.timeZone = try #require( + ["UTC", "Asia/Bangkok"] + .compactMap(TimeZone.init(identifier:)) + .first { $0.identifier != options.calendar.timeZone.identifier }) + CostUsageCacheIO.save( + provider: .claude, + cache: claudeCache, + cacheRoot: env.cacheRoot, + calendar: staleZoneCalendar) + + let cachedSnapshot = await CostUsageFetcher.loadCachedCodexTokenSnapshot( + now: day, + scannerOptions: options) + + #expect(cachedSnapshot == nil) + } +} diff --git a/Tests/CodexBarTests/CostUsageFetcherTests.swift b/Tests/CodexBarTests/CostUsageFetcherTests.swift index 909807640d..7d668c372c 100644 --- a/Tests/CodexBarTests/CostUsageFetcherTests.swift +++ b/Tests/CodexBarTests/CostUsageFetcherTests.swift @@ -2,6 +2,8 @@ import Foundation import Testing @testable import CodexBarCore +// swiftlint:disable file_length + @Suite(.serialized) struct CostUsageFetcherTests { @Test @@ -709,6 +711,619 @@ extension CostUsageFetcherTests { ]) } + @Test + func `claude code proxy usage belongs to codex and keeps route attribution`() async throws { + let env = try CostUsageTestEnvironment() + defer { env.cleanup() } + + let day = try env.makeLocalNoon(year: 2026, month: 7, day: 24) + let proxyAssistant: [String: Any] = [ + "type": "assistant", + "timestamp": env.isoString(for: day), + "sessionId": "session-proxy", + "requestId": "request-proxy", + "message": [ + "id": "message-proxy", + "model": "claude-sonnet-4-6", + "usage": [ + "input_tokens": 100, + "cache_creation_input_tokens": 10, + "cache_read_input_tokens": 20, + "output_tokens": 5, + ], + ], + ] + let nativeClaudeAssistant: [String: Any] = [ + "type": "assistant", + "timestamp": env.isoString(for: day.addingTimeInterval(1)), + "sessionId": "session-proxy", + "requestId": "request-claude", + "message": [ + "id": "message-claude", + "model": "claude-sonnet-4-6", + "usage": [ + "input_tokens": 50, + "cache_creation_input_tokens": 5, + "cache_read_input_tokens": 5, + "output_tokens": 10, + ], + ], + ] + _ = try env.writeClaudeProjectFile( + relativePath: "proxy/session.jsonl", + contents: env.jsonl([proxyAssistant, nativeClaudeAssistant])) + let cliProxyHome = env.root.appendingPathComponent("cli-proxy-api", isDirectory: true) + let cliProxyLogs = cliProxyHome.appendingPathComponent("logs", isDirectory: true) + try FileManager.default.createDirectory(at: cliProxyLogs, withIntermediateDirectories: true) + try Data(#"{"type":"codex","access_token":"must-not-be-exposed"}"#.utf8) + .write(to: cliProxyHome.appendingPathComponent("codex-auth.json")) + let proxyLog = """ + === REQUEST INFO === + URL: /v1/messages + Timestamp: \(env.isoString(for: day)) + === HEADERS === + X-Claude-Code-Session-Id: session-proxy + === REQUEST BODY === + {"model":"claude-sonnet-4-6"} + === API RESPONSE === + """ + try Data(proxyLog.utf8).write(to: cliProxyLogs.appendingPathComponent("request.log")) + CLIProxyAPIUsageCacheIO.merge( + [ + CLIProxyAPIUsageRecord( + timestamp: day, + provider: "codex", + executorType: "CodexExecutor", + model: "gpt-5.5", + alias: "claude-sonnet-4-6", + endpoint: "/v1/messages", + authType: "oauth", + requestID: "cliproxy-request-proxy", + tokens: .init( + input: 100, + output: 5, + cacheRead: 20, + cacheCreation: 10, + total: 135)), + ], + cacheRoot: env.cacheRoot, + now: day) + + let options = CostUsageScanner.Options( + codexSessionsRoot: env.codexSessionsRoot, + claudeProjectsRoots: [env.claudeProjectsRoot], + cacheRoot: env.cacheRoot, + cliProxyAPIHome: cliProxyHome) + let codex = try await CostUsageFetcher.loadTokenSnapshot( + provider: .codex, + now: day, + allowPricingRefresh: false, + includePiSessions: false, + scannerOptions: options) + let scopedHome = env.root.appendingPathComponent("managed-codex-home", isDirectory: true) + try FileManager.default.createDirectory( + at: scopedHome.appendingPathComponent("sessions", isDirectory: true), + withIntermediateDirectories: true) + let scopedCodex = try await CostUsageFetcher.loadTokenSnapshot( + provider: .codex, + now: day, + codexHomePath: scopedHome.path, + allowPricingRefresh: false, + includePiSessions: false, + scannerOptions: options) + let claude = try await CostUsageFetcher.loadTokenSnapshot( + provider: .claude, + now: day, + allowPricingRefresh: false, + includePiSessions: false, + scannerOptions: options) + let cachedCodex = try #require(await CostUsageFetcher.loadCachedCodexTokenSnapshot( + now: day, + scannerOptions: options)) + let cachedScopedCodex = await CostUsageFetcher.loadCachedCodexTokenSnapshot( + now: day, + codexHomePath: scopedHome.path, + allowScopedCodexHome: true, + includePiSessions: false, + includeProjectAndSessionBreakdowns: false, + scannerOptions: options) + + let expectedCodexCost = try #require(CostUsagePricing.claudeProxyCodexCostUSD( + model: "gpt-5.5", + inputTokens: 100, + cacheReadInputTokens: 20, + cacheCreationInputTokens: 10, + outputTokens: 5)) + let codexBreakdown = try #require(codex.daily.first?.modelBreakdowns?.first) + #expect(codex.daily.first?.totalTokens == 135) + #expect(abs((codex.daily.first?.costUSD ?? 0) - expectedCodexCost) < 0.000001) + #expect(codexBreakdown.modelName == "claude-sonnet-4-6") + #expect(codexBreakdown.attribution == CostUsageAttribution( + client: .claudeCode, + route: .cliProxyAPI, + modelProvider: .anthropic, + upstream: .init( + provider: "codex", + authType: .oauth, + model: "gpt-5.5", + executorType: "CodexExecutor"), + evidence: [.cliProxyRequestLog, .cliProxyUsageTelemetry, .modelProvider])) + #expect(codex.projects.map(\.name) == ["Claude Code via CLIProxyAPI"]) + #expect(scopedCodex.daily.isEmpty) + #expect(scopedCodex.projects.isEmpty) + + #expect(claude.daily.first?.totalTokens == 70) + #expect(claude.daily.first?.modelsUsed == ["claude-sonnet-4-6"]) + #expect(claude.daily.first?.modelBreakdowns?.first?.attribution == nil) + + #expect(cachedCodex.daily.first?.totalTokens == 135) + #expect(cachedCodex.daily.first?.modelBreakdowns?.first?.attribution == codexBreakdown.attribution) + #expect(cachedScopedCodex == nil) + + try FileManager.default.removeItem(at: cliProxyLogs) + try FileManager.default.removeItem(at: cliProxyHome.appendingPathComponent("codex-auth.json")) + let cachedAfterLogRotation = try #require(await CostUsageFetcher.loadCachedCodexTokenSnapshot( + now: day, + scannerOptions: options)) + #expect(cachedAfterLogRotation.daily.first?.totalTokens == 135) + #expect(cachedAfterLogRotation.daily.first?.modelBreakdowns?.first?.attribution == codexBreakdown.attribution) + } + + @Test + func `proxy and pi usage keep distinct synthetic projects`() async throws { + let env = try CostUsageTestEnvironment() + defer { env.cleanup() } + + let day = try env.makeLocalNoon(year: 2026, month: 7, day: 24) + _ = try env.writeClaudeProjectFile( + relativePath: "proxy/session.jsonl", + contents: env.jsonl([[ + "type": "assistant", + "timestamp": env.isoString(for: day), + "sessionId": "session-proxy", + "requestId": "request-proxy", + "message": [ + "id": "message-proxy", + "model": "gpt-5.5", + "usage": ["input_tokens": 100, "output_tokens": 5], + ], + ]])) + _ = try env.writePiSessionFile( + relativePath: "2026-07-24T10-00-00-000Z_pi.jsonl", + contents: env.jsonl([[ + "type": "message", + "timestamp": env.isoString(for: day.addingTimeInterval(1)), + "message": [ + "role": "assistant", + "provider": "openai-codex", + "model": "openai/gpt-5.4", + "timestamp": Int(day.addingTimeInterval(1).timeIntervalSince1970 * 1000), + "usage": ["input": 50, "output": 5, "totalTokens": 55], + ], + ]])) + let cliProxyHome = env.root.appendingPathComponent("cli-proxy-api", isDirectory: true) + let cliProxyLogs = cliProxyHome.appendingPathComponent("logs", isDirectory: true) + try FileManager.default.createDirectory(at: cliProxyLogs, withIntermediateDirectories: true) + try Data(#"{"type":"codex"}"#.utf8) + .write(to: cliProxyHome.appendingPathComponent("codex-auth.json")) + let proxyLog = """ + === REQUEST INFO === + URL: /v1/messages + Timestamp: \(env.isoString(for: day)) + === HEADERS === + X-Claude-Code-Session-Id: session-proxy + === REQUEST BODY === + {"model":"gpt-5.5"} + === API RESPONSE === + """ + try Data(proxyLog.utf8).write(to: cliProxyLogs.appendingPathComponent("request.log")) + let options = CostUsageScanner.Options( + codexSessionsRoot: env.codexSessionsRoot, + claudeProjectsRoots: [env.claudeProjectsRoot], + cacheRoot: env.cacheRoot, + cliProxyAPIHome: cliProxyHome) + let piOptions = PiSessionCostScanner.Options( + piSessionsRoot: env.piSessionsRoot, + cacheRoot: env.cacheRoot, + refreshMinIntervalSeconds: 0) + + let snapshot = try await CostUsageFetcher.loadTokenSnapshot( + provider: .codex, + now: day, + allowPricingRefresh: false, + scannerOptions: options, + piScannerOptions: piOptions) + + #expect(snapshot.projects.count == 2) + #expect(Set(snapshot.projects.map(\.name)) == [ + "Claude Code via CLIProxyAPI", + CostUsageProjectBreakdown.unknownProjectName, + ]) + #expect(snapshot.projects.allSatisfy { $0.path == nil }) + #expect(snapshot.projects.allSatisfy { $0.sources.map(\.name) == [$0.name] }) + } + + @Test + func `codex supplemental scan requires proxy evidence`() throws { + let env = try CostUsageTestEnvironment() + defer { env.cleanup() } + let cliProxyHome = env.root.appendingPathComponent("cli-proxy-api", isDirectory: true) + let options = CostUsageScanner.Options( + cacheRoot: env.cacheRoot, + cliProxyAPIHome: cliProxyHome) + + #expect(!CostUsageFetcher.hasCodexProxyEvidence(options: options)) + + let logs = cliProxyHome.appendingPathComponent("logs", isDirectory: true) + try FileManager.default.createDirectory(at: logs, withIntermediateDirectories: true) + try Data().write(to: logs.appendingPathComponent("request.log")) + + #expect(CostUsageFetcher.hasCodexProxyEvidence(options: options)) + } + + @Test + func `explicit disconnect suppresses surviving proxy logs`() throws { + let env = try CostUsageTestEnvironment() + defer { env.cleanup() } + let cliProxyHome = env.root.appendingPathComponent("cli-proxy-api", isDirectory: true) + let options = CostUsageScanner.Options( + cacheRoot: env.cacheRoot, + cliProxyAPIHome: cliProxyHome) + let logs = cliProxyHome.appendingPathComponent("logs", isDirectory: true) + try FileManager.default.createDirectory(at: logs, withIntermediateDirectories: true) + try Data().write(to: logs.appendingPathComponent("request.log")) + #expect(CostUsageFetcher.hasCodexProxyEvidence(options: options)) + + #expect(CostUsageCacheLocations.setCLIProxyAPIExplicitlyDisconnected( + true, + stateRoot: env.cacheRoot)) + #expect(!CostUsageFetcher.hasCodexProxyEvidence(options: options)) + } + + @Test + func `explicit disconnect keeps proxy routed usage in Claude totals`() async throws { + let env = try CostUsageTestEnvironment() + defer { env.cleanup() } + let day = try env.makeLocalNoon(year: 2026, month: 7, day: 24) + _ = try env.writeClaudeProjectFile( + relativePath: "disconnected-proxy/session.jsonl", + contents: env.jsonl([[ + "type": "assistant", + "timestamp": env.isoString(for: day), + "sessionId": "disconnected-proxy-session", + "requestId": "disconnected-proxy-request", + "message": [ + "id": "disconnected-proxy-message", + "model": "claude-sonnet-4-6", + "usage": ["input_tokens": 100, "output_tokens": 5], + ], + ]])) + let cliProxyHome = env.root.appendingPathComponent("cli-proxy-api", isDirectory: true) + let logs = cliProxyHome.appendingPathComponent("logs", isDirectory: true) + try FileManager.default.createDirectory(at: logs, withIntermediateDirectories: true) + try Data(#"{"type":"codex","access_token":"must-not-be-exposed"}"#.utf8) + .write(to: cliProxyHome.appendingPathComponent("codex-auth.json")) + let proxyLog = """ + === REQUEST INFO === + URL: /v1/messages + Timestamp: \(env.isoString(for: day)) + === HEADERS === + X-Claude-Code-Session-Id: disconnected-proxy-session + === REQUEST BODY === + {"model":"claude-sonnet-4-6"} + === API RESPONSE === + """ + try Data(proxyLog.utf8).write(to: logs.appendingPathComponent("request.log")) + let options = CostUsageScanner.Options( + codexSessionsRoot: env.codexSessionsRoot, + claudeProjectsRoots: [env.claudeProjectsRoot], + cacheRoot: env.cacheRoot, + cliProxyAPIHome: cliProxyHome) + #expect(CostUsageCacheLocations.setCLIProxyAPIExplicitlyDisconnected( + true, + stateRoot: env.cacheRoot)) + + let claude = try await CostUsageFetcher.loadTokenSnapshot( + provider: .claude, + now: day, + allowPricingRefresh: false, + includePiSessions: false, + scannerOptions: options) + let codex = try await CostUsageFetcher.loadTokenSnapshot( + provider: .codex, + now: day, + allowPricingRefresh: false, + includePiSessions: false, + scannerOptions: options) + let cachedCodex = await CostUsageFetcher.loadCachedCodexTokenSnapshot( + now: day, + scannerOptions: options) + + #expect(claude.daily.first?.totalTokens == 105) + #expect(claude.daily.first?.modelBreakdowns?.first?.attribution == nil) + #expect(codex.daily.isEmpty) + #expect(cachedCodex == nil) + } + + @Test + func `openai model without proxy evidence stays out of codex totals`() async throws { + let env = try CostUsageTestEnvironment() + defer { env.cleanup() } + + let day = try env.makeLocalNoon(year: 2026, month: 7, day: 24) + _ = try env.writeClaudeProjectFile( + relativePath: "unresolved/session.jsonl", + contents: env.jsonl([[ + "type": "assistant", + "timestamp": env.isoString(for: day), + "sessionId": "unmatched-session", + "requestId": "request-unresolved", + "message": [ + "id": "message-unresolved", + "model": "gpt-5.5", + "usage": ["input_tokens": 100, "output_tokens": 5], + ], + ]])) + let cliProxyHome = env.root.appendingPathComponent("cli-proxy-api", isDirectory: true) + try FileManager.default.createDirectory( + at: cliProxyHome.appendingPathComponent("logs", isDirectory: true), + withIntermediateDirectories: true) + try Data(#"{"type":"codex"}"#.utf8) + .write(to: cliProxyHome.appendingPathComponent("codex-auth.json")) + let options = CostUsageScanner.Options( + codexSessionsRoot: env.codexSessionsRoot, + claudeProjectsRoots: [env.claudeProjectsRoot], + cacheRoot: env.cacheRoot, + cliProxyAPIHome: cliProxyHome) + + let codex = try await CostUsageFetcher.loadTokenSnapshot( + provider: .codex, + now: day, + allowPricingRefresh: false, + includePiSessions: false, + scannerOptions: options) + let claude = try await CostUsageFetcher.loadTokenSnapshot( + provider: .claude, + now: day, + allowPricingRefresh: false, + includePiSessions: false, + scannerOptions: options) + + #expect(codex.daily.isEmpty) + #expect(claude.daily.isEmpty) + } + + @Test + func `proxy route without a confirmed upstream stays out of provider totals`() async throws { + let env = try CostUsageTestEnvironment() + defer { env.cleanup() } + + let day = try env.makeLocalNoon(year: 2026, month: 7, day: 24) + _ = try env.writeClaudeProjectFile( + relativePath: "unresolved-proxy/session.jsonl", + contents: env.jsonl([[ + "type": "assistant", + "timestamp": env.isoString(for: day), + "sessionId": "unresolved-proxy-session", + "requestId": "unresolved-proxy-request", + "message": [ + "id": "unresolved-proxy-message", + "model": "gpt-5.5", + "usage": ["input_tokens": 100, "output_tokens": 5], + ], + ]])) + let cliProxyHome = env.root.appendingPathComponent("cli-proxy-api", isDirectory: true) + let logs = cliProxyHome.appendingPathComponent("logs", isDirectory: true) + try FileManager.default.createDirectory(at: logs, withIntermediateDirectories: true) + let proxyLog = """ + === REQUEST INFO === + URL: /v1/messages + Timestamp: \(env.isoString(for: day)) + === HEADERS === + X-Claude-Code-Session-Id: unresolved-proxy-session + === REQUEST BODY === + {"model":"gpt-5.5"} + === API RESPONSE === + """ + try Data(proxyLog.utf8).write(to: logs.appendingPathComponent("request.log")) + let options = CostUsageScanner.Options( + codexSessionsRoot: env.codexSessionsRoot, + claudeProjectsRoots: [env.claudeProjectsRoot], + cacheRoot: env.cacheRoot, + cliProxyAPIHome: cliProxyHome) + + let codex = try await CostUsageFetcher.loadTokenSnapshot( + provider: .codex, + now: day, + allowPricingRefresh: false, + includePiSessions: false, + scannerOptions: options) + let claude = try await CostUsageFetcher.loadTokenSnapshot( + provider: .claude, + now: day, + allowPricingRefresh: false, + includePiSessions: false, + scannerOptions: options) + + #expect(codex.daily.isEmpty) + #expect(claude.daily.isEmpty) + } + + @Test + func `cliproxy request log does not cover a distant resumed turn`() async throws { + let env = try CostUsageTestEnvironment() + defer { env.cleanup() } + + let day = try env.makeLocalNoon(year: 2026, month: 7, day: 24) + func assistant(seconds: TimeInterval, requestID: String) -> [String: Any] { + [ + "type": "assistant", + "timestamp": env.isoString(for: day.addingTimeInterval(seconds)), + "sessionId": "session-proxy", + "requestId": requestID, + "message": [ + "id": "message-\(requestID)", + "model": "gpt-5.5", + "usage": ["input_tokens": 100, "output_tokens": 5], + ], + ] + } + _ = try env.writeClaudeProjectFile( + relativePath: "stable-proxy/session.jsonl", + contents: env.jsonl([ + assistant(seconds: 0, requestID: "first"), + assistant(seconds: 3 * 60 * 60, requestID: "second"), + ])) + let cliProxyHome = env.root.appendingPathComponent("cli-proxy-api", isDirectory: true) + let cliProxyLogs = cliProxyHome.appendingPathComponent("logs", isDirectory: true) + try FileManager.default.createDirectory(at: cliProxyLogs, withIntermediateDirectories: true) + try Data(#"{"type":"codex","disabled":false,"access_token":"must-not-be-exposed"}"#.utf8) + .write(to: cliProxyHome.appendingPathComponent("codex-auth.json")) + let proxyLog = """ + === REQUEST INFO === + URL: /v1/messages + Timestamp: \(env.isoString(for: day)) + === HEADERS === + X-Claude-Code-Session-Id: session-proxy + === REQUEST BODY === + {"model":"gpt-5.5"} + === API RESPONSE === + """ + try Data(proxyLog.utf8).write(to: cliProxyLogs.appendingPathComponent("request.log")) + let options = CostUsageScanner.Options( + codexSessionsRoot: env.codexSessionsRoot, + claudeProjectsRoots: [env.claudeProjectsRoot], + cacheRoot: env.cacheRoot, + cliProxyAPIHome: cliProxyHome) + + let codex = try await CostUsageFetcher.loadTokenSnapshot( + provider: .codex, + now: day, + allowPricingRefresh: false, + includePiSessions: false, + scannerOptions: options) + let claude = try await CostUsageFetcher.loadTokenSnapshot( + provider: .claude, + now: day, + allowPricingRefresh: false, + includePiSessions: false, + scannerOptions: options) + + let breakdown = try #require(codex.daily.first?.modelBreakdowns?.first) + #expect(codex.daily.first?.totalTokens == 105) + #expect(codex.daily.first?.modelBreakdowns?.count == 1) + #expect(breakdown.modelName == "gpt-5.5") + #expect(breakdown.attribution?.route == .cliProxyAPI) + #expect(breakdown.attribution?.upstream == .init( + provider: "codex", + authType: .oauth, + model: "gpt-5.5")) + #expect(breakdown.attribution?.evidence.contains(.cliProxyAuthInventory) == true) + #expect(claude.daily.isEmpty) + } + + @Test + func `claude report preserves non codex proxy backend attribution`() async throws { + let env = try CostUsageTestEnvironment() + defer { env.cleanup() } + + let day = try env.makeLocalNoon(year: 2026, month: 7, day: 24) + func assistant(sessionID: String, requestID: String, model: String, seconds: TimeInterval) -> [String: Any] { + [ + "type": "assistant", + "timestamp": env.isoString(for: day.addingTimeInterval(seconds)), + "sessionId": sessionID, + "requestId": requestID, + "message": [ + "id": "message-\(requestID)", + "model": model, + "usage": ["input_tokens": 10, "output_tokens": 2], + ], + ] + } + _ = try env.writeClaudeProjectFile( + relativePath: "multi-backend/session.jsonl", + contents: env.jsonl([ + assistant(sessionID: "session-gemini", requestID: "gemini", model: "gemini-3-pro", seconds: 0), + assistant( + sessionID: "session-claude", + requestID: "claude", + model: "claude-sonnet-4-6", + seconds: 1), + ])) + + let cliProxyHome = env.root.appendingPathComponent("cli-proxy-api", isDirectory: true) + let logs = cliProxyHome.appendingPathComponent("logs", isDirectory: true) + try FileManager.default.createDirectory(at: logs, withIntermediateDirectories: true) + try Data(#"{"type":"gemini"}"#.utf8) + .write(to: cliProxyHome.appendingPathComponent("gemini-auth.json")) + try Data(#"{"type":"claude"}"#.utf8) + .write(to: cliProxyHome.appendingPathComponent("claude-auth.json")) + for (name, sessionID, model, seconds) in [ + ("gemini", "session-gemini", "gemini-3-pro", 0.0), + ("claude", "session-claude", "claude-sonnet-4-6", 1.0), + ] { + let log = """ + === REQUEST INFO === + URL: /v1/messages + Timestamp: \(env.isoString(for: day.addingTimeInterval(seconds))) + === HEADERS === + X-Claude-Code-Session-Id: \(sessionID) + === REQUEST BODY === + {"model":"\(model)"} + === API RESPONSE === + """ + try Data(log.utf8).write(to: logs.appendingPathComponent("\(name).log")) + } + CLIProxyAPIUsageCacheIO.merge( + [ + CLIProxyAPIUsageRecord( + timestamp: day, + provider: "gemini", + executorType: "GeminiExecutor", + model: "gemini-3-pro", + alias: "gemini-3-pro", + endpoint: "/v1/messages", + authType: "oauth", + requestID: "cliproxy-gemini", + tokens: .init(input: 10, output: 2, total: 12)), + CLIProxyAPIUsageRecord( + timestamp: day.addingTimeInterval(1), + provider: "claude", + executorType: "ClaudeExecutor", + model: "claude-sonnet-4-6", + alias: "claude-sonnet-4-6", + endpoint: "/v1/messages", + authType: "oauth", + requestID: "cliproxy-claude", + tokens: .init(input: 10, output: 2, total: 12)), + ], + cacheRoot: env.cacheRoot, + now: day.addingTimeInterval(1)) + + let options = CostUsageScanner.Options( + claudeProjectsRoots: [env.claudeProjectsRoot], + cacheRoot: env.cacheRoot, + cliProxyAPIHome: cliProxyHome) + let snapshot = try await CostUsageFetcher.loadTokenSnapshot( + provider: .claude, + now: day, + allowPricingRefresh: false, + includePiSessions: false, + scannerOptions: options) + + let breakdowns = try #require(snapshot.daily.first?.modelBreakdowns) + let claude = try #require(breakdowns.first { $0.modelName == "claude-sonnet-4-6" }) + let gemini = try #require(breakdowns.first { $0.modelName == "gemini-3-pro" }) + #expect(claude.attribution?.route == .cliProxyAPI) + #expect(claude.attribution?.upstream?.provider == "claude") + #expect(claude.attribution?.upstream?.authType == .oauth) + #expect(gemini.attribution?.route == .cliProxyAPI) + #expect(gemini.attribution?.upstream?.provider == "gemini") + #expect(gemini.attribution?.upstream?.authType == .oauth) + #expect(gemini.costUSD == nil) + } + @Test func `fetcher prefers turn context model over token count fallback`() async throws { let env = try CostUsageTestEnvironment() diff --git a/Tests/CodexBarTests/CostUsageFetcherUnknownModelPricingTests.swift b/Tests/CodexBarTests/CostUsageFetcherUnknownModelPricingTests.swift index 4ea858f6c4..1b72fc393c 100644 --- a/Tests/CodexBarTests/CostUsageFetcherUnknownModelPricingTests.swift +++ b/Tests/CodexBarTests/CostUsageFetcherUnknownModelPricingTests.swift @@ -6,6 +6,28 @@ import Testing @testable import CodexBarCore struct CostUsageFetcherUnknownModelPricingTests { + @Test + func `resolved upstream model does not reuse a cached alias price`() { + #expect(CostUsageScanner.resolvedClaudeRowCost( + wasPriced: true, + cachedCostNanos: 1_000_000_000, + cachedPricingModel: "claude-priced-alias", + pricingModel: "unpriced-upstream", + currentCost: nil) == nil) + #expect(CostUsageScanner.resolvedClaudeRowCost( + wasPriced: true, + cachedCostNanos: 1_000_000_000, + cachedPricingModel: "claude-priced-alias", + pricingModel: "claude-priced-alias", + currentCost: nil) == 1) + #expect(CostUsageScanner.resolvedClaudeRowCost( + wasPriced: true, + cachedCostNanos: 1_000_000_000, + cachedPricingModel: "claude-priced-alias", + pricingModel: "priced-upstream", + currentCost: 2) == 2) + } + @Test func `fetcher reprices an unknown model after an on demand catalog refresh`() async throws { let fixture = try UnknownModelPricingFixture() @@ -187,6 +209,210 @@ struct CostUsageFetcherUnknownModelPricingTests { #expect(breakdown.costUSD == nil) #expect(await counter.requestCount == 0) } + + @Test + func `proxy-only fetcher refreshes pricing for the resolved upstream model`() async throws { + let environment = try CostUsageTestEnvironment() + defer { environment.cleanup() } + let day = try environment.makeLocalNoon(year: 2026, month: 7, day: 24) + let alias = "claude-proxy-alias" + let freshCatalog = try JSONDecoder().decode(ModelsDevCatalog.self, from: Data(""" + { + "openai": { + "id": "openai", + "models": { "gpt-old": { "id": "gpt-old", "cost": { "input": 1, "output": 4 } } } + } + } + """.utf8)) + ModelsDevCache.save( + catalog: freshCatalog, + fetchedAt: day.addingTimeInterval(-901), + cacheRoot: environment.cacheRoot) + + _ = try environment.writeClaudeProjectFile( + relativePath: "proxy/unknown-model.jsonl", + contents: environment.jsonl([[ + "type": "assistant", + "timestamp": environment.isoString(for: day), + "sessionId": "session-proxy", + "requestId": "request-proxy", + "message": [ + "id": "message-proxy", + "model": "\(alias)", + "usage": ["input_tokens": 100, "output_tokens": 10], + ], + ]])) + let cliProxyHome = environment.root.appendingPathComponent("cli-proxy-api", isDirectory: true) + let cliProxyLogs = cliProxyHome.appendingPathComponent("logs", isDirectory: true) + try FileManager.default.createDirectory(at: cliProxyLogs, withIntermediateDirectories: true) + try Data(#"{"type":"codex"}"#.utf8) + .write(to: cliProxyHome.appendingPathComponent("codex-auth.json")) + let proxyLog = """ + === REQUEST INFO === + URL: /v1/messages + Timestamp: \(environment.isoString(for: day)) + === HEADERS === + X-Claude-Code-Session-Id: session-proxy + === REQUEST BODY === + {"model":"\(alias)"} + === API RESPONSE === + """ + try Data(proxyLog.utf8).write(to: cliProxyLogs.appendingPathComponent("request.log")) + CLIProxyAPIUsageCacheIO.merge( + [ + CLIProxyAPIUsageRecord( + timestamp: day, + provider: "codex", + executorType: "CodexExecutor", + model: "gpt-new", + alias: alias, + endpoint: "/v1/messages", + authType: "oauth", + requestID: "cliproxy-request", + tokens: .init(input: 100, output: 10, total: 110)), + ], + cacheRoot: environment.cacheRoot, + now: day) + let options = CostUsageScanner.Options( + claudeProjectsRoots: [environment.claudeProjectsRoot], + cacheRoot: environment.cacheRoot, + cliProxyAPIHome: cliProxyHome) + let refreshedCatalog = Data(""" + { + "openai": { + "id": "openai", + "models": { "gpt-new": { "id": "gpt-new", "cost": { "input": 2, "output": 8 } } } + }, + "anthropic": { + "id": "anthropic", + "models": { "claude-new": { "id": "claude-new", "cost": { "input": 3, "output": 15 } } } + } + } + """.utf8) + + let snapshot = try await CostUsageFetcher(scannerOptions: options).loadCodexProxyTokenSnapshot( + now: day, + forceRefresh: true, + refreshPricingInBackground: false, + modelsDevClient: ModelsDevClient(transport: CostUsageFetcherModelsDevTransport( + data: refreshedCatalog))) + + let breakdown = try #require(snapshot.daily.first?.modelBreakdowns?.first) + #expect(breakdown.modelName == alias) + #expect(breakdown.attribution?.upstream?.model == "gpt-new") + #expect(abs((breakdown.costUSD ?? 0) - 0.00028) < 0.0000001) + } + + @Test(arguments: [ + ("gpt-new", "openrouter", "OpenAICompatExecutor", 0.00028), + ("claude-new", "openrouter", "OpenAICompatExecutor", 0.00045), + ("gemma-new", "gemini", nil, 0.00036), + ]) + func `claude fetch resolves proxy pricing across upstream providers`( + upstreamModel: String, + upstreamProvider: String, + executorType: String?, + expectedCost: Double) async throws + { + let environment = try CostUsageTestEnvironment() + defer { environment.cleanup() } + let day = try environment.makeLocalNoon(year: 2026, month: 7, day: 24) + let alias = "claude-proxy-alias" + let staleCatalog = try JSONDecoder().decode(ModelsDevCatalog.self, from: Data(""" + { + "openai": { + "id": "openai", + "models": { "gpt-old": { "id": "gpt-old", "cost": { "input": 1, "output": 4 } } } + }, + "anthropic": { + "id": "anthropic", + "models": { "claude-old": { "id": "claude-old", "cost": { "input": 3, "output": 15 } } } + } + } + """.utf8)) + ModelsDevCache.save( + catalog: staleCatalog, + fetchedAt: day.addingTimeInterval(-901), + cacheRoot: environment.cacheRoot) + + _ = try environment.writeClaudeProjectFile( + relativePath: "proxy/openrouter-unknown-model.jsonl", + contents: environment.jsonl([[ + "type": "assistant", + "timestamp": environment.isoString(for: day), + "sessionId": "session-openrouter", + "requestId": "request-openrouter", + "message": [ + "id": "message-openrouter", + "model": "\(alias)", + "usage": ["input_tokens": 100, "output_tokens": 10], + ], + ]])) + let cliProxyHome = environment.root.appendingPathComponent("cli-proxy-api", isDirectory: true) + let cliProxyLogs = cliProxyHome.appendingPathComponent("logs", isDirectory: true) + try FileManager.default.createDirectory(at: cliProxyLogs, withIntermediateDirectories: true) + let proxyLog = """ + === REQUEST INFO === + URL: /v1/messages + Timestamp: \(environment.isoString(for: day)) + === HEADERS === + X-Claude-Code-Session-Id: session-openrouter + === REQUEST BODY === + {"model":"\(alias)"} + === API RESPONSE === + """ + try Data(proxyLog.utf8).write(to: cliProxyLogs.appendingPathComponent("request.log")) + CLIProxyAPIUsageCacheIO.merge( + [ + CLIProxyAPIUsageRecord( + timestamp: day, + provider: upstreamProvider, + executorType: executorType, + model: upstreamModel, + alias: alias, + endpoint: "/v1/messages", + authType: "api_key", + requestID: "cliproxy-openrouter-\(upstreamModel)", + tokens: .init(input: 100, output: 10, total: 110)), + ], + cacheRoot: environment.cacheRoot, + now: day) + let options = CostUsageScanner.Options( + claudeProjectsRoots: [environment.claudeProjectsRoot], + cacheRoot: environment.cacheRoot, + cliProxyAPIHome: cliProxyHome) + let refreshedCatalog = Data(""" + { + "openai": { + "id": "openai", + "models": { "gpt-new": { "id": "gpt-new", "cost": { "input": 2, "output": 8 } } } + }, + "anthropic": { + "id": "anthropic", + "models": { "claude-new": { "id": "claude-new", "cost": { "input": 3, "output": 15 } } } + }, + "google": { + "id": "google", + "models": { "gemma-new": { "id": "gemma-new", "cost": { "input": 2.5, "output": 11 } } } + } + } + """.utf8) + + let snapshot = try await CostUsageFetcher.loadTokenSnapshot( + provider: .claude, + now: day, + refreshPricingInBackground: false, + includePiSessions: false, + scannerOptions: options, + modelsDevClient: ModelsDevClient(transport: CostUsageFetcherModelsDevTransport( + data: refreshedCatalog))) + + let breakdown = try #require(snapshot.daily.first?.modelBreakdowns?.first) + #expect(breakdown.modelName == alias) + #expect(breakdown.attribution?.upstream?.provider == upstreamProvider) + #expect(breakdown.attribution?.upstream?.model == upstreamModel) + #expect(abs((breakdown.costUSD ?? 0) - expectedCost) < 0.0000001) + } } private struct UnknownModelPricingFixture { diff --git a/Tests/CodexBarTests/CostUsageScannerClaudeFableTests.swift b/Tests/CodexBarTests/CostUsageScannerClaudeFableTests.swift index 0d8beb4b75..ddecf2bff6 100644 --- a/Tests/CodexBarTests/CostUsageScannerClaudeFableTests.swift +++ b/Tests/CodexBarTests/CostUsageScannerClaudeFableTests.swift @@ -43,6 +43,136 @@ struct CostUsageScannerClaudeFableTests { #expect(abs((Double(parsed.rows[0].costNanos) / 1_000_000_000) - expected) < 0.000000001) } + @Test + func `claude proxy google upstream uses models dev pricing`() throws { + let env = try CostUsageTestEnvironment() + defer { env.cleanup() } + + let day = try env.makeLocalNoon(year: 2026, month: 6, day: 9) + let fileURL = try env.writeClaudeProjectFile( + relativePath: "project-a/proxy-google.jsonl", + contents: env.jsonl([ + [ + "message": [ + "model": "proxy-gemini-alias", + "id": "msg_proxy_google", + "type": "message", + "role": "assistant", + "usage": [ + "input_tokens": 100, + "cache_creation_input_tokens": 10, + "cache_read_input_tokens": 20, + "output_tokens": 5, + ], + ], + "requestId": "req_proxy_google", + "type": "assistant", + "timestamp": env.isoString(for: day), + "sessionId": "session_proxy_google", + ], + ])) + let resolver = CLIProxyAPIAttributionResolver( + observations: [ + .init(sessionID: "session_proxy_google", model: "proxy-gemini-alias", timestamp: day), + ], + usageRecords: [ + CLIProxyAPIUsageRecord( + timestamp: day, + provider: "google", + executorType: "GeminiExecutor", + model: "gemini-test-pro", + alias: "proxy-gemini-alias", + endpoint: "POST /v1/messages", + authType: "oauth", + requestID: "req_proxy_google", + tokens: .init( + input: 100, + output: 5, + cacheRead: 20, + cacheCreation: 10, + total: 135)), + ]) + let parsed = try CostUsageScanner.parseClaudeFile( + fileURL: fileURL, + range: CostUsageScanner.CostUsageDayRange(since: day, until: day), + providerFilter: .all, + attributionResolver: resolver, + modelsDevCatalog: Self.googleModelsDevCatalog(model: "gemini-test-pro")) + + let row = try #require(parsed.rows.first) + #expect(row.attribution?.route == .cliProxyAPI) + #expect(row.attribution?.upstream?.provider == "google") + #expect(row.attribution?.upstream?.model == "gemini-test-pro") + let expected = 0.000279 + #expect(abs((Double(row.costNanos) / 1_000_000_000) - expected) < 0.000000001) + #expect(row.costPriced == true) + } + + @Test + func `claude cached proxy row reprices when reconciliation discovers google upstream`() throws { + let env = try CostUsageTestEnvironment() + defer { env.cleanup() } + + let day = try env.makeLocalNoon(year: 2026, month: 6, day: 9) + let cachedRow = CostUsageScanner.ClaudeUsageRow( + dayKey: "2026-06-09", + model: "proxy-gemini-alias", + sessionId: "session_proxy_google_cached", + messageId: "msg_proxy_google_cached", + requestId: "req_proxy_google_cached", + timestampUnixMs: Int64(day.timeIntervalSince1970 * 1000), + isSidechain: false, + pathRole: .parent, + input: 100, + cacheRead: 20, + cacheCreate: 10, + cacheCreate1h: nil, + output: 5, + costNanos: 1_000_000, + costPriced: true, + attribution: nil) + var cache = CostUsageCache() + cache.files["cached-proxy-google.jsonl"] = CostUsageFileUsage( + mtimeUnixMs: 0, + size: 0, + days: [:], + claudeRows: [cachedRow]) + let resolver = CLIProxyAPIAttributionResolver( + observations: [ + .init(sessionID: "session_proxy_google_cached", model: cachedRow.model, timestamp: day), + ], + usageRecords: [ + CLIProxyAPIUsageRecord( + timestamp: day, + provider: "google", + executorType: "GeminiExecutor", + model: "gemini-test-pro", + alias: cachedRow.model, + endpoint: "POST /v1/messages", + authType: "oauth", + requestID: cachedRow.requestId ?? "", + tokens: .init( + input: cachedRow.input, + output: cachedRow.output, + cacheRead: cachedRow.cacheRead, + cacheCreation: cachedRow.cacheCreate, + total: cachedRow.input + cachedRow.output + cachedRow.cacheRead + cachedRow.cacheCreate)), + ]) + + let report = try CostUsageScanner.buildClaudeReportFromCache( + cache: cache, + range: .init(since: day, until: day), + attributionResolver: resolver, + modelsDevCatalog: Self.googleModelsDevCatalog(model: "gemini-test-pro")) + + let expected = 0.000279 + #expect(abs((report.summary?.totalCostUSD ?? 0) - expected) < 0.000000001) + let breakdown = try #require(report.data.first?.modelBreakdowns?.first) + #expect(breakdown.attribution?.upstream?.provider == "google") + #expect(breakdown.attribution?.upstream?.model == "gemini-test-pro") + #expect(abs((breakdown.costUSD ?? 0) - expected) < 0.000000001) + } + @Test func `claude transcript refusal remains priced without billing provenance`() throws { let env = try CostUsageTestEnvironment() @@ -379,6 +509,28 @@ struct CostUsageScannerClaudeFableTests { return try JSONDecoder().decode(ModelsDevCatalog.self, from: Data(json.utf8)) } + private static func googleModelsDevCatalog(model: String) throws -> ModelsDevCatalog { + let json = """ + { + "google": { + "id": "google", + "models": { + "\(model)": { + "id": "\(model)", + "cost": { + "input": 2, + "output": 10, + "cache_read": 0.2, + "cache_write": 2.5 + } + } + } + } + } + """ + return try JSONDecoder().decode(ModelsDevCatalog.self, from: Data(json.utf8)) + } + private static func anthropicThresholdModelsDevCatalog(model: String) throws -> ModelsDevCatalog { let json = """ { diff --git a/Tests/CodexBarTests/ShareStatsTests.swift b/Tests/CodexBarTests/ShareStatsTests.swift index a816eefa2a..401f11f6dd 100644 --- a/Tests/CodexBarTests/ShareStatsTests.swift +++ b/Tests/CodexBarTests/ShareStatsTests.swift @@ -5,6 +5,123 @@ import Testing @testable import CodexBar struct ShareStatsTests { + @Test + func `proxy spend does not inherit an ambient codex subscription`() { + #expect(!spendDashboardShouldUseAmbientCodexSubscription( + rowID: SpendDashboardSource.codexProxySourceID, + codexRowCount: 1)) + #expect(spendDashboardShouldUseAmbientCodexSubscription( + rowID: "codex:managed-account", + codexRowCount: 1)) + #expect(!spendDashboardShouldUseAmbientCodexSubscription( + rowID: "codex:managed-account", + codexRowCount: 2)) + } + + @Test + func `proxy spend is not counted as an account or subscription`() throws { + let rows = [ + SpendDashboardModel.ProviderRow( + id: "codex:managed-account", + rank: 1, + provider: .codex, + displayName: "Codex", + totalTokens: 100, + totalCost: 1, + coveredDayCount: 1), + SpendDashboardModel.ProviderRow( + id: SpendDashboardSource.codexProxySourceID, + rank: 2, + provider: .codex, + displayName: "Codex · CLIProxyAPI", + totalTokens: 100, + totalCost: 1, + coveredDayCount: 1), + SpendDashboardModel.ProviderRow( + id: "cursor", + rank: 3, + provider: .cursor, + displayName: "Cursor", + totalTokens: 100, + totalCost: 1, + coveredDayCount: 1), + ] + + #expect(spendDashboardCodexAccountRowCount(rows) == 1) + #expect(spendDashboardSubscriptionCount(rows) == 2) + + let group = SpendDashboardModel.CurrencyGroup( + currencyCode: "USD", + providers: rows, + models: [ + SpendDashboardModel.ModelRow( + rank: 1, + provider: .codex, + providerName: "Codex · CLIProxyAPI", + modelName: "gpt-5.4", + totalTokens: 100, + totalCost: 1), + ], + dailyPoints: [], + totalTokens: 300, + totalCost: 3, + coveredDayCount: 1, + chartDomain: Self.date...Self.date, + modelHistoryCompleteness: .complete) + let payload = try #require(ShareStatsBuilder.make( + model: SpendDashboardModel(requestedDays: 1, groups: [group]))) + + #expect(payload.providers.map(\.providerName) == ["Codex", "Cursor"]) + #expect(payload.currencies.first?.estimatedCost == 3) + #expect(payload.topModels.first?.estimatedCost == 1) + } + + @Test + func `proxy attributed models share under their verified upstream provider`() throws { + let attribution = CostUsageAttribution( + client: .claudeCode, + route: .cliProxyAPI, + modelProvider: .google, + upstream: .init(provider: "gemini", authType: .oauth), + evidence: [.cliProxyRequestLog, .cliProxyUsageTelemetry, .modelProvider]) + let group = SpendDashboardModel.CurrencyGroup( + currencyCode: "USD", + providers: [ + SpendDashboardModel.ProviderRow( + id: "claude", + rank: 1, + provider: .claude, + displayName: "Claude", + totalTokens: 100, + totalCost: 1, + coveredDayCount: 1), + ], + models: [ + SpendDashboardModel.ModelRow( + rank: 1, + provider: .claude, + providerName: "Claude", + modelName: "gemini-3-pro", + totalTokens: 100, + totalCost: 1, + attribution: attribution), + ], + dailyPoints: [], + totalTokens: 100, + totalCost: 1, + coveredDayCount: 1, + chartDomain: Self.date...Self.date, + modelHistoryCompleteness: .complete) + + let payload = try #require(ShareStatsBuilder.make( + model: SpendDashboardModel(requestedDays: 1, groups: [group]))) + let model = try #require(payload.topModels.first) + + #expect(model.provider == .gemini) + #expect(model.providerName == "Gemini") + #expect(model.modelName == "Gemini") + } + @Test func `builder preserves native currencies and unavailable spend`() throws { let subscriptionNames = try [ diff --git a/Tests/CodexBarTests/SpendDashboardCodexProxySourceTests.swift b/Tests/CodexBarTests/SpendDashboardCodexProxySourceTests.swift new file mode 100644 index 0000000000..9fb5d3c586 --- /dev/null +++ b/Tests/CodexBarTests/SpendDashboardCodexProxySourceTests.swift @@ -0,0 +1,142 @@ +import CodexBarCore +import Foundation +import Testing +@testable import CodexBar + +@MainActor +struct SpendDashboardCodexProxySourceTests { + @Test + func `proxy usage loads once beside account scoped codex snapshots`() async { + let now = Date(timeIntervalSince1970: 1_784_179_200) + let accounts = ["first", "second"].map { id in + CodexSpendScanRequest( + id: id, + displayName: "Codex · \(id)", + source: .profileHome(path: "/synthetic/\(id)"), + homePath: "/synthetic/\(id)", + authFingerprint: nil, + authFileWasReadable: false, + cacheIdentity: "\(id)-cache") + } + let request = SpendDashboardLoadRequest( + configuration: SpendDashboardConfiguration( + costUsageEnabled: true, + providerIDs: [UsageProvider.codex.rawValue], + codexAccountIdentities: accounts.map { "\($0.id)|\($0.cacheIdentity)" }), + capturedInputs: [], + unavailableSourceIDs: [], + codexRequests: accounts, + now: now, + force: false) + let proxyRecorder = SpendDashboardCodexProxyLoadRecorder() + let accountSnapshot = Self.snapshot(cost: 1, now: now) + let proxySnapshot = Self.snapshot(cost: 2, now: now) + + let result = await SpendDashboardSource.load( + request, + codexSnapshotLoader: { _ in accountSnapshot }, + codexProxySnapshotLoader: { context in + await proxyRecorder.record(context) + return proxySnapshot + }) + let proxyContexts = await proxyRecorder.contexts + + #expect(Set(result.inputs.map(\.id)) == [ + "codex:first", + "codex:second", + SpendDashboardSource.codexProxySourceID, + ]) + #expect(result.inputs.count { $0.id == SpendDashboardSource.codexProxySourceID } == 1) + #expect(result.inputs.first { $0.id == SpendDashboardSource.codexProxySourceID }?.displayName == + "Codex · CLIProxyAPI") + #expect(proxyContexts.count == 1) + #expect(proxyContexts.first?.now == now) + } + + @Test + func `proxy usage loads when claude is enabled without codex`() async { + let now = Date(timeIntervalSince1970: 1_784_179_200) + let request = SpendDashboardLoadRequest( + configuration: SpendDashboardConfiguration( + costUsageEnabled: true, + providerIDs: [UsageProvider.claude.rawValue], + codexAccountIdentities: []), + capturedInputs: [], + unavailableSourceIDs: [], + codexRequests: [], + now: now, + force: false) + let proxySnapshot = Self.snapshot(cost: 2, now: now) + let emptySnapshot = Self.snapshot(cost: 0, now: now) + + let result = await SpendDashboardSource.load( + request, + codexSnapshotLoader: { _ in + Issue.record("No account-scoped Codex snapshot should be requested.") + return emptySnapshot + }, + codexProxySnapshotLoader: { _ in proxySnapshot }) + + #expect(result.inputs.map(\.id) == [SpendDashboardSource.codexProxySourceID]) + } + + @Test + func `cancelled proxy load invalidates retained proxy source`() async { + let now = Date(timeIntervalSince1970: 1_784_179_200) + let staleProxyInput = SpendDashboardModel.ProviderInput( + id: SpendDashboardSource.codexProxySourceID, + provider: .codex, + displayName: "Codex · CLIProxyAPI", + modelProviderName: "Codex", + snapshot: Self.snapshot(cost: 2, now: now)) + let request = SpendDashboardLoadRequest( + configuration: SpendDashboardConfiguration( + costUsageEnabled: true, + providerIDs: [UsageProvider.claude.rawValue], + codexAccountIdentities: []), + capturedInputs: [staleProxyInput], + unavailableSourceIDs: [], + codexRequests: [], + now: now, + force: false) + let emptySnapshot = Self.snapshot(cost: 0, now: now) + + let result = await SpendDashboardSource.load( + request, + codexSnapshotLoader: { _ in + Issue.record("No account-scoped Codex snapshot should be requested.") + return emptySnapshot + }, + codexProxySnapshotLoader: { _ in throw CancellationError() }) + + #expect(result.inputs.isEmpty) + #expect(result.failedSourceIDs == [SpendDashboardSource.codexProxySourceID]) + #expect(result.invalidatedSourceIDs == [SpendDashboardSource.codexProxySourceID]) + } + + private static func snapshot(cost: Double, now: Date) -> CostUsageTokenSnapshot { + let entry = CostUsageDailyReport.Entry( + date: "2026-07-15", + inputTokens: nil, + outputTokens: nil, + totalTokens: 10, + costUSD: cost, + modelsUsed: ["gpt-5.5"], + modelBreakdowns: nil) + return CostUsageTokenSnapshot( + sessionTokens: 10, + sessionCostUSD: cost, + last30DaysTokens: 10, + last30DaysCostUSD: cost, + daily: [entry], + updatedAt: now) + } +} + +private actor SpendDashboardCodexProxyLoadRecorder { + private(set) var contexts: [CodexProxySpendSnapshotLoadContext] = [] + + func record(_ context: CodexProxySpendSnapshotLoadContext) { + self.contexts.append(context) + } +} diff --git a/Tests/CodexBarTests/SpendDashboardForceStateMachineTests.swift b/Tests/CodexBarTests/SpendDashboardForceStateMachineTests.swift index e0c1f152c6..d9b8a5e3db 100644 --- a/Tests/CodexBarTests/SpendDashboardForceStateMachineTests.swift +++ b/Tests/CodexBarTests/SpendDashboardForceStateMachineTests.swift @@ -92,6 +92,42 @@ struct SpendDashboardForceStateMachineTests { #expect(controller.model.groups.first?.totalCost == 12) } + @Test + func `forced proxy success carries through a Claude only capture barrier`() async { + let configuration = SpendDashboardConfiguration( + costUsageEnabled: true, + providerIDs: [UsageProvider.claude.rawValue], + codexAccountIdentities: []) + let builder = SpendDashboardBuildScript([ + .init( + mode: .forceRefresh, + request: Self.request(configuration, mode: .forceRefresh)), + .init( + mode: .captureOnly, + request: Self.request(configuration, mode: .captureOnly)), + ]) + let loader = SpendDashboardStateLoaderGate() + let controller = SpendDashboardController( + requestBuilder: { mode in await builder.next(mode) }, + loader: { request in await loader.load(request) }) + + controller.update(configuration: configuration, force: true) + await Self.waitForLoader(loader) + await loader.resume(SpendDashboardLoadResult( + inputs: [Self.input( + id: SpendDashboardSource.codexProxySourceID, + provider: .codex, + cost: 5)], + failedSourceIDs: [])) + await Self.waitUntil { !controller.isRefreshing } + + #expect(builder.modes == [.forceRefresh, .captureOnly]) + #expect(await loader.forces == [true]) + #expect(controller.model.groups.first?.totalCost == 5) + #expect(controller.model.groups.flatMap(\.providers).map(\.id) == + [SpendDashboardSource.codexProxySourceID]) + } + @Test func `C same owner barrier churn repeats capture only and preserves failures`() async { let initial = Self.configuration(owner: "owner", revision: "R") diff --git a/Tests/CodexBarTests/SpendDashboardProxyAttributionTests.swift b/Tests/CodexBarTests/SpendDashboardProxyAttributionTests.swift new file mode 100644 index 0000000000..3292f3d077 --- /dev/null +++ b/Tests/CodexBarTests/SpendDashboardProxyAttributionTests.swift @@ -0,0 +1,200 @@ +import CodexBarCore +import Foundation +import Testing +@testable import CodexBar + +struct SpendDashboardProxyAttributionTests { + @Test + func `missing proxy configuration clears the saved presentation`() { + let presentation = spendDashboardCLIProxyAPIConfigurationPresentation( + loadResult: .missing, + currentBaseURL: "http://localhost:8317", + hasSavedConfiguration: true) + + #expect(presentation.baseURL == "http://localhost:8317") + #expect(!presentation.hasSavedConfiguration) + } + + @Test + func `invalid proxy configuration clears the saved presentation`() { + let presentation = spendDashboardCLIProxyAPIConfigurationPresentation( + loadResult: .invalid, + currentBaseURL: "http://localhost:8317", + hasSavedConfiguration: true) + + #expect(presentation.baseURL == "http://localhost:8317") + #expect(!presentation.hasSavedConfiguration) + } + + @Test + func `temporarily unavailable proxy configuration preserves the presentation`() { + let presentation = spendDashboardCLIProxyAPIConfigurationPresentation( + loadResult: .temporarilyUnavailable, + currentBaseURL: "http://localhost:8317", + hasSavedConfiguration: true) + + #expect(presentation.baseURL == "http://localhost:8317") + #expect(presentation.hasSavedConfiguration) + } + + @Test + func `found proxy configuration refreshes the presentation`() { + let presentation = spendDashboardCLIProxyAPIConfigurationPresentation( + loadResult: .found(CLIProxyAPIConnectionSettings( + baseURL: "http://127.0.0.1:8317", + managementKey: "test-key")), + currentBaseURL: CLIProxyAPIConnectionSettings.defaultBaseURL, + hasSavedConfiguration: false) + + #expect(presentation.baseURL == "http://127.0.0.1:8317") + #expect(presentation.hasSavedConfiguration) + } + + @Test + func `unresolved route describes known facts without an unknown warning`() { + let attribution = CostUsageAttribution( + client: .claudeCode, + route: .unknown, + modelProvider: .openAI, + evidence: [.modelProvider]) + + #expect(spendDashboardModelSourceText( + providerName: "Claude", + attribution: attribution) == "Claude · OpenAI model via Claude Code") + } + + @Test + func `confirmed proxy route without upstream telemetry does not infer a backend`() { + let attribution = CostUsageAttribution( + client: .claudeCode, + route: .cliProxyAPI, + modelProvider: .openAI, + evidence: [.cliProxyRequestLog, .modelProvider]) + + #expect(spendDashboardModelSourceText( + providerName: "Claude", + attribution: attribution) == "CLIProxyAPI via Claude Code") + } + + @Test + func `proxy attribution survives dashboard aggregation and describes the route`() throws { + let attribution = CostUsageAttribution( + client: .claudeCode, + route: .cliProxyAPI, + modelProvider: .openAI, + upstream: .init(provider: "codex", authType: .oauth), + evidence: [.cliProxyRequestLog, .cliProxyUsageTelemetry, .modelProvider]) + let entry = CostUsageDailyReport.Entry( + date: "2026-07-16", + inputTokens: 90, + outputTokens: 10, + totalTokens: 100, + costUSD: 1, + modelsUsed: ["gpt-5.5"], + modelBreakdowns: [ + .init( + modelName: "gpt-5.5", + costUSD: 1, + totalTokens: 100, + attribution: attribution), + ]) + let now = Date(timeIntervalSince1970: 1_784_179_200) + let snapshot = CostUsageTokenSnapshot( + sessionTokens: 100, + sessionCostUSD: 1, + last30DaysTokens: 100, + last30DaysCostUSD: 1, + daily: [entry], + updatedAt: now) + var calendar = Calendar(identifier: .gregorian) + calendar.timeZone = try #require(TimeZone(secondsFromGMT: 0)) + + let model = SpendDashboardModel.build( + inputs: [ + .init( + provider: .codex, + displayName: "Codex", + snapshot: snapshot), + ], + requestedDays: 7, + now: now, + calendar: calendar) + + let row = try #require(model.groups.first?.models.first) + #expect(row.provider == .codex) + #expect(row.attribution == attribution) + #expect(spendDashboardModelSourceText( + providerName: row.providerName, + attribution: row.attribution) == "Codex OAuth · CLIProxyAPI via Claude Code") + } + + @Test + func `model row identity includes complete proxy attribution`() throws { + let inventoryAttribution = CostUsageAttribution( + client: .claudeCode, + route: .cliProxyAPI, + modelProvider: .openAI, + upstream: .init( + provider: "codex", + authType: .oauth, + model: "gpt-5.5"), + evidence: [.cliProxyAuthInventory, .cliProxyRequestLog, .modelProvider]) + let telemetryAttribution = CostUsageAttribution( + client: .claudeCode, + route: .cliProxyAPI, + modelProvider: .openAI, + upstream: .init( + provider: "codex", + authType: .oauth, + model: "openai/gpt-5.5", + executorType: "CodexExecutor"), + evidence: [.cliProxyRequestLog, .cliProxyUsageTelemetry, .modelProvider]) + let entry = CostUsageDailyReport.Entry( + date: "2026-07-16", + inputTokens: 100, + outputTokens: 100, + totalTokens: 200, + costUSD: 2, + modelsUsed: ["gpt-5.5"], + modelBreakdowns: [ + .init( + modelName: "gpt-5.5", + costUSD: 1, + totalTokens: 100, + attribution: inventoryAttribution), + .init( + modelName: "gpt-5.5", + costUSD: 1, + totalTokens: 100, + attribution: telemetryAttribution), + ]) + let now = Date(timeIntervalSince1970: 1_784_179_200) + let snapshot = CostUsageTokenSnapshot( + sessionTokens: 200, + sessionCostUSD: 2, + last30DaysTokens: 200, + last30DaysCostUSD: 2, + daily: [entry], + updatedAt: now) + var calendar = Calendar(identifier: .gregorian) + calendar.timeZone = try #require(TimeZone(secondsFromGMT: 0)) + + let model = SpendDashboardModel.build( + inputs: [ + .init( + provider: .codex, + displayName: "Codex", + snapshot: snapshot), + ], + requestedDays: 7, + now: now, + calendar: calendar) + + let group = try #require(model.groups.first) + let rows = group.models + #expect(rows.count == 2) + #expect(Set(rows.map(\.id)).count == 2) + #expect(rows.map(\.attribution) == [telemetryAttribution, inventoryAttribution]) + #expect(rows.map(\.rank) == [1, 2]) + } +} diff --git a/Tests/CodexBarTests/UsageStoreCachedTokenHydrationTests.swift b/Tests/CodexBarTests/UsageStoreCachedTokenHydrationTests.swift index c120e04079..e5d514f394 100644 --- a/Tests/CodexBarTests/UsageStoreCachedTokenHydrationTests.swift +++ b/Tests/CodexBarTests/UsageStoreCachedTokenHydrationTests.swift @@ -278,6 +278,41 @@ struct UsageStoreCachedTokenHydrationTests { #expect(store.tokenLastAttemptAt(for: .codex) == nil) } + @Test + func `cache clear wins over in flight cached codex hydration`() async throws { + let env = try CostUsageTestEnvironment() + defer { env.cleanup() } + let settings = Self.makeCodexOnlySettings(historyDays: 1) + let options = CostUsageScanner.Options(cacheRoot: env.cacheRoot) + let store = UsageStore( + fetcher: UsageFetcher(), + browserDetection: BrowserDetection(cacheTTL: 0), + costUsageFetcher: CostUsageFetcher(scannerOptions: options), + settings: settings, + startupBehavior: .testing, + environmentBase: [:]) + let gate = CachedTokenHydrationGate() + store._test_cachedCodexTokenSnapshotLoaderOverride = { _, _, _ in + await gate.enter() + return (Self.cachedTokenSnapshot(), Date(), nil) + } + + let hydration = store.hydrateCachedTokenSnapshots() + await gate.waitForStart() + let artifactDirectory = env.cacheRoot.appendingPathComponent("cost-usage", isDirectory: true) + try FileManager.default.createDirectory(at: artifactDirectory, withIntermediateDirectories: true) + let clearResult = CostUsageCacheLocations.clearAllCostUsageCaches( + in: [artifactDirectory], + stateRoot: env.cacheRoot, + fileManager: .default) + await gate.release() + await hydration?.value + + #expect(clearResult.errorDescription == nil) + #expect(store.tokenSnapshot(for: .codex) == nil) + #expect(store.tokenLastAttemptAt(for: .codex) == nil) + } + private static func makeCodexOnlySettings(historyDays: Int) -> SettingsStore { let suite = "UsageStoreCachedTokenHydrationTests-\(UUID().uuidString)" let defaults = UserDefaults(suiteName: suite)! diff --git a/Tests/CodexBarTests/UsageStoreCoverageTests.swift b/Tests/CodexBarTests/UsageStoreCoverageTests.swift index 85af8b03c9..134f874afe 100644 --- a/Tests/CodexBarTests/UsageStoreCoverageTests.swift +++ b/Tests/CodexBarTests/UsageStoreCoverageTests.swift @@ -163,14 +163,12 @@ struct UsageStoreCoverageTests { historyDays: 30, source: .auto, credentialFingerprint: "unresolved") - let revision = store.providerPublicationRevision(for: .cursor) - let providerConfigRevision = settings.providerConfigRevision(for: .cursor) + let publicationGuard = store.tokenRefreshPublicationGuard(for: .cursor) settings.costUsageHistoryDays = 7 #expect(!store.tokenRefreshPublicationIsCurrent( provider: .cursor, - publicationRevision: revision, - providerConfigRevision: providerConfigRevision, + publicationGuard: publicationGuard, historyDays: 30, costScopeSignature: initialSignature, fetchedCredentialScopeFingerprint: fingerprint))