Import Trae CN cookies from the parent trae.cn domain

docs/PROVIDERS.md

@@ -813,10 +813,12 @@ GET https://platform.claude.com/api/organizations/<uuid>/current_spend
 ### 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.
+(trae.cn) account. Sign in to trae.cn in your browser, and CodexBar uses that
+session.
 
-- **Automatic:** CodexBar tries each detected browser holding an api.trae.cn
+- **Automatic:** Trae keeps its session cookies on the parent `trae.cn` domain,
+  so CodexBar reads the `trae.cn` cookies and sends only the ones a browser
+  would send to `api.trae.cn`. It tries each detected browser holding such a
   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

rust/src/providers/traecn/mod.rs

@@ -2,9 +2,10 @@
 //!
 //! 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
+//! browser session, which Trae keeps on the parent `trae.cn` domain (#4732),
+//! 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.
@@ -21,19 +22,29 @@ use reqwest::{
 };
 use serde_json::{Value, json};
 
+use crate::browser::cookies::CookieRequestScope;
 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};
+use crate::providers::{
+    BoundedBodyError, browser_cookie_request_sessions_from, 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";
+/// Automatic browser sessions. Trae keeps its login cookies on the parent
+/// `trae.cn` domain (upstream #4732), so the browser is read for `trae.cn`
+/// and only the cookies a browser sends to the API origin are kept. Pasted
+/// and imported cookies stay on `api.trae.cn` (`ProviderId::cookie_domain`).
+const COOKIE_SCOPE: CookieRequestScope = CookieRequestScope {
+    scan_domain: "trae.cn",
+    host: "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/";
@@ -198,8 +209,9 @@ impl TraeCnProvider {
     }
 
     /// 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`).
+    /// holding a `trae.cn` session that the API origin receives is tried in
+    /// browser order, moving on only after Trae rejects a session (upstream
+    /// `rejectCookie`).
     async fn fetch_web<S: BrowserCookieStore>(
         &self,
         ctx: &FetchContext,
@@ -222,8 +234,7 @@ impl TraeCnProvider {
             return Err(missing_session());
         }
 
-        let mut sessions =
-            crate::providers::browser_cookie_sessions_from(store, ctx, &[COOKIE_DOMAIN]);
+        let mut sessions = browser_cookie_request_sessions_from(store, ctx, COOKIE_SCOPE);
         let mut rejected = false;
         for header in sessions.by_ref() {
             let Some(cookie) = crate::providers::normalize_cookie_header(&header) else {

rust/src/providers/traecn/tests.rs

@@ -7,6 +7,8 @@ use crate::browser::detection::BrowserType;
 use crate::browser::session_scan::test_support::FakeCookieStore;
 use mockito::{Matcher, Mock, Server, ServerGuard};
 
+/// Where Trae keeps automatic browser sessions (upstream #4732).
+const PARENT_DOMAIN: &str = "trae.cn";
 const COOKIE: &str = "trae_session=synthetic-1";
 const OTHER_COOKIE: &str = "trae_session=synthetic-2";
 const TOKEN: &str = "synthetic-token";
@@ -108,6 +110,8 @@ fn metadata_and_source_contract_match_upstream_descriptor() {
         provider.manual_empty_cookie_policy(),
         ManualEmptyCookiePolicy::FailClosedWeb
     ));
+    // Manual cookies and browser import use the API origin; only automatic
+    // browser sessions read the parent `trae.cn` domain (#4732).
     assert_eq!(ProviderId::TraeCn.cookie_domain(), Some("api.trae.cn"));
     assert_eq!(
         ProviderId::from_cli_name("traecn"),
@@ -327,15 +331,15 @@ async fn a_rejected_profile_advances_to_the_next_browser_session() {
     )
     .await;
     let store = FakeCookieStore::default()
-        .cookies(
+        .host_cookies(
             BrowserType::Chrome,
-            COOKIE_DOMAIN,
-            &[("trae_session", "synthetic-1")],
+            PARENT_DOMAIN,
+            &[(".trae.cn", "trae_session", "synthetic-1")],
         )
-        .cookies(
+        .host_cookies(
             BrowserType::Edge,
-            COOKIE_DOMAIN,
-            &[("trae_session", "synthetic-2")],
+            PARENT_DOMAIN,
+            &[(".trae.cn", "trae_session", "synthetic-2")],
         );
     let result = fetch(&server, &FetchContext::default(), &store)
         .await
@@ -350,15 +354,15 @@ async fn all_rejected_browser_sessions_report_an_expired_session() {
     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(
+        .host_cookies(
             BrowserType::Chrome,
-            COOKIE_DOMAIN,
-            &[("trae_session", "synthetic-1")],
+            PARENT_DOMAIN,
+            &[(".trae.cn", "trae_session", "synthetic-1")],
         )
-        .cookies(
+        .host_cookies(
             BrowserType::Edge,
-            COOKIE_DOMAIN,
-            &[("trae_session", "synthetic-2")],
+            PARENT_DOMAIN,
+            &[(".trae.cn", "trae_session", "synthetic-2")],
         );
     let provider = TraeCnProvider::with_base_url(&server.url());
     let error = provider
@@ -386,15 +390,15 @@ async fn transient_failures_keep_the_session_and_stop_the_walk() {
             .await;
         // A second browser would only be tried after a rejected session.
         let store = FakeCookieStore::default()
-            .cookies(
+            .host_cookies(
                 BrowserType::Chrome,
-                COOKIE_DOMAIN,
-                &[("trae_session", "synthetic-1")],
+                PARENT_DOMAIN,
+                &[(".trae.cn", "trae_session", "synthetic-1")],
             )
-            .cookies(
+            .host_cookies(
                 BrowserType::Edge,
-                COOKIE_DOMAIN,
-                &[("trae_session", "synthetic-2")],
+                PARENT_DOMAIN,
+                &[(".trae.cn", "trae_session", "synthetic-2")],
             );
         let result = TraeCnProvider::with_base_url(&server.url())
             .fetch_web(&FetchContext::default(), &store)
@@ -434,10 +438,10 @@ async fn other_http_failures_name_the_status_without_echoing_the_body() {
 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(
+    let store = FakeCookieStore::default().host_cookies(
         BrowserType::Chrome,
-        COOKIE_DOMAIN,
-        &[("trae_session", "synthetic-2")],
+        PARENT_DOMAIN,
+        &[(".trae.cn", "trae_session", "synthetic-2")],
     );
     let result = fetch(&server, &manual(&format!("Cookie: {COOKIE}")), &store)
         .await
@@ -467,10 +471,10 @@ async fn no_browser_session_asks_the_user_to_sign_in() {
 #[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(
+    let store = FakeCookieStore::default().host_cookies(
         BrowserType::Chrome,
-        COOKIE_DOMAIN,
-        &[("trae_session", "synthetic-1")],
+        PARENT_DOMAIN,
+        &[(".trae.cn", "trae_session", "synthetic-1")],
     );
     let ctx = FetchContext {
         source_mode: SourceMode::Auto,
@@ -485,6 +489,65 @@ async fn an_empty_manual_source_never_falls_back_to_a_browser() {
     assert!(store.reads().is_empty());
 }
 
+#[tokio::test]
+async fn automatic_scan_reads_the_parent_domain_and_keeps_only_api_bound_cookies() {
+    let mut server = Server::new_async().await;
+    let mocks = exchange(&mut server, COOKIE, (200, TOKEN_BODY), (200, ENTITLEMENTS)).await;
+    // Only the `.trae.cn` cookie is sent to api.trae.cn. The host-only parent
+    // cookie and the sibling www host are not, so Chrome yields no session.
+    let store = FakeCookieStore::default()
+        .host_cookies(
+            BrowserType::Chrome,
+            PARENT_DOMAIN,
+            &[
+                ("trae.cn", "parent_host_only", "1"),
+                ("www.trae.cn", "www_only", "2"),
+            ],
+        )
+        .host_cookies(
+            BrowserType::Edge,
+            PARENT_DOMAIN,
+            &[
+                (".trae.cn", "trae_session", "synthetic-1"),
+                ("www.trae.cn", "www_only", "3"),
+            ],
+        );
+    let result = fetch(&server, &FetchContext::default(), &store)
+        .await
+        .unwrap();
+    for mock in mocks {
+        mock.assert_async().await;
+    }
+    assert_eq!(result.usage.primary.used_percent, 20.0);
+    assert_eq!(
+        store.reads(),
+        vec![
+            (BrowserType::Chrome, PARENT_DOMAIN.to_string()),
+            (BrowserType::Edge, PARENT_DOMAIN.to_string()),
+        ]
+    );
+    // Requests and manual cookies stay on the API origin.
+    assert_eq!(Some(COOKIE_SCOPE.host), ProviderId::TraeCn.cookie_domain());
+}
+
+#[tokio::test]
+async fn parent_host_only_cookie_is_not_sent_to_the_api_origin() {
+    let server = Server::new_async().await;
+    let store = FakeCookieStore::default().host_cookies(
+        BrowserType::Chrome,
+        PARENT_DOMAIN,
+        &[("trae.cn", "trae_session", "synthetic-1")],
+    );
+    // No request may leave: a sent cookie would reach the mock server, which
+    // has no route and answers with a non-sign-in error.
+    let result = fetch(&server, &FetchContext::default(), &store).await;
+    assert!(matches!(
+        result,
+        Err(ProviderError::BrowserSignInRequired { .. })
+    ));
+    assert_eq!(store.reads().len(), 1);
+}
+
 #[test]
 fn amounts_must_be_non_negative_json_numbers() {
     assert_eq!(amount(Some(&json!(0))).unwrap(), 0.0);