Remove Grok OAuth and legacy BYOK support
This commit is contained in:
+4
-11
@@ -112,7 +112,8 @@ pub struct RequestParams {
|
|||||||
pub planning_enabled: bool,
|
pub planning_enabled: bool,
|
||||||
should_redact_secrets: 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>,
|
pub api_keys: Option<warp_multi_agent_api::request::settings::ApiKeys>,
|
||||||
/// User-provided custom model providers (BYOK endpoints).
|
/// User-provided custom model providers (BYOK endpoints).
|
||||||
pub custom_model_providers:
|
pub custom_model_providers:
|
||||||
@@ -121,7 +122,6 @@ pub struct RequestParams {
|
|||||||
/// `custom_model_providers`: the selected model's `config_key` indexes into this
|
/// `custom_model_providers`: the selected model's `config_key` indexes into this
|
||||||
/// registry. `None` when no custom router is selected.
|
/// registry. `None` when no custom router is selected.
|
||||||
pub custom_model_routers: Option<warp_multi_agent_api::request::settings::CustomModelRouters>,
|
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 autonomy_level: warp_multi_agent_api::AutonomyLevel,
|
||||||
pub isolation_level: warp_multi_agent_api::IsolationLevel,
|
pub isolation_level: warp_multi_agent_api::IsolationLevel,
|
||||||
pub web_search_enabled: bool,
|
pub web_search_enabled: bool,
|
||||||
@@ -203,7 +203,6 @@ impl RequestParams {
|
|||||||
api_keys: None,
|
api_keys: None,
|
||||||
custom_model_providers: None,
|
custom_model_providers: None,
|
||||||
custom_model_routers: None,
|
custom_model_routers: None,
|
||||||
allow_use_of_warp_credits: false,
|
|
||||||
autonomy_level: Default::default(),
|
autonomy_level: Default::default(),
|
||||||
isolation_level: Default::default(),
|
isolation_level: Default::default(),
|
||||||
web_search_enabled: false,
|
web_search_enabled: false,
|
||||||
@@ -305,16 +304,12 @@ impl RequestParams {
|
|||||||
|
|
||||||
let user_workspaces = UserWorkspaces::as_ref(app);
|
let user_workspaces = UserWorkspaces::as_ref(app);
|
||||||
let api_key_manager = ApiKeyManager::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"))]
|
#[cfg(not(target_family = "wasm"))]
|
||||||
let geap_binding = crate::ai::geap_credentials::current_geap_policy(app).mint_binding();
|
let geap_binding = crate::ai::geap_credentials::current_geap_policy(app).mint_binding();
|
||||||
#[cfg(target_family = "wasm")]
|
#[cfg(target_family = "wasm")]
|
||||||
let geap_binding: Option<::ai::api_keys::GeapMintBinding> = None;
|
let geap_binding: Option<::ai::api_keys::GeapMintBinding> = None;
|
||||||
let api_keys = api_key_manager.api_keys_for_request(
|
let api_keys = api_key_manager
|
||||||
is_byo_enabled,
|
.api_keys_for_request(user_workspaces.is_bedrock_enabled(app), geap_binding);
|
||||||
user_workspaces.is_bedrock_enabled(app),
|
|
||||||
geap_binding,
|
|
||||||
);
|
|
||||||
let custom_model_providers = None;
|
let custom_model_providers = None;
|
||||||
let custom_model_routers = FeatureFlag::CustomModelRouters.is_enabled().then(|| {
|
let custom_model_routers = FeatureFlag::CustomModelRouters.is_enabled().then(|| {
|
||||||
LLMPreferences::as_ref(app).custom_model_routers_for_request(
|
LLMPreferences::as_ref(app).custom_model_routers_for_request(
|
||||||
@@ -322,7 +317,6 @@ impl RequestParams {
|
|||||||
&request_input.coding_model_id,
|
&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 app_execution_mode = AppExecutionMode::as_ref(app);
|
||||||
let autonomy_level = if app_execution_mode.is_autonomous() {
|
let autonomy_level = if app_execution_mode.is_autonomous() {
|
||||||
@@ -399,7 +393,6 @@ impl RequestParams {
|
|||||||
api_keys,
|
api_keys,
|
||||||
custom_model_providers,
|
custom_model_providers,
|
||||||
custom_model_routers,
|
custom_model_routers,
|
||||||
allow_use_of_warp_credits,
|
|
||||||
autonomy_level,
|
autonomy_level,
|
||||||
isolation_level,
|
isolation_level,
|
||||||
web_search_enabled,
|
web_search_enabled,
|
||||||
|
|||||||
@@ -57,11 +57,6 @@ pub async fn generate_multi_agent_output(
|
|||||||
redaction::redact_inputs(&mut params.input);
|
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 {
|
let mut request = api::Request {
|
||||||
task_context: Some(api::request::TaskContext {
|
task_context: Some(api::request::TaskContext {
|
||||||
tasks: params.tasks,
|
tasks: params.tasks,
|
||||||
@@ -91,7 +86,7 @@ pub async fn generate_multi_agent_output(
|
|||||||
supports_suggest_prompt: true,
|
supports_suggest_prompt: true,
|
||||||
supports_read_image_files: FeatureFlag::ReadImageFiles.is_enabled(),
|
supports_read_image_files: FeatureFlag::ReadImageFiles.is_enabled(),
|
||||||
supports_reasoning_message: true,
|
supports_reasoning_message: true,
|
||||||
api_keys,
|
api_keys: params.api_keys,
|
||||||
autonomy_level: params.autonomy_level.into(),
|
autonomy_level: params.autonomy_level.into(),
|
||||||
isolation_level: params.isolation_level.into(),
|
isolation_level: params.isolation_level.into(),
|
||||||
web_search_enabled: params.web_search_enabled,
|
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 {
|
fn supports_orchestration_v2(orchestration_enabled: bool) -> bool {
|
||||||
orchestration_enabled
|
orchestration_enabled
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -2,10 +2,7 @@ use galaxy_core::features::FeatureFlag;
|
|||||||
use galaxy_core::HostId;
|
use galaxy_core::HostId;
|
||||||
use warp_multi_agent_api as api;
|
use warp_multi_agent_api as api;
|
||||||
|
|
||||||
use super::{
|
use super::{get_supported_cli_agent_tools, get_supported_tools, supports_orchestration_v2};
|
||||||
api_keys_with_warp_credit_fallback_setting, get_supported_cli_agent_tools, get_supported_tools,
|
|
||||||
supports_orchestration_v2,
|
|
||||||
};
|
|
||||||
use crate::ai::agent::api::RequestParams;
|
use crate::ai::agent::api::RequestParams;
|
||||||
use crate::ai::blocklist::SessionContext;
|
use crate::ai::blocklist::SessionContext;
|
||||||
use crate::ai::llms::LLMId;
|
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,
|
api_keys: None,
|
||||||
custom_model_providers: None,
|
custom_model_providers: None,
|
||||||
custom_model_routers: None,
|
custom_model_routers: None,
|
||||||
allow_use_of_warp_credits: false,
|
|
||||||
autonomy_level: api::AutonomyLevel::Supervised,
|
autonomy_level: api::AutonomyLevel::Supervised,
|
||||||
isolation_level: api::IsolationLevel::None,
|
isolation_level: api::IsolationLevel::None,
|
||||||
web_search_enabled: false,
|
web_search_enabled: false,
|
||||||
@@ -64,47 +60,6 @@ fn request_params_for_remote(host_id: Option<HostId>) -> RequestParams {
|
|||||||
params
|
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]
|
#[test]
|
||||||
fn supports_orchestration_v2_matches_request_orchestration_setting() {
|
fn supports_orchestration_v2_matches_request_orchestration_setting() {
|
||||||
assert!(supports_orchestration_v2(true));
|
assert!(supports_orchestration_v2(true));
|
||||||
|
|||||||
@@ -1,23 +1,14 @@
|
|||||||
use std::collections::HashMap;
|
use std::collections::HashMap;
|
||||||
|
|
||||||
use ai::api_keys::ApiKeyManager;
|
|
||||||
use galaxy_core::features::FeatureFlag;
|
use galaxy_core::features::FeatureFlag;
|
||||||
use warp_multi_agent_api as api;
|
use warp_multi_agent_api as api;
|
||||||
use warpui::{App, SingletonEntity};
|
|
||||||
|
|
||||||
use super::{
|
use super::{
|
||||||
artifact_from_fork_proto, footer_model_token_usage, AIConversation,
|
artifact_from_fork_proto, AIConversation, AIConversationAutoexecuteMode, AIConversationId,
|
||||||
AIConversationAutoexecuteMode, AIConversationId, ConversationStatus, RestoreConversationError,
|
ConversationStatus, RestoreConversationError,
|
||||||
};
|
};
|
||||||
use crate::ai::artifacts::Artifact;
|
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::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 {
|
fn restored_conversation(conversation_data: Option<AgentConversationData>) -> AIConversation {
|
||||||
AIConversation::new_restored(
|
AIConversation::new_restored(
|
||||||
@@ -115,43 +106,6 @@ fn restored_conversation_with_queries(queries: &[&str]) -> AIConversation {
|
|||||||
.unwrap()
|
.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]
|
#[test]
|
||||||
fn latest_user_query_returns_latest_non_empty_user_query() {
|
fn latest_user_query_returns_latest_non_empty_user_query() {
|
||||||
let conversation =
|
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());
|
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
|
/// The legacy `AgentConversationData.root_task_is_optimistic` flag must be
|
||||||
/// ignored on restore. A non-empty task list always produces a real
|
/// ignored on restore. A non-empty task list always produces a real
|
||||||
/// server-backed root regardless of whether the flag is set.
|
/// server-backed root regardless of whether the flag is set.
|
||||||
|
|||||||
@@ -3318,12 +3318,12 @@ fn render_invalid_api_key_error(
|
|||||||
background: Some(internal_colors::fg_overlay_3(theme).into()),
|
background: Some(internal_colors::fg_overlay_3(theme).into()),
|
||||||
..Default::default()
|
..Default::default()
|
||||||
})
|
})
|
||||||
.with_text_label("Edit API Keys".to_string())
|
.with_text_label("Manage Providers".to_string())
|
||||||
.with_cursor(Some(Cursor::PointingHand))
|
.with_cursor(Some(Cursor::PointingHand))
|
||||||
.build()
|
.build()
|
||||||
.on_click(move |ctx, _, _| {
|
.on_click(move |ctx, _, _| {
|
||||||
ctx.dispatch_typed_action(WorkspaceAction::ShowSettingsPageWithSearch {
|
ctx.dispatch_typed_action(WorkspaceAction::ShowSettingsPageWithSearch {
|
||||||
search_query: "api keys".to_string(),
|
search_query: "bedrock openai litellm".to_string(),
|
||||||
section: Some(SettingsSection::WarpAgent),
|
section: Some(SettingsSection::WarpAgent),
|
||||||
});
|
});
|
||||||
})
|
})
|
||||||
|
|||||||
@@ -671,8 +671,7 @@ impl BlocklistAIController {
|
|||||||
OrchestrationEventStreamerEvent::ChildSpawned { .. }
|
OrchestrationEventStreamerEvent::ChildSpawned { .. }
|
||||||
| OrchestrationEventStreamerEvent::ChildStatusChanged { .. } => {}
|
| OrchestrationEventStreamerEvent::ChildStatusChanged { .. } => {}
|
||||||
});
|
});
|
||||||
let crosscheck_reviewer =
|
let crosscheck_reviewer = ctx.add_model(crate::ai::crosscheck::CrosscheckReviewer::new);
|
||||||
ctx.add_model(crate::ai::crosscheck::CrosscheckReviewer::new);
|
|
||||||
ctx.subscribe_to_model(&crosscheck_reviewer, move |me, _, event, ctx| {
|
ctx.subscribe_to_model(&crosscheck_reviewer, move |me, _, event, ctx| {
|
||||||
use crate::ai::crosscheck::{CrosscheckReviewerEvent, ReviewOutcome};
|
use crate::ai::crosscheck::{CrosscheckReviewerEvent, ReviewOutcome};
|
||||||
let CrosscheckReviewerEvent::ReviewCompleted {
|
let CrosscheckReviewerEvent::ReviewCompleted {
|
||||||
@@ -2172,9 +2171,7 @@ impl BlocklistAIController {
|
|||||||
|
|
||||||
match outcome {
|
match outcome {
|
||||||
ReviewOutcome::Approved => {
|
ReviewOutcome::Approved => {
|
||||||
log::info!(
|
log::info!("[crosscheck] Work approved for conversation {conversation_id:?}");
|
||||||
"[crosscheck] Work approved for conversation {conversation_id:?}"
|
|
||||||
);
|
|
||||||
// Nothing to do — the conversation completes normally.
|
// Nothing to do — the conversation completes normally.
|
||||||
}
|
}
|
||||||
ReviewOutcome::MaxIterationsReached { last_feedback } => {
|
ReviewOutcome::MaxIterationsReached { last_feedback } => {
|
||||||
@@ -2191,9 +2188,7 @@ impl BlocklistAIController {
|
|||||||
self.inject_crosscheck_feedback(conversation_id, message, false, ctx);
|
self.inject_crosscheck_feedback(conversation_id, message, false, ctx);
|
||||||
}
|
}
|
||||||
ReviewOutcome::Error { error } => {
|
ReviewOutcome::Error { error } => {
|
||||||
log::error!(
|
log::error!("[crosscheck] Reviewer failed for {conversation_id:?}: {error}");
|
||||||
"[crosscheck] Reviewer failed for {conversation_id:?}: {error}"
|
|
||||||
);
|
|
||||||
// Don't block the conversation on reviewer errors; just log it.
|
// 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();
|
let max_iterations = ai_settings.crosscheck_max_iterations();
|
||||||
|
|
||||||
self.crosscheck_reviewer.update(ctx, |reviewer, ctx| {
|
self.crosscheck_reviewer.update(ctx, |reviewer, ctx| {
|
||||||
reviewer.start_review(
|
reviewer.start_review(conversation_id, max_iterations, agent_output, model_id, ctx);
|
||||||
conversation_id,
|
|
||||||
max_iterations,
|
|
||||||
agent_output,
|
|
||||||
model_id,
|
|
||||||
ctx,
|
|
||||||
);
|
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -2356,9 +2345,7 @@ impl BlocklistAIController {
|
|||||||
AIAgentTextSection::PlainText { text } => {
|
AIAgentTextSection::PlainText { text } => {
|
||||||
Some(text.text().to_string())
|
Some(text.text().to_string())
|
||||||
}
|
}
|
||||||
AIAgentTextSection::Code { code, .. } => {
|
AIAgentTextSection::Code { code, .. } => Some(code.clone()),
|
||||||
Some(code.clone())
|
|
||||||
}
|
|
||||||
_ => None,
|
_ => None,
|
||||||
})
|
})
|
||||||
.collect::<Vec<_>>()
|
.collect::<Vec<_>>()
|
||||||
@@ -2786,20 +2773,15 @@ impl BlocklistAIController {
|
|||||||
&conversation_data.server_conversation_token,
|
&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
|
// nearing or past expiry, kick off a background refresh so upcoming
|
||||||
// requests can authenticate even when the proactive refresh loop
|
// requests can authenticate even when the proactive refresh loop isn't
|
||||||
// isn't running. This request still carries the currently stored
|
// running (e.g. a parked or never-armed refresh chain).
|
||||||
// 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.
|
|
||||||
#[cfg(not(target_family = "wasm"))]
|
#[cfg(not(target_family = "wasm"))]
|
||||||
{
|
{
|
||||||
use ::ai::api_keys::ApiKeyManager;
|
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| {
|
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);
|
crate::ai::geap_credentials::refresh_geap_credentials_if_needed(manager, ctx);
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
@@ -3391,10 +3373,8 @@ impl BlocklistAIController {
|
|||||||
|| error_str.contains("ThrottlingException"));
|
|| error_str.contains("ThrottlingException"));
|
||||||
|
|
||||||
const MAX_ERROR_RETRIES: usize = 2;
|
const MAX_ERROR_RETRIES: usize = 2;
|
||||||
let retry_count = self
|
let retry_count =
|
||||||
.error_retry_counts
|
self.error_retry_counts.entry(conversation_id).or_insert(0);
|
||||||
.entry(conversation_id)
|
|
||||||
.or_insert(0);
|
|
||||||
let should_corrective_retry =
|
let should_corrective_retry =
|
||||||
is_corrective_retry_candidate && *retry_count < MAX_ERROR_RETRIES;
|
is_corrective_retry_candidate && *retry_count < MAX_ERROR_RETRIES;
|
||||||
|
|
||||||
@@ -3445,17 +3425,15 @@ impl BlocklistAIController {
|
|||||||
error_str
|
error_str
|
||||||
);
|
);
|
||||||
|
|
||||||
let inputs = vec![
|
let inputs = vec![AIAgentInput::UserQuery {
|
||||||
AIAgentInput::UserQuery {
|
query: corrective_msg,
|
||||||
query: corrective_msg,
|
context: Arc::from([]),
|
||||||
context: Arc::from([]),
|
static_query_type: None,
|
||||||
static_query_type: None,
|
referenced_attachments: HashMap::new(),
|
||||||
referenced_attachments: HashMap::new(),
|
user_query_mode: UserQueryMode::Normal,
|
||||||
user_query_mode: UserQueryMode::Normal,
|
running_command: None,
|
||||||
running_command: None,
|
intended_agent: None,
|
||||||
intended_agent: None,
|
}];
|
||||||
},
|
|
||||||
];
|
|
||||||
|
|
||||||
let _ = self.send_request_input(
|
let _ = self.send_request_input(
|
||||||
RequestInput::for_task(
|
RequestInput::for_task(
|
||||||
|
|||||||
@@ -324,16 +324,6 @@ impl PromptAlertView {
|
|||||||
};
|
};
|
||||||
text_fragments.push(FormattedTextFragment::hyperlink(label, upgrade_url));
|
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 => {}
|
PromptAlertState::NoAlert => {}
|
||||||
}
|
}
|
||||||
|
|||||||
+16
-77
@@ -3,7 +3,7 @@
|
|||||||
use std::collections::{HashMap, HashSet};
|
use std::collections::{HashMap, HashSet};
|
||||||
use std::sync::{Arc, OnceLock};
|
use std::sync::{Arc, OnceLock};
|
||||||
|
|
||||||
use ai::api_keys::{ApiKeyManager, ApiKeyManagerEvent, CustomEndpoint, CustomEndpointModel};
|
use ai::api_keys::ApiKeyManager;
|
||||||
pub use ai::LLMId;
|
pub use ai::LLMId;
|
||||||
use galaxy_core::features::FeatureFlag;
|
use galaxy_core::features::FeatureFlag;
|
||||||
use galaxy_core::ui::icons::Icon;
|
use galaxy_core::ui::icons::Icon;
|
||||||
@@ -27,23 +27,10 @@ use crate::workspaces::user_workspaces::{UserWorkspaces, UserWorkspacesEvent};
|
|||||||
use crate::{report_error, AISettings};
|
use crate::{report_error, AISettings};
|
||||||
|
|
||||||
/// Checks if a user's' API key is being used for the given provider.
|
/// 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.
|
/// AWS Bedrock is the only supported provider in Galaxy; there is no
|
||||||
/// For xAI, a connected Grok subscription counts: its OAuth access token is
|
/// user-pasted BYO key path for other providers.
|
||||||
/// sent like a BYO key (see `ApiKeyManager::api_keys_for_request`).
|
pub fn is_using_api_key_for_provider(_provider: &LLMProvider, _app: &AppContext) -> bool {
|
||||||
pub fn is_using_api_key_for_provider(provider: &LLMProvider, app: &AppContext) -> bool {
|
false
|
||||||
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,
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn should_show_bedrock_icon_for_model(llm: &LLMInfo, app: &AppContext) -> bool {
|
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 {
|
fn default_llm_info(&self) -> &LLMInfo {
|
||||||
static NO_PROVIDER_FALLBACK: std::sync::LazyLock<LLMInfo> = std::sync::LazyLock::new(|| {
|
static NO_PROVIDER_FALLBACK: std::sync::LazyLock<LLMInfo> =
|
||||||
LLMInfo {
|
std::sync::LazyLock::new(|| LLMInfo {
|
||||||
display_name: "No models configured".to_owned(),
|
display_name: "No models configured".to_owned(),
|
||||||
base_model_name: "No models configured".to_owned(),
|
base_model_name: "No models configured".to_owned(),
|
||||||
id: "none".to_owned().into(),
|
id: "none".to_owned().into(),
|
||||||
@@ -433,9 +420,7 @@ impl AvailableLLMs {
|
|||||||
request_multiplier: 1,
|
request_multiplier: 1,
|
||||||
credit_multiplier: None,
|
credit_multiplier: None,
|
||||||
},
|
},
|
||||||
description: Some(
|
description: Some("Enable Bedrock or OpenAI/LiteLLM in settings".to_string()),
|
||||||
"Enable Bedrock or OpenAI/LiteLLM in settings".to_string(),
|
|
||||||
),
|
|
||||||
disable_reason: Some(DisableReason::Unavailable),
|
disable_reason: Some(DisableReason::Unavailable),
|
||||||
vision_supported: false,
|
vision_supported: false,
|
||||||
spec: None,
|
spec: None,
|
||||||
@@ -443,8 +428,7 @@ impl AvailableLLMs {
|
|||||||
host_configs: HashMap::new(),
|
host_configs: HashMap::new(),
|
||||||
discount_percentage: None,
|
discount_percentage: None,
|
||||||
context_window: LLMContextWindow::default(),
|
context_window: LLMContextWindow::default(),
|
||||||
}
|
});
|
||||||
});
|
|
||||||
|
|
||||||
self.info_for_id(&self.default_id)
|
self.info_for_id(&self.default_id)
|
||||||
.or_else(|| self.choices.first())
|
.or_else(|| self.choices.first())
|
||||||
@@ -675,7 +659,7 @@ impl LLMPreferences {
|
|||||||
});
|
});
|
||||||
|
|
||||||
let base_llm_for_terminal_view = HashMap::new();
|
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 {
|
let mut me = Self {
|
||||||
models_by_feature,
|
models_by_feature,
|
||||||
@@ -758,7 +742,9 @@ impl LLMPreferences {
|
|||||||
.choices
|
.choices
|
||||||
.retain(|m| m.provider != LLMProvider::Bedrock && m.provider != LLMProvider::Unknown);
|
.retain(|m| m.provider != LLMProvider::Bedrock && m.provider != LLMProvider::Unknown);
|
||||||
if let Some(ref mut cli) = self.models_by_feature.cli_agent {
|
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);
|
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
|
/// Fetches model metadata from LiteLLM's `/model/info` endpoint which returns rich
|
||||||
/// metadata including accurate context window sizes, output token limits, and
|
/// metadata including accurate context window sizes, output token limits, and
|
||||||
/// capability flags (vision, function calling).
|
/// capability flags (vision, function calling).
|
||||||
@@ -2144,10 +2084,9 @@ async fn fetch_from_openai_models(
|
|||||||
m,
|
m,
|
||||||
&["max_input_tokens", "input_token_limit", "max_prompt_tokens"],
|
&["max_input_tokens", "input_token_limit", "max_prompt_tokens"],
|
||||||
);
|
);
|
||||||
let context_size =
|
let context_size = u32_from_any(m, &["max_model_len", "context_window", "token_size"])
|
||||||
u32_from_any(m, &["max_model_len", "context_window", "token_size"])
|
.or(max_input_tokens)
|
||||||
.or(max_input_tokens)
|
.unwrap_or(200_000);
|
||||||
.unwrap_or(200_000);
|
|
||||||
let max_output_tokens = u32_from_any(
|
let max_output_tokens = u32_from_any(
|
||||||
m,
|
m,
|
||||||
&[
|
&[
|
||||||
|
|||||||
@@ -138,458 +138,3 @@ fn llm_info_round_trip_serializes_and_deserializes() {
|
|||||||
|
|
||||||
assert_eq!(info, round_tripped);
|
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)
|
|
||||||
);
|
|
||||||
});
|
|
||||||
});
|
|
||||||
}
|
|
||||||
|
|||||||
@@ -1,6 +1,5 @@
|
|||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
|
|
||||||
use ai::api_keys::ApiKeyManager;
|
|
||||||
use chrono::{DateTime, Local, Utc};
|
use chrono::{DateTime, Local, Utc};
|
||||||
use galaxy_core::user_preferences::GetUserPreferences as _;
|
use galaxy_core::user_preferences::GetUserPreferences as _;
|
||||||
pub use galaxy_graphql::billing::BonusGrantType;
|
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)
|
/// 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)
|
/// 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
|
/// 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.
|
/// Use this method as the starting point for AI availability checking.
|
||||||
pub fn has_any_ai_remaining(&self, ctx: &AppContext) -> bool {
|
pub fn has_any_ai_remaining(&self, ctx: &AppContext) -> bool {
|
||||||
let current_workspace = UserWorkspaces::as_ref(ctx).current_workspace();
|
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))
|
.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
|
has_base_plan_ai_requests
|
||||||
|| (user_bonus_credits || workspace_bonus_credits)
|
|| (user_bonus_credits || workspace_bonus_credits)
|
||||||
|| workspace_has_overages
|
|| workspace_has_overages
|
||||||
|| is_payg_enabled
|
|| is_payg_enabled
|
||||||
|| is_enterprise_auto_reload_enabled
|
|| is_enterprise_auto_reload_enabled
|
||||||
|| is_self_serve_auto_reload_enabled
|
|| is_self_serve_auto_reload_enabled
|
||||||
|| has_byo_credentials
|
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn requests_used(&self) -> usize {
|
pub fn requests_used(&self) -> usize {
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
|
|
||||||
use ai::api_keys::{ApiKeyManager, GrokTokens};
|
use ai::api_keys::ApiKeyManager;
|
||||||
use chrono::Duration;
|
use chrono::Duration;
|
||||||
use galaxy_core::features::FeatureFlag;
|
use galaxy_core::features::FeatureFlag;
|
||||||
use galaxy_graphql::billing::{AddonCreditsOption, OveragesPricing, PricingInfo};
|
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::server::server_api::ServerApiProvider;
|
||||||
use crate::workspaces::user_workspaces::UserWorkspaces;
|
use crate::workspaces::user_workspaces::UserWorkspaces;
|
||||||
use crate::workspaces::workspace::{
|
use crate::workspaces::workspace::{
|
||||||
AiOverages, ByoApiKeyPolicy, CustomerType, EnterpriseCreditsAutoReloadPolicy,
|
AiOverages, CustomerType, EnterpriseCreditsAutoReloadPolicy, EnterprisePayAsYouGoPolicy,
|
||||||
EnterprisePayAsYouGoPolicy, PurchaseAddOnCreditsPolicy, Workspace, WorkspaceUid,
|
PurchaseAddOnCreditsPolicy, Workspace, WorkspaceUid,
|
||||||
};
|
};
|
||||||
|
|
||||||
fn create_test_workspace() -> (WorkspaceUid, Workspace) {
|
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]
|
#[test]
|
||||||
fn test_has_any_ai_remaining_false_with_only_ambient_bonus_credits() {
|
fn test_has_any_ai_remaining_false_with_only_ambient_bonus_credits() {
|
||||||
App::test((), |mut app| async move {
|
App::test((), |mut app| async move {
|
||||||
|
|||||||
@@ -1498,23 +1498,6 @@ pub(crate) fn initialize_app(
|
|||||||
if FeatureFlag::GeminiEnterprise.is_enabled() {
|
if FeatureFlag::GeminiEnterprise.is_enabled() {
|
||||||
manager.subscribe_to_geap_settings_changes(ctx);
|
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
|
manager
|
||||||
});
|
});
|
||||||
|
|
||||||
@@ -2393,7 +2376,6 @@ pub(crate) fn app_callbacks(
|
|||||||
if let Some(initialization) = tracing_initialization.as_mut() {
|
if let Some(initialization) = tracing_initialization.as_mut() {
|
||||||
initialization.shutdown();
|
initialization.shutdown();
|
||||||
}
|
}
|
||||||
|
|
||||||
})),
|
})),
|
||||||
on_should_close_window: Some(Box::new(move |window_id, ctx| {
|
on_should_close_window: Some(Box::new(move |window_id, ctx| {
|
||||||
let general_settings = GeneralSettings::as_ref(ctx);
|
let general_settings = GeneralSettings::as_ref(ctx);
|
||||||
|
|||||||
+1
-14
@@ -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 {
|
should_render_use_agent_footer_for_user_commands: ShouldRenderUseAgentToolbarForUserCommands {
|
||||||
type: bool,
|
type: bool,
|
||||||
default: true,
|
default: true,
|
||||||
@@ -1947,8 +1935,7 @@ impl AISettings {
|
|||||||
// Galaxy does not require Warp authentication for AI.
|
// Galaxy does not require Warp authentication for AI.
|
||||||
// AI is enabled as long as the user hasn't explicitly disabled it
|
// AI is enabled as long as the user hasn't explicitly disabled it
|
||||||
// and there's no org policy blocking it.
|
// and there's no org policy blocking it.
|
||||||
*self.is_any_ai_enabled
|
*self.is_any_ai_enabled && !self.is_ai_disabled_due_to_remote_session_org_policy(app)
|
||||||
&& !self.is_ai_disabled_due_to_remote_session_org_policy(app)
|
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn default_session_mode(&self, app: &AppContext) -> DefaultSessionMode {
|
pub fn default_session_mode(&self, app: &AppContext) -> DefaultSessionMode {
|
||||||
|
|||||||
+25
-1197
File diff suppressed because it is too large
Load Diff
@@ -69,7 +69,7 @@ use crate::workspace::Workspace;
|
|||||||
use crate::workspaces::user_workspaces::UserWorkspaces;
|
use crate::workspaces::user_workspaces::UserWorkspaces;
|
||||||
use crate::ToastStack;
|
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)]
|
#[derive(Debug, Clone)]
|
||||||
pub enum MCPServersListPageViewEvent {
|
pub enum MCPServersListPageViewEvent {
|
||||||
|
|||||||
@@ -269,7 +269,7 @@ impl Display for SettingsSection {
|
|||||||
SettingsSection::MCPServers => write!(f, "MCP Servers"),
|
SettingsSection::MCPServers => write!(f, "MCP Servers"),
|
||||||
SettingsSection::Scripting => write!(f, "Scripting"),
|
SettingsSection::Scripting => write!(f, "Scripting"),
|
||||||
SettingsSection::WarpDrive => write!(f, "Galaxy Drive"),
|
SettingsSection::WarpDrive => write!(f, "Galaxy Drive"),
|
||||||
SettingsSection::WarpAgent => write!(f, "Warp Agent"),
|
SettingsSection::WarpAgent => write!(f, "Galaxy Agent"),
|
||||||
SettingsSection::AgentProfiles => write!(f, "Profiles"),
|
SettingsSection::AgentProfiles => write!(f, "Profiles"),
|
||||||
SettingsSection::AgentMCPServers => write!(f, "MCP servers"),
|
SettingsSection::AgentMCPServers => write!(f, "MCP servers"),
|
||||||
SettingsSection::Knowledge => write!(f, "Knowledge"),
|
SettingsSection::Knowledge => write!(f, "Knowledge"),
|
||||||
@@ -544,7 +544,6 @@ pub mod flags {
|
|||||||
pub const SUGGESTED_RULES_FLAG: &str = "Suggested_Rules";
|
pub const SUGGESTED_RULES_FLAG: &str = "Suggested_Rules";
|
||||||
pub const WARP_DRIVE_CONTEXT_FLAG: &str = "Warp_Drive_Context";
|
pub const WARP_DRIVE_CONTEXT_FLAG: &str = "Warp_Drive_Context";
|
||||||
pub const FILE_BASED_MCP_FLAG: &str = "File_Based_MCP";
|
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 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 DEBUG_SHOW_MEMORY_STATS_FLAG: &str = "Debug_Memory_Statistics";
|
||||||
pub const ALLOW_NATIVE_WAYLAND: &str = "Allow_Native_Wayland";
|
pub const ALLOW_NATIVE_WAYLAND: &str = "Allow_Native_Wayland";
|
||||||
|
|||||||
@@ -171,7 +171,7 @@ fn match_data_countable_zero_is_not_truthy() {
|
|||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn subpage_display_names_are_correct() {
|
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::AgentProfiles.to_string(), "Profiles");
|
||||||
assert_eq!(SettingsSection::AgentMCPServers.to_string(), "MCP servers");
|
assert_eq!(SettingsSection::AgentMCPServers.to_string(), "MCP servers");
|
||||||
assert_eq!(SettingsSection::Knowledge.to_string(), "Knowledge");
|
assert_eq!(SettingsSection::Knowledge.to_string(), "Knowledge");
|
||||||
|
|||||||
@@ -630,14 +630,6 @@ impl SearchItem for ModelSearchItem {
|
|||||||
first.make_ascii_uppercase();
|
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![
|
let mut text_fragments = vec![
|
||||||
FormattedTextFragment::plain_text(format!(
|
FormattedTextFragment::plain_text(format!(
|
||||||
"{display_name} is not available for free users. "
|
"{display_name} is not available for free users. "
|
||||||
@@ -645,17 +637,6 @@ impl SearchItem for ModelSearchItem {
|
|||||||
FormattedTextFragment::hyperlink("Upgrade", upgrade_url),
|
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(
|
let upgrade_text = FormattedTextElement::new(
|
||||||
FormattedText::new([FormattedTextLine::Line(text_fragments)]),
|
FormattedText::new([FormattedTextLine::Line(text_fragments)]),
|
||||||
inline_styles::font_size(appearance),
|
inline_styles::font_size(appearance),
|
||||||
|
|||||||
@@ -1,6 +1,5 @@
|
|||||||
use std::future::Future;
|
use std::future::Future;
|
||||||
|
|
||||||
use ai::api_keys::ApiKeyManager;
|
|
||||||
use galaxy_core::features::FeatureFlag;
|
use galaxy_core::features::FeatureFlag;
|
||||||
use galaxy_core::send_telemetry_from_ctx;
|
use galaxy_core::send_telemetry_from_ctx;
|
||||||
use settings::Setting as _;
|
use settings::Setting as _;
|
||||||
@@ -442,7 +441,10 @@ impl OneTimeModalModel {
|
|||||||
.current_workspace()
|
.current_workspace()
|
||||||
.map(|workspace| workspace.billing_metadata.customer_type);
|
.map(|workspace| workspace.billing_metadata.customer_type);
|
||||||
let is_warp_ai_enabled = *AISettings::as_ref(ctx).is_any_ai_enabled;
|
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 completed_new_onboarding = has_completed_local_onboarding(ctx);
|
||||||
let has_zero_base_credits = AIRequestUsageModel::as_ref(ctx).request_limit() == 0;
|
let has_zero_base_credits = AIRequestUsageModel::as_ref(ctx).request_limit() == 0;
|
||||||
|
|
||||||
|
|||||||
@@ -287,9 +287,9 @@ use crate::pane_group::pane::ActionOrigin;
|
|||||||
use crate::pane_group::FilePane;
|
use crate::pane_group::FilePane;
|
||||||
use crate::pane_group::{
|
use crate::pane_group::{
|
||||||
self, AIFactPane, AnyPaneContent, ChildAgentOrigin, CodeDiffPane, CodePane, CodeReviewPanelArg,
|
self, AIFactPane, AnyPaneContent, ChildAgentOrigin, CodeDiffPane, CodePane, CodeReviewPanelArg,
|
||||||
CustomRouterEditorPane, Direction as PaneGroupDirection, Direction,
|
CustomRouterEditorPane, Direction as PaneGroupDirection, Direction, ExecutionProfileEditorPane,
|
||||||
ExecutionProfileEditorPane, NetworkLogPane, NewTerminalOptions, PaneGroup, PaneId, PanesLayout,
|
NetworkLogPane, NewTerminalOptions, PaneGroup, PaneId, PanesLayout, TabBarHoverIndex,
|
||||||
TabBarHoverIndex, TerminalPaneId,
|
TerminalPaneId,
|
||||||
};
|
};
|
||||||
use crate::persistence::ModelEvent;
|
use crate::persistence::ModelEvent;
|
||||||
use crate::projects::ProjectManagementModel;
|
use crate::projects::ProjectManagementModel;
|
||||||
@@ -22858,9 +22858,6 @@ impl Workspace {
|
|||||||
if *ai_settings.file_based_mcp_enabled.value() {
|
if *ai_settings.file_based_mcp_enabled.value() {
|
||||||
context.set.insert(flags::FILE_BASED_MCP_FLAG);
|
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() {
|
if *session_settings.show_model_selectors_in_prompt.value() {
|
||||||
context
|
context
|
||||||
.set
|
.set
|
||||||
|
|||||||
+11
-495
@@ -1,9 +1,4 @@
|
|||||||
use std::time::{Duration, SystemTime};
|
|
||||||
|
|
||||||
use galaxyui_core::{Entity, ModelContext, SingletonEntity};
|
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;
|
use warp_multi_agent_api as api;
|
||||||
|
|
||||||
pub use crate::aws_credentials::{AwsCredentials, AwsCredentialsState};
|
pub use crate::aws_credentials::{AwsCredentials, AwsCredentialsState};
|
||||||
@@ -12,137 +7,13 @@ pub use crate::geap_credentials::{
|
|||||||
LoadGeapCredentialsError, GEAP_REFRESH_LEAD_TIME,
|
LoadGeapCredentialsError, GEAP_REFRESH_LEAD_TIME,
|
||||||
};
|
};
|
||||||
|
|
||||||
const SECURE_STORAGE_KEY: &str = "AiApiKeys";
|
/// Emitted when the manager's stored credentials (AWS Bedrock or Gemini
|
||||||
|
/// Enterprise) are updated in-memory.
|
||||||
/// 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.
|
|
||||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||||
pub enum ApiKeyManagerEvent {
|
pub enum ApiKeyManagerEvent {
|
||||||
KeysUpdated,
|
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<String>,
|
|
||||||
pub anthropic: Option<String>,
|
|
||||||
pub openai: Option<String>,
|
|
||||||
pub open_router: Option<String>,
|
|
||||||
pub custom_endpoints: Vec<CustomEndpoint>,
|
|
||||||
}
|
|
||||||
|
|
||||||
#[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<CustomEndpointModel>,
|
|
||||||
}
|
|
||||||
|
|
||||||
#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)]
|
|
||||||
#[serde(default)]
|
|
||||||
pub struct CustomEndpointModel {
|
|
||||||
pub name: String,
|
|
||||||
pub alias: Option<String>,
|
|
||||||
/// 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<String>,
|
|
||||||
/// Absolute time at which `access_token` expires, if the provider told us.
|
|
||||||
#[serde(default)]
|
|
||||||
pub expires_at: Option<SystemTime>,
|
|
||||||
/// 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<SystemTime>,
|
|
||||||
}
|
|
||||||
|
|
||||||
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`].
|
/// Controls how AWS credentials are refreshed by [`ApiKeyManager`].
|
||||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||||
pub enum AwsCredentialsRefreshStrategy {
|
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 {
|
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<GrokTokens>,
|
|
||||||
/// 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,
|
pub(crate) aws_credentials_state: AwsCredentialsState,
|
||||||
aws_credentials_refresh_strategy: AwsCredentialsRefreshStrategy,
|
aws_credentials_refresh_strategy: AwsCredentialsRefreshStrategy,
|
||||||
/// In-memory Gemini Enterprise (GEAP) credential state.
|
/// In-memory Gemini Enterprise (GEAP) credential state.
|
||||||
pub(crate) geap_credentials_state: GeapCredentialsState,
|
pub(crate) geap_credentials_state: GeapCredentialsState,
|
||||||
secure_storage_write_version: u64,
|
|
||||||
grok_secure_storage_write_version: u64,
|
|
||||||
}
|
}
|
||||||
|
|
||||||
impl ApiKeyManager {
|
impl ApiKeyManager {
|
||||||
pub fn new(ctx: &mut ModelContext<Self>) -> Self {
|
pub fn new(_ctx: &mut ModelContext<Self>) -> Self {
|
||||||
let keys = Self::load_keys_from_secure_storage(ctx);
|
|
||||||
let grok_tokens = Self::load_grok_tokens_from_secure_storage(ctx);
|
|
||||||
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_state: AwsCredentialsState::Missing,
|
||||||
aws_credentials_refresh_strategy: AwsCredentialsRefreshStrategy::default(),
|
aws_credentials_refresh_strategy: AwsCredentialsRefreshStrategy::default(),
|
||||||
geap_credentials_state: GeapCredentialsState::Missing,
|
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<GrokTokens>, ctx: &mut ModelContext<Self>) {
|
|
||||||
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<String>, ctx: &mut ModelContext<Self>) {
|
|
||||||
self.keys.google = key;
|
|
||||||
ctx.emit(ApiKeyManagerEvent::KeysUpdated);
|
|
||||||
self.write_keys_to_secure_storage(ctx);
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn set_anthropic_key(&mut self, key: Option<String>, ctx: &mut ModelContext<Self>) {
|
|
||||||
self.keys.anthropic = key;
|
|
||||||
ctx.emit(ApiKeyManagerEvent::KeysUpdated);
|
|
||||||
self.write_keys_to_secure_storage(ctx);
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn set_openai_key(&mut self, key: Option<String>, ctx: &mut ModelContext<Self>) {
|
|
||||||
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<String>, ctx: &mut ModelContext<Self>) {
|
|
||||||
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<String>, Option<String>)>,
|
|
||||||
ctx: &mut ModelContext<Self>,
|
|
||||||
) {
|
|
||||||
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<String>, Option<String>)>,
|
|
||||||
ctx: &mut ModelContext<Self>,
|
|
||||||
) {
|
|
||||||
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<Self>) {
|
|
||||||
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<Self>) {
|
|
||||||
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(
|
pub fn set_aws_credentials_state(
|
||||||
&mut self,
|
&mut self,
|
||||||
state: AwsCredentialsState,
|
state: AwsCredentialsState,
|
||||||
@@ -380,93 +89,14 @@ impl ApiKeyManager {
|
|||||||
self.aws_credentials_refresh_strategy = strategy;
|
self.aws_credentials_refresh_strategy = strategy;
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Builds the `CustomModelProviders` registry that ships with every agent request.
|
/// Builds the `ApiKeys` request payload carrying AWS Bedrock and/or Gemini
|
||||||
///
|
/// Enterprise (GEAP) credentials, when applicable. Returns `None` when
|
||||||
/// Emits one [`CustomModelProvider`] per configured [`CustomEndpoint`], each populated with
|
/// neither credential type applies to this request.
|
||||||
/// 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<api::request::settings::CustomModelProviders> {
|
|
||||||
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 })
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn api_keys_for_request(
|
pub fn api_keys_for_request(
|
||||||
&self,
|
&self,
|
||||||
include_byo_keys: bool,
|
|
||||||
include_aws_bedrock_credentials: bool,
|
include_aws_bedrock_credentials: bool,
|
||||||
geap_binding: Option<GeapMintBinding>,
|
geap_binding: Option<GeapMintBinding>,
|
||||||
) -> Option<api::request::settings::ApiKeys> {
|
) -> Option<api::request::settings::ApiKeys> {
|
||||||
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
|
// 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).
|
// of the per-user setting flag (which only applies to the local credential chain path).
|
||||||
let include_aws = include_aws_bedrock_credentials
|
let include_aws = include_aws_bedrock_credentials
|
||||||
@@ -506,130 +136,16 @@ impl ApiKeyManager {
|
|||||||
_ => None,
|
_ => None,
|
||||||
});
|
});
|
||||||
|
|
||||||
if anthropic.is_empty()
|
if aws_credentials.is_none() && google_cloud_credentials.is_none() {
|
||||||
&& 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()
|
|
||||||
{
|
|
||||||
None
|
None
|
||||||
} else {
|
} else {
|
||||||
Some(api::request::settings::ApiKeys {
|
Some(api::request::settings::ApiKeys {
|
||||||
anthropic,
|
|
||||||
openai,
|
|
||||||
google,
|
|
||||||
open_router,
|
|
||||||
grok_oauth_access_token,
|
|
||||||
allow_use_of_warp_credits: false,
|
|
||||||
aws_credentials,
|
aws_credentials,
|
||||||
google_cloud_credentials,
|
google_cloud_credentials,
|
||||||
|
..Default::default()
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
fn load_keys_from_secure_storage(ctx: &mut ModelContext<Self>) -> 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<Self>) {
|
|
||||||
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<Self>) -> Option<GrokTokens> {
|
|
||||||
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<Self>) {
|
|
||||||
// `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 {
|
impl Entity for ApiKeyManager {
|
||||||
|
|||||||
+58
-528
@@ -2,41 +2,20 @@ use std::time::{Duration, SystemTime};
|
|||||||
|
|
||||||
use super::*;
|
use super::*;
|
||||||
|
|
||||||
fn make_manager(keys: ApiKeys) -> ApiKeyManager {
|
fn make_manager() -> ApiKeyManager {
|
||||||
make_manager_with_grok(keys, None)
|
|
||||||
}
|
|
||||||
|
|
||||||
fn make_manager_with_grok(keys: ApiKeys, grok_tokens: Option<GrokTokens>) -> ApiKeyManager {
|
|
||||||
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_state: AwsCredentialsState::Missing,
|
||||||
aws_credentials_refresh_strategy: AwsCredentialsRefreshStrategy::default(),
|
aws_credentials_refresh_strategy: AwsCredentialsRefreshStrategy::default(),
|
||||||
geap_credentials_state: GeapCredentialsState::Missing,
|
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 {
|
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.geap_credentials_state = geap_credentials_state;
|
||||||
manager
|
manager
|
||||||
}
|
}
|
||||||
|
|
||||||
fn grok_tokens(access_token: &str, expires_in: Option<u64>) -> 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<u64>) -> GeapCredentials {
|
fn geap_credentials(access_token: &str, expires_in: Option<u64>) -> GeapCredentials {
|
||||||
GeapCredentials::new(
|
GeapCredentials::new(
|
||||||
access_token.into(),
|
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 {
|
fn geap_gate() -> GeapMintBinding {
|
||||||
geap_binding()
|
geap_binding()
|
||||||
}
|
}
|
||||||
@@ -70,487 +47,6 @@ fn geap_loaded(access_token: &str, expires_in: Option<u64>) -> 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::<Vec<_>>()
|
|
||||||
.iter()
|
|
||||||
.map(|(n, a, k)| (*n, *a, k.as_str()))
|
|
||||||
.collect::<Vec<_>>(),
|
|
||||||
)
|
|
||||||
}
|
|
||||||
|
|
||||||
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 ────────────────────────────────────────────
|
// ── geap credentials ────────────────────────────────────────────
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
@@ -587,13 +83,10 @@ fn geap_needs_refresh_lead_time_boundaries() {
|
|||||||
#[test]
|
#[test]
|
||||||
fn api_keys_for_request_includes_geap_token_when_gate_and_binding_match() {
|
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 mgr = make_manager_with_geap(geap_loaded("geap-abc", Some(3600)));
|
||||||
let result = mgr
|
let result = mgr.api_keys_for_request(false, Some(geap_gate())).unwrap();
|
||||||
.api_keys_for_request(false, false, Some(geap_gate()))
|
|
||||||
.unwrap();
|
|
||||||
let credentials = result.google_cloud_credentials.unwrap();
|
let credentials = result.google_cloud_credentials.unwrap();
|
||||||
assert_eq!(credentials.access_token, "geap-abc");
|
assert_eq!(credentials.access_token, "geap-abc");
|
||||||
// The GEAP token is independent of the BYO key gate.
|
assert!(result.aws_credentials.is_none());
|
||||||
assert!(result.anthropic.is_empty());
|
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
@@ -602,9 +95,7 @@ fn api_keys_for_request_includes_expired_geap_token() {
|
|||||||
// rejects truly invalid ones, which surfaces a recoverable error instead
|
// rejects truly invalid ones, which surfaces a recoverable error instead
|
||||||
// of a silent fallback to another route.
|
// of a silent fallback to another route.
|
||||||
let mgr = make_manager_with_geap(geap_loaded("geap-abc", Some(0)));
|
let mgr = make_manager_with_geap(geap_loaded("geap-abc", Some(0)));
|
||||||
let result = mgr
|
let result = mgr.api_keys_for_request(false, Some(geap_gate())).unwrap();
|
||||||
.api_keys_for_request(false, false, Some(geap_gate()))
|
|
||||||
.unwrap();
|
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
result.google_cloud_credentials.unwrap().access_token,
|
result.google_cloud_credentials.unwrap().access_token,
|
||||||
"geap-abc"
|
"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
|
// No gate (policy off at the call site) ⇒ no GEAP credentials, even when
|
||||||
// a token is loaded.
|
// a token is loaded.
|
||||||
let mgr = make_manager_with_geap(geap_loaded("geap-abc", Some(3600)));
|
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]
|
#[test]
|
||||||
@@ -626,19 +117,19 @@ fn api_keys_for_request_omits_geap_token_on_binding_mismatch() {
|
|||||||
// A different user (sign-out/account switch).
|
// A different user (sign-out/account switch).
|
||||||
let mut gate = geap_gate();
|
let mut gate = geap_gate();
|
||||||
gate.user_uid = "someone-else".into();
|
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).
|
// A different audience (admin changed the pool/provider).
|
||||||
let mut gate = geap_gate();
|
let mut gate = geap_gate();
|
||||||
gate.audience = "//iam.googleapis.com/projects/2/locations/global/workloadIdentityPools/other/providers/other".into();
|
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).
|
// A different service account (admin changed impersonation target).
|
||||||
let mut gate = geap_gate();
|
let mut gate = geap_gate();
|
||||||
gate.federation = GeapFederation::ServiceAccount {
|
gate.federation = GeapFederation::ServiceAccount {
|
||||||
email: "other@proj.iam.gserviceaccount.com".into(),
|
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]
|
#[test]
|
||||||
@@ -648,9 +139,7 @@ fn api_keys_for_request_serves_previous_geap_token_while_refreshing() {
|
|||||||
let mgr = make_manager_with_geap(GeapCredentialsState::Refreshing {
|
let mgr = make_manager_with_geap(GeapCredentialsState::Refreshing {
|
||||||
previous: Some((geap_credentials("geap-old", Some(10)), geap_binding())),
|
previous: Some((geap_credentials("geap-old", Some(10)), geap_binding())),
|
||||||
});
|
});
|
||||||
let result = mgr
|
let result = mgr.api_keys_for_request(false, Some(geap_gate())).unwrap();
|
||||||
.api_keys_for_request(false, false, Some(geap_gate()))
|
|
||||||
.unwrap();
|
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
result.google_cloud_credentials.unwrap().access_token,
|
result.google_cloud_credentials.unwrap().access_token,
|
||||||
"geap-old"
|
"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() {
|
fn api_keys_for_request_omits_geap_token_during_first_mint() {
|
||||||
// The very first mint has nothing to serve yet.
|
// The very first mint has nothing to serve yet.
|
||||||
let mgr = make_manager_with_geap(GeapCredentialsState::Refreshing { previous: None });
|
let mgr = make_manager_with_geap(GeapCredentialsState::Refreshing { previous: None });
|
||||||
assert!(mgr
|
assert!(mgr.api_keys_for_request(false, Some(geap_gate())).is_none());
|
||||||
.api_keys_for_request(false, false, Some(geap_gate()))
|
|
||||||
.is_none());
|
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
@@ -679,9 +166,7 @@ fn api_keys_for_request_omits_geap_token_for_non_loaded_states() {
|
|||||||
},
|
},
|
||||||
] {
|
] {
|
||||||
let mgr = make_manager_with_geap(state);
|
let mgr = make_manager_with_geap(state);
|
||||||
assert!(mgr
|
assert!(mgr.api_keys_for_request(false, Some(geap_gate())).is_none());
|
||||||
.api_keys_for_request(false, 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();
|
let mut gate = geap_gate();
|
||||||
gate.user_uid = "someone-else".into();
|
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");
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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<Self>) {
|
|
||||||
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<Self>) {
|
|
||||||
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>,
|
|
||||||
) {
|
|
||||||
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<ApiKeyManager>,
|
|
||||||
) {
|
|
||||||
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<ApiKeyManager>) {
|
|
||||||
// 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<ApiKeyManager>,
|
|
||||||
) {
|
|
||||||
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:#}");
|
|
||||||
}
|
|
||||||
}
|
|
||||||
},
|
|
||||||
);
|
|
||||||
}
|
|
||||||
@@ -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<Self> {
|
|
||||||
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<TokenResponse> {
|
|
||||||
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<TokenResponse> {
|
|
||||||
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<String>,
|
|
||||||
#[serde(default)]
|
|
||||||
pub expires_in: Option<i64>,
|
|
||||||
}
|
|
||||||
|
|
||||||
/// 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<TcpListener> {
|
|
||||||
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<TokenResponse> {
|
|
||||||
// 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<CallbackData> {
|
|
||||||
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<Option<CallbackData>> {
|
|
||||||
// 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<String> {
|
|
||||||
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<TokenResponse> {
|
|
||||||
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<TokenResponse> {
|
|
||||||
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<T: serde::Serialize + ?Sized>(
|
|
||||||
form: &T,
|
|
||||||
) -> anyhow::Result<TokenResponse> {
|
|
||||||
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::<TokenResponse>()
|
|
||||||
.await
|
|
||||||
.context("failed to parse the Grok token response")
|
|
||||||
}
|
|
||||||
|
|
||||||
const SUCCESS_HTML: &str = "<!doctype html><html><head><meta charset=\"utf-8\">\
|
|
||||||
<title>Warp — Grok connected</title></head>\
|
|
||||||
<body style=\"font-family:system-ui,-apple-system,sans-serif;text-align:center;padding:3rem\">\
|
|
||||||
<h1>Grok connected</h1><p>You can close this window and return to Warp.</p></body></html>";
|
|
||||||
|
|
||||||
const FAILURE_HTML: &str = "<!doctype html><html><head><meta charset=\"utf-8\">\
|
|
||||||
<title>Warp — Grok authorization failed</title></head>\
|
|
||||||
<body style=\"font-family:system-ui,-apple-system,sans-serif;text-align:center;padding:3rem\">\
|
|
||||||
<h1>Authorization failed</h1><p>Something went wrong. Return to Warp and try again.</p></body></html>";
|
|
||||||
|
|
||||||
#[cfg(test)]
|
|
||||||
#[path = "oauth_tests.rs"]
|
|
||||||
mod tests;
|
|
||||||
@@ -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());
|
|
||||||
}
|
|
||||||
@@ -2,8 +2,6 @@ pub mod agent;
|
|||||||
pub mod api_keys;
|
pub mod api_keys;
|
||||||
pub mod aws_credentials;
|
pub mod aws_credentials;
|
||||||
pub mod geap_credentials;
|
pub mod geap_credentials;
|
||||||
#[cfg(not(target_family = "wasm"))]
|
|
||||||
pub mod grok_subscription;
|
|
||||||
pub mod llm_id;
|
pub mod llm_id;
|
||||||
|
|
||||||
pub use llm_id::LLMId;
|
pub use llm_id::LLMId;
|
||||||
|
|||||||
Reference in New Issue
Block a user