public/galaxy
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Remove Grok OAuth and legacy BYOK support

Browse Source
This commit is contained in:
Ryan Ward
2026-07-28 02:30:22 -05:00
parent 2faeed7ac5
commit 87e0c83e9e
26 changed files with 153 additions and 4140 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/src/ai/agent/api.rs
+4 -11
View File
@@ -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<warp_multi_agent_api::request::settings::ApiKeys>,
/// 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<warp_multi_agent_api::request::settings::CustomModelRouters>,
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,
app/src/ai/agent/api/impl.rs
+1 -23
View File
@@ -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<api::request::settings::ApiKeys>,
allow_use_of_warp_credits: bool,
) -> Option<api::request::settings::ApiKeys> {
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
}
app/src/ai/agent/api/impl_tests.rs
+1 -46
View File
@@ -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<HostId>) -> 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));
app/src/ai/agent/conversation_tests.rs
+2 -259
View File
@@ -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<AgentConversationData>) -> 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.
app/src/ai/blocklist/block/view_impl/common.rs
+2 -2
View File
@@ -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),
});
})
app/src/ai/blocklist/controller.rs
+19 -41
View File
@@ -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::<Vec<_>>()
@@ -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(
app/src/ai/blocklist/prompt/prompt_alert.rs
-10
View File
@@ -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 => {}
}
app/src/ai/llms.rs
+16 -77
View File
@@ -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<LLMInfo> = std::sync::LazyLock::new(|| {
LLMInfo {
static NO_PROVIDER_FALLBACK: std::sync::LazyLock<LLMInfo> =
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<LLMInfo> {
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,
&[
app/src/ai/llms_tests.rs
-455
View File
@@ -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<CustomEndpointModel>,
) -> 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<u32>,
max_output_tokens: Option<u32>,
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)
);
});
});
}
app/src/ai/request_usage_model.rs
-9
View File
@@ -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 {
app/src/ai/request_usage_model_tests.rs
+3 -180
View File
@@ -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 {
app/src/lib.rs
-18
View File
@@ -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);
app/src/settings/ai.rs
+1 -14
View File
@@ -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 {
app/src/settings_view/ai_page.rs
+25 -1197
View File
File diff suppressed because it is too large Load Diff
app/src/settings_view/mcp_servers/list_page.rs
+1 -1
View File
@@ -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 {
app/src/settings_view/mod.rs
+1 -2
View File
@@ -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";
app/src/settings_view/mod_tests.rs
+1 -1
View File
@@ -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");
app/src/terminal/input/models/data_source.rs
-19
View File
@@ -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),
app/src/workspace/one_time_modal_model.rs
+4 -2
View File
@@ -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;
app/src/workspace/view.rs
+3 -6
View File
@@ -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