From 9c827cbeac55aec9c35ab1d466fa686d18c0b522 Mon Sep 17 00:00:00 2001 From: Factory Agent Date: Fri, 23 Oct 2026 16:14:00 +0700 Subject: [PATCH 1/2] Add Trae CN credits provider Port the Trae CN credit allowance from upstream CodexBar (0703a0c01, #4682). The signed-in api.trae.cn browser session exchanges its cookies for a token at GetUserToken, reads the current entitlement list with that token, and shows credits used against the total, the remaining credits and the plan name. Browser sessions are tried in order and only move on after Trae rejects one. Zero allowance shows details without a gauge. No reset date is inferred. --- README.md | 1 + .../src/commands/provider_settings.rs | 8 + .../providers/icons/ProviderIcon-traecn.svg | 3 + .../src/components/providers/providerIcons.ts | 3 + .../desktop-tauri/src/test/providerCatalog.ts | 1 + docs/PROVIDERS.md | 33 ++ rust/src/core/provider/id.rs | 3 +- rust/src/core/provider/spec.rs | 10 +- rust/src/core/provider_factory.rs | 11 +- rust/src/core/provider_registry_snapshot.txt | 4 + rust/src/core/token_accounts.rs | 3 +- rust/src/providers/mod.rs | 2 + rust/src/providers/traecn/mod.rs | 399 ++++++++++++++ rust/src/providers/traecn/tests.rs | 510 ++++++++++++++++++ rust/src/settings/provider_accessors.rs | 3 +- 15 files changed, 986 insertions(+), 8 deletions(-) create mode 100644 apps/desktop-tauri/src/components/providers/icons/ProviderIcon-traecn.svg create mode 100644 rust/src/providers/traecn/mod.rs create mode 100644 rust/src/providers/traecn/tests.rs diff --git a/README.md b/README.md index b54847e5e1a..76058c64f5e 100755 --- a/README.md +++ b/README.md @@ -151,6 +151,7 @@ For Claude, browser cookies (the `sessionKey`) are preferred, because they match | Future AGI | API Key + secret key | Billing-period AI credits and gateway requests against free allowances, estimated cost | | Munito | Admin API Key | Today's and this month's organization spend against spend limits | | Claude Platform | Browser cookies / manual Cookie header | Console API credit and purchased credit, other grants, month-to-date spend | +| Trae CN | Browser cookies / manual Cookie header | Credit allowance used against total, remaining credits, plan name | diff --git a/apps/desktop-tauri/src-tauri/src/commands/provider_settings.rs b/apps/desktop-tauri/src-tauri/src/commands/provider_settings.rs index 3834413b821..ca13dba6e38 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/provider_settings.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/provider_settings.rs @@ -328,6 +328,7 @@ fn cookie_source_provider(provider_id: &str) -> Option ProviderId::Xapi, "parallel" => ProviderId::Parallel, "claudeplatform" => ProviderId::ClaudePlatform, + "traecn" => ProviderId::TraeCn, _ => return None, }) } @@ -1028,6 +1029,13 @@ pub fn cookie_source_options_for(provider_id: &str, lang: Language) -> Vec vec![ + cookie_option(lang, "auto", None), + cookie_option(lang, "manual", None), + cookie_option(lang, "off", None), + ], _ => Vec::new(), } } diff --git a/apps/desktop-tauri/src/components/providers/icons/ProviderIcon-traecn.svg b/apps/desktop-tauri/src/components/providers/icons/ProviderIcon-traecn.svg new file mode 100644 index 00000000000..84128382709 --- /dev/null +++ b/apps/desktop-tauri/src/components/providers/icons/ProviderIcon-traecn.svg @@ -0,0 +1,3 @@ + + + diff --git a/apps/desktop-tauri/src/components/providers/providerIcons.ts b/apps/desktop-tauri/src/components/providers/providerIcons.ts index 83eba181051..855bee90c00 100644 --- a/apps/desktop-tauri/src/components/providers/providerIcons.ts +++ b/apps/desktop-tauri/src/components/providers/providerIcons.ts @@ -13,6 +13,7 @@ import augment from "./icons/ProviderIcon-augment.svg?raw"; import bedrock from "./icons/ProviderIcon-bedrock.svg?raw"; import claude from "./icons/ProviderIcon-claude.svg?raw"; import claudeplatform from "./icons/ProviderIcon-claudeplatform.svg?raw"; +import traecn from "./icons/ProviderIcon-traecn.svg?raw"; import codebuff from "./icons/ProviderIcon-codebuff.svg?raw"; import coderabbit from "./icons/ProviderIcon-coderabbit.svg?raw"; import codex from "./icons/ProviderIcon-codex.svg?raw"; @@ -130,6 +131,7 @@ const RAW: Record = { bedrock: tint(bedrock), claude: tint(claude), claudeplatform: tint(claudeplatform), + traecn: tint(traecn), codebuff: tint(codebuff), coderabbit: tint(coderabbit), codex: tint(codex), @@ -227,6 +229,7 @@ export const PROVIDER_ICON_REGISTRY: Record = { augment: { id: "augment", brandColor: "#1aa049", fallbackLetter: "A", svgPath: RAW.augment }, claude: { id: "claude", brandColor: "#cc7c5e", fallbackLetter: "◈", svgPath: RAW.claude }, claudeplatform: { id: "claudeplatform", brandColor: "#cc7c5e", fallbackLetter: "◈", svgPath: RAW.claudeplatform }, + traecn: { id: "traecn", brandColor: "#4D7CFE", fallbackLetter: "T", svgPath: RAW.traecn }, pi: { id: "pi", brandColor: "#7c3aed", fallbackLetter: "P" }, codebuff: { id: "codebuff", brandColor: "#00ff95", fallbackLetter: "B", svgPath: RAW.codebuff }, coderabbit: { id: "coderabbit", brandColor: "#ff5c35", fallbackLetter: "C", svgPath: RAW.coderabbit }, diff --git a/apps/desktop-tauri/src/test/providerCatalog.ts b/apps/desktop-tauri/src/test/providerCatalog.ts index a95586ca598..789317e17db 100644 --- a/apps/desktop-tauri/src/test/providerCatalog.ts +++ b/apps/desktop-tauri/src/test/providerCatalog.ts @@ -108,4 +108,5 @@ export const TEST_PROVIDER_CATALOG: Array<[string, string]> = [ ["futureagi", "Future AGI"], ["munito", "Munito"], ["claudeplatform", "Claude Platform"], + ["traecn", "Trae CN"], ]; diff --git a/docs/PROVIDERS.md b/docs/PROVIDERS.md index d55efe6ac07..8378b150887 100644 --- a/docs/PROVIDERS.md +++ b/docs/PROVIDERS.md @@ -764,6 +764,39 @@ GET https://platform.claude.com/api/organizations//current_spend showing $0. 429 reports rate limiting and 5xx reports the Console as unavailable. Response bodies are never echoed in errors. +### Trae CN + +Trae CN is disabled by default and shows the credit allowance of a Trae CN +(trae.cn) account. Sign in to trae.cn in your browser, and CodexBar uses the +`api.trae.cn` session cookies. + +- **Automatic:** CodexBar tries each detected browser holding an api.trae.cn + session. If Trae rejects a session (401, 403 or a redirect), it moves to the + next browser. No session means sign in to trae.cn. +- **Manual:** Paste the Cookie header from an api.trae.cn request. Manual mode + does not fall back to browser sessions. +- **Off:** Cookies are disabled. + +The provider makes two read-only POST requests per session and never makes a +billable call: + +```text +POST https://api.trae.cn/cloudide/api/v3/common/GetUserToken +POST https://api.trae.cn/trae/api/v2/pay/user_current_entitlement_list +``` + +- The token from `GetUserToken` lives only for one refresh and is sent only to + api.trae.cn as a `Cloud-IDE-JWT` authorization header. +- The gauge is consumed over total credits. A zero allowance shows the details + without a gauge. Used, Total and Remaining appear in the **Credits** section. +- The plan name comes from the active plan entitlement and is shown as the + login method. Pack expiry dates do not set a reset: the provider infers no + renewal or reset date. +- Malformed responses (no `is_credits_billing`, non-numeric or negative amounts, + or no token) fail parsing instead of showing zero usage. 429 reports rate + limiting, 408 and 5xx report the service as unavailable, and response bodies + are never echoed in errors. + ### Aerostack Aerostack is disabled by default. In Settings → Providers → Aerostack, enter an diff --git a/rust/src/core/provider/id.rs b/rust/src/core/provider/id.rs index aa48a3144c7..b2c19065c64 100755 --- a/rust/src/core/provider/id.rs +++ b/rust/src/core/provider/id.rs @@ -130,6 +130,7 @@ pub enum ProviderId { FutureAgi, Munito, ClaudePlatform, + TraeCn, } impl ProviderId { @@ -304,7 +305,7 @@ mod tests { #[test] fn test_provider_id_all() { let all = ProviderId::all(); - assert_eq!(all.len(), 111); + assert_eq!(all.len(), 112); assert!(all.contains(&ProviderId::Claude)); assert!(all.contains(&ProviderId::Codex)); assert!(all.contains(&ProviderId::Pi)); diff --git a/rust/src/core/provider/spec.rs b/rust/src/core/provider/spec.rs index 8f119ab2b57..55f0613564e 100644 --- a/rust/src/core/provider/spec.rs +++ b/rust/src/core/provider/spec.rs @@ -148,11 +148,11 @@ const fn spec( } } -const COUNT: usize = 111; +const COUNT: usize = 112; pub(crate) static ALL: [ProviderId; COUNT] = { // Indexing SPECS by discriminant relies on rows following declaration order. - assert!(P::ClaudePlatform as usize + 1 == COUNT); + assert!(P::TraeCn as usize + 1 == COUNT); let mut ids = [P::Codex; COUNT]; let mut i = 0; while i < COUNT { @@ -778,4 +778,10 @@ pub(crate) static SPECS: [ProviderSpec; COUNT] = [ .status("https://status.claude.com/") .aliases(&["claude-platform"]) .parse_aliases(&["claude platform"]), + spec(P::TraeCn, "traecn", "Trae CN", "#4D7CFE") + .labels("Credits", "Window") + .cookie("api.trae.cn") + .dashboard("https://www.trae.cn/dashboard#usage") + .aliases(&["trae-cn"]) + .parse_aliases(&["trae cn"]), ]; diff --git a/rust/src/core/provider_factory.rs b/rust/src/core/provider_factory.rs index 6169a4cdde9..09f52ab7917 100644 --- a/rust/src/core/provider_factory.rs +++ b/rust/src/core/provider_factory.rs @@ -26,6 +26,7 @@ use crate::providers::MuseAIProvider; use crate::providers::ParallelProvider; use crate::providers::SyntheticProvider; use crate::providers::TavilyProvider; +use crate::providers::TraeCnProvider; use crate::providers::WorkBuddyProvider; use crate::providers::{ AbacusProvider, AiAndProvider, AixyProvider, AlibabaProvider, AlibabaTokenPlanProvider, @@ -165,6 +166,7 @@ pub fn instantiate(id: ProviderId) -> Box { ProviderId::FutureAgi => Box::new(FutureAgiProvider::new()), ProviderId::Munito => Box::new(MunitoProvider::new()), ProviderId::ClaudePlatform => Box::new(ClaudePlatformProvider::new()), + ProviderId::TraeCn => Box::new(TraeCnProvider::new()), } } @@ -285,7 +287,8 @@ pub fn pace_capability(id: ProviderId) -> ProviderPaceCapability { | ProviderId::Flora | ProviderId::FutureAgi | ProviderId::Munito - | ProviderId::ClaudePlatform => ProviderPaceCapability::UNSUPPORTED, + | ProviderId::ClaudePlatform + | ProviderId::TraeCn => ProviderPaceCapability::UNSUPPORTED, } } @@ -406,7 +409,8 @@ pub fn limit_reset_policy(id: ProviderId) -> ProviderLimitResetPolicy { | ProviderId::Flora | ProviderId::FutureAgi | ProviderId::Munito - | ProviderId::ClaudePlatform => ProviderLimitResetPolicy::DEFAULT, + | ProviderId::ClaudePlatform + | ProviderId::TraeCn => ProviderLimitResetPolicy::DEFAULT, } } @@ -527,7 +531,8 @@ pub(crate) fn cost_cli_extension_for(id: ProviderId) -> Option<&'static dyn Cost | ProviderId::Flora | ProviderId::FutureAgi | ProviderId::Munito - | ProviderId::ClaudePlatform => None, + | ProviderId::ClaudePlatform + | ProviderId::TraeCn => None, } } diff --git a/rust/src/core/provider_registry_snapshot.txt b/rust/src/core/provider_registry_snapshot.txt index d2be401361f..2192a24a67b 100644 --- a/rust/src/core/provider_registry_snapshot.txt +++ b/rust/src/core/provider_registry_snapshot.txt @@ -442,3 +442,7 @@ ClaudePlatform cli=claudeplatform display="Claude Platform" cookie=Some("platfor map=["claude-platform", "claudeplatform"] metadata=ProviderMetadata { id: ClaudePlatform, display_name: "Claude Platform", session_label: "API credit", weekly_label: "Purchased credit", supports_opus: false, supports_credits: false, default_enabled: false, is_primary: false, dashboard_url: Some("https://platform.claude.com/settings/billing"), status_page_url: Some("https://status.claude.com/"), tertiary_label_key: None } sources=[Auto, Web] oauth=false web=true cli=false +TraeCn cli=traecn display="Trae CN" cookie=Some("api.trae.cn") color=#4D7CFE deprecated=false blocking=None + map=["trae-cn", "traecn"] + metadata=ProviderMetadata { id: TraeCn, display_name: "Trae CN", session_label: "Credits", weekly_label: "Window", supports_opus: false, supports_credits: false, default_enabled: false, is_primary: false, dashboard_url: Some("https://www.trae.cn/dashboard#usage"), status_page_url: None, tertiary_label_key: None } + sources=[Auto, Web] oauth=false web=true cli=false diff --git a/rust/src/core/token_accounts.rs b/rust/src/core/token_accounts.rs index 75c29505800..0c5f8608268 100755 --- a/rust/src/core/token_accounts.rs +++ b/rust/src/core/token_accounts.rs @@ -386,7 +386,8 @@ impl TokenAccountSupport { | ProviderId::Flora | ProviderId::FutureAgi | ProviderId::Munito - | ProviderId::ClaudePlatform => None, + | ProviderId::ClaudePlatform + | ProviderId::TraeCn => None, } } diff --git a/rust/src/providers/mod.rs b/rust/src/providers/mod.rs index d39cd7cfc82..4e96914239e 100755 --- a/rust/src/providers/mod.rs +++ b/rust/src/providers/mod.rs @@ -109,6 +109,7 @@ pub mod tavily; #[cfg(test)] pub(crate) mod test_support; pub mod tinyapi; +pub mod traecn; pub mod typesafe; pub mod v0; pub mod venice; @@ -222,6 +223,7 @@ pub use synthetic::SyntheticProvider; pub use t3chat::T3ChatProvider; pub use tavily::TavilyProvider; pub use tinyapi::TinyApiProvider; +pub use traecn::TraeCnProvider; pub use typesafe::TypeSafeProvider; pub use v0::V0Provider; pub use venice::VeniceProvider; diff --git a/rust/src/providers/traecn/mod.rs b/rust/src/providers/traecn/mod.rs new file mode 100644 index 00000000000..077d4808eed --- /dev/null +++ b/rust/src/providers/traecn/mod.rs @@ -0,0 +1,399 @@ +//! Trae CN credits (trae.cn). +//! +//! Ported from upstream CodexBar (`Resources/Plugins/traecn.js`, +//! `Providers/TraeCN/TraeCNProviderDescriptor.swift`, #4682). The signed-in +//! `api.trae.cn` browser session exchanges its cookies for a token at +//! `GetUserToken`, and that token reads the account's current entitlement +//! list. The token lives only for one refresh and is sent only back to the +//! same API origin. Refresh makes no billable calls, and it infers no reset +//! or renewal date from pack timestamps. Response bodies are never echoed in +//! errors and cookies are never logged. + +#[cfg(test)] +mod tests; + +use std::time::Duration; + +use async_trait::async_trait; +use reqwest::{ + Client, StatusCode, + header::{ACCEPT, AUTHORIZATION, COOKIE, ORIGIN, REFERER}, +}; +use serde_json::{Value, json}; + +use crate::browser::session_scan::{BrowserCookieStore, SystemCookieStore}; +use crate::core::{ + FetchContext, ManualEmptyCookiePolicy, Provider, ProviderDisplayDetail, ProviderError, + ProviderFetchResult, ProviderId, ProviderMetadata, ProviderStateKind, RateWindow, SourceMode, + UsageSnapshot, +}; +use crate::providers::{BoundedBodyError, format, read_bounded_response}; + +const BASE_URL: &str = "https://api.trae.cn"; +const COOKIE_DOMAIN: &str = "api.trae.cn"; +const WEB_ORIGIN: &str = "https://www.trae.cn"; +const WEB_REFERER: &str = "https://www.trae.cn/"; +const SIGN_IN_URL: &str = "https://www.trae.cn/"; +const DASHBOARD_URL: &str = "https://www.trae.cn/dashboard#usage"; +const TOKEN_PATH: &str = "/cloudide/api/v3/common/GetUserToken"; +const ENTITLEMENTS_PATH: &str = "/trae/api/v2/pay/user_current_entitlement_list"; +/// Upstream's fixed request timeout. +const REQUEST_TIMEOUT: Duration = Duration::from_secs(15); +const MAX_RESPONSE_BYTES: usize = 1024 * 1024; +/// Upstream's application code for an expired or invalid login. +const SESSION_EXPIRED_CODE: i64 = 1001; +/// The entitlement group that carries the subscription plan name. +const PLAN_GROUP: i64 = 0; +const ACTIVE_PACK: i64 = 1; +const CREDITS_SECTION: &str = "Credits"; + +const SESSION_EXPIRED: &str = + "Trae CN session expired. Sign in at trae.cn or paste a fresh Cookie header."; +const MISSING_SESSION: &str = + "Sign in at trae.cn in a supported browser, or paste its Cookie header."; +const COOKIES_OFF: &str = "Trae CN cookies are disabled."; +const RATE_LIMITED: &str = "Trae CN is rate limited. Try refreshing later."; +const UNAVAILABLE: &str = "Trae CN is temporarily unavailable. Try refreshing later."; +const MALFORMED: &str = "Trae CN usage response format changed. Refresh to try again, or update CodexBar or report the issue."; +const NO_ALLOWANCE: &str = "No Trae CN credit allowance"; + +static TRAE_CN_METADATA: ProviderMetadata = ProviderMetadata { + id: ProviderId::TraeCn, + display_name: "Trae CN", + session_label: "Credits", + weekly_label: "Window", + supports_opus: false, + supports_credits: false, + default_enabled: false, + is_primary: false, + dashboard_url: Some(DASHBOARD_URL), + status_page_url: None, + tertiary_label_key: None, +}; + +pub struct TraeCnProvider { + client: Option, + base_url: String, +} + +/// Why one browser session did not produce usage. +enum SessionError { + /// Trae rejected the session; the next browser session may still work. + Rejected, + Failed(ProviderError), +} + +impl From for SessionError { + fn from(error: ProviderError) -> Self { + SessionError::Failed(error) + } +} + +/// The subset of the entitlement response that the card shows. +#[derive(Debug, PartialEq)] +struct Entitlements { + used: f64, + total: f64, + plan: Option, +} + +impl TraeCnProvider { + pub fn new() -> Self { + Self::with_base_url(BASE_URL) + } + + fn with_base_url(base_url: &str) -> Self { + let client = crate::core::credentialed_http_client_builder() + .cookie_store(false) + .redirect(reqwest::redirect::Policy::none()) + .timeout(REQUEST_TIMEOUT) + .build() + .ok(); + Self { + client, + base_url: base_url.trim_end_matches('/').to_owned(), + } + } + + /// A pasted header is the only session tried. Otherwise each browser + /// holding a session for api.trae.cn is tried in browser order, moving on + /// only after Trae rejects a session (upstream `rejectCookie`). + async fn fetch_web( + &self, + ctx: &FetchContext, + store: S, + ) -> Result { + let client = self.client.as_ref().ok_or_else(|| { + ProviderError::Other("Could not create a secure Trae CN HTTP client.".into()) + })?; + let timeout = Duration::from_secs(ctx.web_timeout.max(1)).min(REQUEST_TIMEOUT); + if let Some(raw) = ctx.manual_cookie_header.as_deref() { + let cookie = + crate::providers::normalize_cookie_header(raw).ok_or_else(missing_session)?; + return match self.fetch_session(client, &cookie, timeout).await { + Ok(result) => Ok(result), + Err(SessionError::Rejected) => Err(session_expired()), + Err(SessionError::Failed(error)) => Err(error), + }; + } + if ctx.manual_cookie_missing { + return Err(missing_session()); + } + + let mut sessions = + crate::providers::browser_cookie_sessions_from(store, ctx, &[COOKIE_DOMAIN]); + let mut rejected = false; + for header in sessions.by_ref() { + let Some(cookie) = crate::providers::normalize_cookie_header(&header) else { + continue; + }; + match self.fetch_session(client, &cookie, timeout).await { + Ok(result) => return Ok(result), + Err(SessionError::Rejected) => rejected = true, + Err(SessionError::Failed(error)) => return Err(error), + } + } + if rejected { + return Err(session_expired()); + } + match sessions.exhaustion_error() { + ProviderError::NoCookies => Err(missing_session()), + error => { + tracing::debug!(%error, "Trae CN browser sessions are unavailable"); + Err(error) + } + } + } + + async fn fetch_session( + &self, + client: &Client, + cookie: &str, + timeout: Duration, + ) -> Result { + let token_payload = self + .post_json(client, TOKEN_PATH, cookie, &json!({}), None, timeout) + .await?; + let token = token_payload + .get("Result") + .and_then(|result| result.get("Token")) + .and_then(Value::as_str) + .map(str::trim) + .filter(|token| !token.is_empty()) + .ok_or_else(malformed)?; + + let payload = self + .post_json( + client, + ENTITLEMENTS_PATH, + cookie, + &json!({ "require_usage": true, "full_data": true }), + Some(token), + timeout, + ) + .await?; + if payload.get("code").and_then(Value::as_i64) == Some(SESSION_EXPIRED_CODE) { + return Err(SessionError::Rejected); + } + let entitlements = parse_entitlements(&payload)?; + Ok(build_result(&entitlements)) + } + + /// One JSON POST. A 401, 403 or redirect means Trae rejected the session. + async fn post_json( + &self, + client: &Client, + path: &str, + cookie: &str, + body: &Value, + token: Option<&str>, + timeout: Duration, + ) -> Result { + let mut request = client + .post(format!("{}{path}", self.base_url)) + .header(ACCEPT, "application/json") + .header(COOKIE, cookie) + .header(ORIGIN, WEB_ORIGIN) + .header(REFERER, WEB_REFERER) + .json(body) + .timeout(timeout); + if let Some(token) = token { + request = request.header(AUTHORIZATION, format!("Cloud-IDE-JWT {token}")); + } + let response = request + .send() + .await + .map_err(|error| SessionError::Failed(ProviderError::Network(error)))?; + let status = response.status(); + if matches!(status, StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN) + || status.is_redirection() + { + return Err(SessionError::Rejected); + } + match status { + StatusCode::OK => {} + StatusCode::TOO_MANY_REQUESTS => { + return Err(ProviderError::Other(RATE_LIMITED.into()).into()); + } + StatusCode::REQUEST_TIMEOUT => { + return Err(ProviderError::Other(UNAVAILABLE.into()).into()); + } + status if status.is_server_error() => { + return Err(ProviderError::Other(UNAVAILABLE.into()).into()); + } + status => { + return Err(ProviderError::Other(format!( + "Trae CN request failed (HTTP {}). Try refreshing later.", + status.as_u16() + )) + .into()); + } + } + let bytes = read_bounded_response(response, MAX_RESPONSE_BYTES) + .await + .map_err(|error| { + SessionError::Failed(match error { + BoundedBodyError::Read(error) => ProviderError::Network(error), + BoundedBodyError::TooLarge => malformed(), + }) + })?; + serde_json::from_slice(&bytes).map_err(|_| malformed().into()) + } +} + +impl Default for TraeCnProvider { + fn default() -> Self { + Self::new() + } +} + +#[async_trait] +impl Provider for TraeCnProvider { + fn id(&self) -> ProviderId { + ProviderId::TraeCn + } + + fn metadata(&self) -> &ProviderMetadata { + &TRAE_CN_METADATA + } + + async fn fetch_usage(&self, ctx: &FetchContext) -> Result { + match ctx.source_mode { + SourceMode::Auto | SourceMode::Web => self.fetch_web(ctx, SystemCookieStore).await, + // The shell maps an Off cookie source to Cli for web-only + // providers; upstream reports that cookies are disabled. + SourceMode::Cli => Err(ProviderError::Other(COOKIES_OFF.into())), + SourceMode::OAuth => Err(ProviderError::UnsupportedSource(SourceMode::OAuth)), + } + } + + fn available_sources(&self) -> Vec { + vec![SourceMode::Auto, SourceMode::Web] + } + + fn supports_web(&self) -> bool { + true + } + + /// Manual mode sends only the pasted header, so an empty manual source + /// must not fall back to a browser session. + fn manual_empty_cookie_policy(&self) -> ManualEmptyCookiePolicy { + ManualEmptyCookiePolicy::FailClosedWeb + } + + /// The provider walks browsers itself so a rejected session can move on + /// to the next browser. + fn owns_browser_cookie_resolution(&self) -> bool { + true + } + + fn error_state_kind(&self, error: &ProviderError) -> ProviderStateKind { + match error { + ProviderError::Other(message) if message == SESSION_EXPIRED => { + ProviderStateKind::ExpiredSession + } + _ => error.state_kind(), + } + } +} + +fn session_expired() -> ProviderError { + ProviderError::Other(SESSION_EXPIRED.into()) +} + +fn missing_session() -> ProviderError { + ProviderError::BrowserSignInRequired { + message: MISSING_SESSION.into(), + sign_in_url: SIGN_IN_URL.into(), + } +} + +fn malformed() -> ProviderError { + ProviderError::Parse(MALFORMED.into()) +} + +/// A non-negative JSON number. Numeric strings are rejected, as upstream does. +fn amount(value: Option<&Value>) -> Result { + value + .and_then(Value::as_f64) + .filter(|amount| amount.is_finite() && *amount >= 0.0) + .ok_or_else(malformed) +} + +/// Upstream `group_type === 0` with `status` absent or `1`. +fn is_plan_pack(pack: &Value) -> bool { + pack.get("group_type").and_then(Value::as_i64) == Some(PLAN_GROUP) + && pack + .get("status") + .is_none_or(|status| status.as_i64() == Some(ACTIVE_PACK)) +} + +fn parse_entitlements(payload: &Value) -> Result { + if payload.get("is_credits_billing").and_then(Value::as_bool) != Some(true) { + return Err(malformed()); + } + let summary = payload + .get("usage_summary") + .filter(|summary| summary.is_object()) + .ok_or_else(malformed)?; + let total = amount(summary.get("total_amount"))?; + let used = amount(summary.get("consumed_amount"))?; + let plan = payload + .get("user_entitlement_pack_list") + .and_then(Value::as_array) + .and_then(|packs| packs.iter().find(|pack| is_plan_pack(pack))) + .and_then(|pack| pack.get("display_desc")) + .and_then(Value::as_str) + .map(str::trim) + .filter(|plan| !plan.is_empty()) + .map(str::to_owned); + Ok(Entitlements { used, total, plan }) +} + +/// Upstream returns no reset and no window. A zero allowance keeps its +/// details but gets no gauge. +fn build_result(entitlements: &Entitlements) -> ProviderFetchResult { + let Entitlements { used, total, plan } = entitlements; + let primary = if *total > 0.0 { + RateWindow::new((used / total * 100.0).min(100.0)) + } else { + RateWindow::informational(NO_ALLOWANCE).hidden_on_card() + }; + let mut usage = UsageSnapshot::new(primary); + if let Some(plan) = plan { + usage = usage.with_login_method(plan.clone()); + } + + let remaining = (total - used).max(0.0); + let rows = [ + ("credits-used", "Used", *used), + ("credits-total", "Total", *total), + ("credits-remaining", "Remaining", remaining), + ] + .into_iter() + .filter_map(|(id, title, value)| { + ProviderDisplayDetail::new(id, title, format::number(value, 2)) + .and_then(|row| row.with_section_title(CREDITS_SECTION)) + }); + ProviderFetchResult::new(usage, "web").with_display_details(rows) +} diff --git a/rust/src/providers/traecn/tests.rs b/rust/src/providers/traecn/tests.rs new file mode 100644 index 00000000000..7c7dacbcf3f --- /dev/null +++ b/rust/src/providers/traecn/tests.rs @@ -0,0 +1,510 @@ +//! Fixture tests ported from upstream `traecn.js` behavior. Every fetch runs +//! against a local mock of the Trae API with a pasted Cookie header or an +//! in-memory browser store, so no test reads a real browser or account. + +use super::*; +use crate::browser::detection::BrowserType; +use crate::browser::session_scan::test_support::FakeCookieStore; +use mockito::{Matcher, Mock, Server, ServerGuard}; + +const COOKIE: &str = "trae_session=synthetic-1"; +const OTHER_COOKIE: &str = "trae_session=synthetic-2"; +const TOKEN: &str = "synthetic-token"; +const TOKEN_BODY: &str = r#"{"Result":{"Token":"synthetic-token"}}"#; +const ENTITLEMENTS: &str = r#"{"is_credits_billing":true, + "usage_summary":{"total_amount":10000,"consumed_amount":2000}, + "user_entitlement_pack_list":[ + {"group_type":4,"status":1,"display_desc":"Synthetic Add-on"}, + {"group_type":0,"status":1,"display_desc":" Fixture Pro "}]}"#; +const SECRET_BODY: &str = "synthetic-secret-echo"; + +fn manual(cookie: &str) -> FetchContext { + FetchContext { + source_mode: SourceMode::Auto, + manual_cookie_header: Some(cookie.to_owned()), + ..FetchContext::default() + } +} + +/// Token exchange then entitlement read, both bound to one Cookie header. +async fn exchange( + server: &mut ServerGuard, + cookie: &str, + token: (usize, &str), + usage: (usize, &str), +) -> Vec { + vec![ + server + .mock("POST", TOKEN_PATH) + .match_header("cookie", cookie) + .match_header("origin", WEB_ORIGIN) + .match_body(Matcher::Json(json!({}))) + .with_status(token.0) + .with_header("content-type", "application/json") + .with_body(token.1) + .create_async() + .await, + server + .mock("POST", ENTITLEMENTS_PATH) + .match_header("cookie", cookie) + .match_header("authorization", format!("Cloud-IDE-JWT {TOKEN}").as_str()) + .match_body(Matcher::Json( + json!({"require_usage": true, "full_data": true}), + )) + .with_status(usage.0) + .with_header("content-type", "application/json") + .with_body(usage.1) + .create_async() + .await, + ] +} + +async fn fetch( + server: &ServerGuard, + ctx: &FetchContext, + store: &FakeCookieStore, +) -> Result { + TraeCnProvider::with_base_url(&server.url()) + .fetch_web(ctx, store) + .await +} + +fn rows_in<'a>(result: &'a ProviderFetchResult, section: &str) -> Vec<(&'a str, &'a str)> { + result + .display_details + .iter() + .filter(|detail| detail.section_title() == Some(section)) + .map(|detail| (detail.title(), detail.value())) + .collect() +} + +fn assert_malformed(result: Result) { + assert!( + matches!(&result, Err(ProviderError::Parse(message)) if message == MALFORMED), + "expected malformed response, got {result:?}" + ); +} + +#[test] +fn metadata_and_source_contract_match_upstream_descriptor() { + let provider = TraeCnProvider::new(); + assert_eq!(provider.id(), ProviderId::TraeCn); + let metadata = provider.metadata(); + assert_eq!(metadata.display_name, "Trae CN"); + assert_eq!(metadata.session_label, "Credits"); + assert_eq!(metadata.weekly_label, "Window"); + assert_eq!( + metadata.dashboard_url, + Some("https://www.trae.cn/dashboard#usage") + ); + assert!(!metadata.default_enabled); + assert_eq!( + provider.available_sources(), + vec![SourceMode::Auto, SourceMode::Web] + ); + assert!(provider.supports_web()); + assert!(provider.owns_browser_cookie_resolution()); + assert!(matches!( + provider.manual_empty_cookie_policy(), + ManualEmptyCookiePolicy::FailClosedWeb + )); + assert_eq!(ProviderId::TraeCn.cookie_domain(), Some("api.trae.cn")); + assert_eq!( + ProviderId::from_cli_name("traecn"), + Some(ProviderId::TraeCn) + ); + assert_eq!( + ProviderId::from_cli_name("trae-cn"), + Some(ProviderId::TraeCn) + ); +} + +#[tokio::test] +async fn cli_and_oauth_sources_never_reach_the_network() { + let provider = TraeCnProvider::new(); + let cli = FetchContext { + source_mode: SourceMode::Cli, + ..FetchContext::default() + }; + assert!(matches!( + provider.fetch_usage(&cli).await, + Err(ProviderError::Other(message)) if message == COOKIES_OFF + )); + let oauth = FetchContext { + source_mode: SourceMode::OAuth, + ..FetchContext::default() + }; + assert!(matches!( + provider.fetch_usage(&oauth).await, + Err(ProviderError::UnsupportedSource(SourceMode::OAuth)) + )); +} + +#[test] +fn plan_is_the_active_group_zero_pack_and_is_trimmed() { + let payload: Value = serde_json::from_str( + r#"{"is_credits_billing":true, + "usage_summary":{"total_amount":10,"consumed_amount":1}, + "user_entitlement_pack_list":[ + {"group_type":0,"status":0,"display_desc":"Lapsed"}, + {"group_type":4,"status":1,"display_desc":"Add-on"}, + {"group_type":0,"display_desc":" Fixture Pro "}]}"#, + ) + .unwrap(); + let entitlements = parse_entitlements(&payload).unwrap(); + assert_eq!(entitlements.plan.as_deref(), Some("Fixture Pro")); + assert_eq!(entitlements.total, 10.0); + assert_eq!(entitlements.used, 1.0); + + let blank: Value = serde_json::from_str( + r#"{"is_credits_billing":true,"usage_summary":{"total_amount":10,"consumed_amount":1}, + "user_entitlement_pack_list":[{"group_type":0,"status":1,"display_desc":" "}]}"#, + ) + .unwrap(); + assert_eq!(parse_entitlements(&blank).unwrap().plan, None); +} + +#[tokio::test] +async fn entitlement_usage_fills_the_gauge_and_credits_details() { + let mut server = Server::new_async().await; + let mocks = exchange(&mut server, COOKIE, (200, TOKEN_BODY), (200, ENTITLEMENTS)).await; + let store = FakeCookieStore::default(); + let result = fetch(&server, &manual(COOKIE), &store).await.unwrap(); + for mock in mocks { + mock.assert_async().await; + } + + assert_eq!(result.usage.primary.used_percent, 20.0); + assert_eq!(result.usage.primary.reset_description, None); + assert_eq!(result.usage.login_method.as_deref(), Some("Fixture Pro")); + assert!(result.cost.is_none()); + assert_eq!( + rows_in(&result, CREDITS_SECTION), + vec![ + ("Used", "2,000"), + ("Total", "10,000"), + ("Remaining", "8,000"), + ] + ); + // The manual header is the only session, so no browser is read. + assert!(store.reads().is_empty()); +} + +#[tokio::test] +async fn zero_allowance_keeps_details_without_a_gauge() { + let mut server = Server::new_async().await; + let _mocks = exchange( + &mut server, + COOKIE, + (200, TOKEN_BODY), + ( + 200, + r#"{"is_credits_billing":true,"usage_summary":{"total_amount":0,"consumed_amount":0},"user_entitlement_pack_list":[]}"#, + ), + ) + .await; + let result = fetch(&server, &manual(COOKIE), &FakeCookieStore::default()) + .await + .unwrap(); + assert_eq!( + result.usage.primary.reset_description.as_deref(), + Some(NO_ALLOWANCE) + ); + assert_eq!(result.usage.login_method, None); + assert_eq!( + rows_in(&result, CREDITS_SECTION), + vec![("Used", "0"), ("Total", "0"), ("Remaining", "0")] + ); +} + +#[tokio::test] +async fn overspend_clamps_the_gauge_and_keeps_the_true_numbers() { + let mut server = Server::new_async().await; + let _mocks = exchange( + &mut server, + COOKIE, + (200, TOKEN_BODY), + ( + 200, + r#"{"is_credits_billing":true,"usage_summary":{"total_amount":10000,"consumed_amount":12500}}"#, + ), + ) + .await; + let result = fetch(&server, &manual(COOKIE), &FakeCookieStore::default()) + .await + .unwrap(); + assert_eq!(result.usage.primary.used_percent, 100.0); + assert_eq!( + rows_in(&result, CREDITS_SECTION), + vec![("Used", "12,500"), ("Total", "10,000"), ("Remaining", "0")] + ); +} + +#[test] +fn malformed_entitlements_fail_closed() { + for body in [ + r#"{"is_credits_billing":false,"usage_summary":{"total_amount":10,"consumed_amount":1}}"#, + r#"{"is_credits_billing":"true","usage_summary":{"total_amount":10,"consumed_amount":1}}"#, + r#"{"is_credits_billing":true,"usage_summary":{"total_amount":"10000","consumed_amount":2000}}"#, + r#"{"is_credits_billing":true,"usage_summary":{"total_amount":10000,"consumed_amount":-1}}"#, + r#"{"is_credits_billing":true,"usage_summary":{"total_amount":10000}}"#, + r#"{"is_credits_billing":true}"#, + r#"[]"#, + ] { + let payload: Value = serde_json::from_str(body).unwrap(); + assert!( + matches!(parse_entitlements(&payload), Err(ProviderError::Parse(message)) if message == MALFORMED), + "accepted {body}" + ); + } +} + +#[tokio::test] +async fn malformed_bodies_are_rejected_not_reported_as_zero_usage() { + let cases: [(&str, &str); 4] = [ + (TOKEN_BODY, "not-json"), + ( + TOKEN_BODY, + r#"{"is_credits_billing":false,"usage_summary":{"total_amount":10,"consumed_amount":1}}"#, + ), + (r#"{"Result":{}}"#, ENTITLEMENTS), + (r#"{"Result":{"Token":" "}}"#, ENTITLEMENTS), + ]; + for (token_body, usage_body) in cases { + let mut server = Server::new_async().await; + let _mocks = exchange(&mut server, COOKIE, (200, token_body), (200, usage_body)).await; + let result = fetch(&server, &manual(COOKIE), &FakeCookieStore::default()).await; + assert_malformed(result); + } +} + +#[tokio::test] +async fn rejected_token_exchange_reports_an_expired_session() { + for status in [401, 403, 302] { + let mut server = Server::new_async().await; + let _token = server + .mock("POST", TOKEN_PATH) + .match_header("cookie", COOKIE) + .with_status(status) + .with_body("{}") + .create_async() + .await; + let provider = TraeCnProvider::with_base_url(&server.url()); + let error = provider + .fetch_web(&manual(COOKIE), &FakeCookieStore::default()) + .await + .unwrap_err(); + assert!(matches!(&error, ProviderError::Other(message) if message == SESSION_EXPIRED)); + assert_eq!( + provider.error_state_kind(&error), + ProviderStateKind::ExpiredSession + ); + } +} + +#[tokio::test] +async fn rejected_entitlement_read_reports_an_expired_session() { + for (status, body) in [(403, "{}"), (200, r#"{"code":1001,"msg":"synthetic"}"#)] { + let mut server = Server::new_async().await; + let _mocks = exchange(&mut server, COOKIE, (200, TOKEN_BODY), (status, body)).await; + let result = fetch(&server, &manual(COOKIE), &FakeCookieStore::default()).await; + assert!( + matches!(&result, Err(ProviderError::Other(message)) if message == SESSION_EXPIRED), + "expected expired session, got {result:?}" + ); + } +} + +#[tokio::test] +async fn a_rejected_profile_advances_to_the_next_browser_session() { + let mut server = Server::new_async().await; + let _first = exchange(&mut server, COOKIE, (401, "{}"), (200, ENTITLEMENTS)).await; + let _second = exchange( + &mut server, + OTHER_COOKIE, + (200, TOKEN_BODY), + (200, ENTITLEMENTS), + ) + .await; + let store = FakeCookieStore::default() + .cookies( + BrowserType::Chrome, + COOKIE_DOMAIN, + &[("trae_session", "synthetic-1")], + ) + .cookies( + BrowserType::Edge, + COOKIE_DOMAIN, + &[("trae_session", "synthetic-2")], + ); + let result = fetch(&server, &FetchContext::default(), &store) + .await + .unwrap(); + assert_eq!(result.usage.primary.used_percent, 20.0); + assert_eq!(store.reads().len(), 2); +} + +#[tokio::test] +async fn all_rejected_browser_sessions_report_an_expired_session() { + let mut server = Server::new_async().await; + let _first = exchange(&mut server, COOKIE, (401, "{}"), (200, ENTITLEMENTS)).await; + let _second = exchange(&mut server, OTHER_COOKIE, (403, "{}"), (200, ENTITLEMENTS)).await; + let store = FakeCookieStore::default() + .cookies( + BrowserType::Chrome, + COOKIE_DOMAIN, + &[("trae_session", "synthetic-1")], + ) + .cookies( + BrowserType::Edge, + COOKIE_DOMAIN, + &[("trae_session", "synthetic-2")], + ); + let provider = TraeCnProvider::with_base_url(&server.url()); + let error = provider + .fetch_web(&FetchContext::default(), &store) + .await + .unwrap_err(); + assert!(matches!(&error, ProviderError::Other(message) if message == SESSION_EXPIRED)); + assert_eq!( + provider.error_state_kind(&error), + ProviderStateKind::ExpiredSession + ); + assert_eq!(store.reads().len(), 2); +} + +#[tokio::test] +async fn transient_failures_keep_the_session_and_stop_the_walk() { + for (status, expected) in [(429, RATE_LIMITED), (503, UNAVAILABLE), (408, UNAVAILABLE)] { + let mut server = Server::new_async().await; + let _token = server + .mock("POST", TOKEN_PATH) + .match_header("cookie", COOKIE) + .with_status(status) + .with_body("{}") + .create_async() + .await; + // A second browser would only be tried after a rejected session. + let store = FakeCookieStore::default() + .cookies( + BrowserType::Chrome, + COOKIE_DOMAIN, + &[("trae_session", "synthetic-1")], + ) + .cookies( + BrowserType::Edge, + COOKIE_DOMAIN, + &[("trae_session", "synthetic-2")], + ); + let result = TraeCnProvider::with_base_url(&server.url()) + .fetch_web(&FetchContext::default(), &store) + .await; + assert!( + matches!(&result, Err(ProviderError::Other(message)) if message == expected), + "status {status} gave {result:?}" + ); + assert_eq!(store.reads().len(), 1); + } +} + +#[tokio::test] +async fn other_http_failures_name_the_status_without_echoing_the_body() { + let mut server = Server::new_async().await; + let _token = server + .mock("POST", TOKEN_PATH) + .match_header("cookie", COOKIE) + .with_status(404) + .with_body(SECRET_BODY) + .create_async() + .await; + let result = fetch(&server, &manual(COOKIE), &FakeCookieStore::default()).await; + match result { + Err(ProviderError::Other(message)) => { + assert_eq!( + message, + "Trae CN request failed (HTTP 404). Try refreshing later." + ); + assert!(!message.contains(SECRET_BODY)); + } + other => panic!("expected Other, got {other:?}"), + } +} + +#[tokio::test] +async fn pasted_header_is_the_only_session_and_is_normalized() { + let mut server = Server::new_async().await; + let mocks = exchange(&mut server, COOKIE, (200, TOKEN_BODY), (200, ENTITLEMENTS)).await; + let store = FakeCookieStore::default().cookies( + BrowserType::Chrome, + COOKIE_DOMAIN, + &[("trae_session", "synthetic-2")], + ); + let result = fetch(&server, &manual(&format!("Cookie: {COOKIE}")), &store) + .await + .unwrap(); + for mock in mocks { + mock.assert_async().await; + } + assert_eq!(result.usage.primary.used_percent, 20.0); + assert!(store.reads().is_empty()); +} + +#[tokio::test] +async fn no_browser_session_asks_the_user_to_sign_in() { + let server = Server::new_async().await; + let result = fetch( + &server, + &FetchContext::default(), + &FakeCookieStore::default(), + ) + .await; + assert!(matches!( + result, + Err(ProviderError::BrowserSignInRequired { sign_in_url, .. }) if sign_in_url == SIGN_IN_URL + )); +} + +#[tokio::test] +async fn an_empty_manual_source_never_falls_back_to_a_browser() { + let server = Server::new_async().await; + let store = FakeCookieStore::default().cookies( + BrowserType::Chrome, + COOKIE_DOMAIN, + &[("trae_session", "synthetic-1")], + ); + let ctx = FetchContext { + source_mode: SourceMode::Auto, + manual_cookie_missing: true, + ..FetchContext::default() + }; + let result = fetch(&server, &ctx, &store).await; + assert!(matches!( + result, + Err(ProviderError::BrowserSignInRequired { .. }) + )); + assert!(store.reads().is_empty()); +} + +#[test] +fn amounts_must_be_non_negative_json_numbers() { + assert_eq!(amount(Some(&json!(0))).unwrap(), 0.0); + assert_eq!(amount(Some(&json!(12.5))).unwrap(), 12.5); + for value in [json!("1"), json!(-0.5), json!(null), json!([])] { + assert!(amount(Some(&value)).is_err(), "accepted {value}"); + } + assert!(amount(None).is_err()); +} + +#[test] +fn zero_allowance_is_informational_and_hidden_from_the_card() { + let result = build_result(&Entitlements { + used: 0.0, + total: 0.0, + plan: None, + }); + assert!(result.usage.primary.is_informational); + assert_eq!( + result.usage.primary.reset_description.as_deref(), + Some(NO_ALLOWANCE) + ); +} diff --git a/rust/src/settings/provider_accessors.rs b/rust/src/settings/provider_accessors.rs index 5c9f7074f90..da3ebdc397a 100644 --- a/rust/src/settings/provider_accessors.rs +++ b/rust/src/settings/provider_accessors.rs @@ -41,7 +41,8 @@ impl Settings { | ProviderId::TinyApi | ProviderId::Xapi | ProviderId::Parallel - | ProviderId::ClaudePlatform => "auto", + | ProviderId::ClaudePlatform + | ProviderId::TraeCn => "auto", _ => DEFAULT_COOKIE_SOURCE, }) } From 21bc867a4bb5ca29ffa224d7c6e43da258275d73 Mon Sep 17 00:00:00 2001 From: Factory Agent Date: Fri, 23 Oct 2026 17:14:31 +0700 Subject: [PATCH 2/2] Add TRAE_CN_BASE_URL override for proof runs --- docs/PROVIDERS.md | 4 ++ rust/src/providers/traecn/mod.rs | 84 +++++++++++++++++++++++++++++- rust/src/providers/traecn/tests.rs | 54 +++++++++++++++++++ 3 files changed, 141 insertions(+), 1 deletion(-) diff --git a/docs/PROVIDERS.md b/docs/PROVIDERS.md index 8378b150887..d3dad7a0b79 100644 --- a/docs/PROVIDERS.md +++ b/docs/PROVIDERS.md @@ -796,6 +796,10 @@ POST https://api.trae.cn/trae/api/v2/pay/user_current_entitlement_list or no token) fail parsing instead of showing zero usage. 429 reports rate limiting, 408 and 5xx report the service as unavailable, and response bodies are never echoed in errors. +- For proof runs and fake servers, `TRAE_CN_BASE_URL` replaces the API origin + (default `https://api.trae.cn`). It accepts `https://`, or `http://` only for a + loopback host (`127.0.0.1`, `localhost`, `[::1]`), with no credentials, query, + or fragment. Any other value is ignored with a warning. ### Aerostack diff --git a/rust/src/providers/traecn/mod.rs b/rust/src/providers/traecn/mod.rs index 077d4808eed..907bba546e4 100644 --- a/rust/src/providers/traecn/mod.rs +++ b/rust/src/providers/traecn/mod.rs @@ -30,6 +30,9 @@ use crate::core::{ use crate::providers::{BoundedBodyError, format, read_bounded_response}; const BASE_URL: &str = "https://api.trae.cn"; +/// Optional override for the API origin (proof runs and fake servers). +/// Accepts `https://` or loopback `http://` only; anything else is ignored. +const BASE_URL_ENV: &str = "TRAE_CN_BASE_URL"; const COOKIE_DOMAIN: &str = "api.trae.cn"; const WEB_ORIGIN: &str = "https://www.trae.cn"; const WEB_REFERER: &str = "https://www.trae.cn/"; @@ -97,9 +100,88 @@ struct Entitlements { plan: Option, } +/// The API origin: a valid `TRAE_CN_BASE_URL` override, otherwise production. +/// An invalid override is ignored with a warning. +fn base_url_from_override(raw: Option<&str>) -> String { + let Some(raw) = raw.map(clean_override).filter(|raw| !raw.is_empty()) else { + return BASE_URL.to_owned(); + }; + match validated_base_url(&raw) { + Some(base) => base, + None => { + tracing::warn!( + "Ignoring {BASE_URL_ENV}: use an https:// URL or a loopback http:// URL without credentials, query, or fragment" + ); + BASE_URL.to_owned() + } + } +} + +fn clean_override(raw: &str) -> String { + let value = raw.trim(); + let unquoted = value + .strip_prefix('"') + .and_then(|inner| inner.strip_suffix('"')) + .or_else(|| { + value + .strip_prefix('\'') + .and_then(|inner| inner.strip_suffix('\'')) + }) + .unwrap_or(value); + unquoted.trim().to_owned() +} + +/// Same policy as the Sub2API base URL: `https://`, or `http://` only for a +/// loopback host, with no userinfo, query, fragment, or encoded host delimiters. +/// Returns the URL without a trailing slash. +fn validated_base_url(raw: &str) -> Option { + let lower = raw.to_ascii_lowercase(); + if ["%2f", "%5c", "%3f", "%23", "%40", "%3a"] + .iter() + .any(|encoded| lower.contains(encoded)) + { + return None; + } + let url = reqwest::Url::parse(raw).ok()?; + let host = url.host_str()?; + let scheme_ok = match url.scheme() { + "https" => true, + "http" => is_loopback_host(host), + _ => false, + }; + if !scheme_ok + || !url.username().is_empty() + || url.password().is_some() + || url.query().is_some() + || url.fragment().is_some() + || host.contains('%') + || host.chars().any(|c| c.is_control() || c.is_whitespace()) + { + return None; + } + Some(url.as_str().trim_end_matches('/').to_owned()) +} + +fn is_loopback_host(host: &str) -> bool { + let normalized = host.trim_end_matches('.').to_ascii_lowercase(); + if normalized == "localhost" { + return true; + } + let ip_candidate = normalized + .strip_prefix('[') + .and_then(|value| value.strip_suffix(']')) + .unwrap_or(&normalized); + match ip_candidate.parse::() { + Ok(std::net::IpAddr::V4(ip)) => ip.is_loopback(), + Ok(std::net::IpAddr::V6(ip)) => ip.is_loopback(), + Err(_) => false, + } +} + impl TraeCnProvider { pub fn new() -> Self { - Self::with_base_url(BASE_URL) + let override_url = std::env::var(BASE_URL_ENV).ok(); + Self::with_base_url(&base_url_from_override(override_url.as_deref())) } fn with_base_url(base_url: &str) -> Self { diff --git a/rust/src/providers/traecn/tests.rs b/rust/src/providers/traecn/tests.rs index 7c7dacbcf3f..fe181af1ca1 100644 --- a/rust/src/providers/traecn/tests.rs +++ b/rust/src/providers/traecn/tests.rs @@ -508,3 +508,57 @@ fn zero_allowance_is_informational_and_hidden_from_the_card() { Some(NO_ALLOWANCE) ); } + +#[test] +fn base_url_override_accepts_https_and_loopback_http() { + assert_eq!( + validated_base_url("https://api.trae.cn").as_deref(), + Some("https://api.trae.cn") + ); + assert_eq!( + validated_base_url("http://127.0.0.1:4011/").as_deref(), + Some("http://127.0.0.1:4011") + ); + assert_eq!( + validated_base_url("http://localhost:9").as_deref(), + Some("http://localhost:9") + ); + assert_eq!( + validated_base_url("http://[::1]:9").as_deref(), + Some("http://[::1]:9") + ); +} + +#[test] +fn base_url_override_rejects_remote_plain_http_and_unsafe_forms() { + for raw in [ + "http://api.trae.cn", + "http://localhost.attacker.test", + "ftp://127.0.0.1", + "https://user:pass@api.trae.cn", + "https://api.trae.cn?x=1", + "https://api.trae.cn#frag", + "http://127.0.0.1%2f.attacker.test", + "not a url", + "", + ] { + assert!( + validated_base_url(raw).is_none(), + "{raw} should be rejected" + ); + } +} + +#[test] +fn base_url_override_falls_back_to_production_when_unset_or_invalid() { + assert_eq!(base_url_from_override(None), BASE_URL); + assert_eq!(base_url_from_override(Some(" ")), BASE_URL); + assert_eq!( + base_url_from_override(Some("http://attacker.example.com")), + BASE_URL + ); + assert_eq!( + base_url_from_override(Some("'http://127.0.0.1:4011'")), + "http://127.0.0.1:4011" + ); +}