Add provider-specific token account routing

apps/desktop-tauri/src-tauri/src/commands/providers.rs

@@ -84,98 +84,100 @@ pub(crate) fn build_fetch_context(
         .and_then(|override_data| override_data.env_override.as_ref());
     let active_token_api_key = active_token_env.and_then(|env| env.values().next().cloned());
     let usage_source = SourceMode::parse(settings.usage_source(id)).unwrap_or_default();
+    let token_account_kind = token_override.as_ref().map(|account| account.kind);
     // Selected token-account key overrides a stored provider apiKey (upstream #2271 / #1183).
-    let api_key = active_token_api_key.or(stored_api_key);
+    let api_key = active_token_api_key.clone().or(stored_api_key);
     let has_kimi_code_api_key =
         id == ProviderId::Kimi && api_key.as_deref().is_some_and(|key| !key.trim().is_empty());
     let has_opencodego_api_key = id == ProviderId::OpenCodeGo
         && api_key.as_deref().is_some_and(|key| !key.trim().is_empty());
 
-    let (mut source_mode, mut cookie_header, fails_closed_without_cookie) =
-        if id.cookie_domain().is_none() {
-            let source_mode = if active_token_env.is_some() {
-                SourceMode::OAuth
-            } else {
-                usage_source
-            };
-            (source_mode, None, false)
+    let (mut source_mode, mut cookie_header, fails_closed_without_cookie) = if id
+        .cookie_domain()
+        .is_none()
+    {
+        let source_mode = if active_token_env.is_some() {
+            SourceMode::OAuth
         } else {
-            match cookie_source {
-                // #433: an explicitly selected, non-empty Claude manual cookie is
-                // authoritative. Do not let an active OAuth token account silently
-                // replace it; this keeps tray refresh behavior aligned with diagnose,
-                // whose Claude Auto path tries the supplied Web cookie before OAuth.
-                "manual"
-                    if provider.manual_cookie_precedes_token_account()
-                        && stored_cookie
-                            .as_deref()
-                            .is_some_and(|cookie| !cookie.trim().is_empty()) =>
-                {
-                    (SourceMode::Web, stored_cookie.clone(), false)
-                }
-                _ if active_token_env.is_some() => (SourceMode::OAuth, None, false),
-                "off" if provider_uses_oauth_without_cookies(id, usage_source) => {
-                    (SourceMode::OAuth, None, false)
-                }
-                "off"
-                    if (has_kimi_code_api_key || has_opencodego_api_key)
-                        && usage_source == SourceMode::Auto =>
+            usage_source
+        };
+        (source_mode, None, false)
+    } else {
+        match cookie_source {
+            // #433: an explicitly selected, non-empty Claude manual cookie is
+            // authoritative. Do not let an active OAuth token account silently
+            // replace it; this keeps tray refresh behavior aligned with diagnose,
+            // whose Claude Auto path tries the supplied Web cookie before OAuth.
+            "manual"
+                if provider.manual_cookie_precedes_token_account()
+                    && stored_cookie
+                        .as_deref()
+                        .is_some_and(|cookie| !cookie.trim().is_empty()) =>
+            {
+                (SourceMode::Web, stored_cookie.clone(), false)
+            }
+            _ if active_token_env.is_some() => (SourceMode::OAuth, None, false),
+            "off" if provider_uses_oauth_without_cookies(id, usage_source) => {
+                (SourceMode::OAuth, None, false)
+            }
+            "off"
+                if (has_kimi_code_api_key || has_opencodego_api_key)
+                    && usage_source == SourceMode::Auto =>
+            {
+                (SourceMode::Auto, None, false)
+            }
+            // Droid/Factory: cookie-off must never scrape browser cookies. Map to
+            // Cli (API-only in the provider) so Auto does not fall through to web.
+            "off" if id == ProviderId::Factory => (SourceMode::Cli, None, false),
+            "off" => (SourceMode::Cli, None, false),
+            "manual" => {
+                let cookie_header = active_token_cookie.clone().or(stored_cookie);
+                let fails_closed_without_cookie = cookie_header.is_none()
+                    && provider.manual_empty_cookie_policy()
+                        == ManualEmptyCookiePolicy::FailClosedWeb;
+                let source_mode = if (has_kimi_code_api_key || has_opencodego_api_key)
+                    && usage_source == SourceMode::Auto
                 {
-                    (SourceMode::Auto, None, false)
-                }
-                // Droid/Factory: cookie-off must never scrape browser cookies. Map to
-                // Cli (API-only in the provider) so Auto does not fall through to web.
-                "off" if id == ProviderId::Factory => (SourceMode::Cli, None, false),
-                "off" => (SourceMode::Cli, None, false),
-                "manual" => {
-                    let cookie_header = active_token_cookie.or(stored_cookie);
-                    let fails_closed_without_cookie = cookie_header.is_none()
-                        && provider.manual_empty_cookie_policy()
-                            == ManualEmptyCookiePolicy::FailClosedWeb;
-                    let source_mode = if (has_kimi_code_api_key || has_opencodego_api_key)
-                        && usage_source == SourceMode::Auto
-                    {
-                        SourceMode::Auto
-                    } else if let Some(mode) = grok_source_mode_for_manual_cookie(id, usage_source)
-                    {
-                        // Grok Switch writes ~/.grok/auth.json. Leftover grok.com
-                        // cookies must not force Web, or Weekly/notifications keep
-                        // showing the previous browser account.
-                        mode
-                    } else if cookie_header.is_some() {
-                        SourceMode::Web
-                    } else if fails_closed_without_cookie {
-                        // The provider owns this policy; Web with no header means
-                        // it fails closed instead of importing a browser account
-                        // the user did not select.
-                        SourceMode::Web
-                    } else if provider_uses_oauth_without_cookies(id, usage_source) {
-                        SourceMode::OAuth
+                    SourceMode::Auto
+                } else if let Some(mode) = grok_source_mode_for_manual_cookie(id, usage_source) {
+                    // Grok Switch writes ~/.grok/auth.json. Leftover grok.com
+                    // cookies must not force Web, or Weekly/notifications keep
+                    // showing the previous browser account.
+                    mode
+                } else if cookie_header.is_some() {
+                    SourceMode::Web
+                } else if fails_closed_without_cookie {
+                    // The provider owns this policy; Web with no header means
+                    // it fails closed instead of importing a browser account
+                    // the user did not select.
+                    SourceMode::Web
+                } else if provider_uses_oauth_without_cookies(id, usage_source) {
+                    SourceMode::OAuth
+                } else {
+                    SourceMode::Cli
+                };
+                (source_mode, cookie_header, fails_closed_without_cookie)
+            }
+            // `browser` is accepted as a legacy alias from older settings.
+            "auto" | "browser" | "web" => {
+                // Claude resolves its cached cookie and browser fallback inside
+                // the provider; other providers retain the shell fallback.
+                let cookie_header = active_token_cookie.clone().or(stored_cookie).or_else(|| {
+                    if defer_provider_browser_cookie_lookup {
+                        None
                     } else {
-                        SourceMode::Cli
-                    };
-                    (source_mode, cookie_header, fails_closed_without_cookie)
-                }
-                // `browser` is accepted as a legacy alias from older settings.
-                "auto" | "browser" | "web" => {
-                    // Claude resolves its cached cookie and browser fallback inside
-                    // the provider; other providers retain the shell fallback.
-                    let cookie_header = active_token_cookie.or(stored_cookie).or_else(|| {
-                        if defer_provider_browser_cookie_lookup {
-                            None
-                        } else {
-                            provider_cookie_domain(id, settings).and_then(|domain| {
-                                codexbar::browser::cookies::get_cookie_header(domain)
-                                    .ok()
-                                    .filter(|h| !h.is_empty())
-                            })
-                        }
-                    });
-                    (usage_source, cookie_header, false)
-                }
-                _ => (usage_source, stored_cookie, false),
+                        provider_cookie_domain(id, settings).and_then(|domain| {
+                            codexbar::browser::cookies::get_cookie_header(domain)
+                                .ok()
+                                .filter(|h| !h.is_empty())
+                        })
+                    }
+                });
+                (usage_source, cookie_header, false)
             }
-        };
+            _ => (usage_source, stored_cookie, false),
+        }
+    };
 
     // Cookie-web providers (Cursor, OpenCode, …) reject SourceMode::Cli. The shell
     // historically mapped "manual + no cookie" to Cli, which surfaces as
@@ -220,16 +222,51 @@ pub(crate) fn build_fetch_context(
     // token account or manual cookie source scopes the session to web creds.
     let auto_prefer_web = token_override.is_some() || cookie_source == "manual";
 
+    // These upstream account types are explicit identity selections. Keep the
+    // provider's saved region/source settings intact, but project the selected
+    // credential into the route required by that account.
+    let (cookie_header, api_key) = match (id, token_account_kind, usage_source) {
+        (ProviderId::Kimi, Some(_), _) => (active_token_cookie.clone(), None),
+        (ProviderId::Doubao, Some(_), _) => (None, active_token_api_key.clone()),
+        (
+            ProviderId::OpenCodeGo,
+            Some(codexbar::core::TokenAccountKind::ApiKey),
+            SourceMode::Auto,
+        ) => (None, active_token_api_key.clone()),
+        (ProviderId::OpenCodeGo, Some(codexbar::core::TokenAccountKind::ApiKey), _) => {
+            (cookie_header, api_key)
+        }
+        (
+            ProviderId::OpenCodeGo,
+            Some(codexbar::core::TokenAccountKind::Cookie),
+            SourceMode::Auto,
+        ) => (active_token_cookie.clone(), api_key),
+        _ => (cookie_header, api_key),
+    };
+    let source_mode = token_override
+        .as_ref()
+        .and_then(|account| account.effective_source_mode(usage_source))
+        .unwrap_or(source_mode);
+    let token_account_isolated = token_override.is_some()
+        && matches!(
+            id,
+            ProviderId::Kimi | ProviderId::Doubao | ProviderId::OpenCodeGo
+        );
+
     FetchContext {
         source_mode,
         manual_cookie_header: cookie_header,
         manual_cookie_missing: fails_closed_without_cookie,
         api_key,
+        token_account_kind,
+        token_account_isolated,
         workspace_id: (!workspace_id.is_empty()).then_some(workspace_id),
         seat_credit_entitlement: settings.seat_credit_entitlement(id),
         api_region: (!api_region.is_empty()).then_some(api_region),
         gateway_url,
-        auto_prefer_web,
+        auto_prefer_web: auto_prefer_web
+            && !(id == ProviderId::OpenCodeGo
+                && token_account_kind == Some(codexbar::core::TokenAccountKind::ApiKey)),
         ..FetchContext::default()
     }
 }

apps/desktop-tauri/src-tauri/src/commands/tests.rs

@@ -483,6 +483,139 @@ fn fetch_context_opencode_empty_manual_remaps_to_web() {
     assert_eq!(ctx.source_mode, SourceMode::Web);
 }
 
+#[test]
+fn kimi_selected_account_forces_web_and_keeps_saved_region() {
+    let mut settings = Settings::default();
+    settings.set_usage_source(ProviderId::Kimi, "oauth");
+    settings.set_api_region(ProviderId::Kimi, "international");
+    let mut accounts = HashMap::new();
+    let mut data = ProviderAccountData::new();
+    data.add_account(TokenAccount::new("Work", "selected-kimi-session"));
+    accounts.insert(ProviderId::Kimi, data);
+
+    let ctx = super::build_fetch_context(
+        ProviderId::Kimi,
+        &settings,
+        &ManualCookies::default(),
+        &ApiKeys::default(),
+        &accounts,
+    );
+
+    assert_eq!(ctx.source_mode, SourceMode::Web);
+    assert_eq!(
+        ctx.manual_cookie_header.as_deref(),
+        Some("kimi-auth=selected-kimi-session")
+    );
+    assert_eq!(ctx.api_key, None);
+    assert_eq!(ctx.api_region.as_deref(), Some("international"));
+    assert!(ctx.token_account_isolated);
+    assert_eq!(settings.usage_source(ProviderId::Kimi), "oauth");
+    assert_eq!(settings.api_region(ProviderId::Kimi), "international");
+}
+
+#[test]
+fn doubao_selected_account_forces_ark_api_and_ignores_saved_source() {
+    let mut settings = Settings::default();
+    settings.set_usage_source(ProviderId::Doubao, "cli");
+    let mut accounts = HashMap::new();
+    let mut data = ProviderAccountData::new();
+    data.add_account(TokenAccount::new("Work", "selected-ark-key"));
+    accounts.insert(ProviderId::Doubao, data);
+
+    let ctx = super::build_fetch_context(
+        ProviderId::Doubao,
+        &settings,
+        &ManualCookies::default(),
+        &ApiKeys::default(),
+        &accounts,
+    );
+
+    assert_eq!(ctx.source_mode, SourceMode::OAuth);
+    assert_eq!(ctx.api_key.as_deref(), Some("selected-ark-key"));
+    assert!(ctx.token_account_isolated);
+}
+
+#[test]
+fn opencodego_selected_api_account_overrides_global_key_without_changing_explicit_source() {
+    let mut settings = Settings::default();
+    settings.set_usage_source(ProviderId::OpenCodeGo, "auto");
+    let mut keys = ApiKeys::default();
+    keys.set("opencodego", "global-key", None);
+    let mut accounts = HashMap::new();
+    let mut data = ProviderAccountData::new();
+    data.add_account(TokenAccount::new("Work", "selected-account-key"));
+    accounts.insert(ProviderId::OpenCodeGo, data);
+
+    let ctx = super::build_fetch_context(
+        ProviderId::OpenCodeGo,
+        &settings,
+        &ManualCookies::default(),
+        &keys,
+        &accounts,
+    );
+
+    assert_eq!(ctx.source_mode, SourceMode::Auto);
+    assert_eq!(ctx.api_key.as_deref(), Some("selected-account-key"));
+    assert!(!ctx.auto_prefer_web);
+    assert!(ctx.token_account_isolated);
+
+    for cookie_source in ["off", "manual"] {
+        settings.set_cookie_source(ProviderId::OpenCodeGo, cookie_source);
+        settings.set_usage_source(ProviderId::OpenCodeGo, "auto");
+        let auto_ctx = super::build_fetch_context(
+            ProviderId::OpenCodeGo,
+            &settings,
+            &ManualCookies::default(),
+            &keys,
+            &accounts,
+        );
+        assert_eq!(auto_ctx.source_mode, SourceMode::Auto);
+        assert_eq!(auto_ctx.api_key.as_deref(), Some("selected-account-key"));
+        assert!(auto_ctx.manual_cookie_header.is_none());
+    }
+
+    for (saved_source, expected_source) in [("web", SourceMode::Web), ("cli", SourceMode::Cli)] {
+        settings.set_cookie_source(ProviderId::OpenCodeGo, "off");
+        settings.set_usage_source(ProviderId::OpenCodeGo, saved_source);
+        let explicit_ctx = super::build_fetch_context(
+            ProviderId::OpenCodeGo,
+            &settings,
+            &ManualCookies::default(),
+            &keys,
+            &accounts,
+        );
+        assert_eq!(explicit_ctx.source_mode, expected_source);
+    }
+}
+
+#[test]
+fn opencodego_selected_cookie_account_uses_web_route() {
+    let settings = Settings::default();
+    let mut accounts = HashMap::new();
+    let mut data = ProviderAccountData::new();
+    data.add_account(TokenAccount::new("Web", "Cookie: session=selected-session"));
+    accounts.insert(ProviderId::OpenCodeGo, data);
+
+    let ctx = super::build_fetch_context(
+        ProviderId::OpenCodeGo,
+        &settings,
+        &ManualCookies::default(),
+        &ApiKeys::default(),
+        &accounts,
+    );
+
+    assert_eq!(ctx.source_mode, SourceMode::Web);
+    assert_eq!(
+        ctx.manual_cookie_header.as_deref(),
+        Some("Cookie: session=selected-session")
+    );
+    assert_eq!(
+        ctx.token_account_kind,
+        Some(codexbar::core::TokenAccountKind::Cookie)
+    );
+    assert!(ctx.token_account_isolated);
+}
+
 #[test]
 fn fetch_context_replicate_empty_manual_fails_closed_without_browser_import() {
     let settings = Settings::default();

rust/src/cli/diagnose.rs

@@ -177,6 +177,8 @@ async fn collect_provider_diagnostic(
             .map(ToOwned::to_owned),
         manual_cookie_missing: false,
         api_key: api_keys.get(provider_id.cli_name()).map(ToOwned::to_owned),
+        token_account_kind: None,
+        token_account_isolated: false,
         workspace_id: settings
             .provider_config(provider_id)
             .and_then(|config| config.workspace_id.clone()),

rust/src/cli/guard.rs

@@ -318,6 +318,8 @@ async fn fetch_guard_outcome(
         manual_cookie_header: None,
         manual_cookie_missing: false,
         api_key: None,
+        token_account_kind: None,
+        token_account_isolated: false,
         workspace_id: None,
         seat_credit_entitlement: None,
         api_region: None,