From f2be96b90bbe8a93eb3915f1db23a881c5b18ce8 Mon Sep 17 00:00:00 2001 From: claude89757 <138977524+claude89757@users.noreply.github.com> Date: Wed, 4 Feb 2026 04:07:09 +0800 Subject: [PATCH] Add CodeBuddy provider for credit usage tracking Implements CodeBuddy (Tencent Cloud AI assistant) integration with: - Cookie-based authentication from Chrome browser - Enterprise ID auto-detection from API responses - Credit usage fetching from billing API - Daily usage chart visualization - Session keepalive functionality - Unit tests for models, snapshots, and cookie parsing Follows existing provider patterns and architecture. --- .gitignore | 5 + .../CodeBuddyDailyUsageChartMenuView.swift | 310 ++++++++++++++ .../CodeBuddyProviderImplementation.swift | 101 +++++ .../CodeBuddy/CodeBuddyProviderRuntime.swift | 103 +++++ .../CodeBuddy/CodeBuddySettingsStore.swift | 58 +++ .../ProviderImplementationRegistry.swift | 1 + .../CodexBar/StatusItemController+Menu.swift | 53 +++ Sources/CodexBar/UsageStore+Refresh.swift | 12 + .../CodexBar/UsageStore+TokenAccounts.swift | 4 + Sources/CodexBar/UsageStore.swift | 111 +++++ Sources/CodexBarCLI/TokenAccountCLI.swift | 12 +- .../CodexBarCore/Config/CodexBarConfig.swift | 3 + .../CodexBarCore/Logging/LogCategories.swift | 4 + .../CodeBuddy/CodeBuddyAPIError.swift | 30 ++ .../CodeBuddy/CodeBuddyCookieHeader.swift | 120 ++++++ .../CodeBuddy/CodeBuddyCookieImporter.swift | 200 +++++++++ .../Providers/CodeBuddy/CodeBuddyModels.swift | 46 +++ .../CodeBuddyProviderDescriptor.swift | 153 +++++++ .../CodeBuddy/CodeBuddySessionKeepalive.swift | 385 ++++++++++++++++++ .../CodeBuddy/CodeBuddySettingsReader.swift | 25 ++ .../CodeBuddy/CodeBuddyUsageFetcher.swift | 154 +++++++ .../CodeBuddy/CodeBuddyUsageSnapshot.swift | 64 +++ .../Providers/ProviderDescriptor.swift | 1 + .../Providers/ProviderFetchPlan.swift | 5 + .../Providers/ProviderSettingsSnapshot.swift | 29 +- .../CodexBarCore/Providers/Providers.swift | 2 + .../Vendored/CostUsage/CostUsageScanner.swift | 2 + .../CodexBarWidgetProvider.swift | 1 + .../CodexBarWidget/CodexBarWidgetViews.swift | 3 + Tests/CodexBarTests/CodeBuddyTests.swift | 161 ++++++++ Tests/CodexBarTests/SettingsStoreTests.swift | 1 + 31 files changed, 2153 insertions(+), 6 deletions(-) create mode 100644 Sources/CodexBar/CodeBuddyDailyUsageChartMenuView.swift create mode 100644 Sources/CodexBar/Providers/CodeBuddy/CodeBuddyProviderImplementation.swift create mode 100644 Sources/CodexBar/Providers/CodeBuddy/CodeBuddyProviderRuntime.swift create mode 100644 Sources/CodexBar/Providers/CodeBuddy/CodeBuddySettingsStore.swift create mode 100644 Sources/CodexBarCore/Providers/CodeBuddy/CodeBuddyAPIError.swift create mode 100644 Sources/CodexBarCore/Providers/CodeBuddy/CodeBuddyCookieHeader.swift create mode 100644 Sources/CodexBarCore/Providers/CodeBuddy/CodeBuddyCookieImporter.swift create mode 100644 Sources/CodexBarCore/Providers/CodeBuddy/CodeBuddyModels.swift create mode 100644 Sources/CodexBarCore/Providers/CodeBuddy/CodeBuddyProviderDescriptor.swift create mode 100644 Sources/CodexBarCore/Providers/CodeBuddy/CodeBuddySessionKeepalive.swift create mode 100644 Sources/CodexBarCore/Providers/CodeBuddy/CodeBuddySettingsReader.swift create mode 100644 Sources/CodexBarCore/Providers/CodeBuddy/CodeBuddyUsageFetcher.swift create mode 100644 Sources/CodexBarCore/Providers/CodeBuddy/CodeBuddyUsageSnapshot.swift create mode 100644 Tests/CodexBarTests/CodeBuddyTests.swift diff --git a/.gitignore b/.gitignore index 01ae2d5c..1359475b 100644 --- a/.gitignore +++ b/.gitignore @@ -33,3 +33,8 @@ docs/.astro/ # Swift Package Manager metadata (leave sources tracked) # Packages/ # Package.resolved + +# ai code +.claude/ +openspec/ +CLAUDE.md \ No newline at end of file diff --git a/Sources/CodexBar/CodeBuddyDailyUsageChartMenuView.swift b/Sources/CodexBar/CodeBuddyDailyUsageChartMenuView.swift new file mode 100644 index 00000000..583f8fd3 --- /dev/null +++ b/Sources/CodexBar/CodeBuddyDailyUsageChartMenuView.swift @@ -0,0 +1,310 @@ +import Charts +import CodexBarCore +import SwiftUI + +@MainActor +struct CodeBuddyDailyUsageChartMenuView: View { + private struct Point: Identifiable { + let id: String + let date: Date + let creditsUsed: Double + + init(date: Date, creditsUsed: Double) { + self.date = date + self.creditsUsed = creditsUsed + self.id = "\(Int(date.timeIntervalSince1970))-\(creditsUsed)" + } + } + + private let dailyUsage: [CodeBuddyDailyUsageEntry] + private let width: CGFloat + @State private var selectedDayKey: String? + + init(dailyUsage: [CodeBuddyDailyUsageEntry], width: CGFloat) { + self.dailyUsage = dailyUsage + self.width = width + } + + var body: some View { + let model = Self.makeModel(from: self.dailyUsage) + VStack(alignment: .leading, spacing: 10) { + if model.points.isEmpty { + Text("No usage history data.") + .font(.footnote) + .foregroundStyle(.secondary) + } else { + Chart { + ForEach(model.points) { point in + BarMark( + x: .value("Day", point.date, unit: .day), + y: .value("Credits used", point.creditsUsed)) + .foregroundStyle(Self.barColor) + } + if let peak = Self.peakPoint(model: model) { + let capStart = max(peak.creditsUsed - Self.capHeight(maxValue: model.maxCreditsUsed), 0) + BarMark( + x: .value("Day", peak.date, unit: .day), + yStart: .value("Cap start", capStart), + yEnd: .value("Cap end", peak.creditsUsed)) + .foregroundStyle(Color(nsColor: .systemYellow)) + } + } + .chartYAxis(.hidden) + .chartXAxis { + AxisMarks(values: model.axisDates) { _ in + AxisGridLine().foregroundStyle(Color.clear) + AxisTick().foregroundStyle(Color.clear) + AxisValueLabel(format: .dateTime.month(.abbreviated).day()) + .font(.caption2) + .foregroundStyle(Color(nsColor: .tertiaryLabelColor)) + } + } + .chartLegend(.hidden) + .frame(height: 130) + .chartOverlay { proxy in + GeometryReader { geo in + ZStack(alignment: .topLeading) { + if let rect = self.selectionBandRect(model: model, proxy: proxy, geo: geo) { + Rectangle() + .fill(Self.selectionBandColor) + .frame(width: rect.width, height: rect.height) + .position(x: rect.midX, y: rect.midY) + .allowsHitTesting(false) + } + MouseLocationReader { location in + self.updateSelection(location: location, model: model, proxy: proxy, geo: geo) + } + .frame(maxWidth: .infinity, maxHeight: .infinity) + .contentShape(Rectangle()) + } + } + } + + let detail = self.detailLines(model: model) + VStack(alignment: .leading, spacing: 0) { + Text(detail.primary) + .font(.caption) + .foregroundStyle(.secondary) + .lineLimit(1) + .truncationMode(.tail) + .frame(height: 16, alignment: .leading) + Text(detail.secondary ?? " ") + .font(.caption) + .foregroundStyle(.secondary) + .lineLimit(1) + .truncationMode(.tail) + .frame(height: 16, alignment: .leading) + .opacity(detail.secondary == nil ? 0 : 1) + } + + if let total = model.totalCreditsUsed { + Text("Total (30d): \(total.formatted(.number.precision(.fractionLength(0...2)))) credits") + .font(.caption) + .foregroundStyle(.secondary) + } + } + } + .padding(.horizontal, 16) + .padding(.vertical, 10) + .frame(minWidth: self.width, maxWidth: .infinity, alignment: .leading) + } + + private struct Model { + let points: [Point] + let usageByDayKey: [String: CodeBuddyDailyUsageEntry] + let pointsByDayKey: [String: Point] + let dayDates: [(dayKey: String, date: Date)] + let selectableDayDates: [(dayKey: String, date: Date)] + let axisDates: [Date] + let peakKey: String? + let totalCreditsUsed: Double? + let maxCreditsUsed: Double + } + + // CodeBuddy brand color (blue) + private static let barColor = Color(red: 0 / 255, green: 120 / 255, blue: 215 / 255) + private static let selectionBandColor = Color(nsColor: .labelColor).opacity(0.1) + private static func capHeight(maxValue: Double) -> Double { + maxValue * 0.05 + } + + private static func makeModel(from dailyUsage: [CodeBuddyDailyUsageEntry]) -> Model { + let sorted = dailyUsage.sorted { lhs, rhs in lhs.date < rhs.date } + + var points: [Point] = [] + points.reserveCapacity(sorted.count) + + var usageByDayKey: [String: CodeBuddyDailyUsageEntry] = [:] + usageByDayKey.reserveCapacity(sorted.count) + + var pointsByDayKey: [String: Point] = [:] + pointsByDayKey.reserveCapacity(sorted.count) + + var dayDates: [(dayKey: String, date: Date)] = [] + dayDates.reserveCapacity(sorted.count) + + var selectableDayDates: [(dayKey: String, date: Date)] = [] + selectableDayDates.reserveCapacity(sorted.count) + + var totalCreditsUsed: Double = 0 + var peak: (key: String, creditsUsed: Double)? + var maxCreditsUsed: Double = 0 + + for entry in sorted { + guard let date = self.dateFromDayKey(entry.date) else { continue } + usageByDayKey[entry.date] = entry + dayDates.append((dayKey: entry.date, date: date)) + totalCreditsUsed += entry.credit + if entry.credit > 0 { + let point = Point(date: date, creditsUsed: entry.credit) + points.append(point) + pointsByDayKey[entry.date] = point + selectableDayDates.append((dayKey: entry.date, date: date)) + if let cur = peak { + if entry.credit > cur.creditsUsed { peak = (entry.date, entry.credit) } + } else { + peak = (entry.date, entry.credit) + } + maxCreditsUsed = max(maxCreditsUsed, entry.credit) + } + } + + let axisDates: [Date] = { + guard let first = dayDates.first?.date, let last = dayDates.last?.date else { return [] } + if Calendar.current.isDate(first, inSameDayAs: last) { return [first] } + return [first, last] + }() + + return Model( + points: points, + usageByDayKey: usageByDayKey, + pointsByDayKey: pointsByDayKey, + dayDates: dayDates, + selectableDayDates: selectableDayDates, + axisDates: axisDates, + peakKey: peak?.key, + totalCreditsUsed: totalCreditsUsed > 0 ? totalCreditsUsed : nil, + maxCreditsUsed: maxCreditsUsed) + } + + private static func dateFromDayKey(_ key: String) -> Date? { + let parts = key.split(separator: "-") + guard parts.count == 3, + let year = Int(parts[0]), + let month = Int(parts[1]), + let day = Int(parts[2]) + else { + return nil + } + + var comps = DateComponents() + comps.calendar = Calendar.current + comps.timeZone = TimeZone.current + comps.year = year + comps.month = month + comps.day = day + comps.hour = 12 + return comps.date + } + + private static func peakPoint(model: Model) -> Point? { + guard let key = model.peakKey else { return nil } + return model.pointsByDayKey[key] + } + + private func selectionBandRect(model: Model, proxy: ChartProxy, geo: GeometryProxy) -> CGRect? { + guard let key = self.selectedDayKey else { return nil } + guard let plotAnchor = proxy.plotFrame else { return nil } + let plotFrame = geo[plotAnchor] + guard let index = model.dayDates.firstIndex(where: { $0.dayKey == key }) else { return nil } + let date = model.dayDates[index].date + guard let x = proxy.position(forX: date) else { return nil } + + func xForIndex(_ idx: Int) -> CGFloat? { + guard idx >= 0, idx < model.dayDates.count else { return nil } + return proxy.position(forX: model.dayDates[idx].date) + } + + let xPrev = xForIndex(index - 1) + let xNext = xForIndex(index + 1) + + if model.dayDates.count <= 1 { + return CGRect( + x: plotFrame.origin.x, + y: plotFrame.origin.y, + width: plotFrame.width, + height: plotFrame.height) + } + + let leftInPlot: CGFloat = if let xPrev { + (xPrev + x) / 2 + } else if let xNext { + x - (xNext - x) / 2 + } else { + x - 8 + } + + let rightInPlot: CGFloat = if let xNext { + (xNext + x) / 2 + } else if let xPrev { + x + (x - xPrev) / 2 + } else { + x + 8 + } + + let left = plotFrame.origin.x + min(leftInPlot, rightInPlot) + let right = plotFrame.origin.x + max(leftInPlot, rightInPlot) + return CGRect(x: left, y: plotFrame.origin.y, width: right - left, height: plotFrame.height) + } + + private func updateSelection( + location: CGPoint?, + model: Model, + proxy: ChartProxy, + geo: GeometryProxy) + { + guard let location else { + if self.selectedDayKey != nil { self.selectedDayKey = nil } + return + } + + guard let plotAnchor = proxy.plotFrame else { return } + let plotFrame = geo[plotAnchor] + guard plotFrame.contains(location) else { return } + + let xInPlot = location.x - plotFrame.origin.x + guard let date: Date = proxy.value(atX: xInPlot) else { return } + guard let nearest = self.nearestDayKey(to: date, model: model) else { return } + + if self.selectedDayKey != nearest { + self.selectedDayKey = nearest + } + } + + private func nearestDayKey(to date: Date, model: Model) -> String? { + guard !model.selectableDayDates.isEmpty else { return nil } + var best: (key: String, distance: TimeInterval)? + for entry in model.selectableDayDates { + let dist = abs(entry.date.timeIntervalSince(date)) + if let cur = best { + if dist < cur.distance { best = (entry.dayKey, dist) } + } else { + best = (entry.dayKey, dist) + } + } + return best?.key + } + + private func detailLines(model: Model) -> (primary: String, secondary: String?) { + guard let key = self.selectedDayKey, + let entry = model.usageByDayKey[key], + let date = Self.dateFromDayKey(key) + else { + return ("Hover a bar for details", nil) + } + + let dayLabel = date.formatted(.dateTime.month(.abbreviated).day()) + let credits = entry.credit.formatted(.number.precision(.fractionLength(0...2))) + return ("\(dayLabel): \(credits) credits", nil) + } +} diff --git a/Sources/CodexBar/Providers/CodeBuddy/CodeBuddyProviderImplementation.swift b/Sources/CodexBar/Providers/CodeBuddy/CodeBuddyProviderImplementation.swift new file mode 100644 index 00000000..9b773978 --- /dev/null +++ b/Sources/CodexBar/Providers/CodeBuddy/CodeBuddyProviderImplementation.swift @@ -0,0 +1,101 @@ +import AppKit +import CodexBarCore +import CodexBarMacroSupport +import Foundation +import SwiftUI + +@ProviderImplementationRegistration +struct CodeBuddyProviderImplementation: ProviderImplementation { + let id: UsageProvider = .codebuddy + + func makeRuntime() -> (any ProviderRuntime)? { + CodeBuddyProviderRuntime() + } + + @MainActor + func presentation(context _: ProviderPresentationContext) -> ProviderPresentation { + ProviderPresentation { _ in "web" } + } + + @MainActor + func observeSettings(_ settings: SettingsStore) { + _ = settings.codebuddyCookieSource + _ = settings.codebuddyManualCookieHeader + _ = settings.codebuddyEnterpriseID + } + + @MainActor + func settingsSnapshot(context: ProviderSettingsSnapshotContext) -> ProviderSettingsSnapshotContribution? { + .codebuddy(context.settings.codebuddySettingsSnapshot(tokenOverride: context.tokenOverride)) + } + + @MainActor + func settingsPickers(context: ProviderSettingsContext) -> [ProviderSettingsPickerDescriptor] { + let cookieBinding = Binding( + get: { context.settings.codebuddyCookieSource.rawValue }, + set: { raw in + context.settings.codebuddyCookieSource = ProviderCookieSource(rawValue: raw) ?? .auto + }) + let options = ProviderCookieSourceUI.options( + allowsOff: true, + keychainDisabled: context.settings.debugDisableKeychainAccess) + + let subtitle: () -> String? = { + ProviderCookieSourceUI.subtitle( + source: context.settings.codebuddyCookieSource, + keychainDisabled: context.settings.debugDisableKeychainAccess, + auto: "Automatic imports browser cookies.", + manual: "Paste a cookie header from the dashboard.", + off: "CodeBuddy cookies are disabled.") + } + + return [ + ProviderSettingsPickerDescriptor( + id: "codebuddy-cookie-source", + title: "Cookie source", + subtitle: "Automatic imports browser cookies.", + dynamicSubtitle: subtitle, + binding: cookieBinding, + options: options, + isVisible: nil, + onChange: nil), + ] + } + + @MainActor + func settingsFields(context: ProviderSettingsContext) -> [ProviderSettingsFieldDescriptor] { + [ + ProviderSettingsFieldDescriptor( + id: "codebuddy-cookie", + title: "", + subtitle: "", + kind: .secure, + placeholder: "Cookie: session=\u{2026}", + binding: context.stringBinding(\.codebuddyManualCookieHeader), + actions: [ + ProviderSettingsActionDescriptor( + id: "codebuddy-open-dashboard", + title: "Open Dashboard", + style: .link, + isVisible: nil, + perform: { + if let url = URL(string: "https://tencent.sso.codebuddy.cn/profile/usage") { + NSWorkspace.shared.open(url) + } + }), + ], + isVisible: { context.settings.codebuddyCookieSource == .manual }, + onActivate: { context.settings.ensureCodeBuddySessionLoaded() }), + ProviderSettingsFieldDescriptor( + id: "codebuddy-enterprise-id", + title: "Enterprise ID", + subtitle: "Default: etahzsqej0n4 (only change if you have a different enterprise)", + kind: .plain, + placeholder: "etahzsqej0n4", + binding: context.stringBinding(\.codebuddyEnterpriseID), + actions: [], + isVisible: nil, + onActivate: nil), + ] + } +} diff --git a/Sources/CodexBar/Providers/CodeBuddy/CodeBuddyProviderRuntime.swift b/Sources/CodexBar/Providers/CodeBuddy/CodeBuddyProviderRuntime.swift new file mode 100644 index 00000000..1b75a4a5 --- /dev/null +++ b/Sources/CodexBar/Providers/CodeBuddy/CodeBuddyProviderRuntime.swift @@ -0,0 +1,103 @@ +import CodexBarCore +import Foundation + +@MainActor +final class CodeBuddyProviderRuntime: ProviderRuntime { + let id: UsageProvider = .codebuddy + private var keepalive: CodeBuddySessionKeepalive? + private static let log = CodexBarLog.logger(LogCategories.codeBuddyKeepalive) + + func start(context: ProviderRuntimeContext) { + self.updateKeepalive(context: context) + } + + func stop(context: ProviderRuntimeContext) { + self.stopKeepalive(context: context, reason: "provider disabled") + } + + func settingsDidChange(context: ProviderRuntimeContext) { + self.updateKeepalive(context: context) + } + + func providerDidFail(context: ProviderRuntimeContext, provider: UsageProvider, error: Error) { + guard provider == .codebuddy else { return } + let message = error.localizedDescription + // Check for session/cookie related errors + guard message.contains("session expired") || + message.contains("401") || + message.contains("invalid") && message.contains("cookie") + else { return } + Self.log.warning("CodeBuddy session may have expired; triggering recovery") + Task { [weak self] in + guard let self else { return } + await self.forceRefresh(context: context) + } + } + + func perform(action: ProviderRuntimeAction, context: ProviderRuntimeContext) async { + switch action { + case .forceSessionRefresh: + await self.forceRefresh(context: context) + case .openAIWebAccessToggled: + break + } + } + + private func updateKeepalive(context: ProviderRuntimeContext) { + #if os(macOS) + let shouldRun = context.store.isEnabled(.codebuddy) + let isRunning = self.keepalive != nil + + if shouldRun, !isRunning { + self.startKeepalive(context: context) + } else if !shouldRun, isRunning { + self.stopKeepalive(context: context, reason: "provider disabled") + } + #endif + } + + private func startKeepalive(context: ProviderRuntimeContext) { + #if os(macOS) + Self.log.info( + "CodeBuddy keepalive check", + metadata: [ + "enabled": context.store.isEnabled(.codebuddy) ? "1" : "0", + "available": context.store.isProviderAvailable(.codebuddy) ? "1" : "0", + ]) + + guard context.store.isEnabled(.codebuddy) else { + Self.log.warning("CodeBuddy keepalive not started (provider disabled)") + return + } + + let logger: (String) -> Void = { message in + Self.log.verbose(message) + } + + let onSessionRecovered: () async -> Void = { [weak store = context.store] in + guard let store else { return } + Self.log.info("CodeBuddy session recovered; refreshing usage") + await store.refreshProvider(.codebuddy) + } + + self.keepalive = CodeBuddySessionKeepalive(logger: logger, onSessionRecovered: onSessionRecovered) + self.keepalive?.start() + Self.log.info("CodeBuddy keepalive started") + #endif + } + + private func stopKeepalive(context _: ProviderRuntimeContext, reason: String) { + #if os(macOS) + guard self.keepalive != nil else { return } + self.keepalive?.stop() + self.keepalive = nil + Self.log.info("CodeBuddy keepalive stopped (\(reason))") + #endif + } + + private func forceRefresh(context _: ProviderRuntimeContext) async { + #if os(macOS) + await self.keepalive?.forceRefresh() + #endif + } +} diff --git a/Sources/CodexBar/Providers/CodeBuddy/CodeBuddySettingsStore.swift b/Sources/CodexBar/Providers/CodeBuddy/CodeBuddySettingsStore.swift new file mode 100644 index 00000000..28c11109 --- /dev/null +++ b/Sources/CodexBar/Providers/CodeBuddy/CodeBuddySettingsStore.swift @@ -0,0 +1,58 @@ +import CodexBarCore +import Foundation + +extension SettingsStore { + var codebuddyManualCookieHeader: String { + get { self.configSnapshot.providerConfig(for: .codebuddy)?.sanitizedCookieHeader ?? "" } + set { + self.updateProviderConfig(provider: .codebuddy) { entry in + entry.cookieHeader = self.normalizedConfigValue(newValue) + } + self.logSecretUpdate(provider: .codebuddy, field: "cookieHeader", value: newValue) + } + } + + var codebuddyCookieSource: ProviderCookieSource { + get { self.resolvedCookieSource(provider: .codebuddy, fallback: .auto) } + set { + self.updateProviderConfig(provider: .codebuddy) { entry in + entry.cookieSource = newValue + } + self.logProviderModeChange(provider: .codebuddy, field: "cookieSource", value: newValue.rawValue) + } + } + + /// Default enterprise ID for most CodeBuddy users + private static let defaultEnterpriseID = "etahzsqej0n4" + + var codebuddyEnterpriseID: String { + get { + let stored = self.configSnapshot.providerConfig(for: .codebuddy)?.enterpriseID + // Use default if not explicitly set + return stored ?? Self.defaultEnterpriseID + } + set { + self.updateProviderConfig(provider: .codebuddy) { entry in + // Only store if different from default + let normalized = self.normalizedConfigValue(newValue) + entry.enterpriseID = (normalized == Self.defaultEnterpriseID) ? nil : normalized + } + self.logSecretUpdate(provider: .codebuddy, field: "enterpriseID", value: newValue) + } + } + + func ensureCodeBuddySessionLoaded() {} +} + +extension SettingsStore { + func codebuddySettingsSnapshot( + tokenOverride: TokenAccountOverride?) -> ProviderSettingsSnapshot.CodeBuddyProviderSettings + { + _ = tokenOverride + self.ensureCodeBuddySessionLoaded() + return ProviderSettingsSnapshot.CodeBuddyProviderSettings( + cookieSource: self.codebuddyCookieSource, + manualCookieHeader: self.codebuddyManualCookieHeader, + enterpriseID: self.codebuddyEnterpriseID.isEmpty ? nil : self.codebuddyEnterpriseID) + } +} diff --git a/Sources/CodexBar/Providers/Shared/ProviderImplementationRegistry.swift b/Sources/CodexBar/Providers/Shared/ProviderImplementationRegistry.swift index f6a9b2a3..f1306275 100644 --- a/Sources/CodexBar/Providers/Shared/ProviderImplementationRegistry.swift +++ b/Sources/CodexBar/Providers/Shared/ProviderImplementationRegistry.swift @@ -30,6 +30,7 @@ enum ProviderImplementationRegistry { case .kimik2: KimiK2ProviderImplementation() case .amp: AmpProviderImplementation() case .synthetic: SyntheticProviderImplementation() + case .codebuddy: CodeBuddyProviderImplementation() } } diff --git a/Sources/CodexBar/StatusItemController+Menu.swift b/Sources/CodexBar/StatusItemController+Menu.swift index 2167b291..3b3b0cfe 100644 --- a/Sources/CodexBar/StatusItemController+Menu.swift +++ b/Sources/CodexBar/StatusItemController+Menu.swift @@ -4,6 +4,8 @@ import Observation import QuartzCore import SwiftUI +private let codeBuddyMenuLog = CodexBarLog.logger("codebuddy-menu") + // MARK: - NSMenu construction extension StatusItemController { @@ -325,6 +327,18 @@ extension StatusItemController { return true } + // For CodeBuddy, add the menu card with daily usage submenu + if context.currentProvider == .codebuddy { + let submenu = self.makeCodeBuddyDailyUsageSubmenu() + menu.addItem(self.makeMenuCardItem( + UsageMenuCardView(model: model, width: context.menuWidth), + id: "menuCard", + width: context.menuWidth, + submenu: submenu)) + menu.addItem(.separator()) + return false + } + menu.addItem(self.makeMenuCardItem( UsageMenuCardView(model: model, width: context.menuWidth), id: "menuCard", @@ -668,6 +682,7 @@ extension StatusItemController { menu.addItem(self.makeMenuCardItem(headerView, id: "menuCardHeader", width: width)) if hasUsageBlock { + codeBuddyMenuLog.debug("hasUsageBlock=true for provider=\(provider.rawValue)") let usageView = UsageMenuCardUsageSectionView( model: model, showBottomDivider: false, @@ -950,12 +965,16 @@ extension StatusItemController { snapshot: UsageSnapshot?, webItems: OpenAIWebMenuItems) -> NSMenu? { + codeBuddyMenuLog.debug("makeUsageSubmenu called for provider=\(provider.rawValue)") if provider == .codex, webItems.hasUsageBreakdown { return self.makeUsageBreakdownSubmenu() } if provider == .zai { return self.makeZaiUsageDetailsSubmenu(snapshot: snapshot) } + if provider == .codebuddy { + return self.makeCodeBuddyDailyUsageSubmenu() + } return nil } @@ -1059,6 +1078,39 @@ extension StatusItemController { return submenu } + private func makeCodeBuddyDailyUsageSubmenu() -> NSMenu? { + let width = Self.menuCardBaseWidth + let count = self.store.codeBuddyDailyUsage?.count ?? -1 + codeBuddyMenuLog.info("makeCodeBuddyDailyUsageSubmenu: dailyUsage count = \(count)") + guard let dailyUsage = self.store.codeBuddyDailyUsage, !dailyUsage.isEmpty else { return nil } + + if !Self.menuCardRenderingEnabled { + let submenu = NSMenu() + submenu.delegate = self + let chartItem = NSMenuItem() + chartItem.isEnabled = false + chartItem.representedObject = "codeBuddyDailyUsageChart" + submenu.addItem(chartItem) + return submenu + } + + let submenu = NSMenu() + submenu.delegate = self + let chartView = CodeBuddyDailyUsageChartMenuView(dailyUsage: dailyUsage, width: width) + let hosting = MenuHostingView(rootView: chartView) + // Use NSHostingController for efficient size calculation without multiple layout passes + let controller = NSHostingController(rootView: chartView) + let size = controller.sizeThatFits(in: CGSize(width: width, height: .greatestFiniteMagnitude)) + hosting.frame = NSRect(origin: .zero, size: NSSize(width: width, height: size.height)) + + let chartItem = NSMenuItem() + chartItem.view = hosting + chartItem.isEnabled = false + chartItem.representedObject = "codeBuddyDailyUsageChart" + submenu.addItem(chartItem) + return submenu + } + private func makeCostHistorySubmenu(provider: UsageProvider) -> NSMenu? { guard provider == .codex || provider == .claude || provider == .vertexai else { return nil } let width = Self.menuCardBaseWidth @@ -1101,6 +1153,7 @@ extension StatusItemController { "usageBreakdownChart", "creditsHistoryChart", "costHistoryChart", + "codeBuddyDailyUsageChart", ] return menu.items.contains { item in guard let id = item.representedObject as? String else { return false } diff --git a/Sources/CodexBar/UsageStore+Refresh.swift b/Sources/CodexBar/UsageStore+Refresh.swift index 5449d67b..ea55dd67 100644 --- a/Sources/CodexBar/UsageStore+Refresh.swift +++ b/Sources/CodexBar/UsageStore+Refresh.swift @@ -1,6 +1,8 @@ import CodexBarCore import Foundation +private let codeBuddyLog = CodexBarLog.logger("codebuddy-store") + extension UsageStore { /// Force refresh Augment session (called from UI button) func forceRefreshAugmentSession() async { @@ -50,12 +52,22 @@ extension UsageStore { switch outcome.result { case let .success(result): let scoped = result.usage.scoped(to: provider) + // Debug: log CodeBuddy daily usage + if provider == .codebuddy { + let count = result.codeBuddyDailyUsage?.count ?? -1 + codeBuddyLog.info("refreshProvider: result.codeBuddyDailyUsage count = \(count)") + } await MainActor.run { self.handleSessionQuotaTransition(provider: provider, snapshot: scoped) self.snapshots[provider] = scoped self.lastSourceLabels[provider] = result.sourceLabel self.errors[provider] = nil self.failureGates[provider]?.recordSuccess() + // Store CodeBuddy daily usage if available + if provider == .codebuddy, let dailyUsage = result.codeBuddyDailyUsage { + self.codeBuddyDailyUsage = dailyUsage + codeBuddyLog.info("Stored daily usage: \(dailyUsage.count) entries") + } } if let runtime = self.providerRuntimes[provider] { let context = ProviderRuntimeContext(provider: provider, settings: self.settings, store: self) diff --git a/Sources/CodexBar/UsageStore+TokenAccounts.swift b/Sources/CodexBar/UsageStore+TokenAccounts.swift index 7127a123..2aeb98fa 100644 --- a/Sources/CodexBar/UsageStore+TokenAccounts.swift +++ b/Sources/CodexBar/UsageStore+TokenAccounts.swift @@ -161,6 +161,10 @@ extension UsageStore { self.lastSourceLabels[provider] = result.sourceLabel self.errors[provider] = nil self.failureGates[provider]?.recordSuccess() + // Store CodeBuddy daily usage if available + if provider == .codebuddy, let dailyUsage = result.codeBuddyDailyUsage { + self.codeBuddyDailyUsage = dailyUsage + } } case let .failure(error): await MainActor.run { diff --git a/Sources/CodexBar/UsageStore.swift b/Sources/CodexBar/UsageStore.swift index 69491056..c59a79c4 100644 --- a/Sources/CodexBar/UsageStore.swift +++ b/Sources/CodexBar/UsageStore.swift @@ -24,6 +24,7 @@ extension UsageStore { _ = self.openAIDashboardRequiresLogin _ = self.openAIDashboardCookieImportStatus _ = self.openAIDashboardCookieImportDebugLog + _ = self.codeBuddyDailyUsage _ = self.versions _ = self.isRefreshing _ = self.refreshingProviders @@ -154,6 +155,7 @@ final class UsageStore { var openAIDashboardRequiresLogin: Bool = false var openAIDashboardCookieImportStatus: String? var openAIDashboardCookieImportDebugLog: String? + var codeBuddyDailyUsage: [CodeBuddyDailyUsageEntry]? var versions: [UsageProvider: String] = [:] var isRefreshing = false var refreshingProviders: Set = [] @@ -1229,6 +1231,13 @@ extension UsageStore { let text = "JetBrains AI debug log not yet implemented" await MainActor.run { self.probeLogs[.jetbrains] = text } return text + case .codebuddy: + let text = await self.debugCodeBuddyLog( + codebuddyCookieSource: self.settings.codebuddyCookieSource, + codebuddyCookieHeader: self.settings.codebuddyManualCookieHeader, + codebuddyEnterpriseID: self.settings.codebuddyEnterpriseID) + await MainActor.run { self.probeLogs[.codebuddy] = text } + return text } }.value } @@ -1389,6 +1398,108 @@ extension UsageStore { } } + private func debugCodeBuddyLog( + codebuddyCookieSource: ProviderCookieSource, + codebuddyCookieHeader: String, + codebuddyEnterpriseID: String) async -> String + { + await self.runWithTimeout(seconds: 15) { + var lines: [String] = [] + + lines.append("CodeBuddy Debug Log") + lines.append("===================") + lines.append("") + lines.append("Settings:") + lines.append(" cookieSource=\(codebuddyCookieSource.rawValue)") + lines.append(" enterpriseID=\(codebuddyEnterpriseID.isEmpty ? "" : codebuddyEnterpriseID)") + lines.append(" manualCookieHeader=\(codebuddyCookieHeader.isEmpty ? "" : "<\(codebuddyCookieHeader.count) chars>")") + + // Check for manual cookie override + if codebuddyCookieSource == .manual { + if let normalized = CookieHeaderNormalizer.normalize(codebuddyCookieHeader) { + lines.append("") + lines.append("Manual cookie header (normalized):") + lines.append(" length=\(normalized.count)") + let hasSession = normalized.contains("session=") + lines.append(" has_session_cookie=\(hasSession)") + } else { + lines.append("") + lines.append("Manual cookie header: invalid/empty") + } + } + + // Check browser cookies + lines.append("") + lines.append("Browser Cookie Check:") + do { + let sessions = try CodeBuddyCookieImporter.importSessions( + browserDetection: self.browserDetection, + logger: { msg in lines.append(" \(msg)") }) + lines.append(" Found \(sessions.count) session(s)") + for (index, session) in sessions.enumerated() { + lines.append(" [\(index)] \(session.sourceLabel): hasSession=\(session.hasValidSession)") + } + } catch { + lines.append(" Cookie import error: \(error.localizedDescription)") + } + + // Try to make API call if we have credentials + let cookieHeader: String? + if codebuddyCookieSource == .manual, !codebuddyCookieHeader.isEmpty { + cookieHeader = CookieHeaderNormalizer.normalize(codebuddyCookieHeader) + } else { + cookieHeader = try? CodeBuddyCookieImporter.importSession( + browserDetection: self.browserDetection).cookieHeader + } + + if let cookieHeader, !codebuddyEnterpriseID.isEmpty { + lines.append("") + lines.append("API Test:") + do { + let snapshot = try await CodeBuddyUsageFetcher.fetchUsage( + cookieHeader: cookieHeader, + enterpriseID: codebuddyEnterpriseID) + lines.append(" SUCCESS") + lines.append(" creditUsed=\(snapshot.creditUsed)") + lines.append(" creditLimit=\(snapshot.creditLimit)") + lines.append(" cycleStartTime=\(snapshot.cycleStartTime)") + lines.append(" cycleEndTime=\(snapshot.cycleEndTime)") + lines.append(" cycleResetTime=\(snapshot.cycleResetTime)") + } catch { + lines.append(" FAILED: \(error.localizedDescription)") + } + } else { + lines.append("") + lines.append("API Test: skipped (missing cookies or enterpriseID)") + if cookieHeader == nil { + lines.append(" reason: no cookie header available") + } + if codebuddyEnterpriseID.isEmpty { + lines.append(" reason: enterpriseID not set") + } + } + + // Fetch attempts + lines.append("") + lines.append("Last Fetch Attempts:") + let attempts = await MainActor.run { self.lastFetchAttempts[.codebuddy] ?? [] } + if attempts.isEmpty { + lines.append(" ") + } else { + for attempt in attempts { + var line = " \(attempt.strategyID)" + line += attempt.wasAvailable ? " (available)" : " (unavailable)" + if let err = attempt.errorDescription { + line += " error=\(err)" + } + lines.append(line) + } + } + + return lines.joined(separator: "\n") + } + } + private func runWithTimeout(seconds: Double, operation: @escaping @Sendable () async -> String) async -> String { await withTaskGroup(of: String?.self) { group -> String in group.addTask { await operation() } diff --git a/Sources/CodexBarCLI/TokenAccountCLI.swift b/Sources/CodexBarCLI/TokenAccountCLI.swift index a0a88371..bf3c9e21 100644 --- a/Sources/CodexBarCLI/TokenAccountCLI.swift +++ b/Sources/CodexBarCLI/TokenAccountCLI.swift @@ -147,6 +147,12 @@ struct TokenAccountCLIContext { return self.makeSnapshot( jetbrains: ProviderSettingsSnapshot.JetBrainsProviderSettings( ideBasePath: nil)) + case .codebuddy: + return self.makeSnapshot( + codebuddy: ProviderSettingsSnapshot.CodeBuddyProviderSettings( + cookieSource: cookieSource, + manualCookieHeader: cookieHeader, + enterpriseID: config?.enterpriseID)) case .gemini, .antigravity, .copilot, .kiro, .vertexai, .kimik2, .synthetic: return nil } @@ -163,7 +169,8 @@ struct TokenAccountCLIContext { kimi: ProviderSettingsSnapshot.KimiProviderSettings? = nil, augment: ProviderSettingsSnapshot.AugmentProviderSettings? = nil, amp: ProviderSettingsSnapshot.AmpProviderSettings? = nil, - jetbrains: ProviderSettingsSnapshot.JetBrainsProviderSettings? = nil) -> ProviderSettingsSnapshot + jetbrains: ProviderSettingsSnapshot.JetBrainsProviderSettings? = nil, + codebuddy: ProviderSettingsSnapshot.CodeBuddyProviderSettings? = nil) -> ProviderSettingsSnapshot { ProviderSettingsSnapshot.make( codex: codex, @@ -176,7 +183,8 @@ struct TokenAccountCLIContext { kimi: kimi, augment: augment, amp: amp, - jetbrains: jetbrains) + jetbrains: jetbrains, + codebuddy: codebuddy) } func environment( diff --git a/Sources/CodexBarCore/Config/CodexBarConfig.swift b/Sources/CodexBarCore/Config/CodexBarConfig.swift index c4bbbc2c..088f8c07 100644 --- a/Sources/CodexBarCore/Config/CodexBarConfig.swift +++ b/Sources/CodexBarCore/Config/CodexBarConfig.swift @@ -81,6 +81,7 @@ public struct ProviderConfig: Codable, Sendable, Identifiable { public var cookieSource: ProviderCookieSource? public var region: String? public var workspaceID: String? + public var enterpriseID: String? public var tokenAccounts: ProviderTokenAccountData? public init( @@ -92,6 +93,7 @@ public struct ProviderConfig: Codable, Sendable, Identifiable { cookieSource: ProviderCookieSource? = nil, region: String? = nil, workspaceID: String? = nil, + enterpriseID: String? = nil, tokenAccounts: ProviderTokenAccountData? = nil) { self.id = id @@ -102,6 +104,7 @@ public struct ProviderConfig: Codable, Sendable, Identifiable { self.cookieSource = cookieSource self.region = region self.workspaceID = workspaceID + self.enterpriseID = enterpriseID self.tokenAccounts = tokenAccounts } diff --git a/Sources/CodexBarCore/Logging/LogCategories.swift b/Sources/CodexBarCore/Logging/LogCategories.swift index d3d31c8f..8949c265 100644 --- a/Sources/CodexBarCore/Logging/LogCategories.swift +++ b/Sources/CodexBarCore/Logging/LogCategories.swift @@ -9,6 +9,10 @@ public enum LogCategories { public static let claudeCLI = "claude-cli" public static let claudeProbe = "claude-probe" public static let claudeUsage = "claude-usage" + public static let codeBuddyAPI = "codebuddy-api" + public static let codeBuddyCookie = "codebuddy-cookie" + public static let codeBuddyKeepalive = "codebuddy-keepalive" + public static let codeBuddyWeb = "codebuddy-web" public static let codexRPC = "codex-rpc" public static let configMigration = "config-migration" public static let configStore = "config-store" diff --git a/Sources/CodexBarCore/Providers/CodeBuddy/CodeBuddyAPIError.swift b/Sources/CodexBarCore/Providers/CodeBuddy/CodeBuddyAPIError.swift new file mode 100644 index 00000000..ce42155c --- /dev/null +++ b/Sources/CodexBarCore/Providers/CodeBuddy/CodeBuddyAPIError.swift @@ -0,0 +1,30 @@ +import Foundation + +public enum CodeBuddyAPIError: LocalizedError, Sendable, Equatable { + case missingCookies + case missingEnterpriseID + case invalidCookies + case invalidRequest(String) + case networkError(String) + case apiError(String) + case parseFailed(String) + + public var errorDescription: String? { + switch self { + case .missingCookies: + "CodeBuddy cookies are missing. Please import cookies from your browser." + case .missingEnterpriseID: + "CodeBuddy enterprise ID is missing. Please sign in to CodeBuddy dashboard." + case .invalidCookies: + "CodeBuddy cookies are invalid or expired. Please re-import cookies." + case let .invalidRequest(message): + "Invalid request: \(message)" + case let .networkError(message): + "CodeBuddy network error: \(message)" + case let .apiError(message): + "CodeBuddy API error: \(message)" + case let .parseFailed(message): + "Failed to parse CodeBuddy usage data: \(message)" + } + } +} diff --git a/Sources/CodexBarCore/Providers/CodeBuddy/CodeBuddyCookieHeader.swift b/Sources/CodexBarCore/Providers/CodeBuddy/CodeBuddyCookieHeader.swift new file mode 100644 index 00000000..948ba349 --- /dev/null +++ b/Sources/CodexBarCore/Providers/CodeBuddy/CodeBuddyCookieHeader.swift @@ -0,0 +1,120 @@ +import Foundation + +public struct CodeBuddyCookieOverride: Sendable { + public let cookieHeader: String + public let enterpriseID: String? + + public init(cookieHeader: String, enterpriseID: String? = nil) { + self.cookieHeader = cookieHeader + self.enterpriseID = enterpriseID + } +} + +public enum CodeBuddyCookieHeader { + private static let log = CodexBarLog.logger(LogCategories.codeBuddyCookie) + private static let headerPatterns: [String] = [ + #"(?i)session=([A-Za-z0-9._\-+=/|]+)"#, + #"(?i)-H\s*'Cookie:\s*([^']+)'"#, + #"(?i)-H\s*"Cookie:\s*([^"]+)""#, + #"(?i)\bcookie:\s*'([^']+)'"#, + #"(?i)\bcookie:\s*"([^"]+)""#, + #"(?i)\bcookie:\s*([^\r\n]+)"#, + ] + + public static func resolveCookieOverride(context: ProviderFetchContext) -> CodeBuddyCookieOverride? { + if let settings = context.settings?.codebuddy, settings.cookieSource == .manual { + if let manual = settings.manualCookieHeader, !manual.isEmpty { + return self.override(from: manual, enterpriseID: settings.enterpriseID) + } + } + + if let envCookie = self.override(from: context.env["CODEBUDDY_COOKIE"]) { + return envCookie + } + + return nil + } + + public static func override(from raw: String?, enterpriseID: String? = nil) -> CodeBuddyCookieOverride? { + guard let raw = raw?.trimmingCharacters(in: .whitespacesAndNewlines), !raw.isEmpty else { + return nil + } + + // Try to extract session cookie from the raw input + if let sessionValue = self.extractSessionCookie(from: raw) { + // If raw already looks like a cookie header, use it directly + if raw.contains("session=") { + return CodeBuddyCookieOverride(cookieHeader: raw, enterpriseID: enterpriseID) + } + // Otherwise construct a minimal header + return CodeBuddyCookieOverride(cookieHeader: "session=\(sessionValue)", enterpriseID: enterpriseID) + } + + // Try extracting from curl command or header format + if let cookieHeader = self.extractHeader(from: raw) { + return CodeBuddyCookieOverride(cookieHeader: cookieHeader, enterpriseID: enterpriseID) + } + + return nil + } + + private static func extractSessionCookie(from raw: String) -> String? { + let patterns = [ + #"(?i)session=([A-Za-z0-9._\-+=/|]+)"#, + ] + + for pattern in patterns { + guard let regex = try? NSRegularExpression(pattern: pattern, options: []) else { continue } + let range = NSRange(raw.startIndex..= 2, + let captureRange = Range(match.range(at: 1), in: raw) + else { + continue + } + let token = String(raw[captureRange]).trimmingCharacters(in: .whitespacesAndNewlines) + if !token.isEmpty { return token } + } + + return nil + } + + private static func extractHeader(from raw: String) -> String? { + for pattern in self.headerPatterns { + guard let regex = try? NSRegularExpression(pattern: pattern, options: []) else { continue } + let range = NSRange(raw.startIndex..= 2, + let captureRange = Range(match.range(at: 1), in: raw) + else { + continue + } + let captured = String(raw[captureRange]).trimmingCharacters(in: .whitespacesAndNewlines) + if !captured.isEmpty { return captured } + } + return nil + } + + /// Extract enterprise ID from the x-enterprise-id header in a curl command + public static func extractEnterpriseID(from raw: String) -> String? { + let patterns = [ + #"(?i)x-enterprise-id:\s*([A-Za-z0-9]+)"#, + #"(?i)-H\s*'x-enterprise-id:\s*([A-Za-z0-9]+)'"#, + #"(?i)-H\s*"x-enterprise-id:\s*([A-Za-z0-9]+)""#, + ] + + for pattern in patterns { + guard let regex = try? NSRegularExpression(pattern: pattern, options: []) else { continue } + let range = NSRange(raw.startIndex..= 2, + let captureRange = Range(match.range(at: 1), in: raw) + else { + continue + } + let id = String(raw[captureRange]).trimmingCharacters(in: .whitespacesAndNewlines) + if !id.isEmpty { return id } + } + return nil + } +} diff --git a/Sources/CodexBarCore/Providers/CodeBuddy/CodeBuddyCookieImporter.swift b/Sources/CodexBarCore/Providers/CodeBuddy/CodeBuddyCookieImporter.swift new file mode 100644 index 00000000..b83c22b3 --- /dev/null +++ b/Sources/CodexBarCore/Providers/CodeBuddy/CodeBuddyCookieImporter.swift @@ -0,0 +1,200 @@ +import Foundation + +#if os(macOS) +import SweetCookieKit + +public enum CodeBuddyCookieImporter { + private static let log = CodexBarLog.logger(LogCategories.codeBuddyCookie) + private static let cookieClient = BrowserCookieClient() + private static let cookieDomains = ["tencent.sso.codebuddy.cn"] + private static let cookieImportOrder: BrowserCookieImportOrder = + ProviderDefaults.metadata[.codebuddy]?.browserCookieOrder ?? Browser.defaultImportOrder + + public struct SessionInfo: Sendable { + public let cookies: [HTTPCookie] + public let sourceLabel: String + + public init(cookies: [HTTPCookie], sourceLabel: String) { + self.cookies = cookies + self.sourceLabel = sourceLabel + } + + /// The session cookie value + public var sessionCookie: String? { + self.cookies.first(where: { $0.name == "session" })?.value + } + + /// The session_2 cookie value + public var session2Cookie: String? { + self.cookies.first(where: { $0.name == "session_2" })?.value + } + + /// Build a cookie header string from all cookies + public var cookieHeader: String { + self.cookies.map { "\($0.name)=\($0.value)" }.joined(separator: "; ") + } + + /// Check if the session has valid auth cookies + public var hasValidSession: Bool { + self.sessionCookie != nil + } + } + + public static func importSessions( + browserDetection: BrowserDetection = BrowserDetection(), + logger: ((String) -> Void)? = nil) throws -> [SessionInfo] + { + var sessions: [SessionInfo] = [] + let candidates = self.cookieImportOrder.cookieImportCandidates(using: browserDetection) + for browserSource in candidates { + do { + let perSource = try self.importSessions(from: browserSource, logger: logger) + sessions.append(contentsOf: perSource) + } catch { + BrowserCookieAccessGate.recordIfNeeded(error) + self.emit( + "\(browserSource.displayName) cookie import failed: \(error.localizedDescription)", + logger: logger) + } + } + + guard !sessions.isEmpty else { + throw CodeBuddyCookieImportError.noCookies + } + return sessions + } + + public static func importSessions( + from browserSource: Browser, + logger: ((String) -> Void)? = nil) throws -> [SessionInfo] + { + let query = BrowserCookieQuery(domains: self.cookieDomains) + let log: (String) -> Void = { msg in self.emit(msg, logger: logger) } + let sources = try Self.cookieClient.records( + matching: query, + in: browserSource, + logger: log) + + var sessions: [SessionInfo] = [] + let grouped = Dictionary(grouping: sources, by: { $0.store.profile.id }) + let sortedGroups = grouped.values.sorted { lhs, rhs in + self.mergedLabel(for: lhs) < self.mergedLabel(for: rhs) + } + + for group in sortedGroups where !group.isEmpty { + let label = self.mergedLabel(for: group) + let mergedRecords = self.mergeRecords(group) + guard !mergedRecords.isEmpty else { continue } + let httpCookies = BrowserCookieClient.makeHTTPCookies(mergedRecords, origin: query.origin) + guard !httpCookies.isEmpty else { continue } + + // Only include sessions that have the session cookie + guard httpCookies.contains(where: { $0.name == "session" }) else { + continue + } + + log("Found session cookie in \(label)") + sessions.append(SessionInfo(cookies: httpCookies, sourceLabel: label)) + } + return sessions + } + + public static func importSession( + browserDetection: BrowserDetection = BrowserDetection(), + logger: ((String) -> Void)? = nil) throws -> SessionInfo + { + let sessions = try self.importSessions(browserDetection: browserDetection, logger: logger) + guard let first = sessions.first else { + throw CodeBuddyCookieImportError.noCookies + } + return first + } + + public static func hasSession( + browserDetection: BrowserDetection = BrowserDetection(), + logger: ((String) -> Void)? = nil) -> Bool + { + do { + return try !self.importSessions(browserDetection: browserDetection, logger: logger).isEmpty + } catch { + return false + } + } + + private static func cookieNames(from cookies: [HTTPCookie]) -> String { + let names = Set(cookies.map { "\($0.name)@\($0.domain)" }).sorted() + return names.joined(separator: ", ") + } + + private static func emit(_ message: String, logger: ((String) -> Void)?) { + logger?("[codebuddy-cookie] \(message)") + self.log.debug(message) + } + + private static func mergedLabel(for sources: [BrowserCookieStoreRecords]) -> String { + guard let base = sources.map(\.label).min() else { + return "Unknown" + } + if base.hasSuffix(" (Network)") { + return String(base.dropLast(" (Network)".count)) + } + return base + } + + private static func mergeRecords(_ sources: [BrowserCookieStoreRecords]) -> [BrowserCookieRecord] { + let sortedSources = sources.sorted { lhs, rhs in + self.storePriority(lhs.store.kind) < self.storePriority(rhs.store.kind) + } + var mergedByKey: [String: BrowserCookieRecord] = [:] + for source in sortedSources { + for record in source.records { + let key = self.recordKey(record) + if let existing = mergedByKey[key] { + if self.shouldReplace(existing: existing, candidate: record) { + mergedByKey[key] = record + } + } else { + mergedByKey[key] = record + } + } + } + return Array(mergedByKey.values) + } + + private static func storePriority(_ kind: BrowserCookieStoreKind) -> Int { + switch kind { + case .network: 0 + case .primary: 1 + case .safari: 2 + } + } + + private static func recordKey(_ record: BrowserCookieRecord) -> String { + "\(record.name)|\(record.domain)|\(record.path)" + } + + private static func shouldReplace(existing: BrowserCookieRecord, candidate: BrowserCookieRecord) -> Bool { + switch (existing.expires, candidate.expires) { + case let (lhs?, rhs?): + rhs > lhs + case (nil, .some): + true + case (.some, nil): + false + case (nil, nil): + false + } + } +} + +enum CodeBuddyCookieImportError: LocalizedError { + case noCookies + + var errorDescription: String? { + switch self { + case .noCookies: + "No CodeBuddy session cookies found in browsers." + } + } +} +#endif diff --git a/Sources/CodexBarCore/Providers/CodeBuddy/CodeBuddyModels.swift b/Sources/CodexBarCore/Providers/CodeBuddy/CodeBuddyModels.swift new file mode 100644 index 00000000..ca3c24fe --- /dev/null +++ b/Sources/CodexBarCore/Providers/CodeBuddy/CodeBuddyModels.swift @@ -0,0 +1,46 @@ +import Foundation + +/// API response for get-enterprise-user-usage endpoint +struct CodeBuddyUsageResponse: Codable { + let code: Int + let msg: String + let requestId: String + let data: CodeBuddyUsageData +} + +struct CodeBuddyUsageData: Codable { + let credit: Double + let cycleStartTime: String + let cycleEndTime: String + let limitNum: Double + let cycleResetTime: String +} + +/// API response for get-user-daily-usage endpoint (for future use) +struct CodeBuddyDailyUsageResponse: Codable { + let code: Int + let msg: String + let requestId: String + let data: CodeBuddyDailyUsageData +} + +struct CodeBuddyDailyUsageData: Codable { + let total: Int + let data: [CodeBuddyDailyUsage] +} + +struct CodeBuddyDailyUsage: Codable { + let credit: Double + let date: String +} + +/// Public daily usage entry for external access +public struct CodeBuddyDailyUsageEntry: Sendable { + public let date: String // "yyyy-MM-dd" format + public let credit: Double + + public init(date: String, credit: Double) { + self.date = date + self.credit = credit + } +} diff --git a/Sources/CodexBarCore/Providers/CodeBuddy/CodeBuddyProviderDescriptor.swift b/Sources/CodexBarCore/Providers/CodeBuddy/CodeBuddyProviderDescriptor.swift new file mode 100644 index 00000000..79852cb4 --- /dev/null +++ b/Sources/CodexBarCore/Providers/CodeBuddy/CodeBuddyProviderDescriptor.swift @@ -0,0 +1,153 @@ +import CodexBarMacroSupport +import Foundation + +@ProviderDescriptorRegistration +@ProviderDescriptorDefinition +public enum CodeBuddyProviderDescriptor { + static func makeDescriptor() -> ProviderDescriptor { + ProviderDescriptor( + id: .codebuddy, + metadata: ProviderMetadata( + id: .codebuddy, + displayName: "CodeBuddy", + sessionLabel: "Credits", + weeklyLabel: "Cycle", + opusLabel: nil, + supportsOpus: false, + supportsCredits: true, + creditsHint: "Credits used this billing cycle", + toggleTitle: "Show CodeBuddy usage", + cliName: "codebuddy", + defaultEnabled: false, + isPrimaryProvider: false, + usesAccountFallback: false, + browserCookieOrder: nil, + dashboardURL: "https://tencent.sso.codebuddy.cn/profile/usage", + statusPageURL: nil), + branding: ProviderBranding( + iconStyle: .codebuddy, + iconResourceName: "ProviderIcon-codebuddy", + color: ProviderColor(red: 0 / 255, green: 120 / 255, blue: 215 / 255)), + tokenCost: ProviderTokenCostConfig( + supportsTokenCost: false, + noDataMessage: { "CodeBuddy cost summary is not supported." }), + fetchPlan: ProviderFetchPlan( + sourceModes: [.auto, .web], + pipeline: ProviderFetchPipeline(resolveStrategies: { _ in [CodeBuddyWebFetchStrategy()] })), + cli: ProviderCLIConfig( + name: "codebuddy", + aliases: ["cb"], + versionDetector: nil)) + } +} + +struct CodeBuddyWebFetchStrategy: ProviderFetchStrategy { + let id: String = "codebuddy.web" + let kind: ProviderFetchKind = .web + private static let log = CodexBarLog.logger(LogCategories.codeBuddyWeb) + + func isAvailable(_ context: ProviderFetchContext) async -> Bool { + // Check if we have manual cookie override + if CodeBuddyCookieHeader.resolveCookieOverride(context: context) != nil { + // Also need enterprise ID to be available + if context.settings?.codebuddy?.enterpriseID != nil { + return true + } + if CodeBuddySettingsReader.enterpriseID(environment: context.env) != nil { + return true + } + } + + #if os(macOS) + if context.settings?.codebuddy?.cookieSource != .off { + // Need both cookies and enterprise ID + if CodeBuddyCookieImporter.hasSession() { + if context.settings?.codebuddy?.enterpriseID != nil { + return true + } + if CodeBuddySettingsReader.enterpriseID(environment: context.env) != nil { + return true + } + } + } + #endif + + return false + } + + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + let (cookieHeader, enterpriseID) = try self.resolveCookiesAndEnterpriseID(context: context) + + let snapshot = try await CodeBuddyUsageFetcher.fetchUsage( + cookieHeader: cookieHeader, + enterpriseID: enterpriseID) + + // Also fetch daily usage (non-blocking, failure doesn't affect main result) + var dailyUsage: [CodeBuddyDailyUsageEntry]? + do { + dailyUsage = try await CodeBuddyUsageFetcher.fetchDailyUsage( + cookieHeader: cookieHeader, + enterpriseID: enterpriseID) + Self.log.info("CodeBuddy daily usage fetched: \(dailyUsage?.count ?? 0) entries") + } catch { + Self.log.error("CodeBuddy daily usage fetch failed: \(error.localizedDescription)") + } + + let result = self.makeResult( + usage: snapshot.toUsageSnapshot(), + codeBuddyDailyUsage: dailyUsage, + sourceLabel: "web") + Self.log.info("CodeBuddy fetch result created: dailyUsage in result = \(result.codeBuddyDailyUsage?.count ?? -1)") + return result + } + + func shouldFallback(on error: Error, context: ProviderFetchContext) -> Bool { + if case CodeBuddyAPIError.missingCookies = error { return false } + if case CodeBuddyAPIError.missingEnterpriseID = error { return false } + if case CodeBuddyAPIError.invalidCookies = error { return false } + return true + } + + private func resolveCookiesAndEnterpriseID(context: ProviderFetchContext) throws -> (String, String) { + var cookieHeader: String? + var enterpriseID: String? + + // Check manual override first + if let override = CodeBuddyCookieHeader.resolveCookieOverride(context: context) { + cookieHeader = override.cookieHeader + enterpriseID = override.enterpriseID + } + + // Try browser cookie import + #if os(macOS) + if cookieHeader == nil, context.settings?.codebuddy?.cookieSource != .off { + do { + let session = try CodeBuddyCookieImporter.importSession() + cookieHeader = session.cookieHeader + } catch { + // No browser cookies found + } + } + #endif + + // Check settings for enterprise ID + if enterpriseID == nil { + enterpriseID = context.settings?.codebuddy?.enterpriseID + } + + // Check environment for enterprise ID + if enterpriseID == nil { + enterpriseID = CodeBuddySettingsReader.enterpriseID(environment: context.env) + } + + guard let finalCookieHeader = cookieHeader, !finalCookieHeader.isEmpty else { + throw CodeBuddyAPIError.missingCookies + } + + guard let finalEnterpriseID = enterpriseID, !finalEnterpriseID.isEmpty else { + throw CodeBuddyAPIError.missingEnterpriseID + } + + return (finalCookieHeader, finalEnterpriseID) + } +} diff --git a/Sources/CodexBarCore/Providers/CodeBuddy/CodeBuddySessionKeepalive.swift b/Sources/CodexBarCore/Providers/CodeBuddy/CodeBuddySessionKeepalive.swift new file mode 100644 index 00000000..8c5592e5 --- /dev/null +++ b/Sources/CodexBarCore/Providers/CodeBuddy/CodeBuddySessionKeepalive.swift @@ -0,0 +1,385 @@ +import Foundation + +#if os(macOS) +import AppKit +import UserNotifications + +/// Manages automatic session keepalive for CodeBuddy to prevent cookie expiration. +/// +/// This class monitors cookie expiration and proactively refreshes the session +/// before cookies expire, ensuring uninterrupted access to CodeBuddy APIs. +@MainActor +public final class CodeBuddySessionKeepalive { + // MARK: - Configuration + + /// How often to check if session needs refresh (default: 2 minutes) + private let checkInterval: TimeInterval = 120 + + /// Refresh session this many seconds before cookie expiration (default: 10 minutes) + private let refreshBufferSeconds: TimeInterval = 600 + + /// Minimum time between refresh attempts (default: 2 minutes) + private let minRefreshInterval: TimeInterval = 120 + + /// Maximum time to wait for session refresh (default: 30 seconds) + private let refreshTimeout: TimeInterval = 30 + + // MARK: - State + + private var timerTask: Task? + private var lastRefreshAttempt: Date? + private var lastSuccessfulRefresh: Date? + private var isRefreshing = false + private let logger: ((String) -> Void)? + private var onSessionRecovered: (() async -> Void)? + + /// Track consecutive failures to stop retrying after too many failures + private var consecutiveFailures = 0 + private let maxConsecutiveFailures = 3 + private var hasGivenUp = false + + private static let log = CodexBarLog.logger(LogCategories.codeBuddyKeepalive) + + // MARK: - Initialization + + public init(logger: ((String) -> Void)? = nil, onSessionRecovered: (() async -> Void)? = nil) { + self.logger = logger + self.onSessionRecovered = onSessionRecovered + } + + deinit { + self.timerTask?.cancel() + } + + // MARK: - Public API + + /// Start the automatic session keepalive timer + public func start() { + guard self.timerTask == nil else { + self.log("Keepalive already running") + return + } + + self.log("Starting CodeBuddy session keepalive") + self.log(" - Check interval: \(Int(self.checkInterval))s") + self.log(" - Refresh buffer: \(Int(self.refreshBufferSeconds))s before expiry") + self.log(" - Min refresh interval: \(Int(self.minRefreshInterval))s") + + self.timerTask = Task.detached(priority: .utility) { [weak self] in + while !Task.isCancelled { + try? await Task.sleep(for: .seconds(self?.checkInterval ?? 120)) + await self?.checkAndRefreshIfNeeded() + } + } + + self.log("Keepalive timer started successfully") + } + + /// Stop the automatic session keepalive timer + public func stop() { + self.log("Stopping CodeBuddy session keepalive") + self.timerTask?.cancel() + self.timerTask = nil + } + + /// Manually trigger a session refresh (bypasses rate limiting) + public func forceRefresh() async { + self.log("Force refresh requested") + await self.performRefresh(forced: true) + } + + // MARK: - Private Implementation + + private func checkAndRefreshIfNeeded() async { + guard !self.isRefreshing else { + self.log("Refresh already in progress, skipping check") + return + } + + // Stop trying if we've given up + if self.hasGivenUp { + self.log("Keepalive has given up after \(self.maxConsecutiveFailures) consecutive failures") + return + } + + // Rate limit: don't refresh too frequently + if let lastAttempt = self.lastRefreshAttempt { + let timeSinceLastAttempt = Date().timeIntervalSince(lastAttempt) + if timeSinceLastAttempt < self.minRefreshInterval { + return + } + } + + // Check if cookies are about to expire + let shouldRefresh = await self.shouldRefreshSession() + if shouldRefresh { + await self.performRefresh(forced: false) + } + } + + private func shouldRefreshSession() async -> Bool { + do { + let session = try CodeBuddyCookieImporter.importSession() + + self.log("Cookie Status Check:") + self.log(" Total cookies: \(session.cookies.count)") + self.log(" Source: \(session.sourceLabel)") + + // Find the earliest expiration date among session cookies + let expirationDates = session.cookies.compactMap(\.expiresDate) + + guard !expirationDates.isEmpty else { + // Session cookies (no expiration) - refresh periodically + self.log(" All cookies are session cookies (no expiration dates)") + if let lastRefresh = self.lastSuccessfulRefresh { + let timeSinceRefresh = Date().timeIntervalSince(lastRefresh) + // Refresh every 20 minutes for session cookies + if timeSinceRefresh > 1200 { + self.log(" Need periodic refresh (\(Int(timeSinceRefresh))s since last refresh)") + return true + } else { + self.log(" Recently refreshed (\(Int(timeSinceRefresh))s ago)") + return false + } + } else { + // Never refreshed - do it now + self.log(" Never refreshed - doing initial refresh") + return true + } + } + + let earliestExpiration = expirationDates.min()! + let timeUntilExpiration = earliestExpiration.timeIntervalSinceNow + + if timeUntilExpiration < self.refreshBufferSeconds { + self.log(" REFRESH NEEDED: expires in \(Int(timeUntilExpiration))s") + return true + } else { + self.log(" Session healthy: expires in \(Int(timeUntilExpiration))s") + return false + } + } catch { + self.log("Failed to check session: \(error.localizedDescription)") + return false + } + } + + private func performRefresh(forced: Bool) async { + self.isRefreshing = true + self.lastRefreshAttempt = Date() + defer { self.isRefreshing = false } + + self.log(forced ? "Performing forced session refresh..." : "Performing automatic session refresh...") + + // If this is a forced refresh, reset failure tracking + if forced { + self.consecutiveFailures = 0 + self.hasGivenUp = false + } + + do { + // Step 1: Ping the session endpoint to trigger cookie refresh + let refreshed = try await self.pingSessionEndpoint() + + if refreshed { + // Step 2: Re-import cookies from browser + try await Task.sleep(for: .seconds(1)) + let newSession = try CodeBuddyCookieImporter.importSession() + + self.log("Session refresh successful - imported \(newSession.cookies.count) cookies") + self.lastSuccessfulRefresh = Date() + self.consecutiveFailures = 0 + self.hasGivenUp = false + + // Notify callback + if let callback = self.onSessionRecovered { + await callback() + } + } else { + self.log("Session refresh returned no new cookies") + self.consecutiveFailures += 1 + self.checkIfShouldGiveUp() + } + } catch CodeBuddySessionKeepaliveError.sessionExpired { + self.log("Session expired - attempting automatic recovery...") + self.consecutiveFailures += 1 + + if self.consecutiveFailures >= self.maxConsecutiveFailures { + self.log("Too many consecutive failures - giving up") + self.hasGivenUp = true + self.notifyUserLoginRequired() + } else { + await self.attemptSessionRecovery() + } + } catch { + self.log("Session refresh failed: \(error.localizedDescription)") + self.consecutiveFailures += 1 + self.checkIfShouldGiveUp() + } + } + + private func checkIfShouldGiveUp() { + if self.consecutiveFailures >= self.maxConsecutiveFailures { + self.log("Too many consecutive failures - giving up") + self.hasGivenUp = true + self.notifyUserLoginRequired() + } + } + + /// Attempt to recover from an expired session by opening the dashboard + private func attemptSessionRecovery() async { + self.log("Attempting automatic session recovery...") + + #if os(macOS) + // Open the CodeBuddy dashboard in the default browser + if let url = URL(string: "https://tencent.sso.codebuddy.cn/profile/usage") { + _ = await MainActor.run { + NSWorkspace.shared.open(url) + } + self.log("Opened CodeBuddy dashboard in browser") + + // Wait for browser to potentially re-authenticate + try? await Task.sleep(for: .seconds(5)) + + // Try to import cookies again + do { + let newSession = try CodeBuddyCookieImporter.importSession() + self.log("Session recovery successful - imported \(newSession.cookies.count) cookies") + self.lastSuccessfulRefresh = Date() + + // Verify the session is actually valid + let isValid = try await self.pingSessionEndpoint() + if isValid { + self.log("Session verified - recovery complete!") + if let callback = self.onSessionRecovered { + await callback() + } + } else { + self.log("Session imported but not yet valid - may need manual login") + self.notifyUserLoginRequired() + } + } catch { + self.log("Session recovery failed: \(error.localizedDescription)") + self.notifyUserLoginRequired() + } + } + #endif + } + + /// Notify the user that they need to log in to CodeBuddy + private func notifyUserLoginRequired() { + #if os(macOS) + self.log("Sending notification: CodeBuddy session expired") + + Task { + let center = UNUserNotificationCenter.current() + + do { + let granted = try await center.requestAuthorization(options: [.alert, .sound]) + guard granted else { + self.log("Notification permission denied") + return + } + } catch { + self.log("Failed to request notification permission: \(error)") + return + } + + let content = UNMutableNotificationContent() + content.title = "CodeBuddy Session Expired" + content.body = "Please log in to tencent.sso.codebuddy.cn to restore your session." + content.sound = .default + + let request = UNNotificationRequest( + identifier: "codebuddy-session-expired-\(UUID().uuidString)", + content: content, + trigger: nil) + + do { + try await center.add(request) + self.log("Notification delivered successfully") + } catch { + self.log("Failed to deliver notification: \(error)") + } + } + #endif + } + + /// Ping CodeBuddy's API endpoint to check session validity + private func pingSessionEndpoint() async throws -> Bool { + let currentSession = try? CodeBuddyCookieImporter.importSession() + guard let cookieHeader = currentSession?.cookieHeader else { + self.log("No cookies available for session ping") + return false + } + + // We need enterprise ID to make a valid API call + // Try to get it from environment or use a simple endpoint + let profileURL = URL(string: "https://tencent.sso.codebuddy.cn/profile/usage")! + + var request = URLRequest(url: profileURL) + request.timeoutInterval = self.refreshTimeout + request.setValue(cookieHeader, forHTTPHeaderField: "Cookie") + request.setValue("text/html,application/xhtml+xml", forHTTPHeaderField: "Accept") + request.setValue("https://tencent.sso.codebuddy.cn", forHTTPHeaderField: "Origin") + let userAgent = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) " + + "AppleWebKit/537.36 (KHTML, like Gecko) Chrome/144.0.0.0 Safari/537.36" + request.setValue(userAgent, forHTTPHeaderField: "User-Agent") + + do { + let (_, response) = try await URLSession.shared.data(for: request) + + guard let httpResponse = response as? HTTPURLResponse else { + self.log("Invalid response type") + return false + } + + self.log("Session ping response: HTTP \(httpResponse.statusCode)") + + if httpResponse.statusCode == 200 { + self.log("Session is valid") + return true + } else if httpResponse.statusCode == 401 || httpResponse.statusCode == 302 { + // 302 redirect to login page also means session expired + self.log("Session expired (HTTP \(httpResponse.statusCode))") + throw CodeBuddySessionKeepaliveError.sessionExpired + } else { + self.log("Unexpected response: HTTP \(httpResponse.statusCode)") + return false + } + } catch let error as CodeBuddySessionKeepaliveError { + throw error + } catch { + self.log("Request failed: \(error.localizedDescription)") + return false + } + } + + private func log(_ message: String) { + let timestamp = Date().formatted(date: .omitted, time: .standard) + let fullMessage = "[\(timestamp)] [CodeBuddyKeepalive] \(message)" + self.logger?(fullMessage) + Self.log.debug(fullMessage) + } +} + +// MARK: - Errors + +public enum CodeBuddySessionKeepaliveError: LocalizedError, Sendable { + case invalidResponse + case sessionExpired + case networkError(String) + + public var errorDescription: String? { + switch self { + case .invalidResponse: + "Invalid response from session endpoint" + case .sessionExpired: + "Session has expired" + case let .networkError(message): + "Network error: \(message)" + } + } +} + +#endif diff --git a/Sources/CodexBarCore/Providers/CodeBuddy/CodeBuddySettingsReader.swift b/Sources/CodexBarCore/Providers/CodeBuddy/CodeBuddySettingsReader.swift new file mode 100644 index 00000000..e984918d --- /dev/null +++ b/Sources/CodexBarCore/Providers/CodeBuddy/CodeBuddySettingsReader.swift @@ -0,0 +1,25 @@ +import Foundation + +public enum CodeBuddySettingsReader { + public static func cookieHeader(environment: [String: String] = ProcessInfo.processInfo.environment) -> String? { + let raw = environment["CODEBUDDY_COOKIE"] ?? environment["codebuddy_cookie"] + return self.cleaned(raw) + } + + public static func enterpriseID(environment: [String: String] = ProcessInfo.processInfo.environment) -> String? { + let raw = environment["CODEBUDDY_ENTERPRISE_ID"] ?? environment["codebuddy_enterprise_id"] + return self.cleaned(raw) + } + + private static func cleaned(_ raw: String?) -> String? { + guard let raw, !raw.isEmpty else { return nil } + var cleaned = raw.trimmingCharacters(in: .whitespacesAndNewlines) + // Handle quoted strings + if (cleaned.hasPrefix("\"") && cleaned.hasSuffix("\"")) || + (cleaned.hasPrefix("'") && cleaned.hasSuffix("'")) + { + cleaned = String(cleaned.dropFirst().dropLast()) + } + return cleaned.isEmpty ? nil : cleaned + } +} diff --git a/Sources/CodexBarCore/Providers/CodeBuddy/CodeBuddyUsageFetcher.swift b/Sources/CodexBarCore/Providers/CodeBuddy/CodeBuddyUsageFetcher.swift new file mode 100644 index 00000000..67b5c696 --- /dev/null +++ b/Sources/CodexBarCore/Providers/CodeBuddy/CodeBuddyUsageFetcher.swift @@ -0,0 +1,154 @@ +import Foundation + +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif + +public struct CodeBuddyUsageFetcher: Sendable { + private static let log = CodexBarLog.logger(LogCategories.codeBuddyAPI) + private static let usageURL = + URL(string: "https://tencent.sso.codebuddy.cn/billing/meter/get-enterprise-user-usage")! + private static let dailyUsageURL = + URL(string: "https://tencent.sso.codebuddy.cn/billing/meter/get-user-daily-usage")! + private static let baseURL = "https://tencent.sso.codebuddy.cn" + + public static func fetchUsage( + cookieHeader: String, + enterpriseID: String, + now: Date = Date()) async throws -> CodeBuddyUsageSnapshot + { + Self.log.debug("Fetching usage with enterpriseID=\(enterpriseID), cookieLength=\(cookieHeader.count)") + + var request = URLRequest(url: self.usageURL) + request.httpMethod = "POST" + request.setValue("application/json", forHTTPHeaderField: "Content-Type") + request.setValue("application/json, text/plain, */*", forHTTPHeaderField: "Accept") + request.setValue(cookieHeader, forHTTPHeaderField: "Cookie") + request.setValue(enterpriseID, forHTTPHeaderField: "x-enterprise-id") + request.setValue(self.baseURL, forHTTPHeaderField: "Origin") + request.setValue("\(self.baseURL)/profile/usage", forHTTPHeaderField: "Referer") + let userAgent = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) " + + "AppleWebKit/537.36 (KHTML, like Gecko) Chrome/144.0.0.0 Safari/537.36" + request.setValue(userAgent, forHTTPHeaderField: "User-Agent") + + // Empty JSON body as required by the API + request.httpBody = "{}".data(using: .utf8) + + let (data, response) = try await URLSession.shared.data(for: request) + guard let httpResponse = response as? HTTPURLResponse else { + throw CodeBuddyAPIError.networkError("Invalid response") + } + + guard httpResponse.statusCode == 200 else { + let responseBody = String(data: data, encoding: .utf8) ?? "" + Self.log.error("CodeBuddy API returned \(httpResponse.statusCode): \(responseBody)") + + if httpResponse.statusCode == 401 { + throw CodeBuddyAPIError.invalidCookies + } + if httpResponse.statusCode == 403 { + throw CodeBuddyAPIError.invalidCookies + } + if httpResponse.statusCode == 400 { + throw CodeBuddyAPIError.invalidRequest("Bad request") + } + throw CodeBuddyAPIError.apiError("HTTP \(httpResponse.statusCode)") + } + + let usageResponse = try JSONDecoder().decode(CodeBuddyUsageResponse.self, from: data) + + if usageResponse.code != 0 { + throw CodeBuddyAPIError.apiError("API error: \(usageResponse.msg)") + } + + return CodeBuddyUsageSnapshot( + creditUsed: usageResponse.data.credit, + creditLimit: usageResponse.data.limitNum, + cycleStartTime: usageResponse.data.cycleStartTime, + cycleEndTime: usageResponse.data.cycleEndTime, + cycleResetTime: usageResponse.data.cycleResetTime, + updatedAt: now) + } + + /// Attempt to detect enterprise ID by probing the API + /// This can be used when enterprise ID is not explicitly configured + public static func detectEnterpriseID(cookieHeader: String) async throws -> String? { + // For now, enterprise ID must be provided manually or extracted from the dashboard page + // Future enhancement: parse the dashboard HTML to extract the enterprise ID + return nil + } + + /// Fetch daily usage for the last 30 days + public static func fetchDailyUsage( + cookieHeader: String, + enterpriseID: String, + now: Date = Date()) async throws -> [CodeBuddyDailyUsageEntry] + { + // Calculate date range: last 30 days + let calendar = Calendar.current + let shanghaiTZ = TimeZone(identifier: "Asia/Shanghai") ?? .current + var calendarWithTZ = calendar + calendarWithTZ.timeZone = shanghaiTZ + + let endDate = now + guard let startDate = calendarWithTZ.date(byAdding: .day, value: -30, to: endDate) else { + return [] + } + + let dateFormatter = DateFormatter() + dateFormatter.dateFormat = "yyyy-MM-dd HH:mm:ss" + dateFormatter.timeZone = shanghaiTZ + + // Set start time to beginning of day, end time to end of day + let startOfDay = calendarWithTZ.startOfDay(for: startDate) + let endOfDay = calendarWithTZ.date(bySettingHour: 23, minute: 59, second: 59, of: endDate) ?? endDate + + let startTimeStr = dateFormatter.string(from: startOfDay) + let endTimeStr = dateFormatter.string(from: endOfDay) + + var request = URLRequest(url: self.dailyUsageURL) + request.httpMethod = "POST" + request.setValue("application/json", forHTTPHeaderField: "Content-Type") + request.setValue("application/json, text/plain, */*", forHTTPHeaderField: "Accept") + request.setValue(cookieHeader, forHTTPHeaderField: "Cookie") + request.setValue(enterpriseID, forHTTPHeaderField: "x-enterprise-id") + request.setValue(self.baseURL, forHTTPHeaderField: "Origin") + request.setValue("\(self.baseURL)/profile/usage", forHTTPHeaderField: "Referer") + let userAgent = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) " + + "AppleWebKit/537.36 (KHTML, like Gecko) Chrome/144.0.0.0 Safari/537.36" + request.setValue(userAgent, forHTTPHeaderField: "User-Agent") + + // Request body with date range and pagination + let bodyDict: [String: Any] = [ + "startTime": startTimeStr, + "endTime": endTimeStr, + "pageNum": 1, + "pageSize": 31, // Request 31 days to ensure we get all 30 + ] + request.httpBody = try? JSONSerialization.data(withJSONObject: bodyDict) + + let (data, response) = try await URLSession.shared.data(for: request) + guard let httpResponse = response as? HTTPURLResponse else { + Self.log.error("CodeBuddy daily usage API: Invalid response") + return [] + } + + guard httpResponse.statusCode == 200 else { + let responseBody = String(data: data, encoding: .utf8) ?? "" + Self.log.error("CodeBuddy daily usage API returned \(httpResponse.statusCode): \(responseBody)") + return [] + } + + let dailyResponse = try JSONDecoder().decode(CodeBuddyDailyUsageResponse.self, from: data) + + if dailyResponse.code != 0 { + Self.log.error("CodeBuddy daily usage API error: \(dailyResponse.msg)") + return [] + } + + // Convert to public entry type + return dailyResponse.data.data.map { entry in + CodeBuddyDailyUsageEntry(date: entry.date, credit: entry.credit) + } + } +} diff --git a/Sources/CodexBarCore/Providers/CodeBuddy/CodeBuddyUsageSnapshot.swift b/Sources/CodexBarCore/Providers/CodeBuddy/CodeBuddyUsageSnapshot.swift new file mode 100644 index 00000000..166f21cf --- /dev/null +++ b/Sources/CodexBarCore/Providers/CodeBuddy/CodeBuddyUsageSnapshot.swift @@ -0,0 +1,64 @@ +import Foundation + +public struct CodeBuddyUsageSnapshot: Sendable { + public let creditUsed: Double + public let creditLimit: Double + public let cycleStartTime: String + public let cycleEndTime: String + public let cycleResetTime: String + public let updatedAt: Date + + public init( + creditUsed: Double, + creditLimit: Double, + cycleStartTime: String, + cycleEndTime: String, + cycleResetTime: String, + updatedAt: Date) + { + self.creditUsed = creditUsed + self.creditLimit = creditLimit + self.cycleStartTime = cycleStartTime + self.cycleEndTime = cycleEndTime + self.cycleResetTime = cycleResetTime + self.updatedAt = updatedAt + } + + private static func parseDate(_ dateString: String) -> Date? { + // Handle format like: "2026-02-28 23:59:59" + let formatter = DateFormatter() + formatter.dateFormat = "yyyy-MM-dd HH:mm:ss" + formatter.timeZone = TimeZone(identifier: "Asia/Shanghai") + return formatter.date(from: dateString) + } +} + +extension CodeBuddyUsageSnapshot { + public func toUsageSnapshot() -> UsageSnapshot { + let usagePercent = self.creditLimit > 0 ? (self.creditUsed / self.creditLimit) * 100 : 0 + + // Format credit values for display (e.g., "1,121 / 25,000") + let usedInt = Int(self.creditUsed.rounded()) + let limitInt = Int(self.creditLimit.rounded()) + + let creditWindow = RateWindow( + usedPercent: usagePercent, + windowMinutes: nil, + resetsAt: Self.parseDate(self.cycleResetTime), + resetDescription: "\(usedInt.formatted()) / \(limitInt.formatted()) credits") + + let identity = ProviderIdentitySnapshot( + providerID: .codebuddy, + accountEmail: nil, + accountOrganization: nil, + loginMethod: nil) + + return UsageSnapshot( + primary: creditWindow, + secondary: nil, + tertiary: nil, + providerCost: nil, + updatedAt: self.updatedAt, + identity: identity) + } +} diff --git a/Sources/CodexBarCore/Providers/ProviderDescriptor.swift b/Sources/CodexBarCore/Providers/ProviderDescriptor.swift index 6aff8369..a207a9ea 100644 --- a/Sources/CodexBarCore/Providers/ProviderDescriptor.swift +++ b/Sources/CodexBarCore/Providers/ProviderDescriptor.swift @@ -71,6 +71,7 @@ public enum ProviderDescriptorRegistry { .kimik2: KimiK2ProviderDescriptor.descriptor, .amp: AmpProviderDescriptor.descriptor, .synthetic: SyntheticProviderDescriptor.descriptor, + .codebuddy: CodeBuddyProviderDescriptor.descriptor, ] private static let bootstrap: Void = { for provider in UsageProvider.allCases { diff --git a/Sources/CodexBarCore/Providers/ProviderFetchPlan.swift b/Sources/CodexBarCore/Providers/ProviderFetchPlan.swift index cadbdc81..49639996 100644 --- a/Sources/CodexBarCore/Providers/ProviderFetchPlan.swift +++ b/Sources/CodexBarCore/Providers/ProviderFetchPlan.swift @@ -61,6 +61,7 @@ public struct ProviderFetchResult: Sendable { public let usage: UsageSnapshot public let credits: CreditsSnapshot? public let dashboard: OpenAIDashboardSnapshot? + public let codeBuddyDailyUsage: [CodeBuddyDailyUsageEntry]? public let sourceLabel: String public let strategyID: String public let strategyKind: ProviderFetchKind @@ -69,6 +70,7 @@ public struct ProviderFetchResult: Sendable { usage: UsageSnapshot, credits: CreditsSnapshot?, dashboard: OpenAIDashboardSnapshot?, + codeBuddyDailyUsage: [CodeBuddyDailyUsageEntry]? = nil, sourceLabel: String, strategyID: String, strategyKind: ProviderFetchKind) @@ -76,6 +78,7 @@ public struct ProviderFetchResult: Sendable { self.usage = usage self.credits = credits self.dashboard = dashboard + self.codeBuddyDailyUsage = codeBuddyDailyUsage self.sourceLabel = sourceLabel self.strategyID = strategyID self.strategyKind = strategyKind @@ -139,12 +142,14 @@ extension ProviderFetchStrategy { usage: UsageSnapshot, credits: CreditsSnapshot? = nil, dashboard: OpenAIDashboardSnapshot? = nil, + codeBuddyDailyUsage: [CodeBuddyDailyUsageEntry]? = nil, sourceLabel: String) -> ProviderFetchResult { ProviderFetchResult( usage: usage, credits: credits, dashboard: dashboard, + codeBuddyDailyUsage: codeBuddyDailyUsage, sourceLabel: sourceLabel, strategyID: self.id, strategyKind: self.kind) diff --git a/Sources/CodexBarCore/Providers/ProviderSettingsSnapshot.swift b/Sources/CodexBarCore/Providers/ProviderSettingsSnapshot.swift index f83cb9fd..d6e244b3 100644 --- a/Sources/CodexBarCore/Providers/ProviderSettingsSnapshot.swift +++ b/Sources/CodexBarCore/Providers/ProviderSettingsSnapshot.swift @@ -15,7 +15,8 @@ public struct ProviderSettingsSnapshot: Sendable { kimi: KimiProviderSettings? = nil, augment: AugmentProviderSettings? = nil, amp: AmpProviderSettings? = nil, - jetbrains: JetBrainsProviderSettings? = nil) -> ProviderSettingsSnapshot + jetbrains: JetBrainsProviderSettings? = nil, + codebuddy: CodeBuddyProviderSettings? = nil) -> ProviderSettingsSnapshot { ProviderSettingsSnapshot( debugMenuEnabled: debugMenuEnabled, @@ -31,7 +32,8 @@ public struct ProviderSettingsSnapshot: Sendable { kimi: kimi, augment: augment, amp: amp, - jetbrains: jetbrains) + jetbrains: jetbrains, + codebuddy: codebuddy) } public struct CodexProviderSettings: Sendable { @@ -167,6 +169,18 @@ public struct ProviderSettingsSnapshot: Sendable { } } + public struct CodeBuddyProviderSettings: Sendable { + public let cookieSource: ProviderCookieSource + public let manualCookieHeader: String? + public let enterpriseID: String? + + public init(cookieSource: ProviderCookieSource, manualCookieHeader: String?, enterpriseID: String?) { + self.cookieSource = cookieSource + self.manualCookieHeader = manualCookieHeader + self.enterpriseID = enterpriseID + } + } + public let debugMenuEnabled: Bool public let debugKeepCLISessionsAlive: Bool public let codex: CodexProviderSettings? @@ -181,6 +195,7 @@ public struct ProviderSettingsSnapshot: Sendable { public let augment: AugmentProviderSettings? public let amp: AmpProviderSettings? public let jetbrains: JetBrainsProviderSettings? + public let codebuddy: CodeBuddyProviderSettings? public var jetbrainsIDEBasePath: String? { self.jetbrains?.ideBasePath @@ -200,7 +215,8 @@ public struct ProviderSettingsSnapshot: Sendable { kimi: KimiProviderSettings?, augment: AugmentProviderSettings?, amp: AmpProviderSettings?, - jetbrains: JetBrainsProviderSettings? = nil) + jetbrains: JetBrainsProviderSettings? = nil, + codebuddy: CodeBuddyProviderSettings? = nil) { self.debugMenuEnabled = debugMenuEnabled self.debugKeepCLISessionsAlive = debugKeepCLISessionsAlive @@ -216,6 +232,7 @@ public struct ProviderSettingsSnapshot: Sendable { self.augment = augment self.amp = amp self.jetbrains = jetbrains + self.codebuddy = codebuddy } } @@ -232,6 +249,7 @@ public enum ProviderSettingsSnapshotContribution: Sendable { case augment(ProviderSettingsSnapshot.AugmentProviderSettings) case amp(ProviderSettingsSnapshot.AmpProviderSettings) case jetbrains(ProviderSettingsSnapshot.JetBrainsProviderSettings) + case codebuddy(ProviderSettingsSnapshot.CodeBuddyProviderSettings) } public struct ProviderSettingsSnapshotBuilder: Sendable { @@ -249,6 +267,7 @@ public struct ProviderSettingsSnapshotBuilder: Sendable { public var augment: ProviderSettingsSnapshot.AugmentProviderSettings? public var amp: ProviderSettingsSnapshot.AmpProviderSettings? public var jetbrains: ProviderSettingsSnapshot.JetBrainsProviderSettings? + public var codebuddy: ProviderSettingsSnapshot.CodeBuddyProviderSettings? public init(debugMenuEnabled: Bool = false, debugKeepCLISessionsAlive: Bool = false) { self.debugMenuEnabled = debugMenuEnabled @@ -269,6 +288,7 @@ public struct ProviderSettingsSnapshotBuilder: Sendable { case let .augment(value): self.augment = value case let .amp(value): self.amp = value case let .jetbrains(value): self.jetbrains = value + case let .codebuddy(value): self.codebuddy = value } } @@ -287,6 +307,7 @@ public struct ProviderSettingsSnapshotBuilder: Sendable { kimi: self.kimi, augment: self.augment, amp: self.amp, - jetbrains: self.jetbrains) + jetbrains: self.jetbrains, + codebuddy: self.codebuddy) } } diff --git a/Sources/CodexBarCore/Providers/Providers.swift b/Sources/CodexBarCore/Providers/Providers.swift index a267fb95..192d3395 100644 --- a/Sources/CodexBarCore/Providers/Providers.swift +++ b/Sources/CodexBarCore/Providers/Providers.swift @@ -21,6 +21,7 @@ public enum UsageProvider: String, CaseIterable, Sendable, Codable { case kimik2 case amp case synthetic + case codebuddy } // swiftformat:enable sortDeclarations @@ -45,6 +46,7 @@ public enum IconStyle: Sendable, CaseIterable { case amp case synthetic case combined + case codebuddy } public struct ProviderMetadata: Sendable { diff --git a/Sources/CodexBarCore/Vendored/CostUsage/CostUsageScanner.swift b/Sources/CodexBarCore/Vendored/CostUsage/CostUsageScanner.swift index d47d7d55..e7c2b762 100644 --- a/Sources/CodexBarCore/Vendored/CostUsage/CostUsageScanner.swift +++ b/Sources/CodexBarCore/Vendored/CostUsage/CostUsageScanner.swift @@ -99,6 +99,8 @@ enum CostUsageScanner { return CostUsageDailyReport(data: [], summary: nil) case .synthetic: return CostUsageDailyReport(data: [], summary: nil) + case .codebuddy: + return CostUsageDailyReport(data: [], summary: nil) } } diff --git a/Sources/CodexBarWidget/CodexBarWidgetProvider.swift b/Sources/CodexBarWidget/CodexBarWidgetProvider.swift index 1634611e..8d8ffcb6 100644 --- a/Sources/CodexBarWidget/CodexBarWidgetProvider.swift +++ b/Sources/CodexBarWidget/CodexBarWidgetProvider.swift @@ -59,6 +59,7 @@ enum ProviderChoice: String, AppEnum { case .kimik2: return nil // Kimi K2 not yet supported in widgets case .amp: return nil // Amp not yet supported in widgets case .synthetic: return nil // Synthetic not yet supported in widgets + case .codebuddy: return nil // CodeBuddy not yet supported in widgets } } } diff --git a/Sources/CodexBarWidget/CodexBarWidgetViews.swift b/Sources/CodexBarWidget/CodexBarWidgetViews.swift index ed39b450..04ee78eb 100644 --- a/Sources/CodexBarWidget/CodexBarWidgetViews.swift +++ b/Sources/CodexBarWidget/CodexBarWidgetViews.swift @@ -275,6 +275,7 @@ private struct ProviderSwitchChip: View { case .kimik2: "Kimi K2" case .amp: "Amp" case .synthetic: "Synthetic" + case .codebuddy: "CB" } } } @@ -605,6 +606,8 @@ enum WidgetColors { Color(red: 220 / 255, green: 38 / 255, blue: 38 / 255) // Amp red case .synthetic: Color(red: 20 / 255, green: 20 / 255, blue: 20 / 255) // Synthetic charcoal + case .codebuddy: + Color(red: 0 / 255, green: 120 / 255, blue: 215 / 255) // CodeBuddy blue } } } diff --git a/Tests/CodexBarTests/CodeBuddyTests.swift b/Tests/CodexBarTests/CodeBuddyTests.swift new file mode 100644 index 00000000..537fcd99 --- /dev/null +++ b/Tests/CodexBarTests/CodeBuddyTests.swift @@ -0,0 +1,161 @@ +import Foundation +import Testing + +@testable import CodexBarCore + +@Suite +struct CodeBuddyModelsTests { + @Test + func parseUsageResponse_validJSON_parsesCorrectly() throws { + let json = """ + { + "code": 0, + "msg": "OK", + "requestId": "e50a08a9-4c96-45f6-8d54-1b20914d5d29", + "data": { + "credit": 1121.44, + "cycleStartTime": "2026-02-01 00:00:00", + "cycleEndTime": "2026-02-28 23:59:59", + "limitNum": 25000, + "cycleResetTime": "2026-03-01 23:59:59" + } + } + """ + let data = json.data(using: .utf8)! + let response = try JSONDecoder().decode(CodeBuddyUsageResponse.self, from: data) + + #expect(response.code == 0) + #expect(response.msg == "OK") + #expect(response.data.credit == 1121.44) + #expect(response.data.limitNum == 25000) + #expect(response.data.cycleStartTime == "2026-02-01 00:00:00") + #expect(response.data.cycleEndTime == "2026-02-28 23:59:59") + #expect(response.data.cycleResetTime == "2026-03-01 23:59:59") + } + + @Test + func parseDailyUsageResponse_validJSON_parsesCorrectly() throws { + let json = """ + { + "code": 0, + "msg": "OK", + "requestId": "74093760-b510-4955-aa64-2013958f7aba", + "data": { + "total": 7, + "data": [ + {"credit": 387.84, "date": "2026-02-03"}, + {"credit": 278.86, "date": "2026-02-02"}, + {"credit": 129.25, "date": "2026-01-31"} + ] + } + } + """ + let data = json.data(using: .utf8)! + let response = try JSONDecoder().decode(CodeBuddyDailyUsageResponse.self, from: data) + + #expect(response.code == 0) + #expect(response.data.total == 7) + #expect(response.data.data.count == 3) + #expect(response.data.data[0].credit == 387.84) + #expect(response.data.data[0].date == "2026-02-03") + } +} + +@Suite +struct CodeBuddyUsageSnapshotTests { + @Test + func toUsageSnapshot_calculatesPercentageCorrectly() { + let snapshot = CodeBuddyUsageSnapshot( + creditUsed: 1000, + creditLimit: 10000, + cycleStartTime: "2026-02-01 00:00:00", + cycleEndTime: "2026-02-28 23:59:59", + cycleResetTime: "2026-03-01 23:59:59", + updatedAt: Date()) + + let usageSnapshot = snapshot.toUsageSnapshot() + + #expect(usageSnapshot.primary?.usedPercent == 10.0) + #expect(usageSnapshot.identity?.providerID == .codebuddy) + } + + @Test + func toUsageSnapshot_handlesZeroLimit() { + let snapshot = CodeBuddyUsageSnapshot( + creditUsed: 100, + creditLimit: 0, + cycleStartTime: "2026-02-01 00:00:00", + cycleEndTime: "2026-02-28 23:59:59", + cycleResetTime: "2026-03-01 23:59:59", + updatedAt: Date()) + + let usageSnapshot = snapshot.toUsageSnapshot() + + #expect(usageSnapshot.primary?.usedPercent == 0.0) + } + + @Test + func toUsageSnapshot_formatsResetDescription() { + let snapshot = CodeBuddyUsageSnapshot( + creditUsed: 1121.44, + creditLimit: 25000, + cycleStartTime: "2026-02-01 00:00:00", + cycleEndTime: "2026-02-28 23:59:59", + cycleResetTime: "2026-03-01 23:59:59", + updatedAt: Date()) + + let usageSnapshot = snapshot.toUsageSnapshot() + + #expect(usageSnapshot.primary?.resetDescription?.contains("1,121") == true) + #expect(usageSnapshot.primary?.resetDescription?.contains("25,000") == true) + } +} + +@Suite +struct CodeBuddyCookieHeaderTests { + @Test + func extractEnterpriseID_fromCurlCommand() { + let curl = """ + curl 'https://tencent.sso.codebuddy.cn/billing/meter/get-enterprise-user-usage' \\ + -H 'x-enterprise-id: etahzsqej0n4' \\ + --data-raw '{}' + """ + + let enterpriseID = CodeBuddyCookieHeader.extractEnterpriseID(from: curl) + + #expect(enterpriseID == "etahzsqej0n4") + } + + @Test + func extractEnterpriseID_fromHeader() { + let header = "x-enterprise-id: abc123xyz" + + let enterpriseID = CodeBuddyCookieHeader.extractEnterpriseID(from: header) + + #expect(enterpriseID == "abc123xyz") + } + + @Test + func override_fromCookieHeader() { + let cookieHeader = "session=abc123; session_2=xyz789" + + let override = CodeBuddyCookieHeader.override(from: cookieHeader) + + #expect(override != nil) + #expect(override?.cookieHeader == cookieHeader) + } + + @Test + func override_fromEmptyString_returnsNil() { + let override = CodeBuddyCookieHeader.override(from: "") + + #expect(override == nil) + } + + @Test + func override_fromNil_returnsNil() { + let override = CodeBuddyCookieHeader.override(from: nil) + + #expect(override == nil) + } +} diff --git a/Tests/CodexBarTests/SettingsStoreTests.swift b/Tests/CodexBarTests/SettingsStoreTests.swift index fad99a76..919f168d 100644 --- a/Tests/CodexBarTests/SettingsStoreTests.swift +++ b/Tests/CodexBarTests/SettingsStoreTests.swift @@ -364,6 +364,7 @@ struct SettingsStoreTests { .kimik2, .amp, .synthetic, + .codebuddy, ]) // Move one provider; ensure it's persisted across instances.