From 239758a47601c45c1eae78edd135d22d7f56bfe3 Mon Sep 17 00:00:00 2001 From: Meteorkid <2838801110@qq.com> Date: Mon, 4 May 2026 19:40:39 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E6=B7=BB=E5=8A=A0=E4=B8=AD=E5=9B=BD?= =?UTF-8?q?=E5=8E=82=E5=95=86=E6=94=AF=E6=8C=81=20+=20MiMo=20provider=20+?= =?UTF-8?q?=20KimiK2=20=E4=BD=99=E9=A2=9D=20API?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 新增 Provider： - Zhipu（智谱）— Cookie 余额查询 + API key 验证 - Doubao（豆包）— Cookie 余额查询 + API key 验证 - ERNIE（文心）— Cookie 余额查询 + API key 验证 - MiMo（小米）— Cookie 余额查询 + API key 验证（api-key header） 新增功能： - KimiK2 余额 API（api.moonshot.cn/v1/users/me/balance） - 本地代理 Token 追踪（NWListener，支持所有中国厂商） - 中国厂商余额解析单元测试（25 个测试用例） 技术细节： - MiMo 使用 api-key header（非 Bearer）认证 - Cookie 导入使用 SweetCookieKit - 代理路由支持 api.xiaomimimo.com --- .../CodexBar/PreferencesAdvancedPane.swift | 56 ++++ .../Doubao/DoubaoProviderImplementation.swift | 94 ++++++ .../Doubao/DoubaoSettingsStore.swift | 33 +++ .../Ernie/ErnieProviderImplementation.swift | 94 ++++++ .../Providers/Ernie/ErnieSettingsStore.swift | 33 +++ .../MiMo/MiMoProviderImplementation.swift | 94 ++++++ .../Providers/MiMo/MiMoSettingsStore.swift | 33 +++ .../ProviderImplementationRegistry.swift | 4 + .../Zhipu/ZhipuProviderImplementation.swift | 94 ++++++ .../Providers/Zhipu/ZhipuSettingsStore.swift | 33 +++ .../Resources/ProviderIcon-doubao.svg | 4 + .../CodexBar/Resources/ProviderIcon-ernie.svg | 4 + .../CodexBar/Resources/ProviderIcon-mimo.svg | 4 + .../CodexBar/Resources/ProviderIcon-zhipu.svg | 4 + Sources/CodexBar/SettingsStore+Defaults.swift | 19 ++ Sources/CodexBar/SettingsStore.swift | 7 +- Sources/CodexBar/SettingsStoreState.swift | 2 + Sources/CodexBar/UsageStore.swift | 5 +- Sources/CodexBarCLI/TokenAccountCLI.swift | 2 +- .../Config/ProviderConfigEnvironment.swift | 4 + .../CodexBarCore/Logging/LogCategories.swift | 16 ++ .../Providers/Doubao/DoubaoCookieHeader.swift | 64 +++++ .../Doubao/DoubaoCookieImporter.swift | 169 +++++++++++ .../Doubao/DoubaoProviderDescriptor.swift | 128 +++++++++ .../Doubao/DoubaoSettingsReader.swift | 34 +++ .../Providers/Doubao/DoubaoUsageFetcher.swift | 231 +++++++++++++++ .../Providers/Ernie/ErnieCookieHeader.swift | 64 +++++ .../Providers/Ernie/ErnieCookieImporter.swift | 169 +++++++++++ .../Ernie/ErnieProviderDescriptor.swift | 128 +++++++++ .../Providers/Ernie/ErnieSettingsReader.swift | 34 +++ .../Providers/Ernie/ErnieUsageFetcher.swift | 231 +++++++++++++++ .../Providers/KimiK2/KimiK2UsageFetcher.swift | 114 +++++++- .../Providers/MiMo/MiMoCookieHeader.swift | 68 +++++ .../Providers/MiMo/MiMoCookieImporter.swift | 178 ++++++++++++ .../MiMo/MiMoProviderDescriptor.swift | 128 +++++++++ .../Providers/MiMo/MiMoSettingsReader.swift | 34 +++ .../Providers/MiMo/MiMoUsageFetcher.swift | 268 ++++++++++++++++++ .../Providers/ProviderDescriptor.swift | 4 + .../Providers/ProviderSettingsSnapshot.swift | 84 +++++- .../Providers/ProviderTokenResolver.swift | 48 ++++ .../CodexBarCore/Providers/Providers.swift | 8 + .../Shared/OpenAICompatibleVerifier.swift | 103 +++++++ .../Providers/Zhipu/ZhipuCookieHeader.swift | 68 +++++ .../Providers/Zhipu/ZhipuCookieImporter.swift | 178 ++++++++++++ .../Zhipu/ZhipuProviderDescriptor.swift | 128 +++++++++ .../Providers/Zhipu/ZhipuSettingsReader.swift | 34 +++ .../Providers/Zhipu/ZhipuUsageFetcher.swift | 233 +++++++++++++++ .../Proxy/HTTPRequestParser.swift | 64 +++++ .../Proxy/HTTPResponseWriter.swift | 34 +++ .../CodexBarCore/Proxy/LocalProxyServer.swift | 220 ++++++++++++++ .../Proxy/ProxyConfiguration.swift | 43 +++ Sources/CodexBarCore/Proxy/ProxyManager.swift | 55 ++++ .../CodexBarCore/Proxy/ProxyUsageBridge.swift | 46 +++ .../CodexBarCore/Proxy/TokenAccumulator.swift | 88 ++++++ .../Vendored/CostUsage/CostUsageScanner.swift | 2 +- .../CodexBarWidgetProvider.swift | 4 + .../CodexBarWidget/CodexBarWidgetViews.swift | 12 + .../DoubaoBalanceParsingTests.swift | 91 ++++++ .../ErnieBalanceParsingTests.swift | 91 ++++++ .../KimiK2BalanceParsingTests.swift | 92 ++++++ Tests/CodexBarTests/SettingsStoreTests.swift | 4 + .../ZhipuBalanceParsingTests.swift | 108 +++++++ 62 files changed, 4510 insertions(+), 13 deletions(-) create mode 100644 Sources/CodexBar/Providers/Doubao/DoubaoProviderImplementation.swift create mode 100644 Sources/CodexBar/Providers/Doubao/DoubaoSettingsStore.swift create mode 100644 Sources/CodexBar/Providers/Ernie/ErnieProviderImplementation.swift create mode 100644 Sources/CodexBar/Providers/Ernie/ErnieSettingsStore.swift create mode 100644 Sources/CodexBar/Providers/MiMo/MiMoProviderImplementation.swift create mode 100644 Sources/CodexBar/Providers/MiMo/MiMoSettingsStore.swift create mode 100644 Sources/CodexBar/Providers/Zhipu/ZhipuProviderImplementation.swift create mode 100644 Sources/CodexBar/Providers/Zhipu/ZhipuSettingsStore.swift create mode 100644 Sources/CodexBar/Resources/ProviderIcon-doubao.svg create mode 100644 Sources/CodexBar/Resources/ProviderIcon-ernie.svg create mode 100644 Sources/CodexBar/Resources/ProviderIcon-mimo.svg create mode 100644 Sources/CodexBar/Resources/ProviderIcon-zhipu.svg create mode 100644 Sources/CodexBarCore/Providers/Doubao/DoubaoCookieHeader.swift create mode 100644 Sources/CodexBarCore/Providers/Doubao/DoubaoCookieImporter.swift create mode 100644 Sources/CodexBarCore/Providers/Doubao/DoubaoProviderDescriptor.swift create mode 100644 Sources/CodexBarCore/Providers/Doubao/DoubaoSettingsReader.swift create mode 100644 Sources/CodexBarCore/Providers/Doubao/DoubaoUsageFetcher.swift create mode 100644 Sources/CodexBarCore/Providers/Ernie/ErnieCookieHeader.swift create mode 100644 Sources/CodexBarCore/Providers/Ernie/ErnieCookieImporter.swift create mode 100644 Sources/CodexBarCore/Providers/Ernie/ErnieProviderDescriptor.swift create mode 100644 Sources/CodexBarCore/Providers/Ernie/ErnieSettingsReader.swift create mode 100644 Sources/CodexBarCore/Providers/Ernie/ErnieUsageFetcher.swift create mode 100644 Sources/CodexBarCore/Providers/MiMo/MiMoCookieHeader.swift create mode 100644 Sources/CodexBarCore/Providers/MiMo/MiMoCookieImporter.swift create mode 100644 Sources/CodexBarCore/Providers/MiMo/MiMoProviderDescriptor.swift create mode 100644 Sources/CodexBarCore/Providers/MiMo/MiMoSettingsReader.swift create mode 100644 Sources/CodexBarCore/Providers/MiMo/MiMoUsageFetcher.swift create mode 100644 Sources/CodexBarCore/Providers/Shared/OpenAICompatibleVerifier.swift create mode 100644 Sources/CodexBarCore/Providers/Zhipu/ZhipuCookieHeader.swift create mode 100644 Sources/CodexBarCore/Providers/Zhipu/ZhipuCookieImporter.swift create mode 100644 Sources/CodexBarCore/Providers/Zhipu/ZhipuProviderDescriptor.swift create mode 100644 Sources/CodexBarCore/Providers/Zhipu/ZhipuSettingsReader.swift create mode 100644 Sources/CodexBarCore/Providers/Zhipu/ZhipuUsageFetcher.swift create mode 100644 Sources/CodexBarCore/Proxy/HTTPRequestParser.swift create mode 100644 Sources/CodexBarCore/Proxy/HTTPResponseWriter.swift create mode 100644 Sources/CodexBarCore/Proxy/LocalProxyServer.swift create mode 100644 Sources/CodexBarCore/Proxy/ProxyConfiguration.swift create mode 100644 Sources/CodexBarCore/Proxy/ProxyManager.swift create mode 100644 Sources/CodexBarCore/Proxy/ProxyUsageBridge.swift create mode 100644 Sources/CodexBarCore/Proxy/TokenAccumulator.swift create mode 100644 Tests/CodexBarTests/DoubaoBalanceParsingTests.swift create mode 100644 Tests/CodexBarTests/ErnieBalanceParsingTests.swift create mode 100644 Tests/CodexBarTests/KimiK2BalanceParsingTests.swift create mode 100644 Tests/CodexBarTests/ZhipuBalanceParsingTests.swift diff --git a/Sources/CodexBar/PreferencesAdvancedPane.swift b/Sources/CodexBar/PreferencesAdvancedPane.swift index 9b901c9d8..62357c6de 100644 --- a/Sources/CodexBar/PreferencesAdvancedPane.swift +++ b/Sources/CodexBar/PreferencesAdvancedPane.swift @@ -1,3 +1,4 @@ +import CodexBarCore import KeyboardShortcuts import SwiftUI @@ -6,6 +7,7 @@ struct AdvancedPane: View { @Bindable var settings: SettingsStore @State private var isInstallingCLI = false @State private var cliStatus: String? + @StateObject private var proxyManager = ProxyManager() var body: some View { ScrollView(.vertical, showsIndicators: true) { @@ -92,6 +94,60 @@ struct AdvancedPane: View { subtitle: "Prevents any Keychain access while enabled.", binding: self.$settings.debugDisableKeychainAccess) } + + Divider() + + SettingsSection( + title: "Local Proxy", + caption: """ + Run a local HTTP proxy to intercept API responses and track token usage. \ + Set your API client's base URL to http://127.0.0.1:. + """) { + PreferenceToggleRow( + title: "Enable local proxy", + subtitle: self.proxyManager.isRunning + ? "Running on port \(self.proxyManager.activePort)" + : "Not running", + binding: Binding( + get: { self.settings.proxyEnabled }, + set: { newValue in + self.settings.proxyEnabled = newValue + if newValue { + self.proxyManager.start(port: self.settings.proxyPort) + } else { + self.proxyManager.stop() + } + })) + + HStack(spacing: 12) { + Text("Port") + .font(.body) + Spacer() + TextField("9876", value: self.$settings.proxyPort, format: .number) + .textFieldStyle(.roundedBorder) + .frame(width: 80) + .multilineTextAlignment(.trailing) + } + + if self.proxyManager.isRunning { + HStack(spacing: 8) { + Circle() + .fill(.green) + .frame(width: 8, height: 8) + Text("Listening on 127.0.0.1:\(self.proxyManager.activePort)") + .font(.footnote) + .foregroundStyle(.secondary) + Spacer() + Button("Copy URL") { + NSPasteboard.general.clearContents() + NSPasteboard.general.setString( + "http://127.0.0.1:\(self.proxyManager.activePort)", + forType: .string) + } + .controlSize(.small) + } + } + } } .frame(maxWidth: .infinity, alignment: .leading) .padding(.horizontal, 20) diff --git a/Sources/CodexBar/Providers/Doubao/DoubaoProviderImplementation.swift b/Sources/CodexBar/Providers/Doubao/DoubaoProviderImplementation.swift new file mode 100644 index 000000000..b8dbc3303 --- /dev/null +++ b/Sources/CodexBar/Providers/Doubao/DoubaoProviderImplementation.swift @@ -0,0 +1,94 @@ +import AppKit +import CodexBarCore +import CodexBarMacroSupport +import Foundation +import SwiftUI + +@ProviderImplementationRegistration +struct DoubaoProviderImplementation: ProviderImplementation { + let id: UsageProvider = .doubao + + @MainActor + func presentation(context _: ProviderPresentationContext) -> ProviderPresentation { + ProviderPresentation { _ in "web" } + } + + @MainActor + func observeSettings(_ settings: SettingsStore) { + _ = settings.doubaoCookieSource + _ = settings.doubaoManualCookieHeader + } + + @MainActor + func settingsSnapshot(context: ProviderSettingsSnapshotContext) -> ProviderSettingsSnapshotContribution? { + .doubao(context.settings.doubaoSettingsSnapshot(tokenOverride: context.tokenOverride)) + } + + @MainActor + func isAvailable(context: ProviderAvailabilityContext) -> Bool { + if DoubaoSettingsReader.apiKey(environment: context.environment) != nil { + return true + } + return !context.settings.tokenAccounts(for: .doubao).isEmpty + } + + @MainActor + func settingsPickers(context: ProviderSettingsContext) -> [ProviderSettingsPickerDescriptor] { + let cookieBinding = Binding( + get: { context.settings.doubaoCookieSource.rawValue }, + set: { raw in + context.settings.doubaoCookieSource = ProviderCookieSource(rawValue: raw) ?? .auto + }) + let options = ProviderCookieSourceUI.options( + allowsOff: true, + keychainDisabled: context.settings.debugDisableKeychainAccess) + + let subtitle: () -> String? = { + ProviderCookieSourceUI.subtitle( + source: context.settings.doubaoCookieSource, + keychainDisabled: context.settings.debugDisableKeychainAccess, + auto: "Automatic imports browser cookies.", + manual: "Paste the full Cookie header value.", + off: "Doubao cookies are disabled.") + } + + return [ + ProviderSettingsPickerDescriptor( + id: "doubao-cookie-source", + title: "Cookie source", + subtitle: "Automatic imports browser cookies.", + dynamicSubtitle: subtitle, + binding: cookieBinding, + options: options, + isVisible: nil, + onChange: nil), + ] + } + + @MainActor + func settingsFields(context: ProviderSettingsContext) -> [ProviderSettingsFieldDescriptor] { + [ + ProviderSettingsFieldDescriptor( + id: "doubao-cookie", + title: "", + subtitle: "", + kind: .secure, + placeholder: "Cookie: \u{2026}\n\nPaste the full Cookie header from console.volcengine.com", + binding: context.stringBinding(\.doubaoManualCookieHeader), + actions: [ + ProviderSettingsActionDescriptor( + id: "doubao-open-console", + title: "Open Console", + style: .link, + isVisible: nil, + perform: { + if let url = URL(string: "https://console.volcengine.com/ark") { + NSWorkspace.shared.open(url) + } + }), + ], + isVisible: { context.settings.doubaoCookieSource == .manual }, + onActivate: nil), + ] + } +} diff --git a/Sources/CodexBar/Providers/Doubao/DoubaoSettingsStore.swift b/Sources/CodexBar/Providers/Doubao/DoubaoSettingsStore.swift new file mode 100644 index 000000000..58ba342c4 --- /dev/null +++ b/Sources/CodexBar/Providers/Doubao/DoubaoSettingsStore.swift @@ -0,0 +1,33 @@ +import CodexBarCore +import Foundation + +extension SettingsStore { + var doubaoManualCookieHeader: String { + get { self.configSnapshot.providerConfig(for: .doubao)?.sanitizedCookieHeader ?? "" } + set { + self.updateProviderConfig(provider: .doubao) { entry in + entry.cookieHeader = self.normalizedConfigValue(newValue) + } + self.logSecretUpdate(provider: .doubao, field: "cookieHeader", value: newValue) + } + } + + var doubaoCookieSource: ProviderCookieSource { + get { self.resolvedCookieSource(provider: .doubao, fallback: .auto) } + set { + self.updateProviderConfig(provider: .doubao) { entry in + entry.cookieSource = newValue + } + self.logProviderModeChange(provider: .doubao, field: "cookieSource", value: newValue.rawValue) + } + } +} + +extension SettingsStore { + func doubaoSettingsSnapshot(tokenOverride: TokenAccountOverride?) -> ProviderSettingsSnapshot.DoubaoProviderSettings { + _ = tokenOverride + return ProviderSettingsSnapshot.DoubaoProviderSettings( + cookieSource: self.doubaoCookieSource, + manualCookieHeader: self.doubaoManualCookieHeader) + } +} diff --git a/Sources/CodexBar/Providers/Ernie/ErnieProviderImplementation.swift b/Sources/CodexBar/Providers/Ernie/ErnieProviderImplementation.swift new file mode 100644 index 000000000..791effd57 --- /dev/null +++ b/Sources/CodexBar/Providers/Ernie/ErnieProviderImplementation.swift @@ -0,0 +1,94 @@ +import AppKit +import CodexBarCore +import CodexBarMacroSupport +import Foundation +import SwiftUI + +@ProviderImplementationRegistration +struct ErnieProviderImplementation: ProviderImplementation { + let id: UsageProvider = .ernie + + @MainActor + func presentation(context _: ProviderPresentationContext) -> ProviderPresentation { + ProviderPresentation { _ in "web" } + } + + @MainActor + func observeSettings(_ settings: SettingsStore) { + _ = settings.ernieCookieSource + _ = settings.ernieManualCookieHeader + } + + @MainActor + func settingsSnapshot(context: ProviderSettingsSnapshotContext) -> ProviderSettingsSnapshotContribution? { + .ernie(context.settings.ernieSettingsSnapshot(tokenOverride: context.tokenOverride)) + } + + @MainActor + func isAvailable(context: ProviderAvailabilityContext) -> Bool { + if ErnieSettingsReader.apiKey(environment: context.environment) != nil { + return true + } + return !context.settings.tokenAccounts(for: .ernie).isEmpty + } + + @MainActor + func settingsPickers(context: ProviderSettingsContext) -> [ProviderSettingsPickerDescriptor] { + let cookieBinding = Binding( + get: { context.settings.ernieCookieSource.rawValue }, + set: { raw in + context.settings.ernieCookieSource = ProviderCookieSource(rawValue: raw) ?? .auto + }) + let options = ProviderCookieSourceUI.options( + allowsOff: true, + keychainDisabled: context.settings.debugDisableKeychainAccess) + + let subtitle: () -> String? = { + ProviderCookieSourceUI.subtitle( + source: context.settings.ernieCookieSource, + keychainDisabled: context.settings.debugDisableKeychainAccess, + auto: "Automatic imports browser cookies.", + manual: "Paste the full Cookie header value.", + off: "ERNIE cookies are disabled.") + } + + return [ + ProviderSettingsPickerDescriptor( + id: "ernie-cookie-source", + title: "Cookie source", + subtitle: "Automatic imports browser cookies.", + dynamicSubtitle: subtitle, + binding: cookieBinding, + options: options, + isVisible: nil, + onChange: nil), + ] + } + + @MainActor + func settingsFields(context: ProviderSettingsContext) -> [ProviderSettingsFieldDescriptor] { + [ + ProviderSettingsFieldDescriptor( + id: "ernie-cookie", + title: "", + subtitle: "", + kind: .secure, + placeholder: "Cookie: \u{2026}\n\nPaste the full Cookie header from console.bce.baidu.com", + binding: context.stringBinding(\.ernieManualCookieHeader), + actions: [ + ProviderSettingsActionDescriptor( + id: "ernie-open-console", + title: "Open Console", + style: .link, + isVisible: nil, + perform: { + if let url = URL(string: "https://console.bce.baidu.com/qianfan/overview") { + NSWorkspace.shared.open(url) + } + }), + ], + isVisible: { context.settings.ernieCookieSource == .manual }, + onActivate: nil), + ] + } +} diff --git a/Sources/CodexBar/Providers/Ernie/ErnieSettingsStore.swift b/Sources/CodexBar/Providers/Ernie/ErnieSettingsStore.swift new file mode 100644 index 000000000..aafa1c2f8 --- /dev/null +++ b/Sources/CodexBar/Providers/Ernie/ErnieSettingsStore.swift @@ -0,0 +1,33 @@ +import CodexBarCore +import Foundation + +extension SettingsStore { + var ernieManualCookieHeader: String { + get { self.configSnapshot.providerConfig(for: .ernie)?.sanitizedCookieHeader ?? "" } + set { + self.updateProviderConfig(provider: .ernie) { entry in + entry.cookieHeader = self.normalizedConfigValue(newValue) + } + self.logSecretUpdate(provider: .ernie, field: "cookieHeader", value: newValue) + } + } + + var ernieCookieSource: ProviderCookieSource { + get { self.resolvedCookieSource(provider: .ernie, fallback: .auto) } + set { + self.updateProviderConfig(provider: .ernie) { entry in + entry.cookieSource = newValue + } + self.logProviderModeChange(provider: .ernie, field: "cookieSource", value: newValue.rawValue) + } + } +} + +extension SettingsStore { + func ernieSettingsSnapshot(tokenOverride: TokenAccountOverride?) -> ProviderSettingsSnapshot.ErnieProviderSettings { + _ = tokenOverride + return ProviderSettingsSnapshot.ErnieProviderSettings( + cookieSource: self.ernieCookieSource, + manualCookieHeader: self.ernieManualCookieHeader) + } +} diff --git a/Sources/CodexBar/Providers/MiMo/MiMoProviderImplementation.swift b/Sources/CodexBar/Providers/MiMo/MiMoProviderImplementation.swift new file mode 100644 index 000000000..557edd3cd --- /dev/null +++ b/Sources/CodexBar/Providers/MiMo/MiMoProviderImplementation.swift @@ -0,0 +1,94 @@ +import AppKit +import CodexBarCore +import CodexBarMacroSupport +import Foundation +import SwiftUI + +@ProviderImplementationRegistration +struct MiMoProviderImplementation: ProviderImplementation { + let id: UsageProvider = .mimo + + @MainActor + func presentation(context _: ProviderPresentationContext) -> ProviderPresentation { + ProviderPresentation { _ in "web" } + } + + @MainActor + func observeSettings(_ settings: SettingsStore) { + _ = settings.mimoCookieSource + _ = settings.mimoManualCookieHeader + } + + @MainActor + func settingsSnapshot(context: ProviderSettingsSnapshotContext) -> ProviderSettingsSnapshotContribution? { + .mimo(context.settings.mimoSettingsSnapshot(tokenOverride: context.tokenOverride)) + } + + @MainActor + func isAvailable(context: ProviderAvailabilityContext) -> Bool { + if MiMoSettingsReader.apiKey(environment: context.environment) != nil { + return true + } + return !context.settings.tokenAccounts(for: .mimo).isEmpty + } + + @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 options = ProviderCookieSourceUI.options( + allowsOff: true, + keychainDisabled: context.settings.debugDisableKeychainAccess) + + let subtitle: () -> String? = { + ProviderCookieSourceUI.subtitle( + source: context.settings.mimoCookieSource, + keychainDisabled: context.settings.debugDisableKeychainAccess, + auto: "Automatic imports browser cookies.", + manual: "Paste the full Cookie header value.", + off: "MiMo cookies are disabled.") + } + + return [ + ProviderSettingsPickerDescriptor( + id: "mimo-cookie-source", + title: "Cookie source", + subtitle: "Automatic imports browser cookies.", + dynamicSubtitle: subtitle, + binding: cookieBinding, + options: options, + isVisible: nil, + onChange: nil), + ] + } + + @MainActor + func settingsFields(context: ProviderSettingsContext) -> [ProviderSettingsFieldDescriptor] { + [ + ProviderSettingsFieldDescriptor( + id: "mimo-cookie", + title: "", + subtitle: "", + kind: .secure, + placeholder: "Cookie: \u{2026}\n\nPaste the full Cookie header from platform.xiaomimimo.com", + binding: context.stringBinding(\.mimoManualCookieHeader), + actions: [ + ProviderSettingsActionDescriptor( + id: "mimo-open-console", + title: "Open Console", + style: .link, + isVisible: nil, + perform: { + if let url = URL(string: "https://platform.xiaomimimo.com/console/balance") { + NSWorkspace.shared.open(url) + } + }), + ], + isVisible: { context.settings.mimoCookieSource == .manual }, + onActivate: nil), + ] + } +} diff --git a/Sources/CodexBar/Providers/MiMo/MiMoSettingsStore.swift b/Sources/CodexBar/Providers/MiMo/MiMoSettingsStore.swift new file mode 100644 index 000000000..2e84410ba --- /dev/null +++ b/Sources/CodexBar/Providers/MiMo/MiMoSettingsStore.swift @@ -0,0 +1,33 @@ +import CodexBarCore +import Foundation + +extension SettingsStore { + var mimoManualCookieHeader: 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) + } + } +} + +extension SettingsStore { + func mimoSettingsSnapshot(tokenOverride: TokenAccountOverride?) -> ProviderSettingsSnapshot.MiMoProviderSettings { + _ = tokenOverride + return ProviderSettingsSnapshot.MiMoProviderSettings( + cookieSource: self.mimoCookieSource, + manualCookieHeader: self.mimoManualCookieHeader) + } +} diff --git a/Sources/CodexBar/Providers/Shared/ProviderImplementationRegistry.swift b/Sources/CodexBar/Providers/Shared/ProviderImplementationRegistry.swift index 14a9208bd..d896463fe 100644 --- a/Sources/CodexBar/Providers/Shared/ProviderImplementationRegistry.swift +++ b/Sources/CodexBar/Providers/Shared/ProviderImplementationRegistry.swift @@ -42,6 +42,10 @@ enum ProviderImplementationRegistry { case .mistral: MistralProviderImplementation() case .deepseek: DeepSeekProviderImplementation() case .codebuff: CodebuffProviderImplementation() + case .zhipu: ZhipuProviderImplementation() + case .doubao: DoubaoProviderImplementation() + case .ernie: ErnieProviderImplementation() + case .mimo: MiMoProviderImplementation() } } diff --git a/Sources/CodexBar/Providers/Zhipu/ZhipuProviderImplementation.swift b/Sources/CodexBar/Providers/Zhipu/ZhipuProviderImplementation.swift new file mode 100644 index 000000000..9b5d06951 --- /dev/null +++ b/Sources/CodexBar/Providers/Zhipu/ZhipuProviderImplementation.swift @@ -0,0 +1,94 @@ +import AppKit +import CodexBarCore +import CodexBarMacroSupport +import Foundation +import SwiftUI + +@ProviderImplementationRegistration +struct ZhipuProviderImplementation: ProviderImplementation { + let id: UsageProvider = .zhipu + + @MainActor + func presentation(context _: ProviderPresentationContext) -> ProviderPresentation { + ProviderPresentation { _ in "web" } + } + + @MainActor + func observeSettings(_ settings: SettingsStore) { + _ = settings.zhipuCookieSource + _ = settings.zhipuManualCookieHeader + } + + @MainActor + func settingsSnapshot(context: ProviderSettingsSnapshotContext) -> ProviderSettingsSnapshotContribution? { + .zhipu(context.settings.zhipuSettingsSnapshot(tokenOverride: context.tokenOverride)) + } + + @MainActor + func isAvailable(context: ProviderAvailabilityContext) -> Bool { + if ZhipuSettingsReader.apiKey(environment: context.environment) != nil { + return true + } + return !context.settings.tokenAccounts(for: .zhipu).isEmpty + } + + @MainActor + func settingsPickers(context: ProviderSettingsContext) -> [ProviderSettingsPickerDescriptor] { + let cookieBinding = Binding( + get: { context.settings.zhipuCookieSource.rawValue }, + set: { raw in + context.settings.zhipuCookieSource = ProviderCookieSource(rawValue: raw) ?? .auto + }) + let options = ProviderCookieSourceUI.options( + allowsOff: true, + keychainDisabled: context.settings.debugDisableKeychainAccess) + + let subtitle: () -> String? = { + ProviderCookieSourceUI.subtitle( + source: context.settings.zhipuCookieSource, + keychainDisabled: context.settings.debugDisableKeychainAccess, + auto: "Automatic imports browser cookies.", + manual: "Paste the full Cookie header value.", + off: "Zhipu cookies are disabled.") + } + + return [ + ProviderSettingsPickerDescriptor( + id: "zhipu-cookie-source", + title: "Cookie source", + subtitle: "Automatic imports browser cookies.", + dynamicSubtitle: subtitle, + binding: cookieBinding, + options: options, + isVisible: nil, + onChange: nil), + ] + } + + @MainActor + func settingsFields(context: ProviderSettingsContext) -> [ProviderSettingsFieldDescriptor] { + [ + ProviderSettingsFieldDescriptor( + id: "zhipu-cookie", + title: "", + subtitle: "", + kind: .secure, + placeholder: "Cookie: \u{2026}\n\nPaste the full Cookie header from open.bigmodel.cn", + binding: context.stringBinding(\.zhipuManualCookieHeader), + actions: [ + ProviderSettingsActionDescriptor( + id: "zhipu-open-console", + title: "Open Console", + style: .link, + isVisible: nil, + perform: { + if let url = URL(string: "https://open.bigmodel.cn/usercenter") { + NSWorkspace.shared.open(url) + } + }), + ], + isVisible: { context.settings.zhipuCookieSource == .manual }, + onActivate: nil), + ] + } +} diff --git a/Sources/CodexBar/Providers/Zhipu/ZhipuSettingsStore.swift b/Sources/CodexBar/Providers/Zhipu/ZhipuSettingsStore.swift new file mode 100644 index 000000000..5b05e05a4 --- /dev/null +++ b/Sources/CodexBar/Providers/Zhipu/ZhipuSettingsStore.swift @@ -0,0 +1,33 @@ +import CodexBarCore +import Foundation + +extension SettingsStore { + var zhipuManualCookieHeader: String { + get { self.configSnapshot.providerConfig(for: .zhipu)?.sanitizedCookieHeader ?? "" } + set { + self.updateProviderConfig(provider: .zhipu) { entry in + entry.cookieHeader = self.normalizedConfigValue(newValue) + } + self.logSecretUpdate(provider: .zhipu, field: "cookieHeader", value: newValue) + } + } + + var zhipuCookieSource: ProviderCookieSource { + get { self.resolvedCookieSource(provider: .zhipu, fallback: .auto) } + set { + self.updateProviderConfig(provider: .zhipu) { entry in + entry.cookieSource = newValue + } + self.logProviderModeChange(provider: .zhipu, field: "cookieSource", value: newValue.rawValue) + } + } +} + +extension SettingsStore { + func zhipuSettingsSnapshot(tokenOverride: TokenAccountOverride?) -> ProviderSettingsSnapshot.ZhipuProviderSettings { + _ = tokenOverride + return ProviderSettingsSnapshot.ZhipuProviderSettings( + cookieSource: self.zhipuCookieSource, + manualCookieHeader: self.zhipuManualCookieHeader) + } +} diff --git a/Sources/CodexBar/Resources/ProviderIcon-doubao.svg b/Sources/CodexBar/Resources/ProviderIcon-doubao.svg new file mode 100644 index 000000000..2ac75ece5 --- /dev/null +++ b/Sources/CodexBar/Resources/ProviderIcon-doubao.svg @@ -0,0 +1,4 @@ + + + D + diff --git a/Sources/CodexBar/Resources/ProviderIcon-ernie.svg b/Sources/CodexBar/Resources/ProviderIcon-ernie.svg new file mode 100644 index 000000000..e79bd6011 --- /dev/null +++ b/Sources/CodexBar/Resources/ProviderIcon-ernie.svg @@ -0,0 +1,4 @@ + + + E + diff --git a/Sources/CodexBar/Resources/ProviderIcon-mimo.svg b/Sources/CodexBar/Resources/ProviderIcon-mimo.svg new file mode 100644 index 000000000..4a899591c --- /dev/null +++ b/Sources/CodexBar/Resources/ProviderIcon-mimo.svg @@ -0,0 +1,4 @@ + + + M + diff --git a/Sources/CodexBar/Resources/ProviderIcon-zhipu.svg b/Sources/CodexBar/Resources/ProviderIcon-zhipu.svg new file mode 100644 index 000000000..41fb7b0d5 --- /dev/null +++ b/Sources/CodexBar/Resources/ProviderIcon-zhipu.svg @@ -0,0 +1,4 @@ + + + Z + diff --git a/Sources/CodexBar/SettingsStore+Defaults.swift b/Sources/CodexBar/SettingsStore+Defaults.swift index c4705bcc6..a89aa78b8 100644 --- a/Sources/CodexBar/SettingsStore+Defaults.swift +++ b/Sources/CodexBar/SettingsStore+Defaults.swift @@ -532,4 +532,23 @@ extension SettingsStore { } return self.normalizeProviders(providers, maxCount: maxCount) } + + // MARK: - Proxy + + var proxyEnabled: Bool { + get { self.defaultsState.proxyEnabled } + set { + self.defaultsState.proxyEnabled = newValue + self.userDefaults.set(newValue, forKey: "proxyEnabled") + } + } + + var proxyPort: UInt16 { + get { self.defaultsState.proxyPort } + set { + let clamped = max(1024, min(65535, newValue)) + self.defaultsState.proxyPort = clamped + self.userDefaults.set(Int(clamped), forKey: "proxyPort") + } + } } diff --git a/Sources/CodexBar/SettingsStore.swift b/Sources/CodexBar/SettingsStore.swift index d26a48a51..4f814753b 100644 --- a/Sources/CodexBar/SettingsStore.swift +++ b/Sources/CodexBar/SettingsStore.swift @@ -282,6 +282,9 @@ extension SettingsStore { forKey: "mergedOverviewSelectedProviders") as? [String] ?? [] let selectedMenuProviderRaw = userDefaults.string(forKey: "selectedMenuProvider") let providerDetectionCompleted = userDefaults.object(forKey: "providerDetectionCompleted") as? Bool ?? false + let proxyEnabled = userDefaults.object(forKey: "proxyEnabled") as? Bool ?? false + let proxyPortRaw = userDefaults.object(forKey: "proxyPort") as? Int ?? 9876 + let proxyPort = UInt16(max(1024, min(65535, proxyPortRaw))) return SettingsDefaultsState( refreshFrequency: refreshFrequency, @@ -319,7 +322,9 @@ extension SettingsStore { mergedMenuLastSelectedWasOverview: mergedMenuLastSelectedWasOverview, mergedOverviewSelectedProvidersRaw: mergedOverviewSelectedProvidersRaw, selectedMenuProviderRaw: selectedMenuProviderRaw, - providerDetectionCompleted: providerDetectionCompleted) + providerDetectionCompleted: providerDetectionCompleted, + proxyEnabled: proxyEnabled, + proxyPort: proxyPort) } } diff --git a/Sources/CodexBar/SettingsStoreState.swift b/Sources/CodexBar/SettingsStoreState.swift index c28de8f63..0ebbae088 100644 --- a/Sources/CodexBar/SettingsStoreState.swift +++ b/Sources/CodexBar/SettingsStoreState.swift @@ -37,4 +37,6 @@ struct SettingsDefaultsState { var mergedOverviewSelectedProvidersRaw: [String] var selectedMenuProviderRaw: String? var providerDetectionCompleted: Bool + var proxyEnabled: Bool + var proxyPort: UInt16 } diff --git a/Sources/CodexBar/UsageStore.swift b/Sources/CodexBar/UsageStore.swift index 524cf7d21..0db9c1a37 100644 --- a/Sources/CodexBar/UsageStore.swift +++ b/Sources/CodexBar/UsageStore.swift @@ -819,6 +819,9 @@ extension UsageStore { .kimi: "Kimi debug log not yet implemented", .kimik2: "Kimi K2 debug log not yet implemented", .jetbrains: "JetBrains AI debug log not yet implemented", + .zhipu: "Zhipu debug log not yet implemented", + .doubao: "Doubao debug log not yet implemented", + .ernie: "ERNIE debug log not yet implemented", ] let buildText = { switch provider { @@ -892,7 +895,7 @@ extension UsageStore { hasEnvToken: deepSeekHasEnvToken, hasTokenAccount: deepSeekHasTokenAccount) case .gemini, .antigravity, .opencode, .opencodego, .factory, .copilot, .vertexai, .kilo, .kiro, .kimi, - .kimik2, .jetbrains, .perplexity, .abacus, .mistral, .codebuff: + .kimik2, .jetbrains, .perplexity, .abacus, .mistral, .codebuff, .zhipu, .doubao, .ernie, .mimo: return unimplementedDebugLogMessages[provider] ?? "Debug log not yet implemented" } } diff --git a/Sources/CodexBarCLI/TokenAccountCLI.swift b/Sources/CodexBarCLI/TokenAccountCLI.swift index e491fae5f..fe80cc59c 100644 --- a/Sources/CodexBarCLI/TokenAccountCLI.swift +++ b/Sources/CodexBarCLI/TokenAccountCLI.swift @@ -201,7 +201,7 @@ struct TokenAccountCLIContext { cookieSource: cookieSource, manualCookieHeader: cookieHeader)) case .gemini, .antigravity, .copilot, .kiro, .vertexai, .kimik2, .synthetic, .openrouter, .warp, .deepseek, - .codebuff: + .codebuff, .zhipu, .doubao, .ernie, .mimo: return nil } } diff --git a/Sources/CodexBarCore/Config/ProviderConfigEnvironment.swift b/Sources/CodexBarCore/Config/ProviderConfigEnvironment.swift index 80ba60012..2ccd28096 100644 --- a/Sources/CodexBarCore/Config/ProviderConfigEnvironment.swift +++ b/Sources/CodexBarCore/Config/ProviderConfigEnvironment.swift @@ -38,6 +38,10 @@ public enum ProviderConfigEnvironment { if CodebuffSettingsReader.apiKey(environment: base) == nil { env[CodebuffSettingsReader.apiTokenKey] = apiKey } + case .mimo: + if let key = MiMoSettingsReader.apiKeyEnvironmentKeys.first { + env[key] = apiKey + } default: break } diff --git a/Sources/CodexBarCore/Logging/LogCategories.swift b/Sources/CodexBarCore/Logging/LogCategories.swift index ae1d10a34..1166fb550 100644 --- a/Sources/CodexBarCore/Logging/LogCategories.swift +++ b/Sources/CodexBarCore/Logging/LogCategories.swift @@ -22,6 +22,10 @@ public enum LogCategories { public static let cursorLogin = "cursor-login" public static let deepSeekSettings = "deepseek-settings" public static let deepSeekUsage = "deepseek-usage" + public static let doubaoCookie = "doubao-cookie" + public static let doubaoSettings = "doubao-settings" + public static let doubaoUsage = "doubao-usage" + public static let doubaoWeb = "doubao-web" public static let geminiProbe = "gemini-probe" public static let keychainCache = "keychain-cache" public static let keychainMigration = "keychain-migration" @@ -34,6 +38,10 @@ public enum LogCategories { public static let kimiTokenStore = "kimi-token-store" public static let kimiWeb = "kimi-web" public static let kiro = "kiro" + public static let ernieCookie = "ernie-cookie" + public static let ernieWeb = "ernie-web" + public static let qianfanSettings = "qianfan-settings" + public static let qianfanUsage = "qianfan-usage" public static let launchAtLogin = "launch-at-login" public static let login = "login" public static let logging = "logging" @@ -70,4 +78,12 @@ 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 zhipuCookie = "zhipu-cookie" + public static let zhipuSettings = "zhipu-settings" + public static let zhipuUsage = "zhipu-usage" + public static let zhipuWeb = "zhipu-web" + public static let mimoCookie = "mimo-cookie" + public static let mimoSettings = "mimo-settings" + public static let mimoUsage = "mimo-usage" + public static let mimoWeb = "mimo-web" } diff --git a/Sources/CodexBarCore/Providers/Doubao/DoubaoCookieHeader.swift b/Sources/CodexBarCore/Providers/Doubao/DoubaoCookieHeader.swift new file mode 100644 index 000000000..756deae1e --- /dev/null +++ b/Sources/CodexBarCore/Providers/Doubao/DoubaoCookieHeader.swift @@ -0,0 +1,64 @@ +import Foundation + +public struct DoubaoCookieOverride: Sendable { + public let cookieHeader: String + + public init(cookieHeader: String) { + self.cookieHeader = cookieHeader + } +} + +public enum DoubaoCookieHeader { + private static let log = CodexBarLog.logger(LogCategories.doubaoCookie) + private static let headerPatterns: [String] = [ + #"(?i)-H\s*'Cookie:\s*([^']+)'"#, + #"(?i)-H\s*"Cookie:\s*([^"]+)""#, + #"(?i)\bcookie:\s*'([^']+)'"#, + #"(?i)\bcookie:\s*"([^"]+)""#, + #"(?i)\bcookie:\s*([^\r\n]+)"#, + ] + + public static func resolveCookieOverride(context: ProviderFetchContext) -> DoubaoCookieOverride? { + if let settings = context.settings?.doubao, settings.cookieSource == .manual { + if let manual = settings.manualCookieHeader, !manual.isEmpty { + return self.override(from: manual) + } + } + + if let envToken = self.override(from: context.env["DOUBAO_MANUAL_COOKIE"]) { + return envToken + } + + return nil + } + + public static func override(from raw: String?) -> DoubaoCookieOverride? { + guard let raw = raw?.trimmingCharacters(in: .whitespacesAndNewlines), !raw.isEmpty else { + return nil + } + + if raw.contains("=") { + return DoubaoCookieOverride(cookieHeader: raw) + } + + if let cookieHeader = self.extractHeader(from: raw) { + return DoubaoCookieOverride(cookieHeader: cookieHeader) + } + + return nil + } + + private static func extractHeader(from raw: String) -> String? { + for pattern in self.headerPatterns { + guard let regex = try? NSRegularExpression(pattern: pattern, options: []) else { continue } + let range = NSRange(raw.startIndex..= 2, + let captureRange = Range(match.range(at: 1), in: raw) + else { continue } + let captured = String(raw[captureRange]).trimmingCharacters(in: .whitespacesAndNewlines) + if !captured.isEmpty { return captured } + } + return nil + } +} diff --git a/Sources/CodexBarCore/Providers/Doubao/DoubaoCookieImporter.swift b/Sources/CodexBarCore/Providers/Doubao/DoubaoCookieImporter.swift new file mode 100644 index 000000000..8fcbeb0db --- /dev/null +++ b/Sources/CodexBarCore/Providers/Doubao/DoubaoCookieImporter.swift @@ -0,0 +1,169 @@ +import Foundation + +#if os(macOS) +import SweetCookieKit + +public enum DoubaoCookieImporter { + private static let log = CodexBarLog.logger(LogCategories.doubaoCookie) + private static let cookieClient = BrowserCookieClient() + private static let cookieDomains = ["console.volcengine.com"] + private static let cookieImportOrder: BrowserCookieImportOrder = + ProviderDefaults.metadata[.doubao]?.browserCookieOrder ?? Browser.defaultImportOrder + + public struct SessionInfo: Sendable { + public let cookies: [HTTPCookie] + public let sourceLabel: String + + public init(cookies: [HTTPCookie], sourceLabel: String) { + self.cookies = cookies + self.sourceLabel = sourceLabel + } + + public var cookieHeader: String? { + let parts = self.cookies.compactMap { cookie -> String? in + "\(cookie.name)=\(cookie.value)" + } + guard !parts.isEmpty else { return nil } + return parts.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 DoubaoCookieImportError.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.codexBarRecords( + 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 } + + log("Found \(httpCookies.count) cookies for console.volcengine.com 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 DoubaoCookieImportError.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?("[doubao-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 DoubaoCookieImportError: LocalizedError { + case noCookies + + var errorDescription: String? { + switch self { + case .noCookies: "No Doubao session cookies found in browsers." + } + } +} +#endif diff --git a/Sources/CodexBarCore/Providers/Doubao/DoubaoProviderDescriptor.swift b/Sources/CodexBarCore/Providers/Doubao/DoubaoProviderDescriptor.swift new file mode 100644 index 000000000..f7eac2d7a --- /dev/null +++ b/Sources/CodexBarCore/Providers/Doubao/DoubaoProviderDescriptor.swift @@ -0,0 +1,128 @@ +import CodexBarMacroSupport +import Foundation + +@ProviderDescriptorRegistration +@ProviderDescriptorDefinition +public enum DoubaoProviderDescriptor { + static func makeDescriptor() -> ProviderDescriptor { + ProviderDescriptor( + id: .doubao, + metadata: ProviderMetadata( + id: .doubao, + displayName: "Doubao", + sessionLabel: "Status", + weeklyLabel: "Endpoints", + 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", + statusPageURL: nil), + branding: ProviderBranding( + iconStyle: .doubao, + iconResourceName: "ProviderIcon-doubao", + color: ProviderColor(red: 0.0, green: 0.55, blue: 0.95)), + tokenCost: ProviderTokenCostConfig( + supportsTokenCost: false, + noDataMessage: { "Doubao cost summary is not available via API." }), + fetchPlan: ProviderFetchPlan( + sourceModes: [.auto, .web, .api], + pipeline: ProviderFetchPipeline(resolveStrategies: { _ in + [DoubaoWebFetchStrategy(), DoubaoAPIFetchStrategy()] + })), + cli: ProviderCLIConfig( + name: "doubao", + aliases: ["volcengine", "ark"], + 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.verifyAPI(apiKey: apiKey) + return self.makeResult( + usage: usage.toUsageSnapshot(), + sourceLabel: "api") + } + + func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool { + true + } + + private static func resolveToken(environment: [String: String]) -> String? { + ProviderTokenResolver.doubaoToken(environment: environment) + } +} + +struct DoubaoWebFetchStrategy: ProviderFetchStrategy { + let id: String = "doubao.web" + let kind: ProviderFetchKind = .web + private static let log = CodexBarLog.logger(LogCategories.doubaoWeb) + + func isAvailable(_ context: ProviderFetchContext) async -> Bool { + if DoubaoCookieHeader.resolveCookieOverride(context: context) != nil { + return true + } + + #if os(macOS) + if context.settings?.doubao?.cookieSource != .off { + return DoubaoCookieImporter.hasSession() + } + #endif + + return false + } + + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + guard let cookieHeader = self.resolveCookieHeader(context: context) else { + throw DoubaoUsageError.missingCookie + } + + let snapshot = try await DoubaoUsageFetcher.fetchBalance(cookieHeader: cookieHeader) + return self.makeResult( + usage: snapshot.toUsageSnapshot(), + sourceLabel: "web") + } + + func shouldFallback(on error: Error, context: ProviderFetchContext) -> Bool { + if case DoubaoUsageError.missingCookie = error { return false } + return true + } + + private func resolveCookieHeader(context: ProviderFetchContext) -> String? { + if let override = DoubaoCookieHeader.resolveCookieOverride(context: context) { + return override.cookieHeader + } + + #if os(macOS) + if context.settings?.doubao?.cookieSource != .off { + do { + let session = try DoubaoCookieImporter.importSession() + if let header = session.cookieHeader { + return header + } + } catch { + Self.log.debug("Doubao browser cookie import failed: \(error)") + } + } + #endif + + return nil + } +} diff --git a/Sources/CodexBarCore/Providers/Doubao/DoubaoSettingsReader.swift b/Sources/CodexBarCore/Providers/Doubao/DoubaoSettingsReader.swift new file mode 100644 index 000000000..0388b1ea1 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Doubao/DoubaoSettingsReader.swift @@ -0,0 +1,34 @@ +import Foundation + +public struct DoubaoSettingsReader: Sendable { + public static let apiKeyEnvironmentKey = "DOUBAO_API_KEY" + public static let apiKeyEnvironmentKeys = [Self.apiKeyEnvironmentKey, "VOLCENGINE_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..3470b0b71 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Doubao/DoubaoUsageFetcher.swift @@ -0,0 +1,231 @@ +import Foundation +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif + +// MARK: - Domain snapshot + +public struct DoubaoUsageSnapshot: Sendable { + public let isConnected: Bool + public let modelCount: Int + public let modelNames: [String] + public let balanceInfo: DoubaoBalanceInfo? + public let updatedAt: Date + + public init( + isConnected: Bool, + modelCount: Int, + modelNames: [String], + balanceInfo: DoubaoBalanceInfo? = nil, + updatedAt: Date) + { + self.isConnected = isConnected + self.modelCount = modelCount + self.modelNames = modelNames + self.balanceInfo = balanceInfo + self.updatedAt = updatedAt + } + + public func toUsageSnapshot() -> UsageSnapshot { + let detail: String + let usedPercent: Double + if let balance = self.balanceInfo { + let remaining = balance.availableBalance + let total = balance.totalBalance + usedPercent = total > 0 ? max(0, min(100, (total - remaining) / total * 100)) : 0 + detail = String(format: "¥%.2f / ¥%.2f", remaining, total) + } else if self.isConnected { + let names = self.modelNames.joined(separator: ", ") + detail = "API Connected — \(self.modelCount) endpoints (\(names))" + usedPercent = 0 + } else { + detail = "API not connected" + usedPercent = 100 + } + + let identity = ProviderIdentitySnapshot( + providerID: .doubao, + accountEmail: nil, + accountOrganization: nil, + loginMethod: nil) + + let primaryWindow = RateWindow( + usedPercent: usedPercent, + windowMinutes: nil, + resetsAt: nil, + resetDescription: detail) + + var secondaryWindow: RateWindow? + if let balance = self.balanceInfo { + secondaryWindow = RateWindow( + usedPercent: 0, + windowMinutes: nil, + resetsAt: nil, + resetDescription: "Used: ¥\(String(format: "%.2f", balance.usedBalance))") + } + + let providerCost: ProviderCostSnapshot? + if let balance = self.balanceInfo { + providerCost = ProviderCostSnapshot( + used: balance.usedBalance, + limit: balance.totalBalance, + currencyCode: "CNY", + period: nil, + resetsAt: nil, + updatedAt: self.updatedAt) + } else { + providerCost = nil + } + + return UsageSnapshot( + primary: primaryWindow, + secondary: secondaryWindow, + tertiary: nil, + providerCost: providerCost, + updatedAt: self.updatedAt, + identity: identity) + } +} + +public struct DoubaoBalanceInfo: Sendable { + public let availableBalance: Double + public let usedBalance: Double + public let totalBalance: Double + + public init(availableBalance: Double, usedBalance: Double, totalBalance: Double) { + self.availableBalance = availableBalance + self.usedBalance = usedBalance + self.totalBalance = totalBalance + } +} + +// MARK: - Errors + +public enum DoubaoUsageError: LocalizedError, Sendable { + case missingCredentials + case missingCookie + case networkError(String) + case apiError(String) + case parseFailed(String) + + public var errorDescription: String? { + switch self { + case .missingCredentials: + "Missing Doubao API key." + case .missingCookie: + "Missing Doubao session cookie." + case let .networkError(message): + "Doubao network error: \(message)" + case let .apiError(message): + "Doubao API error: \(message)" + case let .parseFailed(message): + "Failed to parse Doubao response: \(message)" + } + } +} + +// MARK: - Fetcher + +public struct DoubaoUsageFetcher: Sendable { + private static let log = CodexBarLog.logger(LogCategories.doubaoUsage) + private static let webLog = CodexBarLog.logger(LogCategories.doubaoWeb) + private static let baseURL = URL(string: "https://ark.cn-beijing.volces.com/api/v3")! + private static let balanceURL = + URL(string: "https://console.volcengine.com/ark/api/open/v1/resourcepack/billing/query")! + + public static func verifyAPI(apiKey: String) async throws -> DoubaoUsageSnapshot { + do { + let result = try await OpenAICompatibleVerifier.verify( + baseURL: self.baseURL, + apiKey: apiKey, + logger: self.log) + return DoubaoUsageSnapshot( + isConnected: result.isConnected, + modelCount: result.modelCount, + modelNames: result.modelNames, + updatedAt: result.verifiedAt) + } catch let error as OpenAICompatibleVerifier.VerificationError { + switch error { + case .missingCredentials: + throw DoubaoUsageError.missingCredentials + case let .networkError(message): + throw DoubaoUsageError.networkError(message) + case let .apiError(message): + throw DoubaoUsageError.apiError(message) + case let .parseFailed(message): + throw DoubaoUsageError.parseFailed(message) + } + } + } + + public static func fetchBalance(cookieHeader: String, now: Date = Date()) async throws -> DoubaoUsageSnapshot { + var request = URLRequest(url: self.balanceURL) + request.httpMethod = "GET" + request.setValue(cookieHeader, forHTTPHeaderField: "Cookie") + request.setValue("application/json", forHTTPHeaderField: "Accept") + let userAgent = "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" + request.setValue(userAgent, forHTTPHeaderField: "User-Agent") + request.setValue("https://console.volcengine.com/ark", forHTTPHeaderField: "Referer") + + let (data, response) = try await URLSession.shared.data(for: request) + guard let httpResponse = response as? HTTPURLResponse else { + throw DoubaoUsageError.networkError("Invalid response") + } + + guard httpResponse.statusCode == 200 else { + let body = String(data: data, encoding: .utf8) ?? "" + Self.webLog.error("Doubao balance API returned \(httpResponse.statusCode): \(body)") + if httpResponse.statusCode == 401 || httpResponse.statusCode == 403 { + throw DoubaoUsageError.apiError("Cookie expired or invalid") + } + throw DoubaoUsageError.apiError("HTTP \(httpResponse.statusCode)") + } + + let balanceInfo = try self.parseBalanceResponse(data: data) + return DoubaoUsageSnapshot( + isConnected: true, + modelCount: 0, + modelNames: [], + balanceInfo: balanceInfo, + updatedAt: now) + } + + static func _parseBalanceForTesting(_ data: Data) throws -> DoubaoBalanceInfo { + try self.parseBalanceResponse(data: data) + } + + private static func parseBalanceResponse(data: Data) throws -> DoubaoBalanceInfo { + guard let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] else { + throw DoubaoUsageError.parseFailed("Invalid JSON") + } + + let root = (json["data"] as? [String: Any]) ?? json + + let available = self.extractDouble(from: root, keys: [ + "available_balance", "availableBalance", "available_quota", + "remain_quota", "balance", "remain"]) + let used = self.extractDouble(from: root, keys: [ + "used_balance", "usedBalance", "used_quota", "used"]) + let total = self.extractDouble(from: root, keys: [ + "total_balance", "totalBalance", "total_quota", "total", "quota"]) + + let totalBalance = total ?? (available ?? 0) + (used ?? 0) + + return DoubaoBalanceInfo( + availableBalance: available ?? 0, + usedBalance: used ?? max(0, totalBalance - (available ?? 0)), + totalBalance: totalBalance) + } + + private static func extractDouble(from dict: [String: Any], keys: [String]) -> Double? { + for key in keys { + if let value = dict[key] { + if let d = value as? Double { return d } + if let i = value as? Int { return Double(i) } + if let s = value as? String, let d = Double(s) { return d } + } + } + return nil + } +} diff --git a/Sources/CodexBarCore/Providers/Ernie/ErnieCookieHeader.swift b/Sources/CodexBarCore/Providers/Ernie/ErnieCookieHeader.swift new file mode 100644 index 000000000..33eba0060 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Ernie/ErnieCookieHeader.swift @@ -0,0 +1,64 @@ +import Foundation + +public struct ErnieCookieOverride: Sendable { + public let cookieHeader: String + + public init(cookieHeader: String) { + self.cookieHeader = cookieHeader + } +} + +public enum ErnieCookieHeader { + private static let log = CodexBarLog.logger(LogCategories.ernieCookie) + private static let headerPatterns: [String] = [ + #"(?i)-H\s*'Cookie:\s*([^']+)'"#, + #"(?i)-H\s*"Cookie:\s*([^"]+)""#, + #"(?i)\bcookie:\s*'([^']+)'"#, + #"(?i)\bcookie:\s*"([^"]+)""#, + #"(?i)\bcookie:\s*([^\r\n]+)"#, + ] + + public static func resolveCookieOverride(context: ProviderFetchContext) -> ErnieCookieOverride? { + if let settings = context.settings?.ernie, settings.cookieSource == .manual { + if let manual = settings.manualCookieHeader, !manual.isEmpty { + return self.override(from: manual) + } + } + + if let envToken = self.override(from: context.env["ERNIE_MANUAL_COOKIE"]) { + return envToken + } + + return nil + } + + public static func override(from raw: String?) -> ErnieCookieOverride? { + guard let raw = raw?.trimmingCharacters(in: .whitespacesAndNewlines), !raw.isEmpty else { + return nil + } + + if raw.contains("=") { + return ErnieCookieOverride(cookieHeader: raw) + } + + if let cookieHeader = self.extractHeader(from: raw) { + return ErnieCookieOverride(cookieHeader: cookieHeader) + } + + return nil + } + + private static func extractHeader(from raw: String) -> String? { + for pattern in self.headerPatterns { + guard let regex = try? NSRegularExpression(pattern: pattern, options: []) else { continue } + let range = NSRange(raw.startIndex..= 2, + let captureRange = Range(match.range(at: 1), in: raw) + else { continue } + let captured = String(raw[captureRange]).trimmingCharacters(in: .whitespacesAndNewlines) + if !captured.isEmpty { return captured } + } + return nil + } +} diff --git a/Sources/CodexBarCore/Providers/Ernie/ErnieCookieImporter.swift b/Sources/CodexBarCore/Providers/Ernie/ErnieCookieImporter.swift new file mode 100644 index 000000000..1a511baa5 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Ernie/ErnieCookieImporter.swift @@ -0,0 +1,169 @@ +import Foundation + +#if os(macOS) +import SweetCookieKit + +public enum ErnieCookieImporter { + private static let log = CodexBarLog.logger(LogCategories.ernieCookie) + private static let cookieClient = BrowserCookieClient() + private static let cookieDomains = ["console.bce.baidu.com"] + private static let cookieImportOrder: BrowserCookieImportOrder = + ProviderDefaults.metadata[.ernie]?.browserCookieOrder ?? Browser.defaultImportOrder + + public struct SessionInfo: Sendable { + public let cookies: [HTTPCookie] + public let sourceLabel: String + + public init(cookies: [HTTPCookie], sourceLabel: String) { + self.cookies = cookies + self.sourceLabel = sourceLabel + } + + public var cookieHeader: String? { + let parts = self.cookies.compactMap { cookie -> String? in + "\(cookie.name)=\(cookie.value)" + } + guard !parts.isEmpty else { return nil } + return parts.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 ErnieCookieImportError.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.codexBarRecords( + 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 } + + log("Found \(httpCookies.count) cookies for console.bce.baidu.com 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 ErnieCookieImportError.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?("[ernie-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 ErnieCookieImportError: LocalizedError { + case noCookies + + var errorDescription: String? { + switch self { + case .noCookies: "No ERNIE session cookies found in browsers." + } + } +} +#endif diff --git a/Sources/CodexBarCore/Providers/Ernie/ErnieProviderDescriptor.swift b/Sources/CodexBarCore/Providers/Ernie/ErnieProviderDescriptor.swift new file mode 100644 index 000000000..d4cb49fbc --- /dev/null +++ b/Sources/CodexBarCore/Providers/Ernie/ErnieProviderDescriptor.swift @@ -0,0 +1,128 @@ +import CodexBarMacroSupport +import Foundation + +@ProviderDescriptorRegistration +@ProviderDescriptorDefinition +public enum ErnieProviderDescriptor { + static func makeDescriptor() -> ProviderDescriptor { + ProviderDescriptor( + id: .ernie, + metadata: ProviderMetadata( + id: .ernie, + displayName: "ERNIE", + sessionLabel: "Status", + weeklyLabel: "Models", + opusLabel: nil, + supportsOpus: false, + supportsCredits: false, + creditsHint: "", + toggleTitle: "Show ERNIE usage", + cliName: "ernie", + defaultEnabled: false, + isPrimaryProvider: false, + usesAccountFallback: false, + browserCookieOrder: nil, + dashboardURL: "https://console.bce.baidu.com/qianfan/overview", + statusPageURL: nil), + branding: ProviderBranding( + iconStyle: .ernie, + iconResourceName: "ProviderIcon-ernie", + color: ProviderColor(red: 0.15, green: 0.45, blue: 0.85)), + tokenCost: ProviderTokenCostConfig( + supportsTokenCost: false, + noDataMessage: { "ERNIE cost summary is not available via API." }), + fetchPlan: ProviderFetchPlan( + sourceModes: [.auto, .web, .api], + pipeline: ProviderFetchPipeline(resolveStrategies: { _ in + [ErnieWebFetchStrategy(), ErnieAPIFetchStrategy()] + })), + cli: ProviderCLIConfig( + name: "ernie", + aliases: ["qianfan", "wenxin", "baidu"], + versionDetector: nil)) + } +} + +struct ErnieAPIFetchStrategy: ProviderFetchStrategy { + let id: String = "ernie.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 ErnieUsageError.missingCredentials + } + let usage = try await ErnieUsageFetcher.verifyAPI(apiKey: apiKey) + return self.makeResult( + usage: usage.toUsageSnapshot(), + sourceLabel: "api") + } + + func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool { + true + } + + private static func resolveToken(environment: [String: String]) -> String? { + ProviderTokenResolver.ernieToken(environment: environment) + } +} + +struct ErnieWebFetchStrategy: ProviderFetchStrategy { + let id: String = "ernie.web" + let kind: ProviderFetchKind = .web + private static let log = CodexBarLog.logger(LogCategories.ernieWeb) + + func isAvailable(_ context: ProviderFetchContext) async -> Bool { + if ErnieCookieHeader.resolveCookieOverride(context: context) != nil { + return true + } + + #if os(macOS) + if context.settings?.ernie?.cookieSource != .off { + return ErnieCookieImporter.hasSession() + } + #endif + + return false + } + + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + guard let cookieHeader = self.resolveCookieHeader(context: context) else { + throw ErnieUsageError.missingCookie + } + + let snapshot = try await ErnieUsageFetcher.fetchBalance(cookieHeader: cookieHeader) + return self.makeResult( + usage: snapshot.toUsageSnapshot(), + sourceLabel: "web") + } + + func shouldFallback(on error: Error, context: ProviderFetchContext) -> Bool { + if case ErnieUsageError.missingCookie = error { return false } + return true + } + + private func resolveCookieHeader(context: ProviderFetchContext) -> String? { + if let override = ErnieCookieHeader.resolveCookieOverride(context: context) { + return override.cookieHeader + } + + #if os(macOS) + if context.settings?.ernie?.cookieSource != .off { + do { + let session = try ErnieCookieImporter.importSession() + if let header = session.cookieHeader { + return header + } + } catch { + Self.log.debug("Ernie browser cookie import failed: \(error)") + } + } + #endif + + return nil + } +} diff --git a/Sources/CodexBarCore/Providers/Ernie/ErnieSettingsReader.swift b/Sources/CodexBarCore/Providers/Ernie/ErnieSettingsReader.swift new file mode 100644 index 000000000..b63d8310d --- /dev/null +++ b/Sources/CodexBarCore/Providers/Ernie/ErnieSettingsReader.swift @@ -0,0 +1,34 @@ +import Foundation + +public struct ErnieSettingsReader: Sendable { + public static let apiKeyEnvironmentKey = "QIANFAN_API_KEY" + public static let apiKeyEnvironmentKeys = [Self.apiKeyEnvironmentKey, "ERNIE_API_KEY", "BAIDU_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/Ernie/ErnieUsageFetcher.swift b/Sources/CodexBarCore/Providers/Ernie/ErnieUsageFetcher.swift new file mode 100644 index 000000000..d639d2030 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Ernie/ErnieUsageFetcher.swift @@ -0,0 +1,231 @@ +import Foundation +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif + +// MARK: - Domain snapshot + +public struct ErnieUsageSnapshot: Sendable { + public let isConnected: Bool + public let modelCount: Int + public let modelNames: [String] + public let balanceInfo: ErnieBalanceInfo? + public let updatedAt: Date + + public init( + isConnected: Bool, + modelCount: Int, + modelNames: [String], + balanceInfo: ErnieBalanceInfo? = nil, + updatedAt: Date) + { + self.isConnected = isConnected + self.modelCount = modelCount + self.modelNames = modelNames + self.balanceInfo = balanceInfo + self.updatedAt = updatedAt + } + + public func toUsageSnapshot() -> UsageSnapshot { + let detail: String + let usedPercent: Double + if let balance = self.balanceInfo { + let remaining = balance.availableBalance + let total = balance.totalBalance + usedPercent = total > 0 ? max(0, min(100, (total - remaining) / total * 100)) : 0 + detail = String(format: "¥%.2f / ¥%.2f", remaining, total) + } else if self.isConnected { + let names = self.modelNames.joined(separator: ", ") + detail = "API Connected — \(self.modelCount) models (\(names))" + usedPercent = 0 + } else { + detail = "API not connected" + usedPercent = 100 + } + + let identity = ProviderIdentitySnapshot( + providerID: .ernie, + accountEmail: nil, + accountOrganization: nil, + loginMethod: nil) + + let primaryWindow = RateWindow( + usedPercent: usedPercent, + windowMinutes: nil, + resetsAt: nil, + resetDescription: detail) + + var secondaryWindow: RateWindow? + if let balance = self.balanceInfo { + secondaryWindow = RateWindow( + usedPercent: 0, + windowMinutes: nil, + resetsAt: nil, + resetDescription: "Used: ¥\(String(format: "%.2f", balance.usedBalance))") + } + + let providerCost: ProviderCostSnapshot? + if let balance = self.balanceInfo { + providerCost = ProviderCostSnapshot( + used: balance.usedBalance, + limit: balance.totalBalance, + currencyCode: "CNY", + period: nil, + resetsAt: nil, + updatedAt: self.updatedAt) + } else { + providerCost = nil + } + + return UsageSnapshot( + primary: primaryWindow, + secondary: secondaryWindow, + tertiary: nil, + providerCost: providerCost, + updatedAt: self.updatedAt, + identity: identity) + } +} + +public struct ErnieBalanceInfo: Sendable { + public let availableBalance: Double + public let usedBalance: Double + public let totalBalance: Double + + public init(availableBalance: Double, usedBalance: Double, totalBalance: Double) { + self.availableBalance = availableBalance + self.usedBalance = usedBalance + self.totalBalance = totalBalance + } +} + +// MARK: - Errors + +public enum ErnieUsageError: LocalizedError, Sendable { + case missingCredentials + case missingCookie + case networkError(String) + case apiError(String) + case parseFailed(String) + + public var errorDescription: String? { + switch self { + case .missingCredentials: + "Missing ERNIE API key." + case .missingCookie: + "Missing ERNIE session cookie." + case let .networkError(message): + "ERNIE network error: \(message)" + case let .apiError(message): + "ERNIE API error: \(message)" + case let .parseFailed(message): + "Failed to parse ERNIE response: \(message)" + } + } +} + +// MARK: - Fetcher + +public struct ErnieUsageFetcher: Sendable { + private static let log = CodexBarLog.logger(LogCategories.qianfanUsage) + private static let webLog = CodexBarLog.logger(LogCategories.ernieWeb) + private static let baseURL = URL(string: "https://qianfan.baidubce.com/v2")! + private static let balanceURL = + URL(string: "https://console.bce.baidu.com/qianfan/api/resourcepack/v1/billing/query")! + + public static func verifyAPI(apiKey: String) async throws -> ErnieUsageSnapshot { + do { + let result = try await OpenAICompatibleVerifier.verify( + baseURL: self.baseURL, + apiKey: apiKey, + logger: self.log) + return ErnieUsageSnapshot( + isConnected: result.isConnected, + modelCount: result.modelCount, + modelNames: result.modelNames, + updatedAt: result.verifiedAt) + } catch let error as OpenAICompatibleVerifier.VerificationError { + switch error { + case .missingCredentials: + throw ErnieUsageError.missingCredentials + case let .networkError(message): + throw ErnieUsageError.networkError(message) + case let .apiError(message): + throw ErnieUsageError.apiError(message) + case let .parseFailed(message): + throw ErnieUsageError.parseFailed(message) + } + } + } + + public static func fetchBalance(cookieHeader: String, now: Date = Date()) async throws -> ErnieUsageSnapshot { + var request = URLRequest(url: self.balanceURL) + request.httpMethod = "GET" + request.setValue(cookieHeader, forHTTPHeaderField: "Cookie") + request.setValue("application/json", forHTTPHeaderField: "Accept") + let userAgent = "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" + request.setValue(userAgent, forHTTPHeaderField: "User-Agent") + request.setValue("https://console.bce.baidu.com/qianfan/overview", forHTTPHeaderField: "Referer") + + let (data, response) = try await URLSession.shared.data(for: request) + guard let httpResponse = response as? HTTPURLResponse else { + throw ErnieUsageError.networkError("Invalid response") + } + + guard httpResponse.statusCode == 200 else { + let body = String(data: data, encoding: .utf8) ?? "" + Self.webLog.error("ERNIE balance API returned \(httpResponse.statusCode): \(body)") + if httpResponse.statusCode == 401 || httpResponse.statusCode == 403 { + throw ErnieUsageError.apiError("Cookie expired or invalid") + } + throw ErnieUsageError.apiError("HTTP \(httpResponse.statusCode)") + } + + let balanceInfo = try self.parseBalanceResponse(data: data) + return ErnieUsageSnapshot( + isConnected: true, + modelCount: 0, + modelNames: [], + balanceInfo: balanceInfo, + updatedAt: now) + } + + static func _parseBalanceForTesting(_ data: Data) throws -> ErnieBalanceInfo { + try self.parseBalanceResponse(data: data) + } + + private static func parseBalanceResponse(data: Data) throws -> ErnieBalanceInfo { + guard let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] else { + throw ErnieUsageError.parseFailed("Invalid JSON") + } + + let root = (json["data"] as? [String: Any]) ?? json + + let available = self.extractDouble(from: root, keys: [ + "available_balance", "availableBalance", "available_quota", + "remain_quota", "balance", "remain"]) + let used = self.extractDouble(from: root, keys: [ + "used_balance", "usedBalance", "used_quota", "used"]) + let total = self.extractDouble(from: root, keys: [ + "total_balance", "totalBalance", "total_quota", "total", "quota"]) + + let totalBalance = total ?? (available ?? 0) + (used ?? 0) + + return ErnieBalanceInfo( + availableBalance: available ?? 0, + usedBalance: used ?? max(0, totalBalance - (available ?? 0)), + totalBalance: totalBalance) + } + + private static func extractDouble(from dict: [String: Any], keys: [String]) -> Double? { + for key in keys { + if let value = dict[key] { + if let d = value as? Double { return d } + if let i = value as? Int { return Double(i) } + if let s = value as? String, let d = Double(s) { return d } + } + } + return nil + } +} diff --git a/Sources/CodexBarCore/Providers/KimiK2/KimiK2UsageFetcher.swift b/Sources/CodexBarCore/Providers/KimiK2/KimiK2UsageFetcher.swift index 775e109bf..4e43c1d60 100644 --- a/Sources/CodexBarCore/Providers/KimiK2/KimiK2UsageFetcher.swift +++ b/Sources/CodexBarCore/Providers/KimiK2/KimiK2UsageFetcher.swift @@ -3,15 +3,29 @@ import Foundation import FoundationNetworking #endif +public struct KimiK2BalanceInfo: Sendable { + public let availableBalance: Double + public let voucherBalance: Double + public let cashBalance: Double + + public init(availableBalance: Double, voucherBalance: Double, cashBalance: Double) { + self.availableBalance = availableBalance + self.voucherBalance = voucherBalance + self.cashBalance = cashBalance + } +} + public struct KimiK2UsageSnapshot: Sendable { public let summary: KimiK2UsageSummary + public let balanceInfo: KimiK2BalanceInfo? - public init(summary: KimiK2UsageSummary) { + public init(summary: KimiK2UsageSummary, balanceInfo: KimiK2BalanceInfo? = nil) { self.summary = summary + self.balanceInfo = balanceInfo } public func toUsageSnapshot() -> UsageSnapshot { - self.summary.toUsageSnapshot() + self.summary.toUsageSnapshot(balanceInfo: self.balanceInfo) } } @@ -28,7 +42,7 @@ public struct KimiK2UsageSummary: Sendable { self.updatedAt = updatedAt } - public func toUsageSnapshot() -> UsageSnapshot { + public func toUsageSnapshot(balanceInfo: KimiK2BalanceInfo? = nil) -> UsageSnapshot { let total = max(0, self.consumed + self.remaining) let usedPercent: Double = if total > 0 { min(100, max(0, (self.consumed / total) * 100)) @@ -42,6 +56,20 @@ public struct KimiK2UsageSummary: Sendable { windowMinutes: nil, resetsAt: nil, resetDescription: total > 0 ? "Credits: \(usedText)/\(totalText)" : nil) + + var balanceWindow: RateWindow? + if let balance = balanceInfo { + let balancePercent = balance.availableBalance > 0 + ? min(100, max(0, (balance.cashBalance / (balance.cashBalance + balance.availableBalance)) * 100)) + : 0 + let balanceText = String(format: "%.2f", balance.availableBalance) + balanceWindow = RateWindow( + usedPercent: balancePercent, + windowMinutes: nil, + resetsAt: nil, + resetDescription: "Balance: ¥\(balanceText)") + } + let identity = ProviderIdentitySnapshot( providerID: .kimik2, accountEmail: nil, @@ -49,7 +77,7 @@ public struct KimiK2UsageSummary: Sendable { loginMethod: nil) return UsageSnapshot( primary: rateWindow, - secondary: nil, + secondary: balanceWindow, tertiary: nil, providerCost: nil, updatedAt: self.updatedAt, @@ -80,6 +108,7 @@ public enum KimiK2UsageError: LocalizedError, Sendable { public struct KimiK2UsageFetcher: Sendable { private static let log = CodexBarLog.logger(LogCategories.kimiK2Usage) private static let creditsURL = URL(string: "https://kimi-k2.ai/api/user/credits")! + private static let balanceURL = URL(string: "https://api.moonshot.cn/v1/users/me/balance")! private static let jsonSerializer = JSONSerialization.self private static let consumedPaths: [[String]] = [ ["total_credits_consumed"], @@ -153,13 +182,88 @@ public struct KimiK2UsageFetcher: Sendable { } let summary = try Self.parseSummary(data: data, headers: httpResponse.allHeaderFields) - return KimiK2UsageSnapshot(summary: summary) + + var balanceInfo: KimiK2BalanceInfo? + do { + balanceInfo = try await Self.fetchBalance(apiKey: apiKey) + } catch { + Self.log.debug("Kimi balance fetch failed (non-fatal): \(error)") + } + + return KimiK2UsageSnapshot(summary: summary, balanceInfo: balanceInfo) + } + + public static func fetchBalance(apiKey: String) async throws -> KimiK2BalanceInfo? { + guard !apiKey.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty else { + return nil + } + + var request = URLRequest(url: self.balanceURL) + request.httpMethod = "GET" + request.setValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization") + request.setValue("application/json", forHTTPHeaderField: "Accept") + + let (data, response) = try await URLSession.shared.data(for: request) + + guard let httpResponse = response as? HTTPURLResponse else { + throw KimiK2UsageError.networkError("Invalid balance response") + } + + guard httpResponse.statusCode == 200 else { + let body = String(data: data, encoding: .utf8) ?? "HTTP \(httpResponse.statusCode)" + Self.log.error("Kimi balance API returned \(httpResponse.statusCode): \(body)") + return nil + } + + return try Self.parseBalance(data: data) + } + + private static func parseBalance(data: Data) throws -> KimiK2BalanceInfo? { + guard let json = try? jsonSerializer.jsonObject(with: data), + let dict = json as? [String: Any] + else { + throw KimiK2UsageError.parseFailed("Balance JSON is not an object") + } + + let dataContext: [String: Any] + if let data = dict["data"] as? [String: Any] { + dataContext = data + } else { + dataContext = dict + } + + let availableBalancePaths: [[String]] = [ + ["available_balance"], + ["availableBalance"], + ] + let voucherBalancePaths: [[String]] = [ + ["voucher_balance"], + ["voucherBalance"], + ] + let cashBalancePaths: [[String]] = [ + ["cash_balance"], + ["cashBalance"], + ] + + let allContexts = [dataContext, dict] + let available = Self.doubleValue(for: availableBalancePaths, in: allContexts) ?? 0 + let voucher = Self.doubleValue(for: voucherBalancePaths, in: allContexts) ?? 0 + let cash = Self.doubleValue(for: cashBalancePaths, in: allContexts) ?? 0 + + return KimiK2BalanceInfo( + availableBalance: available, + voucherBalance: voucher, + cashBalance: cash) } static func _parseSummaryForTesting(_ data: Data, headers: [AnyHashable: Any] = [:]) throws -> KimiK2UsageSummary { try self.parseSummary(data: data, headers: headers) } + static func _parseBalanceForTesting(_ data: Data) throws -> KimiK2BalanceInfo? { + try self.parseBalance(data: data) + } + private static func parseSummary(data: Data, headers: [AnyHashable: Any]) throws -> KimiK2UsageSummary { guard let json = try? jsonSerializer.jsonObject(with: data), let dictionary = json as? [String: Any] diff --git a/Sources/CodexBarCore/Providers/MiMo/MiMoCookieHeader.swift b/Sources/CodexBarCore/Providers/MiMo/MiMoCookieHeader.swift new file mode 100644 index 000000000..bd8ab2b33 --- /dev/null +++ b/Sources/CodexBarCore/Providers/MiMo/MiMoCookieHeader.swift @@ -0,0 +1,68 @@ +import Foundation + +public struct MiMoCookieOverride: Sendable { + public let cookieHeader: String + + public init(cookieHeader: String) { + self.cookieHeader = cookieHeader + } +} + +public enum MiMoCookieHeader { + private static let log = CodexBarLog.logger(LogCategories.mimoCookie) + private static let headerPatterns: [String] = [ + #"(?i)-H\s*'Cookie:\s*([^']+)'"#, + #"(?i)-H\s*"Cookie:\s*([^"]+)""#, + #"(?i)\bcookie:\s*'([^']+)'"#, + #"(?i)\bcookie:\s*"([^"]+)""#, + #"(?i)\bcookie:\s*([^\r\n]+)"#, + ] + + public static func resolveCookieOverride(context: ProviderFetchContext) -> MiMoCookieOverride? { + if let settings = context.settings?.mimo, settings.cookieSource == .manual { + if let manual = settings.manualCookieHeader, !manual.isEmpty { + return self.override(from: manual) + } + } + + if let envToken = self.override(from: context.env["MIMO_MANUAL_COOKIE"]) { + return envToken + } + + return nil + } + + public static func override(from raw: String?) -> MiMoCookieOverride? { + guard let raw = raw?.trimmingCharacters(in: .whitespacesAndNewlines), !raw.isEmpty else { + return nil + } + + // If it looks like a Cookie header value (contains = and ;) + if raw.contains("=") { + return MiMoCookieOverride(cookieHeader: raw) + } + + // Try extracting from curl-style or other patterns + if let cookieHeader = self.extractHeader(from: raw) { + return MiMoCookieOverride(cookieHeader: cookieHeader) + } + + return nil + } + + private static func extractHeader(from raw: String) -> String? { + for pattern in self.headerPatterns { + guard let regex = try? NSRegularExpression(pattern: pattern, options: []) else { continue } + let range = NSRange(raw.startIndex..= 2, + let captureRange = Range(match.range(at: 1), in: raw) + else { + continue + } + let captured = String(raw[captureRange]).trimmingCharacters(in: .whitespacesAndNewlines) + if !captured.isEmpty { return captured } + } + return nil + } +} diff --git a/Sources/CodexBarCore/Providers/MiMo/MiMoCookieImporter.swift b/Sources/CodexBarCore/Providers/MiMo/MiMoCookieImporter.swift new file mode 100644 index 000000000..94d00c3fc --- /dev/null +++ b/Sources/CodexBarCore/Providers/MiMo/MiMoCookieImporter.swift @@ -0,0 +1,178 @@ +import Foundation + +#if os(macOS) +import SweetCookieKit + +public enum MiMoCookieImporter { + private static let log = CodexBarLog.logger(LogCategories.mimoCookie) + private static let cookieClient = BrowserCookieClient() + private static let cookieDomains = ["platform.xiaomimimo.com", "xiaomimimo.com"] + private static let cookieImportOrder: BrowserCookieImportOrder = + ProviderDefaults.metadata[.mimo]?.browserCookieOrder ?? Browser.defaultImportOrder + + public struct SessionInfo: Sendable { + public let cookies: [HTTPCookie] + public let sourceLabel: String + + public init(cookies: [HTTPCookie], sourceLabel: String) { + self.cookies = cookies + self.sourceLabel = sourceLabel + } + + public var cookieHeader: String? { + let parts = self.cookies.compactMap { cookie -> String? in + "\(cookie.name)=\(cookie.value)" + } + guard !parts.isEmpty else { return nil } + return parts.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 MiMoCookieImportError.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.codexBarRecords( + 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 } + + log("Found \(httpCookies.count) cookies for platform.xiaomimimo.com 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 MiMoCookieImportError.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?("[mimo-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 MiMoCookieImportError: LocalizedError { + case noCookies + + var errorDescription: String? { + switch self { + case .noCookies: + "No MiMo session cookies found in browsers." + } + } +} +#endif diff --git a/Sources/CodexBarCore/Providers/MiMo/MiMoProviderDescriptor.swift b/Sources/CodexBarCore/Providers/MiMo/MiMoProviderDescriptor.swift new file mode 100644 index 000000000..724ac13ba --- /dev/null +++ b/Sources/CodexBarCore/Providers/MiMo/MiMoProviderDescriptor.swift @@ -0,0 +1,128 @@ +import CodexBarMacroSupport +import Foundation + +@ProviderDescriptorRegistration +@ProviderDescriptorDefinition +public enum MiMoProviderDescriptor { + static func makeDescriptor() -> ProviderDescriptor { + ProviderDescriptor( + id: .mimo, + metadata: ProviderMetadata( + id: .mimo, + displayName: "MiMo", + sessionLabel: "Status", + weeklyLabel: "Models", + opusLabel: nil, + supportsOpus: false, + supportsCredits: false, + creditsHint: "", + toggleTitle: "Show MiMo usage", + cliName: "mimo", + defaultEnabled: false, + isPrimaryProvider: false, + usesAccountFallback: false, + browserCookieOrder: nil, + dashboardURL: "https://platform.xiaomimimo.com/console/balance", + statusPageURL: nil), + branding: ProviderBranding( + iconStyle: .mimo, + iconResourceName: "ProviderIcon-mimo", + color: ProviderColor(red: 1.0, green: 0.55, blue: 0.0)), + tokenCost: ProviderTokenCostConfig( + supportsTokenCost: false, + noDataMessage: { "MiMo cost summary is not available via API." }), + fetchPlan: ProviderFetchPlan( + sourceModes: [.auto, .web, .api], + pipeline: ProviderFetchPipeline(resolveStrategies: { _ in + [MiMoWebFetchStrategy(), MiMoAPIFetchStrategy()] + })), + cli: ProviderCLIConfig( + name: "mimo", + aliases: ["xiaomi", "mimo-v2"], + versionDetector: nil)) + } +} + +struct MiMoAPIFetchStrategy: ProviderFetchStrategy { + let id: String = "mimo.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 MiMoUsageError.missingCredentials + } + let usage = try await MiMoUsageFetcher.verifyAPI(apiKey: apiKey) + return self.makeResult( + usage: usage.toUsageSnapshot(), + sourceLabel: "api") + } + + func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool { + true + } + + private static func resolveToken(environment: [String: String]) -> String? { + ProviderTokenResolver.mimoToken(environment: environment) + } +} + +struct MiMoWebFetchStrategy: ProviderFetchStrategy { + let id: String = "mimo.web" + let kind: ProviderFetchKind = .web + private static let log = CodexBarLog.logger(LogCategories.mimoWeb) + + func isAvailable(_ context: ProviderFetchContext) async -> Bool { + if MiMoCookieHeader.resolveCookieOverride(context: context) != nil { + return true + } + + #if os(macOS) + if context.settings?.mimo?.cookieSource != .off { + return MiMoCookieImporter.hasSession() + } + #endif + + return false + } + + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + guard let cookieHeader = self.resolveCookieHeader(context: context) else { + throw MiMoUsageError.missingCookie + } + + let snapshot = try await MiMoUsageFetcher.fetchBalance(cookieHeader: cookieHeader) + return self.makeResult( + usage: snapshot.toUsageSnapshot(), + sourceLabel: "web") + } + + func shouldFallback(on error: Error, context: ProviderFetchContext) -> Bool { + if case MiMoUsageError.missingCookie = error { return false } + return true + } + + private func resolveCookieHeader(context: ProviderFetchContext) -> String? { + if let override = MiMoCookieHeader.resolveCookieOverride(context: context) { + return override.cookieHeader + } + + #if os(macOS) + if context.settings?.mimo?.cookieSource != .off { + do { + let session = try MiMoCookieImporter.importSession() + if let header = session.cookieHeader { + return header + } + } catch { + Self.log.debug("MiMo browser cookie import failed: \(error)") + } + } + #endif + + return nil + } +} diff --git a/Sources/CodexBarCore/Providers/MiMo/MiMoSettingsReader.swift b/Sources/CodexBarCore/Providers/MiMo/MiMoSettingsReader.swift new file mode 100644 index 000000000..ec6ad1268 --- /dev/null +++ b/Sources/CodexBarCore/Providers/MiMo/MiMoSettingsReader.swift @@ -0,0 +1,34 @@ +import Foundation + +public struct MiMoSettingsReader: Sendable { + public static let apiKeyEnvironmentKey = "MIMO_API_KEY" + public static let apiKeyEnvironmentKeys = [Self.apiKeyEnvironmentKey, "XIAOMI_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/MiMo/MiMoUsageFetcher.swift b/Sources/CodexBarCore/Providers/MiMo/MiMoUsageFetcher.swift new file mode 100644 index 000000000..c971f1dc9 --- /dev/null +++ b/Sources/CodexBarCore/Providers/MiMo/MiMoUsageFetcher.swift @@ -0,0 +1,268 @@ +import Foundation +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif + +// MARK: - Domain snapshot + +public struct MiMoUsageSnapshot: Sendable { + public let isConnected: Bool + public let modelCount: Int + public let modelNames: [String] + public let balanceInfo: MiMoBalanceInfo? + public let updatedAt: Date + + public init( + isConnected: Bool, + modelCount: Int, + modelNames: [String], + balanceInfo: MiMoBalanceInfo? = nil, + updatedAt: Date) + { + self.isConnected = isConnected + self.modelCount = modelCount + self.modelNames = modelNames + self.balanceInfo = balanceInfo + self.updatedAt = updatedAt + } + + public func toUsageSnapshot() -> UsageSnapshot { + let detail: String + let usedPercent: Double + if let balance = self.balanceInfo { + let remaining = balance.availableBalance + let total = balance.totalBalance + usedPercent = total > 0 ? max(0, min(100, (total - remaining) / total * 100)) : 0 + detail = String(format: "¥%.2f / ¥%.2f", remaining, total) + } else if self.isConnected { + let names = self.modelNames.joined(separator: ", ") + detail = "API Connected — \(self.modelCount) models (\(names))" + usedPercent = 0 + } else { + detail = "API not connected" + usedPercent = 100 + } + + let identity = ProviderIdentitySnapshot( + providerID: .mimo, + accountEmail: nil, + accountOrganization: nil, + loginMethod: nil) + + let primaryWindow = RateWindow( + usedPercent: usedPercent, + windowMinutes: nil, + resetsAt: nil, + resetDescription: detail) + + var secondaryWindow: RateWindow? + if let balance = self.balanceInfo { + secondaryWindow = RateWindow( + usedPercent: 0, + windowMinutes: nil, + resetsAt: nil, + resetDescription: "Used: ¥\(String(format: "%.2f", balance.usedBalance))") + } + + let providerCost: ProviderCostSnapshot? + if let balance = self.balanceInfo { + providerCost = ProviderCostSnapshot( + used: balance.usedBalance, + limit: balance.totalBalance, + currencyCode: "CNY", + period: nil, + resetsAt: nil, + updatedAt: self.updatedAt) + } else { + providerCost = nil + } + + return UsageSnapshot( + primary: primaryWindow, + secondary: secondaryWindow, + tertiary: nil, + providerCost: providerCost, + updatedAt: self.updatedAt, + identity: identity) + } +} + +public struct MiMoBalanceInfo: Sendable { + public let availableBalance: Double + public let usedBalance: Double + public let totalBalance: Double + + public init(availableBalance: Double, usedBalance: Double, totalBalance: Double) { + self.availableBalance = availableBalance + self.usedBalance = usedBalance + self.totalBalance = totalBalance + } +} + +// MARK: - Errors + +public enum MiMoUsageError: LocalizedError, Sendable { + case missingCredentials + case missingCookie + case networkError(String) + case apiError(String) + case parseFailed(String) + + public var errorDescription: String? { + switch self { + case .missingCredentials: + "Missing MiMo API key." + case .missingCookie: + "Missing MiMo session cookie." + case let .networkError(message): + "MiMo network error: \(message)" + case let .apiError(message): + "MiMo API error: \(message)" + case let .parseFailed(message): + "Failed to parse MiMo response: \(message)" + } + } +} + +// MARK: - Fetcher + +public struct MiMoUsageFetcher: Sendable { + private static let log = CodexBarLog.logger(LogCategories.mimoUsage) + private static let webLog = CodexBarLog.logger(LogCategories.mimoWeb) + private static let baseURL = URL(string: "https://api.xiaomimimo.com/v1")! + private static let balanceURL = + URL(string: "https://platform.xiaomimimo.com/api/user/balance")! + + public static func verifyAPI(apiKey: String) async throws -> MiMoUsageSnapshot { + guard !apiKey.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty else { + throw MiMoUsageError.missingCredentials + } + + let modelsURL = self.baseURL.appendingPathComponent("models") + var request = URLRequest(url: modelsURL) + request.httpMethod = "GET" + // MiMo 使用 api-key header 而非 Authorization: Bearer + request.setValue(apiKey, forHTTPHeaderField: "api-key") + request.setValue("application/json", forHTTPHeaderField: "Accept") + request.timeoutInterval = 15 + + let (data, response) = try await URLSession.shared.data(for: request) + + guard let httpResponse = response as? HTTPURLResponse else { + throw MiMoUsageError.networkError("Invalid response") + } + + guard httpResponse.statusCode == 200 else { + let body = String(data: data, encoding: .utf8) ?? "" + Self.log.error("MiMo API returned \(httpResponse.statusCode): \(body)") + throw MiMoUsageError.apiError("HTTP \(httpResponse.statusCode)") + } + + return try self.parseModelsResponse(data: data) + } + + public static func fetchBalance( + cookieHeader: String, + now: Date = Date()) async throws -> MiMoUsageSnapshot + { + var request = URLRequest(url: self.balanceURL) + request.httpMethod = "GET" + request.setValue(cookieHeader, forHTTPHeaderField: "Cookie") + request.setValue("application/json", forHTTPHeaderField: "Accept") + let userAgent = "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" + request.setValue(userAgent, forHTTPHeaderField: "User-Agent") + request.setValue("https://platform.xiaomimimo.com/console/balance", forHTTPHeaderField: "Referer") + + let (data, response) = try await URLSession.shared.data(for: request) + guard let httpResponse = response as? HTTPURLResponse else { + throw MiMoUsageError.networkError("Invalid response") + } + + guard httpResponse.statusCode == 200 else { + let body = String(data: data, encoding: .utf8) ?? "" + Self.webLog.error("MiMo balance API returned \(httpResponse.statusCode): \(body)") + if httpResponse.statusCode == 401 || httpResponse.statusCode == 403 { + throw MiMoUsageError.apiError("Cookie expired or invalid") + } + throw MiMoUsageError.apiError("HTTP \(httpResponse.statusCode)") + } + + let balanceInfo = try self.parseBalanceResponse(data: data) + return MiMoUsageSnapshot( + isConnected: true, + modelCount: 0, + modelNames: [], + balanceInfo: balanceInfo, + updatedAt: now) + } + + static func _parseBalanceForTesting(_ data: Data) throws -> MiMoBalanceInfo { + try self.parseBalanceResponse(data: data) + } + + static func _parseModelsForTesting(_ data: Data) throws -> MiMoUsageSnapshot { + try self.parseModelsResponse(data: data) + } + + private static func parseModelsResponse(data: Data) throws -> MiMoUsageSnapshot { + let decoded: ModelsResponse + do { + decoded = try JSONDecoder().decode(ModelsResponse.self, from: data) + } catch { + throw MiMoUsageError.parseFailed(error.localizedDescription) + } + + let modelNames = decoded.data.prefix(5).map(\.id) + return MiMoUsageSnapshot( + isConnected: true, + modelCount: decoded.data.count, + modelNames: Array(modelNames), + updatedAt: Date()) + } + + private static func parseBalanceResponse(data: Data) throws -> MiMoBalanceInfo { + guard let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] else { + throw MiMoUsageError.parseFailed("Invalid JSON") + } + + // Try "data" wrapper first + let root = (json["data"] as? [String: Any]) ?? json + + let available = self.extractDouble(from: root, keys: [ + "available_balance", "availableBalance", "available_quota", + "remain_quota", "balance", "remain"]) + let used = self.extractDouble(from: root, keys: [ + "used_balance", "usedBalance", "used_quota", "used"]) + let total = self.extractDouble(from: root, keys: [ + "total_balance", "totalBalance", "total_quota", "total", "quota"]) + + let totalBalance = total ?? (available ?? 0) + (used ?? 0) + + return MiMoBalanceInfo( + availableBalance: available ?? 0, + usedBalance: used ?? max(0, totalBalance - (available ?? 0)), + totalBalance: totalBalance) + } + + private static func extractDouble(from dict: [String: Any], keys: [String]) -> Double? { + for key in keys { + if let value = dict[key] { + if let d = value as? Double { return d } + if let i = value as? Int { return Double(i) } + if let s = value as? String, let d = Double(s) { return d } + } + } + return nil + } +} + +// MARK: - API response types + +private struct ModelsResponse: Decodable, Sendable { + let data: [ModelEntry] +} + +private struct ModelEntry: Decodable, Sendable { + let id: String +} diff --git a/Sources/CodexBarCore/Providers/ProviderDescriptor.swift b/Sources/CodexBarCore/Providers/ProviderDescriptor.swift index b70de9c3f..9e0518900 100644 --- a/Sources/CodexBarCore/Providers/ProviderDescriptor.swift +++ b/Sources/CodexBarCore/Providers/ProviderDescriptor.swift @@ -82,6 +82,10 @@ public enum ProviderDescriptorRegistry { .mistral: MistralProviderDescriptor.descriptor, .deepseek: DeepSeekProviderDescriptor.descriptor, .codebuff: CodebuffProviderDescriptor.descriptor, + .zhipu: ZhipuProviderDescriptor.descriptor, + .doubao: DoubaoProviderDescriptor.descriptor, + .ernie: ErnieProviderDescriptor.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 819a6d87e..7ca0d6596 100644 --- a/Sources/CodexBarCore/Providers/ProviderSettingsSnapshot.swift +++ b/Sources/CodexBarCore/Providers/ProviderSettingsSnapshot.swift @@ -22,7 +22,11 @@ public struct ProviderSettingsSnapshot: Sendable { jetbrains: JetBrainsProviderSettings? = nil, perplexity: PerplexityProviderSettings? = nil, abacus: AbacusProviderSettings? = nil, - mistral: MistralProviderSettings? = nil) -> ProviderSettingsSnapshot + mistral: MistralProviderSettings? = nil, + zhipu: ZhipuProviderSettings? = nil, + doubao: DoubaoProviderSettings? = nil, + ernie: ErnieProviderSettings? = nil, + mimo: MiMoProviderSettings? = nil) -> ProviderSettingsSnapshot { ProviderSettingsSnapshot( debugMenuEnabled: debugMenuEnabled, @@ -45,7 +49,11 @@ public struct ProviderSettingsSnapshot: Sendable { jetbrains: jetbrains, perplexity: perplexity, abacus: abacus, - mistral: mistral) + mistral: mistral, + zhipu: zhipu, + doubao: doubao, + ernie: ernie, + mimo: mimo) } public struct CodexProviderSettings: Sendable { @@ -192,6 +200,46 @@ public struct ProviderSettingsSnapshot: Sendable { } } + public struct ZhipuProviderSettings: Sendable { + public let cookieSource: ProviderCookieSource + public let manualCookieHeader: String? + + public init(cookieSource: ProviderCookieSource, manualCookieHeader: String?) { + self.cookieSource = cookieSource + self.manualCookieHeader = manualCookieHeader + } + } + + public struct DoubaoProviderSettings: Sendable { + public let cookieSource: ProviderCookieSource + public let manualCookieHeader: String? + + public init(cookieSource: ProviderCookieSource, manualCookieHeader: String?) { + self.cookieSource = cookieSource + self.manualCookieHeader = manualCookieHeader + } + } + + public struct ErnieProviderSettings: Sendable { + public let cookieSource: ProviderCookieSource + public let manualCookieHeader: String? + + public init(cookieSource: ProviderCookieSource, manualCookieHeader: String?) { + self.cookieSource = cookieSource + self.manualCookieHeader = manualCookieHeader + } + } + + 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? @@ -281,6 +329,10 @@ public struct ProviderSettingsSnapshot: Sendable { public let perplexity: PerplexityProviderSettings? public let abacus: AbacusProviderSettings? public let mistral: MistralProviderSettings? + public let zhipu: ZhipuProviderSettings? + public let doubao: DoubaoProviderSettings? + public let ernie: ErnieProviderSettings? + public let mimo: MiMoProviderSettings? public var jetbrainsIDEBasePath: String? { self.jetbrains?.ideBasePath @@ -307,7 +359,11 @@ public struct ProviderSettingsSnapshot: Sendable { jetbrains: JetBrainsProviderSettings? = nil, perplexity: PerplexityProviderSettings? = nil, abacus: AbacusProviderSettings? = nil, - mistral: MistralProviderSettings? = nil) + mistral: MistralProviderSettings? = nil, + zhipu: ZhipuProviderSettings? = nil, + doubao: DoubaoProviderSettings? = nil, + ernie: ErnieProviderSettings? = nil, + mimo: MiMoProviderSettings? = nil) { self.debugMenuEnabled = debugMenuEnabled self.debugKeepCLISessionsAlive = debugKeepCLISessionsAlive @@ -330,6 +386,10 @@ public struct ProviderSettingsSnapshot: Sendable { self.perplexity = perplexity self.abacus = abacus self.mistral = mistral + self.zhipu = zhipu + self.doubao = doubao + self.ernie = ernie + self.mimo = mimo } } @@ -353,6 +413,10 @@ public enum ProviderSettingsSnapshotContribution: Sendable { case perplexity(ProviderSettingsSnapshot.PerplexityProviderSettings) case abacus(ProviderSettingsSnapshot.AbacusProviderSettings) case mistral(ProviderSettingsSnapshot.MistralProviderSettings) + case zhipu(ProviderSettingsSnapshot.ZhipuProviderSettings) + case doubao(ProviderSettingsSnapshot.DoubaoProviderSettings) + case ernie(ProviderSettingsSnapshot.ErnieProviderSettings) + case mimo(ProviderSettingsSnapshot.MiMoProviderSettings) } public struct ProviderSettingsSnapshotBuilder: Sendable { @@ -377,6 +441,10 @@ public struct ProviderSettingsSnapshotBuilder: Sendable { public var perplexity: ProviderSettingsSnapshot.PerplexityProviderSettings? public var abacus: ProviderSettingsSnapshot.AbacusProviderSettings? public var mistral: ProviderSettingsSnapshot.MistralProviderSettings? + public var zhipu: ProviderSettingsSnapshot.ZhipuProviderSettings? + public var doubao: ProviderSettingsSnapshot.DoubaoProviderSettings? + public var ernie: ProviderSettingsSnapshot.ErnieProviderSettings? + public var mimo: ProviderSettingsSnapshot.MiMoProviderSettings? public init(debugMenuEnabled: Bool = false, debugKeepCLISessionsAlive: Bool = false) { self.debugMenuEnabled = debugMenuEnabled @@ -404,6 +472,10 @@ public struct ProviderSettingsSnapshotBuilder: Sendable { case let .perplexity(value): self.perplexity = value case let .abacus(value): self.abacus = value case let .mistral(value): self.mistral = value + case let .zhipu(value): self.zhipu = value + case let .doubao(value): self.doubao = value + case let .ernie(value): self.ernie = value + case let .mimo(value): self.mimo = value } } @@ -429,6 +501,10 @@ public struct ProviderSettingsSnapshotBuilder: Sendable { jetbrains: self.jetbrains, perplexity: self.perplexity, abacus: self.abacus, - mistral: self.mistral) + mistral: self.mistral, + zhipu: self.zhipu, + doubao: self.doubao, + ernie: self.ernie, + mimo: self.mimo) } } diff --git a/Sources/CodexBarCore/Providers/ProviderTokenResolver.swift b/Sources/CodexBarCore/Providers/ProviderTokenResolver.swift index db70aa069..4ea28a8eb 100644 --- a/Sources/CodexBarCore/Providers/ProviderTokenResolver.swift +++ b/Sources/CodexBarCore/Providers/ProviderTokenResolver.swift @@ -77,6 +77,30 @@ public enum ProviderTokenResolver { self.deepseekResolution(environment: environment)?.token } + public static func zhipuToken( + environment: [String: String] = ProcessInfo.processInfo.environment) -> String? + { + self.zhipuResolution(environment: environment)?.token + } + + public static func doubaoToken( + environment: [String: String] = ProcessInfo.processInfo.environment) -> String? + { + self.doubaoResolution(environment: environment)?.token + } + + public static func ernieToken( + environment: [String: String] = ProcessInfo.processInfo.environment) -> String? + { + self.ernieResolution(environment: environment)?.token + } + + public static func mimoToken( + environment: [String: String] = ProcessInfo.processInfo.environment) -> String? + { + self.mimoResolution(environment: environment)?.token + } + public static func deepseekResolution( environment: [String: String] = ProcessInfo.processInfo.environment) -> ProviderTokenResolution? { @@ -90,6 +114,30 @@ public enum ProviderTokenResolver { self.codebuffResolution(environment: environment, authFileURL: authFileURL)?.token } + public static func zhipuResolution( + environment: [String: String] = ProcessInfo.processInfo.environment) -> ProviderTokenResolution? + { + self.resolveEnv(ZhipuSettingsReader.apiKey(environment: environment)) + } + + public static func doubaoResolution( + environment: [String: String] = ProcessInfo.processInfo.environment) -> ProviderTokenResolution? + { + self.resolveEnv(DoubaoSettingsReader.apiKey(environment: environment)) + } + + public static func ernieResolution( + environment: [String: String] = ProcessInfo.processInfo.environment) -> ProviderTokenResolution? + { + self.resolveEnv(ErnieSettingsReader.apiKey(environment: environment)) + } + + public static func mimoResolution( + environment: [String: String] = ProcessInfo.processInfo.environment) -> ProviderTokenResolution? + { + self.resolveEnv(MiMoSettingsReader.apiKey(environment: environment)) + } + public static func zaiResolution( environment: [String: String] = ProcessInfo.processInfo.environment) -> ProviderTokenResolution? { diff --git a/Sources/CodexBarCore/Providers/Providers.swift b/Sources/CodexBarCore/Providers/Providers.swift index 94e5ef82a..70fb168c0 100644 --- a/Sources/CodexBarCore/Providers/Providers.swift +++ b/Sources/CodexBarCore/Providers/Providers.swift @@ -32,6 +32,10 @@ public enum UsageProvider: String, CaseIterable, Sendable, Codable { case mistral case deepseek case codebuff + case zhipu + case doubao + case ernie + case mimo } // swiftformat:enable sortDeclarations @@ -66,6 +70,10 @@ public enum IconStyle: Sendable, CaseIterable { case mistral case deepseek case codebuff + case zhipu + case doubao + case ernie + case mimo case combined } diff --git a/Sources/CodexBarCore/Providers/Shared/OpenAICompatibleVerifier.swift b/Sources/CodexBarCore/Providers/Shared/OpenAICompatibleVerifier.swift new file mode 100644 index 000000000..454505361 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Shared/OpenAICompatibleVerifier.swift @@ -0,0 +1,103 @@ +import Foundation +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif + +public struct OpenAICompatibleVerifier: Sendable { + public struct Result: Sendable { + public let isConnected: Bool + public let modelCount: Int + public let modelNames: [String] + public let verifiedAt: Date + + public init(isConnected: Bool, modelCount: Int, modelNames: [String], verifiedAt: Date) { + self.isConnected = isConnected + self.modelCount = modelCount + self.modelNames = modelNames + self.verifiedAt = verifiedAt + } + } + + public enum VerificationError: LocalizedError, Sendable { + case missingCredentials + case networkError(String) + case apiError(String) + case parseFailed(String) + + public var errorDescription: String? { + switch self { + case .missingCredentials: + "Missing API key." + case let .networkError(message): + "Network error: \(message)" + case let .apiError(message): + "API error: \(message)" + case let .parseFailed(message): + "Failed to parse response: \(message)" + } + } + } + + private static let timeoutSeconds: TimeInterval = 15 + + public static func verify( + baseURL: URL, + apiKey: String, + logger: CodexBarLogger? = nil + ) async throws -> Result { + guard !apiKey.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty else { + throw VerificationError.missingCredentials + } + + let modelsURL = baseURL.appendingPathComponent("models") + var request = URLRequest(url: modelsURL) + request.httpMethod = "GET" + request.setValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization") + request.setValue("application/json", forHTTPHeaderField: "Accept") + request.timeoutInterval = Self.timeoutSeconds + + let (data, response) = try await URLSession.shared.data(for: request) + + guard let httpResponse = response as? HTTPURLResponse else { + throw VerificationError.networkError("Invalid response") + } + + guard httpResponse.statusCode == 200 else { + let body = String(data: data, encoding: .utf8) ?? "" + logger?.error("API returned \(httpResponse.statusCode): \(body)") + throw VerificationError.apiError("HTTP \(httpResponse.statusCode)") + } + + return try Self.parseResult(data: data) + } + + static func _parseResultForTesting(_ data: Data) throws -> Result { + try self.parseResult(data: data) + } + + private static func parseResult(data: Data) throws -> Result { + let decoded: ModelsResponse + do { + decoded = try JSONDecoder().decode(ModelsResponse.self, from: data) + } catch { + throw VerificationError.parseFailed(error.localizedDescription) + } + + let modelNames = decoded.data.prefix(5).map(\.id) + return Result( + isConnected: true, + modelCount: decoded.data.count, + modelNames: Array(modelNames), + verifiedAt: Date()) + } +} + +// MARK: - API response types + +private struct ModelsResponse: Decodable, Sendable { + let data: [ModelEntry] +} + +private struct ModelEntry: Decodable, Sendable { + let id: String +} diff --git a/Sources/CodexBarCore/Providers/Zhipu/ZhipuCookieHeader.swift b/Sources/CodexBarCore/Providers/Zhipu/ZhipuCookieHeader.swift new file mode 100644 index 000000000..7258793f7 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Zhipu/ZhipuCookieHeader.swift @@ -0,0 +1,68 @@ +import Foundation + +public struct ZhipuCookieOverride: Sendable { + public let cookieHeader: String + + public init(cookieHeader: String) { + self.cookieHeader = cookieHeader + } +} + +public enum ZhipuCookieHeader { + private static let log = CodexBarLog.logger(LogCategories.zhipuCookie) + private static let headerPatterns: [String] = [ + #"(?i)-H\s*'Cookie:\s*([^']+)'"#, + #"(?i)-H\s*"Cookie:\s*([^"]+)""#, + #"(?i)\bcookie:\s*'([^']+)'"#, + #"(?i)\bcookie:\s*"([^"]+)""#, + #"(?i)\bcookie:\s*([^\r\n]+)"#, + ] + + public static func resolveCookieOverride(context: ProviderFetchContext) -> ZhipuCookieOverride? { + if let settings = context.settings?.zhipu, settings.cookieSource == .manual { + if let manual = settings.manualCookieHeader, !manual.isEmpty { + return self.override(from: manual) + } + } + + if let envToken = self.override(from: context.env["ZHIPU_MANUAL_COOKIE"]) { + return envToken + } + + return nil + } + + public static func override(from raw: String?) -> ZhipuCookieOverride? { + guard let raw = raw?.trimmingCharacters(in: .whitespacesAndNewlines), !raw.isEmpty else { + return nil + } + + // If it looks like a Cookie header value (contains = and ;) + if raw.contains("=") { + return ZhipuCookieOverride(cookieHeader: raw) + } + + // Try extracting from curl-style or other patterns + if let cookieHeader = self.extractHeader(from: raw) { + return ZhipuCookieOverride(cookieHeader: cookieHeader) + } + + return nil + } + + private static func extractHeader(from raw: String) -> String? { + for pattern in self.headerPatterns { + guard let regex = try? NSRegularExpression(pattern: pattern, options: []) else { continue } + let range = NSRange(raw.startIndex..= 2, + let captureRange = Range(match.range(at: 1), in: raw) + else { + continue + } + let captured = String(raw[captureRange]).trimmingCharacters(in: .whitespacesAndNewlines) + if !captured.isEmpty { return captured } + } + return nil + } +} diff --git a/Sources/CodexBarCore/Providers/Zhipu/ZhipuCookieImporter.swift b/Sources/CodexBarCore/Providers/Zhipu/ZhipuCookieImporter.swift new file mode 100644 index 000000000..0c3d0ad2a --- /dev/null +++ b/Sources/CodexBarCore/Providers/Zhipu/ZhipuCookieImporter.swift @@ -0,0 +1,178 @@ +import Foundation + +#if os(macOS) +import SweetCookieKit + +public enum ZhipuCookieImporter { + private static let log = CodexBarLog.logger(LogCategories.zhipuCookie) + private static let cookieClient = BrowserCookieClient() + private static let cookieDomains = ["open.bigmodel.cn"] + private static let cookieImportOrder: BrowserCookieImportOrder = + ProviderDefaults.metadata[.zhipu]?.browserCookieOrder ?? Browser.defaultImportOrder + + public struct SessionInfo: Sendable { + public let cookies: [HTTPCookie] + public let sourceLabel: String + + public init(cookies: [HTTPCookie], sourceLabel: String) { + self.cookies = cookies + self.sourceLabel = sourceLabel + } + + public var cookieHeader: String? { + let parts = self.cookies.compactMap { cookie -> String? in + "\(cookie.name)=\(cookie.value)" + } + guard !parts.isEmpty else { return nil } + return parts.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 ZhipuCookieImportError.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.codexBarRecords( + 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 } + + log("Found \(httpCookies.count) cookies for open.bigmodel.cn 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 ZhipuCookieImportError.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?("[zhipu-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 ZhipuCookieImportError: LocalizedError { + case noCookies + + var errorDescription: String? { + switch self { + case .noCookies: + "No Zhipu session cookies found in browsers." + } + } +} +#endif diff --git a/Sources/CodexBarCore/Providers/Zhipu/ZhipuProviderDescriptor.swift b/Sources/CodexBarCore/Providers/Zhipu/ZhipuProviderDescriptor.swift new file mode 100644 index 000000000..abb07649b --- /dev/null +++ b/Sources/CodexBarCore/Providers/Zhipu/ZhipuProviderDescriptor.swift @@ -0,0 +1,128 @@ +import CodexBarMacroSupport +import Foundation + +@ProviderDescriptorRegistration +@ProviderDescriptorDefinition +public enum ZhipuProviderDescriptor { + static func makeDescriptor() -> ProviderDescriptor { + ProviderDescriptor( + id: .zhipu, + metadata: ProviderMetadata( + id: .zhipu, + displayName: "Zhipu", + sessionLabel: "Status", + weeklyLabel: "Models", + opusLabel: nil, + supportsOpus: false, + supportsCredits: false, + creditsHint: "", + toggleTitle: "Show Zhipu usage", + cliName: "zhipu", + defaultEnabled: false, + isPrimaryProvider: false, + usesAccountFallback: false, + browserCookieOrder: nil, + dashboardURL: "https://open.bigmodel.cn/usercenter/apikeys", + statusPageURL: nil), + branding: ProviderBranding( + iconStyle: .zhipu, + iconResourceName: "ProviderIcon-zhipu", + color: ProviderColor(red: 0.0, green: 0.45, blue: 0.85)), + tokenCost: ProviderTokenCostConfig( + supportsTokenCost: false, + noDataMessage: { "Zhipu cost summary is not available via API." }), + fetchPlan: ProviderFetchPlan( + sourceModes: [.auto, .web, .api], + pipeline: ProviderFetchPipeline(resolveStrategies: { _ in + [ZhipuWebFetchStrategy(), ZhipuAPIFetchStrategy()] + })), + cli: ProviderCLIConfig( + name: "zhipu", + aliases: ["glm", "chatglm"], + versionDetector: nil)) + } +} + +struct ZhipuAPIFetchStrategy: ProviderFetchStrategy { + let id: String = "zhipu.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 ZhipuUsageError.missingCredentials + } + let usage = try await ZhipuUsageFetcher.verifyAPI(apiKey: apiKey) + return self.makeResult( + usage: usage.toUsageSnapshot(), + sourceLabel: "api") + } + + func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool { + true + } + + private static func resolveToken(environment: [String: String]) -> String? { + ProviderTokenResolver.zhipuToken(environment: environment) + } +} + +struct ZhipuWebFetchStrategy: ProviderFetchStrategy { + let id: String = "zhipu.web" + let kind: ProviderFetchKind = .web + private static let log = CodexBarLog.logger(LogCategories.zhipuWeb) + + func isAvailable(_ context: ProviderFetchContext) async -> Bool { + if ZhipuCookieHeader.resolveCookieOverride(context: context) != nil { + return true + } + + #if os(macOS) + if context.settings?.zhipu?.cookieSource != .off { + return ZhipuCookieImporter.hasSession() + } + #endif + + return false + } + + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + guard let cookieHeader = self.resolveCookieHeader(context: context) else { + throw ZhipuUsageError.missingCookie + } + + let snapshot = try await ZhipuUsageFetcher.fetchBalance(cookieHeader: cookieHeader) + return self.makeResult( + usage: snapshot.toUsageSnapshot(), + sourceLabel: "web") + } + + func shouldFallback(on error: Error, context: ProviderFetchContext) -> Bool { + if case ZhipuUsageError.missingCookie = error { return false } + return true + } + + private func resolveCookieHeader(context: ProviderFetchContext) -> String? { + if let override = ZhipuCookieHeader.resolveCookieOverride(context: context) { + return override.cookieHeader + } + + #if os(macOS) + if context.settings?.zhipu?.cookieSource != .off { + do { + let session = try ZhipuCookieImporter.importSession() + if let header = session.cookieHeader { + return header + } + } catch { + Self.log.debug("Zhipu browser cookie import failed: \(error)") + } + } + #endif + + return nil + } +} diff --git a/Sources/CodexBarCore/Providers/Zhipu/ZhipuSettingsReader.swift b/Sources/CodexBarCore/Providers/Zhipu/ZhipuSettingsReader.swift new file mode 100644 index 000000000..c19c3d829 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Zhipu/ZhipuSettingsReader.swift @@ -0,0 +1,34 @@ +import Foundation + +public struct ZhipuSettingsReader: Sendable { + public static let apiKeyEnvironmentKey = "ZHIPU_API_KEY" + public static let apiKeyEnvironmentKeys = [Self.apiKeyEnvironmentKey, "GLM_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/Zhipu/ZhipuUsageFetcher.swift b/Sources/CodexBarCore/Providers/Zhipu/ZhipuUsageFetcher.swift new file mode 100644 index 000000000..351cd933e --- /dev/null +++ b/Sources/CodexBarCore/Providers/Zhipu/ZhipuUsageFetcher.swift @@ -0,0 +1,233 @@ +import Foundation +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif + +// MARK: - Domain snapshot + +public struct ZhipuUsageSnapshot: Sendable { + public let isConnected: Bool + public let modelCount: Int + public let modelNames: [String] + public let balanceInfo: ZhipuBalanceInfo? + public let updatedAt: Date + + public init( + isConnected: Bool, + modelCount: Int, + modelNames: [String], + balanceInfo: ZhipuBalanceInfo? = nil, + updatedAt: Date) + { + self.isConnected = isConnected + self.modelCount = modelCount + self.modelNames = modelNames + self.balanceInfo = balanceInfo + self.updatedAt = updatedAt + } + + public func toUsageSnapshot() -> UsageSnapshot { + let detail: String + let usedPercent: Double + if let balance = self.balanceInfo { + let remaining = balance.availableBalance + let total = balance.totalBalance + usedPercent = total > 0 ? max(0, min(100, (total - remaining) / total * 100)) : 0 + detail = String(format: "¥%.2f / ¥%.2f", remaining, total) + } else if self.isConnected { + let names = self.modelNames.joined(separator: ", ") + detail = "API Connected — \(self.modelCount) models (\(names))" + usedPercent = 0 + } else { + detail = "API not connected" + usedPercent = 100 + } + + let identity = ProviderIdentitySnapshot( + providerID: .zhipu, + accountEmail: nil, + accountOrganization: nil, + loginMethod: nil) + + let primaryWindow = RateWindow( + usedPercent: usedPercent, + windowMinutes: nil, + resetsAt: nil, + resetDescription: detail) + + var secondaryWindow: RateWindow? + if let balance = self.balanceInfo { + secondaryWindow = RateWindow( + usedPercent: 0, + windowMinutes: nil, + resetsAt: nil, + resetDescription: "Used: ¥\(String(format: "%.2f", balance.usedBalance))") + } + + let providerCost: ProviderCostSnapshot? + if let balance = self.balanceInfo { + providerCost = ProviderCostSnapshot( + used: balance.usedBalance, + limit: balance.totalBalance, + currencyCode: "CNY", + period: nil, + resetsAt: nil, + updatedAt: self.updatedAt) + } else { + providerCost = nil + } + + return UsageSnapshot( + primary: primaryWindow, + secondary: secondaryWindow, + tertiary: nil, + providerCost: providerCost, + updatedAt: self.updatedAt, + identity: identity) + } +} + +public struct ZhipuBalanceInfo: Sendable { + public let availableBalance: Double + public let usedBalance: Double + public let totalBalance: Double + + public init(availableBalance: Double, usedBalance: Double, totalBalance: Double) { + self.availableBalance = availableBalance + self.usedBalance = usedBalance + self.totalBalance = totalBalance + } +} + +// MARK: - Errors + +public enum ZhipuUsageError: LocalizedError, Sendable { + case missingCredentials + case missingCookie + case networkError(String) + case apiError(String) + case parseFailed(String) + + public var errorDescription: String? { + switch self { + case .missingCredentials: + "Missing Zhipu API key." + case .missingCookie: + "Missing Zhipu session cookie." + case let .networkError(message): + "Zhipu network error: \(message)" + case let .apiError(message): + "Zhipu API error: \(message)" + case let .parseFailed(message): + "Failed to parse Zhipu response: \(message)" + } + } +} + +// MARK: - Fetcher + +public struct ZhipuUsageFetcher: Sendable { + private static let log = CodexBarLog.logger(LogCategories.zhipuUsage) + private static let webLog = CodexBarLog.logger(LogCategories.zhipuWeb) + private static let baseURL = URL(string: "https://open.bigmodel.cn/api/paas/v4")! + private static let balanceURL = + URL(string: "https://open.bigmodel.cn/api/paas/v4/users/balance")! + + public static func verifyAPI(apiKey: String) async throws -> ZhipuUsageSnapshot { + do { + let result = try await OpenAICompatibleVerifier.verify( + baseURL: self.baseURL, + apiKey: apiKey, + logger: self.log) + return ZhipuUsageSnapshot( + isConnected: result.isConnected, + modelCount: result.modelCount, + modelNames: result.modelNames, + updatedAt: result.verifiedAt) + } catch let error as OpenAICompatibleVerifier.VerificationError { + switch error { + case .missingCredentials: + throw ZhipuUsageError.missingCredentials + case let .networkError(message): + throw ZhipuUsageError.networkError(message) + case let .apiError(message): + throw ZhipuUsageError.apiError(message) + case let .parseFailed(message): + throw ZhipuUsageError.parseFailed(message) + } + } + } + + public static func fetchBalance(cookieHeader: String, now: Date = Date()) async throws -> ZhipuUsageSnapshot { + var request = URLRequest(url: self.balanceURL) + request.httpMethod = "GET" + request.setValue(cookieHeader, forHTTPHeaderField: "Cookie") + request.setValue("application/json", forHTTPHeaderField: "Accept") + let userAgent = "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" + request.setValue(userAgent, forHTTPHeaderField: "User-Agent") + request.setValue("https://open.bigmodel.cn/usercenter", forHTTPHeaderField: "Referer") + + let (data, response) = try await URLSession.shared.data(for: request) + guard let httpResponse = response as? HTTPURLResponse else { + throw ZhipuUsageError.networkError("Invalid response") + } + + guard httpResponse.statusCode == 200 else { + let body = String(data: data, encoding: .utf8) ?? "" + Self.webLog.error("Zhipu balance API returned \(httpResponse.statusCode): \(body)") + if httpResponse.statusCode == 401 || httpResponse.statusCode == 403 { + throw ZhipuUsageError.apiError("Cookie expired or invalid") + } + throw ZhipuUsageError.apiError("HTTP \(httpResponse.statusCode)") + } + + let balanceInfo = try self.parseBalanceResponse(data: data) + return ZhipuUsageSnapshot( + isConnected: true, + modelCount: 0, + modelNames: [], + balanceInfo: balanceInfo, + updatedAt: now) + } + + static func _parseBalanceForTesting(_ data: Data) throws -> ZhipuBalanceInfo { + try self.parseBalanceResponse(data: data) + } + + private static func parseBalanceResponse(data: Data) throws -> ZhipuBalanceInfo { + guard let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] else { + throw ZhipuUsageError.parseFailed("Invalid JSON") + } + + // Try "data" wrapper first + let root = (json["data"] as? [String: Any]) ?? json + + // Try common field names + let available = self.extractDouble(from: root, keys: [ + "available_balance", "availableBalance", "available_quota", + "remain_quota", "balance", "remain"]) + let used = self.extractDouble(from: root, keys: [ + "used_balance", "usedBalance", "used_quota", "used"]) + let total = self.extractDouble(from: root, keys: [ + "total_balance", "totalBalance", "total_quota", "total", "quota"]) + + let totalBalance = total ?? (available ?? 0) + (used ?? 0) + + return ZhipuBalanceInfo( + availableBalance: available ?? 0, + usedBalance: used ?? max(0, totalBalance - (available ?? 0)), + totalBalance: totalBalance) + } + + private static func extractDouble(from dict: [String: Any], keys: [String]) -> Double? { + for key in keys { + if let value = dict[key] { + if let d = value as? Double { return d } + if let i = value as? Int { return Double(i) } + if let s = value as? String, let d = Double(s) { return d } + } + } + return nil + } +} diff --git a/Sources/CodexBarCore/Proxy/HTTPRequestParser.swift b/Sources/CodexBarCore/Proxy/HTTPRequestParser.swift new file mode 100644 index 000000000..4fb7b7892 --- /dev/null +++ b/Sources/CodexBarCore/Proxy/HTTPRequestParser.swift @@ -0,0 +1,64 @@ +import Foundation + +public struct HTTPParsedRequest: Sendable { + public let method: String + public let path: String + public let host: String + public let headers: [(String, String)] + public let body: Data? + + public init(method: String, path: String, host: String, headers: [(String, String)], body: Data?) { + self.method = method + self.path = path + self.host = host + self.headers = headers + self.body = body + } +} + +public enum HTTPRequestParser { + public static func parse(data: Data) -> HTTPParsedRequest? { + guard let string = String(data: data, encoding: .utf8) else { return nil } + + let parts = string.components(separatedBy: "\r\n\r\n") + guard let headerSection = parts.first else { return nil } + + let lines = headerSection.components(separatedBy: "\r\n") + guard let requestLine = lines.first else { return nil } + + let requestParts = requestLine.split(separator: " ", maxSplits: 2) + guard requestParts.count >= 2 else { return nil } + + let method = String(requestParts[0]) + let path = String(requestParts[1]) + + var headers: [(String, String)] = [] + var host = "" + + for line in lines.dropFirst() { + if let colonIndex = line.firstIndex(of: ":") { + let key = String(line[line.startIndex.. 1 { + let bodyString = parts.dropFirst().joined(separator: "\r\n\r\n") + body = bodyString.data(using: .utf8) + } else { + body = nil + } + + return HTTPParsedRequest( + method: method, + path: path, + host: host, + headers: headers, + body: body) + } +} diff --git a/Sources/CodexBarCore/Proxy/HTTPResponseWriter.swift b/Sources/CodexBarCore/Proxy/HTTPResponseWriter.swift new file mode 100644 index 000000000..f50dbc5fd --- /dev/null +++ b/Sources/CodexBarCore/Proxy/HTTPResponseWriter.swift @@ -0,0 +1,34 @@ +import Foundation + +public enum HTTPResponseWriter { + public static func write(statusCode: Int, headers: [(String, String)], body: Data) -> Data { + let statusText = HTTPURLResponse.localizedString(forStatusCode: statusCode) + var response = "HTTP/1.1 \(statusCode) \(statusText)\r\n" + + var hasContentLength = false + for (key, value) in headers { + response += "\(key): \(value)\r\n" + if key.lowercased() == "content-length" { + hasContentLength = true + } + } + + if !hasContentLength { + response += "Content-Length: \(body.count)\r\n" + } + + response += "\r\n" + + var data = Data(response.utf8) + data.append(body) + return data + } + + public static func proxyError(statusCode: Int, message: String) -> Data { + let body = "{\"error\":\"\(message)\"}" + return write( + statusCode: statusCode, + headers: [("Content-Type", "application/json")], + body: Data(body.utf8)) + } +} diff --git a/Sources/CodexBarCore/Proxy/LocalProxyServer.swift b/Sources/CodexBarCore/Proxy/LocalProxyServer.swift new file mode 100644 index 000000000..fe7cefb91 --- /dev/null +++ b/Sources/CodexBarCore/Proxy/LocalProxyServer.swift @@ -0,0 +1,220 @@ +import Foundation +import Network + +public final class LocalProxyServer: @unchecked Sendable { + private static let log = CodexBarLog.logger("proxy-server") + + private let configuration: ProxyConfiguration + private let accumulator: TokenAccumulator + private var listener: NWListener? + private let queue = DispatchQueue(label: "com.codexbar.proxy", qos: .utility) + private var isRunning = false + + public init(configuration: ProxyConfiguration, accumulator: TokenAccumulator) { + self.configuration = configuration + self.accumulator = accumulator + } + + public func start() throws { + guard !self.isRunning else { return } + + let params = NWParameters.tcp + params.allowLocalEndpointReuse = true + + let listener = try NWListener(using: params, on: NWEndpoint.Port(rawValue: self.configuration.port)!) + self.listener = listener + + listener.stateUpdateHandler = { [weak self] state in + switch state { + case .ready: + Self.log.info("Proxy server ready on \(self?.configuration.bindAddress ?? ""):\(self?.configuration.port ?? 0)") + self?.isRunning = true + case .failed(let error): + Self.log.error("Proxy server failed: \(error)") + self?.isRunning = false + case .cancelled: + self?.isRunning = false + default: + break + } + } + + listener.newConnectionHandler = { [weak self] connection in + self?.handleConnection(connection) + } + + listener.start(queue: self.queue) + } + + public func stop() { + self.listener?.cancel() + self.listener = nil + self.isRunning = false + Self.log.info("Proxy server stopped") + } + + public var running: Bool { self.isRunning } + + private func handleConnection(_ connection: NWConnection) { + connection.start(queue: self.queue) + + connection.receive(minimumIncompleteLength: 1, maximumLength: 65536) { [weak self] data, _, isComplete, error in + guard let self, let data, !data.isEmpty else { + if let error { Self.log.error("Receive error: \(error)") } + connection.cancel() + return + } + + self.processRequest(data: data, connection: connection) + + if isComplete { + connection.cancel() + } + } + } + + private func processRequest(data: Data, connection: NWConnection) { + guard let request = HTTPRequestParser.parse(data: data) else { + let errorResponse = HTTPResponseWriter.proxyError(statusCode: 400, message: "Bad Request") + self.send(errorResponse, on: connection) + return + } + + let targetHost = request.host.components(separatedBy: ":").first ?? request.host + let provider = self.identifyProvider(host: targetHost) + + var targetURL = "https://\(request.host)\(request.path)" + if !targetURL.hasPrefix("https://") { + targetURL = "https://\(request.host)\(request.path)" + } + + guard let url = URL(string: targetURL) else { + let errorResponse = HTTPResponseWriter.proxyError(statusCode: 400, message: "Invalid URL") + self.send(errorResponse, on: connection) + return + } + + var urlRequest = URLRequest(url: url) + urlRequest.httpMethod = request.method + urlRequest.httpBody = request.body + + for (key, value) in request.headers where key.lowercased() != "host" { + urlRequest.setValue(value, forHTTPHeaderField: key) + } + + let isSSE = request.path.contains("/chat/completions") || + request.headers.contains(where: { $0.0.lowercased() == "accept" && $0.1.contains("text/event-stream") }) + + let task = URLSession.shared.dataTask(with: urlRequest) { [weak self] responseData, response, error in + guard let self else { return } + + if let error { + Self.log.error("Forward error: \(error)") + let errorResponse = HTTPResponseWriter.proxyError(statusCode: 502, message: "Bad Gateway") + self.send(errorResponse, on: connection) + return + } + + guard let httpResponse = response as? HTTPURLResponse, + let responseData else { + let errorResponse = HTTPResponseWriter.proxyError(statusCode: 502, message: "No response") + self.send(errorResponse, on: connection) + return + } + + self.extractUsage(from: responseData, provider: provider, isSSE: isSSE) + + var responseHeaders: [(String, String)] = [] + for (key, value) in httpResponse.allHeaderFields { + if let keyStr = key as? String, let valueStr = value as? String { + if keyStr.lowercased() != "transfer-encoding" { + responseHeaders.append((keyStr, valueStr)) + } + } + } + + let proxyResponse = HTTPResponseWriter.write( + statusCode: httpResponse.statusCode, + headers: responseHeaders, + body: responseData) + self.send(proxyResponse, on: connection) + } + + task.resume() + } + + private func extractUsage(from data: Data, provider: UsageProvider?, isSSE: Bool) { + guard let provider else { return } + + if isSSE { + self.extractUsageFromSSE(data: data, provider: provider) + } else { + self.extractUsageFromJSON(data: data, provider: provider) + } + } + + private func extractUsageFromJSON(data: Data, provider: UsageProvider) { + guard let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any], + let usage = json["usage"] as? [String: Any] + else { return } + + self.recordUsage(usage: usage, provider: provider) + } + + private func extractUsageFromSSE(data: Data, provider: UsageProvider) { + guard let text = String(data: data, encoding: .utf8) else { return } + + let lines = text.components(separatedBy: "\n") + for line in lines { + guard line.hasPrefix("data: ") else { continue } + let jsonString = String(line.dropFirst("data: ".count)) + guard jsonString != "[DONE]", + let jsonData = jsonString.data(using: .utf8), + let json = try? JSONSerialization.jsonObject(with: jsonData) as? [String: Any], + let usage = json["usage"] as? [String: Any] + else { continue } + + self.recordUsage(usage: usage, provider: provider) + } + } + + private func recordUsage(usage: [String: Any], provider: UsageProvider) { + let promptTokens = usage["prompt_tokens"] as? Int ?? 0 + let completionTokens = usage["completion_tokens"] as? Int ?? 0 + let totalTokens = usage["total_tokens"] as? Int ?? (promptTokens + completionTokens) + let model = usage["model"] as? String + + guard promptTokens > 0 || completionTokens > 0 else { return } + + self.accumulator.record( + provider: provider, + promptTokens: promptTokens, + completionTokens: completionTokens, + totalTokens: totalTokens, + model: model) + + Self.log.info("Recorded \(provider): prompt=\(promptTokens) completion=\(completionTokens) total=\(totalTokens)") + } + + private func send(_ data: Data, on connection: NWConnection) { + connection.send(content: data, completion: .contentProcessed { error in + if let error { + Self.log.error("Send error: \(error)") + } + connection.cancel() + }) + } + + private func identifyProvider(host: String) -> UsageProvider? { + switch host { + case "api.deepseek.com": return .deepseek + case "open.bigmodel.cn": return .zhipu + case "ark.cn-beijing.volces.com": return .doubao + case "qianfan.baidubce.com": return .ernie + case "api.moonshot.cn": return .kimi + case "api.minimax.chat": return .minimax + case "api.xiaomimimo.com": return .mimo + default: return nil + } + } +} diff --git a/Sources/CodexBarCore/Proxy/ProxyConfiguration.swift b/Sources/CodexBarCore/Proxy/ProxyConfiguration.swift new file mode 100644 index 000000000..212c0db85 --- /dev/null +++ b/Sources/CodexBarCore/Proxy/ProxyConfiguration.swift @@ -0,0 +1,43 @@ +import Foundation + +public struct ProxyConfiguration: Sendable { + public let port: UInt16 + public let bindAddress: String + public let isEnabled: Bool + + public static let `default` = ProxyConfiguration( + port: 9876, + bindAddress: "127.0.0.1", + isEnabled: false) + + public init(port: UInt16, bindAddress: String, isEnabled: Bool) { + self.port = port + self.bindAddress = bindAddress + self.isEnabled = isEnabled + } +} + +public struct ProxyTokenEntry: Sendable { + public let provider: UsageProvider + public let promptTokens: Int + public let completionTokens: Int + public let totalTokens: Int + public let model: String? + public let timestamp: Date + + public init( + provider: UsageProvider, + promptTokens: Int, + completionTokens: Int, + totalTokens: Int, + model: String? = nil, + timestamp: Date = Date()) + { + self.provider = provider + self.promptTokens = promptTokens + self.completionTokens = completionTokens + self.totalTokens = totalTokens + self.model = model + self.timestamp = timestamp + } +} diff --git a/Sources/CodexBarCore/Proxy/ProxyManager.swift b/Sources/CodexBarCore/Proxy/ProxyManager.swift new file mode 100644 index 000000000..99962f61d --- /dev/null +++ b/Sources/CodexBarCore/Proxy/ProxyManager.swift @@ -0,0 +1,55 @@ +import Foundation + +@MainActor +public final class ProxyManager: ObservableObject { + @Published public private(set) var isRunning = false + @Published public private(set) var activePort: UInt16 = 0 + @Published public private(set) var requestCount: Int = 0 + + public let accumulator = TokenAccumulator() + private var server: LocalProxyServer? + private static let log = CodexBarLog.logger("proxy-manager") + + public init() {} + + public func start(port: UInt16 = 9876) { + guard !self.isRunning else { return } + + let config = ProxyConfiguration(port: port, bindAddress: "127.0.0.1", isEnabled: true) + let server = LocalProxyServer(configuration: config, accumulator: self.accumulator) + + do { + try server.start() + self.server = server + self.isRunning = true + self.activePort = port + Self.log.info("Proxy started on port \(port)") + } catch { + Self.log.error("Failed to start proxy: \(error)") + } + } + + public func stop() { + self.server?.stop() + self.server = nil + self.isRunning = false + self.activePort = 0 + Self.log.info("Proxy stopped") + } + + public func toggle(port: UInt16 = 9876) { + if self.isRunning { + self.stop() + } else { + self.start(port: port) + } + } + + public func snapshot(for provider: UsageProvider) -> ProxyTokenEntry? { + self.accumulator.snapshot(for: provider) + } + + public func allSnapshots() -> [ProxyTokenEntry] { + self.accumulator.allSnapshots() + } +} diff --git a/Sources/CodexBarCore/Proxy/ProxyUsageBridge.swift b/Sources/CodexBarCore/Proxy/ProxyUsageBridge.swift new file mode 100644 index 000000000..0de369ed9 --- /dev/null +++ b/Sources/CodexBarCore/Proxy/ProxyUsageBridge.swift @@ -0,0 +1,46 @@ +import Foundation + +public struct ProxyUsageBridge: Sendable { + public init() {} + + public func toUsageSnapshot(entry: ProxyTokenEntry, provider: UsageProvider) -> UsageSnapshot { + let totalDisplay = Self.formatTokenCount(entry.totalTokens) + let promptDisplay = Self.formatTokenCount(entry.promptTokens) + let completionDisplay = Self.formatTokenCount(entry.completionTokens) + + let primaryWindow = RateWindow( + usedPercent: 0, + windowMinutes: nil, + resetsAt: nil, + resetDescription: "Today: \(totalDisplay) tokens") + + let secondaryWindow = RateWindow( + usedPercent: 0, + windowMinutes: nil, + resetsAt: nil, + resetDescription: "In: \(promptDisplay) / Out: \(completionDisplay)") + + let identity = ProviderIdentitySnapshot( + providerID: provider, + accountEmail: nil, + accountOrganization: nil, + loginMethod: "proxy") + + return UsageSnapshot( + primary: primaryWindow, + secondary: secondaryWindow, + tertiary: nil, + providerCost: nil, + updatedAt: entry.timestamp, + identity: identity) + } + + private static func formatTokenCount(_ count: Int) -> String { + if count >= 1_000_000 { + return String(format: "%.1fM", Double(count) / 1_000_000) + } else if count >= 1_000 { + return String(format: "%.1fK", Double(count) / 1_000) + } + return "\(count)" + } +} diff --git a/Sources/CodexBarCore/Proxy/TokenAccumulator.swift b/Sources/CodexBarCore/Proxy/TokenAccumulator.swift new file mode 100644 index 000000000..2e05df0b2 --- /dev/null +++ b/Sources/CodexBarCore/Proxy/TokenAccumulator.swift @@ -0,0 +1,88 @@ +import Foundation + +public final class TokenAccumulator: @unchecked Sendable { + private let lock = NSLock() + private var entries: [AccumulatorKey: AccumulatorValue] = [:] + + public init() {} + + public func record( + provider: UsageProvider, + promptTokens: Int, + completionTokens: Int, + totalTokens: Int, + model: String?) + { + let key = AccumulatorKey(provider: provider, day: Self.startOfDay(Date())) + self.lock.lock() + defer { self.lock.unlock() } + + if var existing = self.entries[key] { + existing.promptTokens += promptTokens + existing.completionTokens += completionTokens + existing.totalTokens += totalTokens + existing.requestCount += 1 + self.entries[key] = existing + } else { + self.entries[key] = AccumulatorValue( + promptTokens: promptTokens, + completionTokens: completionTokens, + totalTokens: totalTokens, + requestCount: 1) + } + } + + public func snapshot(for provider: UsageProvider, day: Date = Date()) -> ProxyTokenEntry? { + let key = AccumulatorKey(provider: provider, day: Self.startOfDay(day)) + self.lock.lock() + defer { self.lock.unlock() } + + guard let value = self.entries[key] else { return nil } + + return ProxyTokenEntry( + provider: provider, + promptTokens: value.promptTokens, + completionTokens: value.completionTokens, + totalTokens: value.totalTokens, + model: nil, + timestamp: day) + } + + public func reset(for provider: UsageProvider, day: Date = Date()) { + let key = AccumulatorKey(provider: provider, day: Self.startOfDay(day)) + self.lock.lock() + defer { self.lock.unlock() } + self.entries.removeValue(forKey: key) + } + + public func allSnapshots() -> [ProxyTokenEntry] { + self.lock.lock() + defer { self.lock.unlock() } + + return self.entries.map { key, value in + ProxyTokenEntry( + provider: key.provider, + promptTokens: value.promptTokens, + completionTokens: value.completionTokens, + totalTokens: value.totalTokens, + model: nil, + timestamp: key.day) + } + } + + private static func startOfDay(_ date: Date) -> Date { + Calendar.current.startOfDay(for: date) + } + + private struct AccumulatorKey: Hashable { + let provider: UsageProvider + let day: Date + } + + private struct AccumulatorValue { + var promptTokens: Int + var completionTokens: Int + var totalTokens: Int + var requestCount: Int + } +} diff --git a/Sources/CodexBarCore/Vendored/CostUsage/CostUsageScanner.swift b/Sources/CodexBarCore/Vendored/CostUsage/CostUsageScanner.swift index 220f2ec4b..51e04c2d5 100644 --- a/Sources/CodexBarCore/Vendored/CostUsage/CostUsageScanner.swift +++ b/Sources/CodexBarCore/Vendored/CostUsage/CostUsageScanner.swift @@ -236,7 +236,7 @@ enum CostUsageScanner { case .zai, .gemini, .antigravity, .cursor, .opencode, .opencodego, .alibaba, .factory, .copilot, .minimax, .kilo, .kiro, .kimi, .kimik2, .augment, .jetbrains, .amp, .ollama, .synthetic, .openrouter, .warp, .perplexity, .abacus, - .mistral, .deepseek, .codebuff: + .mistral, .deepseek, .codebuff, .zhipu, .doubao, .ernie, .mimo: return emptyReport } } diff --git a/Sources/CodexBarWidget/CodexBarWidgetProvider.swift b/Sources/CodexBarWidget/CodexBarWidgetProvider.swift index a1f4d9323..d3087f6ac 100644 --- a/Sources/CodexBarWidget/CodexBarWidgetProvider.swift +++ b/Sources/CodexBarWidget/CodexBarWidgetProvider.swift @@ -80,6 +80,10 @@ enum ProviderChoice: String, AppEnum { case .mistral: return nil // Mistral not yet supported in widgets case .deepseek: return nil // DeepSeek not yet supported in widgets case .codebuff: return nil // Codebuff not yet supported in widgets + case .zhipu: return nil // Zhipu not yet supported in widgets + case .doubao: return nil // Doubao not yet supported in widgets + case .ernie: return nil // ERNIE 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 04afce0b7..0c603509e 100644 --- a/Sources/CodexBarWidget/CodexBarWidgetViews.swift +++ b/Sources/CodexBarWidget/CodexBarWidgetViews.swift @@ -286,6 +286,10 @@ private struct ProviderSwitchChip: View { case .mistral: "Mistral" case .deepseek: "DeepSeek" case .codebuff: "Codebuff" + case .zhipu: "Zhipu" + case .doubao: "Doubao" + case .ernie: "ERNIE" + case .mimo: "MiMo" } } } @@ -653,6 +657,14 @@ enum WidgetColors { Color(red: 82 / 255, green: 125 / 255, blue: 240 / 255) case .codebuff: Color(red: 68 / 255, green: 255 / 255, blue: 0 / 255) // Codebuff lime + case .zhipu: + Color(red: 0 / 255, green: 115 / 255, blue: 217 / 255) // Zhipu blue + case .doubao: + Color(red: 0 / 255, green: 140 / 255, blue: 243 / 255) // Doubao blue + case .ernie: + Color(red: 38 / 255, green: 115 / 255, blue: 217 / 255) // ERNIE blue + case .mimo: + Color(red: 1.0, green: 140 / 255, blue: 0) // MiMo orange } } } diff --git a/Tests/CodexBarTests/DoubaoBalanceParsingTests.swift b/Tests/CodexBarTests/DoubaoBalanceParsingTests.swift new file mode 100644 index 000000000..0bc299492 --- /dev/null +++ b/Tests/CodexBarTests/DoubaoBalanceParsingTests.swift @@ -0,0 +1,91 @@ +import Foundation +import Testing +@testable import CodexBarCore + +struct DoubaoBalanceParsingTests { + @Test + func `parses standard balance response`() throws { + let json = """ + { + "data": { + "available_balance": 85.50, + "used_balance": 14.50, + "total_balance": 100.00 + } + } + """ + let info = try DoubaoUsageFetcher._parseBalanceForTesting(Data(json.utf8)) + #expect(info.availableBalance == 85.50) + #expect(info.usedBalance == 14.50) + #expect(info.totalBalance == 100.00) + } + + @Test + func `parses camelCase balance response`() throws { + let json = """ + { + "data": { + "availableBalance": 50.0, + "usedBalance": 50.0, + "totalBalance": 100.0 + } + } + """ + let info = try DoubaoUsageFetcher._parseBalanceForTesting(Data(json.utf8)) + #expect(info.availableBalance == 50.0) + #expect(info.totalBalance == 100.0) + } + + @Test + func `falls back to root when no data wrapper`() throws { + let json = """ + { + "available_balance": 200.0, + "used_balance": 0.0, + "total_balance": 200.0 + } + """ + let info = try DoubaoUsageFetcher._parseBalanceForTesting(Data(json.utf8)) + #expect(info.availableBalance == 200.0) + } + + @Test + func `derives total from available plus used`() throws { + let json = """ + { + "data": { + "available_balance": 60.0, + "used_balance": 40.0 + } + } + """ + let info = try DoubaoUsageFetcher._parseBalanceForTesting(Data(json.utf8)) + #expect(info.totalBalance == 100.0) + } + + @Test + func `handles string numeric values`() throws { + let json = """ + { + "data": { + "available_balance": "75.25", + "used_balance": "24.75", + "total_balance": "100.00" + } + } + """ + let info = try DoubaoUsageFetcher._parseBalanceForTesting(Data(json.utf8)) + #expect(info.availableBalance == 75.25) + } + + @Test + func `invalid JSON throws parse error`() { + let data = Data("not json".utf8) + #expect { + _ = try DoubaoUsageFetcher._parseBalanceForTesting(data) + } throws: { error in + guard case DoubaoUsageError.parseFailed = error else { return false } + return true + } + } +} diff --git a/Tests/CodexBarTests/ErnieBalanceParsingTests.swift b/Tests/CodexBarTests/ErnieBalanceParsingTests.swift new file mode 100644 index 000000000..05a20ec20 --- /dev/null +++ b/Tests/CodexBarTests/ErnieBalanceParsingTests.swift @@ -0,0 +1,91 @@ +import Foundation +import Testing +@testable import CodexBarCore + +struct ErnieBalanceParsingTests { + @Test + func `parses standard balance response`() throws { + let json = """ + { + "data": { + "available_balance": 85.50, + "used_balance": 14.50, + "total_balance": 100.00 + } + } + """ + let info = try ErnieUsageFetcher._parseBalanceForTesting(Data(json.utf8)) + #expect(info.availableBalance == 85.50) + #expect(info.usedBalance == 14.50) + #expect(info.totalBalance == 100.00) + } + + @Test + func `parses camelCase balance response`() throws { + let json = """ + { + "data": { + "availableBalance": 50.0, + "usedBalance": 50.0, + "totalBalance": 100.0 + } + } + """ + let info = try ErnieUsageFetcher._parseBalanceForTesting(Data(json.utf8)) + #expect(info.availableBalance == 50.0) + #expect(info.totalBalance == 100.0) + } + + @Test + func `falls back to root when no data wrapper`() throws { + let json = """ + { + "available_balance": 200.0, + "used_balance": 0.0, + "total_balance": 200.0 + } + """ + let info = try ErnieUsageFetcher._parseBalanceForTesting(Data(json.utf8)) + #expect(info.availableBalance == 200.0) + } + + @Test + func `derives total from available plus used`() throws { + let json = """ + { + "data": { + "available_balance": 60.0, + "used_balance": 40.0 + } + } + """ + let info = try ErnieUsageFetcher._parseBalanceForTesting(Data(json.utf8)) + #expect(info.totalBalance == 100.0) + } + + @Test + func `handles string numeric values`() throws { + let json = """ + { + "data": { + "available_balance": "75.25", + "used_balance": "24.75", + "total_balance": "100.00" + } + } + """ + let info = try ErnieUsageFetcher._parseBalanceForTesting(Data(json.utf8)) + #expect(info.availableBalance == 75.25) + } + + @Test + func `invalid JSON throws parse error`() { + let data = Data("not json".utf8) + #expect { + _ = try ErnieUsageFetcher._parseBalanceForTesting(data) + } throws: { error in + guard case ErnieUsageError.parseFailed = error else { return false } + return true + } + } +} diff --git a/Tests/CodexBarTests/KimiK2BalanceParsingTests.swift b/Tests/CodexBarTests/KimiK2BalanceParsingTests.swift new file mode 100644 index 000000000..a24be3973 --- /dev/null +++ b/Tests/CodexBarTests/KimiK2BalanceParsingTests.swift @@ -0,0 +1,92 @@ +import Foundation +import Testing +@testable import CodexBarCore + +struct KimiK2BalanceParsingTests { + @Test + func `parses standard balance response`() throws { + let json = """ + { + "data": { + "available_balance": 100.00, + "voucher_balance": 10.00, + "cash_balance": 90.00 + } + } + """ + let info = try KimiK2UsageFetcher._parseBalanceForTesting(Data(json.utf8)) + #expect(info != nil) + #expect(info?.availableBalance == 100.00) + #expect(info?.voucherBalance == 10.00) + #expect(info?.cashBalance == 90.00) + } + + @Test + func `parses camelCase balance response`() throws { + let json = """ + { + "data": { + "availableBalance": 50.0, + "voucherBalance": 5.0, + "cashBalance": 45.0 + } + } + """ + let info = try KimiK2UsageFetcher._parseBalanceForTesting(Data(json.utf8)) + #expect(info?.availableBalance == 50.0) + #expect(info?.voucherBalance == 5.0) + #expect(info?.cashBalance == 45.0) + } + + @Test + func `falls back to root when no data wrapper`() throws { + let json = """ + { + "available_balance": 200.0, + "voucher_balance": 0.0, + "cash_balance": 200.0 + } + """ + let info = try KimiK2UsageFetcher._parseBalanceForTesting(Data(json.utf8)) + #expect(info?.availableBalance == 200.0) + #expect(info?.cashBalance == 200.0) + } + + @Test + func `handles string numeric values`() throws { + let json = """ + { + "data": { + "available_balance": "75.25", + "voucher_balance": "0", + "cash_balance": "75.25" + } + } + """ + let info = try KimiK2UsageFetcher._parseBalanceForTesting(Data(json.utf8)) + #expect(info?.availableBalance == 75.25) + } + + @Test + func `handles missing fields with defaults`() throws { + let json = """ + { "data": {} } + """ + let info = try KimiK2UsageFetcher._parseBalanceForTesting(Data(json.utf8)) + #expect(info != nil) + #expect(info?.availableBalance == 0) + #expect(info?.voucherBalance == 0) + #expect(info?.cashBalance == 0) + } + + @Test + func `invalid JSON throws parse error`() { + let data = Data("not json".utf8) + #expect { + _ = try KimiK2UsageFetcher._parseBalanceForTesting(data) + } throws: { error in + guard case KimiK2UsageError.parseFailed = error else { return false } + return true + } + } +} diff --git a/Tests/CodexBarTests/SettingsStoreTests.swift b/Tests/CodexBarTests/SettingsStoreTests.swift index c63bcd912..454df755c 100644 --- a/Tests/CodexBarTests/SettingsStoreTests.swift +++ b/Tests/CodexBarTests/SettingsStoreTests.swift @@ -941,6 +941,10 @@ struct SettingsStoreTests { .mistral, .deepseek, .codebuff, + .zhipu, + .doubao, + .ernie, + .mimo, ]) // Move one provider; ensure it's persisted across instances. diff --git a/Tests/CodexBarTests/ZhipuBalanceParsingTests.swift b/Tests/CodexBarTests/ZhipuBalanceParsingTests.swift new file mode 100644 index 000000000..8d441cfac --- /dev/null +++ b/Tests/CodexBarTests/ZhipuBalanceParsingTests.swift @@ -0,0 +1,108 @@ +import Foundation +import Testing +@testable import CodexBarCore + +struct ZhipuBalanceParsingTests { + @Test + func `parses standard balance response`() throws { + let json = """ + { + "data": { + "available_balance": 85.50, + "used_balance": 14.50, + "total_balance": 100.00 + } + } + """ + let info = try ZhipuUsageFetcher._parseBalanceForTesting(Data(json.utf8)) + #expect(info.availableBalance == 85.50) + #expect(info.usedBalance == 14.50) + #expect(info.totalBalance == 100.00) + } + + @Test + func `parses camelCase balance response`() throws { + let json = """ + { + "data": { + "availableBalance": 50.0, + "usedBalance": 50.0, + "totalBalance": 100.0 + } + } + """ + let info = try ZhipuUsageFetcher._parseBalanceForTesting(Data(json.utf8)) + #expect(info.availableBalance == 50.0) + #expect(info.usedBalance == 50.0) + #expect(info.totalBalance == 100.0) + } + + @Test + func `falls back to root when no data wrapper`() throws { + let json = """ + { + "available_balance": 200.0, + "used_balance": 0.0, + "total_balance": 200.0 + } + """ + let info = try ZhipuUsageFetcher._parseBalanceForTesting(Data(json.utf8)) + #expect(info.availableBalance == 200.0) + #expect(info.totalBalance == 200.0) + } + + @Test + func `derives total from available plus used when total missing`() throws { + let json = """ + { + "data": { + "available_balance": 60.0, + "used_balance": 40.0 + } + } + """ + let info = try ZhipuUsageFetcher._parseBalanceForTesting(Data(json.utf8)) + #expect(info.availableBalance == 60.0) + #expect(info.usedBalance == 40.0) + #expect(info.totalBalance == 100.0) + } + + @Test + func `handles string numeric values`() throws { + let json = """ + { + "data": { + "available_balance": "75.25", + "used_balance": "24.75", + "total_balance": "100.00" + } + } + """ + let info = try ZhipuUsageFetcher._parseBalanceForTesting(Data(json.utf8)) + #expect(info.availableBalance == 75.25) + #expect(info.usedBalance == 24.75) + #expect(info.totalBalance == 100.00) + } + + @Test + func `handles missing fields with defaults`() throws { + let json = """ + { "data": {} } + """ + let info = try ZhipuUsageFetcher._parseBalanceForTesting(Data(json.utf8)) + #expect(info.availableBalance == 0) + #expect(info.usedBalance == 0) + #expect(info.totalBalance == 0) + } + + @Test + func `invalid JSON throws parse error`() { + let data = Data("not json".utf8) + #expect { + _ = try ZhipuUsageFetcher._parseBalanceForTesting(data) + } throws: { error in + guard case ZhipuUsageError.parseFailed = error else { return false } + return true + } + } +}