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

Improve provider reliability and usage visibility

Browse Source
This commit is contained in:
rkw6086
2026-08-21 19:12:14 -05:00
parent 19b2c5f687
commit be1dbb600a
28 changed files with 1070 additions and 592 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/src/ai/agent/conversation.rs
+39 -12
View File
@@ -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<String>,
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()
}
app/src/ai/bedrock/response_translator.rs
+3 -3
View File
@@ -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(),
app/src/ai/bedrock/runtime.rs
+32 -9
View File
@@ -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<String>) -> AgentError {
AgentError::new(AgentErrorKind::Protocol, message)
}
app/src/ai/bedrock/runtime_tests.rs
+14 -2
View File
@@ -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);
}
app/src/ai/blocklist/action_model.rs
+15 -17
View File
@@ -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 {
app/src/ai/blocklist/agent_view/agent_input_footer/mod.rs
+146 -1
View File
@@ -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<ActionButton>,
file_button: ViewHandle<ActionButton>,
context_window_button: ViewHandle<ActionButton>,
llm_context_usage_button: ViewHandle<ActionButton>,
llm_cache_details_button: ViewHandle<ActionButton>,
model_selector: ViewHandle<ProfileModelSelector>,
environment_selector: Option<ViewHandle<EnvironmentSelector>>,
handoff_environment_selector: ViewHandle<EnvironmentSelector>,
@@ -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<Self>) {
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;
app/src/ai/blocklist/agent_view/agent_input_footer/toolbar_item.rs
+17 -2
View File
@@ -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<Self> {
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);
app/src/ai/blocklist/agent_view/orchestration_conversation_links.rs
+1 -5
View File
@@ -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))
}
app/src/ai/blocklist/block/status_bar.rs
+6
View File
@@ -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,
app/src/ai/blocklist/block/view_impl/common.rs
+8 -8
View File
@@ -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<String>,
pub secondary_element: Option<Box<dyn Element>>,
/// When an LRC subagent has sent at least one snapshot, the timestamp of the most recent snapshot.
pub last_snapshot_at: Option<instant::Instant>,
@@ -337,14 +338,6 @@ pub fn render_warping_indicator<V: View>(
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<V: View>(
}
};
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);
app/src/ai/blocklist/controller.rs
+521 -272
View File
File diff suppressed because it is too large Load Diff
app/src/ai/blocklist/controller_tests.rs
+21 -12
View File
@@ -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),
app/src/ai/openai/response_translator.rs
+14 -9
View File
@@ -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<i32>,
pub(crate) cost_in_cents: f32,
pub(crate) model_id: String,
pub(crate) max_context_tokens: Option<u32>,
@@ -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(),
app/src/ai/runtime/event_translator.rs
+55 -25
View File
@@ -40,6 +40,7 @@ pub(crate) struct RuntimeResponseTranslator {
activity_message_ids: HashMap<String, String>,
activities: HashMap<String, RuntimeActivity>,
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<Vec<ResponseEvent>, 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<ResponseEvent> {
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<ResponseEvent> {
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<ResponseEvent> {
fn provider_failure(&mut self, message: &str, aggregate_usage: &Usage) -> Vec<ResponseEvent> {
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,
app/src/ai/runtime/event_translator_tests.rs
+1
View File
@@ -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,
app/src/ai/runtime/provider_run_coordinator.rs
+4 -17
View File
@@ -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),
app/src/ai/runtime/provider_run_coordinator_tests.rs
+49 -8
View File
@@ -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
app/src/ai/runtime/rig_request.rs
+10 -1
View File
@@ -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
}
app/src/ai/runtime/rig_request_tests.rs
+26
View File
@@ -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();