diff --git a/Sources/CodexBar/CostHistoryChartMenuView.swift b/Sources/CodexBar/CostHistoryChartMenuView.swift index 210d0a3f52..8efa27dd56 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 @@ -573,7 +582,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 { @@ -797,26 +806,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) } @@ -839,6 +828,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 { @@ -876,6 +878,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 @@ -897,6 +925,7 @@ extension CostHistoryChartMenuView { struct VisibleModelBreakdownFingerprint: Equatable { let modelName: String + let attribution: CostUsageAttribution? let costBitPattern: UInt64? let totalTokens: Int? let standardCostBitPattern: UInt64? @@ -976,6 +1005,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), @@ -995,6 +1025,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 83e7f31038..817b3ae48e 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" + } +} + enum SpendDashboardModelHistoryPresentation: Equatable { case unavailable case empty @@ -48,6 +91,11 @@ 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 init(settings: SettingsStore, store: UsageStore) { self.settings = settings @@ -62,6 +110,7 @@ struct SpendDashboardPane: View { VStack(alignment: .leading, spacing: 18) { self.header self.content + self.cliProxyAPISetup self.provenance self.shareAction } @@ -71,6 +120,7 @@ struct SpendDashboardPane: View { .onAppear { self.controller.refreshDateWindow() self.controller.update(configuration: self.configuration) + self.loadCLIProxyAPIConfiguration() } .onChange(of: self.configuration) { _, configuration in self.controller.update(configuration: configuration) @@ -173,6 +223,150 @@ 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() { + guard let configuration = CLIProxyAPIConnectionSettingsStore.load() else { return } + self.cliProxyAPIBaseURL = configuration.baseURL + self.cliProxyAPIHasSavedConfiguration = true + } + + 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 + } + guard CLIProxyAPIConnectionSettingsStore.save(configuration) else { + self.cliProxyAPIStatus = "Could not save the management key." + return + } + self.store.stopCLIProxyAPIUsageCollector() + + 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)" + } + 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() @@ -194,9 +388,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, @@ -206,7 +399,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)] @@ -263,7 +460,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() } } @@ -366,7 +564,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 744f41ee17..630d32cd36 100644 --- a/Sources/CodexBar/ShareStatsPayload.swift +++ b/Sources/CodexBar/ShareStatsPayload.swift @@ -318,8 +318,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, diff --git a/Sources/CodexBar/SpendDashboardController.swift b/Sources/CodexBar/SpendDashboardController.swift index 8b7de621e3..5d502c74eb 100644 --- a/Sources/CodexBar/SpendDashboardController.swift +++ b/Sources/CodexBar/SpendDashboardController.swift @@ -117,11 +117,28 @@ 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 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 { @@ -143,12 +160,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, @@ -250,14 +272,33 @@ 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 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 @@ -303,6 +344,36 @@ 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) + return SpendDashboardLoadResult( + inputs: [], + failedSourceIDs: failedSourceIDs, + invalidatedSourceIDs: invalidatedSourceIDs) + } catch { + failedSourceIDs.insert(Self.codexProxySourceID) + } + } let lateInvalidatedSourceIDs = Set(request.codexRequests.compactMap { account in self.currentAuthFingerprint(for: account) == account.authFingerprint ? nil @@ -328,7 +399,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 @@ -338,6 +420,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 @@ -398,6 +485,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) @@ -902,7 +999,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..1bd874672b --- /dev/null +++ b/Sources/CodexBar/UsageStore+CLIProxyAPI.swift @@ -0,0 +1,68 @@ +import CodexBarCore +import Foundation + +enum CLIProxyAPIConfigurationRemovalResult: Equatable { + case removed + case configurationRemovalFailed + case telemetryCleanupFailed +} + +@MainActor +extension UsageStore { + private static let cliProxyAPIUsageCollectionInterval: Duration = .seconds(30) + + func startCLIProxyAPIUsageCollector() { + self.stopCLIProxyAPIUsageCollector() + self.cliProxyAPIUsageCollectorTask = Task.detached(priority: .utility) { [weak self] in + while !Task.isCancelled { + guard await self?.collectCLIProxyAPIUsageNow() != nil else { return } + 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( + purgeTelemetry: (() -> Bool)? = nil, + clear: () -> Bool = { CLIProxyAPIConnectionSettingsStore.clear() }) async + -> CLIProxyAPIConfigurationRemovalResult + { + let collectorTask = self.stopCLIProxyAPIUsageCollector() + await collectorTask?.value + guard clear() else { return .configurationRemovalFailed } + let didPurgeTelemetry: Bool = if let purgeTelemetry { + purgeTelemetry() + } else { + await Task.detached(priority: .utility) { + CostUsageCacheLocations.clearCLIProxyAPIArtifacts() + }.value + } + guard didPurgeTelemetry else { return .telemetryCleanupFailed } + return .removed + } + + 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 + }) + } +} diff --git a/Sources/CodexBar/UsageStore+TokenCost.swift b/Sources/CodexBar/UsageStore+TokenCost.swift index a55a14e72c..d42dd91569 100644 --- a/Sources/CodexBar/UsageStore+TokenCost.swift +++ b/Sources/CodexBar/UsageStore+TokenCost.swift @@ -428,31 +428,32 @@ 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 -> String?)? = nil, + fileManager: FileManager = .default) async -> String? + { + let collectorTask = self.stopCLIProxyAPIUsageCollector() + await collectorTask?.value + defer { + if collectorTask != nil { + self.startCLIProxyAPIUsageCollector() } - return nil - }.value + } + + let cacheDirectories = CostUsageCacheLocations.directories(fileManager: fileManager) + let cliProxyAPIStateRoot = cacheDirectories[1].deletingLastPathComponent() + let errorMessage: String? = if let clearDirectories { + await clearDirectories() + } else { + await Task.detached(priority: .utility) { + CostUsageCacheLocations.clearAllCostUsageCaches( + in: cacheDirectories, + stateRoot: cliProxyAPIStateRoot).errorDescription + }.value + } guard errorMessage == nil else { return errorMessage } diff --git a/Sources/CodexBar/UsageStore.swift b/Sources/CodexBar/UsageStore.swift index 19c9488717..77d71fe7b8 100644 --- a/Sources/CodexBar/UsageStore.swift +++ b/Sources/CodexBar/UsageStore.swift @@ -365,6 +365,7 @@ final class UsageStore { /// Background load task; cleared on deinit and on the cancel test seam. @ObservationIgnored var planUtilizationHistoryLoadTask: Task? + @ObservationIgnored var cliProxyAPIUsageCollectorTask: 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 @@ -480,6 +481,7 @@ final class UsageStore { Task { await self.refresh(enrichmentMode: .automatic) } self.startTimer() self.startTokenTimer() + self.startCLIProxyAPIUsageCollector() } var iconStyle: IconStyle { @@ -858,6 +860,7 @@ final class UsageStore { self.codexPlanHistoryBackfillTask?.cancel() self.resetBoundaryRefreshTask?.cancel() self.planUtilizationHistoryLoadTask?.cancel() + self.cliProxyAPIUsageCollectorTask?.cancel() } enum SessionQuotaWindowSource: String { 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..a4bd6fcfbf --- /dev/null +++ b/Sources/CodexBarCore/CLIProxyAPIAttributionResolver.swift @@ -0,0 +1,666 @@ +import Foundation + +struct CLIProxyAPIAttributionResolver: Sendable { + struct Observation: Sendable, Equatable { + let sessionID: String + let model: String + let timestamp: Date? + } + + 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 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 hasConfiguredOpenAIAPIUpstream: Bool + + init( + observations: [Observation], + usageRecords: [CLIProxyAPIUsageRecord] = [], + authProviders: [AuthProvider] = [], + 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.hasConfiguredOpenAIAPIUpstream = hasConfiguredOpenAIAPIUpstream + } + + static func load( + home: URL, + cacheRoot: URL? = nil, + fileManager: FileManager = .default, + forceReload: Bool = false, + 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: CLIProxyAPIUsageCacheIO.load(cacheRoot: cacheRoot), + authProviders: self.loadAuthProviders(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] { + let prepared = requests.map { request in + let routeObservation = self.matchingObservation( + model: request.model, + sessionID: request.sessionID, + timestampUnixMs: request.timestampUnixMs) + let telemetryObservation = self.matchingObservation( + model: request.model, + sessionID: request.sessionID, + timestampUnixMs: request.timestampUnixMs) + let usageRecordMatch = telemetryObservation.flatMap { + self.closestUsageRecordMatch( + observation: $0, + model: request.model, + tokens: request.tokens) + } + return ( + request: request, + routeObservation: routeObservation, + telemetryObservation: telemetryObservation, + usageRecordMatch: usageRecordMatch) + } + 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)) + } + var routeOwnerByObservation: [ObservationKey: Int] = [:] + for (key, candidates) in routeCandidatesByObservation { + if candidates.count == 1, let candidate = candidates.first { + routeOwnerByObservation[key] = candidate.index + } else if let observationTimestamp = candidates.first?.observation.timestamp, + let candidate = Self.uniqueClosest( + candidates, + target: observationTimestamp, + timestamp: { Self.timestamp(for: $0.request) }) + { + routeOwnerByObservation[key] = candidate.index + } + } + let matchCounts = Dictionary( + grouping: prepared.compactMap(\.usageRecordMatch?.key), + by: { $0 }).mapValues(\.count) + 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) + } + } + + 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 inventoryUpstream = routeObservation != nil && usageRecord == nil + ? self.authInventoryUpstream(model: request.model, modelProvider: request.modelProvider) + : nil + let routeConfirmed = routeObservation != nil || inventoryUpstream != nil + var evidence: Set = [.modelProvider] + if routeObservation != nil || inventoryUpstream != nil { + evidence.insert(.cliProxyRequestLog) + } + if usageRecord != nil { + evidence.insert(.cliProxyUsageTelemetry) + } + if inventoryUpstream != nil { + evidence.insert(.cliProxyAuthInventory) + } + + return CostUsageAttribution( + client: .claudeCode, + route: routeConfirmed ? .cliProxyAPI : .unknown, + modelProvider: request.modelProvider, + upstream: usageRecord.map(Self.upstream) ?? inventoryUpstream, + evidence: evidence.sorted { $0.rawValue < $1.rawValue }) + } + + private func matchingObservation( + model: String, + sessionID: String?, + timestampUnixMs: Int64?) -> Observation? + { + guard let sessionID = sessionID?.trimmingCharacters(in: .whitespacesAndNewlines), + !sessionID.isEmpty, + let observations = self.observationsBySessionID[sessionID] + else { return nil } + + let canonicalModel = Self.canonicalModel(model) + let matchingModels = observations.filter { Self.canonicalModel($0.model) == canonicalModel } + guard !matchingModels.isEmpty else { return nil } + guard let timestampUnixMs else { + return matchingModels.count == 1 ? matchingModels[0] : nil + } + let timestamp = Date(timeIntervalSince1970: Double(timestampUnixMs) / 1000) + let candidates = matchingModels.filter { observation in + guard let observationTimestamp = observation.timestamp else { return true } + return abs(observationTimestamp.timeIntervalSince(timestamp)) <= Self.maximumRouteMatchDistance + } + return Self.uniqueClosest( + candidates, + target: timestamp, + timestamp: \.timestamp) + } + + 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( + 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 + && record.endpoint.lowercased().contains("/v1/messages") + { + 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) -> CostUsageAttribution.Upstream? + { + guard modelProvider == .openAI, + !self.observationsBySessionID.isEmpty, + !self.hasConfiguredOpenAIAPIUpstream + else { return nil } + let providers = Array(Set(self.authProviders)) + guard 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)) + } + + private static func hasConfiguredOpenAIAPIUpstream( + home: URL, + fileManager: FileManager) -> 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[.. 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..309d47da99 --- /dev/null +++ b/Sources/CodexBarCore/CLIProxyAPIUsageTelemetry.swift @@ -0,0 +1,840 @@ +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 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 failed + case generate + case tokens + } + + init( + timestamp: Date, + provider: String, + executorType: String? = nil, + model: String, + alias: String, + endpoint: String, + authType: String, + requestID: String, + 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.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.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) + } +} + +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 load( + cacheRoot: URL? = nil, + now: Date = Date()) -> [CLIProxyAPIUsageRecord] + { + let legacyCacheRoot = cacheRoot == nil ? self.defaultLegacyCacheRoot() : nil + return self.load( + cacheRoot: cacheRoot, + legacyCacheRoot: legacyCacheRoot, + now: now) + } + + 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: [String: CLIProxyAPIUsageRecord] = [:] + for record in existingCache.records where record.timestamp >= cutoff { + byKey[self.recordKey(record)] = record + } + let priorCount = byKey.count + for record in records where record.timestamp >= cutoff { + byKey[self.recordKey(record)] = 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 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 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: [String: CLIProxyAPIUsageRecord] = [:] + for record in legacy.records { + byKey[self.recordKey(record)] = record + } + for record in durable?.records ?? [] { + byKey[self.recordKey(record)] = 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 { + private static let key = KeychainCacheStore.Key( + category: "integration", + identifier: "cliproxyapi-management") + + public static func load() -> CLIProxyAPIConnectionSettings? { + self.load( + isDisconnected: { CostUsageCacheLocations.isCLIProxyAPIExplicitlyDisconnected() }, + loadStored: { + switch KeychainCacheStore.load(key: self.key, as: CLIProxyAPIConnectionSettings.self) { + case let .found(settings): settings + case .missing, .temporarilyUnavailable, .invalid: nil + } + }) + } + + static func load( + isDisconnected: () -> Bool, + loadStored: () -> CLIProxyAPIConnectionSettings?) -> CLIProxyAPIConnectionSettings? + { + guard !isDisconnected() else { return nil } + return loadStored() + } + + @discardableResult + public static func save(_ settings: CLIProxyAPIConnectionSettings) -> Bool { + self.save( + settings, + store: { KeychainCacheStore.storeResult(key: self.key, entry: $0) }, + clearDisconnectedState: { CostUsageCacheLocations.setCLIProxyAPIExplicitlyDisconnected(false) }, + rollback: { KeychainCacheStore.clear(key: self.key) }) + } + + static func save( + _ settings: CLIProxyAPIConnectionSettings, + store: (CLIProxyAPIConnectionSettings) -> Bool, + clearDisconnectedState: () -> Bool, + rollback: () -> Bool) -> Bool + { + guard settings.isConfigured else { return false } + guard store(settings) else { return false } + guard clearDisconnectedState() else { + _ = rollback() + return false + } + return true + } + + @discardableResult + public static func clear() -> Bool { + let wasDisconnected = CostUsageCacheLocations.isCLIProxyAPIExplicitlyDisconnected() + guard wasDisconnected || CostUsageCacheLocations.setCLIProxyAPIExplicitlyDisconnected(true) else { + return false + } + guard KeychainCacheStore.clear(key: self.key) else { + if !wasDisconnected { + _ = CostUsageCacheLocations.setCLIProxyAPIExplicitlyDisconnected(false) + } + return false + } + return true + } +} + +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() + + public static func collect( + cacheRoot: URL? = nil, + settings: CLIProxyAPIConnectionSettings? = CLIProxyAPIConnectionSettingsStore.load(), + 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 CLIProxyAPIConnectionSettingsStore.load() == settings + }, + 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..7079362ab2 --- /dev/null +++ b/Sources/CodexBarCore/CostUsageCacheLocations.swift @@ -0,0 +1,221 @@ +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 { + static let cliProxyAPIUsageFileName = "cliproxyapi-usage-v1.json" + static let cliProxyAPIPendingFileName = "cliproxyapi-pending-v1.json" + private static let cliProxyAPIDisconnectedFileName = "cliproxyapi-disconnected-v1" + + 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) + { + 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) } + 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) } + 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) + { + self.clearCLIProxyAPIArtifactsUnserialized( + in: directories, + fileManager: fileManager) + } + } catch { + return false + } + } + + private static func clearCLIProxyAPIArtifactsUnserialized( + in directories: [URL], + fileManager: FileManager) -> Bool + { + var succeeded = true + for directory in directories { + let urls = [ + directory.appendingPathComponent(self.cliProxyAPIUsageFileName, isDirectory: false), + directory.appendingPathComponent(self.cliProxyAPIPendingFileName, isDirectory: false), + CostUsageCacheIO.cacheFileURL( + provider: .claude, + cacheRoot: directory.deletingLastPathComponent()), + ] + for url in urls { + guard fileManager.fileExists(atPath: url.path) else { continue } + do { + try fileManager.removeItem(at: url) + } catch { + succeeded = false + } + } + } + return succeeded + } + + static func isCLIProxyAPIExplicitlyDisconnected( + stateRoot: URL? = nil, + fileManager: FileManager = .default) -> Bool + { + fileManager.fileExists(atPath: self.cliProxyAPIDisconnectedURL( + stateRoot: stateRoot, + fileManager: fileManager).path) + } + + @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 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 226d061364..e1931af495 100644 --- a/Sources/CodexBarCore/CostUsageFetcher.swift +++ b/Sources/CodexBarCore/CostUsageFetcher.swift @@ -36,7 +36,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) { @@ -117,7 +117,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, @@ -131,6 +132,7 @@ public struct CostUsageFetcher: Sendable { allowPricingRefresh: allowPricingRefresh, refreshPricingInBackground: refreshPricingInBackground, includePiSessions: includePiSessions, + includeClaudeProxyUsage: includeClaudeProxyUsage, bypassScannerDebounce: false, scannerOptions: self.scannerOptionsOverride()) } @@ -147,6 +149,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( @@ -161,10 +164,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, @@ -194,12 +233,15 @@ public struct CostUsageFetcher: Sendable { self.scannerOptions } - 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 @@ -210,6 +252,24 @@ 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) + } + + // swiftlint:disable:next function_body_length static func loadTokenSnapshot( provider: UsageProvider, environment: [String: String] = ProcessInfo.processInfo.environment, @@ -222,6 +282,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 @@ -249,10 +310,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, @@ -267,6 +332,7 @@ public struct CostUsageFetcher: Sendable { options.claudeLogProviderFilter = allowVertexClaudeFallback ? .all : .vertexAIOnly } else if provider == .claude { options.claudeLogProviderFilter = .excludeVertexAI + options.claudeAttributionFilter = cliProxyAPIAttributionEnabled ? .excludeCodexBackend : .all } if forceRefresh || bypassScannerDebounce { options.refreshMinIntervalSeconds = 0 @@ -317,24 +383,24 @@ public struct CostUsageFetcher: Sendable { var projects: [CostUsageProjectBreakdown] = [] var sessions: [CostUsageSessionBreakdown] = [] var piDaily: CostUsageDailyReport? + var claudeProxyDaily: CostUsageDailyReport? if provider == .codex { - let roots = CostUsageScanner.codexSessionsRoots(options: scanOptions) - let cache = CostUsageScanner.codexCache( - CostUsageCacheIO.load(provider: .codex, cacheRoot: scanOptions.cacheRoot), - scopedTo: roots) let range = CostUsageScanner.CostUsageDayRange( - since: since, until: now, calendar: scanOptions.calendar) - projects = CostUsageScanner.buildCodexProjectBreakdownsFromCache( - cache: cache, - range: range, - modelsDevCacheRoot: scanOptions.cacheRoot) - sessions = CostUsageScanner.buildCodexSessionBreakdownsFromCache( - cache: cache, + since: since, + until: now, + calendar: scanOptions.calendar) + let supplemental = try Self.loadCodexSupplementalScan( + options: scanOptions, range: range, - modelsDevCacheRoot: scanOptions.cacheRoot, - sessionRoots: roots) + now: now, + includeClaudeProxy: includeClaudeProxyUsage && shouldMergeGlobalCodexUsage, + checkCancellation: checkCancellation) + projects = supplemental.projects + sessions = supplemental.sessions + claudeProxyDaily = supplemental.claudeProxyDaily + daily = claudeProxyDaily.map { daily.merged(with: $0) } ?? daily } - if includePiSessions, provider == .claude || (provider == .codex && shouldMergePiUsage) { + if includePiSessions, provider == .claude || (provider == .codex && shouldMergeGlobalCodexUsage) { let piReport = try PiSessionCostScanner.loadDailyReportCancellable( provider: provider, since: since, @@ -349,11 +415,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 (daily: daily, projects: projects, sessions: sessions) } @@ -380,6 +448,7 @@ public struct CostUsageFetcher: Sendable { allowPricingRefresh: allowPricingRefresh, refreshPricingInBackground: false, includePiSessions: includePiSessions, + includeClaudeProxyUsage: includeClaudeProxyUsage, scannerOptions: options, piScannerOptions: piOptions, modelsDevClient: modelsDevClient, @@ -423,8 +492,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 @@ -438,48 +506,98 @@ 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 } - return breakdown.modelName - } ?? [] - }) - guard !unknownModelIDs.isEmpty else { return nil } + let providerIDs = Self.modelsDevProviderIDs( + pricingModel: pricingModel, + attribution: breakdown.attribution, + fallbackProvider: provider, + cacheRoot: cacheRoot) + for providerID in providerIDs { + modelIDsByProviderID[providerID, default: []].insert(pricingModel) + } + } + } + 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 { + 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( @@ -508,11 +626,16 @@ 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 = overrideScannerOptions ?? CostUsageScanner.Options() + let options = Self.resolvedScannerOptions( + overrideScannerOptions, + provider: .codex, + codexHomePath: nil) let until = now let since = options.calendar.date( byAdding: .day, @@ -534,6 +657,7 @@ public struct CostUsageFetcher: Sendable { var nativeScanAt: Date? var scanTimes: [Date] = [] var piMerged = false + var claudeProxyMerged = false if cache.timeZoneIdentifier == range.calendar.timeZone.identifier, !cache.days.isEmpty, @@ -565,6 +689,49 @@ public struct CostUsageFetcher: Sendable { } } + let claudeCache = CostUsageCacheIO.load( + provider: .claude, + cacheRoot: options.cacheRoot, + calendar: range.calendar) + if Self.isCLIProxyAPIAttributionEnabled(options: options), + !claudeCache.days.isEmpty, + !CostUsageScanner.requestedWindowExpandsCache(range: range, cache: claudeCache) + { + let attributionResolver: CLIProxyAPIAttributionResolver? = if let home = options.cliProxyAPIHome { + try CLIProxyAPIAttributionResolver.load( + home: home, + cacheRoot: options.cacheRoot, + forceReload: options.forceRescan, + checkCancellation: checkCancellation) + } else { + nil + } + let proxyDaily = CostUsageScanner.buildClaudeReportFromCache( + cache: claudeCache, + range: range, + attributionFilter: .codexBackendOnly, + attributionResolver: attributionResolver, + modelsDevCatalog: CostUsagePricing.modelsDevCatalog( + now: now, + cacheRoot: options.cacheRoot), + modelsDevCacheRoot: options.cacheRoot) + if !proxyDaily.data.isEmpty { + reports.append(proxyDaily) + claudeProxyMerged = true + if claudeCache.lastScanUnixMs > 0 { + scanTimes.append(Date( + timeIntervalSince1970: TimeInterval(claudeCache.lastScanUnixMs) / 1000)) + } + if let proxyProject = Self.unknownProjectBreakdown( + from: proxyDaily, + name: "Claude Code via CLIProxyAPI") + { + projects.append(proxyProject) + } + sessions = [] + } + } + if let piResult = PiSessionCostScanner.loadCachedDailyReportResult( provider: .codex, since: since, @@ -600,7 +767,7 @@ public struct CostUsageFetcher: Sendable { projects: Self.mergedProjectBreakdowns(projects), sessions: sessions, updatedAt: scanTimes.min()), - lastRefreshAt: piMerged ? nil : nativeScanAt) + lastRefreshAt: piMerged || claudeProxyMerged ? nil : nativeScanAt) } return cachedSnapshot.flatMap(\.self) } @@ -797,11 +964,216 @@ 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) - private static func unknownProjectBreakdown(from daily: CostUsageDailyReport) -> CostUsageProjectBreakdown? { + 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 } + 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) + } + + 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, @@ -809,7 +1181,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, @@ -821,12 +1193,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 @@ -841,7 +1213,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)) @@ -851,7 +1223,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, @@ -875,15 +1247,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, @@ -921,28 +1303,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 @@ -955,10 +1351,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..e807e023ef 100644 --- a/Sources/CodexBarCore/CostUsageModels.swift +++ b/Sources/CodexBarCore/CostUsageModels.swift @@ -272,6 +272,131 @@ 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 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 +407,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 +420,7 @@ public struct CostUsageDailyReport: Sendable, Decodable { case priorityCostUSD case standardTokens case priorityTokens + case attribution } public init(from decoder: Decoder) throws { @@ -310,6 +437,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 +448,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 +459,7 @@ public struct CostUsageDailyReport: Sendable, Decodable { self.priorityCostUSD = priorityCostUSD self.standardTokens = standardTokens self.priorityTokens = priorityTokens + self.attribution = attribution } } @@ -527,6 +657,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 +703,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 +731,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 +769,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 +796,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 +906,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 ca353e24bb..6ff68f2988 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 = "3aa49b47f4b78e13" + static let value = "be0b8c0e7b2a953b" } 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 3727bf2ce2..f04a574671 100644 --- a/Sources/CodexBarCore/Vendored/CostUsage/CostUsagePricing.swift +++ b/Sources/CodexBarCore/Vendored/CostUsage/CostUsagePricing.swift @@ -484,6 +484,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 @@ -585,6 +601,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, diff --git a/Sources/CodexBarCore/Vendored/CostUsage/CostUsageScanner+CacheHelpers.swift b/Sources/CodexBarCore/Vendored/CostUsage/CostUsageScanner+CacheHelpers.swift index d6c5882db9..a8e6ebaf72 100644 --- a/Sources/CodexBarCore/Vendored/CostUsage/CostUsageScanner+CacheHelpers.swift +++ b/Sources/CodexBarCore/Vendored/CostUsage/CostUsageScanner+CacheHelpers.swift @@ -1639,7 +1639,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 e862bb349e..d3107a0a74 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,91 @@ 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 { 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, @@ -77,6 +171,7 @@ extension CostUsageScanner { range: CostUsageDayRange, providerFilter: ClaudeLogProviderFilter, startOffset: Int64 = 0, + attributionResolver: CLIProxyAPIAttributionResolver? = nil, modelsDevCatalog: ModelsDevCatalog? = nil, modelsDevCacheRoot: URL? = nil) -> ClaudeParseResult { @@ -86,6 +181,7 @@ extension CostUsageScanner { range: range, providerFilter: providerFilter, startOffset: startOffset, + attributionResolver: attributionResolver, modelsDevCatalog: modelsDevCatalog, modelsDevCacheRoot: modelsDevCacheRoot, checkCancellation: nil)) ?? ClaudeParseResult(days: [:], rows: [], parsedBytes: startOffset) @@ -96,6 +192,7 @@ extension CostUsageScanner { range: CostUsageDayRange, providerFilter: ClaudeLogProviderFilter, startOffset: Int64 = 0, + attributionResolver: CLIProxyAPIAttributionResolver? = nil, modelsDevCatalog: ModelsDevCatalog? = nil, modelsDevCacheRoot: URL? = nil, checkCancellation: CancellationCheck? = nil) throws -> ClaudeParseResult @@ -179,17 +276,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, @@ -197,7 +305,7 @@ extension CostUsageScanner { cacheCreate1h: cacheCreate1h, output: output, costNanos: costNanos, - costPriced: cost != nil) + costPriced: modelResolution.cost != nil) guard CostUsageDayRange.isInRange( dayKey: dayKey, @@ -207,18 +315,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, @@ -227,7 +330,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. @@ -346,6 +450,116 @@ 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) + } + 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 { + liveAttribution + } 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]]] = [:] @@ -509,6 +723,7 @@ extension CostUsageScanner { let range: CostUsageDayRange let providerFilter: ClaudeLogProviderFilter let forceFullScan: Bool + let attributionResolver: CLIProxyAPIAttributionResolver? let modelsDevCatalog: ModelsDevCatalog? let modelsDevCacheRoot: URL? let checkCancellation: CancellationCheck? @@ -518,6 +733,7 @@ extension CostUsageScanner { range: CostUsageDayRange, providerFilter: ClaudeLogProviderFilter, forceFullScan: Bool, + attributionResolver: CLIProxyAPIAttributionResolver?, modelsDevCatalog: ModelsDevCatalog?, modelsDevCacheRoot: URL?, checkCancellation: CancellationCheck?) @@ -527,6 +743,7 @@ extension CostUsageScanner { self.range = range self.providerFilter = providerFilter self.forceFullScan = forceFullScan + self.attributionResolver = attributionResolver self.modelsDevCatalog = modelsDevCatalog self.modelsDevCacheRoot = modelsDevCacheRoot self.checkCancellation = checkCancellation @@ -561,6 +778,7 @@ extension CostUsageScanner { range: state.range, providerFilter: state.providerFilter, startOffset: startOffset, + attributionResolver: state.attributionResolver, modelsDevCatalog: state.modelsDevCatalog, modelsDevCacheRoot: state.modelsDevCacheRoot, checkCancellation: state.checkCancellation) @@ -578,6 +796,7 @@ extension CostUsageScanner { fileURL: url, range: state.range, providerFilter: state.providerFilter, + attributionResolver: state.attributionResolver, modelsDevCatalog: state.modelsDevCatalog, modelsDevCacheRoot: state.modelsDevCacheRoot, checkCancellation: state.checkCancellation) @@ -665,13 +884,30 @@ 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 let providerFilter = options.claudeLogProviderFilter + let cliProxyAPIAttributionEnabled = !CostUsageCacheLocations.isCLIProxyAPIExplicitlyDisconnected( + stateRoot: options.cacheRoot) + let attributionResolver: CLIProxyAPIAttributionResolver? = if cliProxyAPIAttributionEnabled, + let home = options.cliProxyAPIHome + { + try CLIProxyAPIAttributionResolver.load( + home: home, + cacheRoot: options.cacheRoot, + forceReload: options.forceRescan, + checkCancellation: checkCancellation) + } else { + nil + } var touched: Set = [] @@ -685,7 +921,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) @@ -706,30 +945,284 @@ 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?.deletingLastPathComponent()) + { + 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) + 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 + } + } return Self.buildClaudeReportFromCache( cache: cache, range: range, + attributionFilter: reportAttributionFilter, + attributionResolver: reportAttributionEnabled ? attributionResolver : nil, + allowCachedCLIProxyAPIAttribution: reportAttributionEnabled, modelsDevCatalog: modelsDevCatalog, modelsDevCacheRoot: options.cacheRoot) } - private static func buildClaudeReportFromCache( + 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 + } + + 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 liveAttribution?.route == .cliProxyAPI { + liveAttribution + } 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 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) + } + 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: modelsDevCacheRoot) + } + + static func buildClaudeReportFromCache( cache: CostUsageCache, range: CostUsageDayRange, + attributionFilter: ClaudeAttributionFilter = .all, + attributionResolver: CLIProxyAPIAttributionResolver? = nil, + allowCachedCLIProxyAPIAttribution: Bool = true, modelsDevCatalog: ModelsDevCatalog? = nil, modelsDevCacheRoot: URL? = nil) -> CostUsageDailyReport { @@ -741,50 +1234,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 @@ -795,8 +1265,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 @@ -810,7 +1281,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 @@ -824,7 +1295,8 @@ extension CostUsageScanner { CostUsageDailyReport.ModelBreakdown( modelName: model, costUSD: cost, - totalTokens: totalTokens)) + totalTokens: totalTokens, + attribution: modelKey.attribution)) if let cost { dayCost += cost dayCostSeen = true @@ -843,7 +1315,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 fe2677734a..e6581b150a 100644 --- a/Sources/CodexBarCore/Vendored/CostUsage/CostUsageScanner.swift +++ b/Sources/CodexBarCore/Vendored/CostUsage/CostUsageScanner.swift @@ -24,6 +24,12 @@ enum CostUsageScanner { case excludeVertexAI } + enum ClaudeAttributionFilter { + case all + case codexBackendOnly + case excludeCodexBackend + } + struct Options { var codexSessionsRoot: URL? var claudeProjectsRoots: [URL]? @@ -32,6 +38,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 @@ -50,6 +58,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, @@ -61,6 +71,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) @@ -983,6 +995,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/CLIProxyAPIAttributionBatchTests.swift b/Tests/CodexBarTests/CLIProxyAPIAttributionBatchTests.swift new file mode 100644 index 0000000000..b35a1e11f7 --- /dev/null +++ b/Tests/CodexBarTests/CLIProxyAPIAttributionBatchTests.swift @@ -0,0 +1,145 @@ +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)) + } + + 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..b68c047a19 --- /dev/null +++ b/Tests/CodexBarTests/CLIProxyAPIAttributionResolverTests.swift @@ -0,0 +1,930 @@ +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)) + } +} + +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..4d1d9abb67 --- /dev/null +++ b/Tests/CodexBarTests/CLIProxyAPIUsageCacheTests.swift @@ -0,0 +1,438 @@ +import Foundation +import Testing +@testable import CodexBarCore + +struct CLIProxyAPIUsageCacheTests { + @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 settings = CLIProxyAPIConnectionSettings( + baseURL: "http://127.0.0.1:8317", + managementKey: "test-management-key") + var didStore = false + var didRollback = false + + let saved = CLIProxyAPIConnectionSettingsStore.save( + settings, + store: { _ in + didStore = true + return true + }, + clearDisconnectedState: { false }, + rollback: { + didRollback = true + return true + }) + + #expect(!saved) + #expect(didStore) + #expect(didRollback) + } + + @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 `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"]) + } + + 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) + } + } + } +} diff --git a/Tests/CodexBarTests/CLIProxyAPIUsageCollectorTests.swift b/Tests/CodexBarTests/CLIProxyAPIUsageCollectorTests.swift new file mode 100644 index 0000000000..f432fa169b --- /dev/null +++ b/Tests/CodexBarTests/CLIProxyAPIUsageCollectorTests.swift @@ -0,0 +1,379 @@ +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 a 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: "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 save CLIProxyAPI usage telemetry.")) + #expect(CLIProxyAPIUsagePendingIO.load(pendingRoot: pendingRoot)?.map(\.requestID) == ["request-1"]) + try fileManager.removeItem(at: cacheRoot) + let retryClient = CLIProxyAPIUsageQueueClient( + settings: .init(managementKey: "management-secret"), + dataLoader: { request in + #expect(CLIProxyAPIUsageCacheIO.load(cacheRoot: cacheRoot).map(\.requestID) == ["request-1"]) + 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(\.requestID) == ["request-1"]) + #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 `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..4674ef024f --- /dev/null +++ b/Tests/CodexBarTests/CLIProxyAPIUsageStoreTests.swift @@ -0,0 +1,253 @@ +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 + } +} + +@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 didClearConfiguration = false + var collectorFinishedBeforePurge = false + + let removed = await store.removeCLIProxyAPIConfiguration( + purgeTelemetry: { + collectorFinishedBeforePurge = collectorFinished.value + return true + }, + clear: { + didClearConfiguration = true + return true + }) + await task.value + + #expect(removed == .removed) + #expect(didClearConfiguration) + #expect(collectorFinishedBeforePurge) + #expect(store.cliProxyAPIUsageCollectorTask == nil) + #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) + var didPurgeTelemetry = false + + let removed = await store.removeCLIProxyAPIConfiguration( + purgeTelemetry: { + didPurgeTelemetry = true + return true + }, + clear: { false }) + + #expect(removed == .configurationRemovalFailed) + #expect(!didPurgeTelemetry) + } + + @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) + + let removed = await store.removeCLIProxyAPIConfiguration( + purgeTelemetry: { false }, + clear: { true }) + + #expect(removed == .telemetryCleanupFailed) + } + + @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 nil + }) + store.stopCLIProxyAPIUsageCollector() + + #expect(error == nil) + #expect(deletionStartedAfterDrain.value) + } + + @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)) + } +} + +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 d7d4b24f00..662abd2fd4 100644 --- a/Tests/CodexBarTests/CostHistoryChartMenuViewTests.swift +++ b/Tests/CodexBarTests/CostHistoryChartMenuViewTests.swift @@ -865,6 +865,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 bcf9d25777..709ac91e71 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..0bbe65a172 --- /dev/null +++ b/Tests/CodexBarTests/CostUsageFetcherCachedProxyDisconnectTests.swift @@ -0,0 +1,177 @@ +import Foundation +import Testing +@testable import CodexBarCore + +struct CostUsageFetcherCachedProxyDisconnectTests { + @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 5db8108693..f4166182d8 100644 --- a/Tests/CodexBarTests/CostUsageFetcherTests.swift +++ b/Tests/CodexBarTests/CostUsageFetcherTests.swift @@ -709,6 +709,611 @@ 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 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) + + 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..64cb9578bc 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,203 @@ 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", 0.00028), + ("claude-new", 0.00045), + ]) + func `claude fetch resolves OpenAI compatible proxy pricing across model vendors`( + upstreamModel: 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: "openrouter", + executorType: "OpenAICompatExecutor", + 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 } } } + } + } + """.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 == "openrouter") + #expect(breakdown.attribution?.upstream?.model == upstreamModel) + #expect(abs((breakdown.costUSD ?? 0) - expectedCost) < 0.0000001) + } } private struct UnknownModelPricingFixture { diff --git a/Tests/CodexBarTests/ShareStatsTests.swift b/Tests/CodexBarTests/ShareStatsTests.swift index adb0a24651..1fe1742d7b 100644 --- a/Tests/CodexBarTests/ShareStatsTests.swift +++ b/Tests/CodexBarTests/ShareStatsTests.swift @@ -5,6 +5,77 @@ 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 `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..1bbc54ca21 --- /dev/null +++ b/Tests/CodexBarTests/SpendDashboardCodexProxySourceTests.swift @@ -0,0 +1,108 @@ +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]) + } + + 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..19f4b99e7c --- /dev/null +++ b/Tests/CodexBarTests/SpendDashboardProxyAttributionTests.swift @@ -0,0 +1,154 @@ +import CodexBarCore +import Foundation +import Testing +@testable import CodexBar + +struct SpendDashboardProxyAttributionTests { + @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]) + } +}