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

first pass of merging in warp (doesn't build)

Browse Source
This commit is contained in:
Ryan Ward
2026-07-01 16:08:58 -05:00
parent 2f64909469
commit 4770ac06b5
3662 changed files with 414574 additions and 89772 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/src/ai/request_usage_model.rs
+78 -23
View File
@@ -1,4 +1,14 @@
#![allow(dead_code, unused_imports, unused_variables)]
use std::sync::Arc;
use ai::api_keys::ApiKeyManager;
use chrono::{DateTime, Local, Utc};
use instant::Instant;
use serde::{Deserialize, Serialize};
use galaxy_core::user_preferences::GetUserPreferences as _;
pub use galaxy_graphql::billing::BonusGrantType;
use galaxy_graphql::scalars::time::ServerTimestamp;
use galaxyui::{AppContext, Entity, ModelContext, SingletonEntity};
use crate::ai::agent::conversation::AIConversationId;
use crate::ai::agent::AIAgentExchangeId;
use crate::auth::AuthStateProvider;
@@ -8,16 +18,6 @@ use crate::settings::AISettings;
use crate::workspaces::user_workspaces::UserWorkspaces;
use crate::workspaces::workspace::WorkspaceUid;
use crate::BlocklistAIHistoryModel;
use ai::api_keys::ApiKeyManager;
use chrono::{DateTime, Utc};
use galaxy_core::user_preferences::GetUserPreferences as _;
use galaxy_graphql::scalars::time::ServerTimestamp;
use galaxyui::{AppContext, Entity, ModelContext, SingletonEntity};
use instant::Instant;
use serde::{Deserialize, Serialize};
use std::sync::Arc;
pub use galaxy_graphql::billing::BonusGrantType;
/// Threshold of ambient-only credits at which we surface upgrade/CTA UI.
pub const AMBIENT_AGENT_TRIAL_CREDIT_THRESHOLD: i32 = 20;
@@ -51,6 +51,7 @@ pub struct BonusGrant {
/// The key for the corresponding entry in UserDefaults.
const REQUEST_LIMIT_INFO_CACHE_KEY: &str = "AIRequestLimitInfo";
const AMBIENT_CREDITS_BANNER_DISMISSED_KEY: &str = "AmbientCreditsBannerDismissed";
#[derive(Copy, Clone, Debug, Serialize, Deserialize)]
pub enum RequestLimitRefreshDuration {
@@ -165,6 +166,21 @@ fn get_cached_request_limit_info(app_mut: &mut AppContext) -> Option<RequestLimi
.and_then(|serialized| serde_json::from_str(serialized.as_str()).ok())
}
fn cache_ambient_credits_banner_dismissed(dismissed: bool, app_mut: &mut AppContext) {
let _ = app_mut
.private_user_preferences()
.write_value(AMBIENT_CREDITS_BANNER_DISMISSED_KEY, dismissed.to_string());
}
fn get_cached_ambient_credits_banner_dismissed(app_mut: &mut AppContext) -> bool {
app_mut
.private_user_preferences()
.read_value(AMBIENT_CREDITS_BANNER_DISMISSED_KEY)
.unwrap_or_default()
.and_then(|value| value.parse::<bool>().ok())
.unwrap_or_default()
}
pub struct AIRequestUsageModel {
ai_client: Arc<dyn AIClient>,
@@ -177,6 +193,9 @@ pub struct AIRequestUsageModel {
/// Whether the buy credits banner has been dismissed by the user.
buy_addon_credits_banner_dismissed: bool,
/// Whether the ambient trial credits banner has been dismissed by the user.
ambient_credits_banner_dismissed: bool,
}
impl Entity for AIRequestUsageModel {
@@ -185,6 +204,7 @@ impl Entity for AIRequestUsageModel {
pub enum AIRequestUsageModelEvent {
RequestUsageUpdated,
AmbientCreditsBannerDismissed,
RequestBonusRefunded {
requests_refunded: i32,
server_conversation_id: String,
@@ -198,6 +218,7 @@ impl AIRequestUsageModel {
// This is only used to show the latest known value before we finish refreshing from the server below.
let cached_request_limit_info = get_cached_request_limit_info(ctx);
let request_limit_info = cached_request_limit_info.unwrap_or_default();
let ambient_credits_banner_dismissed = get_cached_ambient_credits_banner_dismissed(ctx);
Self {
ai_client,
@@ -205,17 +226,19 @@ impl AIRequestUsageModel {
last_update_time: None,
bonus_grants: vec![],
buy_addon_credits_banner_dismissed: false,
ambient_credits_banner_dismissed,
}
}
#[cfg(test)]
pub fn new_for_test(ai_client: Arc<dyn AIClient>, _ctx: &mut ModelContext<Self>) -> Self {
pub fn new_for_test(ai_client: Arc<dyn AIClient>, ctx: &mut ModelContext<Self>) -> Self {
Self {
ai_client,
last_update_time: None,
request_limit_info: RequestLimitInfo::default(),
bonus_grants: vec![],
buy_addon_credits_banner_dismissed: false,
ambient_credits_banner_dismissed: get_cached_ambient_credits_banner_dismissed(ctx),
}
}
@@ -305,7 +328,7 @@ impl AIRequestUsageModel {
if exchange
.input
.iter()
.any(|input| input.user_query().is_some())
.any(|input| input.display_query().is_some())
{
break;
}
@@ -376,8 +399,10 @@ impl AIRequestUsageModel {
/// 2. user has overage enabled
/// 3. user has bonus grants (either team grants or user grants)
/// 4. user's team plan has pay-as-you-go enabled (enterprise only)
/// 5. user's team is on enterprise with bonus grants auto-reload enable (enterprise only)
/// 6. user has BYOK enabled and has provided at least one API key
/// 5. user's team has enterprise bonus grants auto-reload enabled (enterprise only)
/// 6. user's team has self-serve auto-reload enabled within its monthly spend limit
/// 7. user has BYOK enabled and has either provided at least one API key or
/// connected a Grok subscription
/// Use this method as the starting point for AI availability checking.
pub fn has_any_ai_remaining(&self, ctx: &AppContext) -> bool {
let current_workspace = UserWorkspaces::as_ref(ctx).current_workspace();
@@ -394,21 +419,34 @@ impl AIRequestUsageModel {
let is_payg_enabled = current_workspace
.is_some_and(|w| w.billing_metadata.is_enterprise_pay_as_you_go_enabled());
let is_enterprise_auto_reload_enabled = current_workspace
.is_some_and(|w| w.billing_metadata.is_enterprise_auto_reload_enabled());
let is_self_serve_auto_reload_enabled = current_workspace.is_some_and(|workspace| {
workspace
.billing_metadata
.is_purchase_add_on_credits_policy_enabled()
&& workspace
.settings
.addon_credits_settings
.auto_reload_enabled
&& PricingInfoModel::as_ref(ctx)
.addon_credits_options()
.and_then(|options| workspace.get_auto_reload_price_cents(options))
.is_some_and(|price| !workspace.would_addon_purchase_reach_limit(price))
});
// If you have provided your own API key,
// it doesn't matter if you are out of warp-provided requests.
let has_byo_api_key = UserWorkspaces::as_ref(ctx).is_byo_api_key_enabled()
&& ApiKeyManager::as_ref(ctx).keys().has_any_key();
// If you have provided your own API key or connected a Grok
// subscription, it doesn't matter if you are out of warp-provided requests.
let has_byo_credentials = UserWorkspaces::as_ref(ctx).is_byo_api_key_enabled(ctx)
&& ApiKeyManager::as_ref(ctx).has_any_key();
has_base_plan_ai_requests
|| (user_bonus_credits || workspace_bonus_credits)
|| workspace_has_overages
|| is_payg_enabled
|| is_enterprise_auto_reload_enabled
|| has_byo_api_key
|| is_self_serve_auto_reload_enabled
|| has_byo_credentials
}
pub fn requests_used(&self) -> usize {
@@ -456,6 +494,10 @@ impl AIRequestUsageModel {
self.request_limit_info.next_refresh_time.utc()
}
pub fn next_refresh_time_local(&self) -> DateTime<Local> {
self.next_refresh_time().with_timezone(&Local)
}
pub fn is_unlimited(&self) -> bool {
self.request_limit_info.is_unlimited
}
@@ -492,6 +534,19 @@ impl AIRequestUsageModel {
}
}
pub fn is_ambient_credits_banner_dismissed(&self) -> bool {
self.ambient_credits_banner_dismissed
}
pub fn dismiss_ambient_credits_banner(&mut self, ctx: &mut ModelContext<Self>) {
if self.ambient_credits_banner_dismissed {
return;
}
self.ambient_credits_banner_dismissed = true;
cache_ambient_credits_banner_dismissed(true, ctx);
ctx.emit(AIRequestUsageModelEvent::AmbientCreditsBannerDismissed);
}
pub fn total_workspace_bonus_credits_remaining(&self, uid: WorkspaceUid) -> i32 {
let now = Utc::now();
self.bonus_grants
@@ -509,7 +564,7 @@ impl AIRequestUsageModel {
.unwrap_or(0)
}
fn total_user_interactive_bonus_credits_remaining(&self) -> i32 {
pub fn total_user_interactive_bonus_credits_remaining(&self) -> i32 {
let now = Utc::now();
self.bonus_grants
.iter()
@@ -641,5 +696,5 @@ impl AIRequestUsageModel {
impl SingletonEntity for AIRequestUsageModel {}
#[cfg(test)]
#[path = "request_usage_model_test.rs"]
#[path = "request_usage_model_tests.rs"]
mod tests;