diff --git a/app/src/ai/agent/api.rs b/app/src/ai/agent/api.rs index da01567e..72d73b61 100644 --- a/app/src/ai/agent/api.rs +++ b/app/src/ai/agent/api.rs @@ -112,7 +112,8 @@ pub struct RequestParams { pub planning_enabled: bool, should_redact_secrets: bool, - /// User-provided API keys for AI providers (BYO API Key). + /// User-provided API keys for AI providers, currently AWS Bedrock credentials + /// and/or Gemini Enterprise (GEAP) credentials. pub api_keys: Option, /// User-provided custom model providers (BYOK endpoints). pub custom_model_providers: @@ -121,7 +122,6 @@ pub struct RequestParams { /// `custom_model_providers`: the selected model's `config_key` indexes into this /// registry. `None` when no custom router is selected. pub custom_model_routers: Option, - pub allow_use_of_warp_credits: bool, pub autonomy_level: warp_multi_agent_api::AutonomyLevel, pub isolation_level: warp_multi_agent_api::IsolationLevel, pub web_search_enabled: bool, @@ -203,7 +203,6 @@ impl RequestParams { api_keys: None, custom_model_providers: None, custom_model_routers: None, - allow_use_of_warp_credits: false, autonomy_level: Default::default(), isolation_level: Default::default(), web_search_enabled: false, @@ -305,16 +304,12 @@ impl RequestParams { let user_workspaces = UserWorkspaces::as_ref(app); let api_key_manager = ApiKeyManager::as_ref(app); - let is_byo_enabled = user_workspaces.is_byo_api_key_enabled(app); #[cfg(not(target_family = "wasm"))] let geap_binding = crate::ai::geap_credentials::current_geap_policy(app).mint_binding(); #[cfg(target_family = "wasm")] let geap_binding: Option<::ai::api_keys::GeapMintBinding> = None; - let api_keys = api_key_manager.api_keys_for_request( - is_byo_enabled, - user_workspaces.is_bedrock_enabled(app), - geap_binding, - ); + let api_keys = api_key_manager + .api_keys_for_request(user_workspaces.is_bedrock_enabled(app), geap_binding); let custom_model_providers = None; let custom_model_routers = FeatureFlag::CustomModelRouters.is_enabled().then(|| { LLMPreferences::as_ref(app).custom_model_routers_for_request( @@ -322,7 +317,6 @@ impl RequestParams { &request_input.coding_model_id, ) }); - let allow_use_of_warp_credits = *AISettings::as_ref(app).can_use_warp_credits_for_fallback; let app_execution_mode = AppExecutionMode::as_ref(app); let autonomy_level = if app_execution_mode.is_autonomous() { @@ -399,7 +393,6 @@ impl RequestParams { api_keys, custom_model_providers, custom_model_routers, - allow_use_of_warp_credits, autonomy_level, isolation_level, web_search_enabled, diff --git a/app/src/ai/agent/api/impl.rs b/app/src/ai/agent/api/impl.rs index 5c412f51..d3884f28 100644 --- a/app/src/ai/agent/api/impl.rs +++ b/app/src/ai/agent/api/impl.rs @@ -57,11 +57,6 @@ pub async fn generate_multi_agent_output( redaction::redact_inputs(&mut params.input); } - let api_keys = api_keys_with_warp_credit_fallback_setting( - params.api_keys, - params.allow_use_of_warp_credits, - ); - let mut request = api::Request { task_context: Some(api::request::TaskContext { tasks: params.tasks, @@ -91,7 +86,7 @@ pub async fn generate_multi_agent_output( supports_suggest_prompt: true, supports_read_image_files: FeatureFlag::ReadImageFiles.is_enabled(), supports_reasoning_message: true, - api_keys, + api_keys: params.api_keys, autonomy_level: params.autonomy_level.into(), isolation_level: params.isolation_level.into(), web_search_enabled: params.web_search_enabled, @@ -214,23 +209,6 @@ pub async fn generate_multi_agent_output( } } -fn api_keys_with_warp_credit_fallback_setting( - api_keys: Option, - allow_use_of_warp_credits: bool, -) -> Option { - match api_keys { - Some(mut api_keys) => { - api_keys.allow_use_of_warp_credits = allow_use_of_warp_credits; - Some(api_keys) - } - None if allow_use_of_warp_credits => Some(api::request::settings::ApiKeys { - allow_use_of_warp_credits: true, - ..Default::default() - }), - None => None, - } -} - fn supports_orchestration_v2(orchestration_enabled: bool) -> bool { orchestration_enabled } diff --git a/app/src/ai/agent/api/impl_tests.rs b/app/src/ai/agent/api/impl_tests.rs index 66f55ee9..9dce7aeb 100644 --- a/app/src/ai/agent/api/impl_tests.rs +++ b/app/src/ai/agent/api/impl_tests.rs @@ -2,10 +2,7 @@ use galaxy_core::features::FeatureFlag; use galaxy_core::HostId; use warp_multi_agent_api as api; -use super::{ - api_keys_with_warp_credit_fallback_setting, get_supported_cli_agent_tools, get_supported_tools, - supports_orchestration_v2, -}; +use super::{get_supported_cli_agent_tools, get_supported_tools, supports_orchestration_v2}; use crate::ai::agent::api::RequestParams; use crate::ai::blocklist::SessionContext; use crate::ai::llms::LLMId; @@ -36,7 +33,6 @@ fn request_params_with_ask_user_question_enabled(ask_user_question_enabled: bool api_keys: None, custom_model_providers: None, custom_model_routers: None, - allow_use_of_warp_credits: false, autonomy_level: api::AutonomyLevel::Supervised, isolation_level: api::IsolationLevel::None, web_search_enabled: false, @@ -64,47 +60,6 @@ fn request_params_for_remote(host_id: Option) -> RequestParams { params } -#[test] -fn api_keys_with_warp_credit_fallback_setting_returns_none_without_keys_or_fallback() { - let api_keys = api_keys_with_warp_credit_fallback_setting(None, false); - - assert!(api_keys.is_none()); -} - -#[test] -fn api_keys_with_warp_credit_fallback_setting_creates_fallback_only_api_keys() { - let api_keys = api_keys_with_warp_credit_fallback_setting(None, true) - .expect("fallback setting should create ApiKeys"); - - assert!(api_keys.allow_use_of_warp_credits); - assert!(api_keys.anthropic.is_empty()); - assert!(api_keys.openai.is_empty()); - assert!(api_keys.google.is_empty()); - assert!(api_keys.open_router.is_empty()); - assert!(api_keys.aws_credentials.is_none()); -} - -#[test] -fn api_keys_with_warp_credit_fallback_setting_preserves_existing_keys() { - let api_keys = api_keys_with_warp_credit_fallback_setting( - Some(api::request::settings::ApiKeys { - anthropic: "anthropic-key".to_string(), - openai: String::new(), - google: String::new(), - open_router: String::new(), - grok_oauth_access_token: String::new(), - allow_use_of_warp_credits: false, - aws_credentials: None, - google_cloud_credentials: None, - }), - true, - ) - .expect("existing ApiKeys should be preserved"); - - assert_eq!(api_keys.anthropic, "anthropic-key"); - assert!(api_keys.allow_use_of_warp_credits); -} - #[test] fn supports_orchestration_v2_matches_request_orchestration_setting() { assert!(supports_orchestration_v2(true)); diff --git a/app/src/ai/agent/conversation_tests.rs b/app/src/ai/agent/conversation_tests.rs index e113f4a3..d930c12e 100644 --- a/app/src/ai/agent/conversation_tests.rs +++ b/app/src/ai/agent/conversation_tests.rs @@ -1,23 +1,14 @@ use std::collections::HashMap; -use ai::api_keys::ApiKeyManager; use galaxy_core::features::FeatureFlag; use warp_multi_agent_api as api; -use warpui::{App, SingletonEntity}; use super::{ - artifact_from_fork_proto, footer_model_token_usage, AIConversation, - AIConversationAutoexecuteMode, AIConversationId, ConversationStatus, RestoreConversationError, + artifact_from_fork_proto, AIConversation, AIConversationAutoexecuteMode, AIConversationId, + ConversationStatus, RestoreConversationError, }; use crate::ai::artifacts::Artifact; -use crate::ai::llms::LLMPreferences; -use crate::auth::auth_manager::AuthManager; -use crate::auth::AuthStateProvider; -use crate::network::NetworkStatus; use crate::persistence::model::AgentConversationData; -use crate::server::server_api::ServerApiProvider; -use crate::test_util::settings::initialize_settings_for_tests; -use crate::workspaces::user_workspaces::UserWorkspaces; fn restored_conversation(conversation_data: Option) -> AIConversation { AIConversation::new_restored( @@ -115,43 +106,6 @@ fn restored_conversation_with_queries(queries: &[&str]) -> AIConversation { .unwrap() } -fn initialize_custom_endpoint_usage_test_app(app: &mut App) { - initialize_settings_for_tests(app); - app.add_singleton_model(|_| ServerApiProvider::new_for_test()); - app.add_singleton_model(|_| NetworkStatus::new()); - app.add_singleton_model(UserWorkspaces::default_mock); - app.add_singleton_model(|_| AuthStateProvider::new_for_test()); - app.add_singleton_model(AuthManager::new_for_test); -} - -#[allow(deprecated)] -fn custom_endpoint_usage_metadata( - config_key: &str, - total_tokens: u32, -) -> api::response_event::stream_finished::ConversationUsageMetadata { - let category = "primary_agent".to_string(); - api::response_event::stream_finished::ConversationUsageMetadata { - context_window_usage: 0.0, - credits_spent: 0.0, - platform_credits_spent: 0.0, - summarized: false, - token_usage: vec![], - tool_usage_metadata: None, - total_input_tokens: 0, - warp_token_usage: HashMap::new(), - byok_token_usage: HashMap::new(), - context_window_segments: Vec::new(), - custom_endpoint_token_usage: HashMap::from([( - config_key.to_string(), - api::response_event::stream_finished::ModelTokenUsage { - model_id: config_key.to_string(), - total_tokens, - token_usage_by_category: HashMap::from([(category, total_tokens)]), - }, - )]), - } -} - #[test] fn latest_user_query_returns_latest_non_empty_user_query() { let conversation = @@ -291,217 +245,6 @@ fn restored_conversation_with_empty_task_list_creates_in_progress_optimistic_roo assert!(conversation.status_error_message().is_none()); } -#[test] -fn update_cost_and_usage_resolves_custom_endpoint_alias_for_footer_usage() { - App::test((), |mut app| async move { - initialize_custom_endpoint_usage_test_app(&mut app); - ApiKeyManager::handle(&app).update(&mut app, |manager, ctx| { - manager.add_custom_endpoint( - "Endpoint".to_string(), - "https://custom.example".to_string(), - "key".to_string(), - vec![( - "raw-model".to_string(), - Some("Friendly alias".to_string()), - Some("config-key".to_string()), - )], - ctx, - ); - }); - app.add_singleton_model(LLMPreferences::new); - - let mut conversation = AIConversation::new(false, false); - app.read(|ctx| { - conversation - .update_cost_and_usage_for_request( - None, - vec![], - Some(custom_endpoint_usage_metadata("config-key", 6)), - false, - ctx, - ) - .expect("custom endpoint usage should update"); - }); - - let usage = conversation - .token_usage() - .iter() - .find(|usage| usage.model_id == "Friendly alias") - .expect("custom endpoint alias should resolve into footer usage"); - assert_eq!(usage.custom_endpoint_tokens, 6); - assert_eq!(usage.byok_tokens, 0); - assert_eq!( - usage - .custom_endpoint_token_usage_by_category - .get("primary_agent"), - Some(&6) - ); - }); -} - -#[test] -fn update_cost_and_usage_uses_fallback_label_for_unknown_custom_endpoint() { - App::test((), |mut app| async move { - initialize_custom_endpoint_usage_test_app(&mut app); - app.add_singleton_model(LLMPreferences::new); - - let mut conversation = AIConversation::new(false, false); - app.read(|ctx| { - conversation - .update_cost_and_usage_for_request( - None, - vec![], - Some(custom_endpoint_usage_metadata("missing-config-key", 9)), - false, - ctx, - ) - .expect("fallback custom endpoint usage should update"); - }); - - let usage = conversation - .token_usage() - .iter() - .find(|usage| usage.model_id == "Custom endpoint") - .expect("unknown custom endpoint usage should use the fallback label"); - assert_eq!(usage.custom_endpoint_tokens, 9); - assert_eq!(usage.byok_tokens, 0); - assert_eq!( - usage - .custom_endpoint_token_usage_by_category - .get("primary_agent"), - Some(&9) - ); - }); -} - -#[allow(deprecated)] -#[test] -fn footer_model_token_usage_keeps_custom_endpoint_usage_distinct_from_same_labeled_models() { - App::test((), |mut app| async move { - initialize_custom_endpoint_usage_test_app(&mut app); - ApiKeyManager::handle(&app).update(&mut app, |manager, ctx| { - manager.add_custom_endpoint( - "Endpoint".to_string(), - "https://custom.example".to_string(), - "key".to_string(), - vec![( - "raw-model".to_string(), - Some("Resolved custom".to_string()), - Some("config-key".to_string()), - )], - ctx, - ); - }); - app.add_singleton_model(LLMPreferences::new); - - let category = "primary_agent".to_string(); - let usage_metadata = api::response_event::stream_finished::ConversationUsageMetadata { - context_window_usage: 0.0, - credits_spent: 0.0, - platform_credits_spent: 0.0, - summarized: false, - #[allow(deprecated)] - token_usage: vec![], - tool_usage_metadata: None, - total_input_tokens: 0, - warp_token_usage: HashMap::new(), - byok_token_usage: HashMap::from([( - "Resolved custom".to_string(), - api::response_event::stream_finished::ModelTokenUsage { - model_id: "Resolved custom".to_string(), - total_tokens: 4, - token_usage_by_category: HashMap::from([(category.clone(), 4)]), - }, - )]), - custom_endpoint_token_usage: HashMap::from([( - "config-key".to_string(), - api::response_event::stream_finished::ModelTokenUsage { - model_id: "config-key".to_string(), - total_tokens: 6, - token_usage_by_category: HashMap::from([(category.clone(), 6)]), - }, - )]), - context_window_segments: Vec::new(), - }; - - let model_usage = - app.read(|ctx| footer_model_token_usage(&usage_metadata, LLMPreferences::as_ref(ctx))); - let byok_usage = model_usage - .iter() - .find(|usage| usage.model_id == "Resolved custom" && usage.byok_tokens == 4) - .expect("existing model usage should be present"); - let custom_usage = model_usage - .iter() - .find(|usage| usage.model_id == "Resolved custom" && usage.custom_endpoint_tokens == 6) - .expect("custom endpoint usage should remain distinct"); - - assert_eq!(model_usage.len(), 2); - assert_eq!( - byok_usage.byok_token_usage_by_category.get(&category), - Some(&4) - ); - assert_eq!( - custom_usage - .custom_endpoint_token_usage_by_category - .get(&category), - Some(&6) - ); - assert_eq!(byok_usage.warp_tokens, 0); - assert_eq!(custom_usage.warp_tokens, 0); - assert_eq!(custom_usage.byok_tokens, 0); - }); -} - -#[allow(deprecated)] -#[test] -fn footer_model_token_usage_preserves_unresolved_custom_endpoint_usage_with_fallback_label() { - App::test((), |mut app| async move { - initialize_custom_endpoint_usage_test_app(&mut app); - app.add_singleton_model(LLMPreferences::new); - - let category = "primary_agent".to_string(); - let usage_metadata = api::response_event::stream_finished::ConversationUsageMetadata { - context_window_usage: 0.0, - credits_spent: 0.0, - platform_credits_spent: 0.0, - summarized: false, - #[allow(deprecated)] - token_usage: vec![], - tool_usage_metadata: None, - total_input_tokens: 0, - warp_token_usage: HashMap::new(), - byok_token_usage: HashMap::new(), - custom_endpoint_token_usage: HashMap::from([( - "missing-config-key".to_string(), - api::response_event::stream_finished::ModelTokenUsage { - model_id: "missing-config-key".to_string(), - total_tokens: 9, - token_usage_by_category: HashMap::from([(category.clone(), 9)]), - }, - )]), - context_window_segments: Vec::new(), - }; - - let model_usage = - app.read(|ctx| footer_model_token_usage(&usage_metadata, LLMPreferences::as_ref(ctx))); - let custom_usage = model_usage - .iter() - .find(|usage| usage.model_id == "Custom endpoint") - .expect("fallback custom endpoint usage should be present"); - - assert_eq!(model_usage.len(), 1); - assert_eq!(custom_usage.custom_endpoint_tokens, 9); - assert_eq!(custom_usage.byok_tokens, 0); - assert_eq!( - custom_usage - .custom_endpoint_token_usage_by_category - .get(&category), - Some(&9) - ); - assert_eq!(custom_usage.warp_tokens, 0); - }); -} - /// The legacy `AgentConversationData.root_task_is_optimistic` flag must be /// ignored on restore. A non-empty task list always produces a real /// server-backed root regardless of whether the flag is set. diff --git a/app/src/ai/blocklist/block/view_impl/common.rs b/app/src/ai/blocklist/block/view_impl/common.rs index ebb3f167..01c94ce5 100644 --- a/app/src/ai/blocklist/block/view_impl/common.rs +++ b/app/src/ai/blocklist/block/view_impl/common.rs @@ -3318,12 +3318,12 @@ fn render_invalid_api_key_error( background: Some(internal_colors::fg_overlay_3(theme).into()), ..Default::default() }) - .with_text_label("Edit API Keys".to_string()) + .with_text_label("Manage Providers".to_string()) .with_cursor(Some(Cursor::PointingHand)) .build() .on_click(move |ctx, _, _| { ctx.dispatch_typed_action(WorkspaceAction::ShowSettingsPageWithSearch { - search_query: "api keys".to_string(), + search_query: "bedrock openai litellm".to_string(), section: Some(SettingsSection::WarpAgent), }); }) diff --git a/app/src/ai/blocklist/controller.rs b/app/src/ai/blocklist/controller.rs index 4e86f5f0..c7b833af 100644 --- a/app/src/ai/blocklist/controller.rs +++ b/app/src/ai/blocklist/controller.rs @@ -671,8 +671,7 @@ impl BlocklistAIController { OrchestrationEventStreamerEvent::ChildSpawned { .. } | OrchestrationEventStreamerEvent::ChildStatusChanged { .. } => {} }); - let crosscheck_reviewer = - ctx.add_model(crate::ai::crosscheck::CrosscheckReviewer::new); + let crosscheck_reviewer = ctx.add_model(crate::ai::crosscheck::CrosscheckReviewer::new); ctx.subscribe_to_model(&crosscheck_reviewer, move |me, _, event, ctx| { use crate::ai::crosscheck::{CrosscheckReviewerEvent, ReviewOutcome}; let CrosscheckReviewerEvent::ReviewCompleted { @@ -2172,9 +2171,7 @@ impl BlocklistAIController { match outcome { ReviewOutcome::Approved => { - log::info!( - "[crosscheck] Work approved for conversation {conversation_id:?}" - ); + log::info!("[crosscheck] Work approved for conversation {conversation_id:?}"); // Nothing to do — the conversation completes normally. } ReviewOutcome::MaxIterationsReached { last_feedback } => { @@ -2191,9 +2188,7 @@ impl BlocklistAIController { self.inject_crosscheck_feedback(conversation_id, message, false, ctx); } ReviewOutcome::Error { error } => { - log::error!( - "[crosscheck] Reviewer failed for {conversation_id:?}: {error}" - ); + log::error!("[crosscheck] Reviewer failed for {conversation_id:?}: {error}"); // Don't block the conversation on reviewer errors; just log it. } } @@ -2303,13 +2298,7 @@ impl BlocklistAIController { let max_iterations = ai_settings.crosscheck_max_iterations(); self.crosscheck_reviewer.update(ctx, |reviewer, ctx| { - reviewer.start_review( - conversation_id, - max_iterations, - agent_output, - model_id, - ctx, - ); + reviewer.start_review(conversation_id, max_iterations, agent_output, model_id, ctx); }); } @@ -2356,9 +2345,7 @@ impl BlocklistAIController { AIAgentTextSection::PlainText { text } => { Some(text.text().to_string()) } - AIAgentTextSection::Code { code, .. } => { - Some(code.clone()) - } + AIAgentTextSection::Code { code, .. } => Some(code.clone()), _ => None, }) .collect::>() @@ -2786,20 +2773,15 @@ impl BlocklistAIController { &conversation_data.server_conversation_token, ); - // Safety net: if the connected Grok subscription's OAuth token is + // Safety net: if the Gemini Enterprise (GEAP) OIDC/WIF credential is // nearing or past expiry, kick off a background refresh so upcoming - // requests can authenticate even when the proactive refresh loop - // isn't running. This request still carries the currently stored - // token; the server is the authority on its validity. The Gemini - // Enterprise (GEAP) analog re-arms a parked or never-armed WIF - // credential refresh chain the same way. + // requests can authenticate even when the proactive refresh loop isn't + // running (e.g. a parked or never-armed refresh chain). #[cfg(not(target_family = "wasm"))] { use ::ai::api_keys::ApiKeyManager; - let byo_allowed = UserWorkspaces::as_ref(ctx).is_byo_api_key_enabled(ctx); ApiKeyManager::handle(ctx).update(ctx, |manager, ctx| { - manager.refresh_grok_tokens_if_needed(byo_allowed, ctx); crate::ai::geap_credentials::refresh_geap_credentials_if_needed(manager, ctx); }); } @@ -3391,10 +3373,8 @@ impl BlocklistAIController { || error_str.contains("ThrottlingException")); const MAX_ERROR_RETRIES: usize = 2; - let retry_count = self - .error_retry_counts - .entry(conversation_id) - .or_insert(0); + let retry_count = + self.error_retry_counts.entry(conversation_id).or_insert(0); let should_corrective_retry = is_corrective_retry_candidate && *retry_count < MAX_ERROR_RETRIES; @@ -3445,17 +3425,15 @@ impl BlocklistAIController { error_str ); - let inputs = vec![ - AIAgentInput::UserQuery { - query: corrective_msg, - context: Arc::from([]), - static_query_type: None, - referenced_attachments: HashMap::new(), - user_query_mode: UserQueryMode::Normal, - running_command: None, - intended_agent: None, - }, - ]; + let inputs = vec![AIAgentInput::UserQuery { + query: corrective_msg, + context: Arc::from([]), + static_query_type: None, + referenced_attachments: HashMap::new(), + user_query_mode: UserQueryMode::Normal, + running_command: None, + intended_agent: None, + }]; let _ = self.send_request_input( RequestInput::for_task( diff --git a/app/src/ai/blocklist/prompt/prompt_alert.rs b/app/src/ai/blocklist/prompt/prompt_alert.rs index 68370329..6fd5ad2b 100644 --- a/app/src/ai/blocklist/prompt/prompt_alert.rs +++ b/app/src/ai/blocklist/prompt/prompt_alert.rs @@ -324,16 +324,6 @@ impl PromptAlertView { }; text_fragments.push(FormattedTextFragment::hyperlink(label, upgrade_url)); } - if UserWorkspaces::as_ref(app).is_byo_api_key_enabled(app) { - text_fragments.push(FormattedTextFragment::plain_text(" or ")); - text_fragments.push(FormattedTextFragment::hyperlink_action( - "use your own API keys", - WorkspaceAction::ShowSettingsPageWithSearch { - search_query: "api".to_string(), - section: Some(SettingsSection::WarpAgent), - }, - )); - } } PromptAlertState::NoAlert => {} } diff --git a/app/src/ai/llms.rs b/app/src/ai/llms.rs index 39a82e9c..12879148 100644 --- a/app/src/ai/llms.rs +++ b/app/src/ai/llms.rs @@ -3,7 +3,7 @@ use std::collections::{HashMap, HashSet}; use std::sync::{Arc, OnceLock}; -use ai::api_keys::{ApiKeyManager, ApiKeyManagerEvent, CustomEndpoint, CustomEndpointModel}; +use ai::api_keys::ApiKeyManager; pub use ai::LLMId; use galaxy_core::features::FeatureFlag; use galaxy_core::ui::icons::Icon; @@ -27,23 +27,10 @@ use crate::workspaces::user_workspaces::{UserWorkspaces, UserWorkspacesEvent}; use crate::{report_error, AISettings}; /// Checks if a user's' API key is being used for the given provider. -/// Returns `true` if BYO API key is enabled and a key exists for the provider. -/// For xAI, a connected Grok subscription counts: its OAuth access token is -/// sent like a BYO key (see `ApiKeyManager::api_keys_for_request`). -pub fn is_using_api_key_for_provider(provider: &LLMProvider, app: &AppContext) -> bool { - if !UserWorkspaces::as_ref(app).is_byo_api_key_enabled(app) { - return false; - } - let manager = ApiKeyManager::as_ref(app); - - match provider { - LLMProvider::OpenAI => manager.keys().openai.is_some(), - LLMProvider::Anthropic => manager.keys().anthropic.is_some(), - LLMProvider::Google => manager.keys().google.is_some(), - LLMProvider::Xai => manager.grok_tokens().is_some(), - LLMProvider::Bedrock | LLMProvider::LiteLLM => false, - LLMProvider::Unknown => false, - } +/// AWS Bedrock is the only supported provider in Galaxy; there is no +/// user-pasted BYO key path for other providers. +pub fn is_using_api_key_for_provider(_provider: &LLMProvider, _app: &AppContext) -> bool { + false } pub fn should_show_bedrock_icon_for_model(llm: &LLMInfo, app: &AppContext) -> bool { @@ -423,8 +410,8 @@ impl AvailableLLMs { } fn default_llm_info(&self) -> &LLMInfo { - static NO_PROVIDER_FALLBACK: std::sync::LazyLock = std::sync::LazyLock::new(|| { - LLMInfo { + static NO_PROVIDER_FALLBACK: std::sync::LazyLock = + std::sync::LazyLock::new(|| LLMInfo { display_name: "No models configured".to_owned(), base_model_name: "No models configured".to_owned(), id: "none".to_owned().into(), @@ -433,9 +420,7 @@ impl AvailableLLMs { request_multiplier: 1, credit_multiplier: None, }, - description: Some( - "Enable Bedrock or OpenAI/LiteLLM in settings".to_string(), - ), + description: Some("Enable Bedrock or OpenAI/LiteLLM in settings".to_string()), disable_reason: Some(DisableReason::Unavailable), vision_supported: false, spec: None, @@ -443,8 +428,7 @@ impl AvailableLLMs { host_configs: HashMap::new(), discount_percentage: None, context_window: LLMContextWindow::default(), - } - }); + }); self.info_for_id(&self.default_id) .or_else(|| self.choices.first()) @@ -675,7 +659,7 @@ impl LLMPreferences { }); let base_llm_for_terminal_view = HashMap::new(); - let custom_llms = build_custom_llm_infos(ApiKeyManager::as_ref(ctx).keys()); + let custom_llms = Vec::new(); let mut me = Self { models_by_feature, @@ -758,7 +742,9 @@ impl LLMPreferences { .choices .retain(|m| m.provider != LLMProvider::Bedrock && m.provider != LLMProvider::Unknown); if let Some(ref mut cli) = self.models_by_feature.cli_agent { - cli.choices.retain(|m| m.provider != LLMProvider::Bedrock && m.provider != LLMProvider::Unknown); + cli.choices.retain(|m| { + m.provider != LLMProvider::Bedrock && m.provider != LLMProvider::Unknown + }); } let settings = AISettings::as_ref(ctx); @@ -1916,52 +1902,6 @@ fn openai_model_context_window(model: &OpenAIModelConfig) -> LLMContextWindow { } } -/// Builds synthetic [`LLMInfo`]s from the user's persisted custom endpoints. -/// -/// One entry per `CustomEndpointModel`. The display label is the **alias** when present, -/// falling back to the raw model name. The `id` is the model's `config_key`, which is -/// also what flows out to `Request.Settings.custom_model_providers` so the server can map -/// a `ModelConfig.{base,coding,cli_agent,computer_use_agent}` selection back to the -/// user-provided endpoint. -/// -/// Endpoints with empty URL or API key, and models with empty name or config_key, are -/// skipped — they shouldn't surface in the picker until the user finishes configuring them. -fn build_custom_llm_infos(keys: &ai::api_keys::ApiKeys) -> Vec { - keys.custom_endpoints - .iter() - .filter(|ep| !ep.url.trim().is_empty() && !ep.api_key.is_empty()) - .flat_map(|endpoint| { - endpoint - .models - .iter() - .filter(|m| !m.name.trim().is_empty() && !m.config_key.is_empty()) - .map(move |model| custom_llm_info_from(endpoint, model)) - }) - .collect() -} - -fn custom_llm_info_from(endpoint: &CustomEndpoint, model: &CustomEndpointModel) -> LLMInfo { - let label = model.display_label().to_owned(); - LLMInfo { - display_name: label.clone(), - base_model_name: label, - id: model.config_key.clone().into(), - reasoning_level: None, - usage_metadata: LLMUsageMetadata { - request_multiplier: 1, - credit_multiplier: None, - }, - description: Some(format!("Custom · {}", endpoint.name)), - disable_reason: None, - vision_supported: true, - spec: None, - provider: LLMProvider::Unknown, - host_configs: HashMap::new(), - discount_percentage: None, - context_window: LLMContextWindow::default(), - } -} - /// Fetches model metadata from LiteLLM's `/model/info` endpoint which returns rich /// metadata including accurate context window sizes, output token limits, and /// capability flags (vision, function calling). @@ -2144,10 +2084,9 @@ async fn fetch_from_openai_models( m, &["max_input_tokens", "input_token_limit", "max_prompt_tokens"], ); - let context_size = - u32_from_any(m, &["max_model_len", "context_window", "token_size"]) - .or(max_input_tokens) - .unwrap_or(200_000); + let context_size = u32_from_any(m, &["max_model_len", "context_window", "token_size"]) + .or(max_input_tokens) + .unwrap_or(200_000); let max_output_tokens = u32_from_any( m, &[ diff --git a/app/src/ai/llms_tests.rs b/app/src/ai/llms_tests.rs index 77858fa8..a3a11304 100644 --- a/app/src/ai/llms_tests.rs +++ b/app/src/ai/llms_tests.rs @@ -138,458 +138,3 @@ fn llm_info_round_trip_serializes_and_deserializes() { assert_eq!(info, round_tripped); } - -// -- build_custom_llm_infos / display label tests -- - -fn endpoint( - name: &str, - url: &str, - api_key: &str, - models: Vec, -) -> CustomEndpoint { - CustomEndpoint { - name: name.into(), - url: url.into(), - api_key: api_key.into(), - models, - } -} - -fn model(name: &str, alias: Option<&str>, config_key: &str) -> CustomEndpointModel { - CustomEndpointModel { - name: name.into(), - alias: alias.map(|s| s.into()), - config_key: config_key.into(), - } -} - -#[test] -fn custom_llm_infos_built_from_endpoints() { - let keys = ai::api_keys::ApiKeys { - custom_endpoints: vec![endpoint( - "My Endpoint", - "https://x.io", - "k", - vec![ - model("gpt-4", Some("fast"), "uuid-1"), - model("llama", None, "uuid-2"), - ], - )], - ..Default::default() - }; - let infos = build_custom_llm_infos(&keys); - assert_eq!(infos.len(), 2); - assert_eq!(infos[0].display_name, "fast"); - assert_eq!(infos[0].id.as_str(), "uuid-1"); - assert_eq!( - infos[0].description.as_deref(), - Some("Custom · My Endpoint") - ); - assert_eq!(infos[1].display_name, "llama"); - assert_eq!(infos[1].id.as_str(), "uuid-2"); -} - -#[test] -fn custom_llm_display_name_uses_alias_when_present() { - let keys = ai::api_keys::ApiKeys { - custom_endpoints: vec![endpoint( - "ep", - "https://a.io", - "k", - vec![model("raw-name", Some("My Alias"), "uuid-a")], - )], - ..Default::default() - }; - let infos = build_custom_llm_infos(&keys); - assert_eq!(infos[0].display_name, "My Alias"); -} - -#[test] -fn custom_llm_display_name_falls_back_to_name_when_alias_missing() { - let keys = ai::api_keys::ApiKeys { - custom_endpoints: vec![endpoint( - "ep", - "https://a.io", - "k", - vec![model("raw-name", None, "uuid-a")], - )], - ..Default::default() - }; - let infos = build_custom_llm_infos(&keys); - assert_eq!(infos[0].display_name, "raw-name"); -} - -#[test] -fn custom_endpoint_usage_display_label_resolves_alias_name_and_generic_fallback() { - let keys = ai::api_keys::ApiKeys { - custom_endpoints: vec![endpoint( - "ep", - "https://a.io", - "k", - vec![ - model("raw-alias", Some("Alias"), "uuid-alias"), - model("raw-name", None, "uuid-name"), - model("raw~name", None, "uuid-tilde-name"), - ], - )], - ..Default::default() - }; - let preferences = LLMPreferences { - models_by_feature: ModelsByFeature::default(), - last_update: None, - base_llm_for_terminal_view: HashMap::new(), - custom_llms: build_custom_llm_infos(&keys), - custom_model_routers: Vec::new(), - openai_provider_routing: HashMap::new(), - fetched_openai_models: Vec::new(), - }; - - assert_eq!( - preferences.custom_endpoint_usage_display_label("uuid-alias"), - "Alias" - ); - assert_eq!( - preferences.custom_endpoint_usage_display_label("uuid-name"), - "raw-name" - ); - assert_eq!( - preferences.custom_endpoint_usage_display_label("uuid-tilde-name"), - "raw~name" - ); - assert_eq!( - preferences.custom_endpoint_usage_display_label("unknown"), - CUSTOM_ENDPOINT_USAGE_FALLBACK_LABEL - ); -} - -#[cfg(not(target_family = "wasm"))] -fn empty_llm_preferences_for_provider_tests() -> LLMPreferences { - LLMPreferences { - models_by_feature: ModelsByFeature::default(), - last_update: None, - base_llm_for_terminal_view: HashMap::new(), - custom_llms: Vec::new(), - custom_model_routers: Vec::new(), - openai_provider_routing: HashMap::new(), - fetched_openai_models: Vec::new(), - } -} - -#[cfg(not(target_family = "wasm"))] -fn openai_model( - model_id: &str, - display_name: &str, - context_size: u32, - max_input_tokens: Option, - max_output_tokens: Option, - vision_supported: bool, -) -> OpenAIModelConfig { - OpenAIModelConfig { - model_id: model_id.to_string(), - display_name: display_name.to_string(), - vision_supported, - context_size, - max_input_tokens, - max_output_tokens, - provider: Some("openai".to_string()), - } -} - -#[test] -#[cfg(not(target_family = "wasm"))] -fn openai_model_config_accepts_legacy_and_endpoint_field_names() { - let model: OpenAIModelConfig = toml::from_str( - r#" -model_id = "provider/custom-model" -display_name = "Custom Model" -vision_support = true -token_size = 123456 -max_input_tokens = 111111 -max_tokens = 8192 -provider = "openai" -"#, - ) - .expect("model config should parse"); - - assert_eq!(model.model_id, "provider/custom-model"); - assert!(model.vision_supported); - assert_eq!(model.context_size, 123_456); - assert_eq!(model.max_input_tokens, Some(111_111)); - assert_eq!(model.max_output_tokens, Some(8_192)); -} - -#[test] -#[cfg(not(target_family = "wasm"))] -fn inject_openai_models_uses_persisted_model_metadata_and_routing() { - App::test((), |mut app| async move { - initialize_settings_for_tests(&mut app); - let mut preferences = empty_llm_preferences_for_provider_tests(); - - app.update(|ctx| { - AISettings::handle(ctx).update(ctx, |settings, ctx| { - settings.openai_enabled.set_value(true, ctx).unwrap(); - settings - .openai_base_url - .set_value("https://litellm.example/v1".to_string(), ctx) - .unwrap(); - settings - .openai_api_key - .set_value("test-key".to_string(), ctx) - .unwrap(); - settings - .openai_models - .set_value( - vec![openai_model( - "provider/custom-model", - "Custom Model", - 200_000, - Some(128_000), - Some(8_192), - true, - )], - ctx, - ) - .unwrap(); - }); - - preferences.inject_openai_models(ctx); - - let model = preferences - .models_by_feature - .agent_mode - .choices - .iter() - .find(|model| model.id.as_str() == "provider/custom-model") - .expect("configured model should be injected"); - assert_eq!(model.provider, LLMProvider::LiteLLM); - assert_eq!(model.description.as_deref(), Some("LiteLLM")); - assert!(model.vision_supported); - assert_eq!(model.context_window.default_max, 128_000); - assert_eq!(model.context_window.max, 128_000); - - let client_config = preferences - .openai_client_config_for_model("provider/custom-model") - .expect("configured model should have routing"); - assert_eq!(client_config.base_url, "https://litellm.example/v1"); - assert_eq!(client_config.api_key.as_deref(), Some("test-key")); - assert_eq!(client_config.max_input_tokens, Some(128_000)); - assert_eq!(client_config.max_output_tokens, Some(8_192)); - }); - }); -} - -#[test] -#[cfg(not(target_family = "wasm"))] -fn inject_openai_models_uses_multi_provider_models() { - App::test((), |mut app| async move { - initialize_settings_for_tests(&mut app); - let mut preferences = empty_llm_preferences_for_provider_tests(); - - app.update(|ctx| { - AISettings::handle(ctx).update(ctx, |settings, ctx| { - settings.openai_enabled.set_value(true, ctx).unwrap(); - settings - .openai_providers - .set_value( - vec![OpenAIProviderConfig { - name: "Ollama".to_string(), - base_url: "http://localhost:11434/v1".to_string(), - api_key: None, - models: vec![openai_model( - "llama3.2", - "Llama 3.2", - 64_000, - None, - Some(4_096), - false, - )], - }], - ctx, - ) - .unwrap(); - }); - - preferences.inject_openai_models(ctx); - - let model = preferences - .models_by_feature - .agent_mode - .choices - .iter() - .find(|model| model.id.as_str() == "llama3.2") - .expect("provider model should be injected"); - assert_eq!(model.description.as_deref(), Some("Ollama")); - assert!(!model.vision_supported); - assert_eq!(model.context_window.default_max, 64_000); - - let client_config = preferences - .openai_client_config_for_model("llama3.2") - .expect("provider model should have routing"); - assert_eq!(client_config.base_url, "http://localhost:11434/v1"); - assert_eq!(client_config.max_input_tokens, Some(64_000)); - assert_eq!(client_config.max_output_tokens, Some(4_096)); - }); - }); -} - -#[test] -fn custom_llm_infos_skip_endpoints_with_empty_api_key() { - let keys = ai::api_keys::ApiKeys { - custom_endpoints: vec![ - endpoint("bad", "https://a.io", "", vec![model("m", None, "uuid-x")]), - endpoint( - "good", - "https://b.io", - "k", - vec![model("m", None, "uuid-y")], - ), - ], - ..Default::default() - }; - let infos = build_custom_llm_infos(&keys); - assert_eq!(infos.len(), 1); - assert_eq!(infos[0].id.as_str(), "uuid-y"); -} - -#[test] -fn custom_llm_infos_skip_models_without_config_key() { - let keys = ai::api_keys::ApiKeys { - custom_endpoints: vec![endpoint( - "ep", - "https://a.io", - "k", - vec![ - model("unconfigured", None, ""), - model("ready", None, "uuid-a"), - ], - )], - ..Default::default() - }; - let infos = build_custom_llm_infos(&keys); - assert_eq!(infos.len(), 1); - assert_eq!(infos[0].display_name, "ready"); -} - -#[test] -fn removing_model_row_purges_from_custom_llms() { - let before = ai::api_keys::ApiKeys { - custom_endpoints: vec![endpoint( - "ep", - "https://a.io", - "k", - vec![model("a", None, "uuid-a"), model("b", None, "uuid-b")], - )], - ..Default::default() - }; - assert_eq!(build_custom_llm_infos(&before).len(), 2); - - let after = ai::api_keys::ApiKeys { - custom_endpoints: vec![endpoint( - "ep", - "https://a.io", - "k", - vec![model("b", None, "uuid-b")], - )], - ..Default::default() - }; - let infos = build_custom_llm_infos(&after); - assert_eq!(infos.len(), 1); - assert_eq!(infos[0].id.as_str(), "uuid-b"); - assert!(infos.iter().all(|i| i.id.as_str() != "uuid-a")); -} - -#[test] -fn removing_endpoint_purges_all_its_models_from_custom_llms() { - let before = ai::api_keys::ApiKeys { - custom_endpoints: vec![ - endpoint( - "keep", - "https://a.io", - "k", - vec![model("k1", None, "uuid-k1")], - ), - endpoint( - "goner", - "https://b.io", - "k", - vec![model("g1", None, "uuid-g1"), model("g2", None, "uuid-g2")], - ), - ], - ..Default::default() - }; - assert_eq!(build_custom_llm_infos(&before).len(), 3); - - let after = ai::api_keys::ApiKeys { - custom_endpoints: vec![endpoint( - "keep", - "https://a.io", - "k", - vec![model("k1", None, "uuid-k1")], - )], - ..Default::default() - }; - let infos = build_custom_llm_infos(&after); - assert_eq!(infos.len(), 1); - assert_eq!(infos[0].id.as_str(), "uuid-k1"); -} - -#[test] -fn reconcile_preserves_custom_models_saved_on_execution_profile() { - App::test((), |mut app| async move { - let _custom_inference_flag = FeatureFlag::CustomInferenceEndpoints.override_enabled(true); - - initialize_settings_for_tests(&mut app); - app.add_singleton_model(|_| ServerApiProvider::new_for_test()); - app.add_singleton_model(|_| AuthStateProvider::new_for_test()); - app.add_singleton_model(AuthManager::new_for_test); - app.add_singleton_model(|_| NetworkStatus::new()); - app.add_singleton_model(UserWorkspaces::default_mock); - app.add_singleton_model(CloudModel::mock); - app.add_singleton_model(TeamTesterStatus::mock); - app.add_singleton_model(SyncQueue::mock); - app.add_singleton_model(UpdateManager::mock); - app.add_singleton_model(|_| TemplatableMCPServerManager::default()); - - let profiles_model = app.add_singleton_model(|ctx| { - AIExecutionProfilesModel::new(&LaunchMode::new_for_unit_test(), ctx) - }); - let llm_preferences = app.add_singleton_model(LLMPreferences::new); - - let custom_model_id = LLMId::from("custom-model-config-key"); - ApiKeyManager::handle(&app).update(&mut app, |api_key_manager, ctx| { - api_key_manager.add_custom_endpoint( - "local".to_string(), - "https://example.com/v1".to_string(), - "test-key".to_string(), - vec![( - "custom-model".to_string(), - Some("Custom Model".to_string()), - Some(custom_model_id.to_string()), - )], - ctx, - ); - }); - - let default_profile_id = - profiles_model.read(&app, |profiles, _| profiles.default_profile_id()); - profiles_model.update(&mut app, |profiles, ctx| { - profiles.set_base_model(default_profile_id, Some(custom_model_id.clone()), ctx); - profiles.set_coding_model(default_profile_id, Some(custom_model_id.clone()), ctx); - profiles.set_cli_agent_model(default_profile_id, Some(custom_model_id.clone()), ctx); - }); - - llm_preferences.update(&mut app, |preferences, ctx| { - preferences.update_feature_model_choices(Ok(ModelsByFeature::default()), ctx); - }); - - profiles_model.read(&app, |profiles, ctx| { - let profile = profiles.default_profile(ctx); - assert_eq!(profile.data().base_model.as_ref(), Some(&custom_model_id)); - assert_eq!(profile.data().coding_model.as_ref(), Some(&custom_model_id)); - assert_eq!( - profile.data().cli_agent_model.as_ref(), - Some(&custom_model_id) - ); - }); - }); -} diff --git a/app/src/ai/request_usage_model.rs b/app/src/ai/request_usage_model.rs index 05b2a2e0..6603adec 100644 --- a/app/src/ai/request_usage_model.rs +++ b/app/src/ai/request_usage_model.rs @@ -1,6 +1,5 @@ use std::sync::Arc; -use ai::api_keys::ApiKeyManager; use chrono::{DateTime, Local, Utc}; use galaxy_core::user_preferences::GetUserPreferences as _; pub use galaxy_graphql::billing::BonusGrantType; @@ -401,8 +400,6 @@ impl AIRequestUsageModel { /// 4. user's team plan has pay-as-you-go enabled (enterprise only) /// 5. user's team has enterprise bonus grants auto-reload enabled (enterprise only) /// 6. user's team has self-serve auto-reload enabled within its monthly spend limit - /// 7. user has BYOK enabled and has either provided at least one API key or - /// connected a Grok subscription /// Use this method as the starting point for AI availability checking. pub fn has_any_ai_remaining(&self, ctx: &AppContext) -> bool { let current_workspace = UserWorkspaces::as_ref(ctx).current_workspace(); @@ -435,18 +432,12 @@ impl AIRequestUsageModel { .is_some_and(|price| !workspace.would_addon_purchase_reach_limit(price)) }); - // If you have provided your own API key or connected a Grok - // subscription, it doesn't matter if you are out of warp-provided requests. - let has_byo_credentials = UserWorkspaces::as_ref(ctx).is_byo_api_key_enabled(ctx) - && ApiKeyManager::as_ref(ctx).has_any_key(); - has_base_plan_ai_requests || (user_bonus_credits || workspace_bonus_credits) || workspace_has_overages || is_payg_enabled || is_enterprise_auto_reload_enabled || is_self_serve_auto_reload_enabled - || has_byo_credentials } pub fn requests_used(&self) -> usize { diff --git a/app/src/ai/request_usage_model_tests.rs b/app/src/ai/request_usage_model_tests.rs index 95246eec..c0f09b62 100644 --- a/app/src/ai/request_usage_model_tests.rs +++ b/app/src/ai/request_usage_model_tests.rs @@ -1,6 +1,6 @@ use std::sync::Arc; -use ai::api_keys::{ApiKeyManager, GrokTokens}; +use ai::api_keys::ApiKeyManager; use chrono::Duration; use galaxy_core::features::FeatureFlag; use galaxy_graphql::billing::{AddonCreditsOption, OveragesPricing, PricingInfo}; @@ -14,8 +14,8 @@ use crate::server::server_api::workspace::MockWorkspaceClient; use crate::server::server_api::ServerApiProvider; use crate::workspaces::user_workspaces::UserWorkspaces; use crate::workspaces::workspace::{ - AiOverages, ByoApiKeyPolicy, CustomerType, EnterpriseCreditsAutoReloadPolicy, - EnterprisePayAsYouGoPolicy, PurchaseAddOnCreditsPolicy, Workspace, WorkspaceUid, + AiOverages, CustomerType, EnterpriseCreditsAutoReloadPolicy, EnterprisePayAsYouGoPolicy, + PurchaseAddOnCreditsPolicy, Workspace, WorkspaceUid, }; fn create_test_workspace() -> (WorkspaceUid, Workspace) { @@ -681,183 +681,6 @@ fn test_has_any_ai_remaining_false_both_payg_and_autoreload_disabled() { }); } -#[test] -fn test_has_any_ai_remaining_true_with_byok_enabled_and_key_provided() { - App::test((), |mut app| async move { - // Create a workspace with BYOK (Bring Your Own Key) enabled. - let (_uid, mut workspace) = create_test_workspace(); - workspace.billing_metadata.tier.byo_api_key_policy = - Some(ByoApiKeyPolicy { enabled: true }); - - add_user_workspaces_with_workspace(&mut app, workspace); - let request_usage_model = add_request_usage_model(&mut app); - - ApiKeyManager::handle(&app).update(&mut app, |manager, ctx| { - manager.set_openai_key(Some("test-key".to_string()), ctx); - }); - - request_usage_model.update(&mut app, |model, ctx| { - // No standard requests remaining, no bonus credits. - model.request_limit_info = RequestLimitInfo::new_for_test(10, 10); - model.bonus_grants.clear(); - - assert!( - model.has_any_ai_remaining(ctx), - "expected has_any_ai_remaining to be true when BYOK is enabled and a key is provided", - ); - }); - }); -} - -#[test] -fn test_has_any_ai_remaining_false_with_byok_enabled_but_no_key() { - App::test((), |mut app| async move { - // Create a workspace with BYOK enabled but no key provided. - let (_uid, mut workspace) = create_test_workspace(); - workspace.billing_metadata.tier.byo_api_key_policy = - Some(ByoApiKeyPolicy { enabled: true }); - - add_user_workspaces_with_workspace(&mut app, workspace); - let request_usage_model = add_request_usage_model(&mut app); - - request_usage_model.update(&mut app, |model, ctx| { - // No standard requests remaining, no bonus credits. - model.request_limit_info = RequestLimitInfo::new_for_test(10, 10); - model.bonus_grants.clear(); - - assert!( - !model.has_any_ai_remaining(ctx), - "expected has_any_ai_remaining to be false when BYOK is enabled but no key is provided", - ); - }); - }); -} - -#[test] -fn test_has_any_ai_remaining_true_with_grok_subscription_connected() { - App::test((), |mut app| async move { - // Workspace with BYO enabled — the policy a connected Grok - // subscription's OAuth token rides on. - let (_uid, mut workspace) = create_test_workspace(); - workspace.billing_metadata.tier.byo_api_key_policy = - Some(ByoApiKeyPolicy { enabled: true }); - - add_user_workspaces_with_workspace(&mut app, workspace); - let request_usage_model = add_request_usage_model(&mut app); - - // Connect a Grok subscription but provide no pasted API key. - ApiKeyManager::handle(&app).update(&mut app, |manager, ctx| { - manager.set_grok_tokens( - Some(GrokTokens { - access_token: "grok-test-token".to_string(), - ..Default::default() - }), - ctx, - ); - }); - - request_usage_model.update(&mut app, |model, ctx| { - // No standard requests remaining, no bonus credits. - model.request_limit_info = RequestLimitInfo::new_for_test(10, 10); - model.bonus_grants.clear(); - - assert!( - model.has_any_ai_remaining(ctx), - "expected has_any_ai_remaining to be true when a Grok subscription is connected and BYO is enabled", - ); - }); - }); -} - -#[test] -fn test_has_any_ai_remaining_false_with_grok_subscription_but_byo_disabled() { - App::test((), |mut app| async move { - // No BYO policy: the Grok token can't be sent, so it must not count as - // available AI. - let (_uid, workspace) = create_test_workspace(); - - add_user_workspaces_with_workspace(&mut app, workspace); - let request_usage_model = add_request_usage_model(&mut app); - - ApiKeyManager::handle(&app).update(&mut app, |manager, ctx| { - manager.set_grok_tokens( - Some(GrokTokens { - access_token: "grok-test-token".to_string(), - ..Default::default() - }), - ctx, - ); - }); - - request_usage_model.update(&mut app, |model, ctx| { - model.request_limit_info = RequestLimitInfo::new_for_test(10, 10); - model.bonus_grants.clear(); - - assert!( - !model.has_any_ai_remaining(ctx), - "expected has_any_ai_remaining to be false when a Grok subscription is connected but BYO is disabled", - ); - }); - }); -} - -#[test] -fn test_has_any_ai_remaining_true_with_byo_key_and_no_workspace() { - App::test((), |mut app| async move { - let _guard = FeatureFlag::SoloUserByok.override_enabled(true); - - // No workspace — user is not on a team. - app.add_singleton_model(UserWorkspaces::default_mock); - let request_usage_model = add_request_usage_model(&mut app); - - ApiKeyManager::handle(&app).update(&mut app, |manager, ctx| { - manager.set_openai_key(Some("test-key".to_string()), ctx); - }); - - request_usage_model.update(&mut app, |model, ctx| { - // No standard requests remaining, no bonus credits. - model.request_limit_info = RequestLimitInfo::new_for_test(10, 10); - model.bonus_grants.clear(); - - assert!( - model.has_any_ai_remaining(ctx), - "expected has_any_ai_remaining to be true when user has a BYO key but no workspace", - ); - }); - }); -} - -#[test] -fn test_byo_api_key_disabled_for_anonymous_firebase_user() { - App::test((), |mut app| async move { - let _guard = FeatureFlag::SoloUserByok.override_enabled(true); - - app.add_singleton_model(UserWorkspaces::default_mock); - let request_usage_model = add_request_usage_model_for_anonymous_users(&mut app); - - ApiKeyManager::handle(&app).update(&mut app, |manager, ctx| { - manager.set_openai_key(Some("test-key".to_string()), ctx); - }); - - app.read(|ctx| { - assert!( - !UserWorkspaces::as_ref(ctx).is_byo_api_key_enabled(ctx), - "expected is_byo_api_key_enabled to be false for anonymous Firebase users even with SoloUserByok enabled", - ); - }); - - request_usage_model.update(&mut app, |model, ctx| { - model.request_limit_info = RequestLimitInfo::new_for_test(10, 10); - model.bonus_grants.clear(); - - assert!( - !model.has_any_ai_remaining(ctx), - "expected has_any_ai_remaining to be false for anonymous Firebase user even with BYO key and SoloUserByok enabled", - ); - }); - }); -} - #[test] fn test_has_any_ai_remaining_false_with_only_ambient_bonus_credits() { App::test((), |mut app| async move { diff --git a/app/src/lib.rs b/app/src/lib.rs index 0796e112..fef00374 100644 --- a/app/src/lib.rs +++ b/app/src/lib.rs @@ -1498,23 +1498,6 @@ pub(crate) fn initialize_app( if FeatureFlag::GeminiEnterprise.is_enabled() { manager.subscribe_to_geap_settings_changes(ctx); } - // The Grok subscription refresher (`ai::grok_subscription`) has no - // visibility into workspace policy, so wire the BYO API key policy in - // here. The initial value resumes proactive refresh of any tokens - // restored from secure storage; TeamsChanged keeps the policy aligned - // as team data loads or the workspace changes. - #[cfg(not(target_family = "wasm"))] - if FeatureFlag::SuperGrok.is_enabled() { - use crate::workspaces::user_workspaces::UserWorkspacesEvent; - ctx.subscribe_to_model(&UserWorkspaces::handle(ctx), |manager, _, event, ctx| { - if matches!(event, UserWorkspacesEvent::TeamsChanged) { - let allowed = UserWorkspaces::as_ref(ctx).is_byo_api_key_enabled(ctx); - manager.set_grok_refresh_allowed(allowed, ctx); - } - }); - let allowed = UserWorkspaces::as_ref(ctx).is_byo_api_key_enabled(ctx); - manager.set_grok_refresh_allowed(allowed, ctx); - } manager }); @@ -2393,7 +2376,6 @@ pub(crate) fn app_callbacks( if let Some(initialization) = tracing_initialization.as_mut() { initialization.shutdown(); } - })), on_should_close_window: Some(Box::new(move |window_id, ctx| { let general_settings = GeneralSettings::as_ref(ctx); diff --git a/app/src/settings/ai.rs b/app/src/settings/ai.rs index 2913cc4e..96978dd5 100644 --- a/app/src/settings/ai.rs +++ b/app/src/settings/ai.rs @@ -1557,18 +1557,6 @@ define_settings_group!(AISettings, settings: [ - // Whether or not the user has enabled fallback to Warp credits for user-provided models. - can_use_warp_credits_for_fallback: CanUseWarpCreditsForFallback { - type: bool, - default: false, - supported_platforms: SupportedPlatforms::ALL, - sync_to_cloud: SyncToCloud::Never, - private: false, - storage_key: "CanUseWarpCreditsWithByok", - toml_path: "cloud_platform.third_party_api_keys.can_use_warp_credits_with_byok", - description: "Whether Warp credits can be used as a fallback for user-provided models.", - } - should_render_use_agent_footer_for_user_commands: ShouldRenderUseAgentToolbarForUserCommands { type: bool, default: true, @@ -1947,8 +1935,7 @@ impl AISettings { // Galaxy does not require Warp authentication for AI. // AI is enabled as long as the user hasn't explicitly disabled it // and there's no org policy blocking it. - *self.is_any_ai_enabled - && !self.is_ai_disabled_due_to_remote_session_org_policy(app) + *self.is_any_ai_enabled && !self.is_ai_disabled_due_to_remote_session_org_policy(app) } pub fn default_session_mode(&self, app: &AppContext) -> DefaultSessionMode { diff --git a/app/src/settings_view/ai_page.rs b/app/src/settings_view/ai_page.rs index 0d993738..2f5f423f 100644 --- a/app/src/settings_view/ai_page.rs +++ b/app/src/settings_view/ai_page.rs @@ -1,6 +1,4 @@ -use ::ai::api_keys::{ApiKeyManager, ApiKeyManagerEvent, ApiKeys}; -#[cfg(not(target_family = "wasm"))] -use ::ai::grok_subscription::oauth::{self, ManualCodeExchange}; +use ::ai::api_keys::ApiKeyManager; use chrono::{DateTime, Local}; use enum_iterator::all; use galaxy_core::channel::ChannelState; @@ -85,15 +83,14 @@ use crate::settings::{ AIAutoDetectionEnabled, AICommandDenylist, AISettingsChangedEvent, AgentModeCodingPermissionsType, AgentModeCommandExecutionDenylist, AgentModeCommandExecutionPredicate, AgentModeQuerySuggestionsEnabled, BedrockAutoLogin, - BedrockEnabled, CanUseWarpCreditsForFallback, CodeSettings, CodebaseContextEnabled, - CrosscheckEnabled, FileBasedMcpEnabled, GitOperationsAutogenEnabled, - IncludeAgentCommandsInHistory, InputSettings, IntelligentAutosuggestionsEnabled, - LongRunningCommandSubmissionMode, MemoryEnabled, NLDInTerminalEnabled, - NaturalLanguageAutosuggestionsEnabled, OpenAIEnabled, OrchestrationMessageDisplayMode, - PromptSubmissionMode, RuleSuggestionsEnabled, SharedBlockTitleGenerationEnabled, - ShouldRenderCLIAgentToolbar, ShouldRenderUseAgentToolbarForUserCommands, ShowAgentTips, - ShowConversationHistory, ShowHintText, ThinkingDisplayMode, VoiceInputEnabled, - WarpDriveContextEnabled, + BedrockEnabled, CodeSettings, CodebaseContextEnabled, CrosscheckEnabled, FileBasedMcpEnabled, + GitOperationsAutogenEnabled, IncludeAgentCommandsInHistory, InputSettings, + IntelligentAutosuggestionsEnabled, LongRunningCommandSubmissionMode, MemoryEnabled, + NLDInTerminalEnabled, NaturalLanguageAutosuggestionsEnabled, OpenAIEnabled, + OrchestrationMessageDisplayMode, PromptSubmissionMode, RuleSuggestionsEnabled, + SharedBlockTitleGenerationEnabled, ShouldRenderCLIAgentToolbar, + ShouldRenderUseAgentToolbarForUserCommands, ShowAgentTips, ShowConversationHistory, + ShowHintText, ThinkingDisplayMode, VoiceInputEnabled, WarpDriveContextEnabled, }; use crate::terminal::session_settings::{SessionSettings, SessionSettingsChangedEvent}; use crate::terminal::CLIAgent; @@ -111,7 +108,7 @@ use crate::workspaces::user_workspaces::UserWorkspacesEvent; /// When `None`, the page shows all widgets (legacy/full view). #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum AISubpage { - /// The main "WarpAgent" page: global AI toggle + Active AI + Input + Other sections. + /// The main Galaxy Agent page: global AI toggle + Active AI + Input + Other sections. WarpAgent, /// Agent profiles and permissions. Profiles, @@ -187,11 +184,6 @@ const SHARED_BLOCK_TITLE_GENERATION_DESCRIPTION: &str = const GIT_OPERATIONS_AUTOGEN_DESCRIPTION: &str = "Let AI generate commit messages and pull request titles and descriptions."; const WISPR_FLOW_URL: &str = "https://wisprflow.ai/"; -const CUSTOM_INFERENCE_LEARN_MORE_URL: &str = - "https://docs.warp.dev/agent-platform/inference/custom-inference-endpoint/"; -const CUSTOM_INFERENCE_TERMS_URL: &str = "https://www.warp.dev/legal/terms-of-service"; -const CUSTOM_INFERENCE_INFO_TOOLTIP_MAX_WIDTH: f32 = 320.; -const CUSTOM_ENDPOINT_MODAL_MAX_HEIGHT_PERCENTAGE: f32 = 0.8; pub fn init_actions_from_parent_view( app: &mut AppContext, @@ -571,18 +563,6 @@ pub fn init_actions_from_parent_view( ); ToggleSettingActionPair::add_toggle_setting_action_pairs_as_bindings( vec![ - ToggleSettingActionPair::new( - "Warp credit fallback", - builder(SettingsAction::AI( - AISettingsPageAction::ToggleCanUseWarpCreditsForFallback, - )), - &(context.clone() & id!(flags::IS_ANY_AI_ENABLED)), - flags::WARP_CREDIT_FALLBACK_FLAG, - ) - .with_group(bindings::BindingGroup::WarpAi) - .is_supported_on_current_platform( - UserWorkspaces::as_ref(app).is_byo_api_key_enabled(app), - ), ToggleSettingActionPair::new( "auto show or hide Rich Input based on agent status", builder(SettingsAction::AI( @@ -703,21 +683,6 @@ pub struct AISettingsPageView { router_views: Vec>, #[cfg(feature = "local_fs")] add_router_button: ViewHandle, - - // Prompt offering to switch the default Agent Mode model after a BYO key or - // custom endpoint is saved while the default isn't backed by a credential. - set_default_model_modal: ModalViewState>, - // Snapshot of the provider keys from the last `KeysUpdated`, used to detect a - // newly added key and prompt the user to switch their default model. - last_seen_provider_keys: ApiKeys, - - // In-flight fallback exchange for a pasted SuperGrok authorization code. - // This stores only the PKCE verifier clone needed by the manual path while - // `OauthAttempt::finish` owns the full loopback attempt. - #[cfg(not(target_family = "wasm"))] - grok_oauth_attempt: Option, - #[cfg(not(target_family = "wasm"))] - grok_code_editor: ViewHandle, } impl AISettingsPageView { @@ -1128,15 +1093,11 @@ impl AISettingsPageView { }, ); - // Refresh model dropdowns when BYO API keys update so key icons reflect latest state. + // Refresh model dropdowns when Bedrock credentials update so key icons reflect latest state. ctx.subscribe_to_model(&ApiKeyManager::handle(ctx), |me, _model, _event, ctx| { Self::refresh_base_model_menu(&me.base_model_dropdown, ctx); Self::refresh_coding_model_menu(&me.coding_model_dropdown, ctx); me.sync_context_window_editor(ctx, false); - // Driving the prompt off the key-store update (rather than the editor's - // blur/Enter) means it fires reliably however the key was committed — - // clicking outside the field, pressing Enter, or tabbing away. - me.maybe_prompt_for_newly_added_provider_key(ctx); ctx.notify(); }); @@ -1704,38 +1665,6 @@ impl AISettingsPageView { button.set_disabled(!is_any_ai_enabled, ctx); }); - let set_default_model_modal_body = ctx.add_typed_action_view(SetDefaultModelModalBody::new); - ctx.subscribe_to_view(&set_default_model_modal_body, |me, _, event, ctx| { - me.handle_set_default_model_modal_event(event, ctx); - }); - let set_default_model_modal_view = ctx.add_typed_action_view(|ctx| { - Modal::new( - Some("Change your default model?".to_string()), - set_default_model_modal_body.clone(), - ctx, - ) - .with_modal_style(UiComponentStyles { - width: Some(480.), - height: Some(380.), - ..Default::default() - }) - .with_body_style(UiComponentStyles { - height: Some(300.), - ..Default::default() - }) - .with_background_opacity(100) - .with_dismiss_on_click() - .with_dismiss_keystroke(Keystroke::parse("escape").unwrap()) - }); - ctx.subscribe_to_view( - &set_default_model_modal_view, - |me, _, event, ctx| match event { - ModalEvent::Close => me.hide_set_default_model_modal(ctx), - }, - ); - let set_default_model_modal = ModalViewState::new(set_default_model_modal_view); - let last_seen_provider_keys = ApiKeyManager::as_ref(ctx).keys().clone(); - let agent_toolbar_inline_editor = ctx.add_typed_action_view(|ctx| { AgentToolbarInlineEditor::new(AgentToolbarEditorMode::AgentView, ctx) }); @@ -1776,29 +1705,6 @@ impl AISettingsPageView { dropdown }); - #[cfg(not(target_family = "wasm"))] - let grok_code_editor = Self::create_grok_code_editor(ctx); - #[cfg(not(target_family = "wasm"))] - ctx.subscribe_to_view(&grok_code_editor, |me, _, event, ctx| { - if matches!(event, EditorEvent::Enter | EditorEvent::Paste) { - let code = me.grok_code_editor.as_ref(ctx).buffer_text(ctx); - me.submit_grok_code(code, ctx); - } - }); - // Keep the snapshotted editor text colors in sync with theme changes, - // like the API key editors above. - #[cfg(not(target_family = "wasm"))] - { - let grok_code_editor = grok_code_editor.clone(); - ctx.subscribe_to_model(&Appearance::handle(ctx), move |_, _, event, ctx| { - if let AppearanceEvent::ThemeChanged = event { - let colors = editor_text_colors(Appearance::as_ref(ctx)); - grok_code_editor.update(ctx, move |editor, ctx| { - editor.set_text_colors(colors, ctx); - }); - } - }); - } // Subscribe to GalaxyConfig to refresh router views when files change. #[cfg(feature = "local_fs")] ctx.subscribe_to_model( @@ -1867,12 +1773,6 @@ impl AISettingsPageView { router_views, #[cfg(feature = "local_fs")] add_router_button, - set_default_model_modal, - last_seen_provider_keys, - #[cfg(not(target_family = "wasm"))] - grok_oauth_attempt: None, - #[cfg(not(target_family = "wasm"))] - grok_code_editor, } } @@ -1890,401 +1790,7 @@ impl AISettingsPageView { } pub fn get_modal_content(&self, _app: &AppContext) -> Option> { - if self.set_default_model_modal.is_open() { - Some(self.set_default_model_modal.render()) - } else { - None - } - } - - fn handle_set_default_model_modal_event( - &mut self, - event: &SetDefaultModelModalBodyEvent, - ctx: &mut ViewContext, - ) { - match event { - SetDefaultModelModalBodyEvent::Close => self.hide_set_default_model_modal(ctx), - SetDefaultModelModalBodyEvent::SetDefault(id) => { - // Mirror `AISettingsPageAction::SetBaseModel`: set the active - // profile's base model and clear any stale context-window limit. - AIExecutionProfilesModel::handle(ctx).update(ctx, |profiles_model, ctx| { - let profile_id = *profiles_model.active_profile(None, ctx).id(); - profiles_model.set_base_model(profile_id, Some(id.clone()), ctx); - profiles_model.set_context_window_limit(profile_id, None, ctx); - }); - self.sync_context_window_editor(ctx, true); - self.hide_set_default_model_modal(ctx); - - let window_id = ctx.window_id(); - crate::ToastStack::handle(ctx).update(ctx, |toast_stack, ctx| { - let toast = crate::view_components::DismissibleToast::success( - "Default model updated".to_string(), - ); - toast_stack.add_ephemeral_toast(toast, window_id, ctx); - }); - ctx.notify(); - } - } - } - - fn hide_set_default_model_modal(&mut self, ctx: &mut ViewContext) { - self.set_default_model_modal.close(); - ctx.emit(AISettingsPageEvent::HideModal); - ctx.notify(); - } - - fn show_set_default_model_modal( - &mut self, - description: String, - choices: Vec<(LLMId, String)>, - ctx: &mut ViewContext, - ) { - self.set_default_model_modal.view.update(ctx, |modal, ctx| { - modal.body().update(ctx, |body, ctx| { - body.set_choices(description, choices, ctx); - }); - }); - self.set_default_model_modal.open(); - // Focus the modal so Escape closes it (the modal's escape binding only - // fires while something inside the modal holds focus). - ctx.focus(&self.set_default_model_modal.view); - ctx.emit(AISettingsPageEvent::ShowModal); - ctx.notify(); - } - - /// Returns `true` when the active Agent Mode default model is already served - /// by a credential the user has: a BYO key/subscription for its provider, or - /// one of their custom-endpoint models. `auto` models report `false` since - /// they always consume Warp credits. - fn active_base_model_is_byo_covered(ctx: &AppContext) -> bool { - let (active_id, active_provider) = { - let prefs = LLMPreferences::as_ref(ctx); - let active = prefs.get_active_base_model(ctx, None); - (active.id.clone(), active.provider.clone()) - }; - if LLMPreferences::as_ref(ctx) - .custom_llm_info_for_id(&active_id) - .is_some() - { - return true; - } - is_using_api_key_for_provider(&active_provider, ctx) - } - - /// The display name of the user's current default Agent Mode model, used in - /// the prompt copy (e.g. "auto (cost-efficient)"). - fn active_base_model_display_name(ctx: &AppContext) -> String { - LLMPreferences::as_ref(ctx) - .get_active_base_model(ctx, None) - .display_name - .clone() - } - - /// Whether to offer switching the default model. Scoped to free-plan users - /// who are out of monthly (base-plan) credits, since only they hit the - /// "no credits" error with an `auto` model. Also skips when the current - /// default is already served by a BYO credential. - fn should_offer_default_model_switch(ctx: &AppContext) -> bool { - // Exclude only confirmed paid plans. Solo/individual users have no - // `current_workspace`, and billing may not have loaded yet (Unknown), so - // treat both as eligible and rely on the out-of-credits check below to - // filter anyone who can still run Warp-hosted models. (A strict - // `is_free_plan()` check here meant solo free users — the common case — - // never saw the prompt.) - let on_paid_plan = UserWorkspaces::as_ref(ctx) - .current_workspace() - .is_some_and(|workspace| workspace.billing_metadata.is_user_on_paid_plan()); - let out_of_monthly_credits = !AIRequestUsageModel::as_ref(ctx).has_requests_remaining(); - !on_paid_plan && out_of_monthly_credits && !Self::active_base_model_is_byo_covered(ctx) - } - - /// Detects a provider key that was just added (absent -> present) by diffing - /// against the last-seen keys, then offers to switch the default model. Run - /// from `ApiKeyManagerEvent::KeysUpdated` so it fires regardless of how the - /// key editor was committed. - fn maybe_prompt_for_newly_added_provider_key(&mut self, ctx: &mut ViewContext) { - let current = ApiKeyManager::as_ref(ctx).keys().clone(); - let newly_added = [ - ( - LLMProvider::OpenAI, - &self.last_seen_provider_keys.openai, - ¤t.openai, - ), - ( - LLMProvider::Anthropic, - &self.last_seen_provider_keys.anthropic, - ¤t.anthropic, - ), - ( - LLMProvider::Google, - &self.last_seen_provider_keys.google, - ¤t.google, - ), - ] - .into_iter() - .find_map(|(provider, previous_key, current_key)| { - let was_present = previous_key - .as_deref() - .is_some_and(|key| !key.trim().is_empty()); - let now_present = current_key - .as_deref() - .is_some_and(|key| !key.trim().is_empty()); - (!was_present && now_present).then_some(provider) - }); - self.last_seen_provider_keys = current; - if let Some(provider) = newly_added { - self.maybe_prompt_set_default_model_for_provider(provider, ctx); - } - } - - /// After a BYO provider key is added, offer to switch the default Agent Mode - /// model to one from that provider. - fn maybe_prompt_set_default_model_for_provider( - &mut self, - provider: LLMProvider, - ctx: &mut ViewContext, - ) { - // Only prompt when the key is actually usable for requests (BYO enabled). - if !is_using_api_key_for_provider(&provider, ctx) { - return; - } - if !Self::should_offer_default_model_switch(ctx) { - return; - } - let choices: Vec<(LLMId, String)> = LLMPreferences::as_ref(ctx) - .get_base_llm_choices_for_agent_mode(ctx) - .filter(|llm| llm.provider == provider) - .map(|llm| (llm.id.clone(), llm.menu_display_name())) - .collect(); - if choices.is_empty() { - return; - } - let provider_name = provider.display_name(); - let current_default = Self::active_base_model_display_name(ctx); - let description = format!( - "You added your own {provider_name} API key, but your default model is currently set \ - to {current_default}, which won't work without Warp credits. Would you like to change \ - your default model?" - ); - self.show_set_default_model_modal(description, choices, ctx); - } - - #[cfg(not(target_family = "wasm"))] - fn create_grok_code_editor(ctx: &mut ViewContext) -> ViewHandle { - ctx.add_typed_action_view(|ctx| { - let appearance = Appearance::handle(ctx).as_ref(ctx); - let options = SingleLineEditorOptions { - text: TextOptions { - font_size_override: Some(appearance.ui_font_size()), - font_family_override: Some(appearance.monospace_font_family()), - text_colors_override: Some(editor_text_colors(appearance)), - ..Default::default() - }, - ..Default::default() - }; - let mut editor = EditorView::single_line(options, ctx); - editor.set_placeholder_text("Paste sign-in code", ctx); - editor - }) - } - - /// Kicks off the xAI (Grok) subscription OAuth flow: opens the consent - /// screen in the browser, runs a loopback PKCE callback server, exchanges - /// the resulting authorization code for OAuth tokens, and persists them via - /// `ApiKeyManager` (which then proactively refreshes them before expiry). - /// - /// In parallel, this reveals the manual code-entry row so the user can - /// paste the code xAI displays when the browser can't reach the loopback - /// callback. Whichever path completes first connects the subscription; the - /// other completion is ignored once the view-owned attempt state is cleared. - #[cfg(not(target_family = "wasm"))] - fn start_grok_oauth(&mut self, ctx: &mut ViewContext) { - use galaxy_core::safe_error; - - use crate::view_components::{DismissibleToast, ToastLink}; - use crate::workspace::WorkspaceAction; - use crate::ToastStack; - - /// Object id shared by the connect-flow toasts so the completion toast - /// (success or error) automatically replaces the in-progress one. - const CONNECT_TOAST_OBJECT_ID: &str = "grok_oauth_connect_toast"; - - // Record attempt initiation on click (before we attempt to bind the - // loopback server). This ensures every terminal SuperGrokSubscriptionConnectFinished - // (including immediate bind failures) is paired with a preceding Initiated - // for funnel/drop-off analysis. - send_telemetry_from_ctx!(TelemetryEvent::SuperGrokSubscriptionConnectInitiated, ctx); - - // Starting the attempt binds the loopback callback server before the - // browser opens, so a bind failure surfaces immediately, without a - // dangling browser tab. - let attempt = match oauth::OauthAttempt::start() { - Ok(attempt) => attempt, - Err(err) => { - safe_error!( - safe: ("Failed to start Grok OAuth callback server"), - full: ("Failed to start Grok OAuth callback server: {err:#}") - ); - send_telemetry_from_ctx!( - TelemetryEvent::SuperGrokSubscriptionConnectFinished { - error: Some("bind_failed".to_string()), - }, - ctx - ); - let window_id = ctx.window_id(); - ToastStack::handle(ctx).update(ctx, |toast_stack, ctx| { - let toast = - DismissibleToast::error(format!("Couldn't start Grok login: {err}")); - toast_stack.add_ephemeral_toast(toast, window_id, ctx); - }); - return; - } - }; - - // Capture the PKCE verifier so the fallback is ready if xAI shows a - // code instead of redirecting. - self.grok_oauth_attempt = Some(attempt.manual_code_exchange()); - self.grok_code_editor.update(ctx, |editor, ctx| { - editor.clear_buffer(ctx); - }); - ctx.notify(); - // Open xAI's consent screen in the user's default browser. - let authorize_url = attempt.authorize_url(); - ctx.open_url(&authorize_url); - - let window_id = ctx.window_id(); - ToastStack::handle(ctx).update(ctx, |toast_stack, ctx| { - // Persistent rather than ephemeral so the copy-URL fallback stays - // available when the browser fails to open. It can't linger - // forever: the completion toast below replaces it (shared object - // id), and the OAuth attempt itself times out when the callback - // never arrives. - let toast = DismissibleToast::default( - "Opening your browser to connect your SuperGrok subscription…".to_string(), - ) - .with_object_id(CONNECT_TOAST_OBJECT_ID.to_string()) - .with_link( - ToastLink::new("Copy URL".to_string()) - .with_onclick_action(WorkspaceAction::CopyTextToClipboard(authorize_url)), - ); - toast_stack.add_persistent_toast(toast, window_id, ctx); - }); - - ctx.spawn(async move { attempt.finish().await }, |me, result, ctx| { - // Ignore loopback completion after a successful pasted-code path. - if me.grok_oauth_attempt.is_none() { - return; - } - let window_id = ctx.window_id(); - let toast = match result { - Ok(tokens) => { - me.grok_oauth_attempt = None; - me.grok_code_editor.update(ctx, |editor, ctx| { - editor.clear_buffer(ctx); - }); - send_telemetry_from_ctx!( - TelemetryEvent::SuperGrokSubscriptionConnectFinished { error: None }, - ctx - ); - // Persist the tokens to secure storage and kick off the - // proactive refresh loop so subsequent requests can - // authenticate with the connected subscription. - ApiKeyManager::handle(ctx).update(ctx, move |manager, ctx| { - manager.store_grok_tokens(tokens, ctx); - }); - DismissibleToast::success("SuperGrok subscription connected".to_string()) - } - Err(err) => { - me.grok_oauth_attempt = None; - me.grok_code_editor.update(ctx, |editor, ctx| { - editor.clear_buffer(ctx); - }); - safe_error!( - safe: ("Grok OAuth loopback callback failed"), - full: ("Grok OAuth loopback callback failed: {err:#}") - ); - send_telemetry_from_ctx!( - TelemetryEvent::SuperGrokSubscriptionConnectFinished { - error: Some("loopback_failed".to_string()), - }, - ctx - ); - DismissibleToast::error(format!("Couldn't connect SuperGrok: {err}")) - } - }; - ToastStack::handle(ctx).update(ctx, |toast_stack, ctx| { - toast_stack.add_ephemeral_toast( - toast.with_object_id(CONNECT_TOAST_OBJECT_ID.to_string()), - window_id, - ctx, - ); - }); - ctx.notify(); - }); - } - - /// Exchanges a pasted SuperGrok authorization code using the current - /// attempt's PKCE verifier. - #[cfg(not(target_family = "wasm"))] - fn submit_grok_code(&mut self, code: String, ctx: &mut ViewContext) { - use crate::view_components::DismissibleToast; - - // Shared with the browser connect-flow toasts. - const CONNECT_TOAST_OBJECT_ID: &str = "grok_oauth_connect_toast"; - let Some(exchange) = self.grok_oauth_attempt.clone() else { - return; - }; - if code.trim().is_empty() { - return; - } - - ctx.spawn( - async move { exchange.exchange(&code).await }, - |me, result, ctx| { - if me.grok_oauth_attempt.is_none() { - return; - } - let window_id = ctx.window_id(); - let toast = match result { - Ok(tokens) => { - me.grok_oauth_attempt = None; - me.grok_code_editor.update(ctx, |editor, ctx| { - editor.clear_buffer(ctx); - }); - send_telemetry_from_ctx!( - TelemetryEvent::SuperGrokSubscriptionConnectFinished { error: None }, - ctx - ); - ApiKeyManager::handle(ctx).update(ctx, move |manager, ctx| { - manager.store_grok_tokens(tokens, ctx); - }); - DismissibleToast::success("SuperGrok subscription connected".to_string()) - } - Err(err) => { - // Keep the row open so the user can correct the code. - safe_error!( - safe: ("Grok manual code exchange failed"), - full: ("Grok manual code exchange failed: {err:#}") - ); - send_telemetry_from_ctx!( - TelemetryEvent::SuperGrokSubscriptionConnectFinished { - error: Some("manual_code_failed".to_string()), - }, - ctx - ); - DismissibleToast::error(format!("Couldn't connect SuperGrok: {err}")) - } - }; - ToastStack::handle(ctx).update(ctx, |toast_stack, ctx| { - toast_stack.add_ephemeral_toast( - toast.with_object_id(CONNECT_TOAST_OBJECT_ID.to_string()), - window_id, - ctx, - ); - }); - ctx.notify(); - }, - ); + None } /// Set the active subpage and rebuild the widget list to show only relevant widgets. @@ -2365,7 +1871,7 @@ impl AISettingsPageView { widgets.push(Box::new(OtherAIWidget::default())); } Some(AISubpage::WarpAgent) => { - // Oz page: global toggle + Active AI + Input + Other + // Galaxy Agent page: global toggle + Active AI + Input + Other widgets.push(Box::new(GlobalAIWidget::default())); if ai_settings .intelligent_autosuggestions_enabled_internal @@ -2397,8 +1903,6 @@ impl AISettingsPageView { widgets.push(Box::new(VoiceWidget::default())); } widgets.push(Box::new(CloudHandoffWidget::default())); - widgets.push(Box::new(ApiKeysWidget::new(ctx))); - widgets.push(Box::new(BedrockSettingsWidget::new(ctx))); if FeatureFlag::CustomModelRouters.is_enabled() { widgets.push(Box::new(CustomModelRoutersWidget)); } @@ -3142,7 +2646,6 @@ pub enum AISettingsPageAction { ToggleCLIAgentToolbar, ToggleUseAgentToolbar, ToggleVoiceInput, - ToggleCanUseWarpCreditsForFallback, HyperlinkClick(HyperlinkUrl), ToggleCodebaseContext, ToggleShowInputHintText, @@ -3203,9 +2706,6 @@ pub enum AISettingsPageAction { #[cfg(feature = "local_fs")] OpenAddCustomRouter, - ConnectGrokSubscription, - DisconnectGrokSubscription, - #[cfg(feature = "local_fs")] SetConversationLayout(crate::util::file::external_editor::settings::OpenConversationPreference), ToggleCloudHandoff, @@ -3560,14 +3060,6 @@ impl TypedActionView for AISettingsPageView { } ctx.notify(); } - AISettingsPageAction::ToggleCanUseWarpCreditsForFallback => { - AISettings::handle(ctx).update(ctx, |settings, ctx| { - report_if_error!(settings - .can_use_warp_credits_for_fallback - .toggle_and_save_value(ctx)); - }); - ctx.notify(); - } AISettingsPageAction::HyperlinkClick(hyperlink) => { ctx.notify(); ctx.open_url(&hyperlink.url); @@ -4047,31 +3539,6 @@ impl TypedActionView for AISettingsPageView { }); ctx.notify(); } - AISettingsPageAction::ConnectGrokSubscription => { - #[cfg(not(target_family = "wasm"))] - self.start_grok_oauth(ctx); - } - AISettingsPageAction::DisconnectGrokSubscription => { - #[cfg(not(target_family = "wasm"))] - { - self.grok_oauth_attempt = None; - self.grok_code_editor.update(ctx, |editor, ctx| { - editor.clear_buffer(ctx); - }); - } - ApiKeyManager::handle(ctx).update(ctx, |manager, ctx| { - manager.set_grok_tokens(None, ctx); - }); - - let window_id = ctx.window_id(); - crate::ToastStack::handle(ctx).update(ctx, |toast_stack, ctx| { - let toast = crate::view_components::DismissibleToast::default( - "SuperGrok subscription disconnected".to_string(), - ); - toast_stack.add_ephemeral_toast(toast, window_id, ctx); - }); - ctx.notify(); - } } } } @@ -4298,8 +3765,8 @@ impl SettingsWidget for GlobalAIWidget { type View = AISettingsPageView; fn search_terms(&self) -> &str { - "oz warp agent global ai a.i. active next command prompt code diffs suggestion suggested suggestions \ - agent mode natural language detection input hint api keys bring your own byo google anthropic openai" + "galaxy agent global ai a.i. active next command prompt code diffs suggestion suggested suggestions \ + agent mode natural language detection input hint" } fn render( @@ -5584,7 +5051,7 @@ impl AgentsWidget { ); render_ai_list( "Command denylist", - "Regular expressions to match commands that the Warp Agent should always ask permission to execute.", + "Regular expressions to match commands that the Galaxy Agent should always ask permission to execute.", list, view, ai_settings, @@ -5619,7 +5086,7 @@ impl AgentsWidget { render_ai_list( "Command allowlist", - "Regular expressions to match commands that can be automatically executed by the Warp Agent.", + "Regular expressions to match commands that can be automatically executed by the Galaxy Agent.", list, view, ai_settings, @@ -5721,7 +5188,7 @@ impl AgentsWidget { appearance, "Base model", Some( - "This model serves as the primary engine behind the Warp Agent. It powers most interactions and invokes other models for tasks like planning or code generation when necessary. Warp may automatically switch to alternate models based on model availability or for auxiliary tasks such as conversation summarization.", + "This model serves as the primary engine behind the Galaxy Agent. It powers most interactions and invokes other models for tasks like planning or code generation when necessary. Galaxy may automatically switch to alternate models based on model availability or for auxiliary tasks such as conversation summarization.", ), Some(show_in_prompt_checkbox), LocalOnlyIconState::Hidden, @@ -5752,7 +5219,7 @@ impl AgentsWidget { let codebase_context_description = vec![ FormattedTextFragment::plain_text( - "Allow the Warp Agent to generate an outline of your codebase that can be used for context. No code is ever stored on our servers. ", + "Allow the Galaxy Agent to generate an outline of your codebase that can be used for context. No code is ever stored on our servers. ", ), FormattedTextFragment::hyperlink( "Learn more", @@ -5825,7 +5292,7 @@ impl AgentsWidget { let subtext = { let subtext_fragments = vec![ FormattedTextFragment::plain_text( - "You haven't added any MCP servers yet. Once you do, you'll be able to control how much autonomy the Warp Agent has when interacting with them. ", + "You haven't added any MCP servers yet. Once you do, you'll be able to control how much autonomy the Galaxy Agent has when interacting with them. ", ), FormattedTextFragment::hyperlink_action( "Add a server", @@ -5906,7 +5373,7 @@ impl AgentsWidget { { let allowlist = self.render_mcp_list( "MCP allowlist", - "Allow the Warp Agent to call these MCP servers.", + "Allow the Galaxy Agent to call these MCP servers.", &view.mcp_allowlist_dropdown, BlocklistAIPermissions::as_ref(app).get_mcp_allowlist(app, None), view.mcp_allowlist_mouse_state_handles.clone(), @@ -5923,7 +5390,7 @@ impl AgentsWidget { { let denylist = self.render_mcp_list( "MCP denylist", - "The Warp Agent will always ask for permission before calling any MCP servers on this list.", + "The Galaxy Agent will always ask for permission before calling any MCP servers on this list.", &view.mcp_denylist_dropdown, BlocklistAIPermissions::as_ref(app).get_mcp_denylist(app, None), view.mcp_denylist_mouse_state_handles.clone(), @@ -6343,7 +5810,7 @@ impl SettingsWidget for MCPServersWidget { let mcp_description = vec![ FormattedTextFragment::plain_text( - "Add MCP servers to extend the Warp Agent's capabilities. \ + "Add MCP servers to extend the Galaxy Agent's capabilities. \ MCP servers expose data sources or tools to agents through a standardized interface, essentially acting like plugins. ", ), FormattedTextFragment::hyperlink( @@ -6475,7 +5942,7 @@ impl AIFactWidget { let rules_description = vec![ FormattedTextFragment::plain_text( - "Rules help the Warp Agent follow your conventions, whether for codebases or specific workflows. ", + "Rules help the Galaxy Agent follow your conventions, whether for codebases or specific workflows. " ), FormattedTextFragment::hyperlink( "Learn more", @@ -7598,7 +7065,7 @@ impl SettingsWidget for CloudHandoffWidget { ); column.add_child(auto_handoff_on_sleep_row); column.add_child(render_ai_setting_description( - "When macOS is about to sleep, automatically moves the most recently focused running local Warp Agent conversation to Cloud Mode so it can keep working.", + "When macOS is about to sleep, automatically moves the most recently focused running local Galaxy Agent conversation to Cloud Mode so it can keep working.", true, app, )); @@ -7638,645 +7105,6 @@ impl SettingsWidget for CloudHandoffWidget { } } -struct ApiKeysWidget { - openai_api_key_editor: ViewHandle, - anthropic_api_key_editor: ViewHandle, - google_api_key_editor: ViewHandle, - /// Buttons for the SuperGrok (xAI) subscription row; which one renders - /// depends on whether OAuth tokens are stored or a connect attempt is in - /// progress. - grok_connect_button: ViewHandle, - grok_connecting_button: ViewHandle, - grok_disconnect_button: ViewHandle, - - can_use_warp_credits_for_fallback: SwitchStateHandle, - upgrade_highlight_index: HighlightedHyperlink, - - description_learn_more_index: HighlightedHyperlink, -} - -impl ApiKeysWidget { - fn new(ctx: &mut ViewContext<::View>) -> Self { - let ai_settings = AISettings::as_ref(ctx); - let workspace_handle = UserWorkspaces::handle(ctx); - let is_any_ai_enabled = ai_settings.is_any_ai_enabled(ctx); - let is_byo_enabled = workspace_handle.as_ref(ctx).is_byo_api_key_enabled(ctx); - - let ApiKeys { - openai: openai_key, - anthropic: anthropic_key, - google: google_key, - .. - } = ApiKeyManager::as_ref(ctx).keys().clone(); - - // A helper macro to create and configure an API key editor. This avoids a lot - // of code duplication and ensures consistency between the editors. - macro_rules! create_api_key_editor { - ($editor:ident, $key:ident, $set_func:ident, $placeholder:literal) => { - let $editor = ctx.add_typed_action_view(move |ctx| { - let appearance = Appearance::handle(ctx).as_ref(ctx); - let options = SingleLineEditorOptions { - is_password: true, - // Emit Tab/Shift-Tab as navigation events instead of - // inserting whitespace, so focus can move between the - // key fields (see the focus wiring below). - propagate_and_no_op_vertical_navigation_keys: - PropagateAndNoOpNavigationKeys::Always, - text: TextOptions { - font_size_override: Some(appearance.ui_font_size()), - font_family_override: Some(appearance.monospace_font_family()), - text_colors_override: Some(TextColors { - default_color: appearance.theme().active_ui_text_color(), - disabled_color: appearance.theme().disabled_ui_text_color(), - hint_color: appearance.theme().disabled_ui_text_color(), - }), - ..Default::default() - }, - ..Default::default() - }; - let mut editor = EditorView::single_line(options, ctx); - editor.set_placeholder_text($placeholder, ctx); - if let Some(key) = &$key { - editor.set_buffer_text(key, ctx); - } - editor - }); - AISettingsPageView::update_editor_interaction_state( - $editor.clone(), - is_any_ai_enabled && is_byo_enabled, - ctx, - ); - // The default-model prompt is driven off `KeysUpdated` (see the - // `ApiKeyManager` subscription), so this only needs to persist - // the key on commit. - ctx.subscribe_to_view(&$editor, |_, $editor, event, ctx| { - if matches!(event, EditorEvent::Blurred | EditorEvent::Enter) { - let buffer_text = $editor.as_ref(ctx).buffer_text(ctx); - let key = (!buffer_text.is_empty()).then_some(buffer_text); - ApiKeyManager::handle(ctx).update(ctx, |model, ctx| { - model.$set_func(key, ctx); - }); - } - }); - let editor_clone = $editor.clone(); - ctx.subscribe_to_model(&workspace_handle, move |_, workspace, event, ctx| { - if let UserWorkspacesEvent::TeamsChanged = event { - let is_any_ai_enabled = - AISettings::handle(ctx).as_ref(ctx).is_any_ai_enabled(ctx); - let is_byo_enabled = workspace.as_ref(ctx).is_byo_api_key_enabled(ctx); - let is_enabled = is_any_ai_enabled && is_byo_enabled; - let has_key = !editor_clone.as_ref(ctx).is_empty(ctx); - - // If BYO is disabled, clear the API key from the editor and storage - if !is_byo_enabled && has_key { - editor_clone.update(ctx, |editor, ctx| { - editor.set_buffer_text("", ctx); - }); - ApiKeyManager::handle(ctx).update(ctx, |model, ctx| { - model.$set_func(None, ctx); - }); - } - - AISettingsPageView::update_editor_interaction_state( - editor_clone.clone(), - is_enabled, - ctx, - ); - ctx.notify(); - } - }) - }; - } - - create_api_key_editor!(openai_api_key_editor, openai_key, set_openai_key, "sk-..."); - create_api_key_editor!( - anthropic_api_key_editor, - anthropic_key, - set_anthropic_key, - "sk-ant-..." - ); - create_api_key_editor!( - google_api_key_editor, - google_key, - set_google_key, - "AIzaSy..." - ); - - // Tab / Shift-Tab move focus between the provider key fields instead of - // inserting whitespace. - let provider_key_editors = [ - openai_api_key_editor.clone(), - anthropic_api_key_editor.clone(), - google_api_key_editor.clone(), - ]; - for (index, editor) in provider_key_editors.iter().enumerate() { - let next = provider_key_editors.get(index + 1).cloned(); - let previous = index - .checked_sub(1) - .and_then(|prev_index| provider_key_editors.get(prev_index).cloned()); - ctx.subscribe_to_view(editor, move |_, _, event, ctx| match event { - EditorEvent::Navigate(NavigationKey::Tab) => { - if let Some(next) = &next { - ctx.focus(next); - } - } - EditorEvent::Navigate(NavigationKey::ShiftTab) => { - if let Some(previous) = &previous { - ctx.focus(previous); - } - } - _ => {} - }); - } - - // Editor text colors are snapshotted at construction via - // `text_colors_override`, so refresh them whenever the theme changes. - let api_key_editors = [ - openai_api_key_editor.clone(), - anthropic_api_key_editor.clone(), - google_api_key_editor.clone(), - ]; - ctx.subscribe_to_model(&Appearance::handle(ctx), move |_, _, event, ctx| { - if let AppearanceEvent::ThemeChanged = event { - let text_colors = editor_text_colors(Appearance::as_ref(ctx)); - for editor in &api_key_editors { - let colors = text_colors.clone(); - editor.update(ctx, move |editor, ctx| { - editor.set_text_colors(colors, ctx); - }); - } - } - }); - - let grok_connect_button = ctx.add_typed_action_view(|_| { - ActionButton::new("Connect", SecondaryTheme) - .with_size(ButtonSize::Small) - .on_click(|ctx| { - ctx.dispatch_typed_action(AISettingsPageAction::ConnectGrokSubscription); - }) - }); - let grok_connecting_button = ctx.add_typed_action_view(|_| { - ActionButton::new("Connecting", SecondaryTheme).with_size(ButtonSize::Small) - }); - grok_connecting_button.update(ctx, |button, ctx| { - button.set_disabled(true, ctx); - }); - let grok_disconnect_button = ctx.add_typed_action_view(|_| { - ActionButton::new("Disconnect", DangerSecondaryTheme) - .with_size(ButtonSize::Small) - .on_click(|ctx| { - ctx.dispatch_typed_action(AISettingsPageAction::DisconnectGrokSubscription); - }) - }); - for button in [&grok_connect_button, &grok_disconnect_button] { - button.update(ctx, |button, ctx| { - button.set_disabled(!(is_any_ai_enabled && is_byo_enabled), ctx); - }); - } - - // The Grok subscription is BYO auth, so keep the buttons' enablement - // in sync with the BYO API key policy, like the editors above. - let grok_buttons = [grok_connect_button.clone(), grok_disconnect_button.clone()]; - ctx.subscribe_to_model(&workspace_handle, move |_, workspace, event, ctx| { - if let UserWorkspacesEvent::TeamsChanged = event { - let is_any_ai_enabled = AISettings::handle(ctx).as_ref(ctx).is_any_ai_enabled(ctx); - let is_byo_enabled = workspace.as_ref(ctx).is_byo_api_key_enabled(ctx); - for button in &grok_buttons { - button.update(ctx, |button, ctx| { - button.set_disabled(!(is_any_ai_enabled && is_byo_enabled), ctx); - }); - } - ctx.notify(); - } - }); - - // Re-render the SuperGrok row whenever the stored tokens change (the - // connect flow completes, a disconnect, or a background refresh). - ctx.subscribe_to_model(&ApiKeyManager::handle(ctx), |_, _, event, ctx| { - if matches!(event, ApiKeyManagerEvent::KeysUpdated) { - ctx.notify(); - } - }); - - Self { - openai_api_key_editor, - anthropic_api_key_editor, - google_api_key_editor, - - grok_connect_button, - grok_connecting_button, - grok_disconnect_button, - - can_use_warp_credits_for_fallback: Default::default(), - upgrade_highlight_index: Default::default(), - - description_learn_more_index: Default::default(), - } - } - - fn render_api_key_input( - &self, - appearance: &Appearance, - label: &'static str, - editor: ViewHandle, - is_enabled: bool, - app: &AppContext, - ) -> Box { - let padding = Some(Coords { - top: 10., - bottom: 10., - left: 16., - right: 16., - }); - let editor_style = UiComponentStyles { - padding, - background: Some(appearance.theme().surface_2().into()), - ..Default::default() - }; - - let label = Text::new_inline(label, appearance.ui_font_family(), CONTENT_FONT_SIZE) - .with_color(styles::header_font_color(is_enabled, app).into()) - .finish(); - - let input = appearance - .ui_builder() - .text_input(editor) - .with_style(editor_style) - .build() - .finish(); - - Flex::column() - .with_spacing(8.) - .with_child(label) - .with_child(input) - .finish() - } - - fn render_provider_key_editors( - &self, - appearance: &Appearance, - is_enabled: bool, - app: &AppContext, - ) -> Box { - let mut column = Flex::column().with_spacing(16.); - column.add_child(self.render_api_key_input( - appearance, - "OpenAI API key", - self.openai_api_key_editor.clone(), - is_enabled, - app, - )); - column.add_child(self.render_api_key_input( - appearance, - "Anthropic API key", - self.anthropic_api_key_editor.clone(), - is_enabled, - app, - )); - column.add_child(self.render_api_key_input( - appearance, - "Google API key", - self.google_api_key_editor.clone(), - is_enabled, - app, - )); - column.finish() - } - - /// The "Connect SuperGrok subscription" row: label and description on the - /// left, a Connect/Disconnect button on the right, and a "Connected on - /// ..." status line underneath while a subscription is connected. - fn render_grok_subscription_row( - &self, - appearance: &Appearance, - is_enabled: bool, - is_connecting: bool, - app: &AppContext, - ) -> Box { - let grok_tokens = ApiKeyManager::as_ref(app).grok_tokens(); - - let text_color = styles::header_font_color(is_enabled, app); - let label = Flex::row() - .with_cross_axis_alignment(CrossAxisAlignment::Center) - .with_spacing(4.) - .with_child( - Text::new_inline("Use your", appearance.ui_font_family(), CONTENT_FONT_SIZE) - .with_color(text_color.into()) - .finish(), - ) - .with_child( - ConstrainedBox::new(Icon::XLogo.to_galaxyui_icon(text_color).finish()) - .with_width(14.) - .with_height(14.) - .finish(), - ) - .with_child( - Text::new_inline( - "Premium or SuperGrok subscription", - appearance.ui_font_family(), - CONTENT_FONT_SIZE, - ) - .with_color(text_color.into()) - .finish(), - ) - .finish(); - - let button = if grok_tokens.is_some() { - &self.grok_disconnect_button - } else if is_connecting { - &self.grok_connecting_button - } else { - &self.grok_connect_button - }; - - let header_row = Flex::row() - .with_main_axis_size(MainAxisSize::Max) - .with_main_axis_alignment(MainAxisAlignment::SpaceBetween) - .with_cross_axis_alignment(CrossAxisAlignment::Center) - .with_child(Shrinkable::new(1., label).finish()) - .with_child(button.as_ref(app).render(app)) - .finish(); - - let description = Container::new( - Text::new( - "Connect your SuperGrok subscription to use Grok models in the Warp Agent through your xAI account.", - appearance.ui_font_family(), - CONTENT_FONT_SIZE, - ) - .with_color(styles::description_font_color(is_enabled, app).into()) - .soft_wrap(true) - .finish(), - ) - .with_margin_right(styles::TOGGLE_WIDTH_MARGIN) - .finish(); - - let mut column = Flex::column() - .with_cross_axis_alignment(CrossAxisAlignment::Start) - .with_child(header_row) - .with_child(description); - - if let Some(tokens) = grok_tokens { - let connected_text = match tokens.connected_at.map(DateTime::::from) { - Some(connected_at) => format!( - "Connected on {}.", - connected_at.format("%m/%d/%Y at %-I:%M%P") - ), - // Tokens stored before the connection time was tracked. - None => "Connected.".to_string(), - }; - let check = ConstrainedBox::new( - Icon::Check - .to_galaxyui_icon(appearance.theme().ansi_fg_green().into()) - .finish(), - ) - .with_width(12.) - .with_height(12.) - .finish(); - let status_text = Text::new_inline( - connected_text, - appearance.ui_font_family(), - CONTENT_FONT_SIZE, - ) - .with_color(styles::description_font_color(is_enabled, app).into()) - .finish(); - column.add_child( - Flex::row() - .with_cross_axis_alignment(CrossAxisAlignment::Center) - .with_spacing(4.) - .with_child(check) - .with_child(status_text) - .finish(), - ); - } - - column.finish() - } - - /// Paste-the-code fallback for the current SuperGrok connect attempt. - #[cfg(not(target_family = "wasm"))] - fn render_grok_manual_code_entry( - &self, - view: &AISettingsPageView, - appearance: &Appearance, - ) -> Box { - let theme = appearance.theme(); - - let editor_style = UiComponentStyles { - padding: Some(Coords { - top: 10., - bottom: 10., - left: 16., - right: 16., - }), - background: Some(theme.surface_2().into()), - ..Default::default() - }; - let input = appearance - .ui_builder() - .text_input(view.grok_code_editor.clone()) - .with_style(editor_style) - .build() - .finish(); - - let row = Flex::row() - .with_main_axis_size(MainAxisSize::Max) - .with_cross_axis_alignment(CrossAxisAlignment::Center) - .with_spacing(8.) - .with_child(Shrinkable::new(1., input).finish()) - .finish(); - - Flex::column() - .with_cross_axis_alignment(CrossAxisAlignment::Start) - .with_spacing(8.) - .with_child(row) - .finish() - } - - fn render_warp_credit_fallback_toggle( - &self, - view: &AISettingsPageView, - app: &AppContext, - ) -> Box { - let ai_settings = AISettings::as_ref(app); - - let toggle = render_ai_setting_toggle::( - "Warp credit fallback", - AISettingsPageAction::ToggleCanUseWarpCreditsForFallback, - *ai_settings.can_use_warp_credits_for_fallback, - ai_settings.is_any_ai_enabled(app), - self.can_use_warp_credits_for_fallback.clone(), - &view.local_only_icon_tooltip_states, - app, - ); - - let description = render_ai_setting_description( - "When enabled, agent requests may be routed to one of Warp's provided models in the event of an error. Warp will prioritize using your API keys over your Warp credits.", - ai_settings.is_any_ai_enabled(app), - app, - ); - - Flex::column() - .with_child(toggle) - .with_child(description) - .finish() - } -} - -impl SettingsWidget for ApiKeysWidget { - type View = AISettingsPageView; - - fn search_terms(&self) -> &str { - "api keys bring your own byo openai anthropic google claude gemini gpt custom inference endpoint grok supergrok xai subscription" - } - - fn render( - &self, - view: &Self::View, - appearance: &Appearance, - app: &AppContext, - ) -> Box { - let ai_settings = AISettings::as_ref(app); - let is_any_ai_enabled = ai_settings.is_any_ai_enabled(app); - let is_byo_enabled = UserWorkspaces::as_ref(app).is_byo_api_key_enabled(app); - let provider_keys_enabled = is_any_ai_enabled && is_byo_enabled; - - let mut column = Flex::column().with_child(render_separator(appearance)); - - column.add_child( - build_sub_header( - appearance, - "API Keys", - Some(styles::header_font_color(is_any_ai_enabled, app)), - ) - .with_padding_bottom(HEADER_PADDING) - .finish(), - ); - - // Provider key editors (always visible) - column.add_child(self.render_provider_key_editors(appearance, provider_keys_enabled, app)); - - - // Entrypoint for connecting a SuperGrok (xAI) subscription via OAuth. - if FeatureFlag::SuperGrok.is_enabled() { - #[cfg(not(target_family = "wasm"))] - let grok_tokens = ApiKeyManager::as_ref(app).grok_tokens(); - #[cfg(not(target_family = "wasm"))] - let has_grok_oauth_attempt = view.grok_oauth_attempt.is_some(); - #[cfg(not(target_family = "wasm"))] - let is_grok_connecting = grok_tokens.is_none() && has_grok_oauth_attempt; - #[cfg(target_family = "wasm")] - let is_grok_connecting = false; - column.add_child( - Container::new(self.render_grok_subscription_row( - appearance, - provider_keys_enabled, - is_grok_connecting, - app, - )) - .with_margin_top(16.) - .finish(), - ); - - #[cfg(not(target_family = "wasm"))] - if has_grok_oauth_attempt { - column.add_child( - Container::new(self.render_grok_manual_code_entry(view, appearance)) - .with_margin_top(8.) - .finish(), - ); - } - } - - // Warp credit fallback toggle (shown when BYO is enabled) - if is_byo_enabled { - column.add_child( - Container::new(self.render_warp_credit_fallback_toggle(view, app)) - .with_margin_top(16.) - .finish(), - ); - } - - // Upgrade CTA if BYOK not enabled - if !is_byo_enabled { - let auth_state = AuthStateProvider::as_ref(app).get(); - let upgrade_text_fragments = if let Some(team) = - UserWorkspaces::as_ref(app).current_team() - { - if team.billing_metadata.customer_type == CustomerType::Enterprise { - vec![ - FormattedTextFragment::hyperlink("Contact sales", "mailto:sales@warp.dev"), - FormattedTextFragment::plain_text( - " to enable bringing your own API keys on your Enterprise plan.", - ), - ] - } else { - let current_user_email = auth_state.user_email().unwrap_or_default(); - let has_admin_permissions = team.has_admin_permissions(¤t_user_email); - let upgrade_url = UserWorkspaces::upgrade_link_for_team(team.uid); - if has_admin_permissions { - vec![ - FormattedTextFragment::hyperlink( - "Upgrade to the Build plan", - upgrade_url, - ), - FormattedTextFragment::plain_text(" to use your own API keys."), - ] - } else { - vec![FormattedTextFragment::plain_text( - "Ask your team's admin to upgrade to the Build plan to use your own API keys.", - )] - } - } - } else if FeatureFlag::SoloUserByok.is_enabled() - && auth_state.is_anonymous_or_logged_out() - { - vec![ - FormattedTextFragment::hyperlink_action( - "Create an account", - AISettingsPageAction::SignupAnonymousUser, - ), - FormattedTextFragment::plain_text(" to use your own API keys."), - ] - } else { - let user_id = auth_state.user_id().unwrap_or_default(); - let upgrade_url = UserWorkspaces::upgrade_link(user_id); - vec![ - FormattedTextFragment::hyperlink("Upgrade to the Build plan", upgrade_url), - FormattedTextFragment::plain_text(" to use your own API keys."), - ] - }; - - let upgrade_text_element = FormattedTextElement::new( - FormattedText::new([FormattedTextLine::Line(upgrade_text_fragments)]), - appearance.ui_font_size(), - appearance.ui_font_family(), - appearance.ui_font_family(), - blended_colors::text_sub(appearance.theme(), appearance.theme().surface_1()), - self.upgrade_highlight_index.clone(), - ) - .with_hyperlink_font_color(appearance.theme().accent().into_solid()) - .register_default_click_handlers_with_action_support(|hyperlink_lens, event, ctx| { - match hyperlink_lens { - HyperlinkLens::Url(url) => { - ctx.open_url(url); - } - HyperlinkLens::Action(action_ref) => { - if let Some(action) = - action_ref.as_any().downcast_ref::() - { - event.dispatch_typed_action(action.clone()); - } - } - } - }); - - column.add_child(Container::new(upgrade_text_element.finish()).finish()); - } - - column.finish() - } -} - struct BedrockSettingsWidget { enabled_toggle: SwitchStateHandle, auto_login_toggle: SwitchStateHandle, diff --git a/app/src/settings_view/mcp_servers/list_page.rs b/app/src/settings_view/mcp_servers/list_page.rs index dcb915a2..30e2d275 100644 --- a/app/src/settings_view/mcp_servers/list_page.rs +++ b/app/src/settings_view/mcp_servers/list_page.rs @@ -69,7 +69,7 @@ use crate::workspace::Workspace; use crate::workspaces::user_workspaces::UserWorkspaces; use crate::ToastStack; -const DESCRIPTION_TEXT: &str = "Add MCP servers to extend the Warp Agent's capabilities. MCP servers expose data sources or tools to agents through a standardized interface, essentially acting like plugins. Add a custom server, or use the presets to get started with popular servers. You can also find team servers that have been shared with you here. "; +const DESCRIPTION_TEXT: &str = "Add MCP servers to extend the Galaxy Agent's capabilities. MCP servers expose data sources or tools to agents through a standardized interface, essentially acting like plugins. Add a custom server, or use the presets to get started with popular servers. You can also find team servers that have been shared with you here. "; #[derive(Debug, Clone)] pub enum MCPServersListPageViewEvent { diff --git a/app/src/settings_view/mod.rs b/app/src/settings_view/mod.rs index fb65ac56..038c5410 100644 --- a/app/src/settings_view/mod.rs +++ b/app/src/settings_view/mod.rs @@ -269,7 +269,7 @@ impl Display for SettingsSection { SettingsSection::MCPServers => write!(f, "MCP Servers"), SettingsSection::Scripting => write!(f, "Scripting"), SettingsSection::WarpDrive => write!(f, "Galaxy Drive"), - SettingsSection::WarpAgent => write!(f, "Warp Agent"), + SettingsSection::WarpAgent => write!(f, "Galaxy Agent"), SettingsSection::AgentProfiles => write!(f, "Profiles"), SettingsSection::AgentMCPServers => write!(f, "MCP servers"), SettingsSection::Knowledge => write!(f, "Knowledge"), @@ -544,7 +544,6 @@ pub mod flags { pub const SUGGESTED_RULES_FLAG: &str = "Suggested_Rules"; pub const WARP_DRIVE_CONTEXT_FLAG: &str = "Warp_Drive_Context"; pub const FILE_BASED_MCP_FLAG: &str = "File_Based_MCP"; - pub const WARP_CREDIT_FALLBACK_FLAG: &str = "Warp_Credit_Fallback"; pub const SHOW_BASE_MODEL_PICKER_IN_PROMPT_FLAG: &str = "Show_Base_Model_Picker_In_Prompt"; pub const DEBUG_SHOW_MEMORY_STATS_FLAG: &str = "Debug_Memory_Statistics"; pub const ALLOW_NATIVE_WAYLAND: &str = "Allow_Native_Wayland"; diff --git a/app/src/settings_view/mod_tests.rs b/app/src/settings_view/mod_tests.rs index 6bb445b4..14a54e98 100644 --- a/app/src/settings_view/mod_tests.rs +++ b/app/src/settings_view/mod_tests.rs @@ -171,7 +171,7 @@ fn match_data_countable_zero_is_not_truthy() { #[test] fn subpage_display_names_are_correct() { - assert_eq!(SettingsSection::WarpAgent.to_string(), "Warp Agent"); + assert_eq!(SettingsSection::WarpAgent.to_string(), "Galaxy Agent"); assert_eq!(SettingsSection::AgentProfiles.to_string(), "Profiles"); assert_eq!(SettingsSection::AgentMCPServers.to_string(), "MCP servers"); assert_eq!(SettingsSection::Knowledge.to_string(), "Knowledge"); diff --git a/app/src/terminal/input/models/data_source.rs b/app/src/terminal/input/models/data_source.rs index 09ff1fc7..903f739c 100644 --- a/app/src/terminal/input/models/data_source.rs +++ b/app/src/terminal/input/models/data_source.rs @@ -630,14 +630,6 @@ impl SearchItem for ModelSearchItem { first.make_ascii_uppercase(); } - // Show a BYOK option when the user's tier supports it and the provider - // is one that accepts user-supplied API keys. - let byok_available = UserWorkspaces::as_ref(app).is_byo_api_key_enabled(app) - && matches!( - self.provider, - LLMProvider::OpenAI | LLMProvider::Anthropic | LLMProvider::Google - ); - let mut text_fragments = vec![ FormattedTextFragment::plain_text(format!( "{display_name} is not available for free users. " @@ -645,17 +637,6 @@ impl SearchItem for ModelSearchItem { FormattedTextFragment::hyperlink("Upgrade", upgrade_url), ]; - if byok_available { - text_fragments.push(FormattedTextFragment::plain_text(" or ".to_string())); - text_fragments.push(FormattedTextFragment::hyperlink_action( - "bring your own key", - WorkspaceAction::ShowSettingsPageWithSearch { - search_query: "api".to_string(), - section: Some(SettingsSection::WarpAgent), - }, - )); - } - let upgrade_text = FormattedTextElement::new( FormattedText::new([FormattedTextLine::Line(text_fragments)]), inline_styles::font_size(appearance), diff --git a/app/src/workspace/one_time_modal_model.rs b/app/src/workspace/one_time_modal_model.rs index 69fe4486..4a466478 100644 --- a/app/src/workspace/one_time_modal_model.rs +++ b/app/src/workspace/one_time_modal_model.rs @@ -1,6 +1,5 @@ use std::future::Future; -use ai::api_keys::ApiKeyManager; use galaxy_core::features::FeatureFlag; use galaxy_core::send_telemetry_from_ctx; use settings::Setting as _; @@ -442,7 +441,10 @@ impl OneTimeModalModel { .current_workspace() .map(|workspace| workspace.billing_metadata.customer_type); let is_warp_ai_enabled = *AISettings::as_ref(ctx).is_any_ai_enabled; - let has_byok_or_byoe = ApiKeyManager::as_ref(ctx).has_any_key(); + // BYOK/BYOE (bring your own provider key/endpoint) has been removed; Galaxy + // only supports AWS Bedrock and OpenAI-compatible endpoints configured by + // the user in Settings, neither of which participate in this decision. + let has_byok_or_byoe = false; let completed_new_onboarding = has_completed_local_onboarding(ctx); let has_zero_base_credits = AIRequestUsageModel::as_ref(ctx).request_limit() == 0; diff --git a/app/src/workspace/view.rs b/app/src/workspace/view.rs index 763d2903..9f4ba949 100644 --- a/app/src/workspace/view.rs +++ b/app/src/workspace/view.rs @@ -287,9 +287,9 @@ use crate::pane_group::pane::ActionOrigin; use crate::pane_group::FilePane; use crate::pane_group::{ self, AIFactPane, AnyPaneContent, ChildAgentOrigin, CodeDiffPane, CodePane, CodeReviewPanelArg, - CustomRouterEditorPane, Direction as PaneGroupDirection, Direction, - ExecutionProfileEditorPane, NetworkLogPane, NewTerminalOptions, PaneGroup, PaneId, PanesLayout, - TabBarHoverIndex, TerminalPaneId, + CustomRouterEditorPane, Direction as PaneGroupDirection, Direction, ExecutionProfileEditorPane, + NetworkLogPane, NewTerminalOptions, PaneGroup, PaneId, PanesLayout, TabBarHoverIndex, + TerminalPaneId, }; use crate::persistence::ModelEvent; use crate::projects::ProjectManagementModel; @@ -22858,9 +22858,6 @@ impl Workspace { if *ai_settings.file_based_mcp_enabled.value() { context.set.insert(flags::FILE_BASED_MCP_FLAG); } - if *ai_settings.can_use_warp_credits_for_fallback.value() { - context.set.insert(flags::WARP_CREDIT_FALLBACK_FLAG); - } if *session_settings.show_model_selectors_in_prompt.value() { context .set diff --git a/crates/ai/src/api_keys.rs b/crates/ai/src/api_keys.rs index 3b5edb8b..b4931d9a 100644 --- a/crates/ai/src/api_keys.rs +++ b/crates/ai/src/api_keys.rs @@ -1,9 +1,4 @@ -use std::time::{Duration, SystemTime}; - use galaxyui_core::{Entity, ModelContext, SingletonEntity}; -use galaxyui_extras::secure_storage::{self, AppContextExt}; -use serde::{Deserialize, Serialize}; -use uuid::Uuid; use warp_multi_agent_api as api; pub use crate::aws_credentials::{AwsCredentials, AwsCredentialsState}; @@ -12,137 +7,13 @@ pub use crate::geap_credentials::{ LoadGeapCredentialsError, GEAP_REFRESH_LEAD_TIME, }; -const SECURE_STORAGE_KEY: &str = "AiApiKeys"; - -/// Secure-storage key for the connected xAI/Grok subscription's OAuth tokens. -/// Kept separate from [`SECURE_STORAGE_KEY`] because these are OAuth tokens with -/// a refresh lifecycle, not a user-pasted static key. -const GROK_SECURE_STORAGE_KEY: &str = "GrokOAuthTokens"; - -/// Emitted when user-provided API keys are updated in-memory. +/// Emitted when the manager's stored credentials (AWS Bedrock or Gemini +/// Enterprise) are updated in-memory. #[derive(Debug, Clone, PartialEq, Eq)] pub enum ApiKeyManagerEvent { KeysUpdated, } -/// User-provided API keys for AI providers. -/// -/// These are used for "Bring Your Own API Key" functionality, allowing -/// users to use their own API keys instead of Warp's. -#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)] -#[serde(default)] -pub struct ApiKeys { - pub google: Option, - pub anthropic: Option, - pub openai: Option, - pub open_router: Option, - pub custom_endpoints: Vec, -} - -#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)] -#[serde(default)] -pub struct CustomEndpoint { - pub name: String, - pub url: String, - pub api_key: String, - pub models: Vec, -} - -#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)] -#[serde(default)] -pub struct CustomEndpointModel { - pub name: String, - pub alias: Option, - /// Stable identifier used as `ModelConfig.{base,coding,cli_agent,computer_use_agent}` and - /// as the `CustomModelProviders.providers[*].models[*].config_key` on the request wire. - /// Generated as a UUIDv4 at model creation. - pub config_key: String, -} - -impl CustomEndpointModel { - /// Picker label: prefer the user-provided alias; fall back to the raw model name - /// so a row is never blank. - pub fn display_label(&self) -> &str { - match self.alias.as_deref() { - Some(alias) if !alias.trim().is_empty() => alias, - _ => &self.name, - } - } -} - -impl ApiKeys { - pub fn has_any_key(&self) -> bool { - self.openai.is_some() - || self.anthropic.is_some() - || self.google.is_some() - || self.open_router.is_some() - || self - .custom_endpoints - .iter() - .any(|endpoint| !endpoint.api_key.trim().is_empty()) - } - - /// Number of single-provider API keys currently configured (OpenAI, - /// Anthropic, Google, OpenRouter). Custom endpoints are counted separately - /// via `custom_endpoints`. - pub fn provider_key_count(&self) -> usize { - [ - &self.openai, - &self.anthropic, - &self.google, - &self.open_router, - ] - .into_iter() - .filter(|key| key.as_deref().is_some_and(|v| !v.trim().is_empty())) - .count() - } -} - -/// OAuth tokens for a connected xAI / Grok subscription (e.g. SuperGrok). -/// -/// Persisted to secure storage under [`GROK_SECURE_STORAGE_KEY`], separate from -/// the BYO [`ApiKeys`] blob because these are OAuth tokens with a refresh -/// lifecycle rather than a user-pasted static key. `crate::grok_subscription` -/// owns refreshing them; this module is the storage and request-injection -/// source of truth that [`ApiKeyManager::api_keys_for_request`] reads from. -#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)] -pub struct GrokTokens { - pub access_token: String, - #[serde(default)] - pub refresh_token: Option, - /// Absolute time at which `access_token` expires, if the provider told us. - #[serde(default)] - pub expires_at: Option, - /// When the user originally connected the subscription (i.e. when the - /// browser OAuth flow completed). Carried over across token refreshes so - /// it keeps reflecting the initial connection, not the latest refresh; - /// surfaced in the settings UI as "Connected on ...". `None` for tokens - /// stored before this field existed. - #[serde(default)] - pub connected_at: Option, -} - -impl GrokTokens { - /// Returns the access token whenever it is non-empty, regardless of - /// expiry. Possibly-expired tokens are still sent so the server stays the - /// final authority on token validity (it rejects truly invalid tokens); - /// `crate::grok_subscription` refreshes (nearly) expired tokens in the - /// background. - pub fn access_token_for_request(&self) -> Option<&str> { - (!self.access_token.trim().is_empty()).then_some(self.access_token.as_str()) - } - - /// Returns `true` when the token is known to expire within `lead_time` and - /// should be proactively refreshed. Tokens with an unknown expiry never - /// report as needing a refresh (there's no expiry signal to act on). - pub fn needs_refresh(&self, lead_time: Duration) -> bool { - match self.expires_at { - Some(expires_at) => expires_at <= SystemTime::now() + lead_time, - None => false, - } - } -} - /// Controls how AWS credentials are refreshed by [`ApiKeyManager`]. #[derive(Debug, Clone, Default, PartialEq, Eq)] pub enum AwsCredentialsRefreshStrategy { @@ -159,187 +30,25 @@ pub enum AwsCredentialsRefreshStrategy { }, } -/// A structure that manages API keys for AI providers. +/// A structure that manages locally-held credentials used to authenticate AI +/// provider requests: AWS Bedrock credentials and Gemini Enterprise (GEAP) +/// credentials. pub struct ApiKeyManager { - keys: ApiKeys, - /// OAuth tokens for a connected xAI/Grok subscription, if any. Persisted - /// separately from `keys` under [`GROK_SECURE_STORAGE_KEY`]; - /// `crate::grok_subscription` keeps these fresh. - grok_tokens: Option, - /// Whether background refresh of `grok_tokens` is currently allowed. - /// Mirrors the BYO API key policy, which lives in the app layer; wired in - /// via `ApiKeyManager::set_grok_refresh_allowed` (`crate::grok_subscription`). - #[cfg(not(target_family = "wasm"))] - pub(crate) grok_refresh_allowed: bool, - /// Guards against overlapping Grok token refreshes: the proactive refresh - /// timer and the request-time safety net - /// (`ApiKeyManager::refresh_grok_tokens_if_needed`) can otherwise race. - #[cfg(not(target_family = "wasm"))] - pub(crate) grok_refresh_in_flight: bool, pub(crate) aws_credentials_state: AwsCredentialsState, aws_credentials_refresh_strategy: AwsCredentialsRefreshStrategy, /// In-memory Gemini Enterprise (GEAP) credential state. pub(crate) geap_credentials_state: GeapCredentialsState, - secure_storage_write_version: u64, - grok_secure_storage_write_version: u64, } impl ApiKeyManager { - pub fn new(ctx: &mut ModelContext) -> Self { - let keys = Self::load_keys_from_secure_storage(ctx); - let grok_tokens = Self::load_grok_tokens_from_secure_storage(ctx); + pub fn new(_ctx: &mut ModelContext) -> Self { Self { - keys, - grok_tokens, - #[cfg(not(target_family = "wasm"))] - grok_refresh_allowed: false, - #[cfg(not(target_family = "wasm"))] - grok_refresh_in_flight: false, aws_credentials_state: AwsCredentialsState::Missing, aws_credentials_refresh_strategy: AwsCredentialsRefreshStrategy::default(), geap_credentials_state: GeapCredentialsState::Missing, - secure_storage_write_version: 0, - grok_secure_storage_write_version: 0, } } - pub fn keys(&self) -> &ApiKeys { - &self.keys - } - - /// The currently stored xAI/Grok OAuth tokens, if the user has connected a - /// Grok subscription. - pub fn grok_tokens(&self) -> Option<&GrokTokens> { - self.grok_tokens.as_ref() - } - - /// Returns `true` when a Grok subscription is connected with a usable OAuth - /// access token. - pub fn has_grok_subscription(&self) -> bool { - self.grok_tokens - .as_ref() - .and_then(GrokTokens::access_token_for_request) - .is_some() - } - - /// Returns `true` when the user has any usable BYO credential: a pasted - /// provider or custom-endpoint key, or a connected Grok subscription. - pub fn has_any_key(&self) -> bool { - self.keys.has_any_key() || self.has_grok_subscription() - } - - /// Stores (or clears, with `None`) the xAI/Grok OAuth tokens and persists - /// them to secure storage. No-op when the value is unchanged so we don't - /// emit spurious events or schedule redundant keychain writes. - pub fn set_grok_tokens(&mut self, tokens: Option, ctx: &mut ModelContext) { - if self.grok_tokens == tokens { - return; - } - self.grok_tokens = tokens; - ctx.emit(ApiKeyManagerEvent::KeysUpdated); - self.write_grok_tokens_to_secure_storage(ctx); - } - - pub fn set_google_key(&mut self, key: Option, ctx: &mut ModelContext) { - self.keys.google = key; - ctx.emit(ApiKeyManagerEvent::KeysUpdated); - self.write_keys_to_secure_storage(ctx); - } - - pub fn set_anthropic_key(&mut self, key: Option, ctx: &mut ModelContext) { - self.keys.anthropic = key; - ctx.emit(ApiKeyManagerEvent::KeysUpdated); - self.write_keys_to_secure_storage(ctx); - } - - pub fn set_openai_key(&mut self, key: Option, ctx: &mut ModelContext) { - self.keys.openai = key; - ctx.emit(ApiKeyManagerEvent::KeysUpdated); - self.write_keys_to_secure_storage(ctx); - } - - pub fn set_open_router_key(&mut self, key: Option, ctx: &mut ModelContext) { - self.keys.open_router = key; - ctx.emit(ApiKeyManagerEvent::KeysUpdated); - self.write_keys_to_secure_storage(ctx); - } - - pub fn add_custom_endpoint( - &mut self, - name: String, - url: String, - api_key: String, - models: Vec<(String, Option, Option)>, - ctx: &mut ModelContext, - ) { - self.keys.custom_endpoints.push(CustomEndpoint { - name, - url, - api_key, - models: models - .into_iter() - .map(|(name, alias, config_key)| CustomEndpointModel { - name, - alias, - config_key: config_key - .filter(|k| !k.is_empty()) - .unwrap_or_else(|| Uuid::new_v4().to_string()), - }) - .collect(), - }); - ctx.emit(ApiKeyManagerEvent::KeysUpdated); - self.write_keys_to_secure_storage(ctx); - } - - pub fn save_custom_endpoint( - &mut self, - index: usize, - name: String, - url: String, - api_key: String, - models: Vec<(String, Option, Option)>, - ctx: &mut ModelContext, - ) { - if index >= self.keys.custom_endpoints.len() { - return; - } - self.keys.custom_endpoints[index] = CustomEndpoint { - name, - url, - api_key, - models: models - .into_iter() - .map(|(name, alias, config_key)| CustomEndpointModel { - name, - alias, - config_key: config_key - .filter(|k| !k.is_empty()) - .unwrap_or_else(|| Uuid::new_v4().to_string()), - }) - .collect(), - }; - ctx.emit(ApiKeyManagerEvent::KeysUpdated); - self.write_keys_to_secure_storage(ctx); - } - - pub fn remove_custom_endpoint(&mut self, index: usize, ctx: &mut ModelContext) { - if index >= self.keys.custom_endpoints.len() { - return; - } - self.keys.custom_endpoints.remove(index); - ctx.emit(ApiKeyManagerEvent::KeysUpdated); - self.write_keys_to_secure_storage(ctx); - } - - pub fn clear_custom_endpoints(&mut self, ctx: &mut ModelContext) { - if self.keys.custom_endpoints.is_empty() { - return; - } - self.keys.custom_endpoints.clear(); - ctx.emit(ApiKeyManagerEvent::KeysUpdated); - self.write_keys_to_secure_storage(ctx); - } - pub fn set_aws_credentials_state( &mut self, state: AwsCredentialsState, @@ -380,93 +89,14 @@ impl ApiKeyManager { self.aws_credentials_refresh_strategy = strategy; } - /// Builds the `CustomModelProviders` registry that ships with every agent request. - /// - /// Emits one [`CustomModelProvider`] per configured [`CustomEndpoint`], each populated with - /// all of its [`CustomEndpointModel`]s. The per-model `config_key` is what the server uses - /// to map a `ModelConfig.{base,coding,cli_agent,computer_use_agent}` selection back to a - /// user-provided endpoint, so it MUST be the same UUID we store locally. - /// - /// Returns `None` when custom models should not be included or no endpoint has both a - /// non-empty URL and API key. - pub fn custom_model_providers_for_request( - &self, - include_custom_models: bool, - ) -> Option { - if !include_custom_models { - return None; - } - - let providers: Vec<_> = self - .keys - .custom_endpoints - .iter() - .filter(|endpoint| !endpoint.url.trim().is_empty() && !endpoint.api_key.is_empty()) - .map( - |endpoint| api::request::settings::custom_model_providers::CustomModelProvider { - base_url: endpoint.url.clone(), - api_key: endpoint.api_key.clone(), - models: endpoint - .models - .iter() - .filter(|m| !m.name.trim().is_empty() && !m.config_key.is_empty()) - .map( - |m| api::request::settings::custom_model_providers::CustomModel { - slug: m.name.clone(), - config_key: m.config_key.clone(), - }, - ) - .collect(), - }, - ) - .filter(|provider| !provider.models.is_empty()) - .collect(); - - if providers.is_empty() { - None - } else { - Some(api::request::settings::CustomModelProviders { providers }) - } - } - + /// Builds the `ApiKeys` request payload carrying AWS Bedrock and/or Gemini + /// Enterprise (GEAP) credentials, when applicable. Returns `None` when + /// neither credential type applies to this request. pub fn api_keys_for_request( &self, - include_byo_keys: bool, include_aws_bedrock_credentials: bool, geap_binding: Option, ) -> Option { - let anthropic = include_byo_keys - .then(|| self.keys.anthropic.clone()) - .flatten() - .unwrap_or_default(); - let openai = include_byo_keys - .then(|| self.keys.openai.clone()) - .flatten() - .unwrap_or_default(); - let google = include_byo_keys - .then(|| self.keys.google.clone()) - .flatten() - .unwrap_or_default(); - let open_router = include_byo_keys - .then(|| self.keys.open_router.clone()) - .flatten() - .unwrap_or_default(); - - // The connected Grok subscription's OAuth access token is user-provided - // auth, just like a pasted BYO API key, so it respects the same BYO - // policy gate: when BYO keys are disabled (e.g. by workspace policy), - // the token must not be sent. Possibly-expired tokens ARE sent — the - // server is the authority on validity. - let grok_oauth_access_token = include_byo_keys - .then(|| { - self.grok_tokens - .as_ref() - .and_then(GrokTokens::access_token_for_request) - .map(str::to_owned) - }) - .flatten() - .unwrap_or_default(); - // Also include credentials when running with OIDC-managed Bedrock inference, regardless // of the per-user setting flag (which only applies to the local credential chain path). let include_aws = include_aws_bedrock_credentials @@ -506,130 +136,16 @@ impl ApiKeyManager { _ => None, }); - if anthropic.is_empty() - && openai.is_empty() - && google.is_empty() - && open_router.is_empty() - && grok_oauth_access_token.is_empty() - && aws_credentials.is_none() - && google_cloud_credentials.is_none() - { + if aws_credentials.is_none() && google_cloud_credentials.is_none() { None } else { Some(api::request::settings::ApiKeys { - anthropic, - openai, - google, - open_router, - grok_oauth_access_token, - allow_use_of_warp_credits: false, aws_credentials, google_cloud_credentials, + ..Default::default() }) } } - - fn load_keys_from_secure_storage(ctx: &mut ModelContext) -> ApiKeys { - let key_json = match ctx.secure_storage().read_value(SECURE_STORAGE_KEY) { - Ok(json) => json, - Err(e) => { - if !matches!(e, secure_storage::Error::NotFound) { - log::error!("Failed to read API keys from secure storage: {e:#}"); - } - return ApiKeys::default(); - } - }; - - match serde_json::from_str(&key_json) { - Ok(keys) => keys, - Err(e) => { - log::error!("Failed to deserialize API keys: {e:#}"); - ApiKeys::default() - } - } - } - - fn write_keys_to_secure_storage(&mut self, ctx: &mut ModelContext) { - let json = match serde_json::to_string(&self.keys) { - Ok(json) => json, - Err(e) => { - log::error!("Failed to serialize API keys: {e:#}"); - return; - } - }; - self.secure_storage_write_version += 1; - let write_version = self.secure_storage_write_version; - - // Defer the keychain write so it doesn't block the current event - // processing. The in-memory state is already updated and events - // already emitted, so the UI updates immediately while the - // potentially slow platform secure-storage call runs in a - // subsequent main-thread callback. Skip stale callbacks so older - // writes cannot complete after and overwrite a newer payload. - ctx.spawn(async move { json }, move |me, json, ctx| { - if write_version != me.secure_storage_write_version { - return; - } - if let Err(e) = ctx.secure_storage().write_value(SECURE_STORAGE_KEY, &json) { - log::error!("Failed to write API keys to secure storage: {e:#}"); - } - }); - } - - fn load_grok_tokens_from_secure_storage(ctx: &mut ModelContext) -> Option { - let json = match ctx.secure_storage().read_value(GROK_SECURE_STORAGE_KEY) { - Ok(json) => json, - Err(e) => { - if !matches!(e, secure_storage::Error::NotFound) { - log::error!("Failed to read Grok tokens from secure storage: {e:#}"); - } - return None; - } - }; - - match serde_json::from_str(&json) { - Ok(tokens) => Some(tokens), - Err(e) => { - log::error!("Failed to deserialize Grok tokens: {e:#}"); - None - } - } - } - - fn write_grok_tokens_to_secure_storage(&mut self, ctx: &mut ModelContext) { - // `Some(json)` writes the tokens; `None` removes the stored entry (the - // user disconnected). Serialize up front so the deferred callback only - // touches the keychain. - let payload = match self.grok_tokens.as_ref().map(serde_json::to_string) { - Some(Ok(json)) => Some(json), - Some(Err(e)) => { - log::error!("Failed to serialize Grok tokens: {e:#}"); - return; - } - None => None, - }; - self.grok_secure_storage_write_version += 1; - let write_version = self.grok_secure_storage_write_version; - - // Defer the keychain write/remove like `write_keys_to_secure_storage`, - // skipping stale callbacks so an older write can't clobber a newer one. - ctx.spawn(async move { payload }, move |me, payload, ctx| { - if write_version != me.grok_secure_storage_write_version { - return; - } - let result = match payload { - Some(ref json) => ctx - .secure_storage() - .write_value(GROK_SECURE_STORAGE_KEY, json), - None => ctx.secure_storage().remove_value(GROK_SECURE_STORAGE_KEY), - }; - if let Err(e) = result { - if !matches!(e, secure_storage::Error::NotFound) { - log::error!("Failed to persist Grok tokens to secure storage: {e:#}"); - } - } - }); - } } impl Entity for ApiKeyManager { diff --git a/crates/ai/src/api_keys_tests.rs b/crates/ai/src/api_keys_tests.rs index 87c89ed2..a6306198 100644 --- a/crates/ai/src/api_keys_tests.rs +++ b/crates/ai/src/api_keys_tests.rs @@ -2,41 +2,20 @@ use std::time::{Duration, SystemTime}; use super::*; -fn make_manager(keys: ApiKeys) -> ApiKeyManager { - make_manager_with_grok(keys, None) -} - -fn make_manager_with_grok(keys: ApiKeys, grok_tokens: Option) -> ApiKeyManager { +fn make_manager() -> ApiKeyManager { ApiKeyManager { - keys, - grok_tokens, - #[cfg(not(target_family = "wasm"))] - grok_refresh_allowed: false, - #[cfg(not(target_family = "wasm"))] - grok_refresh_in_flight: false, aws_credentials_state: AwsCredentialsState::Missing, aws_credentials_refresh_strategy: AwsCredentialsRefreshStrategy::default(), geap_credentials_state: GeapCredentialsState::Missing, - secure_storage_write_version: 0, - grok_secure_storage_write_version: 0, } } fn make_manager_with_geap(geap_credentials_state: GeapCredentialsState) -> ApiKeyManager { - let mut manager = make_manager(ApiKeys::default()); + let mut manager = make_manager(); manager.geap_credentials_state = geap_credentials_state; manager } -fn grok_tokens(access_token: &str, expires_in: Option) -> GrokTokens { - GrokTokens { - access_token: access_token.into(), - refresh_token: Some("refresh".into()), - expires_at: expires_in.map(|secs| SystemTime::now() + Duration::from_secs(secs)), - connected_at: None, - } -} - fn geap_credentials(access_token: &str, expires_in: Option) -> GeapCredentials { GeapCredentials::new( access_token.into(), @@ -56,8 +35,6 @@ fn geap_binding() -> GeapMintBinding { } } -// The expected binding the request build site passes in is the same type as -// the stored `minted_for`, so the attach check is a plain `==`. fn geap_gate() -> GeapMintBinding { geap_binding() } @@ -70,487 +47,6 @@ fn geap_loaded(access_token: &str, expires_in: Option) -> GeapCredentialsSt } } -fn endpoint( - name: &str, - url: &str, - api_key: &str, - models: &[(&str, Option<&str>)], -) -> CustomEndpoint { - endpoint_with_keys( - name, - url, - api_key, - &models - .iter() - .enumerate() - .map(|(i, (n, a))| (*n, *a, format!("cfg-{i}"))) - .collect::>() - .iter() - .map(|(n, a, k)| (*n, *a, k.as_str())) - .collect::>(), - ) -} - -fn endpoint_with_keys( - name: &str, - url: &str, - api_key: &str, - models: &[(&str, Option<&str>, &str)], -) -> CustomEndpoint { - CustomEndpoint { - name: name.into(), - url: url.into(), - api_key: api_key.into(), - models: models - .iter() - .map(|(n, a, cfg)| CustomEndpointModel { - name: (*n).into(), - alias: a.map(|s| s.into()), - config_key: (*cfg).into(), - }) - .collect(), - } -} - -// ── serde round-trip ──────────────────────────────────────────── - -#[test] -fn serde_round_trip_empty() { - let keys = ApiKeys::default(); - let json = serde_json::to_string(&keys).unwrap(); - let deser: ApiKeys = serde_json::from_str(&json).unwrap(); - assert_eq!(keys, deser); -} - -#[test] -fn serde_round_trip_with_provider_keys() { - let keys = ApiKeys { - openai: Some("sk-openai".into()), - anthropic: Some("sk-ant-abc".into()), - google: Some("AIzaSy123".into()), - open_router: Some("sk-or-xxx".into()), - custom_endpoints: vec![], - }; - let json = serde_json::to_string(&keys).unwrap(); - let deser: ApiKeys = serde_json::from_str(&json).unwrap(); - assert_eq!(keys, deser); -} - -#[test] -fn serde_round_trip_with_custom_endpoints() { - let keys = ApiKeys { - openai: None, - anthropic: None, - google: None, - open_router: None, - custom_endpoints: vec![ - endpoint("ep1", "https://a.io/v1", "key1", &[("gpt-4", Some("fast"))]), - endpoint( - "ep2", - "https://b.io/v1", - "key2", - &[("llama-70b", None), ("mixtral", Some("mix"))], - ), - ], - }; - let json = serde_json::to_string(&keys).unwrap(); - let deser: ApiKeys = serde_json::from_str(&json).unwrap(); - assert_eq!(keys, deser); -} - -#[test] -fn serde_ignores_unknown_fields() { - let json = r#"{"openai":"sk-x","unknown_field":"value","custom_endpoints":[]}"#; - let keys: ApiKeys = serde_json::from_str(json).unwrap(); - assert_eq!(keys.openai, Some("sk-x".into())); - assert!(keys.custom_endpoints.is_empty()); -} - -// ── has_any_key ───────────────────────────────────────────────── - -#[test] -fn has_any_key_false_when_empty() { - assert!(!ApiKeys::default().has_any_key()); -} - -#[test] -fn has_any_key_true_for_openai_only() { - let keys = ApiKeys { - openai: Some("sk-x".into()), - ..Default::default() - }; - assert!(keys.has_any_key()); -} - -#[test] -fn has_any_key_true_for_custom_endpoints_only() { - let keys = ApiKeys { - custom_endpoints: vec![endpoint("ep", "https://a.io", "key", &[("m", None)])], - ..Default::default() - }; - assert!(keys.has_any_key()); -} - -#[test] -fn has_any_key_false_for_endpoint_with_empty_api_key() { - let keys = ApiKeys { - custom_endpoints: vec![endpoint("ep", "https://a.io", "", &[("m", None)])], - ..Default::default() - }; - assert!(!keys.has_any_key()); -} - -// ── provider_key_count ───────────────────────────────────────── - -#[test] -fn provider_key_count_zero_when_empty() { - assert_eq!(ApiKeys::default().provider_key_count(), 0); -} - -#[test] -fn provider_key_count_counts_each_provider_key() { - let keys = ApiKeys { - openai: Some("sk-o".into()), - anthropic: Some("sk-a".into()), - google: Some("AIza".into()), - open_router: Some("sk-or".into()), - custom_endpoints: vec![], - }; - assert_eq!(keys.provider_key_count(), 4); -} - -#[test] -fn provider_key_count_ignores_blank_keys_and_endpoints() { - let keys = ApiKeys { - openai: Some("sk-o".into()), - anthropic: Some(" ".into()), - google: None, - open_router: None, - custom_endpoints: vec![endpoint("ep", "https://a.io", "k", &[("m", None)])], - }; - // Only the non-blank OpenAI key counts; the whitespace Anthropic key and the - // custom endpoint are excluded. - assert_eq!(keys.provider_key_count(), 1); -} - -// ── custom_model_providers_for_request ────────────────────────── - -#[test] -fn custom_model_providers_none_when_empty() { - let mgr = make_manager(ApiKeys::default()); - assert!(mgr.custom_model_providers_for_request(true).is_none()); -} - -#[test] -fn custom_model_providers_none_when_byo_disabled() { - let mgr = make_manager(ApiKeys { - custom_endpoints: vec![endpoint("ep", "https://a.io", "k", &[("m", None)])], - ..Default::default() - }); - assert!(mgr.custom_model_providers_for_request(false).is_none()); -} - -#[test] -fn custom_model_providers_populates_single_endpoint() { - let mgr = make_manager(ApiKeys { - custom_endpoints: vec![endpoint_with_keys( - "My EP", - "https://custom.io/v1", - "ep-key", - &[("big-model", Some("alias"), "uuid-1")], - )], - ..Default::default() - }); - let result = mgr.custom_model_providers_for_request(true).unwrap(); - assert_eq!(result.providers.len(), 1); - let p = &result.providers[0]; - assert_eq!(p.base_url, "https://custom.io/v1"); - assert_eq!(p.api_key, "ep-key"); - assert_eq!(p.models.len(), 1); - assert_eq!(p.models[0].slug, "big-model"); - assert_eq!(p.models[0].config_key, "uuid-1"); -} - -#[test] -fn multiple_endpoints_all_serialize() { - let mgr = make_manager(ApiKeys { - custom_endpoints: vec![ - endpoint_with_keys( - "ep1", - "https://a.io", - "k1", - &[("gpt-4", Some("fast"), "uuid-a")], - ), - endpoint_with_keys( - "ep2", - "https://b.io", - "k2", - &[ - ("llama-70b", None, "uuid-b"), - ("mixtral", Some("mix"), "uuid-c"), - ], - ), - ], - ..Default::default() - }); - let result = mgr.custom_model_providers_for_request(true).unwrap(); - assert_eq!(result.providers.len(), 2); - assert_eq!(result.providers[0].base_url, "https://a.io"); - assert_eq!(result.providers[0].models[0].config_key, "uuid-a"); - assert_eq!(result.providers[1].base_url, "https://b.io"); - assert_eq!(result.providers[1].models.len(), 2); - assert_eq!(result.providers[1].models[0].slug, "llama-70b"); - assert_eq!(result.providers[1].models[0].config_key, "uuid-b"); - assert_eq!(result.providers[1].models[1].config_key, "uuid-c"); -} - -#[test] -fn byok_disabled_returns_none_even_with_endpoints() { - let mgr = make_manager(ApiKeys { - custom_endpoints: vec![endpoint("ep", "https://a.io", "k", &[("m", None)])], - ..Default::default() - }); - assert!(mgr.custom_model_providers_for_request(false).is_none()); -} - -#[test] -fn empty_api_key_endpoints_are_skipped() { - let mgr = make_manager(ApiKeys { - custom_endpoints: vec![ - endpoint_with_keys("empty", "https://a.io", "", &[("m", None, "uuid-x")]), - endpoint_with_keys("ok", "https://b.io", "k", &[("m", None, "uuid-y")]), - ], - ..Default::default() - }); - let result = mgr.custom_model_providers_for_request(true).unwrap(); - assert_eq!(result.providers.len(), 1); - assert_eq!(result.providers[0].base_url, "https://b.io"); -} - -#[test] -fn endpoints_with_only_empty_models_are_skipped() { - let mgr = make_manager(ApiKeys { - custom_endpoints: vec![endpoint_with_keys( - "ep", - "https://a.io", - "k", - &[("", None, "uuid-z")], - )], - ..Default::default() - }); - assert!(mgr.custom_model_providers_for_request(true).is_none()); -} - -// ── display_label fallback ───────────────────────────────────── - -#[test] -fn display_label_uses_alias_when_present() { - let m = CustomEndpointModel { - name: "raw-name".into(), - alias: Some("My Alias".into()), - config_key: "k".into(), - }; - assert_eq!(m.display_label(), "My Alias"); -} - -#[test] -fn display_label_falls_back_to_name_when_alias_missing() { - let m = CustomEndpointModel { - name: "raw-name".into(), - alias: None, - config_key: "k".into(), - }; - assert_eq!(m.display_label(), "raw-name"); -} - -#[test] -fn display_label_falls_back_to_name_when_alias_is_whitespace() { - let m = CustomEndpointModel { - name: "raw-name".into(), - alias: Some(" ".into()), - config_key: "k".into(), - }; - assert_eq!(m.display_label(), "raw-name"); -} - -// ── api_keys_for_request ──────────────────────────────────────── - -#[test] -fn api_keys_for_request_none_when_empty() { - let mgr = make_manager(ApiKeys::default()); - assert!(mgr.api_keys_for_request(true, false, None).is_none()); -} - -#[test] -fn api_keys_for_request_populates_provider_keys() { - let mgr = make_manager(ApiKeys { - openai: Some("sk-o".into()), - anthropic: Some("sk-a".into()), - ..Default::default() - }); - let result = mgr.api_keys_for_request(true, false, None).unwrap(); - assert_eq!(result.openai, "sk-o"); - assert_eq!(result.anthropic, "sk-a"); - assert!(result.google.is_empty()); -} - -#[test] -fn api_keys_for_request_omits_keys_when_byo_disabled() { - let mgr = make_manager(ApiKeys { - openai: Some("sk-o".into()), - ..Default::default() - }); - // With BYO disabled and no other credentials, returns None. - assert!(mgr.api_keys_for_request(false, false, None).is_none()); -} - -#[test] -fn api_keys_for_request_none_for_custom_endpoints_only() { - let mgr = make_manager(ApiKeys { - custom_endpoints: vec![endpoint("ep", "https://a.io", "k", &[("m", None)])], - ..Default::default() - }); - assert!(mgr.api_keys_for_request(true, false, None).is_none()); -} - -// ── grok oauth token ──────────────────────────────────────────── - -#[test] -fn grok_access_token_present_without_expiry() { - let t = GrokTokens { - access_token: "tok".into(), - ..Default::default() - }; - assert_eq!(t.access_token_for_request(), Some("tok")); -} - -#[test] -fn grok_access_token_blank_is_none() { - let t = GrokTokens { - access_token: " ".into(), - ..Default::default() - }; - assert_eq!(t.access_token_for_request(), None); -} - -#[test] -fn grok_access_token_near_expiry_still_sent() { - // Expired tokens are still sent; the server is the authority on validity. - let t = grok_tokens("tok", Some(0)); - assert_eq!(t.access_token_for_request(), Some("tok")); -} - -#[test] -fn grok_access_token_far_future_is_some() { - let t = grok_tokens("tok", Some(3600)); - assert_eq!(t.access_token_for_request(), Some("tok")); -} - -#[test] -fn grok_needs_refresh_within_lead_time() { - assert!(grok_tokens("tok", Some(30)).needs_refresh(Duration::from_secs(300))); - assert!(!grok_tokens("tok", Some(3600)).needs_refresh(Duration::from_secs(300))); - // Expired tokens still need a refresh. - assert!(grok_tokens("tok", Some(0)).needs_refresh(Duration::from_secs(300))); - // Unknown expiry never reports as needing refresh. - assert!(!grok_tokens("tok", None).needs_refresh(Duration::from_secs(300))); -} - -#[test] -fn api_keys_for_request_includes_grok_token() { - let mgr = make_manager_with_grok( - ApiKeys::default(), - Some(grok_tokens("grok-abc", Some(3600))), - ); - let result = mgr.api_keys_for_request(true, false, None).unwrap(); - assert_eq!(result.grok_oauth_access_token, "grok-abc"); - assert!(result.anthropic.is_empty()); -} - -#[test] -fn api_keys_for_request_omits_grok_token_when_byo_disabled() { - // The Grok subscription is user-provided auth, so it follows the BYO - // policy gate: with BYO disabled and no other credentials, returns None. - let mgr = make_manager_with_grok( - ApiKeys::default(), - Some(grok_tokens("grok-abc", Some(3600))), - ); - assert!(mgr.api_keys_for_request(false, false, None).is_none()); -} - -#[test] -fn api_keys_for_request_includes_expired_grok_token() { - // Expired tokens are still sent in requests; the server rejects truly - // invalid ones and the background refresh replaces them. - let mgr = make_manager_with_grok(ApiKeys::default(), Some(grok_tokens("grok-abc", Some(0)))); - let result = mgr.api_keys_for_request(true, false, None).unwrap(); - assert_eq!(result.grok_oauth_access_token, "grok-abc"); -} - -#[test] -fn has_grok_subscription_false_when_not_connected() { - let mgr = make_manager(ApiKeys::default()); - assert!(!mgr.has_grok_subscription()); -} - -#[test] -fn has_grok_subscription_true_when_connected() { - let mgr = make_manager_with_grok( - ApiKeys::default(), - Some(grok_tokens("grok-abc", Some(3600))), - ); - assert!(mgr.has_grok_subscription()); -} - -#[test] -fn has_grok_subscription_true_for_expired_token() { - // A connected subscription still counts even when its token is past expiry: - // the token is sent anyway and the server is the authority on validity. - let mgr = make_manager_with_grok(ApiKeys::default(), Some(grok_tokens("grok-abc", Some(0)))); - assert!(mgr.has_grok_subscription()); -} - -#[test] -fn has_grok_subscription_false_when_token_blank() { - // A blank token can't be sent, so it does not count as a usable credential. - let mgr = make_manager_with_grok(ApiKeys::default(), Some(grok_tokens(" ", None))); - assert!(!mgr.has_grok_subscription()); -} - -// ── ApiKeyManager::has_any_key ────────────────── - -#[test] -fn manager_has_any_key_false_when_no_keys_and_no_grok() { - let mgr = make_manager(ApiKeys::default()); - assert!(!mgr.has_any_key()); -} - -#[test] -fn manager_has_any_key_true_for_pasted_key_without_grok() { - let mgr = make_manager(ApiKeys { - openai: Some("sk-x".into()), - ..Default::default() - }); - assert!(mgr.has_any_key()); -} - -#[test] -fn manager_has_any_key_true_for_connected_grok_without_pasted_key() { - // The crux: a connected Grok subscription counts even with no pasted keys, - // matching how it's sent as a BYO credential on requests. - let mgr = make_manager_with_grok( - ApiKeys::default(), - Some(grok_tokens("grok-abc", Some(3600))), - ); - assert!(mgr.has_any_key()); -} - -#[test] -fn manager_has_any_key_false_for_blank_grok_and_no_keys() { - let mgr = make_manager_with_grok(ApiKeys::default(), Some(grok_tokens(" ", None))); - assert!(!mgr.has_any_key()); -} - // ── geap credentials ──────────────────────────────────────────── #[test] @@ -587,13 +83,10 @@ fn geap_needs_refresh_lead_time_boundaries() { #[test] fn api_keys_for_request_includes_geap_token_when_gate_and_binding_match() { let mgr = make_manager_with_geap(geap_loaded("geap-abc", Some(3600))); - let result = mgr - .api_keys_for_request(false, false, Some(geap_gate())) - .unwrap(); + let result = mgr.api_keys_for_request(false, Some(geap_gate())).unwrap(); let credentials = result.google_cloud_credentials.unwrap(); assert_eq!(credentials.access_token, "geap-abc"); - // The GEAP token is independent of the BYO key gate. - assert!(result.anthropic.is_empty()); + assert!(result.aws_credentials.is_none()); } #[test] @@ -602,9 +95,7 @@ fn api_keys_for_request_includes_expired_geap_token() { // rejects truly invalid ones, which surfaces a recoverable error instead // of a silent fallback to another route. let mgr = make_manager_with_geap(geap_loaded("geap-abc", Some(0))); - let result = mgr - .api_keys_for_request(false, false, Some(geap_gate())) - .unwrap(); + let result = mgr.api_keys_for_request(false, Some(geap_gate())).unwrap(); assert_eq!( result.google_cloud_credentials.unwrap().access_token, "geap-abc" @@ -616,7 +107,7 @@ fn api_keys_for_request_omits_geap_token_without_gate() { // No gate (policy off at the call site) ⇒ no GEAP credentials, even when // a token is loaded. let mgr = make_manager_with_geap(geap_loaded("geap-abc", Some(3600))); - assert!(mgr.api_keys_for_request(false, false, None).is_none()); + assert!(mgr.api_keys_for_request(false, None).is_none()); } #[test] @@ -626,19 +117,19 @@ fn api_keys_for_request_omits_geap_token_on_binding_mismatch() { // A different user (sign-out/account switch). let mut gate = geap_gate(); gate.user_uid = "someone-else".into(); - assert!(mgr.api_keys_for_request(false, false, Some(gate)).is_none()); + assert!(mgr.api_keys_for_request(false, Some(gate)).is_none()); // A different audience (admin changed the pool/provider). let mut gate = geap_gate(); gate.audience = "//iam.googleapis.com/projects/2/locations/global/workloadIdentityPools/other/providers/other".into(); - assert!(mgr.api_keys_for_request(false, false, Some(gate)).is_none()); + assert!(mgr.api_keys_for_request(false, Some(gate)).is_none()); // A different service account (admin changed impersonation target). let mut gate = geap_gate(); gate.federation = GeapFederation::ServiceAccount { email: "other@proj.iam.gserviceaccount.com".into(), }; - assert!(mgr.api_keys_for_request(false, false, Some(gate)).is_none()); + assert!(mgr.api_keys_for_request(false, Some(gate)).is_none()); } #[test] @@ -648,9 +139,7 @@ fn api_keys_for_request_serves_previous_geap_token_while_refreshing() { let mgr = make_manager_with_geap(GeapCredentialsState::Refreshing { previous: Some((geap_credentials("geap-old", Some(10)), geap_binding())), }); - let result = mgr - .api_keys_for_request(false, false, Some(geap_gate())) - .unwrap(); + let result = mgr.api_keys_for_request(false, Some(geap_gate())).unwrap(); assert_eq!( result.google_cloud_credentials.unwrap().access_token, "geap-old" @@ -661,9 +150,7 @@ fn api_keys_for_request_serves_previous_geap_token_while_refreshing() { fn api_keys_for_request_omits_geap_token_during_first_mint() { // The very first mint has nothing to serve yet. let mgr = make_manager_with_geap(GeapCredentialsState::Refreshing { previous: None }); - assert!(mgr - .api_keys_for_request(false, false, Some(geap_gate())) - .is_none()); + assert!(mgr.api_keys_for_request(false, Some(geap_gate())).is_none()); } #[test] @@ -679,9 +166,7 @@ fn api_keys_for_request_omits_geap_token_for_non_loaded_states() { }, ] { let mgr = make_manager_with_geap(state); - assert!(mgr - .api_keys_for_request(false, false, Some(geap_gate())) - .is_none()); + assert!(mgr.api_keys_for_request(false, Some(geap_gate())).is_none()); } } @@ -692,5 +177,50 @@ fn api_keys_for_request_omits_geap_token_when_previous_binding_mismatches() { }); let mut gate = geap_gate(); gate.user_uid = "someone-else".into(); - assert!(mgr.api_keys_for_request(false, false, Some(gate)).is_none()); + assert!(mgr.api_keys_for_request(false, Some(gate)).is_none()); +} + +// ── aws credentials ───────────────────────────────────────────── + +#[test] +fn api_keys_for_request_none_when_nothing_configured() { + let mgr = make_manager(); + assert!(mgr.api_keys_for_request(false, None).is_none()); +} + +#[test] +fn api_keys_for_request_includes_aws_credentials_when_requested() { + let mut mgr = make_manager(); + mgr.aws_credentials_state = AwsCredentialsState::Loaded { + credentials: AwsCredentials::new("ak".into(), "sk".into(), None, None), + loaded_at: SystemTime::now(), + }; + let result = mgr.api_keys_for_request(true, None).unwrap(); + assert_eq!(result.aws_credentials.unwrap().access_key, "ak"); +} + +#[test] +fn api_keys_for_request_omits_aws_credentials_when_not_requested() { + let mut mgr = make_manager(); + mgr.aws_credentials_state = AwsCredentialsState::Loaded { + credentials: AwsCredentials::new("ak".into(), "sk".into(), None, None), + loaded_at: SystemTime::now(), + }; + assert!(mgr.api_keys_for_request(false, None).is_none()); +} + +#[test] +fn api_keys_for_request_includes_aws_credentials_when_oidc_managed_regardless_of_flag() { + let mut mgr = make_manager(); + mgr.aws_credentials_state = AwsCredentialsState::Loaded { + credentials: AwsCredentials::new("ak".into(), "sk".into(), None, None), + loaded_at: SystemTime::now(), + }; + mgr.aws_credentials_refresh_strategy = AwsCredentialsRefreshStrategy::OidcManaged { + task_id: Some("task-1".into()), + role_arn: "arn:aws:iam::123:role/test".into(), + region: "us-east-1".into(), + }; + let result = mgr.api_keys_for_request(false, None).unwrap(); + assert_eq!(result.aws_credentials.unwrap().access_key, "ak"); } diff --git a/crates/ai/src/grok_subscription/mod.rs b/crates/ai/src/grok_subscription/mod.rs deleted file mode 100644 index efb6781d..00000000 --- a/crates/ai/src/grok_subscription/mod.rs +++ /dev/null @@ -1,228 +0,0 @@ -//! Refresh orchestration for a connected xAI / Grok subscription's OAuth -//! tokens. -//! -//! The tokens themselves live in [`ApiKeyManager`] (the request-building -//! source of truth, persisted to secure storage under `GrokOAuthTokens`). -//! This module owns the network-facing refresh lifecycle — converting a -//! [`TokenResponse`] into stored [`GrokTokens`], proactively refreshing the -//! access token shortly before it expires, and rescheduling the next refresh. -//! -//! The Grok subscription is BYO auth, so background refresh follows the BYO -//! API key policy. That policy lives in the app layer (workspace settings), -//! which this crate has no visibility into; the app wires it in via -//! [`ApiKeyManager::set_grok_refresh_allowed`]. -//! -//! The network/protocol side of the connect flow (authorize URL, loopback -//! callback server, token exchange/refresh) lives in the [`oauth`] submodule. - -pub mod oauth; - -use std::time::{Duration, SystemTime}; - -use galaxyui_core::r#async::Timer; -use galaxyui_core::ModelContext; - -use self::oauth::TokenResponse; -use crate::api_keys::{ApiKeyManager, GrokTokens}; - -/// Refresh the access token this long before its hard expiry so a request -/// never races the expiration. Possibly-expired tokens are still sent (the -/// server is the authority on validity), so this lead time is purely about -/// keeping the token fresh, not about when it stops being sent. -const REFRESH_LEAD_TIME: Duration = Duration::from_secs(5 * 60); - -/// Builds [`GrokTokens`] from a token-endpoint [`TokenResponse`], computing the -/// absolute `expires_at` from the relative `expires_in`. Values not present in -/// the response are carried over from `previous`: the refresh token when xAI -/// doesn't return a new one (refresh-token rotation is optional in OAuth 2.0), -/// and `connected_at` so it keeps reflecting the initial connection time -/// (initialized to now when there are no previous tokens, i.e. a fresh -/// connect). -pub fn grok_tokens_from_response( - response: TokenResponse, - previous: Option<&GrokTokens>, -) -> GrokTokens { - let expires_at = response - .expires_in - .and_then(|secs| u64::try_from(secs).ok()) - .and_then(|secs| SystemTime::now().checked_add(Duration::from_secs(secs))); - GrokTokens { - access_token: response.access_token, - refresh_token: response - .refresh_token - .or_else(|| previous.and_then(|tokens| tokens.refresh_token.clone())), - expires_at, - connected_at: previous - .and_then(|tokens| tokens.connected_at) - .or_else(|| Some(SystemTime::now())), - } -} - -impl ApiKeyManager { - /// Persists freshly obtained tokens (e.g. right after the connect flow) and - /// schedules the next proactive refresh. - pub fn store_grok_tokens(&mut self, response: TokenResponse, ctx: &mut ModelContext) { - apply_grok_tokens(self, response, ctx); - } - - /// Updates whether background refresh of the stored Grok tokens is - /// allowed. The Grok subscription is BYO auth, so refresh follows the same - /// policy gate as request injection ([`Self::api_keys_for_request`]): - /// tokens that can never be sent shouldn't be kept fresh. The policy lives - /// in the app layer, which calls this at startup and whenever the policy - /// may have changed (e.g. team data arriving, or a workspace switch). - /// - /// Schedules a refresh on a disabled -> enabled transition (refreshing - /// immediately if the token has already (nearly) expired); in-flight - /// timers re-check the flag when they fire. Repeated calls with an - /// unchanged value are no-ops, so duplicate timers can't pile up. - pub fn set_grok_refresh_allowed(&mut self, allowed: bool, ctx: &mut ModelContext) { - if self.grok_refresh_allowed == allowed { - return; - } - self.grok_refresh_allowed = allowed; - if allowed { - schedule_grok_token_refresh(self, ctx); - } - } - - /// Request-time safety net: kicks off a background refresh of the stored - /// Grok tokens when they are nearing (or already past) expiry, so - /// upcoming requests can authenticate even if the proactive refresh loop - /// never armed or died (e.g. a stale BYO policy at startup, or an earlier - /// failed refresh). The triggering request still carries the currently - /// stored token — the server is the authority on its validity. - /// - /// `byo_allowed` is the BYO API key policy as freshly evaluated by the - /// caller at request time. It also re-syncs the stored policy mirror, - /// which can go stale between `TeamsChanged` events; a disabled -> - /// enabled transition re-arms the proactive refresh loop. - pub fn refresh_grok_tokens_if_needed( - &mut self, - byo_allowed: bool, - ctx: &mut ModelContext, - ) { - self.set_grok_refresh_allowed(byo_allowed, ctx); - if !byo_allowed || self.grok_refresh_in_flight { - return; - } - let Some(tokens) = self.grok_tokens() else { - return; - }; - if !tokens.needs_refresh(REFRESH_LEAD_TIME) { - return; - } - let Some(refresh_token) = tokens.refresh_token.clone() else { - return; - }; - log::info!( - "Grok OAuth token is nearing or past expiry at request time; refreshing in background" - ); - spawn_grok_refresh(self, refresh_token, ctx); - } -} - -/// Stores the tokens from `response` (carrying over the previous refresh token -/// and connection time when absent) and schedules the next proactive refresh. -fn apply_grok_tokens( - manager: &mut ApiKeyManager, - response: TokenResponse, - ctx: &mut ModelContext, -) { - let tokens = grok_tokens_from_response(response, manager.grok_tokens()); - manager.set_grok_tokens(Some(tokens), ctx); - schedule_grok_token_refresh(manager, ctx); -} - -/// Schedules a one-shot proactive refresh [`REFRESH_LEAD_TIME`] before the -/// current token's expiry (immediately if already within that window). -/// -/// No-op when there's nothing to refresh against (no tokens, no refresh token, -/// or no known expiry). Reschedules itself after each successful refresh, so a -/// single call establishes an ongoing refresh loop for the lifetime of the -/// connection. -fn schedule_grok_token_refresh(manager: &mut ApiKeyManager, ctx: &mut ModelContext) { - // When the BYO API key policy is disabled the token is never sent, so - // don't refresh it in the background either. `set_grok_refresh_allowed` - // re-establishes the loop if the policy is later enabled. - if !manager.grok_refresh_allowed { - return; - } - let Some(tokens) = manager.grok_tokens() else { - return; - }; - let Some(refresh_token) = tokens.refresh_token.clone() else { - return; - }; - let Some(expires_at) = tokens.expires_at else { - // No expiry signal, so there's nothing to schedule against. - return; - }; - - let now = SystemTime::now(); - let fire_at = expires_at.checked_sub(REFRESH_LEAD_TIME).unwrap_or(now); - let delay = fire_at.duration_since(now).unwrap_or(Duration::ZERO); - - ctx.spawn( - async move { - Timer::after(delay).await; - }, - move |manager, _output, ctx| { - // The BYO policy may have flipped off while we slept; - // `set_grok_refresh_allowed` restarts the loop if it flips back - // on. - if !manager.grok_refresh_allowed { - return; - } - // The stored token may have changed (reconnect/disconnect) while we - // slept; only refresh if our refresh token is still the current one. - let still_current = manager - .grok_tokens() - .and_then(|t| t.refresh_token.as_deref()) - == Some(refresh_token.as_str()); - if still_current { - spawn_grok_refresh(manager, refresh_token, ctx); - } - }, - ); -} - -/// Kicks off a background token refresh using `refresh_token`, applying the -/// result (which reschedules the next refresh) or logging the failure. -/// -/// No-op when a refresh is already in flight, so the proactive timer and the -/// request-time safety net can't issue overlapping refreshes. -fn spawn_grok_refresh( - manager: &mut ApiKeyManager, - refresh_token: String, - ctx: &mut ModelContext, -) { - if manager.grok_refresh_in_flight { - return; - } - manager.grok_refresh_in_flight = true; - ctx.spawn( - async move { oauth::refresh_access_token(&refresh_token).await }, - |manager, result, ctx| { - manager.grok_refresh_in_flight = false; - match result { - Ok(response) => { - log::info!( - "Refreshed Grok OAuth token (expires_in={:?}, has_refresh_token={})", - response.expires_in, - response.refresh_token.is_some(), - ); - apply_grok_tokens(manager, response, ctx); - } - Err(err) => { - // Leave the existing (possibly expired) token in place; the - // server remains the authority and will reject it if it's - // truly invalid. The request-time safety net - // (`ApiKeyManager::refresh_grok_tokens_if_needed`) retries - // on the next request. - log::error!("Failed to refresh Grok OAuth token: {err:#}"); - } - } - }, - ); -} diff --git a/crates/ai/src/grok_subscription/oauth.rs b/crates/ai/src/grok_subscription/oauth.rs deleted file mode 100644 index 226c9f40..00000000 --- a/crates/ai/src/grok_subscription/oauth.rs +++ /dev/null @@ -1,457 +0,0 @@ -//! OAuth flow for connecting an xAI / Grok subscription (e.g. SuperGrok) to -//! Warp, so users can "plug in" their subscription instead of pasting a -//! pay-as-you-go API key. -//! -//! This mirrors the public Grok-CLI desktop OAuth flow: an OAuth 2.0 -//! Authorization Code grant with PKCE and a fixed loopback redirect URI. xAI's -//! auth server only accepts the loopback redirect for an allowlisted -//! `client_id` bound to a specific port, so we reuse the Grok-CLI client and -//! bind the callback server to that exact port. -//! -//! Some browsers/networks can't reach the loopback callback (e.g. Private -//! Network Access is blocked), in which case xAI's consent screen instead -//! *displays* the authorization code for the user to paste back into the app. -//! [`OauthAttempt::manual_code_exchange`] supports that fallback by capturing -//! the attempt's PKCE verifier so a pasted code can be exchanged directly, -//! without ever observing the loopback redirect. -//! -//! This module owns only the network/protocol side: building the authorize -//! URL, running the loopback callback server, and exchanging/refreshing tokens -//! at xAI's token endpoint. Persistence of the resulting tokens, proactive -//! refresh scheduling, and injection into the request live in the parent -//! [`crate::grok_subscription`] module (refresh orchestration) and -//! [`crate::api_keys::ApiKeyManager`] (storage + request injection). - -use std::io::{ErrorKind, Read, Write}; -use std::net::{Shutdown, TcpListener, TcpStream}; -use std::time::Duration; - -use anyhow::{bail, Context as _}; -use base64::engine::general_purpose::URL_SAFE_NO_PAD; -use base64::Engine as _; -// `std::time::Instant` is disallowed (no wasm support); `instant::Instant` is a -// drop-in that re-exports the std type on native targets. -use instant::Instant; -use rand::RngCore as _; -use serde::Deserialize; -use sha2::{Digest, Sha256}; - -const CLIENT_ID: &str = "b1a00492-073a-47ea-816f-4c329264a828"; -const AUTHORIZE_URL: &str = "https://auth.x.ai/oauth2/authorize"; -const TOKEN_URL: &str = "https://auth.x.ai/oauth2/token"; -const SCOPE: &str = "openid profile email offline_access grok-cli:access api:access"; - -const REDIRECT_HOST: &str = "127.0.0.1"; -const REDIRECT_PORT: u16 = 56121; - -/// How long we keep the loopback server open waiting for the user to approve -/// the consent screen in their browser. -const CALLBACK_TIMEOUT: Duration = Duration::from_secs(300); -/// How long to nap between non-blocking `accept()` attempts. -const POLL_INTERVAL: Duration = Duration::from_millis(100); - -/// xAI's browser consent screen fetches the loopback callback from these -/// origins. Since that request crosses origins (https://accounts.x.ai -> -/// http://127.0.0.1), browsers require CORS and Private Network Access headers -/// before the page can observe the callback response. -const CORS_ALLOWED_ORIGINS: [&str; 2] = ["https://accounts.x.ai", "https://auth.x.ai"]; - -fn redirect_uri() -> String { - format!("http://{REDIRECT_HOST}:{REDIRECT_PORT}/callback") -} - -/// One in-flight OAuth login attempt: the bound loopback callback listener -/// plus the per-attempt PKCE/CSRF secrets, which never leave this module. -/// -/// Construct with [`OauthAttempt::start`], open [`OauthAttempt::authorize_url`] -/// in the browser, then await [`OauthAttempt::finish`] to obtain tokens. Tying -/// the secrets to the attempt guarantees the same PKCE verifier and CSRF state -/// are used for both the authorize URL and the code exchange. -pub struct OauthAttempt { - listener: TcpListener, - pkce: PkceParams, -} - -impl OauthAttempt { - /// Binds the loopback callback server and generates fresh per-attempt - /// secrets. Call this before opening the browser so a bind failure (e.g. - /// another login already in progress, or Grok-CLI holding the port) - /// surfaces before a browser tab opens. - pub fn start() -> anyhow::Result { - Ok(Self { - listener: bind_callback_listener()?, - pkce: PkceParams::generate(), - }) - } - - /// The authorization URL the user's browser should open to begin the flow. - pub fn authorize_url(&self) -> String { - authorize_url(&self.pkce) - } - - /// Runs the rest of the browser-based PKCE flow: waits for the loopback - /// callback, validates the CSRF state, and exchanges the authorization - /// code for tokens. Consumes the attempt so its secrets can't be reused. - pub async fn finish(self) -> anyhow::Result { - run_oauth_flow(self.listener, self.pkce).await - } - - /// Clones the PKCE verifier for the pasted-code fallback while the - /// loopback flow continues racing in parallel. - pub fn manual_code_exchange(&self) -> ManualCodeExchange { - ManualCodeExchange { - verifier: self.pkce.verifier.clone(), - } - } -} - -/// Completes OAuth from a manually-pasted authorization code. -/// -/// There is no redirect `state` to validate in this out-of-band path; PKCE -/// protects the exchange. -#[derive(Clone)] -pub struct ManualCodeExchange { - verifier: String, -} - -impl ManualCodeExchange { - /// Exchanges a user-pasted authorization `code` with the attempt's PKCE verifier. - pub async fn exchange(&self, code: &str) -> anyhow::Result { - let code = code.trim(); - if code.is_empty() { - bail!("enter the code shown in your browser to finish connecting"); - } - exchange_code_for_tokens(code, &self.verifier).await - } -} - -/// The per-attempt secrets for one authorization request: the PKCE -/// verifier/challenge pair and the CSRF `state` value. -struct PkceParams { - verifier: String, - challenge: String, - /// CSRF token echoed back on the redirect and validated against the - /// response before the code is exchanged. - state: String, -} - -impl PkceParams { - /// Generates a fresh PKCE verifier + S256 challenge and a random CSRF state. - fn generate() -> Self { - let verifier = random_url_safe_token(); - let challenge = URL_SAFE_NO_PAD.encode(Sha256::digest(verifier.as_bytes())); - let state = random_url_safe_token(); - Self { - verifier, - challenge, - state, - } - } -} - -/// Returns a URL-safe, unpadded base64 string of 32 random bytes. This is used -/// for both the PKCE code verifier (RFC 7636 allows 43-128 chars from the -/// unreserved set) and the CSRF state. -fn random_url_safe_token() -> String { - let mut bytes = [0u8; 32]; - rand::thread_rng().fill_bytes(&mut bytes); - URL_SAFE_NO_PAD.encode(bytes) -} - -/// Builds the authorization URL the user's browser should open to begin the -/// flow. -fn authorize_url(pkce: &PkceParams) -> String { - let redirect = redirect_uri(); - // `plan=generic` opts the consent screen into xAI's generic OAuth plan tier - // (required for loopback OAuth from non-allowlisted clients); `referrer` - // is best-effort attribution in xAI's OAuth logs. - let params: [(&str, &str); 9] = [ - ("response_type", "code"), - ("client_id", CLIENT_ID), - ("redirect_uri", &redirect), - ("scope", SCOPE), - ("code_challenge", &pkce.challenge), - ("code_challenge_method", "S256"), - ("state", &pkce.state), - ("plan", "generic"), - ("referrer", "warp"), - ]; - let query = - serde_urlencoded::to_string(params).expect("static OAuth params are always serializable"); - format!("{AUTHORIZE_URL}?{query}") -} - -/// The token endpoint's response. Fields beyond `access_token` are optional -/// because xAI does not always return them. Other response fields (e.g. -/// `token_type`, `scope`) are ignored since nothing consumes them. -#[derive(Debug, Deserialize)] -pub struct TokenResponse { - pub access_token: String, - #[serde(default)] - pub refresh_token: Option, - #[serde(default)] - pub expires_in: Option, -} - -/// The authorization code and state captured from the loopback redirect. -struct CallbackData { - code: String, - state: String, -} - -/// Binds the loopback callback server to the fixed redirect address. -fn bind_callback_listener() -> anyhow::Result { - let listener = TcpListener::bind((REDIRECT_HOST, REDIRECT_PORT)).with_context(|| { - format!( - "couldn't bind the Grok OAuth callback server to {REDIRECT_HOST}:{REDIRECT_PORT}. \ - Another login may be in progress, or another app (e.g. Grok CLI) is using the port." - ) - })?; - listener - .set_nonblocking(true) - .context("failed to set the Grok OAuth callback listener to non-blocking mode")?; - Ok(listener) -} - -/// Runs the full browser-based PKCE flow: waits for the loopback callback on a -/// dedicated thread, validates the CSRF state, and exchanges the authorization -/// code for tokens. -async fn run_oauth_flow(listener: TcpListener, pkce: PkceParams) -> anyhow::Result { - // The loopback accept loop is blocking, so run it on a dedicated OS thread - // and bridge the result back through a runtime-agnostic async channel. - let (tx, rx) = async_channel::bounded(1); - std::thread::Builder::new() - .name("grok-oauth-callback".to_owned()) - .spawn(move || { - // `send_blocking` is disallowed (no wasm support); block this - // dedicated thread on the async `send` instead. - let _ = galaxyui_core::r#async::block_on( - tx.send(wait_for_callback(&listener, CALLBACK_TIMEOUT)), - ); - }) - .context("failed to spawn the Grok OAuth callback server thread")?; - - let callback = rx - .recv() - .await - .context("the Grok OAuth callback server stopped unexpectedly")??; - - if callback.state != pkce.state { - bail!("the authorization response state did not match — aborting to prevent CSRF"); - } - - exchange_code_for_tokens(&callback.code, &pkce.verifier).await -} - -/// Blocks (on a non-blocking listener with polling) until the browser hits the -/// redirect URI, returning the captured code and state, or an error on timeout. -fn wait_for_callback(listener: &TcpListener, timeout: Duration) -> anyhow::Result { - let deadline = Instant::now() + timeout; - loop { - if Instant::now() >= deadline { - bail!("timed out waiting for the Grok authorization callback"); - } - match listener.accept() { - Ok((stream, _)) => match handle_callback_connection(stream)? { - Some(data) => return Ok(data), - // Unrelated request (e.g. /favicon.ico); keep waiting. - None => continue, - }, - Err(ref e) if e.kind() == ErrorKind::WouldBlock => { - std::thread::sleep(POLL_INTERVAL); - } - Err(e) => { - return Err(anyhow::Error::new(e).context("Grok OAuth callback accept failed")) - } - } - } -} - -/// Reads a single HTTP request from the callback connection, writes back a -/// minimal HTML response, and extracts the OAuth parameters. -/// -/// Returns `Ok(None)` for requests that aren't the OAuth callback (so the -/// caller keeps listening), `Ok(Some(..))` on a successful callback, and `Err` -/// when the provider reported an error or the callback was malformed. -fn handle_callback_connection(mut stream: TcpStream) -> anyhow::Result> { - // The accepted stream may inherit the listener's non-blocking flag on some - // platforms; force blocking reads with a timeout so we get the full request - // line without spinning. - stream.set_nonblocking(false).ok(); - stream.set_read_timeout(Some(Duration::from_secs(10))).ok(); - - let mut buf = [0u8; 8192]; - let n = stream - .read(&mut buf) - .context("failed to read the Grok OAuth callback request")?; - let request = String::from_utf8_lossy(&buf[..n]); - - let origin = request_header(&request, "Origin"); - - // The request line looks like: "GET /callback?code=...&state=... HTTP/1.1". - let mut request_line_parts = request - .lines() - .next() - .unwrap_or_default() - .split_whitespace(); - let method = request_line_parts.next().unwrap_or_default(); - let path = request_line_parts.next().unwrap_or_default(); - - if method == "OPTIONS" && path.starts_with("/callback") { - write_response(&mut stream, "204 No Content", "", origin.as_deref()); - return Ok(None); - } - - let Some(query) = path - .strip_prefix("/callback") - .and_then(|rest| rest.strip_prefix('?')) - else { - write_response( - &mut stream, - "404 Not Found", - "Not found.", - origin.as_deref(), - ); - return Ok(None); - }; - - let mut code = None; - let mut state = None; - let mut error = None; - let mut error_description = None; - let pairs: Vec<(String, String)> = serde_urlencoded::from_str(query).unwrap_or_default(); - for (key, value) in pairs { - match key.as_str() { - "code" => code = Some(value), - "state" => state = Some(value), - "error" => error = Some(value), - "error_description" => error_description = Some(value), - _ => {} - } - } - - if let Some(error) = error { - write_response( - &mut stream, - "400 Bad Request", - FAILURE_HTML, - origin.as_deref(), - ); - let detail = error_description.unwrap_or(error); - bail!("Grok authorization was denied or failed: {detail}"); - } - - let (Some(code), Some(state)) = (code, state) else { - write_response( - &mut stream, - "400 Bad Request", - FAILURE_HTML, - origin.as_deref(), - ); - bail!("the Grok authorization callback was missing the code or state parameter"); - }; - write_response(&mut stream, "200 OK", SUCCESS_HTML, origin.as_deref()); - Ok(Some(CallbackData { code, state })) -} -fn request_header(request: &str, header_name: &str) -> Option { - request.lines().skip(1).find_map(|line| { - let (name, value) = line.split_once(':')?; - name.eq_ignore_ascii_case(header_name) - .then(|| value.trim().to_owned()) - }) -} - -/// Writes a minimal HTTP/1.1 response and closes the connection. -fn write_response(stream: &mut TcpStream, status: &str, body: &str, origin: Option<&str>) { - let cors_headers = cors_headers(origin); - let response = format!( - "HTTP/1.1 {status}\r\nContent-Type: text/html; charset=utf-8\r\n\ - {cors_headers}Content-Length: {}\r\nConnection: close\r\n\r\n{body}", - body.len() - ); - let _ = stream.write_all(response.as_bytes()); - let _ = stream.flush(); - let _ = stream.shutdown(Shutdown::Both); -} - -fn cors_headers(origin: Option<&str>) -> String { - origin - .filter(|origin| CORS_ALLOWED_ORIGINS.contains(origin)) - .map(|origin| { - format!( - "Access-Control-Allow-Origin: {origin}\r\n\ - Access-Control-Allow-Methods: GET, OPTIONS\r\n\ - Access-Control-Allow-Headers: Content-Type\r\n\ - Access-Control-Allow-Private-Network: true\r\n\ - Vary: Origin\r\n" - ) - }) - .unwrap_or_default() -} - -/// Exchanges the authorization code for OAuth tokens at xAI's token endpoint. -async fn exchange_code_for_tokens(code: &str, verifier: &str) -> anyhow::Result { - let redirect = redirect_uri(); - let form: [(&str, &str); 5] = [ - ("grant_type", "authorization_code"), - ("code", code), - ("redirect_uri", &redirect), - ("client_id", CLIENT_ID), - ("code_verifier", verifier), - ]; - post_token_request(&form).await -} - -/// Exchanges a previously obtained refresh token for a fresh set of tokens via -/// the OAuth 2.0 `refresh_token` grant. Used to keep the connected Grok -/// subscription's access token valid without re-running the browser flow. -/// -/// xAI may or may not return a new `refresh_token`; callers should fall back to -/// the existing one when [`TokenResponse::refresh_token`] is `None` (rotation is -/// optional in OAuth 2.0). -pub async fn refresh_access_token(refresh_token: &str) -> anyhow::Result { - let form: [(&str, &str); 3] = [ - ("grant_type", "refresh_token"), - ("refresh_token", refresh_token), - ("client_id", CLIENT_ID), - ]; - post_token_request(&form).await -} - -/// POSTs a form-encoded body to xAI's token endpoint and parses the -/// [`TokenResponse`]. Shared by the initial code exchange and refresh grants. -async fn post_token_request( - form: &T, -) -> anyhow::Result { - let response = http_client::Client::new() - .post(TOKEN_URL) - .form(form) - .send() - .await - .context("failed to send the Grok token request")?; - - let status = response.status(); - if !status.is_success() { - let body = response.text().await.unwrap_or_default(); - bail!("Grok token request failed ({status}): {body}"); - } - - response - .json::() - .await - .context("failed to parse the Grok token response") -} - -const SUCCESS_HTML: &str = "\ -Warp — Grok connected\ -\ -

Grok connected

You can close this window and return to Warp.

"; - -const FAILURE_HTML: &str = "\ -Warp — Grok authorization failed\ -\ -

Authorization failed

Something went wrong. Return to Warp and try again.

"; - -#[cfg(test)] -#[path = "oauth_tests.rs"] -mod tests; diff --git a/crates/ai/src/grok_subscription/oauth_tests.rs b/crates/ai/src/grok_subscription/oauth_tests.rs deleted file mode 100644 index 5fdbbb9d..00000000 --- a/crates/ai/src/grok_subscription/oauth_tests.rs +++ /dev/null @@ -1,57 +0,0 @@ -use super::*; - -#[test] -fn authorize_url_contains_required_params() { - let pkce = PkceParams::generate(); - let url = authorize_url(&pkce); - - assert!(url.starts_with("https://auth.x.ai/oauth2/authorize?")); - assert!(url.contains("response_type=code")); - assert!(url.contains(&format!("client_id={CLIENT_ID}"))); - assert!(url.contains("code_challenge_method=S256")); - assert!(url.contains("scope=openid")); - assert!(url.contains("plan=generic")); - assert!(url.contains("referrer=warp")); - // The redirect URI must be percent-encoded and match the registered value. - assert!(url.contains("redirect_uri=http%3A%2F%2F127.0.0.1%3A56121%2Fcallback")); - // The CSRF state and PKCE challenge are echoed into the URL verbatim - // (both are URL-safe base64, so no percent-encoding is applied). - assert!(url.contains(&format!("state={}", pkce.state))); - assert!(url.contains(&format!("code_challenge={}", pkce.challenge))); -} - -#[test] -fn token_response_parses_minimal_and_full() { - let minimal: TokenResponse = - serde_json::from_str(r#"{"access_token":"abc"}"#).expect("minimal response should parse"); - assert_eq!(minimal.access_token, "abc"); - assert!(minimal.refresh_token.is_none()); - assert!(minimal.expires_in.is_none()); - - // Unconsumed response fields (token_type, scope) are ignored by serde. - let full: TokenResponse = serde_json::from_str( - r#"{"access_token":"a","refresh_token":"r","token_type":"Bearer","expires_in":3600,"scope":"api:access"}"#, - ) - .expect("full response should parse"); - assert_eq!(full.access_token, "a"); - assert_eq!(full.refresh_token.as_deref(), Some("r")); - assert_eq!(full.expires_in, Some(3600)); -} - -#[test] -fn manual_code_exchange_captures_attempt_verifier() { - let pkce = PkceParams::generate(); - let exchange = ManualCodeExchange { - verifier: pkce.verifier.clone(), - }; - assert_eq!(exchange.verifier, pkce.verifier); -} - -#[test] -fn manual_code_exchange_rejects_blank_code() { - let exchange = ManualCodeExchange { - verifier: "verifier".to_string(), - }; - let result = galaxyui_core::r#async::block_on(exchange.exchange(" ")); - assert!(result.is_err()); -} diff --git a/crates/ai/src/lib.rs b/crates/ai/src/lib.rs index ea06b554..e7c0dfdc 100644 --- a/crates/ai/src/lib.rs +++ b/crates/ai/src/lib.rs @@ -2,8 +2,6 @@ pub mod agent; pub mod api_keys; pub mod aws_credentials; pub mod geap_credentials; -#[cfg(not(target_family = "wasm"))] -pub mod grok_subscription; pub mod llm_id; pub use llm_id::LLMId;