diff --git a/Sources/CodexBar/Config/CodexBarConfigMigrator.swift b/Sources/CodexBar/Config/CodexBarConfigMigrator.swift index a3462976..74ea7aaf 100644 --- a/Sources/CodexBar/Config/CodexBarConfigMigrator.swift +++ b/Sources/CodexBar/Config/CodexBarConfigMigrator.swift @@ -16,6 +16,7 @@ struct CodexBarConfigMigrator { let kimiK2TokenStore: any KimiK2TokenStoring let augmentCookieStore: any CookieHeaderStoring let ampCookieStore: any CookieHeaderStoring + let traeCookieStore: any CookieHeaderStoring let copilotTokenStore: any CopilotTokenStoring let tokenAccountStore: any ProviderTokenAccountStoring } @@ -99,6 +100,7 @@ struct CodexBarConfigMigrator { (.factory, stores.factoryCookieStore.loadCookieHeader), (.augment, stores.augmentCookieStore.loadCookieHeader), (.amp, stores.ampCookieStore.loadCookieHeader), + (.trae, stores.traeCookieStore.loadCookieHeader), ], config: &config, state: &state) @@ -123,6 +125,7 @@ struct CodexBarConfigMigrator { (.kimi, "kimiCookieSource"), (.augment, "augmentCookieSource"), (.amp, "ampCookieSource"), + (.trae, "traeCookieSource"), ] for (provider, key) in sources { @@ -294,6 +297,7 @@ struct CodexBarConfigMigrator { try stores.minimaxCookieStore.storeCookieHeader(nil) try stores.augmentCookieStore.storeCookieHeader(nil) try stores.ampCookieStore.storeCookieHeader(nil) + try stores.traeCookieStore.storeCookieHeader(nil) } catch { log.error("Failed to clear legacy secrets: \(error)") } diff --git a/Sources/CodexBar/KeychainPromptCoordinator.swift b/Sources/CodexBar/KeychainPromptCoordinator.swift index a6add39a..d8a0d211 100644 --- a/Sources/CodexBar/KeychainPromptCoordinator.swift +++ b/Sources/CodexBar/KeychainPromptCoordinator.swift @@ -109,6 +109,11 @@ enum KeychainPromptCoordinator { "CodexBar will ask macOS Keychain for your Amp cookie header", "so it can fetch usage. Click OK to continue.", ].joined(separator: " ")) + case .traeCookie: + return (title, [ + "CodexBar will ask macOS Keychain for your Trae JWT token", + "so it can fetch usage. Click OK to continue.", + ].joined(separator: " ")) } } diff --git a/Sources/CodexBar/PreferencesProvidersPane+Testing.swift b/Sources/CodexBar/PreferencesProvidersPane+Testing.swift index e2dce0a7..b496f8ae 100644 --- a/Sources/CodexBar/PreferencesProvidersPane+Testing.swift +++ b/Sources/CodexBar/PreferencesProvidersPane+Testing.swift @@ -55,6 +55,7 @@ enum ProvidersPaneTestHarness { settings.factoryCookieSource = .manual settings.minimaxCookieSource = .manual settings.augmentCookieSource = .manual + settings.traeCookieSource = .manual } private static func exercisePaneBasics(pane: ProvidersPane) { @@ -68,6 +69,7 @@ enum ProvidersPaneTestHarness { _ = pane._test_providerSubtitle(.minimax) _ = pane._test_providerSubtitle(.kimi) _ = pane._test_providerSubtitle(.gemini) + _ = pane._test_providerSubtitle(.trae) _ = pane._test_menuBarMetricPicker(for: .codex) _ = pane._test_menuBarMetricPicker(for: .gemini) diff --git a/Sources/CodexBar/Providers/Shared/ProviderImplementationRegistry.swift b/Sources/CodexBar/Providers/Shared/ProviderImplementationRegistry.swift index f6a9b2a3..2619e2a0 100644 --- a/Sources/CodexBar/Providers/Shared/ProviderImplementationRegistry.swift +++ b/Sources/CodexBar/Providers/Shared/ProviderImplementationRegistry.swift @@ -29,6 +29,7 @@ enum ProviderImplementationRegistry { case .jetbrains: JetBrainsProviderImplementation() case .kimik2: KimiK2ProviderImplementation() case .amp: AmpProviderImplementation() + case .trae: TraeProviderImplementation() case .synthetic: SyntheticProviderImplementation() } } diff --git a/Sources/CodexBar/Providers/Trae/TraeProviderImplementation.swift b/Sources/CodexBar/Providers/Trae/TraeProviderImplementation.swift new file mode 100644 index 00000000..89c151e7 --- /dev/null +++ b/Sources/CodexBar/Providers/Trae/TraeProviderImplementation.swift @@ -0,0 +1,51 @@ +import AppKit +import CodexBarCore +import CodexBarMacroSupport +import Foundation +import SwiftUI + +@ProviderImplementationRegistration +struct TraeProviderImplementation: ProviderImplementation { + let id: UsageProvider = .trae + + @MainActor + func observeSettings(_ settings: SettingsStore) { + _ = settings.traeCookieHeader + _ = settings.tokenAccounts(for: .trae) + } + + @MainActor + func settingsSnapshot(context: ProviderSettingsSnapshotContext) -> ProviderSettingsSnapshotContribution? { + .trae(context.settings.traeSettingsSnapshot(tokenOverride: context.tokenOverride)) + } + + // No settingsPickers - Trae only supports manual JWT entry + // The JWT is not stored in browser cookies, only in HTTP headers + + @MainActor + func settingsFields(context: ProviderSettingsContext) -> [ProviderSettingsFieldDescriptor] { + [ + ProviderSettingsFieldDescriptor( + id: "trae-jwt-token", + title: "JWT Token", + subtitle: "Paste Authorization header from browser DevTools. Instructions:\n1. Open trae.ai and log in\n2. Open DevTools → Network tab\n3. Find request to 'user_current_entitlement_list'\n4. Copy 'Authorization' header (starts with 'Cloud-IDE-JWT')\n5. Paste here", + kind: .secure, + placeholder: "Cloud-IDE-JWT eyJ...", + binding: context.stringBinding(\.traeCookieHeader), + actions: [ + ProviderSettingsActionDescriptor( + id: "trae-open-settings", + title: "Open Trae Settings", + style: .link, + isVisible: nil, + perform: { + if let url = URL(string: "https://www.trae.ai/account-setting") { + NSWorkspace.shared.open(url) + } + }), + ], + isVisible: { true }, + onActivate: { context.settings.ensureTraeCookieLoaded() }), + ] + } +} diff --git a/Sources/CodexBar/Providers/Trae/TraeSettingsStore.swift b/Sources/CodexBar/Providers/Trae/TraeSettingsStore.swift new file mode 100644 index 00000000..cf604b34 --- /dev/null +++ b/Sources/CodexBar/Providers/Trae/TraeSettingsStore.swift @@ -0,0 +1,62 @@ +import CodexBarCore +import Foundation + +extension SettingsStore { + var traeCookieHeader: String { + get { self.configSnapshot.providerConfig(for: .trae)?.sanitizedCookieHeader ?? "" } + set { + self.updateProviderConfig(provider: .trae) { entry in + entry.cookieHeader = self.normalizedConfigValue(newValue) + } + self.logSecretUpdate(provider: .trae, field: "cookieHeader", value: newValue) + } + } + + var traeCookieSource: ProviderCookieSource { + get { self.resolvedCookieSource(provider: .trae, fallback: .auto) } + set { + self.updateProviderConfig(provider: .trae) { entry in + entry.cookieSource = newValue + } + self.logProviderModeChange(provider: .trae, field: "cookieSource", value: newValue.rawValue) + } + } + + func ensureTraeCookieLoaded() {} +} + +extension SettingsStore { + func traeSettingsSnapshot(tokenOverride: TokenAccountOverride?) -> ProviderSettingsSnapshot.TraeProviderSettings { + ProviderSettingsSnapshot.TraeProviderSettings( + cookieSource: self.traeSnapshotCookieSource(tokenOverride: tokenOverride), + manualCookieHeader: self.traeSnapshotCookieHeader(tokenOverride: tokenOverride)) + } + + private func traeSnapshotCookieHeader(tokenOverride: TokenAccountOverride?) -> String { + let fallback = self.traeCookieHeader + guard let support = TokenAccountSupportCatalog.support(for: .trae), + case .cookieHeader = support.injection + else { + return fallback + } + guard let account = ProviderTokenAccountSelection.selectedAccount( + provider: .trae, + settings: self, + override: tokenOverride) + else { + return fallback + } + return TokenAccountSupportCatalog.normalizedCookieHeader(account.token, support: support) + } + + private func traeSnapshotCookieSource(tokenOverride: TokenAccountOverride?) -> ProviderCookieSource { + let fallback = self.traeCookieSource + guard let support = TokenAccountSupportCatalog.support(for: .trae), + support.requiresManualCookieSource + else { + return fallback + } + if self.tokenAccounts(for: .trae).isEmpty { return fallback } + return .manual + } +} diff --git a/Sources/CodexBar/Resources/ProviderIcon-trae.svg b/Sources/CodexBar/Resources/ProviderIcon-trae.svg new file mode 100644 index 00000000..8f42e96f --- /dev/null +++ b/Sources/CodexBar/Resources/ProviderIcon-trae.svg @@ -0,0 +1 @@ +TRAE \ No newline at end of file diff --git a/Sources/CodexBar/SettingsStore+MenuObservation.swift b/Sources/CodexBar/SettingsStore+MenuObservation.swift index 6c69f6a7..373f6e03 100644 --- a/Sources/CodexBar/SettingsStore+MenuObservation.swift +++ b/Sources/CodexBar/SettingsStore+MenuObservation.swift @@ -36,6 +36,7 @@ extension SettingsStore { _ = self.kimiCookieSource _ = self.augmentCookieSource _ = self.ampCookieSource + _ = self.traeCookieSource _ = self.mergeIcons _ = self.switcherShowsIcons _ = self.zaiAPIToken @@ -52,6 +53,7 @@ extension SettingsStore { _ = self.kimiK2APIToken _ = self.augmentCookieHeader _ = self.ampCookieHeader + _ = self.traeCookieHeader _ = self.copilotAPIToken _ = self.tokenAccountsByProvider _ = self.debugLoadingPattern diff --git a/Sources/CodexBar/SettingsStore.swift b/Sources/CodexBar/SettingsStore.swift index 42514476..2043bc43 100644 --- a/Sources/CodexBar/SettingsStore.swift +++ b/Sources/CodexBar/SettingsStore.swift @@ -107,6 +107,9 @@ final class SettingsStore { ampCookieStore: any CookieHeaderStoring = KeychainCookieHeaderStore( account: "amp-cookie", promptKind: .ampCookie), + traeCookieStore: any CookieHeaderStoring = KeychainCookieHeaderStore( + account: "trae-cookie", + promptKind: .traeCookie), copilotTokenStore: any CopilotTokenStoring = KeychainCopilotTokenStore(), tokenAccountStore: any ProviderTokenAccountStoring = FileTokenAccountStore()) { @@ -124,6 +127,7 @@ final class SettingsStore { kimiK2TokenStore: kimiK2TokenStore, augmentCookieStore: augmentCookieStore, ampCookieStore: ampCookieStore, + traeCookieStore: traeCookieStore, copilotTokenStore: copilotTokenStore, tokenAccountStore: tokenAccountStore) let config = CodexBarConfigMigrator.loadOrMigrate( diff --git a/Sources/CodexBar/UsageStore+Logging.swift b/Sources/CodexBar/UsageStore+Logging.swift index 4a72e352..7f3c0cc3 100644 --- a/Sources/CodexBar/UsageStore+Logging.swift +++ b/Sources/CodexBar/UsageStore+Logging.swift @@ -14,6 +14,7 @@ extension UsageStore { "kimiCookieSource": self.settings.kimiCookieSource.rawValue, "augmentCookieSource": self.settings.augmentCookieSource.rawValue, "ampCookieSource": self.settings.ampCookieSource.rawValue, + "traeCookieSource": self.settings.traeCookieSource.rawValue, "openAIWebAccess": self.settings.openAIWebAccessEnabled ? "1" : "0", "claudeWebExtras": self.settings.claudeWebExtrasEnabled ? "1" : "0", ] diff --git a/Sources/CodexBar/UsageStore.swift b/Sources/CodexBar/UsageStore.swift index 69491056..a6135c13 100644 --- a/Sources/CodexBar/UsageStore.swift +++ b/Sources/CodexBar/UsageStore.swift @@ -1225,6 +1225,10 @@ extension UsageStore { ampCookieHeader: self.settings.ampCookieHeader) await MainActor.run { self.probeLogs[.amp] = text } return text + case .trae: + let text = await TraeUsageFetcher.latestDumps() + await MainActor.run { self.probeLogs[.trae] = text } + return text case .jetbrains: let text = "JetBrains AI debug log not yet implemented" await MainActor.run { self.probeLogs[.jetbrains] = text } diff --git a/Sources/CodexBarCLI/TokenAccountCLI.swift b/Sources/CodexBarCLI/TokenAccountCLI.swift index a0a88371..2a1e4ced 100644 --- a/Sources/CodexBarCLI/TokenAccountCLI.swift +++ b/Sources/CodexBarCLI/TokenAccountCLI.swift @@ -147,6 +147,11 @@ struct TokenAccountCLIContext { return self.makeSnapshot( jetbrains: ProviderSettingsSnapshot.JetBrainsProviderSettings( ideBasePath: nil)) + case .trae: + return self.makeSnapshot( + trae: ProviderSettingsSnapshot.TraeProviderSettings( + cookieSource: cookieSource, + manualCookieHeader: cookieHeader)) case .gemini, .antigravity, .copilot, .kiro, .vertexai, .kimik2, .synthetic: return nil } @@ -163,6 +168,7 @@ struct TokenAccountCLIContext { kimi: ProviderSettingsSnapshot.KimiProviderSettings? = nil, augment: ProviderSettingsSnapshot.AugmentProviderSettings? = nil, amp: ProviderSettingsSnapshot.AmpProviderSettings? = nil, + trae: ProviderSettingsSnapshot.TraeProviderSettings? = nil, jetbrains: ProviderSettingsSnapshot.JetBrainsProviderSettings? = nil) -> ProviderSettingsSnapshot { ProviderSettingsSnapshot.make( @@ -176,6 +182,7 @@ struct TokenAccountCLIContext { kimi: kimi, augment: augment, amp: amp, + trae: trae, jetbrains: jetbrains) } diff --git a/Sources/CodexBarCore/KeychainAccessPreflight.swift b/Sources/CodexBarCore/KeychainAccessPreflight.swift index 866eecb1..36d7415c 100644 --- a/Sources/CodexBarCore/KeychainAccessPreflight.swift +++ b/Sources/CodexBarCore/KeychainAccessPreflight.swift @@ -20,6 +20,7 @@ public struct KeychainPromptContext: Sendable { case minimaxToken case augmentCookie case ampCookie + case traeCookie } public let kind: Kind diff --git a/Sources/CodexBarCore/Logging/LogCategories.swift b/Sources/CodexBarCore/Logging/LogCategories.swift index d3d31c8f..808a449a 100644 --- a/Sources/CodexBarCore/Logging/LogCategories.swift +++ b/Sources/CodexBarCore/Logging/LogCategories.swift @@ -50,6 +50,7 @@ public enum LogCategories { public static let syntheticTokenStore = "synthetic-token-store" public static let syntheticUsage = "synthetic-usage" public static let terminal = "terminal" + public static let trae = "trae" public static let tokenAccounts = "token-accounts" public static let tokenCost = "token-cost" public static let ttyRunner = "tty-runner" diff --git a/Sources/CodexBarCore/Providers/ProviderDescriptor.swift b/Sources/CodexBarCore/Providers/ProviderDescriptor.swift index 6aff8369..5ff3c4fe 100644 --- a/Sources/CodexBarCore/Providers/ProviderDescriptor.swift +++ b/Sources/CodexBarCore/Providers/ProviderDescriptor.swift @@ -70,6 +70,7 @@ public enum ProviderDescriptorRegistry { .jetbrains: JetBrainsProviderDescriptor.descriptor, .kimik2: KimiK2ProviderDescriptor.descriptor, .amp: AmpProviderDescriptor.descriptor, + .trae: TraeProviderDescriptor.descriptor, .synthetic: SyntheticProviderDescriptor.descriptor, ] private static let bootstrap: Void = { diff --git a/Sources/CodexBarCore/Providers/ProviderSettingsSnapshot.swift b/Sources/CodexBarCore/Providers/ProviderSettingsSnapshot.swift index f83cb9fd..069dda2f 100644 --- a/Sources/CodexBarCore/Providers/ProviderSettingsSnapshot.swift +++ b/Sources/CodexBarCore/Providers/ProviderSettingsSnapshot.swift @@ -15,6 +15,7 @@ public struct ProviderSettingsSnapshot: Sendable { kimi: KimiProviderSettings? = nil, augment: AugmentProviderSettings? = nil, amp: AmpProviderSettings? = nil, + trae: TraeProviderSettings? = nil, jetbrains: JetBrainsProviderSettings? = nil) -> ProviderSettingsSnapshot { ProviderSettingsSnapshot( @@ -31,6 +32,7 @@ public struct ProviderSettingsSnapshot: Sendable { kimi: kimi, augment: augment, amp: amp, + trae: trae, jetbrains: jetbrains) } @@ -167,6 +169,16 @@ public struct ProviderSettingsSnapshot: Sendable { } } + public struct TraeProviderSettings: Sendable { + public let cookieSource: ProviderCookieSource + public let manualCookieHeader: String? + + public init(cookieSource: ProviderCookieSource, manualCookieHeader: String?) { + self.cookieSource = cookieSource + self.manualCookieHeader = manualCookieHeader + } + } + public let debugMenuEnabled: Bool public let debugKeepCLISessionsAlive: Bool public let codex: CodexProviderSettings? @@ -180,6 +192,7 @@ public struct ProviderSettingsSnapshot: Sendable { public let kimi: KimiProviderSettings? public let augment: AugmentProviderSettings? public let amp: AmpProviderSettings? + public let trae: TraeProviderSettings? public let jetbrains: JetBrainsProviderSettings? public var jetbrainsIDEBasePath: String? { @@ -200,6 +213,7 @@ public struct ProviderSettingsSnapshot: Sendable { kimi: KimiProviderSettings?, augment: AugmentProviderSettings?, amp: AmpProviderSettings?, + trae: TraeProviderSettings?, jetbrains: JetBrainsProviderSettings? = nil) { self.debugMenuEnabled = debugMenuEnabled @@ -215,6 +229,7 @@ public struct ProviderSettingsSnapshot: Sendable { self.kimi = kimi self.augment = augment self.amp = amp + self.trae = trae self.jetbrains = jetbrains } } @@ -231,6 +246,7 @@ public enum ProviderSettingsSnapshotContribution: Sendable { case kimi(ProviderSettingsSnapshot.KimiProviderSettings) case augment(ProviderSettingsSnapshot.AugmentProviderSettings) case amp(ProviderSettingsSnapshot.AmpProviderSettings) + case trae(ProviderSettingsSnapshot.TraeProviderSettings) case jetbrains(ProviderSettingsSnapshot.JetBrainsProviderSettings) } @@ -248,6 +264,7 @@ public struct ProviderSettingsSnapshotBuilder: Sendable { public var kimi: ProviderSettingsSnapshot.KimiProviderSettings? public var augment: ProviderSettingsSnapshot.AugmentProviderSettings? public var amp: ProviderSettingsSnapshot.AmpProviderSettings? + public var trae: ProviderSettingsSnapshot.TraeProviderSettings? public var jetbrains: ProviderSettingsSnapshot.JetBrainsProviderSettings? public init(debugMenuEnabled: Bool = false, debugKeepCLISessionsAlive: Bool = false) { @@ -268,6 +285,7 @@ public struct ProviderSettingsSnapshotBuilder: Sendable { case let .kimi(value): self.kimi = value case let .augment(value): self.augment = value case let .amp(value): self.amp = value + case let .trae(value): self.trae = value case let .jetbrains(value): self.jetbrains = value } } @@ -287,6 +305,7 @@ public struct ProviderSettingsSnapshotBuilder: Sendable { kimi: self.kimi, augment: self.augment, amp: self.amp, + trae: self.trae, jetbrains: self.jetbrains) } } diff --git a/Sources/CodexBarCore/Providers/Providers.swift b/Sources/CodexBarCore/Providers/Providers.swift index a267fb95..e6516103 100644 --- a/Sources/CodexBarCore/Providers/Providers.swift +++ b/Sources/CodexBarCore/Providers/Providers.swift @@ -20,6 +20,7 @@ public enum UsageProvider: String, CaseIterable, Sendable, Codable { case jetbrains case kimik2 case amp + case trae case synthetic } @@ -43,6 +44,7 @@ public enum IconStyle: Sendable, CaseIterable { case augment case jetbrains case amp + case trae case synthetic case combined } diff --git a/Sources/CodexBarCore/Providers/Trae/TraeProviderDescriptor.swift b/Sources/CodexBarCore/Providers/Trae/TraeProviderDescriptor.swift new file mode 100644 index 00000000..dcea32c9 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Trae/TraeProviderDescriptor.swift @@ -0,0 +1,72 @@ +import CodexBarMacroSupport +import Foundation + +@ProviderDescriptorRegistration +@ProviderDescriptorDefinition +public enum TraeProviderDescriptor { + static func makeDescriptor() -> ProviderDescriptor { + ProviderDescriptor( + id: .trae, + metadata: ProviderMetadata( + id: .trae, + displayName: "Trae", + sessionLabel: "Pro Plan", + weeklyLabel: "Extra Package", + opusLabel: "Extra Package", + supportsOpus: false, + supportsCredits: false, + creditsHint: "", + toggleTitle: "Show Trae usage", + cliName: "trae", + defaultEnabled: false, + isPrimaryProvider: false, + usesAccountFallback: false, + browserCookieOrder: ProviderBrowserCookieDefaults.defaultImportOrder, + dashboardURL: "https://www.trae.ai/account-setting", + statusPageURL: nil), + branding: ProviderBranding( + iconStyle: .trae, + iconResourceName: "ProviderIcon-trae", + color: ProviderColor(red: 59 / 255, green: 130 / 255, blue: 246 / 255)), + tokenCost: ProviderTokenCostConfig( + supportsTokenCost: false, + noDataMessage: { "Trae cost summary is not supported." }), + fetchPlan: ProviderFetchPlan( + sourceModes: [.auto, .web], + pipeline: ProviderFetchPipeline(resolveStrategies: { _ in [TraeStatusFetchStrategy()] })), + cli: ProviderCLIConfig( + name: "trae", + versionDetector: nil)) + } +} + +struct TraeStatusFetchStrategy: ProviderFetchStrategy { + let id: String = "trae.web" + let kind: ProviderFetchKind = .web + + func isAvailable(_ context: ProviderFetchContext) async -> Bool { + guard context.settings?.trae?.cookieSource != .off else { return false } + return true + } + + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + let fetcher = TraeUsageFetcher(browserDetection: context.browserDetection) + let manual = Self.manualJWTToken(from: context) + let logger: ((String) -> Void)? = context.verbose + ? { msg in CodexBarLog.logger(LogCategories.trae).verbose(msg) } + : nil + let snap = try await fetcher.fetch(jwtOverride: manual, logger: logger) + return self.makeResult( + usage: snap.toUsageSnapshot(now: snap.updatedAt), + sourceLabel: "web") + } + + func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool { + false + } + + private static func manualJWTToken(from context: ProviderFetchContext) -> String? { + // Return the JWT token if provided (manual entry only) + return context.settings?.trae?.manualCookieHeader?.trimmingCharacters(in: .whitespacesAndNewlines) + } +} diff --git a/Sources/CodexBarCore/Providers/Trae/TraeUsageFetcher.swift b/Sources/CodexBarCore/Providers/Trae/TraeUsageFetcher.swift new file mode 100644 index 00000000..17720cd2 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Trae/TraeUsageFetcher.swift @@ -0,0 +1,245 @@ +import Foundation +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif + +#if os(macOS) +import SweetCookieKit +#endif + +public enum TraeUsageError: LocalizedError, Sendable { + case notLoggedIn + case invalidCredentials + case parseFailed(String) + case networkError(String) + case noAuthToken + + public var errorDescription: String? { + switch self { + case .notLoggedIn: + "Not logged in to Trae. Please log in via trae.ai." + case .invalidCredentials: + "Trae session expired. Please log in again." + case let .parseFailed(message): + "Could not parse Trae usage: \(message)" + case let .networkError(message): + "Trae request failed: \(message)" + case .noAuthToken: + "No Trae authentication found. Please log in to trae.ai in your browser." + } + } +} + +#if os(macOS) +private let traeCookieImportOrder: BrowserCookieImportOrder = + ProviderDefaults.metadata[.trae]?.browserCookieOrder ?? Browser.defaultImportOrder + +public enum TraeAuthImporter { + private static let cookieClient = BrowserCookieClient() + private static let cookieDomains = ["trae.ai", "www.trae.ai", ".trae.ai", ".byteoversea.com"] + + public struct AuthInfo: Sendable { + public let jwtToken: String + public let sourceLabel: String + + public init(jwtToken: String, sourceLabel: String) { + self.jwtToken = jwtToken + self.sourceLabel = sourceLabel + } + } + + public static func importAuth( + browserDetection: BrowserDetection, + logger: ((String) -> Void)? = nil) throws -> AuthInfo + { + let log: (String) -> Void = { msg in logger?("[trae-auth] \(msg)") } + + let installed = traeCookieImportOrder.cookieImportCandidates(using: browserDetection) + for browserSource in installed { + do { + let query = BrowserCookieQuery(domains: self.cookieDomains) + let sources = try Self.cookieClient.records( + matching: query, + in: browserSource, + logger: log) + for source in sources where !source.records.isEmpty { + let cookies = BrowserCookieClient.makeHTTPCookies(source.records, origin: query.origin) + guard !cookies.isEmpty else { continue } + + // Look for X-Cloudide-Session cookie which contains the JWT + for cookie in cookies { + if cookie.name == "X-Cloudide-Session" { + log("Found Trae session cookie in \(source.label)") + // The JWT is the cookie value itself + return AuthInfo( + jwtToken: "Cloud-IDE-JWT \(cookie.value)", + sourceLabel: source.label) + } + } + } + } catch { + BrowserCookieAccessGate.recordIfNeeded(error) + log("\(browserSource.displayName) import failed: \(error.localizedDescription)") + } + } + + throw TraeUsageError.noAuthToken + } +} +#endif + +public struct TraeUsageFetcher: Sendable { + private static let apiURL = URL(string: "https://api-sg-central.trae.ai/trae/api/v1/pay/user_current_entitlement_list")! + @MainActor private static var recentDumps: [String] = [] + + public let browserDetection: BrowserDetection + + public init(browserDetection: BrowserDetection) { + self.browserDetection = browserDetection + } + + public func fetch( + jwtOverride: String? = nil, + logger: ((String) -> Void)? = nil, + now: Date = Date()) async throws -> TraeUsageSnapshot + { + let log: (String) -> Void = { msg in logger?("[trae] \(msg)") } + + // Resolve JWT token (from override or browser) + let jwtToken = try await self.resolveJWTToken(override: jwtOverride, logger: log) + log("[trae] Using JWT authentication: \(jwtToken.prefix(50))...") + + let (data, response) = try await self.fetchWithJWT(jwtToken: jwtToken) + + guard let httpResponse = response as? HTTPURLResponse else { + throw TraeUsageError.networkError("Invalid response") + } + + guard httpResponse.statusCode == 200 else { + if httpResponse.statusCode == 401 || httpResponse.statusCode == 403 { + throw TraeUsageError.invalidCredentials + } + throw TraeUsageError.networkError("HTTP \(httpResponse.statusCode)") + } + + guard let jsonString = String(data: data, encoding: .utf8) else { + throw TraeUsageError.parseFailed("Response was not UTF-8") + } + + do { + return try TraeUsageParser.parse(json: jsonString, now: now) + } catch { + logger?("[trae] Parse failed: \(error.localizedDescription)") + throw error + } + } + + private func fetchWithJWT(jwtToken: String) async throws -> (Data, URLResponse) { + var request = URLRequest(url: Self.apiURL) + request.httpMethod = "GET" + request.setValue(jwtToken, forHTTPHeaderField: "Authorization") + request.setValue("application/json, text/plain, */*", forHTTPHeaderField: "Accept") + request.setValue("Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/143.0.0.0 Safari/537.36", forHTTPHeaderField: "User-Agent") + request.setValue("https://www.trae.ai", forHTTPHeaderField: "Referer") + request.setValue("https://www.trae.ai", forHTTPHeaderField: "Origin") + + return try await URLSession.shared.data(for: request) + } + + private func resolveJWTToken( + override: String?, + logger: ((String) -> Void)?) async throws -> String + { + // If override provided, use it + if let override = override, !override.isEmpty { + logger?("[trae] Using manual JWT token") + // Ensure it has the Cloud-IDE-JWT prefix + if override.hasPrefix("Cloud-IDE-JWT") { + return override + } else { + return "Cloud-IDE-JWT \(override)" + } + } + + #if os(macOS) + // Try to auto-import from browser + do { + let auth = try TraeAuthImporter.importAuth(browserDetection: self.browserDetection, logger: logger) + logger?("[trae] Using JWT from \(auth.sourceLabel)") + return auth.jwtToken + } catch { + logger?("[trae] Auto-import failed: \(error.localizedDescription)") + } + #endif + + throw TraeUsageError.noAuthToken + } + + public func debugRawProbe(jwtOverride: String? = nil) async -> String { + let stamp = ISO8601DateFormatter().string(from: Date()) + var lines: [String] = [] + lines.append("=== Trae Debug Probe @ \(stamp) ===") + lines.append("") + + do { + let jwtToken = try await self.resolveJWTToken( + override: jwtOverride, + logger: { msg in lines.append("[auth] \(msg)") }) + lines.append("JWT Token: \(jwtToken.prefix(50))...") + + let (data, response) = try await self.fetchWithJWT(jwtToken: jwtToken) + + guard let httpResponse = response as? HTTPURLResponse else { + throw TraeUsageError.networkError("Invalid response") + } + + lines.append("") + lines.append("Fetch Response") + lines.append("Status: \(httpResponse.statusCode)") + + if httpResponse.statusCode == 200 { + if let jsonString = String(data: data, encoding: .utf8) { + do { + let snapshot = try TraeUsageParser.parse(json: jsonString, now: Date()) + lines.append("") + lines.append("Trae Usage:") + lines.append(" total=\(snapshot.totalCredits)") + lines.append(" used=\(snapshot.usedCredits)") + lines.append(" plan=\(snapshot.planName)") + if let expiry = snapshot.expiresAt { + lines.append(" expires=\(expiry)") + } + } catch { + lines.append("") + lines.append("Parse Error: \(error.localizedDescription)") + } + } + } else { + lines.append("") + lines.append("Error: HTTP \(httpResponse.statusCode)") + } + + let output = lines.joined(separator: "\n") + Task { @MainActor in Self.recordDump(output) } + return output + } catch { + lines.append("") + lines.append("Probe Failed: \(error.localizedDescription)") + let output = lines.joined(separator: "\n") + Task { @MainActor in Self.recordDump(output) } + return output + } + } + + public static func latestDumps() async -> String { + await MainActor.run { + let result = Self.recentDumps.joined(separator: "\n\n---\n\n") + return result.isEmpty ? "No Trae probe dumps captured yet." : result + } + } + + @MainActor private static func recordDump(_ text: String) { + if self.recentDumps.count >= 5 { self.recentDumps.removeFirst() } + self.recentDumps.append(text) + } +} diff --git a/Sources/CodexBarCore/Providers/Trae/TraeUsageParser.swift b/Sources/CodexBarCore/Providers/Trae/TraeUsageParser.swift new file mode 100644 index 00000000..bc89c487 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Trae/TraeUsageParser.swift @@ -0,0 +1,210 @@ +import Foundation + +enum TraeUsageParser { + static func parse(json: String, now: Date = Date()) throws -> TraeUsageSnapshot { + guard let data = json.data(using: .utf8) else { + throw TraeUsageError.parseFailed("Invalid UTF-8 encoding") + } + + guard let object = try? JSONSerialization.jsonObject(with: data, options: []), + let dict = object as? [String: Any] + else { + throw TraeUsageError.parseFailed("Invalid JSON structure") + } + + guard let entitlementList = dict["user_entitlement_pack_list"] as? [[String: Any]] + else { + if self.looksSignedOut(dict) { + throw TraeUsageError.notLoggedIn + } + throw TraeUsageError.parseFailed("Missing user_entitlement_pack_list") + } + + // Parse all entitlements with valid data (not just status == 1) + var entitlements: [TraeEntitlementInfo] = [] + var allFeatures = TraeFeatures() + + for entitlement in entitlementList { + guard let baseInfo = entitlement["entitlement_base_info"] as? [String: Any], + let quota = baseInfo["quota"] as? [String: Any], + let usage = entitlement["usage"] as? [String: Any] + else { continue } + + // Get status (0 or 1) - include all for display like web UI does + let status = entitlement["status"] as? Int ?? 0 + + // Parse fast request quota + var totalCredits: Double = 0 + var isUnlimited = false + if let fastLimit = quota["premium_model_fast_request_limit"] as? Double { + if fastLimit == -1 { + isUnlimited = true + } else if fastLimit > 0 { + totalCredits = fastLimit + } + } else if let fastLimitInt = quota["premium_model_fast_request_limit"] as? Int { + if fastLimitInt == -1 { + isUnlimited = true + } else if fastLimitInt > 0 { + totalCredits = Double(fastLimitInt) + } + } + + // Skip entitlements with no quota at all + if totalCredits == 0 && !isUnlimited { + continue + } + + // Parse used amount + var usedCredits: Double = 0 + if let fastUsed = usage["premium_model_fast_amount"] as? Double { + usedCredits = fastUsed + } else if let fastUsedInt = usage["premium_model_fast_amount"] as? Int { + usedCredits = Double(fastUsedInt) + } + + // Determine reset/expire date + // For Pro plan (product_type: 1), use next_billing_time if available + // For packages (product_type: 2), use end_time + let productType = baseInfo["product_type"] as? Int ?? 0 + var resetsAt: Date? + + if productType == 1 { + // Pro plan uses next_billing_time for reset + if let nextBilling = entitlement["next_billing_time"] as? Double, nextBilling > 0 { + resetsAt = Date(timeIntervalSince1970: nextBilling) + } else if let endTime = baseInfo["end_time"] as? Double { + resetsAt = Date(timeIntervalSince1970: endTime) + } else if let endTimeInt = baseInfo["end_time"] as? Int { + resetsAt = Date(timeIntervalSince1970: Double(endTimeInt)) + } + } else { + // Packages use end_time as expiry + if let endTime = baseInfo["end_time"] as? Double { + resetsAt = Date(timeIntervalSince1970: endTime) + } else if let endTimeInt = baseInfo["end_time"] as? Int { + resetsAt = Date(timeIntervalSince1970: Double(endTimeInt)) + } + } + + // Determine name based on product_type and product_extra + let name = extractEntitlementName(baseInfo: baseInfo, productType: productType) + + entitlements.append(TraeEntitlementInfo( + name: name, + totalCredits: totalCredits, + usedCredits: usedCredits, + expiresAt: resetsAt, + productType: productType, + isUnlimited: isUnlimited, + isActive: status == 1)) + + // Collect features only from active entitlements (status == 1) + // to avoid showing features from expired/inactive plans + if status == 1 { + let hasSoloBuilder = quota["enable_solo_builder"] as? Bool ?? false + let hasSoloCoder = quota["enable_solo_coder"] as? Bool ?? false + + var hasUnlimitedSlow = false + if let slowLimit = quota["premium_model_slow_request_limit"] as? Double, slowLimit == -1 { + hasUnlimitedSlow = true + } else if let slowLimit = quota["premium_model_slow_request_limit"] as? Int, slowLimit == -1 { + hasUnlimitedSlow = true + } + + var hasUnlimitedAdvanced = false + if let advancedLimit = quota["premium_model_advanced_request_limit"] as? Double, advancedLimit == -1 { + hasUnlimitedAdvanced = true + } else if let advancedLimit = quota["premium_model_advanced_request_limit"] as? Int, advancedLimit == -1 { + hasUnlimitedAdvanced = true + } + + var hasUnlimitedAutocomplete = false + if let autoLimit = quota["auto_completion_limit"] as? Double, autoLimit == -1 { + hasUnlimitedAutocomplete = true + } else if let autoLimit = quota["auto_completion_limit"] as? Int, autoLimit == -1 { + hasUnlimitedAutocomplete = true + } + + // Merge features - if any active entitlement has it, keep it + if hasSoloBuilder { allFeatures.hasSoloBuilder = true } + if hasSoloCoder { allFeatures.hasSoloCoder = true } + if hasUnlimitedSlow { allFeatures.hasUnlimitedSlow = true } + if hasUnlimitedAdvanced { allFeatures.hasUnlimitedAdvanced = true } + if hasUnlimitedAutocomplete { allFeatures.hasUnlimitedAutocomplete = true } + } + } + + guard !entitlements.isEmpty else { + throw TraeUsageError.parseFailed("No entitlements found with valid quotas") + } + + // Sort entitlements: + // 1. Pro Plan first + // 2. Active entitlements before inactive + // 3. Then by expiry date (soonest first) + let sortedEntitlements = entitlements.sorted { a, b in + // Pro Plan always first + if a.productType == 1 && b.productType != 1 { return true } + if a.productType != 1 && b.productType == 1 { return false } + + // Active before inactive + if a.isActive && !b.isActive { return true } + if !a.isActive && b.isActive { return false } + + // Then by expiry date + if let aDate = a.expiresAt, let bDate = b.expiresAt { + return aDate < bDate + } + return a.expiresAt != nil + } + + return TraeUsageSnapshot( + entitlements: sortedEntitlements, + features: allFeatures, + totalEntitlements: entitlementList.count, + activeEntitlements: sortedEntitlements.filter { $0.isActive }.count, + updatedAt: now) + } + + /// Extract a human-readable name for the entitlement + private static func extractEntitlementName(baseInfo: [String: Any], productType: Int) -> String { + // For Pro plan, always use "Pro Plan" + if productType == 1 { + return "Pro Plan" + } + + // For packages, try to get descriptive name from product_extra + if let productExtra = baseInfo["product_extra"] as? [String: Any], + let packageExtra = productExtra["package_extra"] as? [String: Any] { + + // Get package source type for naming + if let sourceType = packageExtra["package_source_type"] as? Int { + switch sourceType { + case 1: + return "Extra Package" + case 2: + return "Extra Package (Official Bonus)" + case 6: + return "Extra Package (Anniversary Treat)" + default: + break + } + } + } + + return "Extra Package" + } + + private static func looksSignedOut(_ dict: [String: Any]) -> Bool { + if let error = dict["error"] as? String { + let lower = error.lowercased() + return lower.contains("unauthorized") || lower.contains("not logged in") || lower.contains("auth") + } + if let message = dict["message"] as? String { + let lower = message.lowercased() + return lower.contains("login") || lower.contains("sign in") || lower.contains("auth") + } + return false + } +} diff --git a/Sources/CodexBarCore/Providers/Trae/TraeUsageSnapshot.swift b/Sources/CodexBarCore/Providers/Trae/TraeUsageSnapshot.swift new file mode 100644 index 00000000..9979bb16 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Trae/TraeUsageSnapshot.swift @@ -0,0 +1,164 @@ +import Foundation + +public struct TraeEntitlementInfo: Sendable { + public let name: String + public let totalCredits: Double + public let usedCredits: Double + public let expiresAt: Date? + public let productType: Int + public let isUnlimited: Bool + public let isActive: Bool + + public init(name: String, totalCredits: Double, usedCredits: Double, expiresAt: Date?, productType: Int, isUnlimited: Bool, isActive: Bool = true) { + self.name = name + self.totalCredits = totalCredits + self.usedCredits = usedCredits + self.expiresAt = expiresAt + self.productType = productType + self.isUnlimited = isUnlimited + self.isActive = isActive + } + + public var remainingCredits: Double { + if isUnlimited { return -1 } + return max(0, totalCredits - usedCredits) + } + + public var usedPercent: Double { + if isUnlimited || totalCredits <= 0 { return 0 } + return min(100, (usedCredits / totalCredits) * 100) + } + + public var remainingDescription: String { + // Keep simple - just the credit count, no "resets" prefix + if isUnlimited { + return "Unlimited" + } + let remaining = remainingCredits + if remaining == 0 { + return "0 Left" + } + // Format with 0-2 decimal places based on value + if remaining == floor(remaining) { + return "\(Int(remaining)) Left" + } else { + return String(format: "%.2f Left", remaining) + } + } +} + +public struct TraeFeatures: Sendable { + public var hasSoloBuilder: Bool + public var hasSoloCoder: Bool + public var hasUnlimitedSlow: Bool + public var hasUnlimitedAdvanced: Bool + public var hasUnlimitedAutocomplete: Bool + + public init(hasSoloBuilder: Bool = false, hasSoloCoder: Bool = false, hasUnlimitedSlow: Bool = false, hasUnlimitedAdvanced: Bool = false, hasUnlimitedAutocomplete: Bool = false) { + self.hasSoloBuilder = hasSoloBuilder + self.hasSoloCoder = hasSoloCoder + self.hasUnlimitedSlow = hasUnlimitedSlow + self.hasUnlimitedAdvanced = hasUnlimitedAdvanced + self.hasUnlimitedAutocomplete = hasUnlimitedAutocomplete + } + + public var hasAnyFeatures: Bool { + hasSoloBuilder || hasSoloCoder || hasUnlimitedSlow || hasUnlimitedAdvanced || hasUnlimitedAutocomplete + } +} + +public struct TraeUsageSnapshot: Sendable { + public let entitlements: [TraeEntitlementInfo] + public let features: TraeFeatures + public let totalEntitlements: Int + public let activeEntitlements: Int + public let updatedAt: Date + + public init(entitlements: [TraeEntitlementInfo], features: TraeFeatures, totalEntitlements: Int, activeEntitlements: Int, updatedAt: Date) { + self.entitlements = entitlements + self.features = features + self.totalEntitlements = totalEntitlements + self.activeEntitlements = activeEntitlements + self.updatedAt = updatedAt + } + + // Legacy accessors for compatibility + public var totalCredits: Double { + entitlements.reduce(0) { $0 + ($1.isUnlimited ? 0 : $1.totalCredits) } + } + + public var usedCredits: Double { + entitlements.reduce(0) { $0 + $1.usedCredits } + } + + public var planName: String { + let proPlan = entitlements.first { $0.productType == 1 } + return proPlan?.name ?? "Trae" + } + + public var expiresAt: Date? { + entitlements.compactMap { $0.expiresAt }.min() + } +} + +extension TraeUsageSnapshot { + public func toUsageSnapshot(now: Date = Date()) -> UsageSnapshot { + // Map up to 3 entitlements to primary/secondary/tertiary windows + let primary = self.makeWindow(from: entitlements.indices.contains(0) ? entitlements[0] : nil, now: now) + let secondary = self.makeWindow(from: entitlements.indices.contains(1) ? entitlements[1] : nil, now: now) + let tertiary = self.makeWindow(from: entitlements.indices.contains(2) ? entitlements[2] : nil, now: now) + + // Build rich plan description with features + var planDetails: [String] = [] + if features.hasSoloBuilder { + planDetails.append("Solo Builder") + } + if features.hasSoloCoder { + planDetails.append("Solo Coder") + } + if features.hasUnlimitedSlow { + planDetails.append("Slow Requests") + } + if features.hasUnlimitedAdvanced { + planDetails.append("Advanced Model") + } + if features.hasUnlimitedAutocomplete { + planDetails.append("Autocomplete") + } + + let planName = entitlements.first?.name ?? "Trae" + let richPlanName = planDetails.isEmpty ? planName : "\(planName) (\(planDetails.joined(separator: ", ")))" + + let identity = ProviderIdentitySnapshot( + providerID: .trae, + accountEmail: nil, + accountOrganization: nil, + loginMethod: richPlanName) + + return UsageSnapshot( + primary: primary, + secondary: secondary, + tertiary: tertiary, + providerCost: nil, + updatedAt: self.updatedAt, + identity: identity) + } + + private func makeWindow(from entitlement: TraeEntitlementInfo?, now: Date) -> RateWindow? { + guard let entitlement else { return nil } + + let windowMinutes: Int? = { + guard let expiry = entitlement.expiresAt else { return nil } + let minutes = Int(expiry.timeIntervalSince(now) / 60) + return minutes > 0 ? minutes : nil + }() + + // Simple approach: just percentage and reset date + // UsageFormatter.resetLine will show "Resets [date]" or countdown + return RateWindow( + usedPercent: entitlement.usedPercent, + windowMinutes: windowMinutes, + resetsAt: entitlement.expiresAt, + resetDescription: nil) + } +} diff --git a/Sources/CodexBarCore/TokenAccountSupportCatalog+Data.swift b/Sources/CodexBarCore/TokenAccountSupportCatalog+Data.swift index b2bc65d4..4b709597 100644 --- a/Sources/CodexBarCore/TokenAccountSupportCatalog+Data.swift +++ b/Sources/CodexBarCore/TokenAccountSupportCatalog+Data.swift @@ -51,5 +51,12 @@ extension TokenAccountSupportCatalog { injection: .cookieHeader, requiresManualCookieSource: true, cookieName: nil), + .trae: TokenAccountSupport( + title: "Session tokens", + subtitle: "Store multiple Trae Cookie headers.", + placeholder: "Cookie: …", + injection: .cookieHeader, + requiresManualCookieSource: true, + cookieName: nil), ] } diff --git a/Sources/CodexBarCore/Vendored/CostUsage/CostUsageScanner.swift b/Sources/CodexBarCore/Vendored/CostUsage/CostUsageScanner.swift index c50b106c..121ced7f 100644 --- a/Sources/CodexBarCore/Vendored/CostUsage/CostUsageScanner.swift +++ b/Sources/CodexBarCore/Vendored/CostUsage/CostUsageScanner.swift @@ -97,6 +97,8 @@ enum CostUsageScanner { return CostUsageDailyReport(data: [], summary: nil) case .amp: return CostUsageDailyReport(data: [], summary: nil) + case .trae: + return CostUsageDailyReport(data: [], summary: nil) case .synthetic: return CostUsageDailyReport(data: [], summary: nil) } diff --git a/Sources/CodexBarWidget/CodexBarWidgetProvider.swift b/Sources/CodexBarWidget/CodexBarWidgetProvider.swift index 1634611e..703ed8fd 100644 --- a/Sources/CodexBarWidget/CodexBarWidgetProvider.swift +++ b/Sources/CodexBarWidget/CodexBarWidgetProvider.swift @@ -59,6 +59,7 @@ enum ProviderChoice: String, AppEnum { case .kimik2: return nil // Kimi K2 not yet supported in widgets case .amp: return nil // Amp not yet supported in widgets case .synthetic: return nil // Synthetic not yet supported in widgets + case .trae: return nil // Trae not yet supported in widgets } } } diff --git a/Sources/CodexBarWidget/CodexBarWidgetViews.swift b/Sources/CodexBarWidget/CodexBarWidgetViews.swift index ed39b450..2872cfd6 100644 --- a/Sources/CodexBarWidget/CodexBarWidgetViews.swift +++ b/Sources/CodexBarWidget/CodexBarWidgetViews.swift @@ -275,6 +275,7 @@ private struct ProviderSwitchChip: View { case .kimik2: "Kimi K2" case .amp: "Amp" case .synthetic: "Synthetic" + case .trae: "Trae" } } } @@ -605,6 +606,8 @@ enum WidgetColors { Color(red: 220 / 255, green: 38 / 255, blue: 38 / 255) // Amp red case .synthetic: Color(red: 20 / 255, green: 20 / 255, blue: 20 / 255) // Synthetic charcoal + case .trae: + Color(red: 0 / 255, green: 122 / 255, blue: 255 / 255) // Trae blue } } } diff --git a/Tests/CodexBarTests/PreferencesPaneSmokeTests.swift b/Tests/CodexBarTests/PreferencesPaneSmokeTests.swift index ff20e395..2603707c 100644 --- a/Tests/CodexBarTests/PreferencesPaneSmokeTests.swift +++ b/Tests/CodexBarTests/PreferencesPaneSmokeTests.swift @@ -64,6 +64,7 @@ struct PreferencesPaneSmokeTests { kimiK2TokenStore: InMemoryKimiK2TokenStore(), augmentCookieStore: InMemoryCookieHeaderStore(), ampCookieStore: InMemoryCookieHeaderStore(), + traeCookieStore: InMemoryCookieHeaderStore(), copilotTokenStore: InMemoryCopilotTokenStore(), tokenAccountStore: InMemoryTokenAccountStore()) } diff --git a/Tests/CodexBarTests/ProviderIconResourcesTests.swift b/Tests/CodexBarTests/ProviderIconResourcesTests.swift index 7ecbe7b2..4a7a8678 100644 --- a/Tests/CodexBarTests/ProviderIconResourcesTests.swift +++ b/Tests/CodexBarTests/ProviderIconResourcesTests.swift @@ -21,6 +21,7 @@ struct ProviderIconResourcesTests { "antigravity", "factory", "copilot", + "trae", ] for slug in slugs { let url = resources.appending(path: "ProviderIcon-\(slug).svg") diff --git a/Tests/CodexBarTests/ProvidersPaneCoverageTests.swift b/Tests/CodexBarTests/ProvidersPaneCoverageTests.swift index e3f2bd55..8d07d711 100644 --- a/Tests/CodexBarTests/ProvidersPaneCoverageTests.swift +++ b/Tests/CodexBarTests/ProvidersPaneCoverageTests.swift @@ -35,6 +35,7 @@ struct ProvidersPaneCoverageTests { kimiK2TokenStore: InMemoryKimiK2TokenStore(), augmentCookieStore: InMemoryCookieHeaderStore(), ampCookieStore: InMemoryCookieHeaderStore(), + traeCookieStore: InMemoryCookieHeaderStore(), copilotTokenStore: InMemoryCopilotTokenStore(), tokenAccountStore: InMemoryTokenAccountStore()) } diff --git a/Tests/CodexBarTests/SettingsStoreAdditionalTests.swift b/Tests/CodexBarTests/SettingsStoreAdditionalTests.swift index c70eb957..638ca698 100644 --- a/Tests/CodexBarTests/SettingsStoreAdditionalTests.swift +++ b/Tests/CodexBarTests/SettingsStoreAdditionalTests.swift @@ -78,6 +78,7 @@ struct SettingsStoreAdditionalTests { kimiK2TokenStore: InMemoryKimiK2TokenStore(), augmentCookieStore: InMemoryCookieHeaderStore(), ampCookieStore: InMemoryCookieHeaderStore(), + traeCookieStore: InMemoryCookieHeaderStore(), copilotTokenStore: InMemoryCopilotTokenStore(), tokenAccountStore: InMemoryTokenAccountStore()) } diff --git a/Tests/CodexBarTests/SettingsStoreCoverageTests.swift b/Tests/CodexBarTests/SettingsStoreCoverageTests.swift index a84ce945..789b7644 100644 --- a/Tests/CodexBarTests/SettingsStoreCoverageTests.swift +++ b/Tests/CodexBarTests/SettingsStoreCoverageTests.swift @@ -124,6 +124,7 @@ struct SettingsStoreCoverageTests { settings.ensureKimiK2APITokenLoaded() settings.ensureAugmentCookieLoaded() settings.ensureAmpCookieLoaded() + settings.ensureTraeCookieLoaded() settings.ensureCopilotAPITokenLoaded() settings.ensureTokenAccountsLoaded() @@ -177,6 +178,7 @@ struct SettingsStoreCoverageTests { kimiK2TokenStore: InMemoryKimiK2TokenStore(), augmentCookieStore: InMemoryCookieHeaderStore(), ampCookieStore: InMemoryCookieHeaderStore(), + traeCookieStore: InMemoryCookieHeaderStore(), copilotTokenStore: InMemoryCopilotTokenStore(), tokenAccountStore: InMemoryTokenAccountStore()) } diff --git a/Tests/CodexBarTests/SettingsStoreTests.swift b/Tests/CodexBarTests/SettingsStoreTests.swift index 108d31bd..366cf077 100644 --- a/Tests/CodexBarTests/SettingsStoreTests.swift +++ b/Tests/CodexBarTests/SettingsStoreTests.swift @@ -363,6 +363,7 @@ struct SettingsStoreTests { .jetbrains, .kimik2, .amp, + .trae, .synthetic, ]) diff --git a/Tests/CodexBarTests/TraeUsageParserTests.swift b/Tests/CodexBarTests/TraeUsageParserTests.swift new file mode 100644 index 00000000..1bb693a1 --- /dev/null +++ b/Tests/CodexBarTests/TraeUsageParserTests.swift @@ -0,0 +1,284 @@ +import Foundation +import Testing +@testable import CodexBarCore + +@Suite +struct TraeUsageParserTests { + @Test + func parsesActiveEntitlementWithProPlan() throws { + let now = Date(timeIntervalSince1970: 1_700_000_000) + let json = """ + { + "user_entitlement_pack_list": [ + { + "entitlement_base_info": { + "quota": { + "premium_model_fast_request_limit": 600 + }, + "end_time": 1772277112, + "product_type": 1 + }, + "status": 1, + "usage": { + "premium_model_fast_amount": 91.26527 + } + } + ] + } + """ + + let snapshot = try TraeUsageParser.parse(json: json, now: now) + + #expect(snapshot.totalCredits == 600) + #expect(snapshot.usedCredits == 91.26527) + #expect(snapshot.planName == "Pro Plan") + #expect(snapshot.expiresAt != nil) + } + + @Test + func parsesMultipleActiveEntitlements() throws { + let now = Date(timeIntervalSince1970: 1_700_000_000) + let json = """ + { + "user_entitlement_pack_list": [ + { + "entitlement_base_info": { + "quota": { + "premium_model_fast_request_limit": 600 + }, + "end_time": 1772277112, + "product_type": 1 + }, + "status": 1, + "usage": { + "premium_model_fast_amount": 100 + } + }, + { + "entitlement_base_info": { + "quota": { + "premium_model_fast_request_limit": 300 + }, + "end_time": 1772277187, + "product_type": 2 + }, + "status": 1, + "usage": { + "premium_model_fast_amount": 50 + } + } + ] + } + """ + + let snapshot = try TraeUsageParser.parse(json: json, now: now) + + #expect(snapshot.totalCredits == 900) + #expect(snapshot.usedCredits == 150) + #expect(snapshot.planName.contains("Pro Plan")) + #expect(snapshot.planName.contains("Package")) + } + + @Test + func includesInactiveEntitlementsWithQuota() throws { + let now = Date(timeIntervalSince1970: 1_700_000_000) + let json = """ + { + "user_entitlement_pack_list": [ + { + "entitlement_base_info": { + "quota": { + "premium_model_fast_request_limit": 600 + }, + "end_time": 1772277112, + "product_type": 1 + }, + "status": 1, + "usage": { + "premium_model_fast_amount": 100 + } + }, + { + "entitlement_base_info": { + "quota": { + "premium_model_fast_request_limit": 300 + }, + "end_time": 1772277187, + "product_type": 2 + }, + "status": 0, + "usage": { + "premium_model_fast_amount": 50 + } + } + ] + } + """ + + let snapshot = try TraeUsageParser.parse(json: json, now: now) + + // Parser now includes all entitlements with valid quota (not just status=1) + #expect(snapshot.totalCredits == 900) + #expect(snapshot.usedCredits == 150) + #expect(snapshot.activeEntitlements == 1) + #expect(snapshot.totalEntitlements == 2) + } + + @Test + func calculatesUsagePercentCorrectly() throws { + let now = Date(timeIntervalSince1970: 1_700_000_000) + let json = """ + { + "user_entitlement_pack_list": [ + { + "entitlement_base_info": { + "quota": { + "premium_model_fast_request_limit": 100 + }, + "end_time": 1772277112, + "product_type": 1 + }, + "status": 1, + "usage": { + "premium_model_fast_amount": 25 + } + } + ] + } + """ + + let snapshot = try TraeUsageParser.parse(json: json, now: now) + let usage = snapshot.toUsageSnapshot(now: now) + + #expect(usage.primary?.usedPercent == 25) + } + + @Test + func handlesIntegerRequestLimits() throws { + let now = Date(timeIntervalSince1970: 1_700_000_000) + let json = """ + { + "user_entitlement_pack_list": [ + { + "entitlement_base_info": { + "quota": { + "premium_model_fast_request_limit": 600 + }, + "end_time": 1772277112, + "product_type": 1 + }, + "status": 1, + "usage": { + "premium_model_fast_amount": 91 + } + } + ] + } + """ + + let snapshot = try TraeUsageParser.parse(json: json, now: now) + + #expect(snapshot.totalCredits == 600) + #expect(snapshot.usedCredits == 91) + } + + @Test + func missingEntitlementListThrowsParseFailed() { + let json = "{\"error\": \"not found\"}" + + #expect { + try TraeUsageParser.parse(json: json) + } throws: { error in + guard case let TraeUsageError.parseFailed(message) = error else { return false } + return message.contains("user_entitlement_pack_list") + } + } + + @Test + func emptyEntitlementsWithQuotaThrowsParseFailed() { + let json = """ + { + "user_entitlement_pack_list": [ + { + "entitlement_base_info": { + "quota": {"premium_model_fast_request_limit": 0}, + "end_time": 1772277112, + "product_type": 1 + }, + "status": 0, + "usage": {"premium_model_fast_amount": 0} + } + ] + } + """ + + #expect { + try TraeUsageParser.parse(json: json) + } throws: { error in + guard case let TraeUsageError.parseFailed(message) = error else { return false } + return message.contains("No entitlements found with valid quotas") + } + } + + @Test + func unauthorizedResponseThrowsNotLoggedIn() { + let json = """ + { + "error": "unauthorized" + } + """ + + #expect { + try TraeUsageParser.parse(json: json) + } throws: { error in + guard case TraeUsageError.notLoggedIn = error else { return false } + return true + } + } + + @Test + func invalidJsonThrowsParseFailed() { + let json = "not valid json" + + #expect { + try TraeUsageParser.parse(json: json) + } throws: { error in + guard case let TraeUsageError.parseFailed(message) = error else { return false } + return message.contains("Invalid JSON") + } + } + + @Test + func calculatesEarliestExpiry() throws { + let now = Date(timeIntervalSince1970: 1_700_000_000) + let json = """ + { + "user_entitlement_pack_list": [ + { + "entitlement_base_info": { + "quota": {"premium_model_fast_request_limit": 600}, + "end_time": 1772277112, + "product_type": 1 + }, + "status": 1, + "usage": {"premium_model_fast_amount": 100} + }, + { + "entitlement_base_info": { + "quota": {"premium_model_fast_request_limit": 300}, + "end_time": 1700000100, + "product_type": 2 + }, + "status": 1, + "usage": {"premium_model_fast_amount": 50} + } + ] + } + """ + + let snapshot = try TraeUsageParser.parse(json: json, now: now) + + // Should use the earlier expiry (1700000100) + #expect(snapshot.expiresAt == Date(timeIntervalSince1970: 1700000100)) + } +} diff --git a/Tests/CodexBarTests/UsageStoreCoverageTests.swift b/Tests/CodexBarTests/UsageStoreCoverageTests.swift index 888dbbb1..6e9bbedf 100644 --- a/Tests/CodexBarTests/UsageStoreCoverageTests.swift +++ b/Tests/CodexBarTests/UsageStoreCoverageTests.swift @@ -141,6 +141,7 @@ struct UsageStoreCoverageTests { kimiK2TokenStore: InMemoryKimiK2TokenStore(), augmentCookieStore: InMemoryCookieHeaderStore(), ampCookieStore: InMemoryCookieHeaderStore(), + traeCookieStore: InMemoryCookieHeaderStore(), copilotTokenStore: InMemoryCopilotTokenStore(), tokenAccountStore: InMemoryTokenAccountStore()) } diff --git a/docs/providers.md b/docs/providers.md index d25dfb88..b1a2cbb3 100644 --- a/docs/providers.md +++ b/docs/providers.md @@ -1,5 +1,5 @@ --- -summary: "Provider data sources and parsing overview (Codex, Claude, Gemini, Antigravity, Cursor, Droid/Factory, z.ai, Copilot, Kimi, Kimi K2, Kiro, Vertex AI, Augment, Amp, JetBrains AI)." +summary: "Provider data sources and parsing overview (Codex, Claude, Gemini, Antigravity, Cursor, Droid/Factory, z.ai, Copilot, Kimi, Kimi K2, Kiro, Vertex AI, Augment, Amp, Trae, JetBrains AI)." read_when: - Adding or modifying provider fetch/parsing - Adjusting provider labels, toggles, or metadata @@ -34,6 +34,7 @@ until the session is invalid, to avoid repeated Keychain prompts. | Vertex AI | Google ADC OAuth (gcloud) → Cloud Monitoring quota usage (`oauth`). | | JetBrains AI | Local XML quota file (`local`). | | Amp | Web settings page via browser cookies (`web`). | +| Trae | Web API via JWT authentication (`web`). | ## Codex - Web dashboard (when enabled): `https://chatgpt.com/codex/settings/usage` via WebView + browser cookies. @@ -139,4 +140,14 @@ until the session is invalid, to avoid repeated Keychain prompts. - Parses Amp Free usage from the settings HTML. - Status: none yet. - Details: `docs/amp.md`. + +## Trae +- Web API via JWT authentication (`https://api-sg-central.trae.ai/trae/api/v1/pay/user_current_entitlement_list`). +- Uses `Authorization: Cloud-IDE-JWT ` header for authentication. +- **Manual entry only** - JWT is not stored in browser cookies (copied from DevTools Network tab). +- Shows individual entitlement usage (Pro Plan + Extra Packages) in primary/secondary/tertiary windows. +- Tracks `premium_model_fast_request_limit` and `premium_model_fast_amount`. +- Status: none yet. +- Details: `docs/trae.md`. + See also: `docs/provider.md` for architecture notes. diff --git a/docs/trae.md b/docs/trae.md new file mode 100644 index 00000000..c39610ea --- /dev/null +++ b/docs/trae.md @@ -0,0 +1,127 @@ +--- +summary: "Trae provider data sources: Web API via JWT authentication, local usage tracking, and authentication details." +read_when: + - Adding or modifying Trae usage/status parsing + - Updating Trae web endpoints or JWT authentication + - Reviewing local usage scanning +--- + +# Trae provider + +Trae is ByteDance's AI-powered IDE (similar to Cursor), built on VS Code with integrated Claude 3.7 Sonnet and GPT-4o support. This provider tracks your fast request quota usage via the web API. + +## Data source: Web API (JWT) + +The Trae provider uses JWT authentication to fetch your current entitlement/usage data from the Trae API. Unlike other providers, the JWT is not stored in browser cookies or local storage - it exists only in HTTP headers during API requests. + +### API Endpoint +- **Primary:** `https://api-sg-central.trae.ai/trae/api/v1/pay/user_current_entitlement_list` +- **Method:** GET +- **Authentication:** JWT token in `Authorization: Cloud-IDE-JWT ` header + +### JWT Authentication (Manual Entry Required) + +**Important:** The JWT token is not stored in browser cookies or localStorage. It only exists in the HTTP Authorization header during API requests. Therefore, **automatic browser import is not available** - you must manually copy the JWT from your browser's DevTools. + +**How to obtain the JWT token:** + +1. Open https://www.trae.ai in Chrome or Safari and log in to your account +2. Open browser DevTools (Cmd+Option+I on Mac) +3. Go to the **Network** tab +4. Refresh the page or click on the **Usage** menu in Trae +5. Look for a network request to `user_current_entitlement_list` in the list +6. Click on that request +7. In the right panel, find **Headers** → **Request Headers** +8. Locate the `Authorization` header (it starts with `Cloud-IDE-JWT`) +9. Copy the entire header value +10. Open CodexBar → Settings → Providers → Trae +11. Paste the copied value in the "JWT Token" field + +**Token Format:** +- Full format: `Cloud-IDE-JWT eyJhbGciOiJSUzI1NiIs...` +- Or just the JWT payload: `eyJhbGciOiJSUzI1NiIs...` +- Both formats are accepted + +**Session Duration:** +- The JWT typically expires after 24-48 hours +- When it expires, you'll see "Trae session expired" error +- Simply repeat the steps above to get a fresh JWT + +### Response Structure +The API returns a list of user entitlements (`user_entitlement_pack_list`), each containing: + +```json +{ + "user_entitlement_pack_list": [ + { + "entitlement_base_info": { + "quota": { + "premium_model_fast_request_limit": 600 + }, + "end_time": 1772277112, + "product_type": 1 + }, + "status": 1, + "usage": { + "premium_model_fast_amount": 91.26527 + } + } + ] +} +``` + +**Key Fields:** +- `premium_model_fast_request_limit`: Total fast request quota (-1 = unlimited) +- `premium_model_fast_amount`: Amount consumed (in credit units) +- `end_time`: Expiration timestamp (Unix seconds) +- `product_type`: 1 = Subscription, 2 = Package/Add-on +- `status`: 1 = Active, 0 = Inactive + +### Usage Calculation +- **Primary Window:** Pro Plan entitlement (product_type=1) +- **Secondary Window:** Extra Package entitlements (product_type=2) +- **Tertiary Window:** Additional packages if present +- **Used Percent:** Based on `premium_model_fast_amount / premium_model_fast_request_limit` +- **Reset Time:** Pro plan uses `next_billing_time`, packages use `end_time` +- **Unlimited Plans:** When `premium_model_fast_request_limit` is -1, shows unlimited status + +### Settings +- **Authentication:** Manual JWT entry only (no automatic browser import) +- **Dashboard URL:** `https://www.trae.ai/account-setting` +- **Default Enabled:** No (opt-in provider) + +## Error Handling + +Common error scenarios: +- **401/403:** JWT expired or invalid → Re-copy fresh JWT from browser +- **No JWT Token:** No token entered → Follow steps above to obtain JWT +- **Empty Entitlements:** No active plan → Check your Trae account has an active subscription +- **Network Errors:** Retry with exponential backoff + +## Implementation Notes + +- Uses JWT-based authentication (not cookie-based sessions) +- **No automatic browser import** - JWT is not stored in accessible browser storage +- Follows the same pattern as Amp and OpenCode providers +- JWT token cached in Keychain for performance (reused until invalid) +- No status page integration (Trae does not provide a public status API) +- No CLI integration (Trae does not expose usage via CLI) + +## Key Files + +- `Sources/CodexBarCore/Providers/Trae/TraeProviderDescriptor.swift` - Provider metadata +- `Sources/CodexBarCore/Providers/Trae/TraeUsageFetcher.swift` - HTTP client and JWT authentication +- `Sources/CodexBarCore/Providers/Trae/TraeUsageParser.swift` - API response parsing +- `Sources/CodexBarCore/Providers/Trae/TraeUsageSnapshot.swift` - Data model +- `Sources/CodexBar/Providers/Trae/TraeSettingsStore.swift` - Settings extensions +- `Sources/CodexBar/Providers/Trae/TraeProviderImplementation.swift` - UI hooks + +## Testing + +Tests mirror the pattern in `AmpUsageParserTests.swift`: +- Parse valid entitlement list responses +- Handle multiple entitlements (subscription + packages) +- Calculate usage percentages correctly +- Handle unlimited plans (-1 limit) +- Handle expired/inactive entitlements +- Error cases: 401, network failures, malformed JSON