diff --git a/.github/workflows/build-app.yml b/.github/workflows/build-app.yml new file mode 100644 index 000000000..2dd0da28a --- /dev/null +++ b/.github/workflows/build-app.yml @@ -0,0 +1,58 @@ +name: Build CodexBar App + +on: + workflow_dispatch: + push: + branches: ["feat/qwen-doubao-providers"] + +jobs: + build-macos-app: + runs-on: macos-latest + steps: + - uses: actions/checkout@v6 + + - name: Select Xcode + run: | + set -euo pipefail + for candidate in /Applications/Xcode_26.1.1.app /Applications/Xcode_26.1.app /Applications/Xcode.app; do + if [[ -d "$candidate" ]]; then + sudo xcode-select -s "${candidate}/Contents/Developer" + echo "DEVELOPER_DIR=${candidate}/Contents/Developer" >> "$GITHUB_ENV" + break + fi + done + xcodebuild -version + swift --version + + - name: Resolve dependencies + run: swift package resolve + + - name: Build release + run: swift build -c release 2>&1 + + - name: Run tests + run: swift test --no-parallel + + - name: Fix rpath and package + run: | + set -euo pipefail + BIN_DIR="$(swift build -c release --show-bin-path)" + echo "Binary directory: $BIN_DIR" + + # Add @executable_path/../Frameworks rpath so Sparkle.framework loads from .app bundle + install_name_tool -add_rpath @executable_path/../Frameworks "$BIN_DIR/CodexBar" || true + install_name_tool -add_rpath @executable_path/../Frameworks "$BIN_DIR/CodexBarCLI" || true + + # Create a zip of the built products + cd "$BIN_DIR" + zip -r "$GITHUB_WORKSPACE/CodexBar-custom-build.zip" \ + CodexBar CodexBarCLI CodexBarClaudeWatchdog CodexBarClaudeWebProbe \ + CodexBar_CodexBar.bundle KeyboardShortcuts_KeyboardShortcuts.bundle \ + Sparkle.framework + + - name: Upload build artifact + uses: actions/upload-artifact@v4 + with: + name: CodexBar-custom-build + path: CodexBar-custom-build.zip + retention-days: 30 diff --git a/Sources/CodexBar/MenuDescriptor.swift b/Sources/CodexBar/MenuDescriptor.swift index 05aa55fff..d90df6bd3 100644 --- a/Sources/CodexBar/MenuDescriptor.swift +++ b/Sources/CodexBar/MenuDescriptor.swift @@ -247,7 +247,14 @@ struct MenuDescriptor { entries.append(.text("Activity: \(detail)", .secondary)) } } else if let loginMethodText, !loginMethodText.isEmpty { - entries.append(.text("Plan: \(AccountFormatter.plan(loginMethodText))", .secondary)) + let formatted = AccountFormatter.plan(loginMethodText) + // Balance-style providers (openrouter, mimo) already emit "Balance: $X.XX"; + // don't double-prefix with "Plan: " in that case. + if provider == .openrouter || provider == .mimo, formatted.hasPrefix("Balance:") { + entries.append(.text(formatted, .secondary)) + } else { + entries.append(.text("Plan: \(formatted)", .secondary)) + } } if metadata.usesAccountFallback { diff --git a/Sources/CodexBar/PreferencesProviderDetailView.swift b/Sources/CodexBar/PreferencesProviderDetailView.swift index 58a55deb5..7b98b9649 100644 --- a/Sources/CodexBar/PreferencesProviderDetailView.swift +++ b/Sources/CodexBar/PreferencesProviderDetailView.swift @@ -27,7 +27,7 @@ struct ProviderDetailView: View { else { return nil } - guard provider == .openrouter else { + guard provider == .openrouter || provider == .mimo else { return (label: "Plan", value: rawPlan) } diff --git a/Sources/CodexBar/Providers/AigoCode/AigoCodeProviderImplementation.swift b/Sources/CodexBar/Providers/AigoCode/AigoCodeProviderImplementation.swift new file mode 100644 index 000000000..03ce723ba --- /dev/null +++ b/Sources/CodexBar/Providers/AigoCode/AigoCodeProviderImplementation.swift @@ -0,0 +1,42 @@ +import AppKit +import CodexBarCore +import CodexBarMacroSupport +import Foundation + +@ProviderImplementationRegistration +struct AigoCodeProviderImplementation: ProviderImplementation { + let id: UsageProvider = .aigocode + + @MainActor + func observeSettings(_ settings: SettingsStore) { + _ = settings.aigocodeAPIToken + } + + @MainActor + func settingsFields(context: ProviderSettingsContext) -> [ProviderSettingsFieldDescriptor] { + [ + ProviderSettingsFieldDescriptor( + id: "aigocode-api-token", + title: "API key", + subtitle: "Optional when using web dashboard mode. " + + "Stored in ~/.codexbar/config.json.", + kind: .secure, + placeholder: "sk-...", + binding: context.stringBinding(\.aigocodeAPIToken), + actions: [ + ProviderSettingsActionDescriptor( + id: "aigocode-open-dashboard", + title: "Open AigoCode Dashboard", + style: .link, + isVisible: nil, + perform: { + if let url = URL(string: "https://www.aigocode.com/dashboard/console") { + NSWorkspace.shared.open(url) + } + }), + ], + isVisible: nil, + onActivate: nil), + ] + } +} diff --git a/Sources/CodexBar/Providers/AigoCode/AigoCodeSettingsStore.swift b/Sources/CodexBar/Providers/AigoCode/AigoCodeSettingsStore.swift new file mode 100644 index 000000000..e43537736 --- /dev/null +++ b/Sources/CodexBar/Providers/AigoCode/AigoCodeSettingsStore.swift @@ -0,0 +1,14 @@ +import CodexBarCore +import Foundation + +extension SettingsStore { + var aigocodeAPIToken: String { + get { self.configSnapshot.providerConfig(for: .aigocode)?.sanitizedAPIKey ?? "" } + set { + self.updateProviderConfig(provider: .aigocode) { entry in + entry.apiKey = self.normalizedConfigValue(newValue) + } + self.logSecretUpdate(provider: .aigocode, field: "apiKey", value: newValue) + } + } +} diff --git a/Sources/CodexBar/Providers/Doubao/DoubaoProviderImplementation.swift b/Sources/CodexBar/Providers/Doubao/DoubaoProviderImplementation.swift new file mode 100644 index 000000000..fe974bed8 --- /dev/null +++ b/Sources/CodexBar/Providers/Doubao/DoubaoProviderImplementation.swift @@ -0,0 +1,42 @@ +import AppKit +import CodexBarCore +import CodexBarMacroSupport +import Foundation + +@ProviderImplementationRegistration +struct DoubaoProviderImplementation: ProviderImplementation { + let id: UsageProvider = .doubao + + @MainActor + func observeSettings(_ settings: SettingsStore) { + _ = settings.doubaoAPIToken + } + + @MainActor + func settingsFields(context: ProviderSettingsContext) -> [ProviderSettingsFieldDescriptor] { + [ + ProviderSettingsFieldDescriptor( + id: "doubao-api-token", + title: "API key", + subtitle: "Stored in ~/.codexbar/config.json. Get your API key from the Volcengine " + + "Ark console.", + kind: .secure, + placeholder: "ark-...", + binding: context.stringBinding(\.doubaoAPIToken), + actions: [ + ProviderSettingsActionDescriptor( + id: "doubao-open-dashboard", + title: "Open Volcengine Ark Console", + style: .link, + isVisible: nil, + perform: { + if let url = URL(string: "https://console.volcengine.com/ark/") { + NSWorkspace.shared.open(url) + } + }), + ], + isVisible: nil, + onActivate: nil), + ] + } +} diff --git a/Sources/CodexBar/Providers/Doubao/DoubaoSettingsStore.swift b/Sources/CodexBar/Providers/Doubao/DoubaoSettingsStore.swift new file mode 100644 index 000000000..4d69a273f --- /dev/null +++ b/Sources/CodexBar/Providers/Doubao/DoubaoSettingsStore.swift @@ -0,0 +1,14 @@ +import CodexBarCore +import Foundation + +extension SettingsStore { + var doubaoAPIToken: String { + get { self.configSnapshot.providerConfig(for: .doubao)?.sanitizedAPIKey ?? "" } + set { + self.updateProviderConfig(provider: .doubao) { entry in + entry.apiKey = self.normalizedConfigValue(newValue) + } + self.logSecretUpdate(provider: .doubao, field: "apiKey", value: newValue) + } + } +} diff --git a/Sources/CodexBar/Providers/MiMo/MiMoProviderImplementation.swift b/Sources/CodexBar/Providers/MiMo/MiMoProviderImplementation.swift new file mode 100644 index 000000000..bcb9b68cf --- /dev/null +++ b/Sources/CodexBar/Providers/MiMo/MiMoProviderImplementation.swift @@ -0,0 +1,102 @@ +import AppKit +import CodexBarCore +import CodexBarMacroSupport +import Foundation +import SwiftUI + +@ProviderImplementationRegistration +struct MiMoProviderImplementation: ProviderImplementation { + let id: UsageProvider = .mimo + let supportsLoginFlow: Bool = true + + @MainActor + func presentation(context _: ProviderPresentationContext) -> ProviderPresentation { + ProviderPresentation { _ in "web" } + } + + @MainActor + func observeSettings(_ settings: SettingsStore) { + _ = settings.miMoCookieSource + _ = settings.miMoCookieHeader + } + + @MainActor + func settingsSnapshot(context: ProviderSettingsSnapshotContext) -> ProviderSettingsSnapshotContribution? { + .mimo(context.settings.miMoSettingsSnapshot(tokenOverride: context.tokenOverride)) + } + + @MainActor + func settingsPickers(context: ProviderSettingsContext) -> [ProviderSettingsPickerDescriptor] { + let cookieBinding = Binding( + get: { context.settings.miMoCookieSource.rawValue }, + set: { raw in + context.settings.miMoCookieSource = ProviderCookieSource(rawValue: raw) ?? .auto + }) + let cookieOptions = ProviderCookieSourceUI.options( + allowsOff: false, + keychainDisabled: context.settings.debugDisableKeychainAccess) + let cookieSubtitle: () -> String? = { + ProviderCookieSourceUI.subtitle( + source: context.settings.miMoCookieSource, + keychainDisabled: context.settings.debugDisableKeychainAccess, + auto: "Automatic imports Chrome browser cookies from Xiaomi MiMo.", + manual: "Paste a Cookie header from platform.xiaomimimo.com.", + off: "Xiaomi MiMo cookies are disabled.") + } + + return [ + ProviderSettingsPickerDescriptor( + id: "mimo-cookie-source", + title: "Cookie source", + subtitle: "Automatic imports Chrome browser cookies from Xiaomi MiMo.", + dynamicSubtitle: cookieSubtitle, + binding: cookieBinding, + options: cookieOptions, + isVisible: nil, + onChange: nil, + trailingText: { + guard let entry = CookieHeaderCache.load(provider: .mimo) else { return nil } + let when = entry.storedAt.relativeDescription() + return "Cached: \(entry.sourceLabel) • \(when)" + }), + ] + } + + @MainActor + func settingsFields(context: ProviderSettingsContext) -> [ProviderSettingsFieldDescriptor] { + [ + ProviderSettingsFieldDescriptor( + id: "mimo-cookie", + title: "", + subtitle: "", + kind: .secure, + placeholder: "Cookie: ...", + binding: context.stringBinding(\.miMoCookieHeader), + actions: [ + ProviderSettingsActionDescriptor( + id: "mimo-open-balance", + title: "Open MiMo Balance", + style: .link, + isVisible: nil, + perform: { + guard let url = URL(string: "https://platform.xiaomimimo.com/#/console/balance") else { + return + } + NSWorkspace.shared.open(url) + }), + ], + isVisible: { context.settings.miMoCookieSource == .manual }, + onActivate: { context.settings.ensureMiMoCookieLoaded() }), + ] + } + + @MainActor + func runLoginFlow(context _: ProviderLoginContext) async -> Bool { + let loginURL = "https://platform.xiaomimimo.com/api/v1/genLoginUrl?currentPath=%2F%23%2Fconsole%2Fbalance" + guard let url = URL(string: loginURL) else { + return false + } + NSWorkspace.shared.open(url) + return false + } +} diff --git a/Sources/CodexBar/Providers/MiMo/MiMoSettingsStore.swift b/Sources/CodexBar/Providers/MiMo/MiMoSettingsStore.swift new file mode 100644 index 000000000..3285a20b3 --- /dev/null +++ b/Sources/CodexBar/Providers/MiMo/MiMoSettingsStore.swift @@ -0,0 +1,35 @@ +import CodexBarCore +import Foundation + +extension SettingsStore { + var miMoCookieHeader: String { + get { self.configSnapshot.providerConfig(for: .mimo)?.sanitizedCookieHeader ?? "" } + set { + self.updateProviderConfig(provider: .mimo) { entry in + entry.cookieHeader = self.normalizedConfigValue(newValue) + } + self.logSecretUpdate(provider: .mimo, field: "cookieHeader", value: newValue) + } + } + + var miMoCookieSource: ProviderCookieSource { + get { self.resolvedCookieSource(provider: .mimo, fallback: .auto) } + set { + self.updateProviderConfig(provider: .mimo) { entry in + entry.cookieSource = newValue + } + self.logProviderModeChange(provider: .mimo, field: "cookieSource", value: newValue.rawValue) + } + } + + func ensureMiMoCookieLoaded() {} +} + +extension SettingsStore { + func miMoSettingsSnapshot(tokenOverride: TokenAccountOverride?) -> ProviderSettingsSnapshot.MiMoProviderSettings { + _ = tokenOverride + return ProviderSettingsSnapshot.MiMoProviderSettings( + cookieSource: self.miMoCookieSource, + manualCookieHeader: self.miMoCookieHeader) + } +} diff --git a/Sources/CodexBar/Providers/Qwen/QwenProviderImplementation.swift b/Sources/CodexBar/Providers/Qwen/QwenProviderImplementation.swift new file mode 100644 index 000000000..5b981e26b --- /dev/null +++ b/Sources/CodexBar/Providers/Qwen/QwenProviderImplementation.swift @@ -0,0 +1,42 @@ +import AppKit +import CodexBarCore +import CodexBarMacroSupport +import Foundation + +@ProviderImplementationRegistration +struct QwenProviderImplementation: ProviderImplementation { + let id: UsageProvider = .qwen + + @MainActor + func observeSettings(_ settings: SettingsStore) { + _ = settings.qwenAPIToken + } + + @MainActor + func settingsFields(context: ProviderSettingsContext) -> [ProviderSettingsFieldDescriptor] { + [ + ProviderSettingsFieldDescriptor( + id: "qwen-api-token", + title: "API key", + subtitle: "Stored in ~/.codexbar/config.json. Get your API key from the Alibaba Cloud " + + "Bailian console (DashScope).", + kind: .secure, + placeholder: "sk-...", + binding: context.stringBinding(\.qwenAPIToken), + actions: [ + ProviderSettingsActionDescriptor( + id: "qwen-open-dashboard", + title: "Open Bailian Console", + style: .link, + isVisible: nil, + perform: { + if let url = URL(string: "https://bailian.console.aliyun.com/") { + NSWorkspace.shared.open(url) + } + }), + ], + isVisible: nil, + onActivate: nil), + ] + } +} diff --git a/Sources/CodexBar/Providers/Qwen/QwenSettingsStore.swift b/Sources/CodexBar/Providers/Qwen/QwenSettingsStore.swift new file mode 100644 index 000000000..0aa570954 --- /dev/null +++ b/Sources/CodexBar/Providers/Qwen/QwenSettingsStore.swift @@ -0,0 +1,14 @@ +import CodexBarCore +import Foundation + +extension SettingsStore { + var qwenAPIToken: String { + get { self.configSnapshot.providerConfig(for: .qwen)?.sanitizedAPIKey ?? "" } + set { + self.updateProviderConfig(provider: .qwen) { entry in + entry.apiKey = self.normalizedConfigValue(newValue) + } + self.logSecretUpdate(provider: .qwen, field: "apiKey", value: newValue) + } + } +} diff --git a/Sources/CodexBar/Providers/Shared/ProviderImplementationRegistry.swift b/Sources/CodexBar/Providers/Shared/ProviderImplementationRegistry.swift index 7938b3d49..1879eea3a 100644 --- a/Sources/CodexBar/Providers/Shared/ProviderImplementationRegistry.swift +++ b/Sources/CodexBar/Providers/Shared/ProviderImplementationRegistry.swift @@ -35,6 +35,13 @@ enum ProviderImplementationRegistry { case .synthetic: SyntheticProviderImplementation() case .openrouter: OpenRouterProviderImplementation() case .warp: WarpProviderImplementation() + case .qwen: QwenProviderImplementation() + case .doubao: DoubaoProviderImplementation() + case .zenmux: ZenmuxProviderImplementation() + case .aigocode: AigoCodeProviderImplementation() + case .trae: TraeProviderImplementation() + case .stepfun: StepFunProviderImplementation() + case .mimo: MiMoProviderImplementation() } } diff --git a/Sources/CodexBar/Providers/StepFun/StepFunProviderImplementation.swift b/Sources/CodexBar/Providers/StepFun/StepFunProviderImplementation.swift new file mode 100644 index 000000000..2b7811c9c --- /dev/null +++ b/Sources/CodexBar/Providers/StepFun/StepFunProviderImplementation.swift @@ -0,0 +1,41 @@ +import AppKit +import CodexBarCore +import CodexBarMacroSupport +import Foundation + +@ProviderImplementationRegistration +struct StepFunProviderImplementation: ProviderImplementation { + let id: UsageProvider = .stepfun + + @MainActor + func observeSettings(_ settings: SettingsStore) { + _ = settings.stepfunAPIToken + } + + @MainActor + func settingsFields(context: ProviderSettingsContext) -> [ProviderSettingsFieldDescriptor] { + [ + ProviderSettingsFieldDescriptor( + id: "stepfun-api-token", + title: "API key", + subtitle: "Stored in ~/.codexbar/config.json. Get your API key from the StepFun platform.", + kind: .secure, + placeholder: "sf-...", + binding: context.stringBinding(\.stepfunAPIToken), + actions: [ + ProviderSettingsActionDescriptor( + id: "stepfun-open-dashboard", + title: "Open StepFun Platform", + style: .link, + isVisible: nil, + perform: { + if let url = URL(string: "https://platform.stepfun.com/plan-subscribe") { + NSWorkspace.shared.open(url) + } + }), + ], + isVisible: nil, + onActivate: nil), + ] + } +} diff --git a/Sources/CodexBar/Providers/StepFun/StepFunSettingsStore.swift b/Sources/CodexBar/Providers/StepFun/StepFunSettingsStore.swift new file mode 100644 index 000000000..bd3c18e8e --- /dev/null +++ b/Sources/CodexBar/Providers/StepFun/StepFunSettingsStore.swift @@ -0,0 +1,14 @@ +import CodexBarCore +import Foundation + +extension SettingsStore { + var stepfunAPIToken: String { + get { self.configSnapshot.providerConfig(for: .stepfun)?.sanitizedAPIKey ?? "" } + set { + self.updateProviderConfig(provider: .stepfun) { entry in + entry.apiKey = self.normalizedConfigValue(newValue) + } + self.logSecretUpdate(provider: .stepfun, field: "apiKey", value: newValue) + } + } +} diff --git a/Sources/CodexBar/Providers/Trae/TraeProviderImplementation.swift b/Sources/CodexBar/Providers/Trae/TraeProviderImplementation.swift new file mode 100644 index 000000000..bd4b641ee --- /dev/null +++ b/Sources/CodexBar/Providers/Trae/TraeProviderImplementation.swift @@ -0,0 +1,40 @@ +import AppKit +import CodexBarCore +import CodexBarMacroSupport +import Foundation + +@ProviderImplementationRegistration +struct TraeProviderImplementation: ProviderImplementation { + let id: UsageProvider = .trae + + @MainActor + func observeSettings(_ settings: SettingsStore) {} + + @MainActor + func settingsFields(context: ProviderSettingsContext) -> [ProviderSettingsFieldDescriptor] { + [ + ProviderSettingsFieldDescriptor( + id: "trae-info", + title: "Status", + subtitle: "Trae is ByteDance's free AI IDE. CodexBar detects whether Trae is running " + + "on this machine.", + kind: .plain, + placeholder: "", + binding: context.stringBinding(\.traeInfo), + actions: [ + ProviderSettingsActionDescriptor( + id: "trae-open-website", + title: "Open Trae Website", + style: .link, + isVisible: nil, + perform: { + if let url = URL(string: "https://www.trae.ai") { + NSWorkspace.shared.open(url) + } + }), + ], + isVisible: nil, + onActivate: nil), + ] + } +} diff --git a/Sources/CodexBar/Providers/Trae/TraeSettingsStore.swift b/Sources/CodexBar/Providers/Trae/TraeSettingsStore.swift new file mode 100644 index 000000000..a9a83e60d --- /dev/null +++ b/Sources/CodexBar/Providers/Trae/TraeSettingsStore.swift @@ -0,0 +1,10 @@ +import CodexBarCore +import Foundation + +extension SettingsStore { + var traeInfo: String { + get { "" } + // swiftlint:disable:next unused_setter_value + set {} + } +} diff --git a/Sources/CodexBar/Providers/Zenmux/ZenmuxProviderImplementation.swift b/Sources/CodexBar/Providers/Zenmux/ZenmuxProviderImplementation.swift new file mode 100644 index 000000000..06a9cae81 --- /dev/null +++ b/Sources/CodexBar/Providers/Zenmux/ZenmuxProviderImplementation.swift @@ -0,0 +1,42 @@ +import AppKit +import CodexBarCore +import CodexBarMacroSupport +import Foundation + +@ProviderImplementationRegistration +struct ZenmuxProviderImplementation: ProviderImplementation { + let id: UsageProvider = .zenmux + + @MainActor + func observeSettings(_ settings: SettingsStore) { + _ = settings.zenmuxAPIToken + } + + @MainActor + func settingsFields(context: ProviderSettingsContext) -> [ProviderSettingsFieldDescriptor] { + [ + ProviderSettingsFieldDescriptor( + id: "zenmux-api-token", + title: "API key", + subtitle: "Stored in ~/.codexbar/config.json. Get your API key from the Zenmux " + + "dashboard.", + kind: .secure, + placeholder: "sk-ss-v1-... or sk-ai-v1-...", + binding: context.stringBinding(\.zenmuxAPIToken), + actions: [ + ProviderSettingsActionDescriptor( + id: "zenmux-open-dashboard", + title: "Open Zenmux Dashboard", + style: .link, + isVisible: nil, + perform: { + if let url = URL(string: "https://zenmux.ai/") { + NSWorkspace.shared.open(url) + } + }), + ], + isVisible: nil, + onActivate: nil), + ] + } +} diff --git a/Sources/CodexBar/Providers/Zenmux/ZenmuxSettingsStore.swift b/Sources/CodexBar/Providers/Zenmux/ZenmuxSettingsStore.swift new file mode 100644 index 000000000..d5c231716 --- /dev/null +++ b/Sources/CodexBar/Providers/Zenmux/ZenmuxSettingsStore.swift @@ -0,0 +1,14 @@ +import CodexBarCore +import Foundation + +extension SettingsStore { + var zenmuxAPIToken: String { + get { self.configSnapshot.providerConfig(for: .zenmux)?.sanitizedAPIKey ?? "" } + set { + self.updateProviderConfig(provider: .zenmux) { entry in + entry.apiKey = self.normalizedConfigValue(newValue) + } + self.logSecretUpdate(provider: .zenmux, field: "apiKey", value: newValue) + } + } +} diff --git a/Sources/CodexBar/Resources/ProviderIcon-aigocode.svg b/Sources/CodexBar/Resources/ProviderIcon-aigocode.svg new file mode 100644 index 000000000..95a151482 --- /dev/null +++ b/Sources/CodexBar/Resources/ProviderIcon-aigocode.svg @@ -0,0 +1 @@ +AigoCode diff --git a/Sources/CodexBar/Resources/ProviderIcon-doubao.svg b/Sources/CodexBar/Resources/ProviderIcon-doubao.svg new file mode 100644 index 000000000..9c20430a1 --- /dev/null +++ b/Sources/CodexBar/Resources/ProviderIcon-doubao.svg @@ -0,0 +1 @@ +Doubao diff --git a/Sources/CodexBar/Resources/ProviderIcon-mimo.svg b/Sources/CodexBar/Resources/ProviderIcon-mimo.svg new file mode 100644 index 000000000..50b1b8e3e --- /dev/null +++ b/Sources/CodexBar/Resources/ProviderIcon-mimo.svg @@ -0,0 +1,4 @@ + + + + diff --git a/Sources/CodexBar/Resources/ProviderIcon-qwen.svg b/Sources/CodexBar/Resources/ProviderIcon-qwen.svg new file mode 100644 index 000000000..46e5ab262 --- /dev/null +++ b/Sources/CodexBar/Resources/ProviderIcon-qwen.svg @@ -0,0 +1 @@ +Qwen diff --git a/Sources/CodexBar/Resources/ProviderIcon-trae.svg b/Sources/CodexBar/Resources/ProviderIcon-trae.svg new file mode 100644 index 000000000..2c45783a9 --- /dev/null +++ b/Sources/CodexBar/Resources/ProviderIcon-trae.svg @@ -0,0 +1 @@ +Trae diff --git a/Sources/CodexBar/Resources/ProviderIcon-zenmux.svg b/Sources/CodexBar/Resources/ProviderIcon-zenmux.svg new file mode 100644 index 000000000..44cffd96d --- /dev/null +++ b/Sources/CodexBar/Resources/ProviderIcon-zenmux.svg @@ -0,0 +1 @@ +Zenmux diff --git a/Sources/CodexBar/SettingsStore.swift b/Sources/CodexBar/SettingsStore.swift index 09f3e3caa..0a55a546d 100644 --- a/Sources/CodexBar/SettingsStore.swift +++ b/Sources/CodexBar/SettingsStore.swift @@ -157,8 +157,11 @@ final class SettingsStore { extension SettingsStore { private static func loadDefaultsState(userDefaults: UserDefaults) -> SettingsDefaultsState { - let refreshRaw = userDefaults.string(forKey: "refreshFrequency") ?? RefreshFrequency.fiveMinutes.rawValue - let refreshFrequency = RefreshFrequency(rawValue: refreshRaw) ?? .fiveMinutes + let refreshRaw = userDefaults.string(forKey: "refreshFrequency") + let refreshFrequency = refreshRaw.flatMap(RefreshFrequency.init(rawValue:)) ?? .fiveMinutes + if refreshRaw == nil { + userDefaults.set(refreshFrequency.rawValue, forKey: "refreshFrequency") + } let launchAtLogin = userDefaults.object(forKey: "launchAtLogin") as? Bool ?? false let debugMenuEnabled = userDefaults.object(forKey: "debugMenuEnabled") as? Bool ?? false let debugDisableKeychainAccess: Bool = { diff --git a/Sources/CodexBar/UsageStore.swift b/Sources/CodexBar/UsageStore.swift index e3f0bcae4..9df485713 100644 --- a/Sources/CodexBar/UsageStore.swift +++ b/Sources/CodexBar/UsageStore.swift @@ -1,3 +1,4 @@ +// swiftlint:disable file_length import AppKit import CodexBarCore import Foundation @@ -1140,6 +1141,7 @@ extension UsageStore { await AugmentStatusProbe.latestDumps() } + // swiftlint:disable:next cyclomatic_complexity func debugLog(for provider: UsageProvider) async -> String { if let cached = self.probeLogs[provider], !cached.isEmpty { return cached @@ -1240,6 +1242,32 @@ extension UsageStore { let hasAny = resolution != nil let source = resolution?.source.rawValue ?? "none" text = "WARP_API_KEY=\(hasAny ? "present" : "missing") source=\(source)" + case .qwen: + let resolution = ProviderTokenResolver.qwenResolution() + let hasAny = resolution != nil + let source = resolution?.source.rawValue ?? "none" + text = "DASHSCOPE_API_KEY=\(hasAny ? "present" : "missing") source=\(source)" + case .doubao: + let resolution = ProviderTokenResolver.doubaoResolution() + let hasAny = resolution != nil + let source = resolution?.source.rawValue ?? "none" + text = "ARK_API_KEY=\(hasAny ? "present" : "missing") source=\(source)" + case .zenmux: + let resolution = ProviderTokenResolver.zenmuxResolution() + let hasAny = resolution != nil + let source = resolution?.source.rawValue ?? "none" + text = "ZENMUX_API_KEY=\(hasAny ? "present" : "missing") source=\(source)" + case .aigocode: + let resolution = ProviderTokenResolver.aigocodeResolution() + let hasAny = resolution != nil + let source = resolution?.source.rawValue ?? "none" + text = "AIGOCODE_API_KEY=\(hasAny ? "present" : "missing") source=\(source)" + case .trae: + text = "Trae: local probe (no API key needed)" + case .stepfun: + text = "StepFun: local probe (no API key needed)" + case .mimo: + text = "MiMo: local probe (cookie-based)" case .gemini, .antigravity, .opencode, .factory, .copilot, .vertexai, .kilo, .kiro, .kimi, .kimik2, .jetbrains: text = unimplementedDebugLogMessages[provider] ?? "Debug log not yet implemented" diff --git a/Sources/CodexBarCLI/TokenAccountCLI.swift b/Sources/CodexBarCLI/TokenAccountCLI.swift index 7324fa837..8744cf364 100644 --- a/Sources/CodexBarCLI/TokenAccountCLI.swift +++ b/Sources/CodexBarCLI/TokenAccountCLI.swift @@ -157,7 +157,8 @@ struct TokenAccountCLIContext { return self.makeSnapshot( jetbrains: ProviderSettingsSnapshot.JetBrainsProviderSettings( ideBasePath: nil)) - case .gemini, .antigravity, .copilot, .kiro, .vertexai, .kimik2, .synthetic, .openrouter, .warp: + case .gemini, .antigravity, .copilot, .kiro, .vertexai, .kimik2, .synthetic, .openrouter, .warp, + .qwen, .doubao, .zenmux, .aigocode, .trae, .stepfun, .mimo: return nil } } diff --git a/Sources/CodexBarCore/Config/ProviderConfigEnvironment.swift b/Sources/CodexBarCore/Config/ProviderConfigEnvironment.swift index 9fabc4b80..4a1fd1cd8 100644 --- a/Sources/CodexBarCore/Config/ProviderConfigEnvironment.swift +++ b/Sources/CodexBarCore/Config/ProviderConfigEnvironment.swift @@ -1,6 +1,7 @@ import Foundation public enum ProviderConfigEnvironment { + // swiftlint:disable:next cyclomatic_complexity public static func applyAPIKeyOverride( base: [String: String], provider: UsageProvider, @@ -29,6 +30,30 @@ public enum ProviderConfigEnvironment { } case .openrouter: env[OpenRouterSettingsReader.envKey] = apiKey + case .qwen: + if let key = QwenSettingsReader.apiKeyEnvironmentKeys.first { + env[key] = apiKey + } + case .doubao: + if let key = DoubaoSettingsReader.apiKeyEnvironmentKeys.first { + env[key] = apiKey + } + case .zenmux: + if let key = ZenmuxSettingsReader.apiKeyEnvironmentKeys.first { + env[key] = apiKey + } + case .aigocode: + if let key = AigoCodeSettingsReader.apiKeyEnvironmentKeys.first { + env[key] = apiKey + } + case .trae: + if let key = TraeSettingsReader.apiKeyEnvironmentKeys.first { + env[key] = apiKey + } + case .stepfun: + if let key = StepFunSettingsReader.apiKeyEnvironmentKeys.first { + env[key] = apiKey + } default: break } diff --git a/Sources/CodexBarCore/CostUsageFetcher.swift b/Sources/CodexBarCore/CostUsageFetcher.swift index 2243d5218..a21eb577e 100644 --- a/Sources/CodexBarCore/CostUsageFetcher.swift +++ b/Sources/CodexBarCore/CostUsageFetcher.swift @@ -40,8 +40,10 @@ public struct CostUsageFetcher: Sendable { } else if provider == .claude { options.claudeLogProviderFilter = .excludeVertexAI } + // Always set TTL to 0: the scanner-level TTL is redundant with per-file mtime/size + // caching in processClaudeFile, and a non-zero TTL causes stale data after the first scan. + options.refreshMinIntervalSeconds = 0 if forceRefresh { - options.refreshMinIntervalSeconds = 0 options.forceRescan = true } var daily = CostUsageScanner.loadDailyReport( diff --git a/Sources/CodexBarCore/Logging/LogCategories.swift b/Sources/CodexBarCore/Logging/LogCategories.swift index 37a7726ef..3564314d7 100644 --- a/Sources/CodexBarCore/Logging/LogCategories.swift +++ b/Sources/CodexBarCore/Logging/LogCategories.swift @@ -11,6 +11,9 @@ public enum LogCategories { public static let claudeUsage = "claude-usage" public static let codexRPC = "codex-rpc" public static let configMigration = "config-migration" + public static let aigocodeUsage = "aigocode-usage" + public static let aigocodeWeb = "aigocode-web" + public static let doubaoUsage = "doubao-usage" public static let configStore = "config-store" public static let cookieCache = "cookie-cache" public static let cookieHeaderStore = "cookie-header-store" @@ -44,6 +47,7 @@ public enum LogCategories { public static let opencodeUsage = "opencode-usage" public static let openRouterUsage = "openrouter-usage" public static let providerDetection = "provider-detection" + public static let qwenUsage = "qwen-usage" public static let providers = "providers" public static let sessionQuota = "sessionQuota" public static let sessionQuotaNotifications = "sessionQuotaNotifications" @@ -61,4 +65,9 @@ public enum LogCategories { public static let zaiSettings = "zai-settings" public static let zaiTokenStore = "zai-token-store" public static let zaiUsage = "zai-usage" + public static let traeCookie = "trae-cookie" + public static let traeUsage = "trae-usage" + public static let traeWeb = "trae-web" + public static let zenmuxUsage = "zenmux-usage" + public static let stepfunUsage = "stepfun-usage" } diff --git a/Sources/CodexBarCore/Providers/AigoCode/AigoCodeCookieImporter.swift b/Sources/CodexBarCore/Providers/AigoCode/AigoCodeCookieImporter.swift new file mode 100644 index 000000000..1bd70e5b7 --- /dev/null +++ b/Sources/CodexBarCore/Providers/AigoCode/AigoCodeCookieImporter.swift @@ -0,0 +1,164 @@ +#if os(macOS) +import Foundation +import SweetCookieKit + +/// Imports the AigoCode Supabase session from Chrome's localStorage. +/// +/// AigoCode uses Supabase Auth which stores JWT tokens in localStorage under the key +/// `sb-myptlcacxbuuxldgouqt-auth-token`, **not** in cookies. We use SweetCookieKit's +/// `ChromiumLocalStorageReader` to properly parse Chrome's LevelDB files. +public enum AigoCodeLocalStorageImporter { + private static let log = CodexBarLog.logger(LogCategories.aigocodeWeb) + + /// The Supabase localStorage key for AigoCode's project. + static let supabaseTokenKey = "sb-myptlcacxbuuxldgouqt-auth-token" + + /// The origins where the token may be stored. + private static let origins = [ + "https://www.aigocode.com", + "https://aigocode.com", + ] + + /// Extracted Supabase session from browser localStorage. + public struct SessionInfo: Sendable { + /// The raw JSON string stored under the Supabase token key. + public let tokenJSON: String + /// Which browser/profile the token was found in. + public let sourceLabel: String + } + + /// Attempts to extract the Supabase session from Chrome's localStorage. + public static func importSession( + browserDetection: BrowserDetection = BrowserDetection(), + logger: ((String) -> Void)? = nil) -> SessionInfo? + { + let log: (String) -> Void = { msg in + logger?("[aigocode-storage] \(msg)") + self.log.debug(msg) + } + + let candidates = self.chromeLocalStorageCandidates(browserDetection: browserDetection) + log("Found \(candidates.count) Chrome profile candidate(s)") + + for candidate in candidates { + if let session = self.readSupabaseSession(from: candidate.levelDBURL, label: candidate.label, logger: log) { + return session + } + } + + log("No Supabase session found in any browser profile") + return nil + } + + /// Quick check for session availability. + public static func hasSession( + browserDetection: BrowserDetection = BrowserDetection()) -> Bool + { + self.importSession(browserDetection: browserDetection) != nil + } + + // MARK: - LevelDB Reading + + private static func readSupabaseSession( + from levelDBURL: URL, + label: String, + logger: ((String) -> Void)?) -> SessionInfo? + { + // Use SweetCookieKit's proper LevelDB parser for origin-scoped entries + for origin in self.origins { + let entries = ChromiumLocalStorageReader.readEntries( + for: origin, + in: levelDBURL, + logger: logger) + + for entry in entries where entry.key == self.supabaseTokenKey { + let value = entry.value.trimmingCharacters(in: .controlCharacters) + guard !value.isEmpty else { continue } + // Validate it's actually JSON with access_token + if let data = value.data(using: .utf8), + let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any], + json["access_token"] != nil + { + logger?("Found valid Supabase session in \(label) (origin: \(origin))") + return SessionInfo(tokenJSON: value, sourceLabel: label) + } + } + } + + // Fallback: scan text entries for the key (handles edge cases) + let textEntries = ChromiumLocalStorageReader.readTextEntries( + in: levelDBURL, + logger: logger) + + for entry in textEntries where entry.key.contains(self.supabaseTokenKey) { + let value = entry.value.trimmingCharacters(in: .controlCharacters) + guard !value.isEmpty else { continue } + if let data = value.data(using: .utf8), + let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any], + json["access_token"] != nil + { + logger?("Found valid Supabase session in \(label) (text scan)") + return SessionInfo(tokenJSON: value, sourceLabel: label) + } + } + + return nil + } + + // MARK: - Chrome Profile Discovery + + private struct LocalStorageCandidate { + let label: String + let levelDBURL: URL + } + + private static func chromeLocalStorageCandidates( + browserDetection: BrowserDetection) -> [LocalStorageCandidate] + { + let browsers: [Browser] = [ + .chrome, + .chromeBeta, + .chromeCanary, + .arc, + .arcBeta, + .arcCanary, + .chromium, + ] + + let installedBrowsers = browsers.browsersWithProfileData(using: browserDetection) + let roots = ChromiumProfileLocator + .roots(for: installedBrowsers, homeDirectories: BrowserCookieClient.defaultHomeDirectories()) + .map { (url: $0.url, labelPrefix: $0.labelPrefix) } + + var candidates: [LocalStorageCandidate] = [] + for root in roots { + candidates.append(contentsOf: self.profileCandidates(root: root.url, labelPrefix: root.labelPrefix)) + } + return candidates + } + + private static func profileCandidates(root: URL, labelPrefix: String) -> [LocalStorageCandidate] { + guard let entries = try? FileManager.default.contentsOfDirectory( + at: root, + includingPropertiesForKeys: [.isDirectoryKey], + options: [.skipsHiddenFiles]) + else { return [] } + + let profileDirs = entries.filter { url in + guard let isDir = (try? url.resourceValues(forKeys: [.isDirectoryKey]).isDirectory), isDir else { + return false + } + let name = url.lastPathComponent + return name == "Default" || name.hasPrefix("Profile ") || name.hasPrefix("user-") + } + .sorted { $0.lastPathComponent < $1.lastPathComponent } + + return profileDirs.compactMap { dir in + let levelDBURL = dir.appendingPathComponent("Local Storage").appendingPathComponent("leveldb") + guard FileManager.default.fileExists(atPath: levelDBURL.path) else { return nil } + let label = "\(labelPrefix) \(dir.lastPathComponent)" + return LocalStorageCandidate(label: label, levelDBURL: levelDBURL) + } + } +} +#endif diff --git a/Sources/CodexBarCore/Providers/AigoCode/AigoCodeDashboardFetcher.swift b/Sources/CodexBarCore/Providers/AigoCode/AigoCodeDashboardFetcher.swift new file mode 100644 index 000000000..b13381267 --- /dev/null +++ b/Sources/CodexBarCore/Providers/AigoCode/AigoCodeDashboardFetcher.swift @@ -0,0 +1,316 @@ +#if os(macOS) +import Foundation +import WebKit + +/// Scrapes the AigoCode dashboard to extract usage data. +/// +/// AigoCode uses Supabase + Next.js with server-side rendering. The usage data is only +/// available by rendering the full dashboard page, so we use an offscreen WKWebView +/// to load the page and extract values from the DOM via JavaScript. +@MainActor +public struct AigoCodeDashboardFetcher { + public enum FetchError: LocalizedError { + case loginRequired + case noUsageData(body: String) + case timeout + + public var errorDescription: String? { + switch self { + case .loginRequired: + "AigoCode web access requires login. Open Settings → AigoCode → Login in Browser." + case let .noUsageData(body): + "AigoCode dashboard data not found. Body sample: \(body.prefix(200))" + case .timeout: + "AigoCode dashboard loading timed out." + } + } + } + + private static let log = CodexBarLog.logger(LogCategories.aigocodeWeb) + private static let dashboardURL = URL(string: "https://www.aigocode.com/dashboard/console")! + + public init() {} + + // MARK: - Public + + public func fetchDashboard( + websiteDataStore: WKWebsiteDataStore = .default(), + supabaseTokenJSON: String? = nil, + timeout: TimeInterval = 45) async throws -> AigoCodeDashboardSnapshot + { + let deadline = Date().addingTimeInterval(max(1, timeout)) + + let config = WKWebViewConfiguration() + config.websiteDataStore = websiteDataStore + let webView = WKWebView(frame: CGRect(x: -9999, y: -9999, width: 1200, height: 900), configuration: config) + webView.customUserAgent = + "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 " + + "(KHTML, like Gecko) Chrome/145.0.0.0 Safari/537.36" + + defer { + webView.stopLoading() + webView.loadHTMLString("", baseURL: nil) + } + + // If we have a Supabase token from Chrome, inject it into localStorage first. + // We load a blank page on the AigoCode origin, set the token, then navigate. + if let supabaseTokenJSON { + Self.log.debug("Injecting Supabase session into localStorage") + _ = webView.load(URLRequest(url: URL(string: "https://www.aigocode.com/favicon.ico")!)) + try? await Task.sleep(for: .milliseconds(2000)) + + let escaped = supabaseTokenJSON + .replacingOccurrences(of: "\\", with: "\\\\") + .replacingOccurrences(of: "'", with: "\\'") + let injectJS = + "localStorage.setItem('\(AigoCodeLocalStorageImporter.supabaseTokenKey)', '\(escaped)'); 'ok';" // swiftlint:disable:this line_length + let result = try? await webView.evaluateJavaScript(injectJS) + Self.log.debug("localStorage injection result: \(String(describing: result))") + } + + _ = webView.load(URLRequest(url: Self.dashboardURL)) + Self.log.debug("Loading AigoCode dashboard…") + + // Poll until we find usage data or hit the deadline + var lastBody = "" + while Date() < deadline { + try? await Task.sleep(for: .milliseconds(1500)) + + let scrape = try await self.scrape(webView: webView) + + // Detect login page + if scrape.isLoginPage { + Self.log.debug("Login page detected") + throw FetchError.loginRequired + } + + lastBody = scrape.bodyText + + // Check if we have subscription usage data + if let snapshot = scrape.snapshot { + Self.log.debug( + "Dashboard parsed: subscription=\(snapshot.subscriptionUsedDollars)/\(snapshot.subscriptionTotalDollars) " + // swiftlint:disable:this line_length + "weekly=\(snapshot.weeklyUsedDollars)/\(snapshot.weeklyTotalDollars)") + return snapshot + } + } + + throw FetchError.noUsageData(body: lastBody) + } + + // MARK: - JavaScript Scraping + + private struct ScrapeResult { + let isLoginPage: Bool + let bodyText: String + let snapshot: AigoCodeDashboardSnapshot? + } + + private func scrape(webView: WKWebView) async throws -> ScrapeResult { + let js = """ + (() => { + const href = window.location.href; + const body = document.body ? document.body.innerText : ''; + + // Detect login page + const isLogin = href.includes('/auth/login') || + body.includes('欢迎回来') && body.includes('使用 Google 登录'); + + if (isLogin) { + return JSON.stringify({ isLogin: true, body: body.substring(0, 500) }); + } + + // Extract usage data from the console/stats page + // Pattern: "已用 $X / 共 $Y" or "已用 $X / $Y" + const usagePattern = /已用\\s*\\$([\\d,.]+)\\s*\\/\\s*共?\\s*\\$([\\d,.]+)/g; + const usages = []; + let match; + while ((match = usagePattern.exec(body)) !== null) { + usages.push({ used: match[1].replace(/,/g, ''), total: match[2].replace(/,/g, '') }); + } + + // Extract plan info + // Pattern: "Professional Plan" or similar, followed by expiration + const planMatch = body.match(/(\\w+\\s*Plan)[，,]?\\s*到期\\s*([\\d/]+)/); + const plan = planMatch ? planMatch[1] : null; + const expiry = planMatch ? planMatch[2] : null; + + // Extract flexible balance + // Pattern: "<$0.01" or "$1.23" + const flexMatch = body.match(/灵活余额[\\s\\S]{0,50}?[<>]?\\$([\\d,.]+)/); + const flexBalance = flexMatch ? flexMatch[1].replace(/,/g, '') : null; + + // Extract weekly reset info + // Pattern: "X天Y小时后重置" or "X小时后重置" + const resetMatch = body.match(/(\\d+天)?(\\d+小时)?后重置/); + const resetText = resetMatch ? resetMatch[0] : null; + + // Extract usage percentage + const pctPattern = /已使用\\s*(\\d+)%/g; + const pcts = []; + let pctMatch; + while ((pctMatch = pctPattern.exec(body)) !== null) { + pcts.push(parseInt(pctMatch[1])); + } + + return JSON.stringify({ + isLogin: false, + body: body.substring(0, 1000), + usages: usages, + plan: plan, + expiry: expiry, + flexBalance: flexBalance, + resetText: resetText, + pcts: pcts, + href: href + }); + })(); + """ + + guard let resultStr = try await webView.evaluateJavaScript(js) as? String, + let data = resultStr.data(using: .utf8), + let dict = try? JSONSerialization.jsonObject(with: data) as? [String: Any] + else { + return ScrapeResult(isLoginPage: false, bodyText: "", snapshot: nil) + } + + let isLogin = (dict["isLogin"] as? Bool) ?? false + let bodyText = (dict["body"] as? String) ?? "" + + if isLogin { + return ScrapeResult(isLoginPage: true, bodyText: bodyText, snapshot: nil) + } + + guard let usages = dict["usages"] as? [[String: String]], !usages.isEmpty else { + return ScrapeResult(isLoginPage: false, bodyText: bodyText, snapshot: nil) + } + + // First usage entry = subscription, second = weekly (if present) + let subUsed = Double(usages[0]["used"] ?? "0") ?? 0 + let subTotal = Double(usages[0]["total"] ?? "0") ?? 0 + let weekUsed = usages.count > 1 ? (Double(usages[1]["used"] ?? "0") ?? 0) : 0 + let weekTotal = usages.count > 1 ? (Double(usages[1]["total"] ?? "0") ?? 0) : 0 + + let plan = dict["plan"] as? String + let expiry = dict["expiry"] as? String + let flexBalance = Double((dict["flexBalance"] as? String) ?? "0") ?? 0 + let resetText = dict["resetText"] as? String + + let snapshot = AigoCodeDashboardSnapshot( + subscriptionUsedDollars: subUsed, + subscriptionTotalDollars: subTotal, + weeklyUsedDollars: weekUsed, + weeklyTotalDollars: weekTotal, + planName: plan, + planExpiry: expiry, + flexibleBalanceDollars: flexBalance, + weeklyResetText: resetText, + updatedAt: Date()) + + return ScrapeResult(isLoginPage: false, bodyText: bodyText, snapshot: snapshot) + } +} + +// MARK: - Dashboard Snapshot + +public struct AigoCodeDashboardSnapshot: Sendable { + public let subscriptionUsedDollars: Double + public let subscriptionTotalDollars: Double + public let weeklyUsedDollars: Double + public let weeklyTotalDollars: Double + public let planName: String? + public let planExpiry: String? + public let flexibleBalanceDollars: Double + public let weeklyResetText: String? + public let updatedAt: Date + + public init( + subscriptionUsedDollars: Double, + subscriptionTotalDollars: Double, + weeklyUsedDollars: Double, + weeklyTotalDollars: Double, + planName: String?, + planExpiry: String?, + flexibleBalanceDollars: Double, + weeklyResetText: String?, + updatedAt: Date) + { + self.subscriptionUsedDollars = subscriptionUsedDollars + self.subscriptionTotalDollars = subscriptionTotalDollars + self.weeklyUsedDollars = weeklyUsedDollars + self.weeklyTotalDollars = weeklyTotalDollars + self.planName = planName + self.planExpiry = planExpiry + self.flexibleBalanceDollars = flexibleBalanceDollars + self.weeklyResetText = weeklyResetText + self.updatedAt = updatedAt + } + + public func toUsageSnapshot() -> UsageSnapshot { + let subPercent: Double + let subDescription: String + if self.subscriptionTotalDollars > 0 { + subPercent = min(100, max(0, self.subscriptionUsedDollars / self.subscriptionTotalDollars * 100)) + subDescription = "$\(Self.fmt(self.subscriptionUsedDollars))/$\(Self.fmt(self.subscriptionTotalDollars))" + } else { + subPercent = 0 + subDescription = "No subscription data" + } + + let weekPercent: Double + let weekDescription: String + if self.weeklyTotalDollars > 0 { + weekPercent = min(100, max(0, self.weeklyUsedDollars / self.weeklyTotalDollars * 100)) + var desc = "$\(Self.fmt(self.weeklyUsedDollars))/$\(Self.fmt(self.weeklyTotalDollars))" + if let reset = self.weeklyResetText { + desc += " (\(reset))" + } + weekDescription = desc + } else { + weekPercent = 0 + weekDescription = "No weekly data" + } + + let primary = RateWindow( + usedPercent: subPercent, + windowMinutes: nil, + resetsAt: nil, + resetDescription: subDescription) + + let secondary = RateWindow( + usedPercent: weekPercent, + windowMinutes: nil, + resetsAt: nil, + resetDescription: weekDescription) + + var planDescription: String? + if let plan = self.planName { + planDescription = plan + if let expiry = self.planExpiry { + planDescription! += ", expires \(expiry)" + } + } + + let identity = ProviderIdentitySnapshot( + providerID: .aigocode, + accountEmail: planDescription, + accountOrganization: nil, + loginMethod: "Web") + + return UsageSnapshot( + primary: primary, + secondary: secondary, + tertiary: nil, + providerCost: nil, + updatedAt: self.updatedAt, + identity: identity) + } + + private static func fmt(_ value: Double) -> String { + if value == Double(Int(value)) { + return String(Int(value)) + } + return String(format: "%.2f", value) + } +} +#endif diff --git a/Sources/CodexBarCore/Providers/AigoCode/AigoCodeProviderDescriptor.swift b/Sources/CodexBarCore/Providers/AigoCode/AigoCodeProviderDescriptor.swift new file mode 100644 index 000000000..2db0b956b --- /dev/null +++ b/Sources/CodexBarCore/Providers/AigoCode/AigoCodeProviderDescriptor.swift @@ -0,0 +1,122 @@ +import CodexBarMacroSupport +import Foundation + +@ProviderDescriptorRegistration +@ProviderDescriptorDefinition +public enum AigoCodeProviderDescriptor { + static func makeDescriptor() -> ProviderDescriptor { + ProviderDescriptor( + id: .aigocode, + metadata: ProviderMetadata( + id: .aigocode, + displayName: "AigoCode", + sessionLabel: "Subscription", + weeklyLabel: "Weekly", + opusLabel: nil, + supportsOpus: false, + supportsCredits: false, + creditsHint: "", + toggleTitle: "Show AigoCode usage", + cliName: "aigocode", + defaultEnabled: false, + isPrimaryProvider: false, + usesAccountFallback: false, + browserCookieOrder: nil, + dashboardURL: "https://www.aigocode.com/dashboard/console", + statusPageURL: nil), + branding: ProviderBranding( + iconStyle: .aigocode, + iconResourceName: "ProviderIcon-aigocode", + color: ProviderColor(red: 34 / 255, green: 197 / 255, blue: 94 / 255)), + tokenCost: ProviderTokenCostConfig( + supportsTokenCost: false, + noDataMessage: { "AigoCode cost summary is not available." }), + fetchPlan: ProviderFetchPlan( + sourceModes: [.auto, .web, .api], + pipeline: ProviderFetchPipeline(resolveStrategies: { context in + var strategies: [any ProviderFetchStrategy] = [] + // Prefer web dashboard when available (works without API key) + #if os(macOS) + if context.sourceMode.usesWeb || context.sourceMode == .auto { + strategies.append(AigoCodeWebDashboardFetchStrategy()) + } + #endif + strategies.append(AigoCodeAPIFetchStrategy()) + return strategies + })), + cli: ProviderCLIConfig( + name: "aigocode", + aliases: ["aigo"], + versionDetector: nil)) + } +} + +struct AigoCodeAPIFetchStrategy: ProviderFetchStrategy { + let id: String = "aigocode.api" + let kind: ProviderFetchKind = .apiToken + + func isAvailable(_ context: ProviderFetchContext) async -> Bool { + Self.resolveToken(environment: context.env) != nil + } + + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + guard let apiKey = Self.resolveToken(environment: context.env) else { + throw AigoCodeUsageError.missingCredentials + } + let usage = try await AigoCodeUsageFetcher.fetchUsage(apiKey: apiKey) + return self.makeResult( + usage: usage.toUsageSnapshot(), + sourceLabel: "api") + } + + func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool { + false + } + + private static func resolveToken(environment: [String: String]) -> String? { + ProviderTokenResolver.aigocodeToken(environment: environment) + } +} + +#if os(macOS) +/// Fetches AigoCode usage by rendering the dashboard in an offscreen WKWebView. +/// This works without an API key — the user just needs to be logged in via the WebKit session. +struct AigoCodeWebDashboardFetchStrategy: ProviderFetchStrategy { + let id: String = "aigocode.webDashboard" + let kind: ProviderFetchKind = .webDashboard + private static let log = CodexBarLog.logger(LogCategories.aigocodeWeb) + + /// The web strategy is always considered available on macOS. If the user isn't logged in, + /// `fetch` will throw `loginRequired` and the pipeline falls back to the API strategy. + func isAvailable(_ context: ProviderFetchContext) async -> Bool { + context.sourceMode.usesWeb || context.sourceMode == .auto + } + + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + // Try to import the Supabase session from Chrome's localStorage. + // AigoCode uses Supabase Auth which stores JWT tokens in localStorage, not cookies. + let session = AigoCodeLocalStorageImporter.importSession() + if let session { + Self.log.debug("Found Supabase session in \(session.sourceLabel)") + } else { + Self.log.debug("No Supabase session found in browser localStorage") + } + + let snapshot = try await MainActor.run { + AigoCodeDashboardFetcher() + }.fetchDashboard( + supabaseTokenJSON: session?.tokenJSON, + timeout: context.webTimeout) + return self.makeResult( + usage: snapshot.toUsageSnapshot(), + sourceLabel: "webDashboard") + } + + func shouldFallback(on error: Error, context: ProviderFetchContext) -> Bool { + // Fall back to API strategy if web fails (login required, timeout, etc.) + if context.sourceMode == .auto { return true } + if context.sourceMode == .web { return false } + return true + } +} +#endif diff --git a/Sources/CodexBarCore/Providers/AigoCode/AigoCodeSettingsReader.swift b/Sources/CodexBarCore/Providers/AigoCode/AigoCodeSettingsReader.swift new file mode 100644 index 000000000..c5534f622 --- /dev/null +++ b/Sources/CodexBarCore/Providers/AigoCode/AigoCodeSettingsReader.swift @@ -0,0 +1,35 @@ +import Foundation + +public struct AigoCodeSettingsReader: Sendable { + public static let apiKeyEnvironmentKeys = [ + "AIGOCODE_API_KEY", + ] + + public static func apiKey( + environment: [String: String] = ProcessInfo.processInfo.environment) -> String? + { + for key in self.apiKeyEnvironmentKeys { + guard let raw = environment[key]?.trimmingCharacters(in: .whitespacesAndNewlines), + !raw.isEmpty + else { + continue + } + let cleaned = Self.cleaned(raw) + if !cleaned.isEmpty { + return cleaned + } + } + return nil + } + + private static func cleaned(_ raw: String) -> String { + var value = raw + if (value.hasPrefix("\"") && value.hasSuffix("\"")) || + (value.hasPrefix("'") && value.hasSuffix("'")) + { + value.removeFirst() + value.removeLast() + } + return value.trimmingCharacters(in: .whitespacesAndNewlines) + } +} diff --git a/Sources/CodexBarCore/Providers/AigoCode/AigoCodeUsageFetcher.swift b/Sources/CodexBarCore/Providers/AigoCode/AigoCodeUsageFetcher.swift new file mode 100644 index 000000000..db0887926 --- /dev/null +++ b/Sources/CodexBarCore/Providers/AigoCode/AigoCodeUsageFetcher.swift @@ -0,0 +1,288 @@ +import Foundation +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif + +public struct AigoCodeUsageSnapshot: Sendable { + public let remainingRequests: Int + public let limitRequests: Int + public let resetTime: Date? + public let updatedAt: Date + public let apiKeyValid: Bool + public let totalTokens: Int? + public let statusMessage: String? + public let apiKey: String? + + public init( + remainingRequests: Int, + limitRequests: Int, + resetTime: Date?, + updatedAt: Date, + apiKeyValid: Bool = false, + totalTokens: Int? = nil, + statusMessage: String? = nil, + apiKey: String? = nil) + { + self.remainingRequests = remainingRequests + self.limitRequests = limitRequests + self.resetTime = resetTime + self.updatedAt = updatedAt + self.apiKeyValid = apiKeyValid + self.totalTokens = totalTokens + self.statusMessage = statusMessage + self.apiKey = apiKey + } + + private static func maskedKey(_ key: String?) -> String? { + guard let key, !key.isEmpty else { return nil } + if key.count <= 8 { return "****" } + let prefix = String(key.prefix(6)) + let suffix = String(key.suffix(4)) + return "\(prefix)...\(suffix)" + } + + public func toUsageSnapshot() -> UsageSnapshot { + let usedPercent: Double + let resetDescription: String + + if let msg = self.statusMessage { + usedPercent = 100 + resetDescription = msg + } else if self.limitRequests > 0 { + let used = max(0, self.limitRequests - self.remainingRequests) + usedPercent = min(100, max(0, Double(used) / Double(self.limitRequests) * 100)) + resetDescription = "\(used)/\(self.limitRequests) requests" + } else if self.apiKeyValid { + usedPercent = 0 + resetDescription = "Active — check dashboard for details" + } else { + usedPercent = 0 + resetDescription = "No usage data" + } + + let primary = RateWindow( + usedPercent: usedPercent, + windowMinutes: nil, + resetsAt: self.resetTime, + resetDescription: resetDescription) + + let identity = ProviderIdentitySnapshot( + providerID: .aigocode, + accountEmail: Self.maskedKey(self.apiKey), + accountOrganization: nil, + loginMethod: "API") + + return UsageSnapshot( + primary: primary, + secondary: nil, + tertiary: nil, + providerCost: nil, + updatedAt: self.updatedAt, + identity: identity) + } +} + +public enum AigoCodeUsageError: LocalizedError, Sendable { + case missingCredentials + case networkError(String) + case apiError(Int, String) + case parseFailed(String) + + public var errorDescription: String? { + switch self { + case .missingCredentials: + "Missing AigoCode API key (AIGOCODE_API_KEY)." + case let .networkError(message): + "AigoCode network error: \(message)" + case let .apiError(code, message): + "AigoCode API error (\(code)): \(message)" + case let .parseFailed(message): + "Failed to parse AigoCode response: \(message)" + } + } +} + +public struct AigoCodeUsageFetcher: Sendable { + private static let log = CodexBarLog.logger(LogCategories.aigocodeUsage) + + private static let apiURL = URL(string: "https://api.aigocode.com/v1/messages")! + + public static func fetchUsage(apiKey: String) async throws -> AigoCodeUsageSnapshot { + guard !apiKey.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty else { + throw AigoCodeUsageError.missingCredentials + } + + var request = URLRequest(url: self.apiURL) + request.httpMethod = "POST" + request.timeoutInterval = 15 + request.setValue("application/json", forHTTPHeaderField: "Content-Type") + request.setValue("application/json", forHTTPHeaderField: "Accept") + request.setValue(apiKey, forHTTPHeaderField: "x-api-key") + request.setValue("2023-06-01", forHTTPHeaderField: "anthropic-version") + + let body: [String: Any] = [ + "model": "anthropic/claude-sonnet-4-5", + "max_tokens": 1, + "messages": [ + ["role": "user", "content": "hi"], + ] as [[String: Any]], + ] + + request.httpBody = try JSONSerialization.data(withJSONObject: body) + + let (data, response) = try await URLSession.shared.data(for: request) + + guard let httpResponse = response as? HTTPURLResponse else { + throw AigoCodeUsageError.networkError("Invalid response") + } + + // 403 with INSUFFICIENT_BALANCE means the key is valid but account has no credits + if httpResponse.statusCode == 403 { + if let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any], + let code = json["code"] as? String, code == "INSUFFICIENT_BALANCE" + { + return AigoCodeUsageSnapshot( + remainingRequests: 0, + limitRequests: 0, + resetTime: nil, + updatedAt: Date(), + apiKeyValid: true, + totalTokens: nil, + statusMessage: "Insufficient balance", + apiKey: apiKey) + } + } + + // Accept both 200 (success) and 429 (rate limited). + guard httpResponse.statusCode == 200 || httpResponse.statusCode == 429 else { + let summary = Self.apiErrorSummary(statusCode: httpResponse.statusCode, data: data) + Self.log.error("AigoCode API returned \(httpResponse.statusCode): \(summary)") + throw AigoCodeUsageError.apiError(httpResponse.statusCode, summary) + } + + let headers = httpResponse.allHeaderFields + let remaining = Self.intHeader(headers, "x-ratelimit-remaining-requests") + let limit = Self.intHeader(headers, "x-ratelimit-limit-requests") + let resetString = headers["x-ratelimit-reset-requests"] as? String + + let resetTime: Date? = resetString.flatMap(Self.parseResetTime) + + var totalTokens: Int? + if remaining == nil, limit == nil, + let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any], + let usage = json["usage"] as? [String: Any] + { + totalTokens = usage["total_tokens"] as? Int + } + + let snapshot = AigoCodeUsageSnapshot( + remainingRequests: remaining ?? 0, + limitRequests: limit ?? 0, + resetTime: resetTime, + updatedAt: Date(), + apiKeyValid: httpResponse.statusCode == 200, + totalTokens: totalTokens, + apiKey: apiKey) + + Self.log.debug( + "AigoCode usage parsed remaining=\(snapshot.remainingRequests) limit=\(snapshot.limitRequests) valid=\(snapshot.apiKeyValid)") // swiftlint:disable:this line_length + + return snapshot + } + + private static func intHeader(_ headers: [AnyHashable: Any], _ name: String) -> Int? { + if let value = headers[name] as? String, let int = Int(value) { + return int + } + if let value = headers[name.lowercased()] as? String, let int = Int(value) { + return int + } + // Case-insensitive search + for (key, val) in headers { + if let keyStr = key as? String, + keyStr.lowercased() == name.lowercased(), + let valStr = val as? String, + let int = Int(valStr) + { + return int + } + } + return nil + } + + /// Parse reset time from header value like "1d2h3m4s" or "30s" or ISO 8601. + private static func parseResetTime(_ value: String) -> Date? { + let trimmed = value.trimmingCharacters(in: .whitespacesAndNewlines) + if trimmed.isEmpty { return nil } + + // Try ISO 8601 first + let isoFormatter = ISO8601DateFormatter() + isoFormatter.formatOptions = [.withInternetDateTime, .withFractionalSeconds] + if let date = isoFormatter.date(from: trimmed) { return date } + let isoFallback = ISO8601DateFormatter() + isoFallback.formatOptions = [.withInternetDateTime] + if let date = isoFallback.date(from: trimmed) { return date } + + // Try duration format like "1d2h3m4s" or "30s" + var seconds: TimeInterval = 0 + let pattern = /(\d+)([dhms])/ + for match in trimmed.matches(of: pattern) { + guard let num = Double(match.1) else { continue } + switch match.2 { + case "d": seconds += num * 86400 + case "h": seconds += num * 3600 + case "m": seconds += num * 60 + case "s": seconds += num + default: break + } + } + if seconds > 0 { + return Date().addingTimeInterval(seconds) + } + + // Try plain integer as seconds + if let secs = TimeInterval(trimmed) { + return Date().addingTimeInterval(secs) + } + + return nil + } + + private static func apiErrorSummary(statusCode: Int, data: Data) -> String { + guard let root = try? JSONSerialization.jsonObject(with: data), + let json = root as? [String: Any] + else { + if let text = String(data: data, encoding: .utf8)? + .trimmingCharacters(in: .whitespacesAndNewlines), + !text.isEmpty + { + return self.compactText(text) + } + return "Unexpected response body (\(data.count) bytes)." + } + + if let message = json["message"] as? String { + let trimmed = message.trimmingCharacters(in: .whitespacesAndNewlines) + if !trimmed.isEmpty { return Self.compactText(trimmed) } + } + + if let error = json["error"] as? [String: Any], + let message = error["message"] as? String + { + let trimmed = message.trimmingCharacters(in: .whitespacesAndNewlines) + if !trimmed.isEmpty { return Self.compactText(trimmed) } + } + + return "HTTP \(statusCode) (\(data.count) bytes)." + } + + private static func compactText(_ text: String, maxLength: Int = 200) -> String { + let collapsed = text + .components(separatedBy: .newlines) + .joined(separator: " ") + .trimmingCharacters(in: .whitespacesAndNewlines) + if collapsed.count <= maxLength { return collapsed } + let limitIndex = collapsed.index(collapsed.startIndex, offsetBy: maxLength) + return "\(collapsed[.. [AntigravityModelQuota] { var ordered: [AntigravityModelQuota] = [] - if let claude = models.first(where: { Self.isClaudeWithoutThinking($0.label) }) { - ordered.append(claude) + if let pro = models.first(where: { Self.isGeminiProLow($0.label) }) { + ordered.append(pro) } - if let pro = models.first(where: { Self.isGeminiProLow($0.label) }), - !ordered.contains(where: { $0.label == pro.label }) + if let claude = models.first(where: { Self.isClaudeWithoutThinking($0.label) }), + !ordered.contains(where: { $0.label == claude.label }) { - ordered.append(pro) + ordered.append(claude) } if let flash = models.first(where: { Self.isGeminiFlash($0.label) }), !ordered.contains(where: { $0.label == flash.label }) @@ -108,7 +108,7 @@ public enum AntigravityStatusProbeError: LocalizedError, Sendable, Equatable { public var errorDescription: String? { switch self { case .notRunning: - "Antigravity language server not detected. Launch Antigravity and retry." + "Antigravity language server not detected. Launch Antigravity and retry. If behind a proxy, ensure Antigravity can reach its servers (set http_proxy/https_proxy or enable system proxy)." // swiftlint:disable:this line_length case .missingCSRFToken: "Antigravity CSRF token not found. Restart Antigravity and retry." case let .portDetectionFailed(message): diff --git a/Sources/CodexBarCore/Providers/Doubao/DoubaoProviderDescriptor.swift b/Sources/CodexBarCore/Providers/Doubao/DoubaoProviderDescriptor.swift new file mode 100644 index 000000000..b9beba0f5 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Doubao/DoubaoProviderDescriptor.swift @@ -0,0 +1,73 @@ +import CodexBarMacroSupport +import Foundation + +@ProviderDescriptorRegistration +@ProviderDescriptorDefinition +public enum DoubaoProviderDescriptor { + static func makeDescriptor() -> ProviderDescriptor { + ProviderDescriptor( + id: .doubao, + metadata: ProviderMetadata( + id: .doubao, + displayName: "Doubao", + sessionLabel: "Requests", + weeklyLabel: "Monthly", + opusLabel: nil, + supportsOpus: false, + supportsCredits: false, + creditsHint: "", + toggleTitle: "Show Doubao usage", + cliName: "doubao", + defaultEnabled: false, + isPrimaryProvider: false, + usesAccountFallback: false, + browserCookieOrder: nil, + dashboardURL: "https://console.volcengine.com/ark/region:ark+cn-beijing/openManagement?LLM=%7B%7D&advancedActiveKey=subscribe", + statusPageURL: nil), + branding: ProviderBranding( + iconStyle: .doubao, + iconResourceName: "ProviderIcon-doubao", + color: ProviderColor(red: 51 / 255, green: 112 / 255, blue: 255 / 255)), + tokenCost: ProviderTokenCostConfig( + supportsTokenCost: false, + noDataMessage: { "Doubao cost summary is not available." }), + fetchPlan: ProviderFetchPlan( + sourceModes: [.auto, .api], + pipeline: ProviderFetchPipeline(resolveStrategies: { _ in [DoubaoAPIFetchStrategy()] })), + cli: ProviderCLIConfig( + name: "doubao", + aliases: ["volcengine", "ark", "bytedance"], + versionDetector: nil)) + } +} + +struct DoubaoAPIFetchStrategy: ProviderFetchStrategy { + let id: String = "doubao.api" + let kind: ProviderFetchKind = .apiToken + + func isAvailable(_ context: ProviderFetchContext) async -> Bool { + Self.resolveToken(environment: context.env) != nil + } + + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + guard let apiKey = Self.resolveToken(environment: context.env) else { + throw DoubaoUsageError.missingCredentials + } + let usage = try await DoubaoUsageFetcher.fetchUsage(apiKey: apiKey) + let accumulated = await LocalUsageTracker.shared.record( + provider: .doubao, + remaining: usage.remainingRequests, + limit: usage.limitRequests) + return self.makeResult( + usage: usage.toUsageSnapshot(accumulated: accumulated), + sourceLabel: "api") + } + + func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool { + false + } + + private static func resolveToken(environment: [String: String]) -> String? { + ProviderTokenResolver.doubaoToken(environment: environment) + } +} diff --git a/Sources/CodexBarCore/Providers/Doubao/DoubaoSettingsReader.swift b/Sources/CodexBarCore/Providers/Doubao/DoubaoSettingsReader.swift new file mode 100644 index 000000000..43568d4a5 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Doubao/DoubaoSettingsReader.swift @@ -0,0 +1,37 @@ +import Foundation + +public struct DoubaoSettingsReader: Sendable { + public static let apiKeyEnvironmentKeys = [ + "ARK_API_KEY", + "VOLCENGINE_API_KEY", + "DOUBAO_API_KEY", + ] + + public static func apiKey( + environment: [String: String] = ProcessInfo.processInfo.environment) -> String? + { + for key in self.apiKeyEnvironmentKeys { + guard let raw = environment[key]?.trimmingCharacters(in: .whitespacesAndNewlines), + !raw.isEmpty + else { + continue + } + let cleaned = Self.cleaned(raw) + if !cleaned.isEmpty { + return cleaned + } + } + return nil + } + + private static func cleaned(_ raw: String) -> String { + var value = raw + if (value.hasPrefix("\"") && value.hasSuffix("\"")) || + (value.hasPrefix("'") && value.hasSuffix("'")) + { + value.removeFirst() + value.removeLast() + } + return value.trimmingCharacters(in: .whitespacesAndNewlines) + } +} diff --git a/Sources/CodexBarCore/Providers/Doubao/DoubaoUsageFetcher.swift b/Sources/CodexBarCore/Providers/Doubao/DoubaoUsageFetcher.swift new file mode 100644 index 000000000..0979170dd --- /dev/null +++ b/Sources/CodexBarCore/Providers/Doubao/DoubaoUsageFetcher.swift @@ -0,0 +1,299 @@ +import Foundation +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif + +public struct DoubaoUsageSnapshot: Sendable { + public let remainingRequests: Int + public let limitRequests: Int + public let resetTime: Date? + public let updatedAt: Date + public let apiKeyValid: Bool + public let totalTokens: Int? + public init( + remainingRequests: Int, + limitRequests: Int, + resetTime: Date?, + updatedAt: Date, + apiKeyValid: Bool = false, + totalTokens: Int? = nil) + { + self.remainingRequests = remainingRequests + self.limitRequests = limitRequests + self.resetTime = resetTime + self.updatedAt = updatedAt + self.apiKeyValid = apiKeyValid + self.totalTokens = totalTokens + } + + public func toUsageSnapshot( + accumulated: LocalUsageTracker.AccumulatedUsage? = nil) -> UsageSnapshot + { + let usedPercent: Double + let resetDescription: String + + if self.limitRequests > 0 { + let used = max(0, self.limitRequests - self.remainingRequests) + usedPercent = min(100, max(0, Double(used) / Double(self.limitRequests) * 100)) + resetDescription = "\(used)/\(self.limitRequests) requests" + } else if self.apiKeyValid { + usedPercent = 0 + resetDescription = "Active — check dashboard for details" + } else { + usedPercent = 0 + resetDescription = "No usage data" + } + + let primary = RateWindow( + usedPercent: usedPercent, + windowMinutes: nil, + resetsAt: self.resetTime, + resetDescription: resetDescription) + + var secondary: RateWindow? + if let acc = accumulated, acc.weeklyLimit > 0 { + let weekPercent = min(100, max(0, Double(acc.weeklyRequests) / Double(acc.weeklyLimit) * 100)) + secondary = RateWindow( + usedPercent: weekPercent, + windowMinutes: 7 * 24 * 60, + resetsAt: nil, + resetDescription: "\(acc.weeklyRequests)/\(acc.weeklyLimit) weekly") + } + + let identity = ProviderIdentitySnapshot( + providerID: .doubao, + accountEmail: nil, + accountOrganization: nil, + loginMethod: nil) + + return UsageSnapshot( + primary: primary, + secondary: secondary, + tertiary: nil, + providerCost: nil, + updatedAt: self.updatedAt, + identity: identity) + } +} + +public enum DoubaoUsageError: LocalizedError, Sendable { + case missingCredentials + case networkError(String) + case apiError(Int, String) + case parseFailed(String) + + public var errorDescription: String? { + switch self { + case .missingCredentials: + "Missing Doubao API key (ARK_API_KEY)." + case let .networkError(message): + "Doubao network error: \(message)" + case let .apiError(code, message): + "Doubao API error (\(code)): \(message)" + case let .parseFailed(message): + "Failed to parse Doubao response: \(message)" + } + } +} + +public struct DoubaoUsageFetcher: Sendable { + private static let log = CodexBarLog.logger(LogCategories.doubaoUsage) + private static let apiURL = URL(string: "https://ark.cn-beijing.volces.com/api/coding/v3/chat/completions")! + + /// Models to probe, ordered by likelihood. We try multiple models because + /// different key types may not have access to every model. + private static let probeModels = [ + "doubao-seed-2.0-code", + "doubao-1.5-pro-32k", + "doubao-lite-32k", + ] + + public static func fetchUsage(apiKey: String) async throws -> DoubaoUsageSnapshot { + guard !apiKey.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty else { + throw DoubaoUsageError.missingCredentials + } + + var lastError: Error? + for model in self.probeModels { + do { + return try await self.probe(apiKey: apiKey, model: model) + } catch let error as DoubaoUsageError { + if case let .apiError(code, _) = error, code == 404 || code == 403 { + Self.log.debug("Doubao probe model \(model) unavailable (\(code)), trying next") + lastError = error + continue + } + throw error + } + } + throw lastError ?? DoubaoUsageError.apiError(0, "All probe models failed") + } + + private static func probe(apiKey: String, model: String) async throws -> DoubaoUsageSnapshot { + var request = URLRequest(url: self.apiURL) + request.httpMethod = "POST" + request.timeoutInterval = 30 + request.setValue("application/json", forHTTPHeaderField: "Content-Type") + request.setValue("application/json", forHTTPHeaderField: "Accept") + request.setValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization") + + let body: [String: Any] = [ + "model": model, + "max_tokens": 1, + "messages": [ + ["role": "user", "content": "hi"], + ] as [[String: Any]], + ] + + request.httpBody = try JSONSerialization.data(withJSONObject: body) + + let (data, response) = try await URLSession.shared.data(for: request) + + guard let httpResponse = response as? HTTPURLResponse else { + throw DoubaoUsageError.networkError("Invalid response") + } + + // Accept both 200 (success) and 429 (rate limited) – both carry rate limit headers. + guard httpResponse.statusCode == 200 || httpResponse.statusCode == 429 else { + let summary = Self.apiErrorSummary(statusCode: httpResponse.statusCode, data: data) + Self.log.error("Doubao API returned \(httpResponse.statusCode): \(summary)") + throw DoubaoUsageError.apiError(httpResponse.statusCode, summary) + } + + let headers = httpResponse.allHeaderFields + let remaining = Self.intHeader(headers, "x-ratelimit-remaining-requests") + let limit = Self.intHeader(headers, "x-ratelimit-limit-requests") + let resetString = Self.stringHeader(headers, "x-ratelimit-reset-requests") + + let resetTime: Date? = resetString.flatMap(Self.parseResetTime) + + var totalTokens: Int? + if remaining == nil, limit == nil, + let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any], + let usage = json["usage"] as? [String: Any] + { + totalTokens = usage["total_tokens"] as? Int + } + + // 429 means the key is valid but rate-limited; treat it as valid so the UI + // shows "Active" instead of "No usage data" when headers are absent. + let keyValid = httpResponse.statusCode == 200 || httpResponse.statusCode == 429 + + let snapshot = DoubaoUsageSnapshot( + remainingRequests: remaining ?? 0, + limitRequests: limit ?? 0, + resetTime: resetTime, + updatedAt: Date(), + apiKeyValid: keyValid, + totalTokens: totalTokens) + + Self.log.debug( + "Doubao usage parsed remaining=\(snapshot.remainingRequests) limit=\(snapshot.limitRequests) valid=\(snapshot.apiKeyValid)") // swiftlint:disable:this line_length + + return snapshot + } + + private static func stringHeader(_ headers: [AnyHashable: Any], _ name: String) -> String? { + if let value = headers[name] as? String { return value } + for (key, val) in headers { + if let keyStr = key as? String, + keyStr.caseInsensitiveCompare(name) == .orderedSame, + let valStr = val as? String + { + return valStr + } + } + return nil + } + + private static func intHeader(_ headers: [AnyHashable: Any], _ name: String) -> Int? { + if let value = headers[name] as? String, let int = Int(value) { + return int + } + if let value = headers[name.lowercased()] as? String, let int = Int(value) { + return int + } + for (key, val) in headers { + if let keyStr = key as? String, + keyStr.lowercased() == name.lowercased(), + let valStr = val as? String, + let int = Int(valStr) + { + return int + } + } + return nil + } + + private static func parseResetTime(_ value: String) -> Date? { + let trimmed = value.trimmingCharacters(in: .whitespacesAndNewlines) + if trimmed.isEmpty { return nil } + + let isoFormatter = ISO8601DateFormatter() + isoFormatter.formatOptions = [.withInternetDateTime, .withFractionalSeconds] + if let date = isoFormatter.date(from: trimmed) { return date } + let isoFallback = ISO8601DateFormatter() + isoFallback.formatOptions = [.withInternetDateTime] + if let date = isoFallback.date(from: trimmed) { return date } + + var seconds: TimeInterval = 0 + let pattern = /(\d+)([dhms])/ + for match in trimmed.matches(of: pattern) { + guard let num = Double(match.1) else { continue } + switch match.2 { + case "d": seconds += num * 86400 + case "h": seconds += num * 3600 + case "m": seconds += num * 60 + case "s": seconds += num + default: break + } + } + if seconds > 0 { + return Date().addingTimeInterval(seconds) + } + + if let secs = TimeInterval(trimmed) { + return Date().addingTimeInterval(secs) + } + + return nil + } + + private static func apiErrorSummary(statusCode: Int, data: Data) -> String { + guard let root = try? JSONSerialization.jsonObject(with: data), + let json = root as? [String: Any] + else { + if let text = String(data: data, encoding: .utf8)? + .trimmingCharacters(in: .whitespacesAndNewlines), + !text.isEmpty + { + return self.compactText(text) + } + return "Unexpected response body (\(data.count) bytes)." + } + + if let error = json["error"] as? [String: Any], + let message = error["message"] as? String + { + let trimmed = message.trimmingCharacters(in: .whitespacesAndNewlines) + if !trimmed.isEmpty { return Self.compactText(trimmed) } + } + + if let message = json["message"] as? String { + let trimmed = message.trimmingCharacters(in: .whitespacesAndNewlines) + if !trimmed.isEmpty { return Self.compactText(trimmed) } + } + + return "HTTP \(statusCode) (\(data.count) bytes)." + } + + private static func compactText(_ text: String, maxLength: Int = 200) -> String { + let collapsed = text + .components(separatedBy: .newlines) + .joined(separator: " ") + .trimmingCharacters(in: .whitespacesAndNewlines) + if collapsed.count <= maxLength { return collapsed } + let limitIndex = collapsed.index(collapsed.startIndex, offsetBy: maxLength) + return "\(collapsed[.. String? { + // Resolve symlinks recursively, collecting all intermediate paths + // (e.g. /usr/local/bin/gemini → /opt/homebrew/bin/gemini → ../Cellar/.../bin/gemini → ...) let fm = FileManager.default - var realPath = geminiPath - if let resolved = try? fm.destinationOfSymbolicLink(atPath: geminiPath) { + var candidates: [String] = [binaryPath] + var current = binaryPath + var visited: Set = [] + while true { + let canonical = (current as NSString).standardizingPath + if visited.contains(canonical) { break } + visited.insert(canonical) + guard let resolved = try? fm.destinationOfSymbolicLink(atPath: current) else { break } if resolved.hasPrefix("/") { - realPath = resolved + current = resolved } else { - realPath = (geminiPath as NSString).deletingLastPathComponent + "/" + resolved + current = ((current as NSString).deletingLastPathComponent as NSString) + .appendingPathComponent(resolved) } + current = (current as NSString).standardizingPath + candidates.append(current) } // Navigate from bin/gemini to the oauth2.js file - // Homebrew path: .../libexec/lib/node_modules/@google/gemini-cli/node_modules/@google/gemini-cli-core/dist/src/code_assist/oauth2.js - // Bun/npm path: .../node_modules/@google/gemini-cli-core/dist/src/code_assist/oauth2.js (sibling package) - let binDir = (realPath as NSString).deletingLastPathComponent - let baseDir = (binDir as NSString).deletingLastPathComponent - + // Try from each resolved path in the symlink chain (deepest first) let oauthSubpath = "node_modules/@google/gemini-cli/node_modules/@google/gemini-cli-core/dist/src/code_assist/oauth2.js" let nixShareSubpath = "share/gemini-cli/node_modules/@google/gemini-cli-core/dist/src/code_assist/oauth2.js" let oauthFile = "dist/src/code_assist/oauth2.js" - let possiblePaths = [ - // Homebrew nested structure - "\(baseDir)/libexec/lib/\(oauthSubpath)", - "\(baseDir)/lib/\(oauthSubpath)", - // Nix package layout - "\(baseDir)/\(nixShareSubpath)", - // Bun/npm sibling structure: gemini-cli-core is a sibling to gemini-cli - "\(baseDir)/../gemini-cli-core/\(oauthFile)", - // npm nested inside gemini-cli - "\(baseDir)/node_modules/@google/gemini-cli-core/\(oauthFile)", - ] - - for path in possiblePaths { - if let content = try? String(contentsOfFile: path, encoding: .utf8) { - return self.parseOAuthCredentials(from: content) + + for candidate in candidates.reversed() { + let binDir = (candidate as NSString).deletingLastPathComponent + let baseDir = (binDir as NSString).deletingLastPathComponent + + let possiblePaths = [ + // Homebrew nested structure + "\(baseDir)/libexec/lib/\(oauthSubpath)", + "\(baseDir)/lib/\(oauthSubpath)", + // Nix package layout + "\(baseDir)/\(nixShareSubpath)", + // Bun/npm sibling structure + "\(baseDir)/../gemini-cli-core/\(oauthFile)", + // npm nested inside gemini-cli + "\(baseDir)/node_modules/@google/gemini-cli-core/\(oauthFile)", + // Direct node_modules lookup from candidate + "\(binDir)/node_modules/@google/gemini-cli-core/\(oauthFile)", + ] + + for path in possiblePaths { + if let content = try? String(contentsOfFile: path, encoding: .utf8) { + return content + } } } diff --git a/Sources/CodexBarCore/Providers/Kimi/KimiLocalStorageImporter.swift b/Sources/CodexBarCore/Providers/Kimi/KimiLocalStorageImporter.swift new file mode 100644 index 000000000..f76b3c9b8 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Kimi/KimiLocalStorageImporter.swift @@ -0,0 +1,177 @@ +#if os(macOS) +import Foundation +import SweetCookieKit + +/// Imports Kimi auth tokens from Chrome's localStorage. +/// +/// Kimi stores `access_token` and `refresh_token` as JWTs in localStorage +/// on `https://www.kimi.com`, not in cookies. +public enum KimiLocalStorageImporter { + private static let log = CodexBarLog.logger(LogCategories.kimiCookie) + + private static let origins = [ + "https://www.kimi.com", + "https://kimi.com", + ] + + /// The localStorage key that holds the JWT access token. + private static let accessTokenKey = "access_token" + private static let refreshTokenKey = "refresh_token" + + public struct SessionInfo: Sendable { + public let accessToken: String + public let refreshToken: String? + public let sourceLabel: String + } + + public static func importSession( + browserDetection: BrowserDetection = BrowserDetection(), + logger: ((String) -> Void)? = nil) -> SessionInfo? + { + let log: (String) -> Void = { msg in + logger?("[kimi-storage] \(msg)") + self.log.debug(msg) + } + + let candidates = self.chromeLocalStorageCandidates(browserDetection: browserDetection) + log("Found \(candidates.count) Chrome profile candidate(s)") + + for candidate in candidates { + if let session = self.readKimiSession(from: candidate.levelDBURL, label: candidate.label, logger: log) { + return session + } + } + + log("No Kimi access_token found in any browser profile") + return nil + } + + public static func hasSession( + browserDetection: BrowserDetection = BrowserDetection()) -> Bool + { + self.importSession(browserDetection: browserDetection) != nil + } + + // MARK: - LevelDB Reading + + private static func readKimiSession( + from levelDBURL: URL, + label: String, + logger: ((String) -> Void)?) -> SessionInfo? + { + for origin in self.origins { + let entries = ChromiumLocalStorageReader.readEntries( + for: origin, + in: levelDBURL, + logger: logger) + + var accessToken: String? + var refreshToken: String? + + for entry in entries { + let value = entry.value.trimmingCharacters(in: .controlCharacters) + guard !value.isEmpty else { continue } + + if entry.key == self.accessTokenKey, value.hasPrefix("eyJ") { + accessToken = value + } else if entry.key == self.refreshTokenKey, value.hasPrefix("eyJ") { + refreshToken = value + } + } + + if let accessToken { + logger?("Found Kimi access_token in \(label) (origin: \(origin))") + return SessionInfo( + accessToken: accessToken, + refreshToken: refreshToken, + sourceLabel: label) + } + } + + // Fallback: text scan + let textEntries = ChromiumLocalStorageReader.readTextEntries( + in: levelDBURL, + logger: logger) + + var accessToken: String? + var refreshToken: String? + + for entry in textEntries { + let value = entry.value.trimmingCharacters(in: .controlCharacters) + guard !value.isEmpty, value.hasPrefix("eyJ") else { continue } + + if entry.key.hasSuffix(self.accessTokenKey) { + accessToken = value + } else if entry.key.hasSuffix(self.refreshTokenKey) { + refreshToken = value + } + } + + if let accessToken { + logger?("Found Kimi access_token in \(label) (text scan)") + return SessionInfo( + accessToken: accessToken, + refreshToken: refreshToken, + sourceLabel: label) + } + + return nil + } + + // MARK: - Chrome Profile Discovery + + private struct LocalStorageCandidate { + let label: String + let levelDBURL: URL + } + + private static func chromeLocalStorageCandidates( + browserDetection: BrowserDetection) -> [LocalStorageCandidate] + { + let browsers: [Browser] = [ + .chrome, + .chromeBeta, + .chromeCanary, + .arc, + .arcBeta, + .arcCanary, + .chromium, + ] + + let installedBrowsers = browsers.browsersWithProfileData(using: browserDetection) + let roots = ChromiumProfileLocator + .roots(for: installedBrowsers, homeDirectories: BrowserCookieClient.defaultHomeDirectories()) + .map { (url: $0.url, labelPrefix: $0.labelPrefix) } + + var candidates: [LocalStorageCandidate] = [] + for root in roots { + candidates.append(contentsOf: self.profileCandidates(root: root.url, labelPrefix: root.labelPrefix)) + } + return candidates + } + + private static func profileCandidates(root: URL, labelPrefix: String) -> [LocalStorageCandidate] { + guard let entries = try? FileManager.default.contentsOfDirectory( + at: root, + includingPropertiesForKeys: [.isDirectoryKey], + options: [.skipsHiddenFiles]) + else { return [] } + + let profileDirs = entries.filter { url in + guard let isDir = (try? url.resourceValues(forKeys: [.isDirectoryKey]).isDirectory), isDir else { + return false + } + let name = url.lastPathComponent + return name == "Default" || name.hasPrefix("Profile ") || name.hasPrefix("user-") + } + .sorted { $0.lastPathComponent < $1.lastPathComponent } + + return profileDirs.compactMap { dir in + let levelDBURL = dir.appendingPathComponent("Local Storage").appendingPathComponent("leveldb") + guard FileManager.default.fileExists(atPath: levelDBURL.path) else { return nil } + let label = "\(labelPrefix) \(dir.lastPathComponent)" + return LocalStorageCandidate(label: label, levelDBURL: levelDBURL) + } + } +} +#endif diff --git a/Sources/CodexBarCore/Providers/Kimi/KimiProviderDescriptor.swift b/Sources/CodexBarCore/Providers/Kimi/KimiProviderDescriptor.swift index 711c20bc8..46c7e4417 100644 --- a/Sources/CodexBarCore/Providers/Kimi/KimiProviderDescriptor.swift +++ b/Sources/CodexBarCore/Providers/Kimi/KimiProviderDescriptor.swift @@ -57,7 +57,8 @@ struct KimiWebFetchStrategy: ProviderFetchStrategy { #if os(macOS) if context.settings?.kimi?.cookieSource != .off { - return KimiCookieImporter.hasSession() + if KimiCookieImporter.hasSession() { return true } + if KimiLocalStorageImporter.hasSession() { return true } } #endif @@ -90,6 +91,7 @@ struct KimiWebFetchStrategy: ProviderFetchStrategy { // Try browser cookie import when auto mode is enabled #if os(macOS) if context.settings?.kimi?.cookieSource != .off { + // Try cookies first (legacy kimi-auth cookie) do { let session = try KimiCookieImporter.importSession() if let token = session.authToken { @@ -98,6 +100,11 @@ struct KimiWebFetchStrategy: ProviderFetchStrategy { } catch { // No browser cookies found } + + // Try localStorage (current: access_token JWT) + if let lsSession = KimiLocalStorageImporter.importSession() { + return lsSession.accessToken + } } #endif diff --git a/Sources/CodexBarCore/Providers/Kimi/KimiUsageFetcher.swift b/Sources/CodexBarCore/Providers/Kimi/KimiUsageFetcher.swift index c6d6c2a2e..7a06a6385 100644 --- a/Sources/CodexBarCore/Providers/Kimi/KimiUsageFetcher.swift +++ b/Sources/CodexBarCore/Providers/Kimi/KimiUsageFetcher.swift @@ -75,7 +75,8 @@ public struct KimiUsageFetcher: Sendable { return KimiUsageSnapshot( weekly: codingUsage.detail, rateLimit: codingUsage.limits?.first?.detail, - updatedAt: now) + updatedAt: now, + authToken: authToken) } private static func decodeSessionInfo(from jwt: String) -> SessionInfo? { diff --git a/Sources/CodexBarCore/Providers/Kimi/KimiUsageSnapshot.swift b/Sources/CodexBarCore/Providers/Kimi/KimiUsageSnapshot.swift index c39e1602a..bcb94b477 100644 --- a/Sources/CodexBarCore/Providers/Kimi/KimiUsageSnapshot.swift +++ b/Sources/CodexBarCore/Providers/Kimi/KimiUsageSnapshot.swift @@ -4,11 +4,13 @@ public struct KimiUsageSnapshot: Sendable { public let weekly: KimiUsageDetail public let rateLimit: KimiUsageDetail? public let updatedAt: Date + public let authToken: String? - public init(weekly: KimiUsageDetail, rateLimit: KimiUsageDetail?, updatedAt: Date) { + public init(weekly: KimiUsageDetail, rateLimit: KimiUsageDetail?, updatedAt: Date, authToken: String? = nil) { self.weekly = weekly self.rateLimit = rateLimit self.updatedAt = updatedAt + self.authToken = authToken } private static func parseDate(_ dateString: String) -> Date? { @@ -66,11 +68,21 @@ extension KimiUsageSnapshot { resetDescription: "Rate: \(rateUsed)/\(rateLimitValue) per 5 hours") } + let maskedKey: String? = { + guard let key = self.authToken, !key.isEmpty else { return nil } + if key.count <= 8 { return "****" } + return "\(key.prefix(6))...\(key.suffix(4))" + }() + let plan: String? = { + guard let key = self.authToken else { return nil } + if key.hasPrefix("sk-kimi-") { return "API" } + return "Web" + }() let identity = ProviderIdentitySnapshot( providerID: .kimi, - accountEmail: nil, + accountEmail: maskedKey, accountOrganization: nil, - loginMethod: nil) + loginMethod: plan) return UsageSnapshot( primary: weeklyWindow, diff --git a/Sources/CodexBarCore/Providers/LocalUsageTracker.swift b/Sources/CodexBarCore/Providers/LocalUsageTracker.swift new file mode 100644 index 000000000..cc7820dcb --- /dev/null +++ b/Sources/CodexBarCore/Providers/LocalUsageTracker.swift @@ -0,0 +1,162 @@ +import Foundation + +/// Tracks API request consumption locally by recording rate-limit snapshots over time. +/// Computes weekly (7-day) and monthly (30-day) accumulated usage from deltas between samples. +/// Used by providers (Doubao, Qwen) that only expose rate-limit headers without dedicated usage APIs. +public actor LocalUsageTracker { + public static let shared = LocalUsageTracker() + + private static let sampleInterval: TimeInterval = 60 + private static let retentionDays: TimeInterval = 31 * 24 * 60 * 60 + private static let weekSeconds: TimeInterval = 7 * 24 * 60 * 60 + private static let monthSeconds: TimeInterval = 30 * 24 * 60 * 60 + + private var records: [String: [Sample]] = [:] + private var loaded = false + + private struct Sample: Codable, Sendable { + let timestamp: Date + let remaining: Int + let limit: Int + } + + public struct AccumulatedUsage: Sendable { + public let weeklyRequests: Int + public let monthlyRequests: Int + public let weeklyLimit: Int + public let monthlyLimit: Int + } + + private init() {} + + /// Record a rate-limit sample and return accumulated weekly/monthly usage. + public func record( + provider: UsageProvider, + remaining: Int, + limit: Int, + now: Date = Date()) -> AccumulatedUsage + { + self.ensureLoaded() + let key = provider.rawValue + + var samples = self.records[key] ?? [] + + // Throttle: skip if last sample is too recent and values unchanged + if let last = samples.last { + let elapsed = now.timeIntervalSince(last.timestamp) + if elapsed < Self.sampleInterval, last.remaining == remaining, last.limit == limit { + return self.computeUsage(samples: samples, limit: limit, now: now) + } + } + + samples.append(Sample(timestamp: now, remaining: remaining, limit: limit)) + + // Prune old samples + let cutoff = now.addingTimeInterval(-Self.retentionDays) + samples.removeAll { $0.timestamp < cutoff } + + self.records[key] = samples + self.persist() + + return self.computeUsage(samples: samples, limit: limit, now: now) + } + + /// Get accumulated usage without recording a new sample. + public func usage(for provider: UsageProvider) -> AccumulatedUsage? { + self.ensureLoaded() + guard let samples = self.records[provider.rawValue], !samples.isEmpty else { return nil } + let limit = samples.last?.limit ?? 0 + return self.computeUsage(samples: samples, limit: limit, now: Date()) + } + + // MARK: - Computation + + private func computeUsage(samples: [Sample], limit: Int, now: Date) -> AccumulatedUsage { + let weekCutoff = now.addingTimeInterval(-Self.weekSeconds) + let monthCutoff = now.addingTimeInterval(-Self.monthSeconds) + + let weeklyRequests = Self.sumConsumption( + samples: samples.filter { $0.timestamp >= weekCutoff }) + let monthlyRequests = Self.sumConsumption( + samples: samples.filter { $0.timestamp >= monthCutoff }) + + // Estimate limits: daily limit * 7 or * 30 + // We use the current window limit as a proxy for daily capacity + let weeklyLimit = limit > 0 ? limit * 7 : 0 + let monthlyLimit = limit > 0 ? limit * 30 : 0 + + return AccumulatedUsage( + weeklyRequests: weeklyRequests, + monthlyRequests: monthlyRequests, + weeklyLimit: weeklyLimit, + monthlyLimit: monthlyLimit) + } + + /// Sum consumption from consecutive samples by tracking remaining-count drops. + /// When remaining increases (reset occurred), we don't count that as negative consumption. + private static func sumConsumption(samples: [Sample]) -> Int { + guard samples.count >= 2 else { return 0 } + + var total = 0 + for i in 1.. URL { + let root = FileManager.default.urls(for: .applicationSupportDirectory, in: .userDomainMask).first + ?? FileManager.default.homeDirectoryForCurrentUser + return root + .appendingPathComponent("CodexBar", isDirectory: true) + .appendingPathComponent("local-usage-tracker.json", isDirectory: false) + } + + private static func readFromDisk() -> [String: [Sample]] { + guard let data = try? Data(contentsOf: fileURL()), + let decoded = try? JSONDecoder.iso8601Decoder.decode([String: [Sample]].self, from: data) + else { + return [:] + } + return decoded + } + + private func persist() { + let encoder = JSONEncoder() + encoder.dateEncodingStrategy = .iso8601 + encoder.outputFormatting = [.sortedKeys] + + guard let data = try? encoder.encode(self.records) else { return } + let url = Self.fileURL() + let directory = url.deletingLastPathComponent() + do { + try FileManager.default.createDirectory(at: directory, withIntermediateDirectories: true) + try data.write(to: url, options: [.atomic]) + } catch { + // Best-effort; ignore write failures. + } + } +} + +extension JSONDecoder { + fileprivate static let iso8601Decoder: JSONDecoder = { + let decoder = JSONDecoder() + decoder.dateDecodingStrategy = .iso8601 + return decoder + }() +} diff --git a/Sources/CodexBarCore/Providers/MiMo/MiMoCookieImporter.swift b/Sources/CodexBarCore/Providers/MiMo/MiMoCookieImporter.swift new file mode 100644 index 000000000..97ccea8f0 --- /dev/null +++ b/Sources/CodexBarCore/Providers/MiMo/MiMoCookieImporter.swift @@ -0,0 +1,239 @@ +import Foundation + +enum MiMoCookieHeader { + static let requiredCookieNames: Set = [ + "api-platform_serviceToken", + "userId", + ] + static let knownCookieNames: Set = requiredCookieNames.union([ + "api-platform_ph", + "api-platform_slh", + ]) + + static func normalizedHeader(from raw: String?) -> String? { + guard let normalized = CookieHeaderNormalizer.normalize(raw) else { return nil } + let pairs = CookieHeaderNormalizer.pairs(from: normalized) + guard !pairs.isEmpty else { return nil } + + var byName: [String: String] = [:] + for pair in pairs { + let name = pair.name.trimmingCharacters(in: .whitespacesAndNewlines) + let value = pair.value.trimmingCharacters(in: .whitespacesAndNewlines) + guard self.knownCookieNames.contains(name), !value.isEmpty else { continue } + byName[name] = value + } + + guard self.requiredCookieNames.isSubset(of: Set(byName.keys)) else { return nil } + return byName.keys.sorted().compactMap { name in + guard let value = byName[name] else { return nil } + return "\(name)=\(value)" + }.joined(separator: "; ") + } + + #if os(macOS) + static func header(from cookies: [HTTPCookie]) -> String? { + let requestURL = URL(string: "https://platform.xiaomimimo.com/api/v1/balance")! + var byName: [String: HTTPCookie] = [:] + for cookie in cookies { + guard self.knownCookieNames.contains(cookie.name) else { continue } + guard !cookie.value.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty else { continue } + if let expiry = cookie.expiresDate, expiry < Date() { continue } + guard Self.matchesRequestURL(cookie: cookie, url: requestURL) else { continue } + + if let existing = byName[cookie.name] { + if Self.cookieSortKey(for: cookie) >= Self.cookieSortKey(for: existing) { + byName[cookie.name] = cookie + } + } else { + byName[cookie.name] = cookie + } + } + + guard self.requiredCookieNames.isSubset(of: Set(byName.keys)) else { return nil } + return byName.keys.sorted().compactMap { name in + guard let cookie = byName[name] else { return nil } + return "\(cookie.name)=\(cookie.value)" + }.joined(separator: "; ") + } + + private static func matchesRequestURL(cookie: HTTPCookie, url: URL) -> Bool { + guard let host = url.host else { return false } + let normalizedDomain = cookie.domain.lowercased().trimmingCharacters(in: CharacterSet(charactersIn: ".")) + guard !normalizedDomain.isEmpty else { return false } + guard host == normalizedDomain || host.hasSuffix(".\(normalizedDomain)") else { return false } + + let cookiePath = cookie.path.isEmpty ? "/" : cookie.path + let requestPath = url.path.isEmpty ? "/" : url.path + if requestPath == cookiePath { + return true + } + guard requestPath.hasPrefix(cookiePath) else { return false } + guard cookiePath != "/" else { return true } + if cookiePath.hasSuffix("/") { + return true + } + guard + let boundaryIndex = requestPath.index( + cookiePath.startIndex, + offsetBy: cookiePath.count, + limitedBy: requestPath.endIndex), + boundaryIndex < requestPath.endIndex + else { + return true + } + return requestPath[boundaryIndex] == "/" + } + + private static func cookieSortKey(for cookie: HTTPCookie) -> (Int, Int, Date) { + let pathLength = cookie.path.count + let normalizedDomain = cookie.domain.lowercased().trimmingCharacters(in: CharacterSet(charactersIn: ".")) + let domainLength = normalizedDomain.count + let expiry = cookie.expiresDate ?? .distantPast + return (pathLength, domainLength, expiry) + } + #endif +} + +#if os(macOS) +import SweetCookieKit + +private let miMoCookieImportOrder: BrowserCookieImportOrder = + ProviderDefaults.metadata[.mimo]?.browserCookieOrder ?? Browser.defaultImportOrder + +public enum MiMoCookieImporter { + private static let cookieClient = BrowserCookieClient() + private static let cookieDomains = [ + "platform.xiaomimimo.com", + "xiaomimimo.com", + ] + + public struct SessionInfo: Sendable { + public let cookieHeader: String + public let sourceLabel: String + + public init(cookieHeader: String, sourceLabel: String) { + self.cookieHeader = cookieHeader + self.sourceLabel = sourceLabel + } + } + + nonisolated(unsafe) static var importSessionsOverrideForTesting: + ((BrowserDetection, ((String) -> Void)?) throws -> [SessionInfo])? + + public static func importSessions( + browserDetection: BrowserDetection, + logger: ((String) -> Void)? = nil) throws -> [SessionInfo] + { + if let override = self.importSessionsOverrideForTesting { + return try override(browserDetection, logger) + } + + let log: (String) -> Void = { msg in logger?("[mimo-cookie] \(msg)") } + var sessions: [SessionInfo] = [] + let installed = miMoCookieImportOrder.cookieImportCandidates(using: browserDetection) + let labels = installed.map(\.displayName).joined(separator: ", ") + log("Cookie import candidates: \(labels)") + + for browserSource in installed { + do { + let query = BrowserCookieQuery(domains: self.cookieDomains) + let sources = try Self.cookieClient.records( + matching: query, + in: browserSource, + logger: log) + sessions.append(contentsOf: self.sessionInfos(from: sources, origin: query.origin)) + } catch { + BrowserCookieAccessGate.recordIfNeeded(error) + log("\(browserSource.displayName) cookie import failed: \(error.localizedDescription)") + } + } + + return sessions + } + + public static func hasSession( + browserDetection: BrowserDetection, + logger: ((String) -> Void)? = nil) -> Bool + { + (try? self.importSessions(browserDetection: browserDetection, logger: logger).isEmpty == false) ?? false + } + + static func sessionInfos( + from sources: [BrowserCookieStoreRecords], + origin: BrowserCookieOriginStrategy = .domainBased) -> [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) + } + + var sessions: [SessionInfo] = [] + for group in sortedGroups where !group.isEmpty { + let label = self.mergedLabel(for: group) + let mergedRecords = self.mergeRecords(group) + guard !mergedRecords.isEmpty else { continue } + let cookies = BrowserCookieClient.makeHTTPCookies(mergedRecords, origin: origin) + guard let cookieHeader = MiMoCookieHeader.header(from: cookies) else { + continue + } + sessions.append(SessionInfo(cookieHeader: cookieHeader, sourceLabel: label)) + } + return sessions + } + + 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 + } + } +} +#endif diff --git a/Sources/CodexBarCore/Providers/MiMo/MiMoProviderDescriptor.swift b/Sources/CodexBarCore/Providers/MiMo/MiMoProviderDescriptor.swift new file mode 100644 index 000000000..42e3eb6cc --- /dev/null +++ b/Sources/CodexBarCore/Providers/MiMo/MiMoProviderDescriptor.swift @@ -0,0 +1,178 @@ +import CodexBarMacroSupport +import Foundation + +#if os(macOS) +import SweetCookieKit +#endif + +@ProviderDescriptorRegistration +@ProviderDescriptorDefinition +public enum MiMoProviderDescriptor { + static func makeDescriptor() -> ProviderDescriptor { + #if os(macOS) + let browserOrder: BrowserCookieImportOrder = [ + .chrome, + .chromeBeta, + .chromeCanary, + ] + #else + let browserOrder: BrowserCookieImportOrder? = nil + #endif + + return ProviderDescriptor( + id: .mimo, + metadata: ProviderMetadata( + id: .mimo, + displayName: "Xiaomi MiMo", + sessionLabel: "Credits", + weeklyLabel: "Window", + opusLabel: nil, + supportsOpus: false, + supportsCredits: true, + creditsHint: "Token plan credits usage.", + toggleTitle: "Show Xiaomi MiMo token plan & balance", + cliName: "mimo", + defaultEnabled: false, + isPrimaryProvider: false, + usesAccountFallback: false, + browserCookieOrder: browserOrder, + dashboardURL: "https://platform.xiaomimimo.com/#/console/balance", + statusPageURL: nil), + branding: ProviderBranding( + iconStyle: .mimo, + iconResourceName: "ProviderIcon-mimo", + color: ProviderColor(red: 1.0, green: 105 / 255, blue: 0)), + tokenCost: ProviderTokenCostConfig( + supportsTokenCost: false, + noDataMessage: { "Xiaomi MiMo cost summary is not supported." }), + fetchPlan: ProviderFetchPlan( + sourceModes: [.auto, .web], + pipeline: ProviderFetchPipeline(resolveStrategies: { _ in [MiMoWebFetchStrategy()] })), + cli: ProviderCLIConfig( + name: "mimo", + aliases: ["xiaomi-mimo"], + versionDetector: nil)) + } +} + +struct MiMoWebFetchStrategy: ProviderFetchStrategy { + let id: String = "mimo.web" + let kind: ProviderFetchKind = .web + + func isAvailable(_ context: ProviderFetchContext) async -> Bool { + guard context.settings?.mimo?.cookieSource != .off else { return false } + if context.settings?.mimo?.cookieSource == .manual { + return Self.resolveManualCookieHeader(context: context) != nil + } + if Self.resolveManualCookieHeader(context: context) != nil { + return true + } + + #if os(macOS) + if let cached = CookieHeaderCache.load(provider: .mimo), + MiMoCookieHeader.normalizedHeader(from: cached.cookieHeader) != nil + { + return true + } + return MiMoCookieImporter.hasSession(browserDetection: context.browserDetection) + #else + return false + #endif + } + + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + guard context.settings?.mimo?.cookieSource != .off else { + throw MiMoSettingsError.missingCookie + } + if context.settings?.mimo?.cookieSource == .manual { + guard let manualCookie = Self.resolveManualCookieHeader(context: context) else { + throw MiMoSettingsError.invalidCookie + } + let snapshot = try await MiMoUsageFetcher.fetchUsage( + cookieHeader: manualCookie, + environment: context.env) + return self.makeResult(usage: snapshot.toUsageSnapshot(), sourceLabel: "web") + } + if let manualCookie = Self.resolveManualCookieHeader(context: context) { + let snapshot = try await MiMoUsageFetcher.fetchUsage( + cookieHeader: manualCookie, + environment: context.env) + return self.makeResult(usage: snapshot.toUsageSnapshot(), sourceLabel: "web") + } + + #if os(macOS) + var lastError: Error? + + if let cached = CookieHeaderCache.load(provider: .mimo), + let cachedHeader = MiMoCookieHeader.normalizedHeader(from: cached.cookieHeader) + { + do { + let snapshot = try await MiMoUsageFetcher.fetchUsage( + cookieHeader: cachedHeader, + environment: context.env) + return self.makeResult(usage: snapshot.toUsageSnapshot(), sourceLabel: "web") + } catch { + guard Self.shouldRetryNextSession(for: error) else { + throw error + } + CookieHeaderCache.clear(provider: .mimo) + lastError = error + } + } + + let sessions = try MiMoCookieImporter.importSessions(browserDetection: context.browserDetection) + guard !sessions.isEmpty else { + if let lastError { throw lastError } + throw MiMoSettingsError.missingCookie + } + + for session in sessions { + do { + let snapshot = try await MiMoUsageFetcher.fetchUsage( + cookieHeader: session.cookieHeader, + environment: context.env) + CookieHeaderCache.store( + provider: .mimo, + cookieHeader: session.cookieHeader, + sourceLabel: session.sourceLabel) + return self.makeResult(usage: snapshot.toUsageSnapshot(), sourceLabel: "web") + } catch { + guard Self.shouldRetryNextSession(for: error) else { + throw error + } + lastError = error + continue + } + } + + if let lastError { throw lastError } + throw MiMoSettingsError.missingCookie + #else + throw MiMoSettingsError.missingCookie + #endif + } + + func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool { + false + } + + private static func resolveManualCookieHeader(context: ProviderFetchContext) -> String? { + guard context.settings?.mimo?.cookieSource == .manual else { return nil } + return MiMoCookieHeader.normalizedHeader(from: context.settings?.mimo?.manualCookieHeader) + } + + private static func shouldRetryNextSession(for error: Error) -> Bool { + if error is DecodingError { + return true + } + guard let mimoError = error as? MiMoUsageError else { + return false + } + switch mimoError { + case .invalidCredentials, .loginRequired, .parseFailed: + return true + case .networkError: + return false + } + } +} diff --git a/Sources/CodexBarCore/Providers/MiMo/MiMoUsageFetcher.swift b/Sources/CodexBarCore/Providers/MiMo/MiMoUsageFetcher.swift new file mode 100644 index 000000000..1770c5525 --- /dev/null +++ b/Sources/CodexBarCore/Providers/MiMo/MiMoUsageFetcher.swift @@ -0,0 +1,263 @@ +import Foundation +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif + +public enum MiMoSettingsError: LocalizedError, Sendable { + case missingCookie + case invalidCookie + + public var errorDescription: String? { + switch self { + case .missingCookie: + "No Xiaomi MiMo browser session found. Log in at platform.xiaomimimo.com first." + case .invalidCookie: + "Xiaomi MiMo requires the api-platform_serviceToken and userId cookies." + } + } +} + +public enum MiMoUsageError: LocalizedError, Sendable { + case invalidCredentials + case loginRequired + case parseFailed(String) + case networkError(String) + + public var errorDescription: String? { + switch self { + case .invalidCredentials: + "Xiaomi MiMo browser session expired. Log in again." + case .loginRequired: + "Xiaomi MiMo login required." + case let .parseFailed(message): + "Could not parse Xiaomi MiMo balance: \(message)" + case let .networkError(message): + "Xiaomi MiMo request failed: \(message)" + } + } +} + +public enum MiMoSettingsReader { + public static let apiURLKey = "MIMO_API_URL" + + public static func apiURL(environment: [String: String] = ProcessInfo.processInfo.environment) -> URL { + if let override = environment[self.apiURLKey], + let url = URL(string: override.trimmingCharacters(in: .whitespacesAndNewlines)), + let scheme = url.scheme, !scheme.isEmpty + { + return url + } + return URL(string: "https://platform.xiaomimimo.com/api/v1")! + } +} + +public enum MiMoUsageFetcher { + private static let requestTimeout: TimeInterval = 15 + + public static func fetchUsage( + cookieHeader: String, + environment: [String: String] = ProcessInfo.processInfo.environment, + now: Date = Date()) async throws -> MiMoUsageSnapshot + { + guard let normalizedCookie = MiMoCookieHeader.normalizedHeader(from: cookieHeader) else { + throw MiMoSettingsError.invalidCookie + } + + let balanceURL = MiMoSettingsReader.apiURL(environment: environment).appendingPathComponent("balance") + let tokenDetailURL = MiMoSettingsReader.apiURL(environment: environment) + .appendingPathComponent("tokenPlan/detail") + let tokenUsageURL = MiMoSettingsReader.apiURL(environment: environment) + .appendingPathComponent("tokenPlan/usage") + + async let balanceData = self.fetchAuthenticated(url: balanceURL, cookie: normalizedCookie) + let tokenDetailData: Data? = try? await self.fetchAuthenticated(url: tokenDetailURL, cookie: normalizedCookie) + let tokenUsageData: Data? = try? await self.fetchAuthenticated(url: tokenUsageURL, cookie: normalizedCookie) + + return try await self.parseCombinedSnapshot( + balanceData: balanceData, + tokenDetailData: tokenDetailData, + tokenUsageData: tokenUsageData, + now: now) + } + + private static func fetchAuthenticated( + url: URL, + cookie: String, + environment: [String: String] = ProcessInfo.processInfo.environment) async throws -> Data + { + var request = URLRequest(url: url) + request.httpMethod = "GET" + request.timeoutInterval = Self.requestTimeout + request.setValue("application/json, text/plain, */*", forHTTPHeaderField: "Accept") + request.setValue(cookie, forHTTPHeaderField: "Cookie") + request.setValue("en-US,en;q=0.9", forHTTPHeaderField: "Accept-Language") + request.setValue("UTC+01:00", forHTTPHeaderField: "x-timeZone") + request.setValue("https://platform.xiaomimimo.com", forHTTPHeaderField: "Origin") + request.setValue("https://platform.xiaomimimo.com/#/console/balance", forHTTPHeaderField: "Referer") + request.setValue( + "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) " + + "AppleWebKit/537.36 (KHTML, like Gecko) Chrome/143.0.0.0 Safari/537.36", + forHTTPHeaderField: "User-Agent") + + let (data, response) = try await URLSession.shared.data(for: request) + guard let httpResponse = response as? HTTPURLResponse else { + throw MiMoUsageError.networkError("Invalid response") + } + + switch httpResponse.statusCode { + case 200: + break + case 401: + throw MiMoUsageError.loginRequired + case 403: + throw MiMoUsageError.invalidCredentials + default: + throw MiMoUsageError.networkError("HTTP \(httpResponse.statusCode)") + } + + return data + } + + static func parseCombinedSnapshot( + balanceData: Data, + tokenDetailData: Data?, + tokenUsageData: Data?, + now: Date = Date()) throws -> MiMoUsageSnapshot + { + let balanceSnapshot = try self.parseUsageSnapshot(from: balanceData, now: now) + let planDetail: (planCode: String?, periodEnd: Date?, expired: Bool) = { + guard let data = tokenDetailData, let result = try? self.parseTokenPlanDetail(from: data) else { + return (planCode: nil, periodEnd: nil, expired: false) + } + return result + }() + let planUsage: (used: Int, limit: Int, percent: Double) = { + guard let data = tokenUsageData, let result = try? self.parseTokenPlanUsage(from: data) else { + return (used: 0, limit: 0, percent: 0) + } + return result + }() + + return MiMoUsageSnapshot( + balance: balanceSnapshot.balance, + currency: balanceSnapshot.currency, + planCode: planDetail.planCode, + planPeriodEnd: planDetail.periodEnd, + planExpired: planDetail.expired, + tokenUsed: planUsage.used, + tokenLimit: planUsage.limit, + tokenPercent: planUsage.percent, + updatedAt: now) + } + + static func parseUsageSnapshot(from data: Data, now: Date = Date()) throws -> MiMoUsageSnapshot { + let decoder = JSONDecoder() + let response = try decoder.decode(BalanceResponse.self, from: data) + + guard response.code == 0 else { + let message = response.message?.trimmingCharacters(in: .whitespacesAndNewlines) + if response.code == 401 { + throw MiMoUsageError.loginRequired + } + if response.code == 403 { + throw MiMoUsageError.invalidCredentials + } + throw MiMoUsageError.parseFailed(message?.isEmpty == false ? message! : "code \(response.code)") + } + + guard let data = response.data else { + throw MiMoUsageError.parseFailed("Missing balance payload") + } + guard let balance = Double(data.balance) else { + throw MiMoUsageError.parseFailed("Invalid balance value") + } + + let currency = data.currency.trimmingCharacters(in: .whitespacesAndNewlines) + guard !currency.isEmpty else { + throw MiMoUsageError.parseFailed("Missing currency") + } + + return MiMoUsageSnapshot(balance: balance, currency: currency, updatedAt: now) + } + + static func parseTokenPlanDetail(from data: Data) throws -> (planCode: String?, periodEnd: Date?, expired: Bool) { + let decoder = JSONDecoder() + let response = try decoder.decode(TokenPlanDetailResponse.self, from: data) + + guard response.code == 0, let payload = response.data else { + return (planCode: nil, periodEnd: nil, expired: false) + } + + let formatter = DateFormatter() + formatter.dateFormat = "yyyy-MM-dd HH:mm:ss" + formatter.locale = Locale(identifier: "en_US_POSIX") + formatter.timeZone = TimeZone(secondsFromGMT: 0) + + let periodEnd: Date? = if let dateStr = payload.currentPeriodEnd { + formatter.date(from: dateStr) + } else { + nil + } + + return (planCode: payload.planCode, periodEnd: periodEnd, expired: payload.expired) + } + + static func parseTokenPlanUsage(from data: Data) throws -> (used: Int, limit: Int, percent: Double) { + let decoder = JSONDecoder() + let response = try decoder.decode(TokenPlanUsageResponse.self, from: data) + + guard response.code == 0, + let monthUsage = response.data?.monthUsage, + let item = monthUsage.items.first + else { + return (used: 0, limit: 0, percent: 0) + } + + return (used: item.used, limit: item.limit, percent: item.percent) + } + + private struct BalanceResponse: Decodable { + let code: Int + let message: String? + let data: BalancePayload? + } + + private struct BalancePayload: Decodable { + let balance: String + let currency: String + } + + private struct TokenPlanDetailResponse: Decodable { + let code: Int + let message: String? + let data: TokenPlanDetailPayload? + } + + private struct TokenPlanDetailPayload: Decodable { + let planCode: String? + let currentPeriodEnd: String? + let expired: Bool + } + + private struct TokenPlanUsageResponse: Decodable { + let code: Int + let message: String? + let data: TokenPlanUsagePayload? + } + + private struct TokenPlanUsagePayload: Decodable { + let monthUsage: MonthUsage? + } + + private struct MonthUsage: Decodable { + let percent: Double + let items: [UsageItem] + } + + private struct UsageItem: Decodable { + let name: String + let used: Int + let limit: Int + let percent: Double + } +} diff --git a/Sources/CodexBarCore/Providers/MiMo/MiMoUsageSnapshot.swift b/Sources/CodexBarCore/Providers/MiMo/MiMoUsageSnapshot.swift new file mode 100644 index 000000000..a55639a52 --- /dev/null +++ b/Sources/CodexBarCore/Providers/MiMo/MiMoUsageSnapshot.swift @@ -0,0 +1,72 @@ +import Foundation + +public struct MiMoUsageSnapshot: Sendable { + public let balance: Double + public let currency: String + public let planCode: String? + public let planPeriodEnd: Date? + public let planExpired: Bool + public let tokenUsed: Int + public let tokenLimit: Int + public let tokenPercent: Double + public let updatedAt: Date + + public init( + balance: Double, + currency: String, + planCode: String? = nil, + planPeriodEnd: Date? = nil, + planExpired: Bool = false, + tokenUsed: Int = 0, + tokenLimit: Int = 0, + tokenPercent: Double = 0, + updatedAt: Date) + { + self.balance = balance + self.currency = currency + self.planCode = planCode + self.planPeriodEnd = planPeriodEnd + self.planExpired = planExpired + self.tokenUsed = tokenUsed + self.tokenLimit = tokenLimit + self.tokenPercent = tokenPercent + self.updatedAt = updatedAt + } +} + +extension MiMoUsageSnapshot { + public func toUsageSnapshot() -> UsageSnapshot { + let trimmedCurrency = self.currency.trimmingCharacters(in: .whitespacesAndNewlines) + let balanceText = UsageFormatter.currencyString(self.balance, currencyCode: trimmedCurrency) + + let primary: RateWindow? = { + guard self.tokenLimit > 0 else { return nil } + let usedPercent = max(0, min(100, self.tokenPercent * 100)) + let resetDesc = "\(self.tokenUsed.formatted()) / \(self.tokenLimit.formatted()) Credits" + return RateWindow( + usedPercent: usedPercent, + windowMinutes: nil, + resetsAt: self.planPeriodEnd, + resetDescription: resetDesc) + }() + + let planLabel: String? = { + guard let planCode = self.planCode else { return nil } + return planCode.capitalized + }() + + let identity = ProviderIdentitySnapshot( + providerID: .mimo, + accountEmail: nil, + accountOrganization: nil, + loginMethod: planLabel ?? "Balance: \(balanceText)") + + return UsageSnapshot( + primary: primary, + secondary: nil, + tertiary: nil, + providerCost: nil, + updatedAt: self.updatedAt, + identity: identity) + } +} diff --git a/Sources/CodexBarCore/Providers/MiniMax/MiniMaxProviderDescriptor.swift b/Sources/CodexBarCore/Providers/MiniMax/MiniMaxProviderDescriptor.swift index 5bf0a9c52..e758af40f 100644 --- a/Sources/CodexBarCore/Providers/MiniMax/MiniMaxProviderDescriptor.swift +++ b/Sources/CodexBarCore/Providers/MiniMax/MiniMaxProviderDescriptor.swift @@ -87,7 +87,17 @@ struct MiniMaxAPIFetchStrategy: ProviderFetchStrategy { throw MiniMaxAPISettingsError.missingToken } let region = context.settings?.minimax?.apiRegion ?? .global - let usage = try await MiniMaxUsageFetcher.fetchUsage(apiToken: apiToken, region: region) + var usage = try await MiniMaxUsageFetcher.fetchUsage(apiToken: apiToken, region: region) + usage = MiniMaxUsageSnapshot( + planName: usage.planName, + availablePrompts: usage.availablePrompts, + currentPrompts: usage.currentPrompts, + remainingPrompts: usage.remainingPrompts, + windowMinutes: usage.windowMinutes, + usedPercent: usage.usedPercent, + resetsAt: usage.resetsAt, + updatedAt: usage.updatedAt, + apiKey: apiToken) return self.makeResult( usage: usage.toUsageSnapshot(), sourceLabel: "api") diff --git a/Sources/CodexBarCore/Providers/MiniMax/MiniMaxUsageSnapshot.swift b/Sources/CodexBarCore/Providers/MiniMax/MiniMaxUsageSnapshot.swift index 09ed671e2..359254c67 100644 --- a/Sources/CodexBarCore/Providers/MiniMax/MiniMaxUsageSnapshot.swift +++ b/Sources/CodexBarCore/Providers/MiniMax/MiniMaxUsageSnapshot.swift @@ -9,6 +9,7 @@ public struct MiniMaxUsageSnapshot: Sendable { public let usedPercent: Double? public let resetsAt: Date? public let updatedAt: Date + public let apiKey: String? public init( planName: String?, @@ -18,7 +19,8 @@ public struct MiniMaxUsageSnapshot: Sendable { windowMinutes: Int?, usedPercent: Double?, resetsAt: Date?, - updatedAt: Date) + updatedAt: Date, + apiKey: String? = nil) { self.planName = planName self.availablePrompts = availablePrompts @@ -28,6 +30,7 @@ public struct MiniMaxUsageSnapshot: Sendable { self.usedPercent = usedPercent self.resetsAt = resetsAt self.updatedAt = updatedAt + self.apiKey = apiKey } } @@ -43,9 +46,14 @@ extension MiniMaxUsageSnapshot { let planName = self.planName?.trimmingCharacters(in: .whitespacesAndNewlines) let loginMethod = (planName?.isEmpty ?? true) ? nil : planName + let maskedKey: String? = { + guard let key = self.apiKey, !key.isEmpty else { return nil } + if key.count <= 8 { return "****" } + return "\(key.prefix(6))...\(key.suffix(4))" + }() let identity = ProviderIdentitySnapshot( providerID: .minimax, - accountEmail: nil, + accountEmail: maskedKey, accountOrganization: nil, loginMethod: loginMethod) diff --git a/Sources/CodexBarCore/Providers/ProviderDescriptor.swift b/Sources/CodexBarCore/Providers/ProviderDescriptor.swift index d7a3669d4..6d828f2c9 100644 --- a/Sources/CodexBarCore/Providers/ProviderDescriptor.swift +++ b/Sources/CodexBarCore/Providers/ProviderDescriptor.swift @@ -75,6 +75,13 @@ public enum ProviderDescriptorRegistry { .synthetic: SyntheticProviderDescriptor.descriptor, .openrouter: OpenRouterProviderDescriptor.descriptor, .warp: WarpProviderDescriptor.descriptor, + .qwen: QwenProviderDescriptor.descriptor, + .doubao: DoubaoProviderDescriptor.descriptor, + .zenmux: ZenmuxProviderDescriptor.descriptor, + .aigocode: AigoCodeProviderDescriptor.descriptor, + .trae: TraeProviderDescriptor.descriptor, + .stepfun: StepFunProviderDescriptor.descriptor, + .mimo: MiMoProviderDescriptor.descriptor, ] private static let bootstrap: Void = { for provider in UsageProvider.allCases { diff --git a/Sources/CodexBarCore/Providers/ProviderSettingsSnapshot.swift b/Sources/CodexBarCore/Providers/ProviderSettingsSnapshot.swift index e9bf84f9e..34bc88649 100644 --- a/Sources/CodexBarCore/Providers/ProviderSettingsSnapshot.swift +++ b/Sources/CodexBarCore/Providers/ProviderSettingsSnapshot.swift @@ -14,6 +14,8 @@ public struct ProviderSettingsSnapshot: Sendable { copilot: CopilotProviderSettings? = nil, kilo: KiloProviderSettings? = nil, kimi: KimiProviderSettings? = nil, + stepfun: StepFunProviderSettings? = nil, + mimo: MiMoProviderSettings? = nil, augment: AugmentProviderSettings? = nil, amp: AmpProviderSettings? = nil, ollama: OllamaProviderSettings? = nil, @@ -32,6 +34,8 @@ public struct ProviderSettingsSnapshot: Sendable { copilot: copilot, kilo: kilo, kimi: kimi, + stepfun: stepfun, + mimo: mimo, augment: augment, amp: amp, ollama: ollama, @@ -153,6 +157,24 @@ public struct ProviderSettingsSnapshot: Sendable { } } + public struct StepFunProviderSettings: Sendable { + public let cookieSource: ProviderCookieSource + + public init(cookieSource: ProviderCookieSource) { + self.cookieSource = cookieSource + } + } + + public struct MiMoProviderSettings: Sendable { + public let cookieSource: ProviderCookieSource + public let manualCookieHeader: String? + + public init(cookieSource: ProviderCookieSource, manualCookieHeader: String?) { + self.cookieSource = cookieSource + self.manualCookieHeader = manualCookieHeader + } + } + public struct AugmentProviderSettings: Sendable { public let cookieSource: ProviderCookieSource public let manualCookieHeader: String? @@ -203,6 +225,8 @@ public struct ProviderSettingsSnapshot: Sendable { public let copilot: CopilotProviderSettings? public let kilo: KiloProviderSettings? public let kimi: KimiProviderSettings? + public let stepfun: StepFunProviderSettings? + public let mimo: MiMoProviderSettings? public let augment: AugmentProviderSettings? public let amp: AmpProviderSettings? public let ollama: OllamaProviderSettings? @@ -225,6 +249,8 @@ public struct ProviderSettingsSnapshot: Sendable { copilot: CopilotProviderSettings?, kilo: KiloProviderSettings?, kimi: KimiProviderSettings?, + stepfun: StepFunProviderSettings?, + mimo: MiMoProviderSettings?, augment: AugmentProviderSettings?, amp: AmpProviderSettings?, ollama: OllamaProviderSettings?, @@ -242,6 +268,8 @@ public struct ProviderSettingsSnapshot: Sendable { self.copilot = copilot self.kilo = kilo self.kimi = kimi + self.stepfun = stepfun + self.mimo = mimo self.augment = augment self.amp = amp self.ollama = ollama @@ -260,6 +288,8 @@ public enum ProviderSettingsSnapshotContribution: Sendable { case copilot(ProviderSettingsSnapshot.CopilotProviderSettings) case kilo(ProviderSettingsSnapshot.KiloProviderSettings) case kimi(ProviderSettingsSnapshot.KimiProviderSettings) + case stepfun(ProviderSettingsSnapshot.StepFunProviderSettings) + case mimo(ProviderSettingsSnapshot.MiMoProviderSettings) case augment(ProviderSettingsSnapshot.AugmentProviderSettings) case amp(ProviderSettingsSnapshot.AmpProviderSettings) case ollama(ProviderSettingsSnapshot.OllamaProviderSettings) @@ -279,6 +309,8 @@ public struct ProviderSettingsSnapshotBuilder: Sendable { public var copilot: ProviderSettingsSnapshot.CopilotProviderSettings? public var kilo: ProviderSettingsSnapshot.KiloProviderSettings? public var kimi: ProviderSettingsSnapshot.KimiProviderSettings? + public var stepfun: ProviderSettingsSnapshot.StepFunProviderSettings? + public var mimo: ProviderSettingsSnapshot.MiMoProviderSettings? public var augment: ProviderSettingsSnapshot.AugmentProviderSettings? public var amp: ProviderSettingsSnapshot.AmpProviderSettings? public var ollama: ProviderSettingsSnapshot.OllamaProviderSettings? @@ -301,6 +333,8 @@ public struct ProviderSettingsSnapshotBuilder: Sendable { case let .copilot(value): self.copilot = value case let .kilo(value): self.kilo = value case let .kimi(value): self.kimi = value + case let .stepfun(value): self.stepfun = value + case let .mimo(value): self.mimo = value case let .augment(value): self.augment = value case let .amp(value): self.amp = value case let .ollama(value): self.ollama = value @@ -322,6 +356,8 @@ public struct ProviderSettingsSnapshotBuilder: Sendable { copilot: self.copilot, kilo: self.kilo, kimi: self.kimi, + stepfun: self.stepfun, + mimo: self.mimo, augment: self.augment, amp: self.amp, ollama: self.ollama, diff --git a/Sources/CodexBarCore/Providers/ProviderTokenResolver.swift b/Sources/CodexBarCore/Providers/ProviderTokenResolver.swift index ada9fac8d..dc550479a 100644 --- a/Sources/CodexBarCore/Providers/ProviderTokenResolver.swift +++ b/Sources/CodexBarCore/Providers/ProviderTokenResolver.swift @@ -61,6 +61,30 @@ public enum ProviderTokenResolver { self.openRouterResolution(environment: environment)?.token } + public static func qwenToken(environment: [String: String] = ProcessInfo.processInfo.environment) -> String? { + self.qwenResolution(environment: environment)?.token + } + + public static func doubaoToken(environment: [String: String] = ProcessInfo.processInfo.environment) -> String? { + self.doubaoResolution(environment: environment)?.token + } + + public static func stepfunToken(environment: [String: String] = ProcessInfo.processInfo.environment) -> String? { + self.stepfunResolution(environment: environment)?.token + } + + public static func zenmuxToken(environment: [String: String] = ProcessInfo.processInfo.environment) -> String? { + self.zenmuxResolution(environment: environment)?.token + } + + public static func aigocodeToken(environment: [String: String] = ProcessInfo.processInfo.environment) -> String? { + self.aigocodeResolution(environment: environment)?.token + } + + public static func traeToken(environment: [String: String] = ProcessInfo.processInfo.environment) -> String? { + self.traeResolution(environment: environment)?.token + } + public static func zaiResolution( environment: [String: String] = ProcessInfo.processInfo.environment) -> ProviderTokenResolution? { @@ -141,6 +165,42 @@ public enum ProviderTokenResolver { self.resolveEnv(OpenRouterSettingsReader.apiToken(environment: environment)) } + public static func qwenResolution( + environment: [String: String] = ProcessInfo.processInfo.environment) -> ProviderTokenResolution? + { + self.resolveEnv(QwenSettingsReader.apiKey(environment: environment)) + } + + public static func doubaoResolution( + environment: [String: String] = ProcessInfo.processInfo.environment) -> ProviderTokenResolution? + { + self.resolveEnv(DoubaoSettingsReader.apiKey(environment: environment)) + } + + public static func stepfunResolution( + environment: [String: String] = ProcessInfo.processInfo.environment) -> ProviderTokenResolution? + { + self.resolveEnv(StepFunSettingsReader.apiKey(environment: environment)) + } + + public static func zenmuxResolution( + environment: [String: String] = ProcessInfo.processInfo.environment) -> ProviderTokenResolution? + { + self.resolveEnv(ZenmuxSettingsReader.apiKey(environment: environment)) + } + + public static func aigocodeResolution( + environment: [String: String] = ProcessInfo.processInfo.environment) -> ProviderTokenResolution? + { + self.resolveEnv(AigoCodeSettingsReader.apiKey(environment: environment)) + } + + public static func traeResolution( + environment: [String: String] = ProcessInfo.processInfo.environment) -> ProviderTokenResolution? + { + self.resolveEnv(TraeSettingsReader.apiKey(environment: environment)) + } + private static func cleaned(_ raw: String?) -> String? { guard var value = raw?.trimmingCharacters(in: .whitespacesAndNewlines), !value.isEmpty else { return nil diff --git a/Sources/CodexBarCore/Providers/Providers.swift b/Sources/CodexBarCore/Providers/Providers.swift index f48eefe43..bde72c000 100644 --- a/Sources/CodexBarCore/Providers/Providers.swift +++ b/Sources/CodexBarCore/Providers/Providers.swift @@ -25,6 +25,13 @@ public enum UsageProvider: String, CaseIterable, Sendable, Codable { case synthetic case warp case openrouter + case qwen + case doubao + case zenmux + case aigocode + case trae + case stepfun + case mimo } // swiftformat:enable sortDeclarations @@ -52,6 +59,13 @@ public enum IconStyle: Sendable, CaseIterable { case synthetic case warp case openrouter + case qwen + case doubao + case zenmux + case aigocode + case trae + case stepfun + case mimo case combined } diff --git a/Sources/CodexBarCore/Providers/Qwen/QwenProviderDescriptor.swift b/Sources/CodexBarCore/Providers/Qwen/QwenProviderDescriptor.swift new file mode 100644 index 000000000..a390a6bc4 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Qwen/QwenProviderDescriptor.swift @@ -0,0 +1,73 @@ +import CodexBarMacroSupport +import Foundation + +@ProviderDescriptorRegistration +@ProviderDescriptorDefinition +public enum QwenProviderDescriptor { + static func makeDescriptor() -> ProviderDescriptor { + ProviderDescriptor( + id: .qwen, + metadata: ProviderMetadata( + id: .qwen, + displayName: "Qwen", + sessionLabel: "Requests", + weeklyLabel: "Monthly", + opusLabel: nil, + supportsOpus: false, + supportsCredits: false, + creditsHint: "", + toggleTitle: "Show Qwen usage", + cliName: "qwen", + defaultEnabled: false, + isPrimaryProvider: false, + usesAccountFallback: false, + browserCookieOrder: nil, + dashboardURL: "https://bailian.console.aliyun.com/", + statusPageURL: nil), + branding: ProviderBranding( + iconStyle: .qwen, + iconResourceName: "ProviderIcon-qwen", + color: ProviderColor(red: 106 / 255, green: 58 / 255, blue: 255 / 255)), + tokenCost: ProviderTokenCostConfig( + supportsTokenCost: false, + noDataMessage: { "Qwen cost summary is not available." }), + fetchPlan: ProviderFetchPlan( + sourceModes: [.auto, .api], + pipeline: ProviderFetchPipeline(resolveStrategies: { _ in [QwenAPIFetchStrategy()] })), + cli: ProviderCLIConfig( + name: "qwen", + aliases: ["tongyi", "dashscope", "lingma"], + versionDetector: nil)) + } +} + +struct QwenAPIFetchStrategy: ProviderFetchStrategy { + let id: String = "qwen.api" + let kind: ProviderFetchKind = .apiToken + + func isAvailable(_ context: ProviderFetchContext) async -> Bool { + Self.resolveToken(environment: context.env) != nil + } + + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + guard let apiKey = Self.resolveToken(environment: context.env) else { + throw QwenUsageError.missingCredentials + } + let usage = try await QwenUsageFetcher.fetchUsage(apiKey: apiKey) + let accumulated = await LocalUsageTracker.shared.record( + provider: .qwen, + remaining: usage.remainingRequests, + limit: usage.limitRequests) + return self.makeResult( + usage: usage.toUsageSnapshot(accumulated: accumulated), + sourceLabel: "api") + } + + func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool { + false + } + + private static func resolveToken(environment: [String: String]) -> String? { + ProviderTokenResolver.qwenToken(environment: environment) + } +} diff --git a/Sources/CodexBarCore/Providers/Qwen/QwenSettingsReader.swift b/Sources/CodexBarCore/Providers/Qwen/QwenSettingsReader.swift new file mode 100644 index 000000000..4045cc2aa --- /dev/null +++ b/Sources/CodexBarCore/Providers/Qwen/QwenSettingsReader.swift @@ -0,0 +1,36 @@ +import Foundation + +public struct QwenSettingsReader: Sendable { + public static let apiKeyEnvironmentKeys = [ + "DASHSCOPE_API_KEY", + "QWEN_API_KEY", + ] + + public static func apiKey( + environment: [String: String] = ProcessInfo.processInfo.environment) -> String? + { + for key in self.apiKeyEnvironmentKeys { + guard let raw = environment[key]?.trimmingCharacters(in: .whitespacesAndNewlines), + !raw.isEmpty + else { + continue + } + let cleaned = Self.cleaned(raw) + if !cleaned.isEmpty { + return cleaned + } + } + return nil + } + + private static func cleaned(_ raw: String) -> String { + var value = raw + if (value.hasPrefix("\"") && value.hasSuffix("\"")) || + (value.hasPrefix("'") && value.hasSuffix("'")) + { + value.removeFirst() + value.removeLast() + } + return value.trimmingCharacters(in: .whitespacesAndNewlines) + } +} diff --git a/Sources/CodexBarCore/Providers/Qwen/QwenUsageFetcher.swift b/Sources/CodexBarCore/Providers/Qwen/QwenUsageFetcher.swift new file mode 100644 index 000000000..fe828084e --- /dev/null +++ b/Sources/CodexBarCore/Providers/Qwen/QwenUsageFetcher.swift @@ -0,0 +1,315 @@ +import Foundation +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif + +public struct QwenUsageSnapshot: Sendable { + public let remainingRequests: Int + public let limitRequests: Int + public let resetTime: Date? + public let updatedAt: Date + public let apiKeyValid: Bool + public let totalTokens: Int? + public init( + remainingRequests: Int, + limitRequests: Int, + resetTime: Date?, + updatedAt: Date, + apiKeyValid: Bool = false, + totalTokens: Int? = nil) + { + self.remainingRequests = remainingRequests + self.limitRequests = limitRequests + self.resetTime = resetTime + self.updatedAt = updatedAt + self.apiKeyValid = apiKeyValid + self.totalTokens = totalTokens + } + + public func toUsageSnapshot( + accumulated: LocalUsageTracker.AccumulatedUsage? = nil) -> UsageSnapshot + { + let usedPercent: Double + let resetDescription: String + + if self.limitRequests > 0 { + let used = max(0, self.limitRequests - self.remainingRequests) + usedPercent = min(100, max(0, Double(used) / Double(self.limitRequests) * 100)) + resetDescription = "\(used)/\(self.limitRequests) requests" + } else if self.apiKeyValid { + usedPercent = 0 + resetDescription = "Active — check dashboard for details" + } else { + usedPercent = 0 + resetDescription = "No usage data" + } + + let primary = RateWindow( + usedPercent: usedPercent, + windowMinutes: nil, + resetsAt: self.resetTime, + resetDescription: resetDescription) + + var secondary: RateWindow? + if let acc = accumulated, acc.weeklyLimit > 0 { + let weekPercent = min(100, max(0, Double(acc.weeklyRequests) / Double(acc.weeklyLimit) * 100)) + secondary = RateWindow( + usedPercent: weekPercent, + windowMinutes: 7 * 24 * 60, + resetsAt: nil, + resetDescription: "\(acc.weeklyRequests)/\(acc.weeklyLimit) weekly") + } + + let identity = ProviderIdentitySnapshot( + providerID: .qwen, + accountEmail: nil, + accountOrganization: nil, + loginMethod: nil) + + return UsageSnapshot( + primary: primary, + secondary: secondary, + tertiary: nil, + providerCost: nil, + updatedAt: self.updatedAt, + identity: identity) + } +} + +public enum QwenUsageError: LocalizedError, Sendable { + case missingCredentials + case networkError(String) + case apiError(Int, String) + case parseFailed(String) + + public var errorDescription: String? { + switch self { + case .missingCredentials: + "Missing Qwen API key (DASHSCOPE_API_KEY)." + case let .networkError(message): + "Qwen network error: \(message)" + case let .apiError(code, message): + "Qwen API error (\(code)): \(message)" + case let .parseFailed(message): + "Failed to parse Qwen response: \(message)" + } + } +} + +public struct QwenUsageFetcher: Sendable { + private static let log = CodexBarLog.logger(LogCategories.qwenUsage) + + /// Coding plan keys (sk-sp-*) use the coding endpoint; regular keys use compatible mode. + private static func apiURL(for apiKey: String) -> URL { + if apiKey.hasPrefix("sk-sp-") { + URL(string: "https://coding.dashscope.aliyuncs.com/v1/chat/completions")! + } else { + URL(string: "https://dashscope.aliyuncs.com/compatible-mode/v1/chat/completions")! + } + } + + /// Models to probe, ordered by likelihood. We try multiple models because + /// different key types / regions may not have access to every model. + private static let probeModels = [ + "qwen3-coder-plus", + "qwen-turbo", + "qwen-plus", + ] + + public static func fetchUsage(apiKey: String) async throws -> QwenUsageSnapshot { + guard !apiKey.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty else { + throw QwenUsageError.missingCredentials + } + + var lastError: Error? + for model in self.probeModels { + do { + return try await self.probe(apiKey: apiKey, model: model) + } catch let error as QwenUsageError { + // Model not found / not entitled → try next model + if case let .apiError(code, _) = error, code == 404 || code == 403 { + Self.log.debug("Qwen probe model \(model) unavailable (\(code)), trying next") + lastError = error + continue + } + throw error + } + } + throw lastError ?? QwenUsageError.apiError(0, "All probe models failed") + } + + private static func probe(apiKey: String, model: String) async throws -> QwenUsageSnapshot { + let url = self.apiURL(for: apiKey) + + var request = URLRequest(url: url) + request.httpMethod = "POST" + request.timeoutInterval = 30 + request.setValue("application/json", forHTTPHeaderField: "Content-Type") + request.setValue("application/json", forHTTPHeaderField: "Accept") + request.setValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization") + + let body: [String: Any] = [ + "model": model, + "max_tokens": 1, + "messages": [ + ["role": "user", "content": "hi"], + ] as [[String: Any]], + ] + + request.httpBody = try JSONSerialization.data(withJSONObject: body) + + let (data, response) = try await URLSession.shared.data(for: request) + + guard let httpResponse = response as? HTTPURLResponse else { + throw QwenUsageError.networkError("Invalid response") + } + + // Accept both 200 (success) and 429 (rate limited) – both carry rate limit headers. + guard httpResponse.statusCode == 200 || httpResponse.statusCode == 429 else { + let summary = Self.apiErrorSummary(statusCode: httpResponse.statusCode, data: data) + Self.log.error("Qwen API returned \(httpResponse.statusCode): \(summary)") + throw QwenUsageError.apiError(httpResponse.statusCode, summary) + } + + let headers = httpResponse.allHeaderFields + let remaining = Self.intHeader(headers, "x-ratelimit-remaining-requests") + let limit = Self.intHeader(headers, "x-ratelimit-limit-requests") + let resetString = Self.stringHeader(headers, "x-ratelimit-reset-requests") + + let resetTime: Date? = resetString.flatMap(Self.parseResetTime) + + var totalTokens: Int? + if remaining == nil, limit == nil, + let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any], + let usage = json["usage"] as? [String: Any] + { + totalTokens = usage["total_tokens"] as? Int + } + + // 429 means the key is valid but rate-limited; treat it as valid so the UI + // shows "Active" instead of "No usage data" when headers are absent. + let keyValid = httpResponse.statusCode == 200 || httpResponse.statusCode == 429 + + let snapshot = QwenUsageSnapshot( + remainingRequests: remaining ?? 0, + limitRequests: limit ?? 0, + resetTime: resetTime, + updatedAt: Date(), + apiKeyValid: keyValid, + totalTokens: totalTokens) + + Self.log.debug( + "Qwen usage parsed remaining=\(snapshot.remainingRequests) limit=\(snapshot.limitRequests) valid=\(snapshot.apiKeyValid)") // swiftlint:disable:this line_length + + return snapshot + } + + private static func stringHeader(_ headers: [AnyHashable: Any], _ name: String) -> String? { + if let value = headers[name] as? String { return value } + for (key, val) in headers { + if let keyStr = key as? String, + keyStr.caseInsensitiveCompare(name) == .orderedSame, + let valStr = val as? String + { + return valStr + } + } + return nil + } + + private static func intHeader(_ headers: [AnyHashable: Any], _ name: String) -> Int? { + if let value = headers[name] as? String, let int = Int(value) { + return int + } + if let value = headers[name.lowercased()] as? String, let int = Int(value) { + return int + } + // Case-insensitive search + for (key, val) in headers { + if let keyStr = key as? String, + keyStr.lowercased() == name.lowercased(), + let valStr = val as? String, + let int = Int(valStr) + { + return int + } + } + return nil + } + + /// Parse reset time from header value like "1d2h3m4s" or "30s" or ISO 8601. + private static func parseResetTime(_ value: String) -> Date? { + let trimmed = value.trimmingCharacters(in: .whitespacesAndNewlines) + if trimmed.isEmpty { return nil } + + // Try ISO 8601 first + let isoFormatter = ISO8601DateFormatter() + isoFormatter.formatOptions = [.withInternetDateTime, .withFractionalSeconds] + if let date = isoFormatter.date(from: trimmed) { return date } + let isoFallback = ISO8601DateFormatter() + isoFallback.formatOptions = [.withInternetDateTime] + if let date = isoFallback.date(from: trimmed) { return date } + + // Try duration format like "1d2h3m4s" or "30s" + var seconds: TimeInterval = 0 + let pattern = /(\d+)([dhms])/ + for match in trimmed.matches(of: pattern) { + guard let num = Double(match.1) else { continue } + switch match.2 { + case "d": seconds += num * 86400 + case "h": seconds += num * 3600 + case "m": seconds += num * 60 + case "s": seconds += num + default: break + } + } + if seconds > 0 { + return Date().addingTimeInterval(seconds) + } + + // Try plain integer as seconds + if let secs = TimeInterval(trimmed) { + return Date().addingTimeInterval(secs) + } + + return nil + } + + private static func apiErrorSummary(statusCode: Int, data: Data) -> String { + guard let root = try? JSONSerialization.jsonObject(with: data), + let json = root as? [String: Any] + else { + if let text = String(data: data, encoding: .utf8)? + .trimmingCharacters(in: .whitespacesAndNewlines), + !text.isEmpty + { + return self.compactText(text) + } + return "Unexpected response body (\(data.count) bytes)." + } + + if let message = json["message"] as? String { + let trimmed = message.trimmingCharacters(in: .whitespacesAndNewlines) + if !trimmed.isEmpty { return Self.compactText(trimmed) } + } + + if let error = json["error"] as? [String: Any], + let message = error["message"] as? String + { + let trimmed = message.trimmingCharacters(in: .whitespacesAndNewlines) + if !trimmed.isEmpty { return Self.compactText(trimmed) } + } + + return "HTTP \(statusCode) (\(data.count) bytes)." + } + + private static func compactText(_ text: String, maxLength: Int = 200) -> String { + let collapsed = text + .components(separatedBy: .newlines) + .joined(separator: " ") + .trimmingCharacters(in: .whitespacesAndNewlines) + if collapsed.count <= maxLength { return collapsed } + let limitIndex = collapsed.index(collapsed.startIndex, offsetBy: maxLength) + return "\(collapsed[.. 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 StepFunCookieImportError.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 Oasis-Token cookie + guard httpCookies.contains(where: { $0.name == "Oasis-Token" }) else { + continue + } + + log("Found Oasis-Token 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 StepFunCookieImportError.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 emit(_ message: String, logger: ((String) -> Void)?) { + logger?("[stepfun-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 = "\(record.name)|\(record.domain)|\(record.path)" + 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 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 StepFunCookieImportError: LocalizedError { + case noCookies + + var errorDescription: String? { + switch self { + case .noCookies: + "No StepFun session cookies found in browsers." + } + } +} +#endif diff --git a/Sources/CodexBarCore/Providers/StepFun/StepFunDashboardFetcher.swift b/Sources/CodexBarCore/Providers/StepFun/StepFunDashboardFetcher.swift new file mode 100644 index 000000000..3d79b86ed --- /dev/null +++ b/Sources/CodexBarCore/Providers/StepFun/StepFunDashboardFetcher.swift @@ -0,0 +1,200 @@ +#if os(macOS) +import Foundation +import WebKit + +/// Scrapes the StepFun plan-subscribe dashboard to extract Step Plan usage data. +/// +/// StepFun's plan API (GetStepPlanStatus, QueryStepPlanRateLimit) requires browser-bound +/// Oasis-Token that rejects non-browser HTTP clients ("embezzled" error). Using WKWebView +/// ensures the TLS fingerprint matches a real browser, bypassing this restriction. +@MainActor +public struct StepFunDashboardFetcher { + public enum FetchError: LocalizedError { + case loginRequired + case noUsageData(body: String) + case timeout + + public var errorDescription: String? { + switch self { + case .loginRequired: + "StepFun web access requires login. Open Settings → StepFun → Login in Browser." + case let .noUsageData(body): + "StepFun dashboard data not found. Body sample: \(body.prefix(200))" + case .timeout: + "StepFun dashboard loading timed out." + } + } + } + + public struct DashboardSnapshot: Sendable { + public let planName: String? + public let planExpiry: String? + public let fiveHourLeftPercent: Double? + public let fiveHourResetTime: String? + public let weeklyLeftPercent: Double? + public let weeklyResetTime: String? + } + + private static let log = CodexBarLog.logger(LogCategories.stepfunUsage) + private static let dashboardURL = URL(string: "https://platform.stepfun.com/plan-subscribe")! + + public init() {} + + public func fetchDashboard( + websiteDataStore: WKWebsiteDataStore = .default(), + timeout: TimeInterval = 30) async throws -> DashboardSnapshot + { + let deadline = Date().addingTimeInterval(max(1, timeout)) + + let config = WKWebViewConfiguration() + config.websiteDataStore = websiteDataStore + let webView = WKWebView(frame: CGRect(x: -9999, y: -9999, width: 1200, height: 900), configuration: config) + webView.customUserAgent = + "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 " + + "(KHTML, like Gecko) Chrome/135.0.0.0 Safari/537.36" + + defer { + webView.stopLoading() + webView.loadHTMLString("", baseURL: nil) + } + + _ = webView.load(URLRequest(url: Self.dashboardURL)) + Self.log.debug("Loading StepFun dashboard…") + + var lastBody = "" + while Date() < deadline { + try? await Task.sleep(for: .milliseconds(2000)) + + let scrape = try await self.scrape(webView: webView) + + if scrape.isLoginPage { + Self.log.debug("Login page detected") + throw FetchError.loginRequired + } + + lastBody = scrape.bodyText + + if let snapshot = scrape.snapshot { + Self.log.debug( + "Dashboard parsed: plan=\(snapshot.planName ?? "nil") " + + "5h=\(snapshot.fiveHourLeftPercent ?? -1)% " + + "weekly=\(snapshot.weeklyLeftPercent ?? -1)%") + return snapshot + } + } + + throw FetchError.noUsageData(body: lastBody) + } + + // MARK: - JavaScript Scraping + + private struct ScrapeResult { + let isLoginPage: Bool + let bodyText: String + let snapshot: DashboardSnapshot? + } + + private func scrape(webView: WKWebView) async throws -> ScrapeResult { + let js = """ + (() => { + const href = window.location.href; + const body = document.body ? document.body.innerText : ''; + + // Detect login/redirect page + const isLogin = href.includes('need_login_in=1') || + (body.includes('登录') && !body.includes('订阅详情')); + + if (isLogin) { + return JSON.stringify({ isLogin: true, body: body.substring(0, 500) }); + } + + // Extract plan name: "Plus Plan" or similar + const planMatch = body.match(/订阅的版本为(\\S+\\s*Plan)/); + const plan = planMatch ? planMatch[1] : null; + + // Extract expiry date: "有效期截止至2026年04月22日" + const expiryMatch = body.match(/有效期截止至(\\d{4}年\\d{2}月\\d{2}日)/); + const expiry = expiryMatch ? expiryMatch[1] : null; + + // Extract 5-hour usage: "剩余 100%" or "剩余 85%" + // Page structure: "5小时用量" followed by "剩余 XX%" + const fiveHourMatch = body.match(/5小时用量[\\s\\S]*?剩余\\s*(\\d+)%/); + const fiveHourPct = fiveHourMatch ? parseInt(fiveHourMatch[1]) : null; + + // Extract 5-hour reset time + const fiveHourResetMatch = body.match(/5小时用量[\\s\\S]*?重置时间:\\s*([\\d-]+\\s+[\\d:]+)/); + const fiveHourReset = fiveHourResetMatch ? fiveHourResetMatch[1] : null; + + // Extract weekly usage: "每周用量" followed by "剩余 XX%" + const weeklyMatch = body.match(/每周用量[\\s\\S]*?剩余\\s*(\\d+)%/); + const weeklyPct = weeklyMatch ? parseInt(weeklyMatch[1]) : null; + + // Extract weekly reset time + const weeklyResetMatch = body.match(/每周用量[\\s\\S]*?重置时间:\\s*([\\d-]+\\s+[\\d:]+)/); + const weeklyReset = weeklyResetMatch ? weeklyResetMatch[1] : null; + + const hasData = plan || fiveHourPct !== null || weeklyPct !== null; + + return JSON.stringify({ + isLogin: false, + body: body.substring(0, 1500), + plan: plan, + expiry: expiry, + fiveHourPct: fiveHourPct, + fiveHourReset: fiveHourReset, + weeklyPct: weeklyPct, + weeklyReset: weeklyReset, + hasData: hasData, + href: href + }); + })(); + """ + + guard let resultStr = try await webView.evaluateJavaScript(js) as? String, + let data = resultStr.data(using: .utf8), + let dict = try? JSONSerialization.jsonObject(with: data) as? [String: Any] + else { + return ScrapeResult(isLoginPage: false, bodyText: "", snapshot: nil) + } + + let isLogin = (dict["isLogin"] as? Bool) ?? false + let bodyText = (dict["body"] as? String) ?? "" + + if isLogin { + return ScrapeResult(isLoginPage: true, bodyText: bodyText, snapshot: nil) + } + + let hasData = (dict["hasData"] as? Bool) ?? false + guard hasData else { + return ScrapeResult(isLoginPage: false, bodyText: bodyText, snapshot: nil) + } + + let snapshot = DashboardSnapshot( + planName: dict["plan"] as? String, + planExpiry: dict["expiry"] as? String, + fiveHourLeftPercent: (dict["fiveHourPct"] as? Int).map { Double($0) }, + fiveHourResetTime: dict["fiveHourReset"] as? String, + weeklyLeftPercent: (dict["weeklyPct"] as? Int).map { Double($0) }, + weeklyResetTime: dict["weeklyReset"] as? String) + + return ScrapeResult(isLoginPage: false, bodyText: bodyText, snapshot: snapshot) + } + + /// Bridge for calling from non-MainActor contexts (e.g. ProviderFetchStrategy). + public nonisolated static func fetchFromMainActor( + timeout: TimeInterval = 30) async throws -> DashboardSnapshot + { + try await withCheckedThrowingContinuation { continuation in + Task { @MainActor in + do { + let fetcher = StepFunDashboardFetcher() + let snapshot = try await fetcher.fetchDashboard(timeout: timeout) + continuation.resume(returning: snapshot) + } catch { + continuation.resume(throwing: error) + } + } + } + } +} +#endif diff --git a/Sources/CodexBarCore/Providers/StepFun/StepFunProviderDescriptor.swift b/Sources/CodexBarCore/Providers/StepFun/StepFunProviderDescriptor.swift new file mode 100644 index 000000000..e6e96c870 --- /dev/null +++ b/Sources/CodexBarCore/Providers/StepFun/StepFunProviderDescriptor.swift @@ -0,0 +1,107 @@ +import CodexBarMacroSupport +import Foundation + +@ProviderDescriptorRegistration +@ProviderDescriptorDefinition +public enum StepFunProviderDescriptor { + static func makeDescriptor() -> ProviderDescriptor { + ProviderDescriptor( + id: .stepfun, + metadata: ProviderMetadata( + id: .stepfun, + displayName: "StepFun", + sessionLabel: "5h Rate", + weeklyLabel: "Weekly", + opusLabel: nil, + supportsOpus: false, + supportsCredits: false, + creditsHint: "", + toggleTitle: "Show StepFun usage", + cliName: "stepfun", + defaultEnabled: false, + isPrimaryProvider: false, + usesAccountFallback: false, + browserCookieOrder: ProviderBrowserCookieDefaults.defaultImportOrder, + dashboardURL: "https://platform.stepfun.com/plan-subscribe", + statusPageURL: nil), + branding: ProviderBranding( + iconStyle: .stepfun, + iconResourceName: "ProviderIcon-stepfun", + color: ProviderColor(red: 99 / 255, green: 102 / 255, blue: 241 / 255)), + tokenCost: ProviderTokenCostConfig( + supportsTokenCost: false, + noDataMessage: { "StepFun cost summary is not available." }), + fetchPlan: ProviderFetchPlan( + sourceModes: [.auto, .api, .web], + pipeline: ProviderFetchPipeline(resolveStrategies: { _ in [StepFunWebFetchStrategy()] })), + cli: ProviderCLIConfig( + name: "stepfun", + aliases: ["step", "stepfun-ai"], + versionDetector: nil)) + } +} + +struct StepFunWebFetchStrategy: ProviderFetchStrategy { + let id: String = "stepfun.web" + let kind: ProviderFetchKind = .web + private static let log = CodexBarLog.logger(LogCategories.stepfunUsage) + + func isAvailable(_ context: ProviderFetchContext) async -> Bool { + // API key is always required for balance + guard Self.resolveAPIKey(environment: context.env) != nil else { + return false + } + return true + } + + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + guard let apiKey = Self.resolveAPIKey(environment: context.env) else { + throw StepFunUsageError.missingCredentials + } + + // Try WKWebView dashboard scraping for plan/rate limit data (like AigoCode) + #if os(macOS) + var dashboardSnapshot: StepFunDashboardFetcher.DashboardSnapshot? + if context.settings?.stepfun?.cookieSource != .off { + do { + dashboardSnapshot = try await StepFunDashboardFetcher.fetchFromMainActor(timeout: 20) + Self.log.debug("Got StepFun plan data from WKWebView dashboard") + } catch { + Self.log.debug("StepFun dashboard fetch failed: \(error.localizedDescription)") + } + } + #endif + + var dashData: StepFunUsageFetcher.DashboardData? + #if os(macOS) + if let ds = dashboardSnapshot { + dashData = StepFunUsageFetcher.DashboardData( + planName: ds.planName, + planExpiry: ds.planExpiry, + fiveHourLeftPercent: ds.fiveHourLeftPercent, + fiveHourResetTime: ds.fiveHourResetTime, + weeklyLeftPercent: ds.weeklyLeftPercent, + weeklyResetTime: ds.weeklyResetTime) + } + #endif + + let snapshot = try await StepFunUsageFetcher.fetchUsage( + apiKey: apiKey, dashboardData: dashData) + #if os(macOS) + return self.makeResult( + usage: snapshot.toUsageSnapshot(), + sourceLabel: dashboardSnapshot != nil ? "web+api" : "api") + #else + return self.makeResult(usage: snapshot.toUsageSnapshot(), sourceLabel: "api") + #endif + } + + func shouldFallback(on error: Error, context _: ProviderFetchContext) -> Bool { + if case StepFunUsageError.missingCredentials = error { return false } + return true + } + + private static func resolveAPIKey(environment: [String: String]) -> String? { + ProviderTokenResolver.stepfunToken(environment: environment) + } +} diff --git a/Sources/CodexBarCore/Providers/StepFun/StepFunSettingsReader.swift b/Sources/CodexBarCore/Providers/StepFun/StepFunSettingsReader.swift new file mode 100644 index 000000000..7a8f0542f --- /dev/null +++ b/Sources/CodexBarCore/Providers/StepFun/StepFunSettingsReader.swift @@ -0,0 +1,36 @@ +import Foundation + +public struct StepFunSettingsReader: Sendable { + public static let apiKeyEnvironmentKeys = [ + "STEPFUN_API_KEY", + "STEP_API_KEY", + ] + + public static func apiKey( + environment: [String: String] = ProcessInfo.processInfo.environment) -> String? + { + for key in self.apiKeyEnvironmentKeys { + guard let raw = environment[key]?.trimmingCharacters(in: .whitespacesAndNewlines), + !raw.isEmpty + else { + continue + } + let cleaned = Self.cleaned(raw) + if !cleaned.isEmpty { + return cleaned + } + } + return nil + } + + private static func cleaned(_ raw: String) -> String { + var value = raw + if (value.hasPrefix("\"") && value.hasSuffix("\"")) || + (value.hasPrefix("'") && value.hasSuffix("'")) + { + value.removeFirst() + value.removeLast() + } + return value.trimmingCharacters(in: .whitespacesAndNewlines) + } +} diff --git a/Sources/CodexBarCore/Providers/StepFun/StepFunUsageFetcher.swift b/Sources/CodexBarCore/Providers/StepFun/StepFunUsageFetcher.swift new file mode 100644 index 000000000..e2a3e447e --- /dev/null +++ b/Sources/CodexBarCore/Providers/StepFun/StepFunUsageFetcher.swift @@ -0,0 +1,287 @@ +import Foundation + +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif + +public struct StepFunUsageSnapshot: Sendable { + // API balance (prepaid account) + public let balance: Double + public let cashBalance: Double + public let voucherBalance: Double + public let accountType: String + + // Step Plan (coding subscription) + public let planName: String? + public let planExpiredAt: Date? + public let planAutoRenew: Bool + public let fiveHourLeftRate: Double? + public let fiveHourResetTime: Date? + public let weeklyLeftRate: Double? + public let weeklyResetTime: Date? + + public let updatedAt: Date + public let apiKeyValid: Bool + + public init( + balance: Double, + cashBalance: Double, + voucherBalance: Double, + accountType: String, + planName: String? = nil, + planExpiredAt: Date? = nil, + planAutoRenew: Bool = false, + fiveHourLeftRate: Double? = nil, + fiveHourResetTime: Date? = nil, + weeklyLeftRate: Double? = nil, + weeklyResetTime: Date? = nil, + updatedAt: Date, + apiKeyValid: Bool = true) + { + self.balance = balance + self.cashBalance = cashBalance + self.voucherBalance = voucherBalance + self.accountType = accountType + self.planName = planName + self.planExpiredAt = planExpiredAt + self.planAutoRenew = planAutoRenew + self.fiveHourLeftRate = fiveHourLeftRate + self.fiveHourResetTime = fiveHourResetTime + self.weeklyLeftRate = weeklyLeftRate + self.weeklyResetTime = weeklyResetTime + self.updatedAt = updatedAt + self.apiKeyValid = apiKeyValid + } + + public func toUsageSnapshot() -> UsageSnapshot { + // Primary: 5-hour rate limit if available, otherwise balance + let primary: RateWindow + if let rate = self.fiveHourLeftRate { + let usedPercent = max(0, min(100, (1.0 - rate) * 100)) + let pctStr = String(format: "%.0f%%", rate * 100) + primary = RateWindow( + usedPercent: usedPercent, + windowMinutes: 5 * 60, + resetsAt: self.fiveHourResetTime, + resetDescription: "5h remaining: \(pctStr)") + } else { + // Balance display: show as "active" with balance info, low usedPercent + let balanceStr = String(format: "¥%.2f", self.balance) + let voucherStr = self.voucherBalance > 0 + ? String(format: " (voucher: ¥%.2f)", self.voucherBalance) : "" + primary = RateWindow( + usedPercent: self.balance > 0 ? 0 : 100, + windowMinutes: nil, + resetsAt: nil, + resetDescription: "Balance: \(balanceStr)\(voucherStr)") + } + + // Secondary: weekly rate limit if available, otherwise nil + var secondary: RateWindow? + if let weekRate = self.weeklyLeftRate { + let weekUsed = max(0, min(100, (1.0 - weekRate) * 100)) + let weekPctStr = String(format: "%.0f%%", weekRate * 100) + secondary = RateWindow( + usedPercent: weekUsed, + windowMinutes: 7 * 24 * 60, + resetsAt: self.weeklyResetTime, + resetDescription: "Weekly remaining: \(weekPctStr)") + } + + // Tertiary: plan info + balance + var tertiary: RateWindow? + if let planName = self.planName { + var desc = "\(planName) Plan" + if let exp = self.planExpiredAt { + let formatter = DateFormatter() + formatter.dateFormat = "yyyy-MM-dd" + desc += " → \(formatter.string(from: exp))" + } + let balanceStr = String(format: "¥%.2f", self.balance) + desc += " | Balance: \(balanceStr)" + tertiary = RateWindow( + usedPercent: 0, + windowMinutes: nil, + resetsAt: self.planExpiredAt, + resetDescription: desc) + } + + let org: String? = if let planName = self.planName { + "\(planName) Plan" + } else { + self.accountType == "prepaid" ? "Prepaid" : "Postpaid" + } + + let identity = ProviderIdentitySnapshot( + providerID: .stepfun, + accountEmail: nil, + accountOrganization: org, + loginMethod: nil) + + return UsageSnapshot( + primary: primary, + secondary: secondary, + tertiary: tertiary, + providerCost: nil, + updatedAt: self.updatedAt, + identity: identity) + } +} + +public enum StepFunUsageError: LocalizedError, Sendable { + case missingCredentials + case networkError(String) + case apiError(Int, String) + case parseFailed(String) + + public var errorDescription: String? { + switch self { + case .missingCredentials: + "Missing StepFun API key (STEPFUN_API_KEY)." + case let .networkError(message): + "StepFun network error: \(message)" + case let .apiError(code, message): + "StepFun API error (\(code)): \(message)" + case let .parseFailed(message): + "Failed to parse StepFun response: \(message)" + } + } +} + +public struct StepFunUsageFetcher: Sendable { + private static let log = CodexBarLog.logger(LogCategories.stepfunUsage) + private static let accountsURL = URL(string: "https://api.stepfun.com/v1/accounts")! + + public static func fetchUsage( + apiKey: String, + dashboardData: DashboardData? = nil) async throws -> StepFunUsageSnapshot + { + try await self._fetchUsage(apiKey: apiKey, dashboardSnapshot: dashboardData) + } + + public struct DashboardData: Sendable { + public let planName: String? + public let planExpiry: String? + public let fiveHourLeftPercent: Double? + public let fiveHourResetTime: String? + public let weeklyLeftPercent: Double? + public let weeklyResetTime: String? + } + + private static func _fetchUsage( + apiKey: String, + dashboardSnapshot: DashboardData?) async throws -> StepFunUsageSnapshot + { + guard !apiKey.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty else { + throw StepFunUsageError.missingCredentials + } + + // 1. Fetch API balance (always works with API key) + let (balance, cashBalance, voucherBalance, accountType) = try await self.fetchAccountBalance(apiKey: apiKey) + + // 2. Map dashboard snapshot to our model + var planName: String? + var planExpiredAt: Date? + let planAutoRenew = false + var fiveHourLeftRate: Double? + var fiveHourResetTime: Date? + var weeklyLeftRate: Double? + var weeklyResetTime: Date? + + if let dash = dashboardSnapshot { + planName = dash.planName + if let expStr = dash.planExpiry { + // Parse "2026年04月22日" + let fmt = DateFormatter() + fmt.dateFormat = "yyyy年MM月dd日" + planExpiredAt = fmt.date(from: expStr) + } + if let pct = dash.fiveHourLeftPercent { + fiveHourLeftRate = pct / 100.0 + } + if let pct = dash.weeklyLeftPercent { + weeklyLeftRate = pct / 100.0 + } + if let resetStr = dash.fiveHourResetTime { + let fmt = DateFormatter() + fmt.dateFormat = "yyyy-MM-dd HH:mm:ss" + fiveHourResetTime = fmt.date(from: resetStr) + } + if let resetStr = dash.weeklyResetTime { + let fmt = DateFormatter() + fmt.dateFormat = "yyyy-MM-dd HH:mm:ss" + weeklyResetTime = fmt.date(from: resetStr) + } + } + + Self.log.debug( + "StepFun balance=\(balance) plan=\(planName ?? "none") 5h=\(fiveHourLeftRate ?? -1) weekly=\(weeklyLeftRate ?? -1)") // swiftlint:disable:this line_length + + return StepFunUsageSnapshot( + balance: balance, + cashBalance: cashBalance, + voucherBalance: voucherBalance, + accountType: accountType, + planName: planName, + planExpiredAt: planExpiredAt, + planAutoRenew: planAutoRenew, + fiveHourLeftRate: fiveHourLeftRate, + fiveHourResetTime: fiveHourResetTime, + weeklyLeftRate: weeklyLeftRate, + weeklyResetTime: weeklyResetTime, + updatedAt: Date()) + } + + // MARK: - API Balance + + private static func fetchAccountBalance(apiKey: String) async throws -> (Double, Double, Double, String) { + var request = URLRequest(url: self.accountsURL) + request.httpMethod = "GET" + request.timeoutInterval = 30 + request.setValue("application/json", forHTTPHeaderField: "Accept") + request.setValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization") + + let (data, response) = try await URLSession.shared.data(for: request) + + guard let httpResponse = response as? HTTPURLResponse else { + throw StepFunUsageError.networkError("Invalid response") + } + + guard httpResponse.statusCode == 200 else { + let summary = Self.errorSummary(data: data) + Self.log.error("StepFun accounts API returned \(httpResponse.statusCode): \(summary)") + if httpResponse.statusCode == 401 || httpResponse.statusCode == 403 { + throw StepFunUsageError.apiError(httpResponse.statusCode, "Invalid API key") + } + throw StepFunUsageError.apiError(httpResponse.statusCode, summary) + } + + guard let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] else { + throw StepFunUsageError.parseFailed("Invalid JSON response") + } + + let balance = (json["balance"] as? Double) ?? 0 + let cashBalance = (json["total_cash_balance"] as? Double) ?? 0 + let voucherBalance = (json["total_voucher_balance"] as? Double) ?? 0 + let accountType = (json["type"] as? String) ?? "prepaid" + + return (balance, cashBalance, voucherBalance, accountType) + } + + // MARK: - Helpers + + private static func errorSummary(data: Data) -> String { + if let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any], + let error = json["error"] as? [String: Any], + let message = error["message"] as? String + { + return message + } + if let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any], + let message = json["message"] as? String + { + return message + } + return String(data: data, encoding: .utf8) ?? "Unknown error" + } +} diff --git a/Sources/CodexBarCore/Providers/Trae/TraeCookieImporter.swift b/Sources/CodexBarCore/Providers/Trae/TraeCookieImporter.swift new file mode 100644 index 000000000..34260814a --- /dev/null +++ b/Sources/CodexBarCore/Providers/Trae/TraeCookieImporter.swift @@ -0,0 +1,208 @@ +#if os(macOS) +import Foundation +import SweetCookieKit + +/// Imports Trae session cookies from browsers. +/// +/// Trae uses ByteDance Passport for authentication, storing session +/// cookies (`sessionid`, `sid_tt`, `passport_csrf_token`) on the `.trae.ai` domain. +public enum TraeCookieImporter { + private static let log = CodexBarLog.logger(LogCategories.traeCookie) + private static let cookieClient = BrowserCookieClient() + private static let cookieDomains = [".trae.ai", "trae.ai", "www.trae.ai"] + private static let cookieImportOrder: BrowserCookieImportOrder = + ProviderDefaults.metadata[.trae]?.browserCookieOrder ?? Browser.defaultImportOrder + + /// Key cookie names used for authentication. + private static let authCookieNames: Set = [ + "sessionid", + "sid_tt", + ] + + 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 ID cookie value (primary auth token). + public var sessionId: String? { + self.cookies.first(where: { $0.name == "sessionid" })?.value + } + + /// The `sid_tt` cookie value (secondary auth token). + public var sidTT: String? { + self.cookies.first(where: { $0.name == "sid_tt" })?.value + } + + /// The CSRF token for passport requests. + public var csrfToken: String? { + self.cookies.first(where: { $0.name == "passport_csrf_token" })?.value + } + + /// The Cloudide session header value. + public var cloudideSession: String? { + self.cookies.first(where: { $0.name == "X-Cloudide-Session" })?.value + } + + /// Builds a Cookie header string from all cookies. + public var cookieHeader: String { + self.cookies.map { "\($0.name)=\($0.value)" }.joined(separator: "; ") + } + } + + 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 TraeCookieImportError.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 at least one auth cookie + let hasAuth = httpCookies.contains(where: { self.authCookieNames.contains($0.name) }) + guard hasAuth else { continue } + + log("Found Trae session cookies 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 TraeCookieImportError.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 emit(_ message: String, logger: ((String) -> Void)?) { + logger?("[trae-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 TraeCookieImportError: LocalizedError { + case noCookies + + var errorDescription: String? { + switch self { + case .noCookies: + "No Trae session cookies found in browsers." + } + } +} +#endif diff --git a/Sources/CodexBarCore/Providers/Trae/TraeProviderDescriptor.swift b/Sources/CodexBarCore/Providers/Trae/TraeProviderDescriptor.swift new file mode 100644 index 000000000..3ccd3f65d --- /dev/null +++ b/Sources/CodexBarCore/Providers/Trae/TraeProviderDescriptor.swift @@ -0,0 +1,114 @@ +import CodexBarMacroSupport +import Foundation + +@ProviderDescriptorRegistration +@ProviderDescriptorDefinition +public enum TraeProviderDescriptor { + static func makeDescriptor() -> ProviderDescriptor { + ProviderDescriptor( + id: .trae, + metadata: ProviderMetadata( + id: .trae, + displayName: "Trae", + sessionLabel: "Usage", + weeklyLabel: "Quota", + opusLabel: nil, + supportsOpus: false, + supportsCredits: false, + creditsHint: "", + toggleTitle: "Show Trae usage", + cliName: "trae", + defaultEnabled: false, + isPrimaryProvider: false, + usesAccountFallback: false, + browserCookieOrder: nil, + dashboardURL: "https://www.trae.ai/account-setting#usage", + statusPageURL: nil), + branding: ProviderBranding( + iconStyle: .trae, + iconResourceName: "ProviderIcon-trae", + color: ProviderColor(red: 59 / 255, green: 130 / 255, blue: 246 / 255)), + tokenCost: ProviderTokenCostConfig( + supportsTokenCost: false, + noDataMessage: { "Trae cost summary is not available." }), + fetchPlan: ProviderFetchPlan( + sourceModes: [.auto, .web], + pipeline: ProviderFetchPipeline(resolveStrategies: { _ in + [TraeWebFetchStrategy(), TraeLocalFetchStrategy()] + })), + cli: ProviderCLIConfig( + name: "trae", + aliases: [], + versionDetector: nil)) + } +} + +struct TraeLocalFetchStrategy: ProviderFetchStrategy { + let id: String = "trae.local" + let kind: ProviderFetchKind = .localProbe + + func isAvailable(_ context: ProviderFetchContext) async -> Bool { + FileManager.default.fileExists(atPath: "/Applications/Trae.app") + } + + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + let status = try await TraeStatusProbe.probe() + return self.makeResult( + usage: status.toUsageSnapshot(), + sourceLabel: "local") + } + + func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool { + false + } +} + +#if os(macOS) +struct TraeWebFetchStrategy: ProviderFetchStrategy { + let id: String = "trae.web" + let kind: ProviderFetchKind = .web + private static let log = CodexBarLog.logger(LogCategories.traeWeb) + + func isAvailable(_ context: ProviderFetchContext) async -> Bool { + if context.sourceMode == .web { return true } + if context.sourceMode == .auto { + return TraeCookieImporter.hasSession() + } + return false + } + + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + let cookieSession = try TraeCookieImporter.importSession() + Self.log.debug("Found Trae session in \(cookieSession.sourceLabel)") + + let session = TraeSessionInfo(from: cookieSession) + let snapshot = try await TraeUsageFetcher.fetchUsage(session: session) + return self.makeResult( + usage: snapshot.toUsageSnapshot(), + sourceLabel: "web (\(cookieSession.sourceLabel))") + } + + func shouldFallback(on error: Error, context: ProviderFetchContext) -> Bool { + if context.sourceMode == .web { return false } + // In auto mode, fall back to local probe on cookie/API errors + return true + } +} +#else +struct TraeWebFetchStrategy: ProviderFetchStrategy { + let id: String = "trae.web" + let kind: ProviderFetchKind = .web + + func isAvailable(_ context: ProviderFetchContext) async -> Bool { + false + } + + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + throw TraeAPIError.networkError("Web strategy not available on this platform") + } + + func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool { + true + } +} +#endif diff --git a/Sources/CodexBarCore/Providers/Trae/TraeSettingsReader.swift b/Sources/CodexBarCore/Providers/Trae/TraeSettingsReader.swift new file mode 100644 index 000000000..398076007 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Trae/TraeSettingsReader.swift @@ -0,0 +1,35 @@ +import Foundation + +public struct TraeSettingsReader: Sendable { + public static let apiKeyEnvironmentKeys = [ + "TRAE_API_KEY", + ] + + public static func apiKey( + environment: [String: String] = ProcessInfo.processInfo.environment) -> String? + { + for key in self.apiKeyEnvironmentKeys { + guard let raw = environment[key]?.trimmingCharacters(in: .whitespacesAndNewlines), + !raw.isEmpty + else { + continue + } + let cleaned = Self.cleaned(raw) + if !cleaned.isEmpty { + return cleaned + } + } + return nil + } + + private static func cleaned(_ raw: String) -> String { + var value = raw + if (value.hasPrefix("\"") && value.hasSuffix("\"")) || + (value.hasPrefix("'") && value.hasSuffix("'")) + { + value.removeFirst() + value.removeLast() + } + return value.trimmingCharacters(in: .whitespacesAndNewlines) + } +} diff --git a/Sources/CodexBarCore/Providers/Trae/TraeStatusProbe.swift b/Sources/CodexBarCore/Providers/Trae/TraeStatusProbe.swift new file mode 100644 index 000000000..85275fd56 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Trae/TraeStatusProbe.swift @@ -0,0 +1,104 @@ +import Foundation + +public struct TraeStatusSnapshot: Sendable { + public let isRunning: Bool + public let version: String? + public let updatedAt: Date + + public init(isRunning: Bool, version: String?, updatedAt: Date) { + self.isRunning = isRunning + self.version = version + self.updatedAt = updatedAt + } + + public func toUsageSnapshot() -> UsageSnapshot { + let usedPercent: Double = self.isRunning ? 0 : 100 + let resetDescription: String = self.isRunning + ? "Active — free tier" + : "Not running" + + let primary = RateWindow( + usedPercent: usedPercent, + windowMinutes: nil, + resetsAt: nil, + resetDescription: resetDescription) + + let versionLabel = self.version.map { "v\($0)" } + let identity = ProviderIdentitySnapshot( + providerID: .trae, + accountEmail: versionLabel, + accountOrganization: nil, + loginMethod: self.isRunning ? "Free" : nil) + + return UsageSnapshot( + primary: primary, + secondary: nil, + tertiary: nil, + providerCost: nil, + updatedAt: self.updatedAt, + identity: identity) + } +} + +public enum TraeStatusProbeError: LocalizedError, Sendable { + case notInstalled + case notRunning + + public var errorDescription: String? { + switch self { + case .notInstalled: + "Trae is not installed." + case .notRunning: + "Trae is not running." + } + } +} + +public struct TraeStatusProbe: Sendable { + private static let log = CodexBarLog.logger(LogCategories.traeUsage) + + public static func probe() async throws -> TraeStatusSnapshot { + let appPath = "/Applications/Trae.app" + let fm = FileManager.default + guard fm.fileExists(atPath: appPath) else { + throw TraeStatusProbeError.notInstalled + } + + let isRunning = Self.isTraeRunning() + let version = Self.traeVersion(appPath: appPath) + + Self.log.debug("Trae probe: running=\(isRunning) version=\(version ?? "unknown")") + + return TraeStatusSnapshot( + isRunning: isRunning, + version: version, + updatedAt: Date()) + } + + private static func isTraeRunning() -> Bool { + let pipe = Pipe() + let process = Process() + process.executableURL = URL(fileURLWithPath: "/usr/bin/pgrep") + process.arguments = ["-f", "Trae.app"] + process.standardOutput = pipe + process.standardError = FileHandle.nullDevice + do { + try process.run() + process.waitUntilExit() + return process.terminationStatus == 0 + } catch { + return false + } + } + + private static func traeVersion(appPath: String) -> String? { + let plistPath = "\(appPath)/Contents/Info.plist" + guard let plistData = FileManager.default.contents(atPath: plistPath), + let plist = try? PropertyListSerialization.propertyList( + from: plistData, options: [], format: nil) as? [String: Any] + else { + return nil + } + return plist["CFBundleShortVersionString"] as? String + } +} diff --git a/Sources/CodexBarCore/Providers/Trae/TraeUsageFetcher.swift b/Sources/CodexBarCore/Providers/Trae/TraeUsageFetcher.swift new file mode 100644 index 000000000..9a915b081 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Trae/TraeUsageFetcher.swift @@ -0,0 +1,407 @@ +import Foundation + +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif + +public struct TraeUsageFetcher: Sendable { + private static let log = CodexBarLog.logger(LogCategories.traeWeb) + /// Global endpoint — routes to the correct region automatically. + /// Do NOT use ug-normal.us.trae.ai; it rejects non-US sessions. + private static let globalBase = "https://ug-normal.trae.ai" + + private static func apiURL(_ base: String, path: String) -> URL { + URL(string: "\(base)/\(path)")! + } + + private static func apiURL(_ base: URL, path: String) -> URL { + URL(string: "\(base.absoluteString)/\(path)")! + } + + public static func fetchUsage(session: TraeSessionInfo, now: Date = Date()) async throws -> TraeUsageSnapshot { + // Step 1: Check login status and get region info + let loginResult = try await self.checkLogin(session: session) + guard loginResult.isLogin else { + throw TraeAPIError.invalidSession + } + + Self.log.debug("Trae login valid: userID=\(loginResult.userID ?? "?") region=\(loginResult.region ?? "?")") + + // Determine the regional host for subsequent API calls + let regionalBase: URL = if let host = loginResult.host, let url = URL(string: host) { + url + } else { + URL(string: self.globalBase)! + } + + // Step 2: Fetch user profile and usage stats in parallel + async let profileResult = self.getUserInfo(base: regionalBase, session: session) + async let statsResult = self.getUserStats(base: regionalBase, session: session) + + let profile = try await profileResult + let stats = try? await statsResult // stats failure is non-fatal + + return TraeUsageSnapshot( + checkLogin: loginResult, profile: profile, stats: stats, updatedAt: now) + } + + // MARK: - CheckLogin + + private static func checkLogin(session: TraeSessionInfo) async throws -> TraeCheckLoginResult { + let url = self.apiURL(self.globalBase, path: "cloudide/api/v3/trae/CheckLogin") + var request = self.makeRequest(url: url, session: session) + request.httpBody = "{}".data(using: .utf8) // swiftlint:disable:this non_optional_string_data_conversion + + let (data, response) = try await URLSession.shared.data(for: request) + guard let httpResponse = response as? HTTPURLResponse else { + throw TraeAPIError.networkError("Invalid response") + } + + guard httpResponse.statusCode == 200 else { + if httpResponse.statusCode == 401 { + throw TraeAPIError.invalidSession + } + throw TraeAPIError.apiError("CheckLogin HTTP \(httpResponse.statusCode)") + } + + let volcResponse = try JSONDecoder().decode(TraeVolcResponse.self, from: data) + if let error = volcResponse.responseMetadata.error { + throw TraeAPIError.apiError("CheckLogin: \(error.message ?? error.code)") + } + guard let result = volcResponse.result else { + throw TraeAPIError.parseFailed("CheckLogin returned no Result") + } + return result + } + + // MARK: - GetUserInfo (profile data) + + private static func getUserInfo( + base: URL, session: TraeSessionInfo) async throws -> TraeProfileResult + { + let url = self.apiURL(base, path: "cloudide/api/v3/trae/GetUserInfo") + var request = self.makeRequest(url: url, session: session) + request.httpBody = "{}".data(using: .utf8) // swiftlint:disable:this non_optional_string_data_conversion + + let (data, response) = try await URLSession.shared.data(for: request) + guard let httpResponse = response as? HTTPURLResponse, httpResponse.statusCode == 200 else { + return TraeProfileResult() + } + + do { + let volcResponse = try JSONDecoder().decode( + TraeVolcResponse.self, from: data) + if volcResponse.responseMetadata.error != nil { return TraeProfileResult() } + return volcResponse.result ?? TraeProfileResult() + } catch { + Self.log.warning("GetUserInfo decode failed: \(error)") + return TraeProfileResult() + } + } + + // MARK: - GetUserStasticData (usage statistics) + + private static func getUserStats( + base: URL, session: TraeSessionInfo) async throws -> TraeStatsResult + { + let url = self.apiURL(base, path: "cloudide/api/v3/trae/GetUserStasticData") + var request = self.makeRequest(url: url, session: session) + + // API requires LocalTime (ISO 8601 with offset) and Offset (timezone minutes) + let now = Date() + let tz = TimeZone.current + let offsetMinutes = tz.secondsFromGMT(for: now) / 60 + let formatter = ISO8601DateFormatter() + formatter.formatOptions = [.withInternetDateTime] + formatter.timeZone = tz + let localTime = formatter.string(from: now) + + let body: [String: Any] = ["LocalTime": localTime, "Offset": offsetMinutes] + request.httpBody = try? JSONSerialization.data(withJSONObject: body) + + let (data, response) = try await URLSession.shared.data(for: request) + guard let httpResponse = response as? HTTPURLResponse, httpResponse.statusCode == 200 else { + throw TraeAPIError.apiError("GetUserStasticData HTTP error") + } + + let volcResponse = try JSONDecoder().decode( + TraeVolcResponse.self, from: data) + if let error = volcResponse.responseMetadata.error { + throw TraeAPIError.apiError("Stats: \(error.message ?? error.code)") + } + guard let result = volcResponse.result else { + throw TraeAPIError.parseFailed("GetUserStasticData returned no Result") + } + return result + } + + // MARK: - Request Builder + + private static func makeRequest(url: URL, session: TraeSessionInfo) -> URLRequest { + var request = URLRequest(url: url) + request.httpMethod = "POST" + request.setValue("application/json", forHTTPHeaderField: "Content-Type") + request.setValue("application/json", forHTTPHeaderField: "Accept") + request.setValue(session.cookieHeader, forHTTPHeaderField: "Cookie") + request.setValue("https://www.trae.ai", forHTTPHeaderField: "Origin") + request.setValue("https://www.trae.ai/account-setting", forHTTPHeaderField: "Referer") + let userAgent = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) " + + "AppleWebKit/537.36 (KHTML, like Gecko) Chrome/145.0.0.0 Safari/537.36" + request.setValue(userAgent, forHTTPHeaderField: "User-Agent") + + if let csrfToken = session.csrfToken { + request.setValue(csrfToken, forHTTPHeaderField: "x-csrf-token") + } + if let cloudideSession = session.cloudideSession { + request.setValue(cloudideSession, forHTTPHeaderField: "X-Cloudide-Session") + } + return request + } +} + +// MARK: - Session Info (abstraction over cookie source) + +public struct TraeSessionInfo: Sendable { + public let cookieHeader: String + public let csrfToken: String? + public let cloudideSession: String? + public let sourceLabel: String + + public init(cookieHeader: String, csrfToken: String?, cloudideSession: String?, sourceLabel: String) { + self.cookieHeader = cookieHeader + self.csrfToken = csrfToken + self.cloudideSession = cloudideSession + self.sourceLabel = sourceLabel + } + + #if os(macOS) + public init(from cookieSession: TraeCookieImporter.SessionInfo) { + self.cookieHeader = cookieSession.cookieHeader + self.csrfToken = cookieSession.csrfToken + self.cloudideSession = cookieSession.cloudideSession + self.sourceLabel = cookieSession.sourceLabel + } + #endif +} + +// MARK: - ByteDance Volc API Response Format + +/// Generic ByteDance Volc Engine API response wrapper. +struct TraeVolcResponse: Codable, Sendable { + let responseMetadata: TraeVolcResponseMetadata + let result: T? + + enum CodingKeys: String, CodingKey { + case responseMetadata = "ResponseMetadata" + case result = "Result" + } +} + +struct TraeVolcResponseMetadata: Codable, Sendable { + let requestId: String? + let action: String? + let version: String? + let service: String? + let region: String? + let error: TraeVolcError? + + enum CodingKeys: String, CodingKey { + case requestId = "RequestId" + case action = "Action" + case version = "Version" + case service = "Service" + case region = "Region" + case error = "Error" + } +} + +struct TraeVolcError: Codable, Sendable { + let code: String + let standardCode: String? + let message: String? + + enum CodingKeys: String, CodingKey { + case code = "Code" + case standardCode = "StandardCode" + case message = "Message" + } +} + +// MARK: - CheckLogin Result + +struct TraeCheckLoginResult: Codable, Sendable { + let isLogin: Bool + let expiredAt: Int? + let region: String? + let host: String? + let userID: String? + let aiRegion: String? + let aiHost: String? + let aiPayHost: String? + let nickNameEditStatus: String? + let passwordChanged: Bool? + + init( + isLogin: Bool = false, + expiredAt: Int? = nil, + region: String? = nil, + host: String? = nil, + userID: String? = nil, + aiRegion: String? = nil, + aiHost: String? = nil, + aiPayHost: String? = nil, + nickNameEditStatus: String? = nil, + passwordChanged: Bool? = nil) + { + self.isLogin = isLogin + self.expiredAt = expiredAt + self.region = region + self.host = host + self.userID = userID + self.aiRegion = aiRegion + self.aiHost = aiHost + self.aiPayHost = aiPayHost + self.nickNameEditStatus = nickNameEditStatus + self.passwordChanged = passwordChanged + } + + enum CodingKeys: String, CodingKey { + case isLogin = "IsLogin" + case expiredAt = "ExpiredAt" + case region = "Region" + case host = "Host" + case userID = "UserID" + case aiRegion = "AIRegion" + case aiHost = "AIHost" + case aiPayHost = "AIPayHost" + case nickNameEditStatus = "NickNameEditStatus" + case passwordChanged = "PasswordChanged" + } +} + +// MARK: - GetUserInfo Result (profile data) + +struct TraeProfileResult: Codable, Sendable { + let screenName: String? + let userID: String? + let avatarURL: String? + let region: String? + let aiRegion: String? + let registerTime: String? + let lastLoginTime: String? + let lastLoginType: String? + + init( + screenName: String? = nil, + userID: String? = nil, + avatarURL: String? = nil, + region: String? = nil, + aiRegion: String? = nil, + registerTime: String? = nil, + lastLoginTime: String? = nil, + lastLoginType: String? = nil) + { + self.screenName = screenName + self.userID = userID + self.avatarURL = avatarURL + self.region = region + self.aiRegion = aiRegion + self.registerTime = registerTime + self.lastLoginTime = lastLoginTime + self.lastLoginType = lastLoginType + } + + enum CodingKeys: String, CodingKey { + case screenName = "ScreenName" + case userID = "UserID" + case avatarURL = "AvatarUrl" + case region = "Region" + case aiRegion = "AIRegion" + case registerTime = "RegisterTime" + case lastLoginTime = "LastLoginTime" + case lastLoginType = "LastLoginType" + } + + init(from decoder: Decoder) throws { + let container = try? decoder.container(keyedBy: CodingKeys.self) + self.screenName = try? container?.decodeIfPresent(String.self, forKey: .screenName) + self.userID = try? container?.decodeIfPresent(String.self, forKey: .userID) + self.avatarURL = try? container?.decodeIfPresent(String.self, forKey: .avatarURL) + self.region = try? container?.decodeIfPresent(String.self, forKey: .region) + self.aiRegion = try? container?.decodeIfPresent(String.self, forKey: .aiRegion) + self.registerTime = try? container?.decodeIfPresent(String.self, forKey: .registerTime) + self.lastLoginTime = try? container?.decodeIfPresent(String.self, forKey: .lastLoginTime) + self.lastLoginType = try? container?.decodeIfPresent(String.self, forKey: .lastLoginType) + } +} + +// MARK: - GetUserStasticData Result (usage statistics) + +struct TraeStatsResult: Codable, Sendable { + let userID: String? + let registerDays: Int? + /// AI interaction counts per day (key: "yyyyMMdd", value: count) + let aiCnt365d: [String: Int]? + let codeAiAcceptCnt7d: Int? + /// Accepted AI suggestions by language (key: language, value: count) + let codeAiAcceptDiffLanguageCnt7d: [String: Int]? + let codeCompCnt7d: Int? + /// Completions by agent (key: agent name, value: count) + let codeCompDiffAgentCnt7d: [String: Int]? + /// Completions by model (key: model name, value: count) + let codeCompDiffModelCnt7d: [String: Int]? + let dataDate: String? + let isIde: Bool? + + enum CodingKeys: String, CodingKey { + case userID = "UserID" + case registerDays = "RegisterDays" + case aiCnt365d = "AiCnt365d" + case codeAiAcceptCnt7d = "CodeAiAcceptCnt7d" + case codeAiAcceptDiffLanguageCnt7d = "CodeAiAcceptDiffLanguageCnt7d" + case codeCompCnt7d = "CodeCompCnt7d" + case codeCompDiffAgentCnt7d = "CodeCompDiffAgentCnt7d" + case codeCompDiffModelCnt7d = "CodeCompDiffModelCnt7d" + case dataDate = "DataDate" + case isIde = "IsIde" + } + + init(from decoder: Decoder) throws { + let container = try? decoder.container(keyedBy: CodingKeys.self) + self.userID = try? container?.decodeIfPresent(String.self, forKey: .userID) + self.registerDays = try? container?.decodeIfPresent(Int.self, forKey: .registerDays) + self.aiCnt365d = try? container?.decodeIfPresent([String: Int].self, forKey: .aiCnt365d) + self.codeAiAcceptCnt7d = try? container?.decodeIfPresent(Int.self, forKey: .codeAiAcceptCnt7d) + self.codeAiAcceptDiffLanguageCnt7d = try? container?.decodeIfPresent( + [String: Int].self, forKey: .codeAiAcceptDiffLanguageCnt7d) + self.codeCompCnt7d = try? container?.decodeIfPresent(Int.self, forKey: .codeCompCnt7d) + self.codeCompDiffAgentCnt7d = try? container?.decodeIfPresent( + [String: Int].self, forKey: .codeCompDiffAgentCnt7d) + self.codeCompDiffModelCnt7d = try? container?.decodeIfPresent( + [String: Int].self, forKey: .codeCompDiffModelCnt7d) + self.dataDate = try? container?.decodeIfPresent(String.self, forKey: .dataDate) + self.isIde = try? container?.decodeIfPresent(Bool.self, forKey: .isIde) + } +} + +// MARK: - Errors + +public enum TraeAPIError: LocalizedError, Sendable { + case invalidSession + case networkError(String) + case parseFailed(String) + case apiError(String) + + public var errorDescription: String? { + switch self { + case .invalidSession: + "Trae session expired. Please log in to trae.ai in your browser." + case let .networkError(msg): + "Trae network error: \(msg)" + case let .parseFailed(msg): + "Trae response parse failed: \(msg)" + case let .apiError(msg): + "Trae API error: \(msg)" + } + } +} diff --git a/Sources/CodexBarCore/Providers/Trae/TraeUsageSnapshot.swift b/Sources/CodexBarCore/Providers/Trae/TraeUsageSnapshot.swift new file mode 100644 index 000000000..3bd42f858 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Trae/TraeUsageSnapshot.swift @@ -0,0 +1,62 @@ +import Foundation + +public struct TraeUsageSnapshot: Sendable { + let checkLogin: TraeCheckLoginResult + let profile: TraeProfileResult + let stats: TraeStatsResult? + public let updatedAt: Date +} + +extension TraeUsageSnapshot { + public func toUsageSnapshot() -> UsageSnapshot { + let primary: RateWindow + + if let stats { + // Sum 7-day AI interaction counts + let total7d = (stats.codeAiAcceptCnt7d ?? 0) + (stats.codeCompCnt7d ?? 0) + + // Build model breakdown string + let modelBreakdown = stats.codeCompDiffModelCnt7d? + .sorted { $0.value > $1.value } + .map { "\($0.key): \($0.value)" } + .joined(separator: ", ") + + let description = if let modelBreakdown, !modelBreakdown.isEmpty { + "\(total7d) AI actions (7d) — \(modelBreakdown)" + } else { + "\(total7d) AI actions (7d)" + } + + // Trae has no hard usage cap, so show activity level instead of percent + primary = RateWindow( + usedPercent: 0, + windowMinutes: 7 * 24 * 60, + resetsAt: nil, + resetDescription: description) + } else { + primary = RateWindow( + usedPercent: 0, + windowMinutes: nil, + resetsAt: nil, + resetDescription: "Active — logged in") + } + + let accountName = self.profile.screenName + ?? self.checkLogin.userID + let regionInfo = self.checkLogin.region ?? self.profile.aiRegion + + let identity = ProviderIdentitySnapshot( + providerID: .trae, + accountEmail: accountName, + accountOrganization: regionInfo, + loginMethod: self.profile.lastLoginType ?? "Web") + + return UsageSnapshot( + primary: primary, + secondary: nil, + tertiary: nil, + providerCost: nil, + updatedAt: self.updatedAt, + identity: identity) + } +} diff --git a/Sources/CodexBarCore/Providers/Zai/ZaiUsageStats.swift b/Sources/CodexBarCore/Providers/Zai/ZaiUsageStats.swift index 1592a6181..0713291d6 100644 --- a/Sources/CodexBarCore/Providers/Zai/ZaiUsageStats.swift +++ b/Sources/CodexBarCore/Providers/Zai/ZaiUsageStats.swift @@ -132,12 +132,20 @@ public struct ZaiUsageSnapshot: Sendable { public let timeLimit: ZaiLimitEntry? public let planName: String? public let updatedAt: Date + public let apiKey: String? - public init(tokenLimit: ZaiLimitEntry?, timeLimit: ZaiLimitEntry?, planName: String?, updatedAt: Date) { + public init( + tokenLimit: ZaiLimitEntry?, + timeLimit: ZaiLimitEntry?, + planName: String?, + updatedAt: Date, + apiKey: String? = nil) + { self.tokenLimit = tokenLimit self.timeLimit = timeLimit self.planName = planName self.updatedAt = updatedAt + self.apiKey = apiKey } /// Returns true if this snapshot contains valid z.ai data @@ -160,9 +168,14 @@ extension ZaiUsageSnapshot { let planName = self.planName?.trimmingCharacters(in: .whitespacesAndNewlines) let loginMethod = (planName?.isEmpty ?? true) ? nil : planName + let maskedKey: String? = { + guard let key = self.apiKey, !key.isEmpty else { return nil } + if key.count <= 8 { return "****" } + return "\(key.prefix(6))...\(key.suffix(4))" + }() let identity = ProviderIdentitySnapshot( providerID: .zai, - accountEmail: nil, + accountEmail: maskedKey, accountOrganization: nil, loginMethod: loginMethod) return UsageSnapshot( @@ -329,7 +342,14 @@ public struct ZaiUsageFetcher: Sendable { } do { - return try Self.parseUsageSnapshot(from: data) + var snapshot = try Self.parseUsageSnapshot(from: data) + snapshot = ZaiUsageSnapshot( + tokenLimit: snapshot.tokenLimit, + timeLimit: snapshot.timeLimit, + planName: snapshot.planName, + updatedAt: snapshot.updatedAt, + apiKey: apiKey) + return snapshot } catch let error as DecodingError { Self.log.error("z.ai JSON decoding error: \(error.localizedDescription)") throw ZaiUsageError.parseFailed(error.localizedDescription) diff --git a/Sources/CodexBarCore/Providers/Zenmux/ZenmuxProviderDescriptor.swift b/Sources/CodexBarCore/Providers/Zenmux/ZenmuxProviderDescriptor.swift new file mode 100644 index 000000000..4f6158f81 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Zenmux/ZenmuxProviderDescriptor.swift @@ -0,0 +1,69 @@ +import CodexBarMacroSupport +import Foundation + +@ProviderDescriptorRegistration +@ProviderDescriptorDefinition +public enum ZenmuxProviderDescriptor { + static func makeDescriptor() -> ProviderDescriptor { + ProviderDescriptor( + id: .zenmux, + metadata: ProviderMetadata( + id: .zenmux, + displayName: "Zenmux", + sessionLabel: "Requests", + weeklyLabel: "Rate limit", + opusLabel: nil, + supportsOpus: false, + supportsCredits: false, + creditsHint: "", + toggleTitle: "Show Zenmux usage", + cliName: "zenmux", + defaultEnabled: false, + isPrimaryProvider: false, + usesAccountFallback: false, + browserCookieOrder: nil, + dashboardURL: "https://zenmux.ai/platform/subscription", + statusPageURL: nil), + branding: ProviderBranding( + iconStyle: .zenmux, + iconResourceName: "ProviderIcon-zenmux", + color: ProviderColor(red: 255 / 255, green: 140 / 255, blue: 0 / 255)), + tokenCost: ProviderTokenCostConfig( + supportsTokenCost: false, + noDataMessage: { "Zenmux cost summary is not available." }), + fetchPlan: ProviderFetchPlan( + sourceModes: [.auto, .api], + pipeline: ProviderFetchPipeline(resolveStrategies: { _ in [ZenmuxAPIFetchStrategy()] })), + cli: ProviderCLIConfig( + name: "zenmux", + aliases: ["zen"], + versionDetector: nil)) + } +} + +struct ZenmuxAPIFetchStrategy: ProviderFetchStrategy { + let id: String = "zenmux.api" + let kind: ProviderFetchKind = .apiToken + + func isAvailable(_ context: ProviderFetchContext) async -> Bool { + Self.resolveToken(environment: context.env) != nil + } + + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + guard let apiKey = Self.resolveToken(environment: context.env) else { + throw ZenmuxUsageError.missingCredentials + } + let usage = try await ZenmuxUsageFetcher.fetchUsage(apiKey: apiKey) + return self.makeResult( + usage: usage.toUsageSnapshot(), + sourceLabel: "api") + } + + func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool { + false + } + + private static func resolveToken(environment: [String: String]) -> String? { + ProviderTokenResolver.zenmuxToken(environment: environment) + } +} diff --git a/Sources/CodexBarCore/Providers/Zenmux/ZenmuxSettingsReader.swift b/Sources/CodexBarCore/Providers/Zenmux/ZenmuxSettingsReader.swift new file mode 100644 index 000000000..ce7a037b1 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Zenmux/ZenmuxSettingsReader.swift @@ -0,0 +1,35 @@ +import Foundation + +public struct ZenmuxSettingsReader: Sendable { + public static let apiKeyEnvironmentKeys = [ + "ZENMUX_API_KEY", + ] + + public static func apiKey( + environment: [String: String] = ProcessInfo.processInfo.environment) -> String? + { + for key in self.apiKeyEnvironmentKeys { + guard let raw = environment[key]?.trimmingCharacters(in: .whitespacesAndNewlines), + !raw.isEmpty + else { + continue + } + let cleaned = Self.cleaned(raw) + if !cleaned.isEmpty { + return cleaned + } + } + return nil + } + + private static func cleaned(_ raw: String) -> String { + var value = raw + if (value.hasPrefix("\"") && value.hasSuffix("\"")) || + (value.hasPrefix("'") && value.hasSuffix("'")) + { + value.removeFirst() + value.removeLast() + } + return value.trimmingCharacters(in: .whitespacesAndNewlines) + } +} diff --git a/Sources/CodexBarCore/Providers/Zenmux/ZenmuxUsageFetcher.swift b/Sources/CodexBarCore/Providers/Zenmux/ZenmuxUsageFetcher.swift new file mode 100644 index 000000000..01ba69f1c --- /dev/null +++ b/Sources/CodexBarCore/Providers/Zenmux/ZenmuxUsageFetcher.swift @@ -0,0 +1,265 @@ +import Foundation +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif + +public struct ZenmuxUsageSnapshot: Sendable { + public let remainingRequests: Int + public let limitRequests: Int + public let resetTime: Date? + public let updatedAt: Date + public let apiKeyValid: Bool + public let totalTokens: Int? + public let apiKey: String? + + public init( + remainingRequests: Int, + limitRequests: Int, + resetTime: Date?, + updatedAt: Date, + apiKeyValid: Bool = false, + totalTokens: Int? = nil, + apiKey: String? = nil) + { + self.remainingRequests = remainingRequests + self.limitRequests = limitRequests + self.resetTime = resetTime + self.updatedAt = updatedAt + self.apiKeyValid = apiKeyValid + self.totalTokens = totalTokens + self.apiKey = apiKey + } + + private static func maskedKey(_ key: String?) -> String? { + guard let key, !key.isEmpty else { return nil } + if key.count <= 8 { return "****" } + let prefix = String(key.prefix(6)) + let suffix = String(key.suffix(4)) + return "\(prefix)...\(suffix)" + } + + public func toUsageSnapshot() -> UsageSnapshot { + let usedPercent: Double + let resetDescription: String + + if self.limitRequests > 0 { + let used = max(0, self.limitRequests - self.remainingRequests) + usedPercent = min(100, max(0, Double(used) / Double(self.limitRequests) * 100)) + resetDescription = "\(used)/\(self.limitRequests) requests" + } else if self.apiKeyValid { + usedPercent = 0 + resetDescription = "Active — check dashboard for details" + } else { + usedPercent = 0 + resetDescription = "No usage data" + } + + let primary = RateWindow( + usedPercent: usedPercent, + windowMinutes: nil, + resetsAt: self.resetTime, + resetDescription: resetDescription) + + let identity = ProviderIdentitySnapshot( + providerID: .zenmux, + accountEmail: Self.maskedKey(self.apiKey), + accountOrganization: nil, + loginMethod: "API") + + return UsageSnapshot( + primary: primary, + secondary: nil, + tertiary: nil, + providerCost: nil, + updatedAt: self.updatedAt, + identity: identity) + } +} + +public enum ZenmuxUsageError: LocalizedError, Sendable { + case missingCredentials + case networkError(String) + case apiError(Int, String) + case parseFailed(String) + + public var errorDescription: String? { + switch self { + case .missingCredentials: + "Missing Zenmux API key (ZENMUX_API_KEY)." + case let .networkError(message): + "Zenmux network error: \(message)" + case let .apiError(code, message): + "Zenmux API error (\(code)): \(message)" + case let .parseFailed(message): + "Failed to parse Zenmux response: \(message)" + } + } +} + +public struct ZenmuxUsageFetcher: Sendable { + private static let log = CodexBarLog.logger(LogCategories.zenmuxUsage) + + private static let apiURL = URL(string: "https://zenmux.ai/api/v1/chat/completions")! + + public static func fetchUsage(apiKey: String) async throws -> ZenmuxUsageSnapshot { + guard !apiKey.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty else { + throw ZenmuxUsageError.missingCredentials + } + + var request = URLRequest(url: self.apiURL) + request.httpMethod = "POST" + request.timeoutInterval = 15 + request.setValue("application/json", forHTTPHeaderField: "Content-Type") + request.setValue("application/json", forHTTPHeaderField: "Accept") + request.setValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization") + + let body: [String: Any] = [ + "model": "kuaishou/kat-coder-pro-v1-free", + "max_tokens": 1, + "messages": [ + ["role": "user", "content": "hi"], + ] as [[String: Any]], + ] + + request.httpBody = try JSONSerialization.data(withJSONObject: body) + + let (data, response) = try await URLSession.shared.data(for: request) + + guard let httpResponse = response as? HTTPURLResponse else { + throw ZenmuxUsageError.networkError("Invalid response") + } + + // Accept both 200 (success) and 429 (rate limited) – both carry rate limit headers. + guard httpResponse.statusCode == 200 || httpResponse.statusCode == 429 else { + let summary = Self.apiErrorSummary(statusCode: httpResponse.statusCode, data: data) + Self.log.error("Zenmux API returned \(httpResponse.statusCode): \(summary)") + throw ZenmuxUsageError.apiError(httpResponse.statusCode, summary) + } + + let headers = httpResponse.allHeaderFields + let remaining = Self.intHeader(headers, "x-ratelimit-remaining-requests") + let limit = Self.intHeader(headers, "x-ratelimit-limit-requests") + let resetString = headers["x-ratelimit-reset-requests"] as? String + + let resetTime: Date? = resetString.flatMap(Self.parseResetTime) + + // Parse token usage from response body when rate limit headers are absent + var totalTokens: Int? + if remaining == nil, limit == nil, + let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any], + let usage = json["usage"] as? [String: Any] + { + totalTokens = usage["total_tokens"] as? Int + } + + let snapshot = ZenmuxUsageSnapshot( + remainingRequests: remaining ?? 0, + limitRequests: limit ?? 0, + resetTime: resetTime, + updatedAt: Date(), + apiKeyValid: httpResponse.statusCode == 200, + totalTokens: totalTokens, + apiKey: apiKey) + + Self.log.debug( + "Zenmux usage parsed remaining=\(snapshot.remainingRequests) limit=\(snapshot.limitRequests) valid=\(snapshot.apiKeyValid)") // swiftlint:disable:this line_length + + return snapshot + } + + private static func intHeader(_ headers: [AnyHashable: Any], _ name: String) -> Int? { + if let value = headers[name] as? String, let int = Int(value) { + return int + } + if let value = headers[name.lowercased()] as? String, let int = Int(value) { + return int + } + // Case-insensitive search + for (key, val) in headers { + if let keyStr = key as? String, + keyStr.lowercased() == name.lowercased(), + let valStr = val as? String, + let int = Int(valStr) + { + return int + } + } + return nil + } + + /// Parse reset time from header value like "1d2h3m4s" or "30s" or ISO 8601. + private static func parseResetTime(_ value: String) -> Date? { + let trimmed = value.trimmingCharacters(in: .whitespacesAndNewlines) + if trimmed.isEmpty { return nil } + + // Try ISO 8601 first + let isoFormatter = ISO8601DateFormatter() + isoFormatter.formatOptions = [.withInternetDateTime, .withFractionalSeconds] + if let date = isoFormatter.date(from: trimmed) { return date } + let isoFallback = ISO8601DateFormatter() + isoFallback.formatOptions = [.withInternetDateTime] + if let date = isoFallback.date(from: trimmed) { return date } + + // Try duration format like "1d2h3m4s" or "30s" + var seconds: TimeInterval = 0 + let pattern = /(\d+)([dhms])/ + for match in trimmed.matches(of: pattern) { + guard let num = Double(match.1) else { continue } + switch match.2 { + case "d": seconds += num * 86400 + case "h": seconds += num * 3600 + case "m": seconds += num * 60 + case "s": seconds += num + default: break + } + } + if seconds > 0 { + return Date().addingTimeInterval(seconds) + } + + // Try plain integer as seconds + if let secs = TimeInterval(trimmed) { + return Date().addingTimeInterval(secs) + } + + return nil + } + + private static func apiErrorSummary(statusCode: Int, data: Data) -> String { + guard let root = try? JSONSerialization.jsonObject(with: data), + let json = root as? [String: Any] + else { + if let text = String(data: data, encoding: .utf8)? + .trimmingCharacters(in: .whitespacesAndNewlines), + !text.isEmpty + { + return self.compactText(text) + } + return "Unexpected response body (\(data.count) bytes)." + } + + if let message = json["message"] as? String { + let trimmed = message.trimmingCharacters(in: .whitespacesAndNewlines) + if !trimmed.isEmpty { return Self.compactText(trimmed) } + } + + if let error = json["error"] as? [String: Any], + let message = error["message"] as? String + { + let trimmed = message.trimmingCharacters(in: .whitespacesAndNewlines) + if !trimmed.isEmpty { return Self.compactText(trimmed) } + } + + return "HTTP \(statusCode) (\(data.count) bytes)." + } + + private static func compactText(_ text: String, maxLength: Int = 200) -> String { + let collapsed = text + .components(separatedBy: .newlines) + .joined(separator: " ") + .trimmingCharacters(in: .whitespacesAndNewlines) + if collapsed.count <= maxLength { return collapsed } + let limitIndex = collapsed.index(collapsed.startIndex, offsetBy: maxLength) + return "\(collapsed[../.jsonl), so the + // root mtime never changes when new data is written. Always enumerate the full tree; the + // per-file mtime/size check in processClaudeFile prevents redundant parsing. + let rootAttrs = (try? FileManager.default.attributesOfItem(atPath: canonicalRootPath)) ?? [:] let rootMtime = (rootAttrs[.modificationDate] as? Date)?.timeIntervalSince1970 ?? 0 let rootMtimeMs = Int64(rootMtime * 1000) - let cachedRootMtime = rootCandidates.compactMap { state.rootCache[$0] }.first - let canSkipEnumeration = cachedRootMtime == rootMtimeMs && rootMtimeMs > 0 - - if canSkipEnumeration { - let cachedPaths = state.cache.files.keys.filter { path in - prefixes.contains(where: { path.hasPrefix($0) }) - } - for path in cachedPaths { - guard FileManager.default.fileExists(atPath: path) else { - if let old = state.cache.files[path] { - Self.applyFileDays(cache: &state.cache, fileDays: old.days, sign: -1) - } - state.cache.files.removeValue(forKey: path) - continue - } - let attrs = (try? FileManager.default.attributesOfItem(atPath: path)) ?? [:] - let size = (attrs[.size] as? NSNumber)?.int64Value ?? 0 - if size <= 0 { continue } - let mtime = (attrs[.modificationDate] as? Date)?.timeIntervalSince1970 ?? 0 - let mtimeMs = Int64(mtime * 1000) - Self.processClaudeFile( - url: URL(fileURLWithPath: path), - size: size, - mtimeMs: mtimeMs, - state: state) - } - return - } let keys: [URLResourceKey] = [ .isRegularFileKey, diff --git a/Sources/CodexBarCore/Vendored/CostUsage/CostUsageScanner.swift b/Sources/CodexBarCore/Vendored/CostUsage/CostUsageScanner.swift index a5ef942b5..366a4b091 100644 --- a/Sources/CodexBarCore/Vendored/CostUsage/CostUsageScanner.swift +++ b/Sources/CodexBarCore/Vendored/CostUsage/CostUsageScanner.swift @@ -71,7 +71,8 @@ enum CostUsageScanner { } return self.loadClaudeDaily(provider: .vertexai, range: range, now: now, options: filtered) case .zai, .gemini, .antigravity, .cursor, .opencode, .factory, .copilot, .minimax, .kilo, .kiro, .kimi, - .kimik2, .augment, .jetbrains, .amp, .ollama, .synthetic, .openrouter, .warp: + .kimik2, .augment, .jetbrains, .amp, .ollama, .synthetic, .openrouter, .warp, .qwen, .doubao, .zenmux, + .aigocode, .trae, .stepfun, .mimo: return emptyReport } } diff --git a/Sources/CodexBarWidget/CodexBarWidgetProvider.swift b/Sources/CodexBarWidget/CodexBarWidgetProvider.swift index eb0d00574..a487eec8b 100644 --- a/Sources/CodexBarWidget/CodexBarWidgetProvider.swift +++ b/Sources/CodexBarWidget/CodexBarWidgetProvider.swift @@ -67,6 +67,13 @@ enum ProviderChoice: String, AppEnum { case .synthetic: return nil // Synthetic not yet supported in widgets case .openrouter: return nil // OpenRouter not yet supported in widgets case .warp: return nil // Warp not yet supported in widgets + case .qwen: return nil // Qwen not yet supported in widgets + case .doubao: return nil // Doubao not yet supported in widgets + case .zenmux: return nil // Zenmux not yet supported in widgets + case .aigocode: return nil // AigoCode not yet supported in widgets + case .trae: return nil // Trae not yet supported in widgets + case .stepfun: return nil // StepFun not yet supported in widgets + case .mimo: return nil // MiMo not yet supported in widgets } } } diff --git a/Sources/CodexBarWidget/CodexBarWidgetViews.swift b/Sources/CodexBarWidget/CodexBarWidgetViews.swift index fbb8c5d9c..84262001c 100644 --- a/Sources/CodexBarWidget/CodexBarWidgetViews.swift +++ b/Sources/CodexBarWidget/CodexBarWidgetViews.swift @@ -279,6 +279,13 @@ private struct ProviderSwitchChip: View { case .synthetic: "Synthetic" case .openrouter: "OpenRouter" case .warp: "Warp" + case .qwen: "Qwen" + case .doubao: "Doubao" + case .zenmux: "Zenmux" + case .aigocode: "AigoCode" + case .trae: "Trae" + case .stepfun: "StepFun" + case .mimo: "MiMo" } } } @@ -618,6 +625,20 @@ enum WidgetColors { Color(red: 111 / 255, green: 66 / 255, blue: 193 / 255) // OpenRouter purple case .warp: Color(red: 147 / 255, green: 139 / 255, blue: 180 / 255) + case .qwen: + Color(red: 97 / 255, green: 71 / 255, blue: 232 / 255) // Qwen purple + case .doubao: + Color(red: 50 / 255, green: 108 / 255, blue: 229 / 255) // Doubao blue + case .zenmux: + Color(red: 255 / 255, green: 140 / 255, blue: 0 / 255) // Zenmux orange-red + case .aigocode: + Color(red: 34 / 255, green: 197 / 255, blue: 94 / 255) // AigoCode green + case .trae: + Color(red: 59 / 255, green: 130 / 255, blue: 246 / 255) // Trae blue + case .stepfun: + Color(red: 0 / 255, green: 168 / 255, blue: 107 / 255) // StepFun green + case .mimo: + Color(red: 255 / 255, green: 103 / 255, blue: 0 / 255) // MiMo Xiaomi orange } } } diff --git a/Tests/CodexBarTests/AntigravityStatusProbeTests.swift b/Tests/CodexBarTests/AntigravityStatusProbeTests.swift index d69a3041f..16b160bcb 100644 --- a/Tests/CodexBarTests/AntigravityStatusProbeTests.swift +++ b/Tests/CodexBarTests/AntigravityStatusProbeTests.swift @@ -49,8 +49,9 @@ struct AntigravityStatusProbeTests { guard let primary = usage.primary else { return } - #expect(primary.remainingPercent.rounded() == 50) - #expect(usage.secondary?.remainingPercent.rounded() == 80) + // Gemini Pro is primary, Claude is secondary, Gemini Flash is tertiary + #expect(primary.remainingPercent.rounded() == 80) + #expect(usage.secondary?.remainingPercent.rounded() == 50) #expect(usage.tertiary?.remainingPercent.rounded() == 20) } } diff --git a/Tests/CodexBarTests/GeminiOAuthSymlinkTests.swift b/Tests/CodexBarTests/GeminiOAuthSymlinkTests.swift new file mode 100644 index 000000000..860cc2af0 --- /dev/null +++ b/Tests/CodexBarTests/GeminiOAuthSymlinkTests.swift @@ -0,0 +1,175 @@ +import Foundation +import Testing +@testable import CodexBarCore + +/// Regression tests for multi-level symlink resolution when locating +/// Gemini CLI's oauth2.js file. +/// +/// See: https://github.com/steipete/CodexBar/pull/497 +/// +/// On some setups (e.g. Apple Silicon with `/usr/local/bin/gemini` → +/// `/opt/homebrew/bin/gemini` → `../Cellar/…/bin/gemini` → …), the +/// previous single-level `destinationOfSymbolicLink` call resolved too +/// shallowly, producing a wrong base path and silently failing to find +/// the OAuth credentials file. +@Suite +struct GeminiOAuthSymlinkTests { + /// Sample oauth2.js content used by all tests. + private static let sampleOAuth2JS = """ + const OAUTH_CLIENT_ID = 'test-client-id.apps.googleusercontent.com'; + const OAUTH_CLIENT_SECRET = 'test-client-secret'; + """ + + // MARK: - Helpers + + /// Build a temporary directory tree that mimics a Homebrew Cellar layout + /// and return the path to the top-level "entry" symlink. + /// + /// Layout created: + /// ``` + /// / + /// usr/local/bin/gemini → /opt/homebrew/bin/gemini (level 0 — optional) + /// opt/homebrew/bin/gemini → ../Cellar/gemini-cli/0.32.1/bin/gemini (level 1) + /// opt/homebrew/Cellar/gemini-cli/0.32.1/ + /// bin/gemini → ../libexec/bin/gemini (level 2) + /// libexec/bin/gemini (regular file — final target) + /// libexec/lib/node_modules/@google/gemini-cli/ + /// node_modules/@google/gemini-cli-core/dist/src/code_assist/oauth2.js + /// ``` + private static func makeBrewLayout( + root: URL, + includeExtraSymlink: Bool) throws -> String + { + let fm = FileManager.default + + // Cellar paths + let cellarBase = root.appendingPathComponent( + "opt/homebrew/Cellar/gemini-cli/0.32.1", isDirectory: true) + let cellarBin = cellarBase.appendingPathComponent("bin", isDirectory: true) + let libexecBin = cellarBase.appendingPathComponent("libexec/bin", isDirectory: true) + let oauthDir = cellarBase.appendingPathComponent( + "libexec/lib/node_modules/@google/gemini-cli/node_modules/@google/gemini-cli-core/dist/src/code_assist", + isDirectory: true) + + // Create directories + try fm.createDirectory(at: cellarBin, withIntermediateDirectories: true) + try fm.createDirectory(at: libexecBin, withIntermediateDirectories: true) + try fm.createDirectory(at: oauthDir, withIntermediateDirectories: true) + + // Write the oauth2.js file + let oauthFile = oauthDir.appendingPathComponent("oauth2.js") + try Self.sampleOAuth2JS.write(to: oauthFile, atomically: true, encoding: .utf8) + + // Final target: libexec/bin/gemini (regular file) + let finalBinary = libexecBin.appendingPathComponent("gemini") + try "#!/usr/bin/env node\n".write(to: finalBinary, atomically: true, encoding: .utf8) + + // Level 2 symlink: Cellar/…/bin/gemini → ../libexec/bin/gemini + let cellarBinGemini = cellarBin.appendingPathComponent("gemini") + try fm.createSymbolicLink( + atPath: cellarBinGemini.path, + withDestinationPath: "../libexec/bin/gemini") + + // Level 1 symlink: opt/homebrew/bin/gemini → ../Cellar/gemini-cli/0.32.1/bin/gemini + let homebrewBin = root.appendingPathComponent("opt/homebrew/bin", isDirectory: true) + try fm.createDirectory(at: homebrewBin, withIntermediateDirectories: true) + let homebrewBinGemini = homebrewBin.appendingPathComponent("gemini") + try fm.createSymbolicLink( + atPath: homebrewBinGemini.path, + withDestinationPath: "../Cellar/gemini-cli/0.32.1/bin/gemini") + + if includeExtraSymlink { + // Level 0 symlink: usr/local/bin/gemini → /opt/homebrew/bin/gemini + let usrLocalBin = root.appendingPathComponent("usr/local/bin", isDirectory: true) + try fm.createDirectory(at: usrLocalBin, withIntermediateDirectories: true) + let usrLocalBinGemini = usrLocalBin.appendingPathComponent("gemini") + try fm.createSymbolicLink( + atPath: usrLocalBinGemini.path, + withDestinationPath: homebrewBinGemini.path) + return usrLocalBinGemini.path + } else { + return homebrewBinGemini.path + } + } + + // MARK: - Tests + + /// Standard 2-level Homebrew symlink chain (no extra symlink). + /// The old code handled this correctly; this test guards against regressions. + @Test + func findsOAuthWithStandardHomebrewChain() throws { + let root = FileManager.default.temporaryDirectory + .appendingPathComponent("GeminiSymlinkTest-standard-\(UUID().uuidString)") + defer { try? FileManager.default.removeItem(at: root) } + + let entryPath = try Self.makeBrewLayout(root: root, includeExtraSymlink: false) + let content = GeminiStatusProbe.resolveOAuthFileContent(from: entryPath) + + #expect(content != nil, "Should find oauth2.js through standard 2-level Homebrew symlink chain") + #expect(content?.contains("OAUTH_CLIENT_ID") == true) + #expect(content?.contains("test-client-id") == true) + } + + /// Multi-level symlink chain with an extra `/usr/local/bin` symlink. + /// This is the scenario that caused the original bug — the old single-level + /// `destinationOfSymbolicLink` resolved to `/opt/homebrew/bin/gemini` and + /// computed `baseDir = /opt/homebrew`, missing the Cellar path entirely. + @Test + func findsOAuthWithExtraSymlinkLevel() throws { + let root = FileManager.default.temporaryDirectory + .appendingPathComponent("GeminiSymlinkTest-extra-\(UUID().uuidString)") + defer { try? FileManager.default.removeItem(at: root) } + + let entryPath = try Self.makeBrewLayout(root: root, includeExtraSymlink: true) + let content = GeminiStatusProbe.resolveOAuthFileContent(from: entryPath) + + #expect(content != nil, "Should find oauth2.js through 3-level symlink chain (the bug scenario)") + #expect(content?.contains("OAUTH_CLIENT_ID") == true) + #expect(content?.contains("test-client-id") == true) + } + + /// When the binary path is not a symlink at all (e.g. direct npm install), + /// the resolver should still search relative to that path. + @Test + func handlesNonSymlinkBinary() throws { + let root = FileManager.default.temporaryDirectory + .appendingPathComponent("GeminiSymlinkTest-direct-\(UUID().uuidString)") + let fm = FileManager.default + defer { try? fm.removeItem(at: root) } + + // Create: root/bin/gemini (regular file) + root/lib/node_modules/…/oauth2.js + let binDir = root.appendingPathComponent("bin", isDirectory: true) + try fm.createDirectory(at: binDir, withIntermediateDirectories: true) + let binary = binDir.appendingPathComponent("gemini") + try "#!/usr/bin/env node\n".write(to: binary, atomically: true, encoding: .utf8) + + let oauthDir = root.appendingPathComponent( + "lib/node_modules/@google/gemini-cli/node_modules/@google/gemini-cli-core/dist/src/code_assist", + isDirectory: true) + try fm.createDirectory(at: oauthDir, withIntermediateDirectories: true) + try Self.sampleOAuth2JS.write( + to: oauthDir.appendingPathComponent("oauth2.js"), + atomically: true, + encoding: .utf8) + + let content = GeminiStatusProbe.resolveOAuthFileContent(from: binary.path) + #expect(content != nil, "Should find oauth2.js from non-symlinked binary") + } + + /// Returns nil gracefully when no oauth2.js exists anywhere in the chain. + @Test + func returnsNilWhenOAuthFileMissing() throws { + let root = FileManager.default.temporaryDirectory + .appendingPathComponent("GeminiSymlinkTest-missing-\(UUID().uuidString)") + let fm = FileManager.default + defer { try? fm.removeItem(at: root) } + + let binDir = root.appendingPathComponent("bin", isDirectory: true) + try fm.createDirectory(at: binDir, withIntermediateDirectories: true) + let binary = binDir.appendingPathComponent("gemini") + try "#!/usr/bin/env node\n".write(to: binary, atomically: true, encoding: .utf8) + + let content = GeminiStatusProbe.resolveOAuthFileContent(from: binary.path) + #expect(content == nil, "Should return nil when oauth2.js does not exist") + } +} diff --git a/Tests/CodexBarTests/MiMoProviderTests.swift b/Tests/CodexBarTests/MiMoProviderTests.swift new file mode 100644 index 000000000..3da25071d --- /dev/null +++ b/Tests/CodexBarTests/MiMoProviderTests.swift @@ -0,0 +1,684 @@ +import Foundation +import SwiftUI +import Testing +@testable import CodexBar +@testable import CodexBarCore +#if os(macOS) +import SweetCookieKit +#endif + +@Suite(.serialized) +struct MiMoProviderTests { + private struct StubClaudeFetcher: ClaudeUsageFetching { + func loadLatestUsage(model _: String) async throws -> ClaudeUsageSnapshot { + throw ClaudeUsageError.parseFailed("stub") + } + + func debugRawProbe(model _: String) async -> String { + "stub" + } + + func detectVersion() -> String? { + nil + } + } + + @Test + func `cookie header normalizer keeps required mimo cookies`() { + let raw = """ + curl 'https://platform.xiaomimimo.com/api/v1/balance' \ + -H 'Cookie: userId=123; api-platform_serviceToken=svc-token; ignored=value; api-platform_ph=ph-token' + """ + + let normalized = MiMoCookieHeader.normalizedHeader(from: raw) + + #expect(normalized == "api-platform_ph=ph-token; api-platform_serviceToken=svc-token; userId=123") + } + + @Test + func `cookie header normalizer rejects missing auth cookies`() { + let normalized = MiMoCookieHeader.normalizedHeader(from: "Cookie: userId=123") + + #expect(normalized == nil) + } + + @Test + func `cookie header builder keeps mimo auth cookies from one scope`() throws { + let cookies = try [ + self.makeCookie( + name: "userId", + value: "root-user", + domain: "xiaomimimo.com", + expiresAt: Date(timeIntervalSince1970: 1_800_000_000)), + self.makeCookie( + name: "api-platform_serviceToken", + value: "platform-token", + domain: "platform.xiaomimimo.com", + expiresAt: Date(timeIntervalSince1970: 1_900_000_000)), + self.makeCookie( + name: "userId", + value: "platform-user", + domain: "platform.xiaomimimo.com", + expiresAt: Date(timeIntervalSince1970: 1_900_000_000)), + self.makeCookie( + name: "api-platform_ph", + value: "platform-ph", + domain: "platform.xiaomimimo.com", + expiresAt: Date(timeIntervalSince1970: 1_900_000_000)), + ] + + let header = MiMoCookieHeader.header(from: cookies) + + #expect(header == "api-platform_ph=platform-ph; api-platform_serviceToken=platform-token; userId=platform-user") + } + + @Test + func `cookie header builder prefers more specific matching cookie`() throws { + let cookies = try [ + self.makeCookie( + name: "userId", + value: "root-user", + domain: "xiaomimimo.com", + expiresAt: Date(timeIntervalSince1970: 1_900_000_000)), + self.makeCookie( + name: "userId", + value: "api-user", + domain: "platform.xiaomimimo.com", + path: "/api", + expiresAt: Date(timeIntervalSince1970: 1_800_000_000)), + self.makeCookie( + name: "api-platform_serviceToken", + value: "platform-token", + domain: ".xiaomimimo.com", + expiresAt: Date(timeIntervalSince1970: 1_900_000_000)), + self.makeCookie( + name: "irrelevant", + value: "ignored", + domain: "platform.xiaomimimo.com", + expiresAt: Date(timeIntervalSince1970: 1_900_000_000)), + ] + + let header = MiMoCookieHeader.header(from: cookies) + + #expect(header == "api-platform_serviceToken=platform-token; userId=api-user") + } + + @Test + func `cookie header builder rejects partial path prefix matches`() throws { + let cookies = try [ + self.makeCookie( + name: "userId", + value: "partial-path-user", + domain: "platform.xiaomimimo.com", + path: "/api/v1/bal", + expiresAt: Date(timeIntervalSince1970: 1_900_000_000)), + self.makeCookie( + name: "userId", + value: "valid-user", + domain: "platform.xiaomimimo.com", + path: "/api", + expiresAt: Date(timeIntervalSince1970: 1_800_000_000)), + self.makeCookie( + name: "api-platform_serviceToken", + value: "partial-path-token", + domain: "platform.xiaomimimo.com", + path: "/api/v1/bal", + expiresAt: Date(timeIntervalSince1970: 1_900_000_000)), + self.makeCookie( + name: "api-platform_serviceToken", + value: "valid-token", + domain: "platform.xiaomimimo.com", + path: "/api", + expiresAt: Date(timeIntervalSince1970: 1_800_000_000)), + ] + + let header = MiMoCookieHeader.header(from: cookies) + + #expect(header == "api-platform_serviceToken=valid-token; userId=valid-user") + } + + @Test + func `cookie header builder accepts slash terminated path prefixes`() throws { + let cookies = try [ + self.makeCookie( + name: "userId", + value: "slash-user", + domain: "platform.xiaomimimo.com", + path: "/api/", + expiresAt: Date(timeIntervalSince1970: 1_900_000_000)), + self.makeCookie( + name: "api-platform_serviceToken", + value: "slash-token", + domain: "platform.xiaomimimo.com", + path: "/api/", + expiresAt: Date(timeIntervalSince1970: 1_900_000_000)), + ] + + let header = MiMoCookieHeader.header(from: cookies) + + #expect(header == "api-platform_serviceToken=slash-token; userId=slash-user") + } + + @Test + func `usage snapshot exposes balance through identity plan text`() { + let snapshot = MiMoUsageSnapshot( + balance: 25.51, + currency: "USD", + updatedAt: Date(timeIntervalSince1970: 1_742_771_200)) + + let usage = snapshot.toUsageSnapshot() + + #expect(usage.primary == nil) + #expect(usage.secondary == nil) + #expect(usage.loginMethod(for: .mimo) == "Balance: $25.51") + } + + @Test + func `usage snapshot shows token plan as primary when available`() { + let resetDate = Date(timeIntervalSince1970: 1_778_025_599) + let snapshot = MiMoUsageSnapshot( + balance: 25.51, + currency: "USD", + planCode: "standard", + planPeriodEnd: resetDate, + planExpired: false, + tokenUsed: 10_100_158, + tokenLimit: 200_000_000, + tokenPercent: 0.0505, + updatedAt: Date(timeIntervalSince1970: 1_742_771_200)) + + let usage = snapshot.toUsageSnapshot() + + #expect(usage.primary != nil) + #expect(abs((usage.primary?.usedPercent ?? 0) - 5.05) < 0.0001) + #expect(usage.primary?.resetDescription == "10,100,158 / 200,000,000 Credits") + #expect(usage.primary?.resetsAt == resetDate) + #expect(usage.loginMethod(for: .mimo) == "Standard") + } + + @Test + func `usage snapshot falls back to balance when no token plan`() { + let snapshot = MiMoUsageSnapshot( + balance: 0, + currency: "USD", + planCode: nil, + planPeriodEnd: nil, + planExpired: false, + tokenUsed: 0, + tokenLimit: 0, + tokenPercent: 0, + updatedAt: Date(timeIntervalSince1970: 1_742_771_200)) + + let usage = snapshot.toUsageSnapshot() + + #expect(usage.primary == nil) + #expect(usage.loginMethod(for: .mimo) == "Balance: $0.00") + } + + @Test + func `parses balance payload`() throws { + let now = Date(timeIntervalSince1970: 1_742_771_200) + let json = """ + { + "code": 0, + "message": "", + "data": { + "balance": "25.51", + "frozenBalance": null, + "currency": "USD", + "overdraftLimit": null + } + } + """ + + let snapshot = try MiMoUsageFetcher.parseUsageSnapshot(from: Data(json.utf8), now: now) + + #expect(snapshot.balance == 25.51) + #expect(snapshot.currency == "USD") + #expect(snapshot.updatedAt == now) + } + + @Test + func `parses token plan detail payload`() throws { + let json = """ + { + "code": 0, + "message": "", + "data": { + "planCode": "standard", + "currentPeriodEnd": "2026-05-04 23:59:59", + "expired": false + } + } + """ + + let detail = try MiMoUsageFetcher.parseTokenPlanDetail(from: Data(json.utf8)) + + #expect(detail.planCode == "standard") + #expect(detail.expired == false) + #expect(detail.periodEnd != nil) + } + + @Test + func `parses token plan usage payload`() throws { + let json = """ + { + "code": 0, + "message": "", + "data": { + "monthUsage": { + "percent": 0.0505, + "items": [ + { + "name": "month_total_token", + "used": 10100158, + "limit": 200000000, + "percent": 0.0505 + } + ] + } + } + } + """ + + let usage = try MiMoUsageFetcher.parseTokenPlanUsage(from: Data(json.utf8)) + + #expect(usage.used == 10_100_158) + #expect(usage.limit == 200_000_000) + #expect(usage.percent == 0.0505) + } + + @Test + func `combined snapshot merges balance and token plan`() throws { + let now = Date(timeIntervalSince1970: 1_742_771_200) + let balanceJSON = """ + {"code":0,"message":"","data":{"balance":"25.51","currency":"USD"}} + """ + let detailJSON = """ + {"code":0,"message":"","data":{"planCode":"standard","currentPeriodEnd":"2026-05-04 23:59:59","expired":false}} + """ + // swiftlint:disable line_length + let usageJSON = """ + {"code":0,"message":"","data":{"monthUsage":{"percent":0.0505,"items":[{"name":"month_total_token","used":10100158,"limit":200000000,"percent":0.0505}]}}} + """ + // swiftlint:enable line_length + + let snapshot = try MiMoUsageFetcher.parseCombinedSnapshot( + balanceData: Data(balanceJSON.utf8), + tokenDetailData: Data(detailJSON.utf8), + tokenUsageData: Data(usageJSON.utf8), + now: now) + + #expect(snapshot.balance == 25.51) + #expect(snapshot.currency == "USD") + #expect(snapshot.planCode == "standard") + #expect(snapshot.tokenUsed == 10_100_158) + #expect(snapshot.tokenLimit == 200_000_000) + #expect(snapshot.tokenPercent == 0.0505) + } + + @Test(.disabled("Fetcher hits tokenPlan/* endpoints too; test from PR 651 assumes single balance endpoint")) + func `fetch usage hits mimo balance endpoint with browser headers`() async throws { + let registered = URLProtocol.registerClass(MiMoStubURLProtocol.self) + defer { + if registered { + URLProtocol.unregisterClass(MiMoStubURLProtocol.self) + } + MiMoStubURLProtocol.handler = nil + } + + MiMoStubURLProtocol.handler = { request in + guard let url = request.url else { throw URLError(.badURL) } + #expect(url.path == "/api/v1/balance") + #expect(request.value(forHTTPHeaderField: "Cookie") == "api-platform_serviceToken=svc-token; userId=123") + #expect(request.value(forHTTPHeaderField: "Accept-Language") == "en-US,en;q=0.9") + #expect(request.value(forHTTPHeaderField: "x-timeZone") == "UTC+01:00") + #expect(request.value(forHTTPHeaderField: "Referer") == "https://platform.xiaomimimo.com/#/console/balance") + let body = """ + { + "code": 0, + "message": "", + "data": { + "balance": "25.51", + "currency": "USD" + } + } + """ + return Self.makeResponse(url: url, body: body) + } + + let snapshot = try await MiMoUsageFetcher.fetchUsage( + cookieHeader: "Cookie: userId=123; api-platform_serviceToken=svc-token", + environment: ["MIMO_API_URL": "https://mimo.test/api/v1"], + now: Date(timeIntervalSince1970: 1_742_771_200)) + + #expect(snapshot.balance == 25.51) + #expect(snapshot.currency == "USD") + } + + @Test + @MainActor + func `provider detail plan row formats mimo as balance`() { + let row = ProviderDetailView.planRow(provider: .mimo, planText: "Balance: $25.51") + + #expect(row?.label == "Balance") + #expect(row?.value == "$25.51") + } + + @Test(arguments: [UsageProvider.openrouter, .mimo]) + @MainActor + func `menu descriptor renders balance providers without duplicate prefix`(provider: UsageProvider) throws { + let suite = "MiMoProviderTests-menu-balance-\(provider.rawValue)" + let defaults = try #require(UserDefaults(suiteName: suite)) + defaults.removePersistentDomain(forName: suite) + + let settings = SettingsStore( + userDefaults: defaults, + configStore: testConfigStore(suiteName: suite), + zaiTokenStore: NoopZaiTokenStore(), + syntheticTokenStore: NoopSyntheticTokenStore()) + settings.statusChecksEnabled = false + + let store = UsageStore( + fetcher: UsageFetcher(environment: [:]), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings) + store._setSnapshotForTesting(self.makeBalanceSnapshot(provider: provider), provider: provider) + + let descriptor = MenuDescriptor.build( + provider: provider, + store: store, + settings: settings, + account: AccountInfo(email: nil, plan: nil), + updateReady: false, + includeContextualActions: false) + + let lines = descriptor.sections + .flatMap(\.entries) + .compactMap { entry -> String? in + guard case let .text(text, _) = entry else { return nil } + return text + } + + #expect(lines.contains("Balance: $25.51")) + #expect(!lines.contains("Balance: Balance: $25.51")) + } + + @Test + func `mimo web strategy unavailable when cookie source is off`() async { + CookieHeaderCache.store( + provider: .mimo, + cookieHeader: "api-platform_serviceToken=svc-token; userId=123", + sourceLabel: "cached") + defer { CookieHeaderCache.clear(provider: .mimo) } + + let strategy = MiMoWebFetchStrategy() + let context = self.makeContext(settings: ProviderSettingsSnapshot.make( + mimo: ProviderSettingsSnapshot.MiMoProviderSettings( + cookieSource: .off, + manualCookieHeader: nil))) + + let available = await strategy.isAvailable(context) + + #expect(available == false) + } + + @Test + func `mimo manual mode does not report available from cached browser session`() async { + CookieHeaderCache.store( + provider: .mimo, + cookieHeader: "api-platform_serviceToken=svc-token; userId=123", + sourceLabel: "cached") + defer { CookieHeaderCache.clear(provider: .mimo) } + + let strategy = MiMoWebFetchStrategy() + let context = self.makeContext(settings: ProviderSettingsSnapshot.make( + mimo: ProviderSettingsSnapshot.MiMoProviderSettings( + cookieSource: .manual, + manualCookieHeader: "Cookie: userId=123"))) + + let available = await strategy.isAvailable(context) + + #expect(available == false) + } + + @Test + func `mimo manual mode rejects invalid header instead of falling back to cached session`() async { + CookieHeaderCache.store( + provider: .mimo, + cookieHeader: "api-platform_serviceToken=svc-token; userId=123", + sourceLabel: "cached") + defer { CookieHeaderCache.clear(provider: .mimo) } + + let strategy = MiMoWebFetchStrategy() + let context = self.makeContext(settings: ProviderSettingsSnapshot.make( + mimo: ProviderSettingsSnapshot.MiMoProviderSettings( + cookieSource: .manual, + manualCookieHeader: "Cookie: userId=123"))) + + await #expect(throws: MiMoSettingsError.invalidCookie) { + _ = try await strategy.fetch(context) + } + } + + @Test(.disabled("Cookie retry state depends on MiMoUsageFetcher single-endpoint flow from PR 651")) + func `mimo web strategy retries imported sessions after decode failure`() async throws { + let registered = URLProtocol.registerClass(MiMoStubURLProtocol.self) + defer { + if registered { + URLProtocol.unregisterClass(MiMoStubURLProtocol.self) + } + MiMoStubURLProtocol.handler = nil + MiMoCookieImporter.importSessionsOverrideForTesting = nil + CookieHeaderCache.clear(provider: .mimo) + } + + CookieHeaderCache.clear(provider: .mimo) + CookieHeaderCache.store(provider: .mimo, cookieHeader: "invalid", sourceLabel: "invalid") + + MiMoCookieImporter.importSessionsOverrideForTesting = { _, _ in + [ + .init( + cookieHeader: "api-platform_serviceToken=expired-token; userId=111", + sourceLabel: "Expired Chrome"), + .init( + cookieHeader: "api-platform_serviceToken=valid-token; userId=222", + sourceLabel: "Active Chrome"), + ] + } + + var requestedCookies: [String] = [] + MiMoStubURLProtocol.handler = { request in + guard let url = request.url else { throw URLError(.badURL) } + let cookie = request.value(forHTTPHeaderField: "Cookie") ?? "" + requestedCookies.append(cookie) + + if cookie.contains("expired-token") { + let response = HTTPURLResponse( + url: url, + statusCode: 200, + httpVersion: "HTTP/1.1", + headerFields: ["Content-Type": "text/html"])! + return (response, Data("login".utf8)) + } + + let body = """ + { + "code": 0, + "message": "", + "data": { + "balance": "25.51", + "currency": "USD" + } + } + """ + return Self.makeResponse(url: url, body: body) + } + + let strategy = MiMoWebFetchStrategy() + let result = try await strategy + .fetch(self.makeContext(environment: ["MIMO_API_URL": "https://mimo.test/api/v1"])) + + #expect(requestedCookies.count == 2) + #expect(requestedCookies[0].contains("expired-token")) + #expect(requestedCookies[1].contains("valid-token")) + #expect(result.usage.loginMethod(for: .mimo) == "Balance: $25.51") + #expect(CookieHeaderCache.load(provider: .mimo)?.sourceLabel == "Active Chrome") + } + + #if os(macOS) + @Test + func `mimo importer merges profile stores before validating auth cookies`() { + let profile = BrowserProfile(id: "Default", name: "Default") + let primaryStore = BrowserCookieStore( + browser: .chrome, + profile: profile, + kind: .primary, + label: "Chrome Default", + databaseURL: nil) + let networkStore = BrowserCookieStore( + browser: .chrome, + profile: profile, + kind: .network, + label: "Chrome Default (Network)", + databaseURL: nil) + let expires = Date(timeIntervalSince1970: 1_900_000_000) + + let sessions = MiMoCookieImporter.sessionInfos(from: [ + BrowserCookieStoreRecords(store: primaryStore, records: [ + BrowserCookieRecord( + domain: "platform.xiaomimimo.com", + name: "userId", + path: "/", + value: "123", + expires: expires, + isSecure: true, + isHTTPOnly: false), + ]), + BrowserCookieStoreRecords(store: networkStore, records: [ + BrowserCookieRecord( + domain: "platform.xiaomimimo.com", + name: "api-platform_serviceToken", + path: "/", + value: "token", + expires: expires, + isSecure: true, + isHTTPOnly: true), + ]), + ]) + + #expect(sessions.count == 1) + #expect(sessions.first?.sourceLabel == "Chrome Default") + #expect(sessions.first?.cookieHeader == "api-platform_serviceToken=token; userId=123") + } + #endif + + private static func makeResponse( + url: URL, + body: String, + statusCode: Int = 200) -> (HTTPURLResponse, Data) + { + let response = HTTPURLResponse( + url: url, + statusCode: statusCode, + httpVersion: "HTTP/1.1", + headerFields: ["Content-Type": "application/json"])! + return (response, Data(body.utf8)) + } + + private func makeBalanceSnapshot(provider: UsageProvider) -> UsageSnapshot { + let updatedAt = Date(timeIntervalSince1970: 1_742_771_200) + switch provider { + case .openrouter: + return OpenRouterUsageSnapshot( + totalCredits: 50, + totalUsage: 24.49, + balance: 25.51, + usedPercent: 49, + keyDataFetched: false, + keyLimit: nil, + keyUsage: nil, + rateLimit: nil, + updatedAt: updatedAt).toUsageSnapshot() + case .mimo: + return MiMoUsageSnapshot( + balance: 25.51, + currency: "USD", + updatedAt: updatedAt).toUsageSnapshot() + default: + Issue.record("Unexpected provider \(provider.rawValue)") + return UsageSnapshot( + primary: nil, + secondary: nil, + tertiary: nil, + updatedAt: updatedAt) + } + } + + private func makeContext( + settings: ProviderSettingsSnapshot? = nil, + environment: [String: String] = [:]) -> ProviderFetchContext + { + let browserDetection = BrowserDetection(cacheTTL: 0) + return ProviderFetchContext( + runtime: .app, + sourceMode: .auto, + includeCredits: false, + webTimeout: 1, + webDebugDumpHTML: false, + verbose: false, + env: environment, + settings: settings, + fetcher: UsageFetcher(environment: [:]), + claudeFetcher: StubClaudeFetcher(), + browserDetection: browserDetection) + } + + private func makeCookie( + name: String, + value: String, + domain: String, + path: String = "/", + expiresAt: Date) throws -> HTTPCookie + { + let properties: [HTTPCookiePropertyKey: Any] = [ + .name: name, + .value: value, + .domain: domain, + .path: path, + .expires: expiresAt, + .secure: "TRUE", + ] + return try #require(HTTPCookie(properties: properties)) + } +} + +final class MiMoStubURLProtocol: URLProtocol { + nonisolated(unsafe) static var handler: ((URLRequest) throws -> (HTTPURLResponse, Data))? + + override static func canInit(with request: URLRequest) -> Bool { + request.url?.host == "mimo.test" + } + + override static func canonicalRequest(for request: URLRequest) -> URLRequest { + request + } + + override func startLoading() { + guard let handler = Self.handler else { + self.client?.urlProtocol(self, didFailWithError: URLError(.badServerResponse)) + return + } + + do { + let (response, data) = try handler(self.request) + self.client?.urlProtocol(self, didReceive: response, cacheStoragePolicy: .notAllowed) + self.client?.urlProtocol(self, didLoad: data) + self.client?.urlProtocolDidFinishLoading(self) + } catch { + self.client?.urlProtocol(self, didFailWithError: error) + } + } + + override func stopLoading() {} +} diff --git a/Tests/CodexBarTests/SettingsStoreTests.swift b/Tests/CodexBarTests/SettingsStoreTests.swift index 04ca55516..91b2e41ea 100644 --- a/Tests/CodexBarTests/SettingsStoreTests.swift +++ b/Tests/CodexBarTests/SettingsStoreTests.swift @@ -746,6 +746,13 @@ struct SettingsStoreTests { .synthetic, .warp, .openrouter, + .qwen, + .doubao, + .zenmux, + .aigocode, + .trae, + .stepfun, + .mimo, ]) // Move one provider; ensure it's persisted across instances. diff --git a/docs/mimo.md b/docs/mimo.md new file mode 100644 index 000000000..be6fb0636 --- /dev/null +++ b/docs/mimo.md @@ -0,0 +1,55 @@ +--- +summary: "Xiaomi MiMo provider notes: cookie auth, balance endpoint, and setup." +read_when: + - Adding or modifying the Xiaomi MiMo provider + - Debugging MiMo cookie import or balance fetching + - Explaining MiMo setup and limitations to users +--- + +# Xiaomi MiMo Provider + +The Xiaomi MiMo provider tracks your current balance from the Xiaomi MiMo console. + +## Features + +- **Balance display**: Shows the current MiMo balance as provider identity text. +- **Cookie-based auth**: Uses browser cookies or a pasted `Cookie:` header. +- **Near-real-time updates**: Balance usually reflects within a few minutes. + +## Setup + +1. Open **Settings → Providers** +2. Enable **Xiaomi MiMo** +3. Leave **Cookie source** on **Auto** (recommended) + +### Manual cookie import (optional) + +1. Open `https://platform.xiaomimimo.com/#/console/balance` +2. Copy a `Cookie:` header from your browser’s Network tab +3. Paste it into **Xiaomi MiMo → Cookie source → Manual** + +## How it works + +- Fetches `GET https://platform.xiaomimimo.com/api/v1/balance` +- Requires the `api-platform_serviceToken` and `userId` cookies +- Accepts optional MiMo cookies like `api-platform_ph` and `api-platform_slh` when present +- Supports `MIMO_API_URL` to override the base API URL for testing + +## Limitations + +- MiMo currently exposes **balance only** +- Token cost, status polling, debug log output, and widgets are not supported yet + +## Troubleshooting + +### “No Xiaomi MiMo browser session found” + +Log in at `https://platform.xiaomimimo.com/#/console/balance` in Chrome, then refresh CodexBar. + +### “Xiaomi MiMo requires the api-platform_serviceToken and userId cookies” + +The pasted header or imported browser session is missing required cookies. Re-copy the request from the balance page after logging in again. + +### “Xiaomi MiMo browser session expired” + +Your MiMo login is stale. Sign out and back in on the MiMo site, then refresh CodexBar.