From f4dc6231d6c3c784ae107436e4b3ab3cb5df60c2 Mon Sep 17 00:00:00 2001 From: Ryan Ward Date: Sun, 23 Aug 2026 22:22:32 -0500 Subject: [PATCH] Improve agent retries and tool progress --- AGENTS.md | 2 +- app/src/ai/agent/conversation.rs | 23 +- app/src/ai/blocklist/block/status_bar.rs | 21 +- app/src/ai/blocklist/block/view_impl.rs | 9 + .../ai/blocklist/block/view_impl/common.rs | 8 - .../ai/blocklist/block/view_impl/output.rs | 24 + app/src/ai/blocklist/controller.rs | 648 ++++++++++++++++-- app/src/ai/blocklist/controller_tests.rs | 86 +-- .../inline_action/requested_command.rs | 63 +- .../ai/runtime/provider_run_coordinator.rs | 3 +- crates/galaxy_agent_core/src/provider_run.rs | 78 +++ .../src/provider_run_tests.rs | 54 ++ crates/galaxy_agent_rig/src/stream.rs | 19 + crates/galaxy_agent_rig/src/stream_tests.rs | 16 + 14 files changed, 913 insertions(+), 141 deletions(-) create mode 100644 crates/galaxy_agent_rig/src/stream_tests.rs diff --git a/AGENTS.md b/AGENTS.md index 2b92b3bb..00315f7f 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -154,7 +154,7 @@ Key invariants: - Before progressive summarization drains messages, `ConversationMessage::archive_tool_results()` moves tool-use/result pairs into `tool_result_archive` - Bedrock prompt caching uses three cache points: system prompt, second-to-last history message, and tool configuration - `ensure_tool_results_paired()` enforces Bedrock's invariant that every `tool_use` has a matching `tool_result` -- Progressive summaries are prepended to provider requests as a user/assistant pair; background summarization remains independent of the active provider run +- Progressive summaries are prepended to provider requests as a user/assistant pair; direct-provider runs compact their live transcript at a checkpointed `ReadyToCallModel` boundary before another model call, while background summarization remains independent for session-owned runtimes - Loop prevention in `controller.rs` detects repeated tool failures (3+ identical) and injects a corrective instruction - Direct-provider long-running shell follow-ups retain CLI tasks as UI projections while command monitoring and completion stay owned by the same provider run - A direct-provider command completion is only queued when the terminal reports it; the CLI task remains active until the provider run applies that completion at a safe boundary and deactivates it diff --git a/app/src/ai/agent/conversation.rs b/app/src/ai/agent/conversation.rs index 38a96ef3..438d43e2 100644 --- a/app/src/ai/agent/conversation.rs +++ b/app/src/ai/agent/conversation.rs @@ -2490,16 +2490,19 @@ 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 = 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() - }); + let live_input = match usage_metadata.as_ref() { + // When a failed follow-up has no usage of its own, the provider emits + // zero here. Preserve the most recent live occupancy rather than + // falling back to the run's cumulative billing totals. + Some(metadata) => metadata.total_input_tokens, + None if was_user_initiated_request => token_usage + .iter() + .map(|u| u.total_input + u.input_cache_read + u.input_cache_write) + .sum(), + // Auxiliary requests such as progressive summarization count toward + // cumulative /usage totals but do not become the agent's live context. + None => 0, + }; if live_input > 0 { self.current_context_tokens = live_input; } diff --git a/app/src/ai/blocklist/block/status_bar.rs b/app/src/ai/blocklist/block/status_bar.rs index 2bdb7e87..d108074e 100644 --- a/app/src/ai/blocklist/block/status_bar.rs +++ b/app/src/ai/blocklist/block/status_bar.rs @@ -125,7 +125,6 @@ pub struct BlocklistAIStatusBar { /// The random loading message chosen for the current conversation run, stable across /// re-renders and tool-call follow-up exchanges. warping_message: Option<&'static str>, - warping_message_ticks: u8, /// Handle for the periodic timer that updates the warping elapsed timer UI. warping_timer_handle: Option, @@ -435,7 +434,6 @@ impl BlocklistAIStatusBar { last_read_refresh_handle: None, warping_start_time: None, warping_message: None, - warping_message_ticks: 0, warping_timer_handle: None, ambient_agent_view_model, current_tip: None, @@ -805,7 +803,6 @@ impl BlocklistAIStatusBar { self.warping_start_time.get_or_insert_with(Instant::now); self.warping_message .get_or_insert_with(random_load_output_message); - self.warping_message_ticks = 0; // Don't start a new timer if one is already running if self.warping_timer_handle.is_some() { return; @@ -817,11 +814,6 @@ impl BlocklistAIStatusBar { |me, _, ctx| { me.warping_timer_handle = None; if me.warping_start_time.is_some() { - me.warping_message_ticks = me.warping_message_ticks.saturating_add(1); - if me.warping_message_ticks >= 5 { - me.warping_message_ticks = 0; - me.warping_message = Some(random_load_output_message()); - } ctx.notify(); me.restart_warping_timer(ctx); } @@ -842,11 +834,6 @@ impl BlocklistAIStatusBar { |me, _, ctx| { me.warping_timer_handle = None; if me.warping_start_time.is_some() { - me.warping_message_ticks = me.warping_message_ticks.saturating_add(1); - if me.warping_message_ticks >= 5 { - me.warping_message_ticks = 0; - me.warping_message = Some(random_load_output_message()); - } ctx.notify(); me.restart_warping_timer(ctx); } @@ -859,7 +846,6 @@ impl BlocklistAIStatusBar { fn stop_warping_timer(&mut self) { self.warping_start_time = None; self.warping_message = None; - self.warping_message_ticks = 0; if let Some(handle) = self.warping_timer_handle.take() { handle.abort(); } @@ -1011,17 +997,13 @@ impl BlocklistAIStatusBar { let default_warping_text = fallback_warping_text .as_deref() .or(self.warping_message) - .unwrap_or_else(|| random_load_output_message()) + .unwrap_or("Exploring the Galaxy...") .to_owned(); let retry_status_text = self .controller .as_ref(app) .provider_retry_status(conversation.id()) .map(|status| status.label()); - let model_progress_text = self - .controller - .as_ref(app) - .provider_tool_call_progress_label(conversation.id()); let secondary_element = if fallback_warping_text.is_some() { Some(render_fallback_explanation(model.as_ref(), app)) } else { @@ -1081,7 +1063,6 @@ impl BlocklistAIStatusBar { force_refresh_button, default_warping_text, retry_status_text, - model_progress_text, secondary_element, last_snapshot_at, warping_start_time: self.warping_start_time, diff --git a/app/src/ai/blocklist/block/view_impl.rs b/app/src/ai/blocklist/block/view_impl.rs index 6f2fa430..04f937cc 100644 --- a/app/src/ai/blocklist/block/view_impl.rs +++ b/app/src/ai/blocklist/block/view_impl.rs @@ -1077,6 +1077,14 @@ impl View for AIBlock { .as_ref() .is_some_and(|model| model.as_ref(app).is_ambient_agent()); + let streaming_tool_calls = if self.model.is_latest_visible_exchange_in_root_task(app) { + self.controller + .as_ref(app) + .provider_tool_call_progress(self.client_ids.conversation_id) + } else { + Vec::new() + }; + contents.add_child(output::render( output::Props { conversation_id: self.client_ids.conversation_id, @@ -1150,6 +1158,7 @@ impl View for AIBlock { && self.has_imported_comments_in_current_thread(app), ask_user_question_view: self.ask_user_question_view.as_ref(), is_cloud_agent_pre_first_exchange, + streaming_tool_calls: &streaming_tool_calls, }, app, )); diff --git a/app/src/ai/blocklist/block/view_impl/common.rs b/app/src/ai/blocklist/block/view_impl/common.rs index 37f4e4ec..20038973 100644 --- a/app/src/ai/blocklist/block/view_impl/common.rs +++ b/app/src/ai/blocklist/block/view_impl/common.rs @@ -219,7 +219,6 @@ pub struct WarpingProps<'a, V> { pub terminal_model: &'a TerminalModel, pub default_warping_text: String, pub retry_status_text: Option, - pub model_progress_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, @@ -480,13 +479,6 @@ pub fn render_warping_indicator( } }; - if let Some(model_progress_text) = props.model_progress_text.as_deref() { - non_shimmering_text = Some(match non_shimmering_text { - Some(text) if !text.is_empty() => format!("{text} · {model_progress_text}"), - Some(_) | None => format!(" · {model_progress_text}"), - }); - } - 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}"), diff --git a/app/src/ai/blocklist/block/view_impl/output.rs b/app/src/ai/blocklist/block/view_impl/output.rs index a900ec8d..0a1e775d 100644 --- a/app/src/ai/blocklist/block/view_impl/output.rs +++ b/app/src/ai/blocklist/block/view_impl/output.rs @@ -81,6 +81,7 @@ use crate::ai::blocklist::block::{ CollapsibleElementState, CollapsibleExpansionState, EmbeddedCodeEditorView, FinishReason, ImportedCommentGroup, RequestedEdit, TextLocation, TodoListElementState, }; +use crate::ai::blocklist::controller::ProviderToolCallProgressStatus; use crate::ai::blocklist::history_model::BlocklistAIHistoryModel; use crate::ai::blocklist::inline_action::ask_user_question_view::AskUserQuestionView; use crate::ai::blocklist::inline_action::aws_bedrock_credentials_error::AwsBedrockCredentialsErrorView; @@ -212,6 +213,9 @@ pub(crate) struct Props<'a> { /// `true` when this block belongs to a cloud agent pane that is still in its setup phase /// (running environment startup commands before the first agent turn). pub(super) is_cloud_agent_pre_first_exchange: bool, + /// Tool calls whose arguments are still arriving from the model. These are + /// rendered as provisional panes until the complete calls replace them. + pub(super) streaming_tool_calls: &'a [ProviderToolCallProgressStatus], } pub(super) fn render(props: Props, app: &AppContext) -> Box { @@ -1211,6 +1215,10 @@ pub(super) fn render(props: Props, app: &AppContext) -> Box { } } + for tool_call in props.streaming_tool_calls { + output_items.add_child(render_streaming_tool_call(tool_call, app)); + } + if request_type.is_active() { if let AIBlockOutputStatus::Failed { error, .. } = &status { // While an automatic resume is still in flight, keep the failed exchange @@ -1287,6 +1295,22 @@ pub(super) fn render(props: Props, app: &AppContext) -> Box { output_items.finish() } +fn render_streaming_tool_call( + progress: &ProviderToolCallProgressStatus, + app: &AppContext, +) -> Box { + let appearance = Appearance::as_ref(app); + let loading_icon = Icon::Loading.to_galaxyui_icon(galaxy_core::ui::theme::Fill::Solid( + internal_colors::neutral_6(appearance.theme()), + )); + let header = HeaderConfig::new(progress.label(), app) + .with_icon(loading_icon) + .with_corner_radius_override(CornerRadius::with_all(Radius::Pixels(8.))) + .render(app); + + render_tool_pane_shell(Clipped::new(header).finish(), false, false, false, app) +} + fn render_runtime_activity( output_message: &AIAgentOutputMessage, activity: &RuntimeActivity, diff --git a/app/src/ai/blocklist/controller.rs b/app/src/ai/blocklist/controller.rs index ce43a1cc..e247bacf 100644 --- a/app/src/ai/blocklist/controller.rs +++ b/app/src/ai/blocklist/controller.rs @@ -116,6 +116,13 @@ 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); +const PROGRESSIVE_SUMMARY_PROMPT: &str = "Summarize the following conversation history. Preserve:\n\ + - All decisions made and their rationale\n\ + - All file paths modified and what was changed\n\ + - All tool calls with their significant results (commands run, files read, errors encountered)\n\ + - Current task state and any pending work\n\ + - Technical details, code patterns, and architecture discussed\n\n\ + Be comprehensive. This summary will be the only record of these exchanges."; #[derive(Clone)] struct ProgressiveSummaryCandidate { @@ -124,6 +131,21 @@ struct ProgressiveSummaryCandidate { provider_config: ProviderConfig, } +struct ActiveProgressiveSummaryPlan { + split_point: usize, + retained_messages_count: usize, + active_context_limit: u32, + persistent_context_tokens: u32, + candidates: Vec, + summarize_messages: Vec, +} + +struct ActiveProviderRunIdentity { + conversation_id: AIConversationId, + stream_id: ResponseStreamId, + run_id: ProviderRunId, +} + fn configured_llm_context_limit(info: &LLMInfo) -> u32 { [ info.context_window.default_max, @@ -218,6 +240,79 @@ fn progressive_summary_split_point( split_point } +fn progressive_summary_content( + messages: &[ConversationMessage], + existing_summary: Option<&str>, +) -> String { + fn safe_truncate(text: &str, max_chars: usize) -> String { + if text.len() <= max_chars { + text.to_string() + } else { + let truncated = text.chars().take(max_chars).collect::(); + format!( + "{truncated}... [truncated, {total_chars} total chars]", + total_chars = text.len() + ) + } + } + + let mut content = String::new(); + if let Some(prior) = existing_summary { + content.push_str("\n"); + content.push_str(prior); + content.push_str("\n\n\n"); + } + content.push_str("\n"); + for message in messages { + let role = match message.role { + MessageRole::User => "User", + MessageRole::Assistant => "Assistant", + }; + let message_content = match &message.content { + MessageContent::Text(text) => text.clone(), + MessageContent::ToolUse { name, input, .. } => { + format!("[Tool Call: {name}] {input}") + } + MessageContent::ToolResult { content, .. } => safe_truncate(content, 2_000), + MessageContent::MultiPart(parts) => parts + .iter() + .map(|part| match part { + ContentPart::Text(text) => text.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, 2_000), + }) + .collect::>() + .join("\n"), + }; + content.push_str(&format!("[{role}]: {message_content}\n")); + } + content.push_str(""); + content +} + +fn progressive_summary_messages(summary: String) -> Vec { + vec![ + ConversationMessage { + role: MessageRole::User, + content: MessageContent::Text(format!( + "\n{summary}\n\n\n\ + The above summarizes earlier conversation history. The detailed messages below are the most recent exchanges." + )), + }, + ConversationMessage { + role: MessageRole::Assistant, + content: MessageContent::Text( + "Understood, I have the prior context. Continuing with the recent conversation." + .to_string(), + ), + }, + ] +} + async fn collect_progressive_summary( provider_config: ProviderConfig, request: TurnRequest, @@ -850,16 +945,19 @@ impl ProviderRetryStatus { } #[derive(Clone, Debug, PartialEq, Eq)] -struct ProviderToolCallProgressStatus { +pub(crate) struct ProviderToolCallProgressStatus { call_id: String, name: Option, arguments_bytes: u64, } impl ProviderToolCallProgressStatus { - fn label(&self) -> String { + pub(crate) fn label(&self) -> String { let activity = match self.name.as_deref() { Some("apply_file_diffs") => "Preparing file edit", + Some("run_shell_command") => "Preparing shell command", + Some("read_files") => "Preparing file read", + Some("grep" | "file_glob") => "Preparing file search", Some(name) if name.starts_with("mcp__") => "Preparing MCP tool call", Some(_) | None => "Preparing tool call", }; @@ -895,10 +993,10 @@ fn provider_tool_call_progress_update( arguments_bytes: *arguments_bytes, }), ProviderRunProjection::ModelTurnRequested { .. } - | ProviderRunProjection::ModelTurnFinished { .. } | ProviderRunProjection::ModelRetry { .. } | ProviderRunProjection::ToolBatchReady { .. } => ProviderToolCallProgressUpdate::Clear, ProviderRunProjection::ModelTurnStarted { .. } + | ProviderRunProjection::ModelTurnFinished { .. } | ProviderRunProjection::ModelEvent { .. } => ProviderToolCallProgressUpdate::Unchanged, } } @@ -922,7 +1020,7 @@ struct ActiveProviderRunSlot { pending_command_completion: Option, monitor_prose_continuations: usize, retry_status: Option, - tool_call_progress: Option, + tool_call_progress: Vec, } struct QueuedProviderRun { @@ -1512,15 +1610,22 @@ fn provider_command_completion_matches( .is_some_and(|execution_ref| execution_ref.run_id == *slot_run_id) } +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum ProviderCompletionSnapshotMatch { + Absent, + Match, + Mismatch, +} + fn reconcile_provider_completion_with_snapshot( completion: Option<&mut PendingProviderCommandCompletion>, block_id: &BlockId, expected_initial_action_id: &AIAgentActionId, snapshot_command: Option<&str>, fallback_command: Option<&str>, -) -> Result { +) -> ProviderCompletionSnapshotMatch { let Some(completion) = completion else { - return Ok(false); + return ProviderCompletionSnapshotMatch::Absent; }; if completion.block_id != *block_id || completion @@ -1528,7 +1633,7 @@ fn reconcile_provider_completion_with_snapshot( .as_ref() .is_some_and(|action_id| action_id != expected_initial_action_id) { - return Err("provider command completion did not match committed snapshot".to_owned()); + return ProviderCompletionSnapshotMatch::Mismatch; } if completion.command.is_empty() { completion.command = snapshot_command @@ -1536,7 +1641,7 @@ fn reconcile_provider_completion_with_snapshot( .unwrap_or_default() .to_owned(); } - Ok(true) + ProviderCompletionSnapshotMatch::Match } fn classify_provider_command_result( @@ -2155,14 +2260,14 @@ impl BlocklistAIController { .and_then(|slot| slot.retry_status) } - pub(crate) fn provider_tool_call_progress_label( + pub(crate) fn provider_tool_call_progress( &self, conversation_id: AIConversationId, - ) -> Option { + ) -> Vec { self.active_provider_runs .get(&conversation_id) - .and_then(|slot| slot.tool_call_progress.as_ref()) - .map(ProviderToolCallProgressStatus::label) + .map(|slot| slot.tool_call_progress.clone()) + .unwrap_or_default() } fn has_unresolved_ask_user_question( @@ -5204,7 +5309,7 @@ impl BlocklistAIController { pending_command_completion: None, monitor_prose_continuations: 0, retry_status: None, - tool_call_progress: None, + tool_call_progress: Vec::new(), }; match self.active_provider_runs.entry(conversation_data.id) { Entry::Occupied(_) => { @@ -5827,7 +5932,7 @@ impl BlocklistAIController { pending_command_completion, monitor_prose_continuations, retry_status: None, - tool_call_progress: None, + tool_call_progress: Vec::new(), }, ); if let Err(error) = @@ -6032,7 +6137,7 @@ impl BlocklistAIController { pending_command_completion: None, monitor_prose_continuations: 0, retry_status: None, - tool_call_progress: None, + tool_call_progress: Vec::new(), }, base_provider_config, cli_provider_config, @@ -6341,11 +6446,396 @@ impl BlocklistAIController { }) } + fn progressive_summary_candidates( + estimated_summary_input_tokens: u32, + terminal_surface_id: EntityId, + ctx: &AppContext, + ) -> Vec { + 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(terminal_surface_id)); + let base_model = preferences.get_active_base_model(ctx, Some(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); + 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, + 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, + context_limit: base_context_limit, + provider_config: base_provider_config, + }); + } + candidates + } + + fn begin_active_provider_progressive_summary( + &mut self, + conversation_id: AIConversationId, + ctx: &mut ModelContext, + ) -> bool { + let Some(active_run) = self + .active_provider_runs + .get(&conversation_id) + .and_then(|slot| slot.run.as_ref()) + else { + return false; + }; + if !matches!( + active_run.coordinator.run().state(), + ProviderRunState::ReadyToCallModel + ) { + return false; + } + + let active_context_limit = active_run + .response_config + .max_context_tokens + .unwrap_or_else(|| { + let preferences = LLMPreferences::as_ref(ctx); + configured_llm_context_limit( + preferences.get_active_base_model(ctx, Some(self.terminal_surface_id)), + ) + }) + .max(1); + let (current_context_tokens, summary_pending) = { + let history = BlocklistAIHistoryModel::as_ref(ctx); + let Some(conversation) = history.conversation(&conversation_id) else { + return false; + }; + ( + conversation + .current_context_tokens() + .min(active_context_limit), + conversation.has_pending_progressive_summary(), + ) + }; + if summary_pending + || current_context_tokens as f32 / (active_context_limit as f32) + < PROGRESSIVE_SUMMARY_TRIGGER_USAGE + { + return false; + } + + let transcript = active_run.coordinator.run().transcript(); + let transcript_tokens = estimated_history_tokens(transcript); + let persistent_context_tokens = current_context_tokens.saturating_sub(transcript_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(transcript, retained_history_budget); + if split_point == 0 { + return false; + } + + let summarize_content = progressive_summary_content(&transcript[..split_point], None); + let estimated_summary_input_tokens = estimated_text_tokens(PROGRESSIVE_SUMMARY_PROMPT) + .saturating_add(estimated_text_tokens(&summarize_content)); + let candidates = Self::progressive_summary_candidates( + estimated_summary_input_tokens, + self.terminal_surface_id, + ctx, + ); + if candidates.is_empty() { + log::warn!( + "[progressive-summary] No configured model can fit ~{} input tokens for active conversation {:?}", + estimated_summary_input_tokens, + conversation_id, + ); + return false; + } + let plan = ActiveProgressiveSummaryPlan { + split_point, + retained_messages_count: transcript.len().saturating_sub(split_point), + active_context_limit, + persistent_context_tokens, + candidates, + summarize_messages: vec![ConversationMessage { + role: MessageRole::User, + content: MessageContent::Text(summarize_content), + }], + }; + + let Some(slot) = self.active_provider_runs.get_mut(&conversation_id) else { + return false; + }; + let Some(run) = slot.run.take() else { + return false; + }; + let checkpoint = match ActiveProviderRunCheckpoint::from_active_run(&run) { + Ok(checkpoint) => checkpoint, + Err(error) => { + slot.run = Some(run); + self.fail_active_provider_run(conversation_id, error, ctx); + return true; + } + }; + slot.checkpoint = Some(checkpoint); + let stream_id = slot.stream_id.clone(); + let run_id = slot.run_id.clone(); + let identity = ActiveProviderRunIdentity { + conversation_id, + stream_id, + run_id, + }; + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history, _| { + if let Some(conversation) = history.conversation_mut(&conversation_id) { + conversation.set_has_pending_progressive_summary(true); + } + }); + if let Err(error) = self.persist_active_provider_run(conversation_id, ctx) { + if let Some(slot) = self.active_provider_runs.get_mut(&conversation_id) { + slot.run = Some(run); + } + self.fail_active_provider_run( + conversation_id, + format!("failed to checkpoint provider run before summarization: {error}"), + ctx, + ); + return true; + } + + let candidate_names = plan + .candidates + .iter() + .map(|candidate| format!("{}({})", candidate.model_id, candidate.context_limit)) + .join(", "); + log::info!( + "[progressive-summary] Pausing active run {:?}: summarizing {} messages, retaining {}, candidates=[{}]", + conversation_id, + plan.split_point, + plan.retained_messages_count, + candidate_names, + ); + ctx.spawn( + async move { + let mut errors = Vec::new(); + for candidate in plan.candidates.clone() { + let mut request = TurnRequest::new( + candidate.model_id.clone(), + plan.summarize_messages.clone(), + ); + request.system_prompt = Some(PROGRESSIVE_SUMMARY_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 (run, plan, Ok((summary, usage, candidate.model_id))); + } + Err(error) => errors.push(format!("{}: {error}", candidate.model_id)), + } + } + ( + run, + plan, + Err(anyhow!( + "all progressive summary candidates failed: {}", + errors.join("; ") + )), + ) + }, + move |me, (run, plan, result), ctx| { + me.finish_active_provider_progressive_summary(identity, run, plan, result, ctx); + }, + ); + true + } + + fn finish_active_provider_progressive_summary( + &mut self, + identity: ActiveProviderRunIdentity, + mut run: ActiveProviderRun, + plan: ActiveProgressiveSummaryPlan, + result: anyhow::Result<(String, Usage, String)>, + ctx: &mut ModelContext, + ) { + let ActiveProviderRunIdentity { + conversation_id, + stream_id, + run_id, + } = identity; + let Some(slot) = self.active_provider_runs.get(&conversation_id) else { + return; + }; + if slot.stream_id != stream_id || slot.run_id != run_id { + return; + } + + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history, _| { + if let Some(conversation) = history.conversation_mut(&conversation_id) { + conversation.set_has_pending_progressive_summary(false); + } + }); + + match result { + Ok((summary, usage, summarizer_model_id)) => { + let transcript = run.coordinator.run().transcript(); + if plan.split_point > transcript.len() { + if let Err(error) = run.coordinator.run_mut().fail( + ProviderRunFailureKind::Protocol, + "provider transcript changed while progressive summarization was running", + ) { + log::error!("Failed to record progressive summary protocol error: {error}"); + } + } else { + let drained = transcript[..plan.split_point].to_vec(); + let retained = transcript[plan.split_point..].to_vec(); + let summary_prefix = progressive_summary_messages(summary.clone()); + let summary_prefix_tokens = estimated_history_tokens(&summary_prefix); + if let Err(error) = run + .coordinator + .run_mut() + .compact_transcript_at_model_boundary(plan.split_point, summary_prefix) + { + let _ = run.coordinator.run_mut().fail( + ProviderRunFailureKind::Protocol, + format!("failed to compact provider transcript: {error}"), + ); + } else { + run.persistence_offset = 2; + if let Ok(mut messages_sent) = run.messages_sent.lock() { + *messages_sent = retained.clone(); + } + let new_context_tokens = plan + .persistent_context_tokens + .saturating_add(summary_prefix_tokens) + .saturating_add(estimated_history_tokens(&retained)); + let new_usage = + new_context_tokens as f32 / plan.active_context_limit.max(1) as f32; + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history, _| { + let Some(conversation) = history.conversation_mut(&conversation_id) + else { + return; + }; + conversation.archive_tool_results(drained); + *conversation.bedrock_message_history_mut() = retained; + conversation.set_progressive_summary(Some(summary.clone()), 0); + conversation.set_current_context_tokens(new_context_tokens); + conversation.set_context_window_usage(new_usage.clamp(0.0, 1.0)); + }); + log::info!( + "[progressive-summary] Resumed active run {:?} with {}: ~{} tokens ({:.1}% of {} context), {} messages retained", + conversation_id, + summarizer_model_id, + new_context_tokens, + new_usage * 100.0, + plan.active_context_limit, + plan.retained_messages_count, + ); + self.record_progressive_summary_usage( + conversation_id, + usage, + summarizer_model_id, + ctx, + ); + } + } + } + Err(error) => { + log::error!( + "[progressive-summary] Failed for active conversation {:?}: {error:#}", + conversation_id, + ); + } + } + + let Some(slot) = self.active_provider_runs.get_mut(&conversation_id) else { + return; + }; + slot.run = Some(run); + slot.checkpoint = slot + .run + .as_ref() + .and_then(|run| ActiveProviderRunCheckpoint::from_active_run(run).ok()); + if let Err(error) = self.persist_active_provider_run(conversation_id, ctx) { + self.fail_active_provider_run( + conversation_id, + format!("failed to persist provider run after summarization: {error}"), + ctx, + ); + return; + } + self.drive_active_provider_run(conversation_id, ctx); + } + + fn record_progressive_summary_usage( + &self, + conversation_id: AIConversationId, + usage: Usage, + summarizer_model_id: String, + ctx: &mut ModelContext, + ) { + 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, + ); + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history, ctx| { + use warp_multi_agent_api::response_event::stream_finished; + history.update_conversation_cost_and_usage_for_request( + conversation_id, + None, + vec![stream_finished::TokenUsage { + 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, + }], + None, + false, + ctx, + ); + }); + } + fn drive_active_provider_run( &mut self, conversation_id: AIConversationId, ctx: &mut ModelContext, ) { + if self.begin_active_provider_progressive_summary(conversation_id, ctx) { + return; + } let Some(slot) = self.active_provider_runs.get_mut(&conversation_id) else { return; }; @@ -6503,15 +6993,23 @@ impl BlocklistAIController { match tool_call_progress { ProviderToolCallProgressUpdate::Unchanged => false, ProviderToolCallProgressUpdate::Set(progress) => { - let changed = slot + if let Some(current) = slot .tool_call_progress - .as_ref() - .is_none_or(|current| current.label() != progress.label()); - slot.tool_call_progress = Some(progress); - changed + .iter_mut() + .find(|current| current.call_id == progress.call_id) + { + let changed = current.label() != progress.label(); + *current = progress; + changed + } else { + slot.tool_call_progress.push(progress); + true + } } ProviderToolCallProgressUpdate::Clear => { - slot.tool_call_progress.take().is_some() + let changed = !slot.tool_call_progress.is_empty(); + slot.tool_call_progress.clear(); + changed } } }; @@ -6793,6 +7291,18 @@ impl BlocklistAIController { } }; match block { + ProviderRunBlock::ReadyToCallModel => { + slot.run = Some(run); + if let Err(error) = self.persist_active_provider_run(conversation_id, ctx) { + self.fail_active_provider_run( + conversation_id, + format!("failed to persist provider model boundary: {error}"), + ctx, + ); + return; + } + self.drive_active_provider_run(conversation_id, ctx); + } ProviderRunBlock::Tools(batch) => { slot.committed_provider_batch = None; slot.finished_provider_batch = None; @@ -7098,7 +7608,7 @@ impl BlocklistAIController { match command_result { ProviderCommandResult::Snapshot { block_id, command } => { - let completion_already_pending = { + let (completion_match, monitored_block_id) = { let Some(slot) = self.active_provider_runs.get_mut(&conversation_id) else { return Ok(()); }; @@ -7111,15 +7621,33 @@ impl BlocklistAIController { .unwrap_or_else(|| result.id.clone()); let fallback_command = existing_monitor.map(|monitor| monitor.command.clone()); - reconcile_provider_completion_with_snapshot( - slot.pending_command_completion.as_mut(), - &block_id, - &expected_initial_action_id, - command.as_deref(), - fallback_command.as_deref(), - )? + ( + reconcile_provider_completion_with_snapshot( + slot.pending_command_completion.as_mut(), + &block_id, + &expected_initial_action_id, + command.as_deref(), + fallback_command.as_deref(), + ), + slot.command_monitor + .as_ref() + .map(|monitor| monitor.block_id.clone()), + ) }; - if completion_already_pending { + match completion_match { + ProviderCompletionSnapshotMatch::Match => continue, + ProviderCompletionSnapshotMatch::Mismatch => { + log::debug!( + "Ignoring provider command snapshot for block {block_id:?}; another command completion already owns the continuation boundary" + ); + continue; + } + ProviderCompletionSnapshotMatch::Absent => {} + } + if let Some(monitored_block_id) = monitored_block_id { + log::debug!( + "Ignoring provider command snapshot for block {block_id:?}; block {monitored_block_id:?} already owns automatic monitoring" + ); continue; } let cli_task_id = @@ -8222,17 +8750,8 @@ impl BlocklistAIController { warp_multi_agent_api::response_event::Type::Finished( finished_event, ) => { - self.handle_response_stream_finished( - &stream_id, - finished_event, - conversation_id, - did_input_contain_user_query, - ctx, - ); - - // After the stream finishes, persist the full message - // history (input + assistant response) from the Arc back - // into the conversation for the next request cycle. + // Persist provider-owned history before finalization so any + // end-of-stream compaction decision sees the completed turn. let new_history = response_stream .as_ref(ctx) .host_manages_history() @@ -8275,6 +8794,13 @@ impl BlocklistAIController { } }); } + self.handle_response_stream_finished( + &stream_id, + finished_event, + conversation_id, + did_input_contain_user_query, + ctx, + ); } warp_multi_agent_api::response_event::Type::ClientActions(actions) => { let client_actions = actions.actions; @@ -8977,27 +9503,29 @@ impl BlocklistAIController { ctx.emit(BlocklistAIControllerEvent::FreeTierLimitCheckTriggered); } - // 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. + // Session-owned runtimes summarize after the stream. Direct-provider runs + // compact at their durable ReadyToCallModel boundary instead, including + // immediately before the next user turn, so a background summary cannot + // race the live ProviderRun transcript. + let is_active_provider_run = self.active_provider_runs.contains_key(&conversation_id); let should_progressive_summarize = { let history_model = BlocklistAIHistoryModel::as_ref(ctx); - history_model - .conversation(&conversation_id) - .is_some_and(|conversation| { - let is_summarization_request = - conversation.latest_exchange().is_some_and(|exchange| { - exchange - .input - .iter() - .any(|i| matches!(i, AIAgentInput::SummarizeConversation { .. })) - }); - conversation.context_window_usage() >= PROGRESSIVE_SUMMARY_TRIGGER_USAGE - && !conversation.has_pending_progressive_summary() - && !is_summarization_request - && conversation.bedrock_message_history().len() - > PROGRESSIVE_SUMMARY_MIN_RECENT_MESSAGES - }) + !is_active_provider_run + && history_model + .conversation(&conversation_id) + .is_some_and(|conversation| { + let is_summarization_request = + conversation.latest_exchange().is_some_and(|exchange| { + exchange.input.iter().any(|i| { + matches!(i, AIAgentInput::SummarizeConversation { .. }) + }) + }); + conversation.context_window_usage() >= PROGRESSIVE_SUMMARY_TRIGGER_USAGE + && !conversation.has_pending_progressive_summary() + && !is_summarization_request + && conversation.bedrock_message_history().len() + > PROGRESSIVE_SUMMARY_MIN_RECENT_MESSAGES + }) }; if should_progressive_summarize { diff --git a/app/src/ai/blocklist/controller_tests.rs b/app/src/ai/blocklist/controller_tests.rs index 4898cb01..bf50f9a7 100644 --- a/app/src/ai/blocklist/controller_tests.rs +++ b/app/src/ai/blocklist/controller_tests.rs @@ -468,7 +468,7 @@ fn queued_provider_follow_up_persists_while_active_slot_is_unprepared() { pending_command_completion: None, monitor_prose_continuations: 0, retry_status: None, - tool_call_progress: None, + tool_call_progress: Vec::new(), }, ); controller @@ -495,7 +495,7 @@ fn queued_provider_follow_up_persists_while_active_slot_is_unprepared() { pending_command_completion: None, monitor_prose_continuations: 0, retry_status: None, - tool_call_progress: None, + tool_call_progress: Vec::new(), }, base_provider_config: crate::ai::provider::ProviderConfig::None, cli_provider_config: crate::ai::provider::ProviderConfig::None, @@ -1046,7 +1046,7 @@ fn cancelling_provider_startup_keeps_slot_and_durable_cancellation() { pending_command_completion: None, monitor_prose_continuations: 0, retry_status: None, - tool_call_progress: None, + tool_call_progress: Vec::new(), }, ); @@ -1146,7 +1146,7 @@ fn same_conversation_follow_up_waits_for_cancelled_provider_generation_cleanup() pending_command_completion: None, monitor_prose_continuations: 0, retry_status: None, - tool_call_progress: None, + tool_call_progress: Vec::new(), }, ); controller @@ -1179,7 +1179,7 @@ fn same_conversation_follow_up_waits_for_cancelled_provider_generation_cleanup() pending_command_completion: None, monitor_prose_continuations: 0, retry_status: None, - tool_call_progress: None, + tool_call_progress: Vec::new(), }, base_provider_config: crate::ai::provider::ProviderConfig::None, cli_provider_config: crate::ai::provider::ProviderConfig::None, @@ -1290,7 +1290,7 @@ fn non_follow_up_provider_cancellation_does_not_admit_an_overlapping_generation( pending_command_completion: None, monitor_prose_continuations: 0, retry_status: None, - tool_call_progress: None, + tool_call_progress: Vec::new(), }, ); @@ -2473,27 +2473,31 @@ fn pending_provider_completion_reconciles_with_its_committed_snapshot() { exit_code: 0, }; - assert!(super::reconcile_provider_completion_with_snapshot( - Some(&mut completion), - &block_id, - &action_id, - Some("sleep 10"), - None, - ) - .unwrap()); + assert_eq!( + super::reconcile_provider_completion_with_snapshot( + Some(&mut completion), + &block_id, + &action_id, + Some("sleep 10"), + None, + ), + super::ProviderCompletionSnapshotMatch::Match + ); assert_eq!(completion.command, "sleep 10"); - assert!(super::reconcile_provider_completion_with_snapshot( - None, - &block_id, - &action_id, - Some("sleep 10"), - None, - ) - .is_ok_and(|matched| !matched)); + assert_eq!( + super::reconcile_provider_completion_with_snapshot( + None, + &block_id, + &action_id, + Some("sleep 10"), + None, + ), + super::ProviderCompletionSnapshotMatch::Absent + ); } #[test] -fn pending_provider_completion_rejects_a_different_snapshot() { +fn pending_provider_completion_defers_a_different_snapshot() { let block_id = BlockId::new(); let action_id = AIAgentActionId::from("command-1".to_owned()); let mut completion = super::PendingProviderCommandCompletion { @@ -2504,23 +2508,27 @@ fn pending_provider_completion_rejects_a_different_snapshot() { exit_code: 0, }; - assert!(super::reconcile_provider_completion_with_snapshot( - Some(&mut completion), - &BlockId::new(), - &action_id, - Some("sleep 10"), - None, - ) - .is_err()); + assert_eq!( + super::reconcile_provider_completion_with_snapshot( + Some(&mut completion), + &BlockId::new(), + &action_id, + Some("sleep 10"), + None, + ), + super::ProviderCompletionSnapshotMatch::Mismatch + ); let completion_block_id = completion.block_id.clone(); - assert!(super::reconcile_provider_completion_with_snapshot( - Some(&mut completion), - &completion_block_id, - &AIAgentActionId::from("command-2".to_owned()), - Some("sleep 10"), - None, - ) - .is_err()); + assert_eq!( + super::reconcile_provider_completion_with_snapshot( + Some(&mut completion), + &completion_block_id, + &AIAgentActionId::from("command-2".to_owned()), + Some("sleep 10"), + None, + ), + super::ProviderCompletionSnapshotMatch::Mismatch + ); } #[test] diff --git a/app/src/ai/blocklist/inline_action/requested_command.rs b/app/src/ai/blocklist/inline_action/requested_command.rs index a15fa7ba..70571f48 100644 --- a/app/src/ai/blocklist/inline_action/requested_command.rs +++ b/app/src/ai/blocklist/inline_action/requested_command.rs @@ -1,4 +1,5 @@ use std::borrow::Cow; +use std::cell::Cell; use std::cmp::{Ordering, PartialEq}; use std::collections::HashMap; use std::rc::Rc; @@ -16,6 +17,7 @@ use galaxyui::elements::{ }; use galaxyui::keymap::{Context, EditableBinding, FixedBinding, Keystroke}; use galaxyui::ui_components::components::UiComponent as _; +use galaxyui::units::IntoPixels; use galaxyui::{ AppContext, Element, Entity, EntityId, EventContext, ModelHandle, SingletonEntity, TypedActionView, UpdateView, View, ViewContext, ViewHandle, @@ -79,8 +81,11 @@ const REQUESTED_COMMAND_EDIT_LABEL: &str = "Edit"; const REQUESTED_COMMAND_MINIMIZE_LABEL: &str = "Done"; const LOADING_MESSAGE: &str = "Generating command..."; -const COMMAND_WAITING_FOR_USER_MESSAGE: &str = "OK if I run this command and read the output?"; -const MCP_TOOL_WAITING_FOR_USER_MESSAGE: &str = "OK if I call this MCP tool?"; +const COMMAND_WAITING_FOR_USER_MESSAGE: &str = "Approve this command to continue"; +const MCP_TOOL_WAITING_FOR_USER_MESSAGE: &str = "Approve this MCP tool call to continue"; +const APPROVAL_REQUIRED_BADGE: &str = "Approval required"; +const WAITING_FOR_APPROVAL_BADGE: &str = "Waiting for approval above"; +const QUEUED_TOOL_BADGE: &str = "Queued behind an earlier tool"; const MONITORING_COMMAND_MESSAGE: &str = "Agent is monitoring command..."; const AGENT_NEEDS_INPUT_MESSAGE: &str = "Agent needs your input to continue"; const USER_TOOK_CONTROL_COMMAND_MESSAGE: &str = "User is in control."; @@ -329,6 +334,9 @@ pub struct RequestedCommandView { is_user_expanded: bool, header_mouse_state: MouseStateHandle, output_scroll_state: ClippedScrollStateHandle, + follow_command_output: Cell, + last_command_output_len: Cell, + last_output_scroll_start: Cell, is_editing: bool, // A requested command can either be copied directly off of one citation (such as a Warp Drive @@ -588,6 +596,9 @@ impl RequestedCommandView { is_user_expanded: false, header_mouse_state: Default::default(), output_scroll_state: ClippedScrollStateHandle::new(), + follow_command_output: Cell::new(true), + last_command_output_len: Cell::new(0), + last_output_scroll_start: Cell::new(0.), copied_from_citation: None, derived_from_citations: Default::default(), citation_state_handles: Default::default(), @@ -701,6 +712,12 @@ impl RequestedCommandView { return; } self.is_header_expanded = value; + if value && self.action_type.is_requested_command() { + self.follow_command_output.set(true); + self.last_command_output_len.set(0); + self.last_output_scroll_start + .set(self.output_scroll_state.scroll_start().as_f32()); + } if is_user_initiated { self.is_user_expanded = value; } @@ -1147,6 +1164,15 @@ impl RequestedCommandView { let mut title: Cow<'static, str>; let mut font_override = None; let mut font_color_override = None; + let mut badge = None; + + let conversation_is_blocked = + self.block_model + .as_ref() + .conversation(app) + .is_some_and(|conversation| { + matches!(conversation.status(), ConversationStatus::Blocked { .. }) + }); let terminal_model = self.terminal_model.lock(); let requested_command_block = match &self.action_type { @@ -1177,12 +1203,21 @@ impl RequestedCommandView { appearance.theme(), appearance.theme().surface_2(), )); + badge = Some( + if conversation_is_blocked { + WAITING_FOR_APPROVAL_BADGE + } else { + QUEUED_TOOL_BADGE + } + .to_string(), + ); } Some(AIActionStatus::Blocked) => { title = match &self.action_type { RequestedActionViewType::Command => COMMAND_WAITING_FOR_USER_MESSAGE.into(), RequestedActionViewType::McpTool => MCP_TOOL_WAITING_FOR_USER_MESSAGE.into(), }; + badge = Some(APPROVAL_REQUIRED_BADGE.to_string()); } Some(AIActionStatus::RunningAsync) | Some(AIActionStatus::Finished(..)) if self.is_header_expanded => @@ -1329,6 +1364,9 @@ impl RequestedCommandView { if let Some(font_color_override) = font_color_override { config = config.with_font_color(font_color_override); } + if let Some(badge) = badge { + config = config.with_badge(badge); + } match action_status { Some(AIActionStatus::Blocked) => { @@ -1564,6 +1602,27 @@ impl View for RequestedCommandView { } else { None }; + + if let Some(output) = command_output.as_ref() { + let scroll_start = self.output_scroll_state.scroll_start().as_f32(); + let last_scroll_start = self.last_output_scroll_start.get(); + + // Keep following live command output until the user deliberately scrolls upward. + // Collapsing and reopening the pane resumes following from the bottom. + if self.follow_command_output.get() && scroll_start + 0.5 < last_scroll_start { + self.follow_command_output.set(false); + } + + if self.follow_command_output.get() + && output.len() != self.last_command_output_len.get() + { + // Clipped scrollables clamp this sentinel to their actual maximum during layout. + self.output_scroll_state.scroll_to(f32::MAX.into_pixels()); + } + + self.last_command_output_len.set(output.len()); + self.last_output_scroll_start.set(scroll_start); + } let should_render_command_output = command_output.is_some(); let has_citations_footer = diff --git a/app/src/ai/runtime/provider_run_coordinator.rs b/app/src/ai/runtime/provider_run_coordinator.rs index 14ee6078..ce5f22c2 100644 --- a/app/src/ai/runtime/provider_run_coordinator.rs +++ b/app/src/ai/runtime/provider_run_coordinator.rs @@ -81,6 +81,7 @@ pub(crate) enum ProviderRunProjection { #[derive(Clone, Debug, PartialEq)] pub(crate) enum ProviderRunBlock { Tools(PendingToolBatch), + ReadyToCallModel, AwaitingDriver { work_id: ExternalWorkId, stop_reason: StopReason, @@ -383,7 +384,7 @@ impl ProviderRunCoordinator { } if batch.is_complete() { self.run.commit_tool_batch(&batch.work_id)?; - continue; + return Ok(ProviderRunBlock::ReadyToCallModel); } return Ok(ProviderRunBlock::Tools(batch)); } diff --git a/crates/galaxy_agent_core/src/provider_run.rs b/crates/galaxy_agent_core/src/provider_run.rs index 5d0b0363..64b0f317 100644 --- a/crates/galaxy_agent_core/src/provider_run.rs +++ b/crates/galaxy_agent_core/src/provider_run.rs @@ -353,6 +353,9 @@ pub enum ProviderRunProtocolError { InvalidDriverObservation { message: String, }, + InvalidTranscriptCompaction { + message: String, + }, InvalidRestoredState { message: String, }, @@ -415,6 +418,9 @@ impl fmt::Display for ProviderRunProtocolError { Self::InvalidDriverObservation { message } => { write!(f, "invalid driver observation: {message}") } + Self::InvalidTranscriptCompaction { message } => { + write!(f, "invalid transcript compaction: {message}") + } Self::InvalidRestoredState { message } => { write!(f, "invalid restored provider run: {message}") } @@ -490,6 +496,36 @@ impl ProviderRun { self.tool_result_archive = archive; } + /// Replaces a summarized prefix while the run is parked between model calls. + /// + /// Tool calls and results removed from the live transcript remain available to + /// `recall_tool_history` through the run-owned archive. + pub fn compact_transcript_at_model_boundary( + &mut self, + summarized_messages: usize, + summary_prefix: Vec, + ) -> Result<(), ProviderRunProtocolError> { + if !matches!(self.state, ProviderRunState::ReadyToCallModel) { + return Err(self.unexpected_state(ProviderRunPhase::ReadyToCallModel)); + } + if summarized_messages == 0 || summarized_messages > self.transcript.len() { + return Err(ProviderRunProtocolError::InvalidTranscriptCompaction { + message: format!( + "cannot summarize {summarized_messages} of {} messages", + self.transcript.len() + ), + }); + } + + let drained = self + .transcript + .drain(0..summarized_messages) + .collect::>(); + archive_tool_results(&mut self.tool_result_archive, drained); + self.transcript.splice(0..0, summary_prefix); + Ok(()) + } + pub fn usage(&self) -> &Usage { &self.usage } @@ -1662,6 +1698,48 @@ fn add_usage(total: &mut Usage, turn: &Usage) { .saturating_add(turn.cache_creation_input_tokens); } +fn archive_tool_results( + archive: &mut Vec, + messages: Vec, +) { + const MAX_TOOL_RESULT_ARCHIVE_ENTRIES: usize = 400; + + let mut pending_tool_uses = Vec::new(); + for message in messages { + match &message.content { + MessageContent::ToolUse { .. } => pending_tool_uses.push(message), + MessageContent::ToolResult { .. } => { + if let Some(tool_use) = pending_tool_uses.pop() { + archive.push(tool_use); + } + archive.push(message); + } + MessageContent::MultiPart(parts) => { + for part in parts { + match part { + ContentPart::ToolUse { .. } => pending_tool_uses.push(message.clone()), + ContentPart::ToolResult { .. } => { + if let Some(tool_use) = pending_tool_uses.pop() { + archive.push(tool_use); + } + archive.push(message.clone()); + } + ContentPart::Text(_) + | ContentPart::Reasoning { .. } + | ContentPart::Image { .. } => {} + } + } + } + MessageContent::Text(_) => {} + } + } + archive.extend(pending_tool_uses); + + if archive.len() > MAX_TOOL_RESULT_ARCHIVE_ENTRIES { + archive.drain(0..archive.len() - MAX_TOOL_RESULT_ARCHIVE_ENTRIES); + } +} + #[cfg(test)] #[path = "provider_run_tests.rs"] mod tests; diff --git a/crates/galaxy_agent_core/src/provider_run_tests.rs b/crates/galaxy_agent_core/src/provider_run_tests.rs index eaf67695..dd43336e 100644 --- a/crates/galaxy_agent_core/src/provider_run_tests.rs +++ b/crates/galaxy_agent_core/src/provider_run_tests.rs @@ -1211,6 +1211,60 @@ fn recovery_pending_tool_survives_another_restore_and_completes_once() { assert_eq!(restored.state().phase(), ProviderRunPhase::ReadyToCallModel); } +#[test] +fn transcript_compaction_at_model_boundary_archives_drained_tool_history() { + let mut run = run(); + let batch = accept_tool_turn( + &mut run, + tool_turn(vec![tool_call("read", "read_files")], &["read_files"]), + ); + run.complete_tool( + &batch.work_id, + successful_result("read", "important file contents"), + ) + .unwrap(); + run.commit_tool_batch(&batch.work_id).unwrap(); + let original_len = run.transcript().len(); + let summary_prefix = vec![ConversationMessage { + role: MessageRole::User, + content: MessageContent::Text("summary".to_string()), + }]; + + run.compact_transcript_at_model_boundary(original_len, summary_prefix.clone()) + .unwrap(); + + assert_eq!(run.transcript(), summary_prefix); + assert_eq!(run.tool_result_archive().len(), 2); + assert!(matches!( + run.tool_result_archive()[0].content, + MessageContent::MultiPart(_) + )); + assert!(matches!( + run.tool_result_archive()[1].content, + MessageContent::MultiPart(_) + )); +} + +#[test] +fn transcript_compaction_is_rejected_during_a_model_call() { + let mut run = run(); + let _ = next_model_call(&mut run); + + assert!(matches!( + run.compact_transcript_at_model_boundary( + 1, + vec![ConversationMessage { + role: MessageRole::User, + content: MessageContent::Text("summary".to_string()), + }], + ), + Err(ProviderRunProtocolError::UnexpectedState { + expected: ProviderRunPhase::ReadyToCallModel, + actual: ProviderRunPhase::AwaitingModel, + }) + )); +} + #[test] fn restore_normalization_commits_a_fully_resolved_batch() { let mut run = run(); diff --git a/crates/galaxy_agent_rig/src/stream.rs b/crates/galaxy_agent_rig/src/stream.rs index c57729ef..aaec5715 100644 --- a/crates/galaxy_agent_rig/src/stream.rs +++ b/crates/galaxy_agent_rig/src/stream.rs @@ -383,13 +383,28 @@ fn text_indicates_context_window_exceeded(text: &str) -> bool { || normalized.contains("exceeds the context") } +fn completion_error_indicates_recoverable_transport(error: &CompletionError) -> bool { + match error { + // Rig flattens non-status SSE transport failures into ProviderError. + // Preserve their transport semantics so the durable provider run can + // retry a truncated or otherwise interrupted response stream. + CompletionError::ProviderError(message) => message + .to_ascii_lowercase() + .starts_with("http client error:"), + _ => false, + } +} + fn map_completion_error(error: CompletionError) -> AgentError { let is_context_window_exceeded = completion_error_indicates_context_window_exceeded(&error); + let is_recoverable_transport = completion_error_indicates_recoverable_transport(&error); let status = error .provider_response_status() .map(|status| status.as_u16()); let kind = if is_context_window_exceeded { AgentErrorKind::ContextWindowExceeded + } else if is_recoverable_transport { + AgentErrorKind::Transport } else { match status { Some(401 | 403) => AgentErrorKind::Authentication, @@ -522,3 +537,7 @@ mod tests { assert!(mapped.recoverable); } } + +#[cfg(test)] +#[path = "stream_tests.rs"] +mod regression_tests; diff --git a/crates/galaxy_agent_rig/src/stream_tests.rs b/crates/galaxy_agent_rig/src/stream_tests.rs new file mode 100644 index 00000000..72fff657 --- /dev/null +++ b/crates/galaxy_agent_rig/src/stream_tests.rs @@ -0,0 +1,16 @@ +use galaxy_agent_core::AgentErrorKind; +use rig_core::completion::CompletionError; + +use super::map_completion_error; + +#[test] +fn flattened_sse_http_client_error_is_recoverable_transport() { + let error = CompletionError::ProviderError( + "Http client error: error decoding response body".to_string(), + ); + + let mapped = map_completion_error(error); + + assert_eq!(mapped.kind, AgentErrorKind::Transport); + assert!(mapped.recoverable); +}