From be1dbb600a21b0183a0ba829ad18f6546b1e3003 Mon Sep 17 00:00:00 2001 From: Ryan Ward Date: Fri, 21 Aug 2026 19:12:14 -0500 Subject: [PATCH] Improve provider reliability and usage visibility --- app/src/ai/agent/conversation.rs | 51 +- app/src/ai/bedrock/response_translator.rs | 6 +- app/src/ai/bedrock/runtime.rs | 41 +- app/src/ai/bedrock/runtime_tests.rs | 16 +- app/src/ai/blocklist/action_model.rs | 32 +- .../agent_view/agent_input_footer/mod.rs | 147 +++- .../agent_input_footer/toolbar_item.rs | 19 +- .../orchestration_conversation_links.rs | 6 +- app/src/ai/blocklist/block/status_bar.rs | 6 + .../ai/blocklist/block/view_impl/common.rs | 16 +- app/src/ai/blocklist/controller.rs | 793 ++++++++++++------ app/src/ai/blocklist/controller_tests.rs | 33 +- app/src/ai/openai/response_translator.rs | 23 +- app/src/ai/runtime/event_translator.rs | 80 +- app/src/ai/runtime/event_translator_tests.rs | 1 + .../ai/runtime/provider_run_coordinator.rs | 21 +- .../runtime/provider_run_coordinator_tests.rs | 57 +- app/src/ai/runtime/rig_request.rs | 11 +- app/src/ai/runtime/rig_request_tests.rs | 26 + app/src/terminal/input/agent.rs | 171 +--- crates/galaxy_agent_core/src/provider_run.rs | 12 +- crates/galaxy_agent_core/src/types.rs | 11 +- crates/galaxy_agent_core/src/types_tests.rs | 4 +- crates/galaxy_agent_rig/src/bedrock_tests.rs | 2 +- crates/galaxy_agent_rig/src/chatgpt.rs | 1 + .../src/openai_compatible_tests.rs | 2 +- crates/galaxy_agent_rig/src/stream.rs | 70 +- crates/galaxy_core/src/paths.rs | 4 +- 28 files changed, 1070 insertions(+), 592 deletions(-) diff --git a/app/src/ai/agent/conversation.rs b/app/src/ai/agent/conversation.rs index 146ada28..38a96ef3 100644 --- a/app/src/ai/agent/conversation.rs +++ b/app/src/ai/agent/conversation.rs @@ -5,6 +5,7 @@ use ai::agent::orchestration_config::{OrchestrationConfig, OrchestrationConfigSt use ai::document::AIDocumentId; use ai::skills::SkillPathOrigin; use chrono::{DateTime, Local, TimeZone}; +use galaxy_agent_core::Usage as ProviderUsage; use galaxy_cli::agent::Harness; use galaxy_core::command::ExitCode; use galaxy_core::execution_mode::AppExecutionMode; @@ -338,6 +339,7 @@ pub struct AIConversation { progressive_summary: Option, messages_summarized_up_to: usize, current_context_tokens: u32, + latest_model_call_usage: ProviderUsage, has_pending_progressive_summary: bool, } @@ -413,6 +415,7 @@ impl AIConversation { progressive_summary: None, messages_summarized_up_to: 0, current_context_tokens: 0, + latest_model_call_usage: ProviderUsage::default(), has_pending_progressive_summary: false, } } @@ -701,6 +704,7 @@ impl AIConversation { progressive_summary, messages_summarized_up_to, current_context_tokens: 0, + latest_model_call_usage: ProviderUsage::default(), has_pending_progressive_summary: false, }) } @@ -765,6 +769,29 @@ impl AIConversation { self.current_context_tokens = val; } + pub fn set_latest_model_call_usage(&mut self, usage: ProviderUsage) { + self.current_context_tokens = u32::try_from( + usage + .input_tokens + .saturating_add(usage.cached_input_tokens) + .saturating_add(usage.cache_creation_input_tokens), + ) + .unwrap_or(u32::MAX); + self.latest_model_call_usage = usage; + } + + pub fn latest_model_call_cache_read_tokens(&self) -> u32 { + u32::try_from(self.latest_model_call_usage.cached_input_tokens).unwrap_or(u32::MAX) + } + + pub fn latest_model_call_cache_write_tokens(&self) -> u32 { + u32::try_from(self.latest_model_call_usage.cache_creation_input_tokens).unwrap_or(u32::MAX) + } + + pub fn latest_model_call_cache_miss_tokens(&self) -> u32 { + u32::try_from(self.latest_model_call_usage.input_tokens).unwrap_or(u32::MAX) + } + pub fn has_pending_progressive_summary(&self) -> bool { self.has_pending_progressive_summary } @@ -2463,10 +2490,16 @@ impl AIConversation { // Update live context token count from this response's input tokens. // This represents the actual current context window size (not cumulative). - let live_input: u32 = token_usage - .iter() - .map(|u| u.total_input + u.input_cache_read + u.input_cache_write) - .sum(); + let live_input = usage_metadata + .as_ref() + .map(|metadata| metadata.total_input_tokens) + .filter(|tokens| *tokens > 0) + .unwrap_or_else(|| { + token_usage + .iter() + .map(|u| u.total_input + u.input_cache_read + u.input_cache_write) + .sum() + }); if live_input > 0 { self.current_context_tokens = live_input; } @@ -4269,10 +4302,7 @@ impl AIConversation { pub fn cache_miss_tokens(&self) -> u32 { self.total_token_usage_by_model .values() - .map(|u| { - u.total_input - .saturating_sub(u.input_cache_read + u.input_cache_write) - }) + .map(|u| u.total_input) .sum() } @@ -4293,10 +4323,7 @@ impl AIConversation { pub fn last_block_cache_miss_tokens(&self) -> u32 { self.last_block_token_usage_by_model .values() - .map(|u| { - u.total_input - .saturating_sub(u.input_cache_read + u.input_cache_write) - }) + .map(|u| u.total_input) .sum() } diff --git a/app/src/ai/bedrock/response_translator.rs b/app/src/ai/bedrock/response_translator.rs index 96ebb82b..92521ce3 100644 --- a/app/src/ai/bedrock/response_translator.rs +++ b/app/src/ai/bedrock/response_translator.rs @@ -668,9 +668,9 @@ pub fn build_stream_finished( }]; let max_context_tokens = context_window_for_model(model_id); + let effective_input = input_tokens + cache_read_input_tokens + cache_write_input_tokens; let context_usage = if max_context_tokens > 0 { - (input_tokens as f32 + cache_read_input_tokens as f32 + cache_write_input_tokens as f32) - / max_context_tokens as f32 + effective_input as f32 / max_context_tokens as f32 } else { 0.0 } @@ -682,7 +682,7 @@ pub fn build_stream_finished( summarized: is_summarization, credits_spent: 0.0, platform_credits_spent: 0.0, - total_input_tokens: input_tokens as u32, + total_input_tokens: effective_input.max(0) as u32, token_usage: vec![], tool_usage_metadata: None, warp_token_usage: std::collections::HashMap::new(), diff --git a/app/src/ai/bedrock/runtime.rs b/app/src/ai/bedrock/runtime.rs index fc377a7a..a6638c05 100644 --- a/app/src/ai/bedrock/runtime.rs +++ b/app/src/ai/bedrock/runtime.rs @@ -340,16 +340,27 @@ impl BedrockStreamTranslator { let Some(usage) = metadata.usage() else { return Ok(Vec::new()); }; + let output_tokens = nonnegative_tokens(usage.output_tokens()); + let cached_input_tokens = + nonnegative_tokens(usage.cache_read_input_tokens().unwrap_or(0)); + let cache_creation_input_tokens = + nonnegative_tokens(usage.cache_write_input_tokens().unwrap_or(0)); + let reported_total_tokens = nonnegative_tokens(usage.total_tokens()); + let total_input_tokens = if reported_total_tokens > 0 { + reported_total_tokens.saturating_sub(output_tokens) + } else { + nonnegative_tokens(usage.input_tokens()) + .saturating_add(cached_input_tokens) + .saturating_add(cache_creation_input_tokens) + }; Ok(vec![AgentEvent::UsageUpdated { usage: Usage { - input_tokens: nonnegative_tokens(usage.input_tokens()), - output_tokens: nonnegative_tokens(usage.output_tokens()), - cached_input_tokens: nonnegative_tokens( - usage.cache_read_input_tokens().unwrap_or(0), - ), - cache_creation_input_tokens: nonnegative_tokens( - usage.cache_write_input_tokens().unwrap_or(0), - ), + input_tokens: total_input_tokens + .saturating_sub(cached_input_tokens) + .saturating_sub(cache_creation_input_tokens), + output_tokens, + cached_input_tokens, + cache_creation_input_tokens, }, }]) } @@ -450,10 +461,22 @@ fn map_bedrock_error(error: impl std::fmt::Display + std::fmt::Debug) -> AgentEr error.recoverable = matches!( kind, AgentErrorKind::RateLimited | AgentErrorKind::Transport - ); + ) || (kind == AgentErrorKind::Provider + && is_transient_provider_error(&normalized)); error } +fn is_transient_provider_error(normalized: &str) -> bool { + normalized.contains("modelnotready") + || normalized.contains("model not ready") + || normalized.contains("serviceunavailable") + || normalized.contains("service unavailable") + || normalized.contains("internalserver") + || normalized.contains("internal server") + || normalized.contains("temporarily unavailable") + || normalized.contains("overloaded") +} + fn protocol_error(message: impl Into) -> AgentError { AgentError::new(AgentErrorKind::Protocol, message) } diff --git a/app/src/ai/bedrock/runtime_tests.rs b/app/src/ai/bedrock/runtime_tests.rs index 7310d1fe..4689609c 100644 --- a/app/src/ai/bedrock/runtime_tests.rs +++ b/app/src/ai/bedrock/runtime_tests.rs @@ -163,9 +163,13 @@ fn metadata(usage: Usage) -> AwsStreamEvent { ConverseStreamMetadataEvent::builder() .usage( TokenUsage::builder() - .input_tokens(usage.input_tokens as i32) + .input_tokens( + (usage.input_tokens + + usage.cached_input_tokens + + usage.cache_creation_input_tokens) as i32, + ) .output_tokens(usage.output_tokens as i32) - .total_tokens((usage.input_tokens + usage.output_tokens) as i32) + .total_tokens(usage.total_tokens() as i32) .cache_read_input_tokens(usage.cached_input_tokens as i32) .cache_write_input_tokens(usage.cache_creation_input_tokens as i32) .build() @@ -348,3 +352,11 @@ fn bedrock_stop_reasons_map_to_domain_reasons() { StopReason::Refusal ); } + +#[test] +fn transient_bedrock_provider_failure_is_recoverable() { + let error = map_bedrock_error("ServiceUnavailableException: model temporarily unavailable"); + + assert_eq!(error.kind, AgentErrorKind::Provider); + assert!(error.recoverable); +} diff --git a/app/src/ai/blocklist/action_model.rs b/app/src/ai/blocklist/action_model.rs index 0c12f839..a0856aca 100644 --- a/app/src/ai/blocklist/action_model.rs +++ b/app/src/ai/blocklist/action_model.rs @@ -857,26 +857,24 @@ impl BlocklistAIActionModel { }); let history_model = BlocklistAIHistoryModel::handle(ctx); - ctx.subscribe_to_model(&history_model, |me, _, event, ctx| { - match event { - BlocklistAIHistoryEvent::RemoveConversation { - conversation_id, .. - } - | BlocklistAIHistoryEvent::DeletedConversation { - conversation_id, .. - } => { + ctx.subscribe_to_model(&history_model, |me, _, event, ctx| match event { + BlocklistAIHistoryEvent::RemoveConversation { + conversation_id, .. + } + | BlocklistAIHistoryEvent::DeletedConversation { + conversation_id, .. + } => { + me.cleanup_conversation_state(*conversation_id, ctx); + } + BlocklistAIHistoryEvent::ClearedConversationsForTerminalSurface { + cleared_conversation_ids, + .. + } => { + for conversation_id in cleared_conversation_ids { me.cleanup_conversation_state(*conversation_id, ctx); } - BlocklistAIHistoryEvent::ClearedConversationsForTerminalSurface { - cleared_conversation_ids, - .. - } => { - for conversation_id in cleared_conversation_ids { - me.cleanup_conversation_state(*conversation_id, ctx); - } - } - _ => {} } + _ => {} }); Self { diff --git a/app/src/ai/blocklist/agent_view/agent_input_footer/mod.rs b/app/src/ai/blocklist/agent_view/agent_input_footer/mod.rs index 3216fa5b..21045899 100644 --- a/app/src/ai/blocklist/agent_view/agent_input_footer/mod.rs +++ b/app/src/ai/blocklist/agent_view/agent_input_footer/mod.rs @@ -55,7 +55,9 @@ use crate::ai::blocklist::prompt::prompt_alert::{PromptAlertEvent, PromptAlertVi use crate::ai::blocklist::usage::icon_for_context_window_usage; use crate::ai::blocklist::BlocklistAIInputModel; use crate::ai::execution_profiles::profiles::AIExecutionProfilesModel; +use crate::ai::execution_profiles::AIExecutionProfileAppExt; use crate::ai::harness_availability::HarnessAvailabilityModel; +use crate::ai::llms::LLMPreferences; use crate::ai::AIRequestUsageModel; use crate::appearance::Appearance; use crate::auth::{AuthManager, AuthStateProvider}; @@ -123,6 +125,16 @@ const FAST_FORWARD_LOCKED_TOOLTIP: &str = const CLOUD_MODE_V2_FOOTER_GAP: f32 = 4.; +fn format_llm_token_count(tokens: u32) -> String { + if tokens >= 1_000_000 { + format!("{:.1}M", tokens as f64 / 1_000_000.0) + } else if tokens >= 1_000 { + format!("{:.1}k", tokens as f64 / 1_000.0) + } else { + tokens.to_string() + } +} + /// Voice input state for the CLI agent footer. Unlike the editor-based voice /// flow (which goes through Input → EditorView), this state is self-contained /// so that transcribed text can be written directly to the PTY. @@ -194,6 +206,8 @@ pub struct AgentInputFooter { nld_button: ViewHandle, file_button: ViewHandle, context_window_button: ViewHandle, + llm_context_usage_button: ViewHandle, + llm_cache_details_button: ViewHandle, model_selector: ViewHandle, environment_selector: Option>, handoff_environment_selector: ViewHandle, @@ -610,6 +624,22 @@ impl AgentInputFooter { .with_tooltip_alignment(TooltipAlignment::Left) }); + let llm_context_usage_button = ctx.add_typed_action_view(|_ctx| { + ActionButton::new("Ctx —", AgentInputButtonTheme) + .with_icon(Icon::ContextRemaining100) + .with_tooltip("LLM context usage is available after the first model call") + .with_size(button_size) + .with_tooltip_alignment(TooltipAlignment::Left) + }); + + let llm_cache_details_button = ctx.add_typed_action_view(|_ctx| { + ActionButton::new("Cache —", AgentInputButtonTheme) + .with_icon(Icon::LayersThree01) + .with_tooltip("LLM cache details are available after the first model call") + .with_size(button_size) + .with_tooltip_alignment(TooltipAlignment::Left) + }); + let profile_model_selector_full = ctx.add_typed_action_view(|ctx| { let mut selector = ProfileModelSelector::new( menu_positioning_provider.clone(), @@ -752,7 +782,12 @@ impl AgentInputFooter { }, ); // Subscribe to AIExecutionProfilesModel to potentially show/hide the profile selector button when profiles are added/removed - ctx.subscribe_to_model(&AIExecutionProfilesModel::handle(ctx), |_, _, _, ctx| { + ctx.subscribe_to_model(&AIExecutionProfilesModel::handle(ctx), |me, _, _, ctx| { + me.update_llm_usage_buttons(ctx); + ctx.notify(); + }); + ctx.subscribe_to_model(&LLMPreferences::handle(ctx), |me, _, _, ctx| { + me.update_llm_usage_buttons(ctx); ctx.notify(); }); @@ -775,6 +810,7 @@ impl AgentInputFooter { | BlocklistAIHistoryEvent::UpdatedAutoexecuteOverride { .. } => { me.sync_fast_forward_button(ctx); me.update_context_window_button(ctx); + me.update_llm_usage_buttons(ctx); me.model_selector.update(ctx, |_, ctx| ctx.notify()); ctx.notify(); } @@ -783,6 +819,7 @@ impl AgentInputFooter { | BlocklistAIHistoryEvent::AppendedExchange { .. } | BlocklistAIHistoryEvent::UpdatedStreamingExchange { .. } => { me.update_context_window_button(ctx); + me.update_llm_usage_buttons(ctx); me.model_selector.update(ctx, |_, ctx| ctx.notify()); ctx.notify(); } @@ -836,6 +873,8 @@ impl AgentInputFooter { plugin_operation_in_progress: false, plugin_chip_ready: false, context_window_button, + llm_context_usage_button, + llm_cache_details_button, model_selector: profile_model_selector_full, environment_selector, handoff_environment_selector, @@ -860,6 +899,7 @@ impl AgentInputFooter { }; me.sync_fast_forward_button(ctx); me.update_context_window_button(ctx); + me.update_llm_usage_buttons(ctx); me.update_display_chips(&prompt, ctx); me } @@ -1463,6 +1503,8 @@ impl AgentInputFooter { AgentToolbarItemKind::ModelSelector | AgentToolbarItemKind::NLDToggle | AgentToolbarItemKind::ContextWindowUsage + | AgentToolbarItemKind::LLMContextUsage + | AgentToolbarItemKind::LLMCacheDetails | AgentToolbarItemKind::FastForwardToggle | AgentToolbarItemKind::HandoffToCloud => None, } @@ -1970,6 +2012,101 @@ impl AgentInputFooter { } } + fn update_llm_usage_buttons(&self, ctx: &mut ViewContext) { + let metrics = BlocklistAIHistoryModel::as_ref(ctx) + .active_conversation(self.terminal_view_id) + .map(|conversation| { + ( + conversation.current_context_tokens(), + conversation.context_window_usage(), + conversation.latest_model_call_cache_read_tokens(), + conversation.latest_model_call_cache_write_tokens(), + conversation.latest_model_call_cache_miss_tokens(), + ) + }); + + let Some((current_context, reported_usage, cache_read, cache_write, cache_miss)) = metrics + else { + self.llm_context_usage_button.update(ctx, |button, ctx| { + button.set_label("Ctx —", ctx); + button.set_icon(Some(Icon::ContextRemaining100), ctx); + button.set_tooltip( + Some("LLM context usage is available after the first model call"), + ctx, + ); + }); + self.llm_cache_details_button.update(ctx, |button, ctx| { + button.set_label("Cache —", ctx); + button.set_tooltip( + Some("LLM cache details are available after the first model call"), + ctx, + ); + }); + return; + }; + + let active_model = + LLMPreferences::as_ref(ctx).get_active_base_model(ctx, Some(self.terminal_view_id)); + let profile_context = AIExecutionProfilesModel::as_ref(ctx) + .active_profile(Some(self.terminal_view_id), ctx) + .data() + .context_window_display_value(ctx); + let model_max_context = active_model + .context_window + .default_max + .max(active_model.context_window.max); + let max_context = profile_context + .or((model_max_context > 0).then_some(model_max_context)) + .unwrap_or(200_000); + let context_usage = if max_context > 0 && current_context > 0 { + (current_context as f32 / max_context as f32).clamp(0.0, 1.0) + } else { + reported_usage.clamp(0.0, 1.0) + }; + let context_pct = context_usage * 100.0; + let context_label = format!( + "Ctx {context_pct:.1}% · {}/{}", + format_llm_token_count(current_context), + format_llm_token_count(max_context), + ); + let context_tooltip = format!( + "Current LLM context: {} of {} tokens ({context_pct:.1}%). This is not cumulative.", + format_llm_token_count(current_context), + format_llm_token_count(max_context), + ); + let context_icon = icon_for_context_window_usage(context_usage); + self.llm_context_usage_button.update(ctx, |button, ctx| { + button.set_label(context_label, ctx); + button.set_icon(Some(context_icon), ctx); + button.set_tooltip(Some(context_tooltip), ctx); + }); + + let cache_total = cache_read + .saturating_add(cache_write) + .saturating_add(cache_miss); + let cache_hit_pct = if cache_total > 0 { + cache_read as f64 / cache_total as f64 * 100.0 + } else { + 0.0 + }; + let cache_label = format!( + "Cache {cache_hit_pct:.1}% · R {} · W {} · M {}", + format_llm_token_count(cache_read), + format_llm_token_count(cache_write), + format_llm_token_count(cache_miss), + ); + let cache_tooltip = format!( + "Latest model call: {} cache-read, {} cache-write, {} cache-miss tokens.", + format_llm_token_count(cache_read), + format_llm_token_count(cache_write), + format_llm_token_count(cache_miss), + ); + self.llm_cache_details_button.update(ctx, |button, ctx| { + button.set_label(cache_label, ctx); + button.set_tooltip(Some(cache_tooltip), ctx); + }); + } + /// Schedules a refresh of the context-window button at the prompt-cache /// expiry instant so the notification dot appears while the conversation is idle. fn reschedule_prompt_cache_expiry_timer( @@ -2101,6 +2238,14 @@ impl AgentInputFooter { stack.finish() }) } + AgentToolbarItemKind::LLMContextUsage => BlocklistAIHistoryModel::as_ref(app) + .active_conversation(self.terminal_view_id) + .is_some() + .then(|| ChildView::new(&self.llm_context_usage_button).finish()), + AgentToolbarItemKind::LLMCacheDetails => BlocklistAIHistoryModel::as_ref(app) + .active_conversation(self.terminal_view_id) + .is_some() + .then(|| ChildView::new(&self.llm_cache_details_button).finish()), AgentToolbarItemKind::ShareSession => { if is_conversation_transcript_context { return None; diff --git a/app/src/ai/blocklist/agent_view/agent_input_footer/toolbar_item.rs b/app/src/ai/blocklist/agent_view/agent_input_footer/toolbar_item.rs index ac5b1d4d..063ca0cc 100644 --- a/app/src/ai/blocklist/agent_view/agent_input_footer/toolbar_item.rs +++ b/app/src/ai/blocklist/agent_view/agent_input_footer/toolbar_item.rs @@ -52,6 +52,8 @@ pub enum AgentToolbarItemKind { ModelSelector, NLDToggle, ContextWindowUsage, + LLMContextUsage, + LLMCacheDetails, // CLI agent only FileExplorer, @@ -83,6 +85,8 @@ impl AgentToolbarItemKind { Self::ModelSelector | Self::NLDToggle | Self::ContextWindowUsage + | Self::LLMContextUsage + | Self::LLMCacheDetails | Self::FastForwardToggle | Self::HandoffToCloud | Self::ShareSession => ToolbarAvailability::AgentViewOnly, @@ -111,6 +115,8 @@ impl AgentToolbarItemKind { | Self::ModelSelector | Self::NLDToggle | Self::ContextWindowUsage + | Self::LLMContextUsage + | Self::LLMCacheDetails | Self::RichInput | Self::VoiceInput => true, } @@ -123,7 +129,9 @@ impl AgentToolbarItemKind { Self::NLDToggle => "Autodetection", Self::VoiceInput => "Voice Input", Self::FileAttach => "Attach File", - Self::ContextWindowUsage => "Context Usage", + Self::ContextWindowUsage => "Warp Context Indicator", + Self::LLMContextUsage => "LLM Context Usage", + Self::LLMCacheDetails => "LLM Cache Details", Self::FileExplorer => "File Explorer", Self::RichInput => "Rich Input", Self::Settings => "Settings", @@ -141,6 +149,8 @@ impl AgentToolbarItemKind { Self::VoiceInput => Some(Icon::Microphone), Self::FileAttach => Some(Icon::Plus), Self::ContextWindowUsage => Some(Icon::ContextRemaining100), + Self::LLMContextUsage => Some(Icon::ContextRemaining100), + Self::LLMCacheDetails => Some(Icon::LayersThree01), Self::FileExplorer => Some(Icon::FileCopy), Self::RichInput => Some(Icon::TextInput), Self::Settings => Some(Icon::Settings), @@ -163,6 +173,8 @@ impl AgentToolbarItemKind { Self::ContextChip(_) | Self::NLDToggle | Self::ContextWindowUsage + | Self::LLMContextUsage + | Self::LLMCacheDetails | Self::FastForwardToggle | Self::HandoffToCloud | Self::ShareSession @@ -214,7 +226,8 @@ impl AgentToolbarItemKind { pub fn default_right() -> Vec { let mut items = vec![ Self::ContextChip(ContextChipKind::AgentPlanAndTodoList), - Self::ContextWindowUsage, + Self::LLMContextUsage, + Self::LLMCacheDetails, Self::ModelSelector, ]; if FeatureFlag::OzHandoff.is_enabled() @@ -240,6 +253,8 @@ impl AgentToolbarItemKind { Self::VoiceInput, Self::FileAttach, Self::ContextWindowUsage, + Self::LLMContextUsage, + Self::LLMCacheDetails, ]); if FeatureFlag::FastForwardAutoexecuteButton.is_enabled() { items.push(Self::FastForwardToggle); diff --git a/app/src/ai/blocklist/agent_view/orchestration_conversation_links.rs b/app/src/ai/blocklist/agent_view/orchestration_conversation_links.rs index 55f3a2ee..3f73f2a1 100644 --- a/app/src/ai/blocklist/agent_view/orchestration_conversation_links.rs +++ b/app/src/ai/blocklist/agent_view/orchestration_conversation_links.rs @@ -81,11 +81,7 @@ pub(crate) fn conversation_id_for_agent_id_in_orchestration( history, orchestrator_id, )) - .find(|conversation_id| { - history - .conversation(conversation_id) - .is_some_and(|conversation| matches(conversation)) - }) + .find(|conversation_id| history.conversation(conversation_id).is_some_and(&matches)) .or_else(|| conversation_id_for_agent_id(canonical_id, app)) } diff --git a/app/src/ai/blocklist/block/status_bar.rs b/app/src/ai/blocklist/block/status_bar.rs index 00209cc7..660a90b7 100644 --- a/app/src/ai/blocklist/block/status_bar.rs +++ b/app/src/ai/blocklist/block/status_bar.rs @@ -1013,6 +1013,11 @@ impl BlocklistAIStatusBar { .or(self.warping_message) .unwrap_or_else(|| random_load_output_message()) .to_owned(); + let retry_status_text = self + .controller + .as_ref(app) + .provider_retry_status(conversation.id()) + .map(|status| status.label()); let secondary_element = if fallback_warping_text.is_some() { Some(render_fallback_explanation(model.as_ref(), app)) } else { @@ -1071,6 +1076,7 @@ impl BlocklistAIStatusBar { )), force_refresh_button, default_warping_text, + retry_status_text, secondary_element, last_snapshot_at, warping_start_time: self.warping_start_time, diff --git a/app/src/ai/blocklist/block/view_impl/common.rs b/app/src/ai/blocklist/block/view_impl/common.rs index 3ebaf119..20038973 100644 --- a/app/src/ai/blocklist/block/view_impl/common.rs +++ b/app/src/ai/blocklist/block/view_impl/common.rs @@ -218,6 +218,7 @@ pub struct WarpingProps<'a, V> { pub action_model: &'a BlocklistAIActionModel, pub terminal_model: &'a TerminalModel, pub default_warping_text: String, + pub retry_status_text: Option, pub secondary_element: Option>, /// When an LRC subagent has sent at least one snapshot, the timestamp of the most recent snapshot. pub last_snapshot_at: Option, @@ -337,14 +338,6 @@ pub fn render_warping_indicator( let mut should_render_waiting_icon = false; let mut non_shimmering_text = None; - if let Some(status_message) = props - .model - .conversation(app) - .and_then(|conversation| conversation.status_error_message()) - .filter(|message| message.starts_with("Retrying LLM request")) - { - non_shimmering_text = Some(format!(" • {status_message}")); - } let message = if let Some(summarization_type) = summarization_type { // Choose the appropriate message based on summarization type let base_message = match summarization_type { @@ -486,6 +479,13 @@ pub fn render_warping_indicator( } }; + if let Some(retry_status_text) = props.retry_status_text.as_deref() { + non_shimmering_text = Some(match non_shimmering_text { + Some(text) if !text.is_empty() => format!("{text} {retry_status_text}"), + Some(_) | None => format!(" {retry_status_text}"), + }); + } + let appearance = Appearance::as_ref(app); let mut buttons_row = Flex::row().with_cross_axis_alignment(CrossAxisAlignment::Center); diff --git a/app/src/ai/blocklist/controller.rs b/app/src/ai/blocklist/controller.rs index 5ca37e77..d5f98a58 100644 --- a/app/src/ai/blocklist/controller.rs +++ b/app/src/ai/blocklist/controller.rs @@ -10,6 +10,7 @@ mod pending_response_streams; pub mod response_stream; pub(super) mod shared_session; mod slash_command; +use std::collections::hash_map::Entry; use std::collections::{BTreeMap, HashMap, HashSet, VecDeque}; #[cfg(not(target_family = "wasm"))] use std::path::PathBuf; @@ -20,11 +21,12 @@ use ai::skills::SkillPathOrigin; use anyhow::anyhow; use chrono::{DateTime, Local}; use futures::channel::oneshot; +use futures::{FutureExt, StreamExt}; use galaxy_agent_core::{ - turn_control, ExternalWorkId, PendingToolBatch, PendingToolCallState, ProviderRun, + turn_control, AgentEvent, ExternalWorkId, PendingToolBatch, PendingToolCallState, ProviderRun, ProviderRunFailureKind, ProviderRunId, ProviderRunLimits, ProviderRunOutcome, ProviderRunState, StopReason, ToolLoopGuard, ToolResult, ToolResultStatus, TurnCommand, TurnCommandSender, - TurnRequest, + TurnRequest, Usage, }; use galaxy_core::assertions::safe_assert; use input_context::{input_context_for_request, parse_context_attachments}; @@ -73,8 +75,9 @@ use crate::ai::ambient_agents::AmbientAgentTaskId; use crate::ai::document::ai_document_model::{ AIDocumentId, AIDocumentModel, AIDocumentUserEditStatus, }; -use crate::ai::llms::{LLMId, LLMPreferences}; -use crate::ai::provider::types::{ContentPart, ConversationMessage, MessageContent}; +use crate::ai::llms::{LLMId, LLMInfo, LLMPreferences}; +use crate::ai::provider::types::{ContentPart, ConversationMessage, MessageContent, MessageRole}; +use crate::ai::provider::ProviderConfig; #[cfg(not(target_family = "wasm"))] use crate::ai::remote_logging::{self, RemoteLogLevel, RemoteLogRecord}; use crate::ai::runtime::{ @@ -106,6 +109,182 @@ use crate::terminal::ShellLaunchData; use crate::workspaces::update_manager::TeamUpdateManager; use crate::workspaces::user_workspaces::UserWorkspaces; +const PROGRESSIVE_SUMMARY_TRIGGER_USAGE: f32 = 0.70; +const PROGRESSIVE_SUMMARY_RETAINED_USAGE: f32 = 0.50; +const PROGRESSIVE_SUMMARY_MIN_RECENT_MESSAGES: usize = 20; +const PROGRESSIVE_SUMMARY_OUTPUT_TOKENS: u32 = 8_000; +const PROGRESSIVE_SUMMARY_CONTEXT_SAFETY_TOKENS: u32 = 2_000; +const PROGRESSIVE_SUMMARY_START_TIMEOUT: Duration = Duration::from_secs(120); +const PROGRESSIVE_SUMMARY_EVENT_TIMEOUT: Duration = Duration::from_secs(300); + +#[derive(Clone)] +struct ProgressiveSummaryCandidate { + model_id: String, + context_limit: u32, + provider_config: ProviderConfig, +} + +fn configured_llm_context_limit(info: &LLMInfo) -> u32 { + [ + info.context_window.default_max, + info.context_window.max, + info.context_window.min, + ] + .into_iter() + .find(|limit| *limit > 0) + .unwrap_or(128_000) +} + +fn estimated_text_tokens(text: &str) -> u32 { + u32::try_from(text.chars().count().div_ceil(4)).unwrap_or(u32::MAX) +} + +fn estimated_message_tokens(message: &ConversationMessage) -> u32 { + let content_tokens = match &message.content { + MessageContent::Text(text) => estimated_text_tokens(text), + MessageContent::ToolUse { name, input, .. } => estimated_text_tokens(name) + .saturating_add(estimated_text_tokens(&input.to_string())) + .saturating_add(20), + MessageContent::ToolResult { content, .. } => { + estimated_text_tokens(content).saturating_add(20) + } + MessageContent::MultiPart(parts) => parts.iter().fold(0_u32, |total, part| { + let part_tokens = match part { + ContentPart::Text(text) | ContentPart::Reasoning { text, .. } => { + estimated_text_tokens(text) + } + ContentPart::Image { .. } => 1_600, + ContentPart::ToolUse { name, input, .. } => estimated_text_tokens(name) + .saturating_add(estimated_text_tokens(&input.to_string())) + .saturating_add(20), + ContentPart::ToolResult { content, .. } => { + estimated_text_tokens(content).saturating_add(20) + } + }; + total.saturating_add(part_tokens) + }), + }; + content_tokens.saturating_add(8) +} + +fn estimated_history_tokens(messages: &[ConversationMessage]) -> u32 { + messages.iter().fold(0_u32, |total, message| { + total.saturating_add(estimated_message_tokens(message)) + }) +} + +fn content_contains_tool_result(content: &MessageContent) -> bool { + match content { + MessageContent::ToolResult { .. } => true, + MessageContent::MultiPart(parts) => parts + .iter() + .any(|part| matches!(part, ContentPart::ToolResult { .. })), + MessageContent::Text(_) | MessageContent::ToolUse { .. } => false, + } +} + +fn progressive_summary_split_point( + messages: &[ConversationMessage], + retained_token_budget: u32, +) -> usize { + if messages.len() <= PROGRESSIVE_SUMMARY_MIN_RECENT_MESSAGES { + return 0; + } + + let mut retained_tokens = 0_u32; + let mut split_point = messages.len(); + for index in (0..messages.len()).rev() { + let message_tokens = estimated_message_tokens(&messages[index]); + let retained_count = messages.len() - index; + if retained_count > PROGRESSIVE_SUMMARY_MIN_RECENT_MESSAGES + && retained_tokens.saturating_add(message_tokens) > retained_token_budget + { + break; + } + retained_tokens = retained_tokens.saturating_add(message_tokens); + split_point = index; + } + + // Do not leave a tool result at the start of retained history without its + // preceding assistant tool call. + while split_point > 0 + && messages + .get(split_point) + .is_some_and(|message| content_contains_tool_result(&message.content)) + { + split_point -= 1; + } + + split_point +} + +async fn collect_progressive_summary( + provider_config: ProviderConfig, + request: TurnRequest, +) -> anyhow::Result<(String, Usage)> { + let (command_sender, control) = turn_control(); + let startup = async move { + let runtime = provider_runtime_for_request(provider_config, &request).await?; + runtime + .start_turn(request, control) + .await + .map_err(|error| anyhow!(error.to_string())) + } + .fuse(); + let startup_timeout = FutureExt::fuse(Timer::after(PROGRESSIVE_SUMMARY_START_TIMEOUT)); + futures::pin_mut!(startup, startup_timeout); + let mut stream = futures::select_biased! { + result = startup => result?, + _ = startup_timeout => { + let _ = command_sender.send(TurnCommand::Cancel).await; + anyhow::bail!("summary model did not start within {} seconds", PROGRESSIVE_SUMMARY_START_TIMEOUT.as_secs()); + } + }; + + let mut summary = String::new(); + let mut usage = Usage::default(); + let stop_reason = loop { + let next_event = stream.next().fuse(); + let event_timeout = FutureExt::fuse(Timer::after(PROGRESSIVE_SUMMARY_EVENT_TIMEOUT)); + futures::pin_mut!(next_event, event_timeout); + let event = futures::select_biased! { + event = next_event => event, + _ = event_timeout => { + let _ = command_sender.send(TurnCommand::Cancel).await; + anyhow::bail!("summary model was idle for {} seconds", PROGRESSIVE_SUMMARY_EVENT_TIMEOUT.as_secs()); + } + }; + let Some(event) = event else { + anyhow::bail!("summary model stream ended without a completion event"); + }; + match event.map_err(|error| anyhow!(error.to_string()))? { + AgentEvent::TextDelta { text } => summary.push_str(&text), + AgentEvent::UsageUpdated { usage: event_usage } => usage = event_usage, + AgentEvent::TurnStopped { reason } => break reason, + AgentEvent::Tool { .. } => { + anyhow::bail!("summary model unexpectedly requested a tool") + } + AgentEvent::TurnStarted { .. } + | AgentEvent::KeepAlive + | AgentEvent::ReasoningDelta { .. } + | AgentEvent::ReasoningCompleted { .. } + | AgentEvent::RuntimeActivityUpdated { .. } + | AgentEvent::ContextUsageUpdated { .. } + | AgentEvent::UserInputAccepted { .. } + | AgentEvent::RuntimeNotice { .. } => {} + } + }; + + if stop_reason != StopReason::Completed { + anyhow::bail!("summary model stopped with {stop_reason:?}"); + } + if summary.trim().is_empty() { + anyhow::bail!("summary model returned an empty response"); + } + + Ok((summary, usage)) +} + #[derive(Debug, Clone)] pub struct SessionContext { session_type: Option, @@ -657,6 +836,18 @@ fn convert_provider_tool_batch( (actions, invalid_results) } +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) struct ProviderRetryStatus { + attempt: u32, + max_retries: u32, +} + +impl ProviderRetryStatus { + pub(crate) fn label(self) -> String { + format!("(Retry {}/{})", self.attempt, self.max_retries) + } +} + struct ActiveProviderRunSlot { stream_id: ResponseStreamId, response_stream: ModelHandle, @@ -675,6 +866,7 @@ struct ActiveProviderRunSlot { pending_monitor_observation: Option, pending_command_completion: Option, monitor_prose_continuations: usize, + retry_status: Option, } struct QueuedProviderRun { @@ -1497,6 +1689,7 @@ struct ProviderLlmLifecycle { model_id: String, runtime_request_id: Option, retry_attempt: u32, + max_retries: Option, elapsed_ms: Option, stop_reason: Option, tool_call_count: Option, @@ -1521,6 +1714,7 @@ fn provider_llm_lifecycle(projection: &ProviderRunProjection) -> Option Option Option Option ProviderLlmLifecycle { @@ -1591,6 +1788,7 @@ fn provider_llm_lifecycle(projection: &ProviderRunProjection) -> Option, + latest_usage: Option, events: Vec, acknowledgement: oneshot::Sender>, }, @@ -1889,6 +2089,15 @@ impl InputQuery { } impl BlocklistAIController { + pub(crate) fn provider_retry_status( + &self, + conversation_id: AIConversationId, + ) -> Option { + self.active_provider_runs + .get(&conversation_id) + .and_then(|slot| slot.retry_status) + } + fn has_unresolved_ask_user_question( &self, conversation_id: AIConversationId, @@ -2783,10 +2992,10 @@ impl BlocklistAIController { if self .in_flight_response_streams .has_active_stream_for_conversation(conversation_id, ctx) - && !self + && self .active_provider_runs .get(&conversation_id) - .is_some_and(|slot| slot.cancellation_reason.is_some()) + .is_none_or(|slot| slot.cancellation_reason.is_none()) || self .action_model .as_ref(ctx) @@ -4927,33 +5136,35 @@ impl BlocklistAIController { pending_monitor_observation: None, pending_command_completion: None, monitor_prose_continuations: 0, + retry_status: None, }; - if self - .active_provider_runs - .contains_key(&conversation_data.id) - { - self.queued_provider_runs - .entry(conversation_data.id) - .or_default() - .push_back(QueuedProviderRun { - slot, + match self.active_provider_runs.entry(conversation_data.id) { + Entry::Occupied(_) => { + self.queued_provider_runs + .entry(conversation_data.id) + .or_default() + .push_back(QueuedProviderRun { + slot, + base_provider_config, + cli_provider_config, + request_params: request_params.clone(), + }); + if let Err(error) = self.persist_active_provider_run(conversation_data.id, ctx) + { + log::error!("Failed to persist queued provider follow-up: {error}"); + } + } + Entry::Vacant(entry) => { + entry.insert(slot); + self.prepare_active_provider_run( + conversation_data.id, + response_stream_id.clone(), base_provider_config, cli_provider_config, - request_params: request_params.clone(), - }); - if let Err(error) = self.persist_active_provider_run(conversation_data.id, ctx) { - log::error!("Failed to persist queued provider follow-up: {error}"); + request_params.clone(), + ctx, + ); } - } else { - self.active_provider_runs.insert(conversation_data.id, slot); - self.prepare_active_provider_run( - conversation_data.id, - response_stream_id.clone(), - base_provider_config, - cli_provider_config, - request_params.clone(), - ctx, - ); } } @@ -5547,6 +5758,7 @@ impl BlocklistAIController { pending_monitor_observation, pending_command_completion, monitor_prose_continuations, + retry_status: None, }, ); if let Err(error) = @@ -5750,6 +5962,7 @@ impl BlocklistAIController { pending_monitor_observation: None, pending_command_completion: None, monitor_prose_continuations: 0, + retry_status: None, }, base_provider_config, cli_provider_config, @@ -6103,6 +6316,18 @@ impl BlocklistAIController { turn_control, |projection| { let lifecycle = provider_llm_lifecycle(&projection); + let latest_usage = match &projection { + ProviderRunProjection::ModelEvent { + event: AgentEvent::UsageUpdated { usage }, + .. + } => Some(usage.clone()), + ProviderRunProjection::ModelTurnRequested { .. } + | ProviderRunProjection::ModelTurnStarted { .. } + | ProviderRunProjection::ModelTurnFinished { .. } + | ProviderRunProjection::ModelEvent { .. } + | ProviderRunProjection::ModelRetry { .. } + | ProviderRunProjection::ToolBatchReady { .. } => None, + }; let events = run.projector.project(projection); let projection_sender = projection_sender.clone(); Box::pin(async move { @@ -6111,6 +6336,7 @@ impl BlocklistAIController { projection_sender .send(ProviderDriveMessage::Projection { lifecycle, + latest_usage, events, acknowledgement, }) @@ -6169,9 +6395,17 @@ impl BlocklistAIController { match message { ProviderDriveMessage::Projection { lifecycle, + latest_usage, events, acknowledgement, } => { + if let Some(usage) = latest_usage { + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history, _| { + if let Some(conversation) = history.conversation_mut(&conversation_id) { + conversation.set_latest_model_call_usage(usage); + } + }); + } let response_stream = slot.response_stream.clone(); let did_input_contain_user_query = slot.did_input_contain_user_query; let mut result = Ok(()); @@ -6187,9 +6421,37 @@ impl BlocklistAIController { break; } } - #[cfg(not(target_family = "wasm"))] if let Some(lifecycle) = lifecycle.as_ref() { - if lifecycle.phase == ProviderLlmLifecyclePhase::Requested { + let should_refresh_status = { + let slot = self + .active_provider_runs + .get_mut(&conversation_id) + .expect("provider projection retained its active run slot"); + match lifecycle.phase { + ProviderLlmLifecyclePhase::Requested => { + if lifecycle.retry_attempt == 0 { + slot.retry_status = None; + } + true + } + ProviderLlmLifecyclePhase::Started => false, + ProviderLlmLifecyclePhase::RetryScheduled => { + slot.retry_status = + lifecycle + .max_retries + .map(|max_retries| ProviderRetryStatus { + attempt: lifecycle.retry_attempt, + max_retries, + }); + true + } + ProviderLlmLifecyclePhase::Finished => { + slot.retry_status = None; + true + } + } + }; + if should_refresh_status { BlocklistAIHistoryModel::handle(ctx).update(ctx, |history, ctx| { history.update_conversation_status( self.terminal_surface_id, @@ -6199,25 +6461,7 @@ impl BlocklistAIController { ); }); } - if lifecycle.phase == ProviderLlmLifecyclePhase::RetryScheduled { - let retry_message = format!( - "Retrying LLM request ({}/3): {}", - lifecycle.retry_attempt, - lifecycle - .error - .as_deref() - .unwrap_or("temporary provider error") - ); - BlocklistAIHistoryModel::handle(ctx).update(ctx, |history, ctx| { - history.update_conversation_status_with_error( - self.terminal_surface_id, - conversation_id, - ConversationStatus::TransientError, - Some(RenderableAIError::other(retry_message, false)), - ctx, - ); - }); - } + #[cfg(not(target_family = "wasm"))] remote_logging::log_model_event( ctx, provider_llm_lifecycle_remote_log_record( @@ -6227,29 +6471,6 @@ impl BlocklistAIController { ), ); } - #[cfg(target_family = "wasm")] - if let Some(lifecycle) = lifecycle.as_ref() { - if lifecycle.phase == ProviderLlmLifecyclePhase::Requested { - BlocklistAIHistoryModel::handle(ctx).update(ctx, |history, ctx| { - history.update_conversation_status( - self.terminal_surface_id, - conversation_id, - ConversationStatus::InProgress, - ctx, - ); - }); - } - if lifecycle.phase == ProviderLlmLifecyclePhase::RetryScheduled { - BlocklistAIHistoryModel::handle(ctx).update(ctx, |history, ctx| { - history.update_conversation_status( - self.terminal_surface_id, - conversation_id, - ConversationStatus::TransientError, - ctx, - ); - }); - } - } let _ = acknowledgement.send(result); } ProviderDriveMessage::Checkpoint { @@ -7103,7 +7324,8 @@ impl BlocklistAIController { let offset = run.persistence_offset.min(transcript.len()); *messages_sent = transcript[offset..].to_vec(); } - let events = match run.projector.finish(&outcome) { + let aggregate_usage = run.coordinator.run().usage().clone(); + let events = match run.projector.finish(&outcome, &aggregate_usage) { Ok(events) => events, Err(message) => { self.fail_provider_startup( @@ -8660,10 +8882,9 @@ impl BlocklistAIController { ctx.emit(BlocklistAIControllerEvent::FreeTierLimitCheckTriggered); } - // Progressive summarization: when context window usage >= 85% and we have - // more than 100 messages, summarize the oldest messages while keeping the - // most recent 100 verbatim. This runs as a background Bedrock call — no UI, - // no exchange created, no tool execution shown. + // Progressive summarization runs in the background with no UI exchange or + // tool execution. Token-budget planning below decides how much recent + // history to retain and leaves enough hysteresis before the next summary. let should_progressive_summarize = { let history_model = BlocklistAIHistoryModel::as_ref(ctx); history_model @@ -8676,10 +8897,11 @@ impl BlocklistAIController { .iter() .any(|i| matches!(i, AIAgentInput::SummarizeConversation { .. })) }); - conversation.context_window_usage() >= 0.85 + conversation.context_window_usage() >= PROGRESSIVE_SUMMARY_TRIGGER_USAGE && !conversation.has_pending_progressive_summary() && !is_summarization_request - && conversation.bedrock_message_history().len() > 100 + && conversation.bedrock_message_history().len() + > PROGRESSIVE_SUMMARY_MIN_RECENT_MESSAGES }) }; @@ -8693,82 +8915,79 @@ impl BlocklistAIController { conversation_id: AIConversationId, ctx: &mut ModelContext, ) { - use settings::Setting; - - use crate::ai::bedrock::client::{BedrockClient, BedrockClientConfig}; - use crate::ai::bedrock::convert::{ConversationMessage, MessageContent, MessageRole}; - use crate::ai::bedrock::response_translator::{ - context_window_for_model, estimate_cost_cents, - }; - use crate::settings::ai::AISettings; - - let settings = AISettings::as_ref(ctx); - if !*settings.bedrock_enabled.value() { - return; - } - - let api_key_manager = ::ai::api_keys::ApiKeyManager::as_ref(ctx); - let mut config = BedrockClientConfig { - auth_method: *settings.bedrock_auth_method.value(), - profile: settings.bedrock_profile.value().clone(), - region: settings.bedrock_region.value().clone(), - access_key_id: settings.bedrock_access_key_id.value().clone(), - secret_access_key: settings.bedrock_secret_access_key.value().clone(), - session_token: None, - cross_region_inference: *settings.bedrock_cross_region_inference.value(), - use_rig: false, - }; - - if let ::ai::api_keys::AwsCredentialsState::Loaded { credentials, .. } = - api_key_manager.aws_credentials_state() - { - config.auth_method = crate::settings::BedrockAuthMethod::StaticKeys; - config.access_key_id = credentials.access_key().to_string(); - config.secret_access_key = credentials.secret_key().to_string(); - config.session_token = credentials.session_token().map(|s| s.to_string()); - } - - let config = config.with_external_fallbacks(); - - let cross_region = config.cross_region_inference; - // Use Sonnet for summarization — cheaper and fast enough for this task - let model_id = "us.anthropic.claude-sonnet-4-6-20250514-v1:0".to_string(); - let history_model = BlocklistAIHistoryModel::handle(ctx); + let (terminal_model_id, terminal_context_limit, base_model_id, base_context_limit) = { + let preferences = LLMPreferences::as_ref(ctx); + let terminal_model = + preferences.get_active_cli_agent_model(ctx, Some(self.terminal_surface_id)); + let base_model = preferences.get_active_base_model(ctx, Some(self.terminal_surface_id)); + ( + terminal_model.id.to_string(), + configured_llm_context_limit(terminal_model), + base_model.id.to_string(), + configured_llm_context_limit(base_model), + ) + }; + let terminal_provider_config = + ResponseStream::resolve_provider_config(&terminal_model_id, ctx); + let base_provider_config = ResponseStream::resolve_provider_config(&base_model_id, ctx); - // Extract the messages to summarize and set the guard flag - let (messages_to_summarize, existing_summary, messages_count) = { + let ( + messages_to_summarize, + existing_summary, + messages_count, + retained_messages_count, + active_context_limit, + persistent_context_tokens, + ) = { let history = history_model.as_ref(ctx); let Some(conversation) = history.conversation(&conversation_id) else { return; }; - let history_len = conversation.bedrock_message_history().len(); - let split_point = history_len.saturating_sub(100); + let messages = conversation.bedrock_message_history(); + let history_tokens = estimated_history_tokens(messages); + let existing_summary = conversation.progressive_summary().map(str::to_string); + let existing_summary_tokens = existing_summary + .as_deref() + .map(estimated_text_tokens) + .unwrap_or_default(); + let current_tokens = conversation.current_context_tokens(); + let reported_usage = conversation.context_window_usage(); + let inferred_context_limit = if current_tokens > 0 && reported_usage > 0.0 { + ((current_tokens as f64 / reported_usage as f64).round() as u64) + .min(u64::from(u32::MAX)) as u32 + } else { + 0 + }; + let active_context_limit = if inferred_context_limit > 0 { + inferred_context_limit + } else { + base_context_limit + }; + let persistent_context_tokens = current_tokens + .saturating_sub(history_tokens) + .saturating_sub(existing_summary_tokens); + let retained_total_budget = + (active_context_limit as f32 * PROGRESSIVE_SUMMARY_RETAINED_USAGE) as u32; + let retained_history_budget = retained_total_budget + .saturating_sub(persistent_context_tokens) + .saturating_sub(PROGRESSIVE_SUMMARY_OUTPUT_TOKENS); + let split_point = progressive_summary_split_point(messages, retained_history_budget); if split_point == 0 { return; } - let msgs: Vec = - conversation.bedrock_message_history()[..split_point].to_vec(); - let existing = conversation.progressive_summary().map(str::to_string); - - (msgs, existing, split_point) + ( + messages[..split_point].to_vec(), + existing_summary, + split_point, + messages.len().saturating_sub(split_point), + active_context_limit, + persistent_context_tokens, + ) }; - history_model.update(ctx, |history_model, _| { - if let Some(conversation) = history_model.conversation_mut(&conversation_id) { - conversation.set_has_pending_progressive_summary(true); - } - }); - - log::info!( - "[progressive-summary] Triggering for conversation {:?}: summarizing {} messages, keeping last 100", - conversation_id, - messages_count - ); - - // Build the summarization input let mut summarize_content = String::new(); if let Some(ref prior) = existing_summary { summarize_content.push_str("\n"); @@ -8793,29 +9012,26 @@ impl BlocklistAIController { let content_str = match &msg.content { MessageContent::Text(t) => t.clone(), MessageContent::ToolUse { name, input, .. } => { - format!("[Tool Call: {}] {}", name, input) + format!("[Tool Call: {name}] {input}") } MessageContent::ToolResult { content, .. } => safe_truncate(content, 2000), - MessageContent::MultiPart(parts) => { - use crate::ai::bedrock::convert::ContentPart; - parts - .iter() - .map(|p| match p { - ContentPart::Text(t) => t.clone(), - ContentPart::Reasoning { text, .. } => { - format!("[Reasoning] {text}") - } - ContentPart::Image { .. } => "[Image attachment]".to_string(), - ContentPart::ToolUse { name, input, .. } => { - format!("[Tool: {}] {}", name, input) - } - ContentPart::ToolResult { content, .. } => safe_truncate(content, 2000), - }) - .collect::>() - .join("\n") - } + MessageContent::MultiPart(parts) => parts + .iter() + .map(|p| match p { + ContentPart::Text(t) => t.clone(), + ContentPart::Reasoning { text, .. } => { + format!("[Reasoning] {text}") + } + ContentPart::Image { .. } => "[Image attachment]".to_string(), + ContentPart::ToolUse { name, input, .. } => { + format!("[Tool: {name}] {input}") + } + ContentPart::ToolResult { content, .. } => safe_truncate(content, 2000), + }) + .collect::>() + .join("\n"), }; - summarize_content.push_str(&format!("[{}]: {}\n", role_str, content_str)); + summarize_content.push_str(&format!("[{role_str}]: {content_str}\n")); } summarize_content.push_str(""); @@ -8827,135 +9043,170 @@ impl BlocklistAIController { - Technical details, code patterns, and architecture discussed\n\n\ Be comprehensive. This summary will be the only record of these exchanges."; + let estimated_summary_input_tokens = estimated_text_tokens(summarize_prompt) + .saturating_add(estimated_text_tokens(&summarize_content)); + let candidate_fits = |context_limit: u32| { + estimated_summary_input_tokens + .saturating_add(PROGRESSIVE_SUMMARY_OUTPUT_TOKENS) + .saturating_add(PROGRESSIVE_SUMMARY_CONTEXT_SAFETY_TOKENS) + <= context_limit + }; + let model_is_usable = |model_id: &str, provider_config: &ProviderConfig| { + !model_id.trim().is_empty() + && !model_id.eq_ignore_ascii_case("placeholder") + && !matches!(provider_config, ProviderConfig::None) + }; + let mut candidates = Vec::new(); + if model_is_usable(&terminal_model_id, &terminal_provider_config) + && candidate_fits(terminal_context_limit) + { + candidates.push(ProgressiveSummaryCandidate { + model_id: terminal_model_id.clone(), + context_limit: terminal_context_limit, + provider_config: terminal_provider_config, + }); + } + if model_is_usable(&base_model_id, &base_provider_config) + && candidate_fits(base_context_limit) + && !candidates + .iter() + .any(|candidate| candidate.model_id == base_model_id) + { + candidates.push(ProgressiveSummaryCandidate { + model_id: base_model_id.clone(), + context_limit: base_context_limit, + provider_config: base_provider_config, + }); + } + if candidates.is_empty() { + log::warn!( + "[progressive-summary] No configured model can fit ~{} input tokens for conversation {:?}", + estimated_summary_input_tokens, + conversation_id, + ); + return; + } + + history_model.update(ctx, |history_model, _| { + if let Some(conversation) = history_model.conversation_mut(&conversation_id) { + conversation.set_has_pending_progressive_summary(true); + } + }); + + let candidate_names = candidates + .iter() + .map(|candidate| format!("{}({})", candidate.model_id, candidate.context_limit)) + .join(", "); + log::info!( + "[progressive-summary] Triggering for {:?}: summarizing {} messages, retaining {}, candidates=[{}]", + conversation_id, + messages_count, + retained_messages_count, + candidate_names, + ); + let summarize_messages = vec![ConversationMessage { role: MessageRole::User, - content: MessageContent::Text(format!("{}\n\n{}", summarize_prompt, summarize_content)), + content: MessageContent::Text(summarize_content), }]; - // Spawn the background Bedrock call - let model_id_clone = model_id.clone(); ctx.spawn( async move { - let client = BedrockClient::from_config(config).await?; - client - .converse_collect( - &model_id_clone, - summarize_messages, - None, - 16000, - cross_region, - ) - .await + let mut errors = Vec::new(); + for candidate in candidates { + let mut request = + TurnRequest::new(candidate.model_id.clone(), summarize_messages.clone()); + request.system_prompt = Some(summarize_prompt.to_string()); + request.max_output_tokens = Some(u64::from(PROGRESSIVE_SUMMARY_OUTPUT_TOKENS)); + match collect_progressive_summary(candidate.provider_config, request).await { + Ok((summary, usage)) => { + return Ok((summary, usage, candidate.model_id)); + } + Err(error) => errors.push(format!("{}: {error}", candidate.model_id)), + } + } + Err(anyhow!( + "all progressive summary candidates failed: {}", + errors.join("; ") + )) }, - move |me, result, ctx| { + move |_, result, ctx| { let history_model = BlocklistAIHistoryModel::handle(ctx); match result { - Ok((summary_text, input_tokens, output_tokens)) => { + Ok((summary_text, usage, summarizer_model_id)) => { log::info!( - "[progressive-summary] Completed for {:?}: {} chars, input={} output={} tokens", + "[progressive-summary] Completed for {:?} with {}: {} chars, input={} output={} tokens", conversation_id, + summarizer_model_id, summary_text.len(), - input_tokens, - output_tokens, + usage.input_tokens, + usage.output_tokens, ); - let cost_cents = estimate_cost_cents( - input_tokens, - output_tokens, - 0, - 0, - &model_id, - ); - - // Use the conversation's active model for context window sizing, - // not the summarizer model. - let active_model_id = crate::ai::llms::LLMPreferences::as_ref(ctx) - .get_active_base_model(ctx, Some(me.terminal_surface_id)) - .id - .to_string(); - + let usage_u32 = |tokens: u64| { + u32::try_from(tokens).unwrap_or(u32::MAX) + }; + let cost_cents = + crate::ai::bedrock::response_translator::estimate_cost_cents( + usage_u32(usage.input_tokens), + usage_u32(usage.output_tokens), + usage_u32(usage.cached_input_tokens), + usage_u32(usage.cache_creation_input_tokens), + &summarizer_model_id, + ); history_model.update(ctx, |history_model, _| { - if let Some(conversation) = + let Some(conversation) = history_model.conversation_mut(&conversation_id) - { - // Drain the summarized messages from history - let drain_count = - messages_count.min(conversation.bedrock_message_history().len()); - let drained: Vec<_> = conversation - .bedrock_message_history() - .iter() - .take(drain_count) - .cloned() - .collect(); - conversation.archive_tool_results(drained); - conversation - .bedrock_message_history_mut() - .drain(0..drain_count); + else { + return; + }; - conversation - .set_progressive_summary(Some(summary_text.clone()), drain_count); - conversation.set_has_pending_progressive_summary(false); + let drain_count = + messages_count.min(conversation.bedrock_message_history().len()); + let drained = conversation + .bedrock_message_history() + .iter() + .take(drain_count) + .cloned() + .collect(); + conversation.archive_tool_results(drained); + conversation + .bedrock_message_history_mut() + .drain(0..drain_count); + conversation + .set_progressive_summary(Some(summary_text.clone()), drain_count); + conversation.set_has_pending_progressive_summary(false); - // Estimate new context window usage - let summary_tokens = (summary_text.len() / 4) as u32; - let remaining_msgs_tokens: u32 = conversation - .bedrock_message_history() - .iter() - .map(|m| match &m.content { - MessageContent::Text(t) => (t.len() / 4) as u32, - MessageContent::ToolUse { input, .. } => { - (input.to_string().len() / 4) as u32 + 20 - } - MessageContent::ToolResult { content, .. } => { - (content.len() / 4) as u32 - } - MessageContent::MultiPart(parts) => { - parts - .iter() - .map(|p| match p { - ContentPart::Text(t) => (t.len() / 4) as u32, - ContentPart::Reasoning { text, .. } => { - (text.len() / 4) as u32 - } - ContentPart::Image { .. } => 1_600, - ContentPart::ToolUse { input, .. } => { - (input.to_string().len() / 4) as u32 - } - ContentPart::ToolResult { content, .. } => { - (content.len() / 4) as u32 - } - }) - .sum() - } - }) - .sum(); + let summary_tokens = estimated_text_tokens(&summary_text); + let remaining_messages_tokens = + estimated_history_tokens(conversation.bedrock_message_history()); + let new_context_tokens = persistent_context_tokens + .saturating_add(summary_tokens) + .saturating_add(remaining_messages_tokens); + let new_usage = + new_context_tokens as f32 / active_context_limit.max(1) as f32; + conversation.set_context_window_usage(new_usage.clamp(0.0, 1.0)); + conversation.set_current_context_tokens(new_context_tokens); - let max_ctx = context_window_for_model(&active_model_id); - let new_usage = - (summary_tokens + remaining_msgs_tokens) as f32 / max_ctx as f32; - conversation.set_context_window_usage(new_usage.clamp(0.0, 1.0)); - conversation - .set_current_context_tokens(summary_tokens + remaining_msgs_tokens); - - log::info!( - "[progressive-summary] Post-summary: ~{} tokens ({:.1}% of {} context), {} messages retained", - summary_tokens + remaining_msgs_tokens, - new_usage * 100.0, - active_model_id, - conversation.bedrock_message_history().len() - ); - } + log::info!( + "[progressive-summary] Post-summary: ~{} tokens ({:.1}% of {} context), {} messages retained", + new_context_tokens, + new_usage * 100.0, + active_context_limit, + conversation.bedrock_message_history().len(), + ); }); - // Update cost tracking history_model.update(ctx, |history_model, ctx| { use warp_multi_agent_api::response_event::stream_finished; let token_usage = vec![stream_finished::TokenUsage { - model_id: "bedrock".to_string(), - total_input: input_tokens, - output: output_tokens, - input_cache_read: 0, - input_cache_write: 0, + model_id: summarizer_model_id, + total_input: usage_u32(usage.input_tokens), + output: usage_u32(usage.output_tokens), + input_cache_read: usage_u32(usage.cached_input_tokens), + input_cache_write: usage_u32( + usage.cache_creation_input_tokens, + ), cost_in_cents: cost_cents, }]; history_model.update_conversation_cost_and_usage_for_request( @@ -8968,11 +9219,10 @@ impl BlocklistAIController { ); }); } - Err(e) => { + Err(error) => { log::error!( - "[progressive-summary] Failed for {:?}: {:?}", + "[progressive-summary] Failed for {:?}: {error:#}", conversation_id, - e ); history_model.update(ctx, |history_model, _| { if let Some(conversation) = @@ -8983,7 +9233,6 @@ impl BlocklistAIController { }); } } - let _ = me; }, ); } diff --git a/app/src/ai/blocklist/controller_tests.rs b/app/src/ai/blocklist/controller_tests.rs index 8cadc63e..2ffeb9b0 100644 --- a/app/src/ai/blocklist/controller_tests.rs +++ b/app/src/ai/blocklist/controller_tests.rs @@ -75,6 +75,7 @@ fn provider_lifecycle_logs_expose_llm_completion_and_sanitize_errors() { runtime_id: "rig:openai".to_owned(), model_id: "test-model".to_owned(), retry_attempt: 1, + max_retries: 3, elapsed_ms: 120_000, error: retry_error, }, @@ -112,6 +113,7 @@ fn provider_lifecycle_logs_expose_llm_completion_and_sanitize_errors() { assert_eq!(records[0].context["llm_finished"], false); assert_eq!(records[1].context["llm_finished"], false); assert_eq!(records[2].context["llm_finished"], false); + assert_eq!(records[2].context["max_retries"], 3); assert_eq!(records[3].context["llm_finished"], true); assert_eq!(records[2].level, RemoteLogLevel::Warn); assert_eq!(records[2].context["error"], "[redacted] connection failed"); @@ -124,6 +126,16 @@ fn provider_lifecycle_logs_expose_llm_completion_and_sanitize_errors() { assert_eq!(records[3].context["tool_call_count"], 2); } +#[test] +fn provider_retry_status_uses_configured_budget_in_label() { + let status = super::ProviderRetryStatus { + attempt: 1, + max_retries: 5, + }; + + assert_eq!(status.label(), "(Retry 1/5)"); +} + #[test] fn provider_terminal_logs_distinguish_clean_completion_from_failure() { let conversation_id = AIConversationId::new(); @@ -455,6 +467,7 @@ fn queued_provider_follow_up_persists_while_active_slot_is_unprepared() { pending_monitor_observation: None, pending_command_completion: None, monitor_prose_continuations: 0, + retry_status: None, }, ); controller @@ -480,6 +493,7 @@ fn queued_provider_follow_up_persists_while_active_slot_is_unprepared() { pending_monitor_observation: None, pending_command_completion: None, monitor_prose_continuations: 0, + retry_status: None, }, base_provider_config: crate::ai::provider::ProviderConfig::None, cli_provider_config: crate::ai::provider::ProviderConfig::None, @@ -1029,6 +1043,7 @@ fn cancelling_provider_startup_keeps_slot_and_durable_cancellation() { pending_monitor_observation: None, pending_command_completion: None, monitor_prose_continuations: 0, + retry_status: None, }, ); @@ -1127,6 +1142,7 @@ fn same_conversation_follow_up_waits_for_cancelled_provider_generation_cleanup() pending_monitor_observation: None, pending_command_completion: None, monitor_prose_continuations: 0, + retry_status: None, }, ); controller @@ -1158,6 +1174,7 @@ fn same_conversation_follow_up_waits_for_cancelled_provider_generation_cleanup() pending_monitor_observation: None, pending_command_completion: None, monitor_prose_continuations: 0, + retry_status: None, }, base_provider_config: crate::ai::provider::ProviderConfig::None, cli_provider_config: crate::ai::provider::ProviderConfig::None, @@ -1267,6 +1284,7 @@ fn non_follow_up_provider_cancellation_does_not_admit_an_overlapping_generation( pending_monitor_observation: None, pending_command_completion: None, monitor_prose_continuations: 0, + retry_status: None, }, ); @@ -2033,18 +2051,9 @@ fn restored_evidence_is_ignored_without_a_command_monitor() { #[test] fn restored_projection_accepts_empty_or_complete_and_rejects_partial_state() { - assert_eq!( - super::restored_projection_was_initialized(false, false, false).unwrap(), - false - ); - assert_eq!( - super::restored_projection_was_initialized(true, true, false).unwrap(), - true - ); - assert_eq!( - super::restored_projection_was_initialized(true, true, true).unwrap(), - true - ); + assert!(!super::restored_projection_was_initialized(false, false, false).unwrap()); + assert!(super::restored_projection_was_initialized(true, true, false).unwrap()); + assert!(super::restored_projection_was_initialized(true, true, true).unwrap()); for state in [ (false, false, true), (false, true, false), diff --git a/app/src/ai/openai/response_translator.rs b/app/src/ai/openai/response_translator.rs index 6354941d..4a36ca66 100644 --- a/app/src/ai/openai/response_translator.rs +++ b/app/src/ai/openai/response_translator.rs @@ -39,6 +39,9 @@ pub(crate) struct StreamUsage { pub(crate) output_tokens: i32, pub(crate) cache_read_tokens: i32, pub(crate) cache_write_tokens: i32, + /// Input currently occupying the model context. When absent, the + /// per-request token fields above are used. + pub(crate) current_context_tokens: Option, pub(crate) cost_in_cents: f32, pub(crate) model_id: String, pub(crate) max_context_tokens: Option, @@ -376,14 +379,13 @@ pub fn openai_stream_to_response_events( } } - // If we got cache_read but no explicit cache_write, infer it: - // cache_write = prompt_tokens - cache_read (the non-cached input that will be cached) - if cache_read_tokens > 0 && cache_write_tokens == 0 { - cache_write_tokens = (input_tokens - cache_read_tokens).max(0); - } + let current_context_tokens = input_tokens; + let cache_miss_tokens = input_tokens + .saturating_sub(cache_read_tokens) + .saturating_sub(cache_write_tokens); let cost = estimate_cost_cents( - input_tokens as u32, + cache_miss_tokens as u32, output_tokens as u32, cache_read_tokens as u32, cache_write_tokens as u32, @@ -392,10 +394,11 @@ pub fn openai_stream_to_response_events( let finished_event = build_stream_finished( stop_reason, StreamUsage { - input_tokens, + input_tokens: cache_miss_tokens, output_tokens, cache_read_tokens, cache_write_tokens, + current_context_tokens: Some(current_context_tokens), cost_in_cents: cost, model_id: model_id.clone(), max_context_tokens, @@ -545,6 +548,7 @@ pub(crate) fn build_stream_finished( output_tokens, cache_read_tokens, cache_write_tokens, + current_context_tokens, cost_in_cents, model_id, max_context_tokens, @@ -577,7 +581,8 @@ pub(crate) fn build_stream_finished( let max_context_tokens = max_context_tokens.unwrap_or_else(|| context_window_for_model(&model_id)); // Context usage should reflect the full input including cached tokens - let effective_input = input_tokens + cache_read_tokens + cache_write_tokens; + let effective_input = + current_context_tokens.unwrap_or(input_tokens + cache_read_tokens + cache_write_tokens); let context_usage = if max_context_tokens > 0 { effective_input as f32 / max_context_tokens as f32 } else { @@ -591,7 +596,7 @@ pub(crate) fn build_stream_finished( summarized: false, credits_spent: 0.0, platform_credits_spent: 0.0, - total_input_tokens: input_tokens as u32, + total_input_tokens: effective_input.max(0) as u32, token_usage: vec![], tool_usage_metadata: None, warp_token_usage: std::collections::HashMap::new(), diff --git a/app/src/ai/runtime/event_translator.rs b/app/src/ai/runtime/event_translator.rs index 948d4e42..b0b31111 100644 --- a/app/src/ai/runtime/event_translator.rs +++ b/app/src/ai/runtime/event_translator.rs @@ -40,6 +40,7 @@ pub(crate) struct RuntimeResponseTranslator { activity_message_ids: HashMap, activities: HashMap, has_visible_output: bool, + /// Usage for the most recent model call. usage: Usage, context_usage: Option<(u64, u64)>, } @@ -107,25 +108,22 @@ impl ProviderRunResponseProjector { pub(crate) fn finish( &mut self, outcome: &ProviderRunOutcome, + aggregate_usage: &Usage, ) -> Result, String> { if self.finished { return Err("provider run projection is already finished".to_string()); } self.finished = true; match outcome { - ProviderRunOutcome::Completed(completion) => { - self.translator.translate(AgentEvent::TurnStopped { - reason: completion.stop_reason.clone(), - }) - } - ProviderRunOutcome::Failed(failure) => { - Ok(self.translator.provider_failure(&failure.message)) - } - ProviderRunOutcome::Cancelled { .. } => { - self.translator.translate(AgentEvent::TurnStopped { - reason: StopReason::Cancelled, - }) - } + ProviderRunOutcome::Completed(completion) => Ok(self + .translator + .finish_provider_run(completion.stop_reason.clone(), aggregate_usage)), + ProviderRunOutcome::Failed(failure) => Ok(self + .translator + .provider_failure(&failure.message, aggregate_usage)), + ProviderRunOutcome::Cancelled { .. } => Ok(self + .translator + .finish_provider_run(StopReason::Cancelled, aggregate_usage)), } } } @@ -165,6 +163,7 @@ impl RuntimeResponseTranslator { let mut events = Vec::new(); match event { AgentEvent::TurnStarted { .. } => self.initialize(&mut events), + AgentEvent::KeepAlive => {} AgentEvent::TextDelta { text } => { self.initialize(&mut events); self.add_or_append_text(&text, &mut events); @@ -249,13 +248,31 @@ impl RuntimeResponseTranslator { events } + fn finish_provider_run( + &mut self, + reason: StopReason, + aggregate_usage: &Usage, + ) -> Vec { + let mut events = Vec::new(); + self.initialize(&mut events); + if !self.has_visible_output && reason != StopReason::Cancelled { + if let Some(message) = self.config.empty_output_message.clone() { + self.add_or_append_text(&message, &mut events); + } + } + events.push(self.finished_with_usage(map_stop_reason(reason), aggregate_usage)); + events + } + pub(crate) fn begin_followup_turn(&mut self) { self.text_message_id = None; self.reasoning_message_id = None; + self.usage = Usage::default(); } fn discard_failed_turn_output(&mut self) -> Vec { let mut events = Vec::new(); + self.usage = Usage::default(); if let Some(message_id) = self.text_message_id.take() { events.push(build_replace_text_message( &self.config.task_id, @@ -404,20 +421,27 @@ impl RuntimeResponseTranslator { self.finished_with_reason(map_stop_reason(reason)) } - fn provider_failure(&mut self, message: &str) -> Vec { + fn provider_failure(&mut self, message: &str, aggregate_usage: &Usage) -> Vec { let mut events = Vec::new(); self.initialize(&mut events); - events.push( - self.finished_with_reason(stream_finished::Reason::InternalError( - stream_finished::InternalError { - message: message.to_owned(), - }, - )), - ); + events.push(self.finished_with_usage( + stream_finished::Reason::InternalError(stream_finished::InternalError { + message: message.to_owned(), + }), + aggregate_usage, + )); events } fn finished_with_reason(&self, reason: stream_finished::Reason) -> ResponseEvent { + self.finished_with_usage(reason, &self.usage) + } + + fn finished_with_usage( + &self, + reason: stream_finished::Reason, + aggregate_usage: &Usage, + ) -> ResponseEvent { if !self.config.capabilities.host_managed_history { let (used_tokens, context_size) = self.context_usage.unwrap_or_default(); return build_context_finished( @@ -430,10 +454,16 @@ impl RuntimeResponseTranslator { build_stream_finished( reason, StreamUsage { - input_tokens: saturating_i32(self.usage.input_tokens), - output_tokens: saturating_i32(self.usage.output_tokens), - cache_read_tokens: saturating_i32(self.usage.cached_input_tokens), - cache_write_tokens: saturating_i32(self.usage.cache_creation_input_tokens), + input_tokens: saturating_i32(aggregate_usage.input_tokens), + output_tokens: saturating_i32(aggregate_usage.output_tokens), + cache_read_tokens: saturating_i32(aggregate_usage.cached_input_tokens), + cache_write_tokens: saturating_i32(aggregate_usage.cache_creation_input_tokens), + current_context_tokens: Some(saturating_i32( + self.usage + .input_tokens + .saturating_add(self.usage.cached_input_tokens) + .saturating_add(self.usage.cache_creation_input_tokens), + )), cost_in_cents: 0.0, model_id: self.config.model_id.clone(), max_context_tokens: self.config.max_context_tokens, diff --git a/app/src/ai/runtime/event_translator_tests.rs b/app/src/ai/runtime/event_translator_tests.rs index 5f9cef2d..b3ad938e 100644 --- a/app/src/ai/runtime/event_translator_tests.rs +++ b/app/src/ai/runtime/event_translator_tests.rs @@ -404,6 +404,7 @@ fn provider_retry_clears_failed_attempt_output_before_new_messages() { runtime_id: "runtime".to_owned(), model_id: "model".to_owned(), retry_attempt: 1, + max_retries: 3, elapsed_ms: 2, error: galaxy_agent_core::AgentError::new( galaxy_agent_core::AgentErrorKind::Transport, diff --git a/app/src/ai/runtime/provider_run_coordinator.rs b/app/src/ai/runtime/provider_run_coordinator.rs index 2c751376..4b1d1853 100644 --- a/app/src/ai/runtime/provider_run_coordinator.rs +++ b/app/src/ai/runtime/provider_run_coordinator.rs @@ -67,6 +67,7 @@ pub(crate) enum ProviderRunProjection { runtime_id: String, model_id: String, retry_attempt: u32, + max_retries: u32, elapsed_ms: u64, error: AgentError, }, @@ -639,6 +640,7 @@ impl ProviderRunCoordinator { return Ok(()); } } + AgentEvent::KeepAlive => {} AgentEvent::TextDelta { text } => { if !self .ensure_model_started_acknowledged( @@ -732,14 +734,11 @@ impl ProviderRunCoordinator { return Ok(()); } buffer.usage.clone_from(&usage); - let cumulative_usage = combined_usage(self.run.usage(), &usage); if !self .project_or_fail_acknowledged( ProviderRunProjection::ModelEvent { work_id: call.work_id.clone(), - event: AgentEvent::UsageUpdated { - usage: cumulative_usage, - }, + event: AgentEvent::UsageUpdated { usage }, }, project, ) @@ -890,6 +889,7 @@ impl ProviderRunCoordinator { runtime_id: profile.runtime.descriptor().id.clone(), model_id: profile.request.model.as_str().to_string(), retry_attempt, + max_retries: self.run.max_model_retries_per_turn(), elapsed_ms: elapsed_millis(started_at), error, }, @@ -1018,19 +1018,6 @@ fn request_for_model_call(mut template: TurnRequest, call: &ProviderModelCall) - template } -fn combined_usage(previous: &Usage, current: &Usage) -> Usage { - Usage { - input_tokens: previous.input_tokens.saturating_add(current.input_tokens), - output_tokens: previous.output_tokens.saturating_add(current.output_tokens), - cached_input_tokens: previous - .cached_input_tokens - .saturating_add(current.cached_input_tokens), - cache_creation_input_tokens: previous - .cache_creation_input_tokens - .saturating_add(current.cache_creation_input_tokens), - } -} - fn tool_event_call_id(event: &ToolEvent) -> Result<&str, ProviderRunCoordinatorError> { match event { ToolEvent::Proposed { call } => Ok(&call.id), diff --git a/app/src/ai/runtime/provider_run_coordinator_tests.rs b/app/src/ai/runtime/provider_run_coordinator_tests.rs index 6992a587..8def1ec3 100644 --- a/app/src/ai/runtime/provider_run_coordinator_tests.rs +++ b/app/src/ai/runtime/provider_run_coordinator_tests.rs @@ -363,7 +363,7 @@ async fn one_run_drives_model_tool_and_followup_turns_with_atomic_history() { ProviderRunProjection::ModelEvent { event: AgentEvent::UsageUpdated { usage }, .. - } if usage.input_tokens == 30 && usage.output_tokens == 9 + } if usage.input_tokens == 20 && usage.output_tokens == 5 ))); coordinator.run_mut().complete(&work_id).unwrap(); @@ -761,8 +761,9 @@ async fn recoverable_start_failure_retries_the_same_work_identity() { ProviderRunProjection::ModelRetry { work_id, retry_attempt, + max_retries, .. - } => Some((work_id.clone(), *retry_attempt)), + } => Some((work_id.clone(), *retry_attempt, *max_retries)), ProviderRunProjection::ModelTurnRequested { .. } | ProviderRunProjection::ModelTurnStarted { .. } | ProviderRunProjection::ModelTurnFinished { .. } @@ -788,6 +789,7 @@ async fn recoverable_start_failure_retries_the_same_work_identity() { .expect("retried model start"); assert_eq!(retry.0, started); assert_eq!(retry.1, 1); + assert_eq!(retry.2, 3); } #[tokio::test] @@ -969,6 +971,38 @@ async fn model_event_idle_timeout_retries_the_same_work_identity() { assert_single_retry_lifecycle(&projections, "event timed out", true); } +#[tokio::test] +async fn model_keepalive_preserves_the_active_turn_without_rendering_output() { + let runtime = Arc::new(ScriptedRuntime::new(vec![Ok(vec![ + started("request-keepalive"), + Ok(AgentEvent::KeepAlive), + Ok(AgentEvent::TextDelta { + text: "finished".to_string(), + }), + stopped(StopReason::Completed), + ])])); + let mut coordinator = coordinator(runtime); + let mut projections = Vec::new(); + let (_sender, control) = turn_control(); + + let block = coordinator + .drive_until_blocked(control, collect_projection(&mut projections)) + .await + .unwrap(); + + assert!(matches!(block, ProviderRunBlock::AwaitingDriver { .. })); + assert_eq!(coordinator.run().model_retries(), 0); + assert!(!projections.iter().any(|projection| { + matches!( + projection, + ProviderRunProjection::ModelEvent { + event: AgentEvent::KeepAlive, + .. + } + ) + })); +} + #[tokio::test] async fn persistent_checkpoint_failure_terminates_without_redrive() { let runtime = Arc::new(ScriptedRuntime::new(vec![answer_turn()])); @@ -1168,7 +1202,11 @@ async fn transcript_projector_emits_one_ui_stream_for_the_whole_run() { else { panic!("expected terminal run"); }; - ui_events.extend(projector.finish(&outcome).unwrap()); + ui_events.extend( + projector + .finish(&outcome, coordinator.run().usage()) + .unwrap(), + ); assert_eq!( count_response_events(&ui_events, ResponseEventKind::Finished), 1 @@ -1188,11 +1226,14 @@ fn transcript_projector_preserves_provider_failure_message() { empty_output_message: None, }); let events = projector - .finish(&ProviderRunOutcome::Failed(ProviderRunFailure { - kind: ProviderRunFailureKind::ModelCall, - message: "upstream provider rejected the request".to_string(), - source: None, - })) + .finish( + &ProviderRunOutcome::Failed(ProviderRunFailure { + kind: ProviderRunFailureKind::ModelCall, + message: "upstream provider rejected the request".to_string(), + source: None, + }), + &Usage::default(), + ) .unwrap(); let finished = events diff --git a/app/src/ai/runtime/rig_request.rs b/app/src/ai/runtime/rig_request.rs index a0fb926e..4240699a 100644 --- a/app/src/ai/runtime/rig_request.rs +++ b/app/src/ai/runtime/rig_request.rs @@ -651,6 +651,7 @@ fn build_system_prompt( ); let contexts = inputs.iter().filter_map(AIAgentInput::context).flatten(); let mut environment = Vec::new(); + let mut request_time = None; let mut project_rules = Vec::new(); let mut available_skills = Vec::new(); let mut attached_context = Vec::new(); @@ -724,7 +725,7 @@ fn build_system_prompt( attached_context.push(("Selected text".to_string(), text.clone())); } AIAgentContext::CurrentTime { current_time } => { - environment.push(format!("Current time: {current_time}")); + request_time = Some(*current_time); } AIAgentContext::Codebase { path, name } => { environment.push(format!("Indexed codebase: {name} ({path})")); @@ -851,6 +852,14 @@ fn build_system_prompt( ); } } + // Keep volatile request data at the end of the system prompt. Provider + // prompt caches match the longest exact prefix, so putting the current + // timestamp ahead of rules and tool instructions invalidates that stable + // prefix on every model call. + if let Some(request_time) = request_time { + prompt.push_str("\n## Request Time\n"); + prompt.push_str(&format!("- Current time: {request_time}\n")); + } prompt } diff --git a/app/src/ai/runtime/rig_request_tests.rs b/app/src/ai/runtime/rig_request_tests.rs index 0a70c4b4..2355938e 100644 --- a/app/src/ai/runtime/rig_request_tests.rs +++ b/app/src/ai/runtime/rig_request_tests.rs @@ -81,6 +81,32 @@ fn native_context_reaches_rig_without_a_proto_context_conversion() { assert!(prompt.contains("Indexed codebase: galaxy (/repo)")); } +#[test] +fn volatile_request_time_follows_the_cacheable_system_prompt_prefix() { + let mut params = RequestParams::new_for_test(); + params.global_rules = vec![( + "Stable rule".to_string(), + "Preserve this cacheable instruction.".to_string(), + )]; + params.input = vec![user_query_with_context( + "Inspect the cache layout", + vec![AIAgentContext::CurrentTime { + current_time: chrono::Local::now(), + }], + )]; + + let prepared = prepare_rig_turn(&config(), params, vec![ToolType::ReadFiles], Vec::new()); + let prompt = prepared.request.system_prompt.expect("system prompt"); + let rule_position = prompt + .find("Preserve this cacheable instruction.") + .expect("global rule"); + let tools_position = prompt.find("## Available Tools").expect("tool contract"); + let time_position = prompt.find("## Request Time").expect("request time"); + + assert!(rule_position < time_position); + assert!(tools_position < time_position); +} + #[test] fn builds_a_rig_turn_directly_from_galaxy_request_state() { let mut params = RequestParams::new_for_test(); diff --git a/app/src/terminal/input/agent.rs b/app/src/terminal/input/agent.rs index 20d26a9b..ce338195 100644 --- a/app/src/terminal/input/agent.rs +++ b/app/src/terminal/input/agent.rs @@ -5,11 +5,10 @@ use galaxyui::elements::{ Align, AnchorPair, Border, ConstrainedBox, Container, CornerRadius, CrossAxisAlignment, DispatchEventResult, DropTarget, Element, Empty, EventHandler, Expanded, Flex, Hoverable, MainAxisSize, OffsetPositioning, OffsetType, ParentElement, PositionedElementOffsetBounds, - PositioningAxis, Radius, SavePosition, Stack, Text, XAxisAnchor, YAxisAnchor, + PositioningAxis, Radius, SavePosition, Stack, XAxisAnchor, YAxisAnchor, }; use galaxyui::presenter::ChildView; -use galaxyui::{AppContext, EntityId, SingletonEntity as _}; -use pathfinder_color::ColorU; +use galaxyui::{AppContext, SingletonEntity as _}; use super::common::{ add_command_xray_overlay, add_input_suggestions_overlays, add_voltron_overlay, @@ -22,10 +21,7 @@ use crate::ai::blocklist::agent_view::shortcuts::{ }; use crate::ai::blocklist::agent_view::{agent_view_bg_fill, AgentViewState}; use crate::ai::blocklist::InputType; -use crate::ai::execution_profiles::profiles::AIExecutionProfilesModel; -use crate::ai::execution_profiles::AIExecutionProfileAppExt; use crate::ai::harness_availability::HarnessAvailabilityModel; -use crate::ai::llms::LLMPreferences; use crate::appearance::Appearance; use crate::context_chips::spacing::{self}; use crate::editor::position_id_for_cursor; @@ -164,19 +160,6 @@ impl Input { .finish(), ); - if let Some(conv_id) = self - .agent_view_controller - .as_ref(app) - .agent_view_state() - .active_conversation_id() - { - if let Some(status_bar) = - render_session_status_bar(appearance, app, self.terminal_view_id, conv_id) - { - column.add_child(status_bar); - } - } - stack.add_child(wrap_input_with_terminal_padding_and_focus_handler( self.is_active_session(app), column.finish(), @@ -735,156 +718,6 @@ impl Input { } } -fn format_token_count(tokens: u32) -> String { - if tokens >= 1_000_000 { - format!("{:.1}M", tokens as f64 / 1_000_000.0) - } else if tokens >= 1_000 { - format!("{:.1}k", tokens as f64 / 1_000.0) - } else { - format!("{tokens}") - } -} - -fn cache_hit_color(pct: f64, theme: &galaxy_core::ui::theme::GalaxyTheme) -> ColorU { - if pct >= 90.0 { - theme.ansi_fg_green() - } else if pct >= 50.0 { - theme.ansi_fg_yellow() - } else { - theme.ansi_fg_red() - } -} - -fn render_session_status_bar( - appearance: &Appearance, - app: &AppContext, - terminal_view_id: EntityId, - conversation_id: crate::ai::agent::conversation::AIConversationId, -) -> Option> { - let (cache_read, cache_write, cache_miss, cost_cents, context_usage, current_context) = - if let Some(conversation) = - BlocklistAIHistoryModel::as_ref(app).conversation(&conversation_id) - { - ( - conversation.last_block_cache_read_tokens(), - conversation.last_block_cache_write_tokens(), - conversation.last_block_cache_miss_tokens(), - conversation.total_cost_cents(), - conversation.context_window_usage(), - conversation.current_context_tokens(), - ) - } else { - (0, 0, 0, 0.0, 0.0, 0) - }; - - let cache_total_ops = cache_read + cache_write + cache_miss; - let cache_hit_pct = if cache_total_ops > 0 { - (cache_read as f64 / cache_total_ops as f64) * 100.0 - } else { - 0.0 - }; - - let active_model = - LLMPreferences::as_ref(app).get_active_base_model(app, Some(terminal_view_id)); - let profile_context = AIExecutionProfilesModel::as_ref(app) - .active_profile(Some(terminal_view_id), app) - .data() - .context_window_display_value(app); - let model_max_context = active_model - .context_window - .default_max - .max(active_model.context_window.max); - let max_context = profile_context - .or((model_max_context > 0).then_some(model_max_context)) - .unwrap_or(200_000); - let context_pct = context_usage.clamp(0.0, 1.0) * 100.0; - - let theme = appearance.theme(); - let font_family = appearance.ui_font_family(); - let font_size = appearance.monospace_font_size() - 1.0; - let dim_color: ColorU = theme.sub_text_color(theme.background()).into(); - let cache_color = cache_hit_color(cache_hit_pct, theme); - - let mut row = Flex::row() - .with_cross_axis_alignment(CrossAxisAlignment::Center) - .with_main_axis_size(MainAxisSize::Min); - - // Context: XX.X% (Xk / Xk) - let context_text = format!( - "\u{25a0} Ctx: {:.1}% ({}/{})", - context_pct, - format_token_count(current_context), - format_token_count(max_context), - ); - row.add_child( - Text::new_inline(context_text, font_family, font_size) - .with_color(dim_color) - .finish(), - ); - - // Only show cache stats when the provider actually reports them - // (Bedrock reports cache data; LiteLLM/OpenAI typically does not) - if cache_read > 0 || cache_write > 0 { - // Separator - row.add_child( - Container::new( - Text::new_inline(" \u{2502} ".to_string(), font_family, font_size) - .with_color(dim_color) - .finish(), - ) - .finish(), - ); - - // Cache Hit: XX.X% (R: Xk W: Xk M: Xk) - let cache_label = format!("\u{25c6} Cache: {:.1}%", cache_hit_pct); - row.add_child( - Text::new_inline(cache_label, font_family, font_size) - .with_color(cache_color) - .finish(), - ); - - let cache_detail = format!( - " (R:{} W:{} M:{})", - format_token_count(cache_read), - format_token_count(cache_write), - format_token_count(cache_miss), - ); - row.add_child( - Text::new_inline(cache_detail, font_family, font_size) - .with_color(dim_color) - .finish(), - ); - } - - // Separator - row.add_child( - Container::new( - Text::new_inline(" \u{2502} ".to_string(), font_family, font_size) - .with_color(dim_color) - .finish(), - ) - .finish(), - ); - - // Cost: $X.XX - let cost_text = format!("\u{25b2} ${:.2}", cost_cents / 100.0); - row.add_child( - Text::new_inline(cost_text, font_family, font_size) - .with_color(theme.ansi_fg_green()) - .finish(), - ); - - Some( - Container::new(row.finish()) - .with_padding_left(12.) - .with_padding_right(12.) - .with_padding_top(2.) - .with_padding_bottom(2.) - .with_background(theme.background()) - .finish(), - ) -} - pub mod styles { use galaxy_core::ui::theme::GalaxyTheme; use pathfinder_color::ColorU; diff --git a/crates/galaxy_agent_core/src/provider_run.rs b/crates/galaxy_agent_core/src/provider_run.rs index 207acbf4..5d0b0363 100644 --- a/crates/galaxy_agent_core/src/provider_run.rs +++ b/crates/galaxy_agent_core/src/provider_run.rs @@ -502,6 +502,10 @@ impl ProviderRun { self.model_retries } + pub fn max_model_retries_per_turn(&self) -> u32 { + self.limits.max_model_retries_per_turn + } + pub fn is_terminal(&self) -> bool { matches!( self.state, @@ -665,20 +669,20 @@ impl ProviderRun { } ProviderRunState::Failed { failure } => match failure.kind { ProviderRunFailureKind::ModelCall - if !failure + if failure .source .as_ref() - .is_some_and(|source| !source.recoverable) => + .is_none_or(|source| source.recoverable) => { return Err(invalid( "model-call failure lacks a non-recoverable source".to_string(), )); } ProviderRunFailureKind::RetryLimitExceeded - if !failure + if failure .source .as_ref() - .is_some_and(|source| source.recoverable) => + .is_none_or(|source| !source.recoverable) => { return Err(invalid( "retry-limit failure lacks a recoverable source".to_string(), diff --git a/crates/galaxy_agent_core/src/types.rs b/crates/galaxy_agent_core/src/types.rs index 19d3c9bc..2dc8dc12 100644 --- a/crates/galaxy_agent_core/src/types.rs +++ b/crates/galaxy_agent_core/src/types.rs @@ -223,15 +223,21 @@ impl ToolEvent { #[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] pub struct Usage { + /// Input tokens not served by or written to a provider prompt cache. pub input_tokens: u64, pub output_tokens: u64, + /// Input tokens served from the provider prompt cache. pub cached_input_tokens: u64, + /// Input tokens written to the provider prompt cache. pub cache_creation_input_tokens: u64, } impl Usage { pub fn total_tokens(&self) -> u64 { - self.input_tokens.saturating_add(self.output_tokens) + self.input_tokens + .saturating_add(self.cached_input_tokens) + .saturating_add(self.cache_creation_input_tokens) + .saturating_add(self.output_tokens) } } @@ -268,6 +274,9 @@ pub enum AgentEvent { TurnStarted { runtime_request_id: String, }, + /// A transport heartbeat proving that the current model stream is still connected. + /// Consumers should use this to refresh idle timeouts without rendering output. + KeepAlive, TextDelta { text: String, }, diff --git a/crates/galaxy_agent_core/src/types_tests.rs b/crates/galaxy_agent_core/src/types_tests.rs index 1ea3749d..6984c060 100644 --- a/crates/galaxy_agent_core/src/types_tests.rs +++ b/crates/galaxy_agent_core/src/types_tests.rs @@ -26,9 +26,9 @@ fn truncates_large_tool_results_for_provider_request() { } #[test] -fn usage_total_excludes_cached_breakdown_to_avoid_double_counting() { +fn usage_total_combines_disjoint_cache_and_miss_buckets() { let usage = Usage { - input_tokens: 100, + input_tokens: 10, output_tokens: 25, cached_input_tokens: 80, cache_creation_input_tokens: 10, diff --git a/crates/galaxy_agent_rig/src/bedrock_tests.rs b/crates/galaxy_agent_rig/src/bedrock_tests.rs index 7c314918..6982a487 100644 --- a/crates/galaxy_agent_rig/src/bedrock_tests.rs +++ b/crates/galaxy_agent_rig/src/bedrock_tests.rs @@ -83,7 +83,7 @@ fn normalizes_bedrock_usage_and_max_token_stop() { assert_eq!( map_usage((&response).into()), Usage { - input_tokens: 100, + input_tokens: 50, output_tokens: 25, cached_input_tokens: 40, cache_creation_input_tokens: 10, diff --git a/crates/galaxy_agent_rig/src/chatgpt.rs b/crates/galaxy_agent_rig/src/chatgpt.rs index fb1e7b89..98986c4b 100644 --- a/crates/galaxy_agent_rig/src/chatgpt.rs +++ b/crates/galaxy_agent_rig/src/chatgpt.rs @@ -116,6 +116,7 @@ impl ChatGPTSubscriptionRuntime { AgentEvent::ReasoningDelta { .. } | AgentEvent::ReasoningCompleted { .. } | AgentEvent::TurnStarted { .. } + | AgentEvent::KeepAlive | AgentEvent::UsageUpdated { .. } | AgentEvent::RuntimeActivityUpdated { .. } | AgentEvent::ContextUsageUpdated { .. } diff --git a/crates/galaxy_agent_rig/src/openai_compatible_tests.rs b/crates/galaxy_agent_rig/src/openai_compatible_tests.rs index 520caf20..795d632c 100644 --- a/crates/galaxy_agent_rig/src/openai_compatible_tests.rs +++ b/crates/galaxy_agent_rig/src/openai_compatible_tests.rs @@ -73,7 +73,7 @@ async fn rig_stream_maps_reasoning_text_usage_and_stop() { }, AgentEvent::UsageUpdated { usage: Usage { - input_tokens: 4, + input_tokens: 2, output_tokens: 6, cached_input_tokens: 2, cache_creation_input_tokens: 0, diff --git a/crates/galaxy_agent_rig/src/stream.rs b/crates/galaxy_agent_rig/src/stream.rs index eeb378e2..1719151c 100644 --- a/crates/galaxy_agent_rig/src/stream.rs +++ b/crates/galaxy_agent_rig/src/stream.rs @@ -125,11 +125,15 @@ where }); } Ok(StreamedAssistantContent::Unknown(value)) => { - yield Err(AgentError::new( - AgentErrorKind::Protocol, - format!("Rig returned an unsupported provider event: {value}"), - )); - return; + if is_keepalive_event(&value) { + yield Ok(AgentEvent::KeepAlive); + } else { + yield Err(AgentError::new( + AgentErrorKind::Protocol, + format!("Rig returned an unsupported provider event: {value}"), + )); + return; + } } Err(error) => { if let Some(reason) = completion_error_stop_reason(&error) { @@ -255,8 +259,24 @@ fn stopped_with_reason(runtime_request_id: String, reason: StopReason) -> AgentE } pub(crate) fn map_usage(usage: rig_core::completion::Usage) -> Usage { + // Rig preserves provider-native input semantics: OpenAI includes cached + // tokens in `input_tokens`, while Anthropic reports cache reads/writes + // separately. `total_tokens - output_tokens` gives the normalized prompt + // size for both shapes, so store only the uncached portion in + // `input_tokens` and keep cache reads/writes disjoint. + let total_input_tokens = if usage.total_tokens > 0 && usage.total_tokens >= usage.output_tokens + { + usage.total_tokens - usage.output_tokens + } else { + usage + .input_tokens + .saturating_add(usage.cached_input_tokens) + .saturating_add(usage.cache_creation_input_tokens) + }; Usage { - input_tokens: usage.input_tokens, + input_tokens: total_input_tokens + .saturating_sub(usage.cached_input_tokens) + .saturating_sub(usage.cache_creation_input_tokens), output_tokens: usage.output_tokens, cached_input_tokens: usage.cached_input_tokens, cache_creation_input_tokens: usage.cache_creation_input_tokens, @@ -307,6 +327,10 @@ fn json_value_indicates_context_window_exceeded(value: &serde_json::Value) -> bo } } +fn is_keepalive_event(value: &serde_json::Value) -> bool { + value.get("type").and_then(serde_json::Value::as_str) == Some("keepalive") +} + fn text_indicates_context_window_exceeded(text: &str) -> bool { let normalized = text.to_ascii_lowercase(); normalized.contains("modelcontextwindowexceeded") @@ -352,7 +376,7 @@ fn map_completion_error(error: CompletionError) -> AgentError { mapped.recoverable = matches!( kind, AgentErrorKind::RateLimited | AgentErrorKind::Transport - ); + ) || status.is_some_and(|status| (500..=599).contains(&status)); mapped } @@ -361,7 +385,9 @@ mod tests { use galaxy_agent_core::{AgentErrorKind, StopReason}; use rig_core::completion::CompletionError; - use super::{completion_error_stop_reason, domain_tool_call, map_completion_error}; + use super::{ + completion_error_stop_reason, domain_tool_call, is_keepalive_event, map_completion_error, + }; #[test] fn domain_tool_call_prefers_responses_call_id() { @@ -428,4 +454,32 @@ mod tests { Some(StopReason::ContextWindowExceeded) ); } + + #[test] + fn provider_keepalive_event_is_a_transport_heartbeat() { + assert!(is_keepalive_event(&serde_json::json!({ + "type": "keepalive", + "sequence_number": 3, + }))); + assert!(!is_keepalive_event(&serde_json::json!({ + "type": "unsupported", + }))); + } + + #[test] + fn provider_server_error_is_recoverable() { + let status = rig_core::http_client::Response::builder() + .status(503) + .body(()) + .unwrap() + .status(); + let error = CompletionError::from_http_response( + status, + r#"{"error":{"message":"Service temporarily unavailable"}}"#, + ); + + let mapped = map_completion_error(error); + assert_eq!(mapped.kind, AgentErrorKind::Provider); + assert!(mapped.recoverable); + } } diff --git a/crates/galaxy_core/src/paths.rs b/crates/galaxy_core/src/paths.rs index a8d0a37b..57da7b60 100644 --- a/crates/galaxy_core/src/paths.rs +++ b/crates/galaxy_core/src/paths.rs @@ -389,9 +389,7 @@ fn migrate_directory_contents(source_dir: &Path, target_dir: &Path) { } } - if std::fs::read_dir(source_dir) - .is_ok_and(|mut entries| entries.next().is_none()) - { + if std::fs::read_dir(source_dir).is_ok_and(|mut entries| entries.next().is_none()) { let _ = std::fs::remove_dir(source_dir); } }