From 2a89d038660b28867a9f88291b4e090d3d378a56 Mon Sep 17 00:00:00 2001 From: Ryan Ward Date: Thu, 24 Sep 2026 21:54:20 -0500 Subject: [PATCH] Bound agent context and preprocess oversized image attachments --- app/src/ai/blocklist/action_model/execute.rs | 17 +- app/src/ai/blocklist/controller.rs | 503 +++++++++++++----- app/src/ai/blocklist/controller_tests.rs | 285 ++++++++-- app/src/ai/blocklist/permissions_tests.rs | 55 ++ .../ai/runtime/provider_run_coordinator.rs | 136 ++++- .../runtime/provider_run_coordinator_tests.rs | 171 +++++- app/src/ai/runtime/rig.rs | 16 +- app/src/ai/runtime/rig_request.rs | 76 ++- app/src/ai/runtime/rig_request_tests.rs | 34 +- app/src/editor/view/mod.rs | 36 +- app/src/terminal/view.rs | 1 + app/src/util/image.rs | 224 ++++++-- app/src/util/image_tests.rs | 84 ++- crates/galaxy_agent_core/src/types.rs | 145 ++++- crates/galaxy_agent_core/src/types_tests.rs | 106 +++- 15 files changed, 1606 insertions(+), 283 deletions(-) diff --git a/app/src/ai/blocklist/action_model/execute.rs b/app/src/ai/blocklist/action_model/execute.rs index 24878e55..5fa25008 100644 --- a/app/src/ai/blocklist/action_model/execute.rs +++ b/app/src/ai/blocklist/action_model/execute.rs @@ -104,7 +104,8 @@ use crate::terminal::shell::ShellType; use crate::terminal::{ShellLaunchData, TerminalModel}; #[cfg(feature = "local_fs")] use crate::util::image::{ - is_supported_image_mime_type, process_image_for_agent, ProcessImageResult, + infer_mime_type, is_supported_image_mime_type, process_image_for_agent, ProcessImageResult, + MIME_SNIFF_BYTES, }; #[cfg(feature = "local_fs")] use crate::util::openable_file_type::is_binary_file; @@ -1554,7 +1555,19 @@ async fn read_binary_file_context( let mime_type = from_path(path).first_or_octet_stream().to_string(); let processed_content = if is_supported_image_mime_type(&mime_type) { match process_image_for_agent(&content) { - ProcessImageResult::Success { data } => Some(data), + ProcessImageResult::Success { data, mime_type } => { + let sniff_len = data.len().min(MIME_SNIFF_BYTES); + let inferred_mime_type = infer_mime_type(path, &data[..sniff_len]); + if inferred_mime_type != mime_type { + log::warn!( + "Processed image MIME mismatch for {}: declared {}, sniffed {}", + path.display(), + mime_type, + inferred_mime_type + ); + } + Some(data) + } ProcessImageResult::TooLarge => { log::warn!("Image file too large after processing: {}", path.display()); return Ok(BinaryFileReadResult::Missing); diff --git a/app/src/ai/blocklist/controller.rs b/app/src/ai/blocklist/controller.rs index 1ec88265..fb4d679f 100644 --- a/app/src/ai/blocklist/controller.rs +++ b/app/src/ai/blocklist/controller.rs @@ -23,10 +23,10 @@ use chrono::{DateTime, Local}; use futures::channel::oneshot; use futures::{FutureExt, StreamExt}; use galaxy_agent_core::{ - turn_control, AgentEvent, ExternalWorkId, PendingToolBatch, PendingToolCallState, ProviderRun, - ProviderRunFailureKind, ProviderRunId, ProviderRunLimits, ProviderRunOutcome, ProviderRunState, - StopReason, ToolLoopGuard, ToolResult, ToolResultStatus, TurnCommand, TurnCommandSender, - TurnRequest, Usage, + bounded_terminal_output_for_provider_request, turn_control, AgentEvent, ExternalWorkId, + PendingToolBatch, PendingToolCallState, ProviderRun, ProviderRunFailureKind, ProviderRunId, + ProviderRunLimits, ProviderRunOutcome, ProviderRunState, StopReason, ToolLoopGuard, ToolResult, + ToolResultStatus, TurnCommand, TurnCommandSender, TurnRequest, Usage, }; use galaxy_core::assertions::safe_assert; use input_context::{input_context_for_request, parse_context_attachments}; @@ -118,6 +118,7 @@ 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 PROVIDER_COMMAND_OUTPUT_POLL_INTERVAL: Duration = Duration::from_secs(2); const PROGRESSIVE_SUMMARY_PROMPT: &str = "Create a concise, factual task-state checkpoint from the following conversation history. Preserve:\n\ - The user's current goal, constraints, and durable preferences\n\ - Decisions made and the rationale needed to continue\n\ @@ -201,7 +202,7 @@ fn progressive_summary_candidate( } fn estimated_text_tokens(text: &str) -> u32 { - u32::try_from(text.len().div_ceil(3)).unwrap_or(u32::MAX) + u32::try_from(text.chars().count().div_ceil(3)).unwrap_or(u32::MAX) } fn estimated_message_tokens(message: &ConversationMessage) -> u32 { @@ -299,16 +300,15 @@ fn progressive_summary_split_point( fn progressive_summary_content( messages: &[ConversationMessage], existing_summary: Option<&str>, + max_chars: usize, ) -> String { fn safe_truncate(text: &str, max_chars: usize) -> String { - if text.len() <= max_chars { + let total_chars = text.chars().count(); + if total_chars <= 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() - ) + format!("{truncated}... [truncated, {total_chars} total chars]") } } @@ -347,7 +347,28 @@ fn progressive_summary_content( content.push_str(&format!("[{role}]: {message_content}\n")); } content.push_str(""); - content + if content.chars().count() <= max_chars { + return content; + } + + let marker = "\n... [summary input truncated to fit model context; omitted history remains in the conversation archive] ...\n"; + let marker_chars = marker.chars().count(); + if max_chars <= marker_chars { + return marker.chars().take(max_chars).collect(); + } + let retained_chars = max_chars - marker_chars; + let head_chars = retained_chars / 2; + let tail_chars = retained_chars - head_chars; + let head = content.chars().take(head_chars).collect::(); + let tail = content + .chars() + .rev() + .take(tail_chars) + .collect::() + .chars() + .rev() + .collect::(); + format!("{head}{marker}{tail}") } fn progressive_summary_messages(summary: String) -> Vec { @@ -896,10 +917,46 @@ struct RestoredProviderCommandEvidence { cli_task_id: Option, command: String, state: BlockState, + finished: bool, output: String, exit_code: i32, } +struct ProviderCommandOutputSnapshot { + state: BlockState, + finished: bool, + output: String, +} + +fn provider_monitor_matches_scheduled( + active: &ProviderCommandMonitorState, + scheduled: &ProviderCommandMonitorState, +) -> bool { + active.run_id == scheduled.run_id + && active.originating_work_id == scheduled.originating_work_id + && active.originating_call_id == scheduled.originating_call_id + && active.initial_requested_command_action_id + == scheduled.initial_requested_command_action_id + && active.block_id == scheduled.block_id + && active.cli_task_id == scheduled.cli_task_id +} + +fn provider_command_snapshot_needs_poll(snapshot: &ProviderCommandOutputSnapshot) -> bool { + matches!( + snapshot.state, + BlockState::BeforeExecution | BlockState::Executing | BlockState::Background + ) && !snapshot.finished +} + +fn reconcile_live_command_output_tokens(live_tokens: &mut u32, output: &str) -> bool { + let estimated_output_tokens = estimated_text_tokens(output); + if estimated_output_tokens <= *live_tokens { + return false; + } + *live_tokens = estimated_output_tokens; + true +} + #[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)] pub(super) struct PendingProviderCommandCompletion { block_id: BlockId, @@ -926,11 +983,11 @@ impl PendingProviderCommandCompletion { } } - fn observation(&self) -> MessageContent { + fn observation(&self, max_context_tokens: Option) -> MessageContent { let output = if self.output.is_empty() { - "(no output)" + "(no output)".to_string() } else { - self.output.as_str() + bounded_terminal_output_for_provider_request(&self.output, max_context_tokens) }; MessageContent::Text(format!( "The monitored command has finished with exit code {}. Continue the original objective \ @@ -1067,6 +1124,8 @@ struct ActiveProviderRunSlot { pending_monitor_observation: Option, pending_command_completion: Option, monitor_prose_continuations: usize, + live_command_output_tokens: u32, + command_output_poll_scheduled: bool, retry_status: Option, tool_call_progress: Vec, } @@ -1215,6 +1274,8 @@ struct ActiveProviderRunSnapshot { pending_command_completion: Option, monitor_prose_continuations: usize, #[serde(default)] + live_command_output_tokens: u32, + #[serde(default)] queued_follow_ups: Vec, } @@ -1258,6 +1319,7 @@ impl ActiveProviderRunSnapshot { pending_monitor_observation: slot.pending_monitor_observation.clone(), pending_command_completion: slot.pending_command_completion.clone(), monitor_prose_continuations: slot.monitor_prose_continuations, + live_command_output_tokens: slot.live_command_output_tokens, queued_follow_ups: Vec::new(), }) } @@ -1527,6 +1589,12 @@ fn apply_restored_provider_command_evidence( .to_owned(), exit_code: 130, }); + if let Some(completion) = snapshot.pending_command_completion.as_ref() { + reconcile_live_command_output_tokens( + &mut snapshot.live_command_output_tokens, + &completion.output, + ); + } return Ok(()); }; if evidence.conversation_id != Some(conversation_id) @@ -1539,15 +1607,24 @@ fn apply_restored_provider_command_evidence( if monitor.command.trim().is_empty() || evidence.command != monitor.command { return Err("restored provider command text does not match".to_string()); } - match evidence.state { - BlockState::BeforeExecution | BlockState::Executing => { + match (evidence.state, evidence.finished) { + (BlockState::BeforeExecution | BlockState::Executing | BlockState::Background, false) => { snapshot.pending_command_completion = None; snapshot.pending_monitor_observation = Some(PendingProviderMonitorObservation { block_id: monitor.block_id.clone(), cli_task_id: monitor.cli_task_id.clone(), }); + reconcile_live_command_output_tokens( + &mut snapshot.live_command_output_tokens, + &evidence.output, + ); } - BlockState::DoneWithExecution | BlockState::DoneWithNoExecution => { + ( + BlockState::DoneWithExecution + | BlockState::DoneWithNoExecution + | BlockState::Background, + true, + ) => { snapshot.pending_monitor_observation = None; snapshot.pending_command_completion = Some(PendingProviderCommandCompletion { block_id: monitor.block_id.clone(), @@ -1558,8 +1635,16 @@ fn apply_restored_provider_command_evidence( output: evidence.output, exit_code: evidence.exit_code, }); + if let Some(completion) = snapshot.pending_command_completion.as_ref() { + reconcile_live_command_output_tokens( + &mut snapshot.live_command_output_tokens, + &completion.output, + ); + } } - BlockState::Background | BlockState::Static => { + (BlockState::DoneWithExecution | BlockState::DoneWithNoExecution, false) + | (BlockState::BeforeExecution | BlockState::Executing | BlockState::Static, true) + | (BlockState::Static, false) => { return Err("restored provider command block has an invalid state".to_string()); } } @@ -1574,8 +1659,15 @@ fn merge_completion_offered_during_restore( return; } prepared.pending_command_completion = latest.pending_command_completion; - if prepared.pending_command_completion.is_some() { + prepared.live_command_output_tokens = prepared + .live_command_output_tokens + .max(latest.live_command_output_tokens); + if let Some(completion) = prepared.pending_command_completion.as_ref() { prepared.pending_monitor_observation = None; + reconcile_live_command_output_tokens( + &mut prepared.live_command_output_tokens, + &completion.output, + ); } } @@ -5453,6 +5545,8 @@ impl BlocklistAIController { pending_monitor_observation: None, pending_command_completion: None, monitor_prose_continuations: 0, + live_command_output_tokens: 0, + command_output_poll_scheduled: false, retry_status: None, tool_call_progress: Vec::new(), }; @@ -5778,6 +5872,7 @@ impl BlocklistAIController { snapshot.base_request.model.as_str(), snapshot.response_config.max_context_tokens, ); + let base_context_limit = snapshot.response_config.max_context_tokens; if let Err(error) = self.persist_provider_run_snapshot(conversation_id, &snapshot, ctx) { self.fail_restored_provider_run(conversation_id, error, ctx); return; @@ -5794,18 +5889,25 @@ impl BlocklistAIController { let mut profiles = BTreeMap::new(); profiles.insert( BASE_PROVIDER_PROFILE.to_string(), - ProviderRunProfile::new(base_runtime, snapshot.base_request.clone()), + ProviderRunProfile::new(base_runtime, snapshot.base_request.clone()) + .with_max_context_tokens(base_context_limit), ); if let Some(cli_monitor_request) = snapshot.cli_monitor_request.as_ref() { let cli_provider_config = cli_provider_config.ok_or_else(|| { anyhow!("restored CLI provider request is missing its provider config") })?; + let cli_context_limit = provider_context_window_tokens( + &cli_provider_config, + cli_monitor_request.model.as_str(), + base_context_limit, + ); let cli_runtime = provider_runtime_for_request(cli_provider_config, cli_monitor_request) .await?; profiles.insert( CLI_MONITOR_PROVIDER_PROFILE.to_string(), - ProviderRunProfile::new(cli_runtime, cli_monitor_request.clone()), + ProviderRunProfile::new(cli_runtime, cli_monitor_request.clone()) + .with_max_context_tokens(cli_context_limit), ); } Ok::<_, anyhow::Error>(PreparedRestoredProviderRun { @@ -5896,6 +5998,7 @@ impl BlocklistAIController { cli_task_id: block.cli_subagent_task_id().cloned(), command: block.command_to_string(), state: block.state(), + finished: block.finished(), output: block.output_to_string(), exit_code: block.exit_code().value(), }) @@ -5956,6 +6059,10 @@ impl BlocklistAIController { .restoring_provider_command_completions .remove(&conversation_id) { + reconcile_live_command_output_tokens( + &mut snapshot.live_command_output_tokens, + &completion.output, + ); snapshot.pending_command_completion = Some(completion); snapshot.pending_monitor_observation = None; } @@ -5979,6 +6086,7 @@ impl BlocklistAIController { pending_monitor_observation, pending_command_completion, monitor_prose_continuations, + live_command_output_tokens, queued_follow_ups, } = snapshot; let run_id = provider_run.id().clone(); @@ -6086,6 +6194,8 @@ impl BlocklistAIController { pending_monitor_observation, pending_command_completion, monitor_prose_continuations, + live_command_output_tokens, + command_output_poll_scheduled: false, retry_status: None, tool_call_progress: Vec::new(), }, @@ -6291,6 +6401,8 @@ impl BlocklistAIController { pending_monitor_observation: None, pending_command_completion: None, monitor_prose_continuations: 0, + live_command_output_tokens: 0, + command_output_poll_scheduled: false, retry_status: None, tool_call_progress: Vec::new(), }, @@ -6483,8 +6595,7 @@ impl BlocklistAIController { } = prepared; let mut coordinator = match ProviderRunCoordinator::from_request( provider_run_id, - base_profile.runtime, - base_profile.request, + base_profile, tool_result_archive, ProviderRunLimits::default(), ) { @@ -6496,11 +6607,7 @@ impl BlocklistAIController { }; coordinator.set_max_context_tokens(response_config.max_context_tokens); if let Some(profile) = cli_monitor_profile { - if let Err(error) = coordinator.insert_profile( - CLI_MONITOR_PROVIDER_PROFILE, - profile.runtime, - profile.request, - ) { + if let Err(error) = coordinator.insert_profile(CLI_MONITOR_PROVIDER_PROFILE, profile) { self.fail_provider_startup(conversation_id, stream_id, error.to_string(), ctx); return; } @@ -6656,7 +6763,6 @@ impl BlocklistAIController { } fn progressive_summary_candidates( - estimated_summary_input_tokens: u32, terminal_surface_id: EntityId, ctx: &AppContext, ) -> Vec { @@ -6675,12 +6781,6 @@ impl BlocklistAIController { 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") @@ -6688,9 +6788,7 @@ impl BlocklistAIController { }; let mut candidates = Vec::new(); - if model_is_usable(&terminal_model_id, &terminal_provider_config) - && candidate_fits(terminal_context_limit) - { + if model_is_usable(&terminal_model_id, &terminal_provider_config) { candidates.push(progressive_summary_candidate( terminal_model_id, terminal_context_limit, @@ -6698,7 +6796,6 @@ impl BlocklistAIController { )); } if model_is_usable(&base_model_id, &base_provider_config) - && candidate_fits(base_context_limit) && !candidates .iter() .any(|candidate| candidate.selection_model_id == base_model_id) @@ -6712,6 +6809,50 @@ impl BlocklistAIController { candidates } + fn live_provider_command_context_tokens(slot: &ActiveProviderRunSlot) -> u32 { + if slot.command_monitor.is_some() || slot.pending_command_completion.is_some() { + slot.live_command_output_tokens + } else { + 0 + } + } + + fn active_provider_context_estimate( + &self, + conversation_id: AIConversationId, + ctx: &mut ModelContext, + ) -> Option<(u32, u32)> { + let history = BlocklistAIHistoryModel::as_ref(ctx); + let conversation = history.conversation(&conversation_id)?; + let current_context_tokens = conversation.current_context_tokens(); + let active_run = self + .active_provider_runs + .get(&conversation_id)? + .run + .as_ref()?; + let transcript = active_run.coordinator.run().transcript(); + let (estimated_context_tokens, persistent_context_tokens) = active_run + .coordinator + .estimated_context_tokens() + .unwrap_or_else(|| { + ( + current_context_tokens, + current_context_tokens.saturating_sub(estimated_history_tokens(transcript)), + ) + }); + let live_command_output_tokens = self + .active_provider_runs + .get(&conversation_id) + .map(Self::live_provider_command_context_tokens) + .unwrap_or(0); + Some(( + current_context_tokens + .max(estimated_context_tokens) + .saturating_add(live_command_output_tokens), + persistent_context_tokens, + )) + } + fn begin_active_provider_progressive_summary( &mut self, conversation_id: AIConversationId, @@ -6750,27 +6891,19 @@ impl BlocklistAIController { .max(1); let last_failure_context_tokens = active_run.last_progressive_summary_failure_context_tokens; - let (current_context_tokens, summary_pending) = { + let summary_pending = { let history = BlocklistAIHistoryModel::as_ref(ctx); let Some(conversation) = history.conversation(&conversation_id) else { return false; }; - ( - conversation.current_context_tokens(), - conversation.has_pending_progressive_summary(), - ) + conversation.has_pending_progressive_summary() }; let transcript = active_run.coordinator.run().transcript(); - let (estimated_context_tokens, persistent_context_tokens) = active_run - .coordinator - .estimated_context_tokens() - .unwrap_or_else(|| { - ( - current_context_tokens, - current_context_tokens.saturating_sub(estimated_history_tokens(transcript)), - ) - }); - let current_context_tokens = current_context_tokens.max(estimated_context_tokens); + let Some((current_context_tokens, persistent_context_tokens)) = + self.active_provider_context_estimate(conversation_id, ctx) + else { + return false; + }; // Trigger and retain against the space history can actually use. Fixed // prompts and a large output reservation cannot be reduced by summarizing. let history_capacity = active_context_limit.saturating_sub(persistent_context_tokens); @@ -6793,17 +6926,44 @@ impl BlocklistAIController { 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, - ); + let candidates = Self::progressive_summary_candidates(self.terminal_surface_id, ctx); + let prompt_tokens = estimated_text_tokens(PROGRESSIVE_SUMMARY_PROMPT); + let Some(max_summary_input_tokens) = candidates + .iter() + .map(|candidate| { + candidate + .context_limit + .saturating_sub(PROGRESSIVE_SUMMARY_OUTPUT_TOKENS) + .saturating_sub(PROGRESSIVE_SUMMARY_CONTEXT_SAFETY_TOKENS) + .saturating_sub(prompt_tokens) + }) + .max() + else { + log::warn!( + "[progressive-summary] No configured model is available for active conversation {:?}", + conversation_id, + ); + return false; + }; + let max_summary_chars = usize::try_from(max_summary_input_tokens) + .unwrap_or(usize::MAX) + .saturating_mul(3); + let summarize_content = + progressive_summary_content(&transcript[..split_point], None, max_summary_chars); + let estimated_summary_input_tokens = + prompt_tokens.saturating_add(estimated_text_tokens(&summarize_content)); + let candidates = candidates + .into_iter() + .filter(|candidate| { + estimated_summary_input_tokens + .saturating_add(PROGRESSIVE_SUMMARY_OUTPUT_TOKENS) + .saturating_add(PROGRESSIVE_SUMMARY_CONTEXT_SAFETY_TOKENS) + <= candidate.context_limit + }) + .collect::>(); if candidates.is_empty() { log::warn!( - "[progressive-summary] No configured model can fit ~{} input tokens for active conversation {:?}", + "[progressive-summary] No configured model can fit bounded ~{} token summary input for active conversation {:?}", estimated_summary_input_tokens, conversation_id, ); @@ -7104,6 +7264,110 @@ impl BlocklistAIController { false } + fn snapshot_provider_command_output( + &self, + block_id: &BlockId, + ) -> Option { + let terminal_model = self.terminal_model.lock(); + let block = terminal_model.block_list().block_with_id(block_id)?; + Some(ProviderCommandOutputSnapshot { + state: block.state(), + finished: block.finished(), + output: block.output_to_string(), + }) + } + + fn schedule_provider_command_output_poll( + &mut self, + conversation_id: AIConversationId, + ctx: &mut ModelContext, + ) { + let monitor = { + let Some(slot) = self.active_provider_runs.get_mut(&conversation_id) else { + return; + }; + if slot.command_output_poll_scheduled { + return; + } + let Some(monitor) = slot.command_monitor.clone() else { + return; + }; + slot.command_output_poll_scheduled = true; + monitor + }; + ctx.spawn( + async move { Timer::after(PROVIDER_COMMAND_OUTPUT_POLL_INTERVAL).await }, + move |me, _, ctx| { + me.poll_provider_command_output(conversation_id, monitor, ctx); + }, + ); + } + + fn poll_provider_command_output( + &mut self, + conversation_id: AIConversationId, + scheduled_monitor: ProviderCommandMonitorState, + ctx: &mut ModelContext, + ) { + let (monitor, should_schedule_replacement_poll) = { + let Some(slot) = self.active_provider_runs.get_mut(&conversation_id) else { + return; + }; + slot.command_output_poll_scheduled = false; + let Some(active_monitor) = slot.command_monitor.clone() else { + return; + }; + if !provider_monitor_matches_scheduled(&active_monitor, &scheduled_monitor) { + (None, slot.pending_command_completion.is_none()) + } else { + (Some(active_monitor), false) + } + }; + if should_schedule_replacement_poll { + self.schedule_provider_command_output_poll(conversation_id, ctx); + } + let Some(monitor) = monitor else { + return; + }; + let Some(snapshot) = self.snapshot_provider_command_output(&monitor.block_id) else { + return; + }; + let mut should_drive = false; + let mut should_schedule_next_poll = provider_command_snapshot_needs_poll(&snapshot); + if let Some(slot) = self.active_provider_runs.get_mut(&conversation_id) { + let monitor_is_current = slot.command_monitor.as_ref().is_some_and(|active_monitor| { + provider_monitor_matches_scheduled(active_monitor, &monitor) + }); + if !monitor_is_current { + return; + } + if reconcile_live_command_output_tokens( + &mut slot.live_command_output_tokens, + &snapshot.output, + ) { + should_drive = true; + } + if slot.pending_command_completion.is_some() { + should_schedule_next_poll = false; + } + } + if should_drive { + self.drive_active_provider_run(conversation_id, ctx); + } + if should_schedule_next_poll { + let monitor_still_current = self + .active_provider_runs + .get(&conversation_id) + .and_then(|slot| slot.command_monitor.as_ref()) + .is_some_and(|active_monitor| { + provider_monitor_matches_scheduled(active_monitor, &monitor) + }); + if monitor_still_current { + self.schedule_provider_command_output_poll(conversation_id, ctx); + } + } + } + fn record_progressive_summary_usage( &self, conversation_id: AIConversationId, @@ -7147,6 +7411,7 @@ impl BlocklistAIController { if self.apply_ready_active_provider_progressive_summary(conversation_id, ctx) { return; } + self.schedule_provider_command_output_poll(conversation_id, ctx); self.begin_active_provider_progressive_summary(conversation_id, false, ctx); if matches!( self.active_provider_progressive_summaries @@ -7493,7 +7758,7 @@ impl BlocklistAIController { ); } run.set_task_id(&slot.root_task_id); - let observation = completion.observation(); + let observation = completion.observation(run.response_config.max_context_tokens); if ready_work_id.is_some() { run.coordinator .run_mut() @@ -7511,6 +7776,7 @@ impl BlocklistAIController { slot.command_monitor = None; slot.pending_monitor_observation = None; slot.monitor_prose_continuations = 0; + slot.live_command_output_tokens = 0; return Ok(ProviderBoundaryDisposition::Advance { completed_block_id: Some(completed_block_id), }); @@ -8149,6 +8415,7 @@ impl BlocklistAIController { block_id, cli_task_id, }); + self.schedule_provider_command_output_poll(conversation_id, ctx); } ProviderCommandResult::Finished { block_id, @@ -8167,6 +8434,10 @@ impl BlocklistAIController { continue; }; if slot.pending_command_completion.is_none() { + reconcile_live_command_output_tokens( + &mut slot.live_command_output_tokens, + &output, + ); slot.pending_command_completion = Some(PendingProviderCommandCompletion { block_id, initial_requested_command_action_id: Some( @@ -8841,6 +9112,10 @@ impl BlocklistAIController { .map(|monitor| monitor.command.clone()) .unwrap_or_default(); } + reconcile_live_command_output_tokens( + &mut slot.live_command_output_tokens, + &completion.output, + ); slot.pending_command_completion = Some(completion); slot.committed_provider_batch.is_none() && slot.run.as_ref().is_some_and(|run| { @@ -8950,6 +9225,10 @@ impl BlocklistAIController { } } snapshot.pending_monitor_observation = None; + reconcile_live_command_output_tokens( + &mut snapshot.live_command_output_tokens, + &completion.output, + ); snapshot.pending_command_completion = Some(completion); match self.persist_provider_run_snapshot(conversation_id, &snapshot, ctx) { Ok(()) => true, @@ -10149,68 +10428,54 @@ impl BlocklistAIController { ) }; - let mut summarize_content = String::new(); - if let Some(ref prior) = existing_summary { - summarize_content.push_str("\n"); - summarize_content.push_str(prior); - summarize_content.push_str("\n\n\n"); - } - summarize_content.push_str("\n"); - fn safe_truncate(s: &str, max_chars: usize) -> String { - if s.len() <= max_chars { - s.to_string() - } else { - let trunc = s.chars().take(max_chars).collect::(); - format!("{trunc}... [truncated, {len} total chars]", len = s.len()) - } - } - - for msg in &messages_to_summarize { - let role_str = match msg.role { - MessageRole::User => "User", - MessageRole::Assistant => "Assistant", - }; - let content_str = match &msg.content { - MessageContent::Text(t) => t.clone(), - MessageContent::ToolUse { name, input, .. } => { - format!("[Tool Call: {name}] {input}") - } - MessageContent::ToolResult { content, .. } => safe_truncate(content, 2000), - MessageContent::MultiPart(parts) => parts - .iter() - .map(|p| match p { - ContentPart::Text(t) => t.clone(), - ContentPart::Reasoning { text, .. } => { - format!("[Reasoning] {text}") - } - ContentPart::Image { .. } => "[Image attachment]".to_string(), - ContentPart::ToolUse { name, input, .. } => { - format!("[Tool: {name}] {input}") - } - ContentPart::ToolResult { content, .. } => safe_truncate(content, 2000), - }) - .collect::>() - .join("\n"), - }; - summarize_content.push_str(&format!("[{role_str}]: {content_str}\n")); - } - summarize_content.push_str(""); - let summarize_prompt = PROGRESSIVE_SUMMARY_PROMPT; - - let estimated_summary_input_tokens = estimated_text_tokens(summarize_prompt) - .saturating_add(estimated_text_tokens(&summarize_content)); + let candidates = [ + ( + &terminal_model_id, + terminal_context_limit, + &terminal_provider_config, + ), + (&base_model_id, base_context_limit, &base_provider_config), + ]; + 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 prompt_tokens = estimated_text_tokens(summarize_prompt); + let Some(max_summary_input_tokens) = candidates + .iter() + .filter(|(model_id, _, provider_config)| model_is_usable(model_id, provider_config)) + .map(|(_, context_limit, _)| { + context_limit + .saturating_sub(PROGRESSIVE_SUMMARY_OUTPUT_TOKENS) + .saturating_sub(PROGRESSIVE_SUMMARY_CONTEXT_SAFETY_TOKENS) + .saturating_sub(prompt_tokens) + }) + .max() + else { + log::warn!( + "[progressive-summary] No configured model is available for conversation {:?}", + conversation_id, + ); + return; + }; + let max_summary_chars = usize::try_from(max_summary_input_tokens) + .unwrap_or(usize::MAX) + .saturating_mul(3); + let summarize_content = progressive_summary_content( + &messages_to_summarize, + existing_summary.as_deref(), + max_summary_chars, + ); + let estimated_summary_input_tokens = + prompt_tokens.saturating_add(estimated_text_tokens(&summarize_content)); let candidate_fits = |context_limit: u32| { estimated_summary_input_tokens .saturating_add(PROGRESSIVE_SUMMARY_OUTPUT_TOKENS) .saturating_add(PROGRESSIVE_SUMMARY_CONTEXT_SAFETY_TOKENS) <= context_limit }; - let model_is_usable = |model_id: &str, provider_config: &ProviderConfig| { - !model_id.trim().is_empty() - && !model_id.eq_ignore_ascii_case("placeholder") - && !matches!(provider_config, ProviderConfig::None) - }; let mut candidates = Vec::new(); if model_is_usable(&terminal_model_id, &terminal_provider_config) && candidate_fits(terminal_context_limit) diff --git a/app/src/ai/blocklist/controller_tests.rs b/app/src/ai/blocklist/controller_tests.rs index d70bff37..f92652be 100644 --- a/app/src/ai/blocklist/controller_tests.rs +++ b/app/src/ai/blocklist/controller_tests.rs @@ -289,6 +289,7 @@ fn provider_snapshot(conversation_id: AIConversationId) -> super::ActiveProvider pending_monitor_observation: None, pending_command_completion: None, monitor_prose_continuations: 0, + live_command_output_tokens: 0, queued_follow_ups: Vec::new(), } } @@ -311,6 +312,21 @@ fn provider_snapshot_persists_cancellation_reason() { assert!(!restored.run.is_terminal()); } +#[test] +fn provider_snapshot_defaults_live_command_output_tokens_for_restore() { + let conversation_id = AIConversationId::new(); + let snapshot = provider_snapshot(conversation_id); + let mut value = serde_json::to_value(&snapshot).unwrap(); + value + .as_object_mut() + .unwrap() + .remove("live_command_output_tokens"); + + let restored = super::ActiveProviderRunSnapshot::parse(&value.to_string()).unwrap(); + + assert_eq!(restored.live_command_output_tokens, 0); +} + #[test] fn queued_provider_follow_up_snapshot_survives_restore_round_trip() { let conversation_id = AIConversationId::new(); @@ -470,6 +486,8 @@ fn queued_provider_follow_up_persists_while_active_slot_is_unprepared() { pending_monitor_observation: None, pending_command_completion: None, monitor_prose_continuations: 0, + live_command_output_tokens: 0, + command_output_poll_scheduled: false, retry_status: None, tool_call_progress: Vec::new(), }, @@ -497,6 +515,8 @@ fn queued_provider_follow_up_persists_while_active_slot_is_unprepared() { pending_monitor_observation: None, pending_command_completion: None, monitor_prose_continuations: 0, + live_command_output_tokens: 0, + command_output_poll_scheduled: false, retry_status: None, tool_call_progress: Vec::new(), }, @@ -1050,6 +1070,8 @@ fn cancelling_provider_startup_keeps_slot_and_durable_cancellation() { pending_monitor_observation: None, pending_command_completion: None, monitor_prose_continuations: 0, + live_command_output_tokens: 0, + command_output_poll_scheduled: false, retry_status: None, tool_call_progress: Vec::new(), }, @@ -1152,6 +1174,8 @@ fn same_conversation_follow_up_waits_for_cancelled_provider_generation_cleanup() pending_monitor_observation: None, pending_command_completion: None, monitor_prose_continuations: 0, + live_command_output_tokens: 0, + command_output_poll_scheduled: false, retry_status: None, tool_call_progress: Vec::new(), }, @@ -1185,6 +1209,8 @@ fn same_conversation_follow_up_waits_for_cancelled_provider_generation_cleanup() pending_monitor_observation: None, pending_command_completion: None, monitor_prose_continuations: 0, + live_command_output_tokens: 0, + command_output_poll_scheduled: false, retry_status: None, tool_call_progress: Vec::new(), }, @@ -1296,6 +1322,8 @@ fn non_follow_up_provider_cancellation_does_not_admit_an_overlapping_generation( pending_monitor_observation: None, pending_command_completion: None, monitor_prose_continuations: 0, + live_command_output_tokens: 0, + command_output_poll_scheduled: false, retry_status: None, tool_call_progress: Vec::new(), }, @@ -1794,6 +1822,7 @@ fn restored_active_command_rebuilds_monitor_observation() { cli_task_id: Some(cli_task_id.clone()), command: "sleep 10".to_owned(), state: BlockState::Executing, + finished: false, output: "running".to_owned(), exit_code: 0, }), @@ -1806,46 +1835,74 @@ fn restored_active_command_rebuilds_monitor_observation() { .expect("active command should restore monitoring"); assert_eq!(observation.block_id, block_id); assert_eq!(observation.cli_task_id, cli_task_id); + assert!(snapshot.live_command_output_tokens > 0); } #[test] fn restored_completed_command_rebuilds_exact_completion_evidence() { - let conversation_id = AIConversationId::new(); - let mut snapshot = provider_snapshot(conversation_id); - let (action_id, block_id, cli_task_id) = - attach_snapshot_command_monitor(&mut snapshot, conversation_id); - snapshot.pending_monitor_observation = Some(super::PendingProviderMonitorObservation { - block_id: block_id.clone(), - cli_task_id: cli_task_id.clone(), - }); + for (state, finished) in [ + (BlockState::DoneWithExecution, true), + (BlockState::Background, true), + ] { + let conversation_id = AIConversationId::new(); + let mut snapshot = provider_snapshot(conversation_id); + let (action_id, block_id, cli_task_id) = + attach_snapshot_command_monitor(&mut snapshot, conversation_id); + snapshot.pending_monitor_observation = Some(super::PendingProviderMonitorObservation { + block_id: block_id.clone(), + cli_task_id: cli_task_id.clone(), + }); - super::apply_restored_provider_command_evidence( - conversation_id, - &mut snapshot, - Some(super::RestoredProviderCommandEvidence { - conversation_id: Some(conversation_id), - requested_command_action_id: Some(action_id.clone()), - cli_task_id: Some(cli_task_id), - command: "sleep 10".to_owned(), - state: BlockState::DoneWithExecution, - output: "done".to_owned(), - exit_code: 17, - }), - ) - .unwrap(); + super::apply_restored_provider_command_evidence( + conversation_id, + &mut snapshot, + Some(super::RestoredProviderCommandEvidence { + conversation_id: Some(conversation_id), + requested_command_action_id: Some(action_id.clone()), + cli_task_id: Some(cli_task_id), + command: "sleep 10".to_owned(), + state, + finished, + output: "done".to_owned(), + exit_code: 17, + }), + ) + .unwrap(); - assert!(snapshot.pending_monitor_observation.is_none()); - let completion = snapshot - .pending_command_completion - .expect("completed command should restore final evidence"); - assert_eq!(completion.block_id, block_id); - assert_eq!( - completion.initial_requested_command_action_id, - Some(action_id) + assert!(snapshot.pending_monitor_observation.is_none()); + let completion = snapshot + .pending_command_completion + .expect("completed command should restore final evidence"); + assert_eq!(completion.block_id, block_id); + assert_eq!( + completion.initial_requested_command_action_id, + Some(action_id) + ); + assert_eq!(completion.command, "sleep 10"); + assert_eq!(completion.output, "done"); + assert_eq!(completion.exit_code, 17); + assert!(snapshot.live_command_output_tokens > 0); + } +} + +#[test] +fn monitored_completion_bounds_model_observation_without_mutating_durable_output() { + let output = format!("first:{}:last", "界".repeat(20_000)); + let completion = super::PendingProviderCommandCompletion::new( + BlockId::new(), + None, + "generate output".to_owned(), + output.clone(), + 0, ); - assert_eq!(completion.command, "sleep 10"); - assert_eq!(completion.output, "done"); - assert_eq!(completion.exit_code, 17); + + let MessageContent::Text(observation) = completion.observation(Some(8_000)) else { + panic!("expected text observation"); + }; + assert!(observation.contains("terminal output truncated for model context")); + assert!(observation.contains("full output remains available in the terminal")); + assert!(!observation.contains(&output)); + assert_eq!(completion.output, output); } #[test] @@ -1910,6 +1967,7 @@ fn completion_offered_during_restore_stays_provider_owned_and_durable() { .and_then(|json| super::ActiveProviderRunSnapshot::parse(json).ok()) .unwrap(); assert!(restored.pending_monitor_observation.is_none()); + assert!(restored.live_command_output_tokens > 0); let completion = restored.pending_command_completion.unwrap(); assert_eq!(completion.block_id, block_id); assert_eq!( @@ -1944,6 +2002,7 @@ fn prepared_restore_merges_a_later_durable_completion_once() { super::merge_completion_offered_during_restore(&mut prepared, latest); assert!(prepared.pending_monitor_observation.is_none()); + assert!(prepared.live_command_output_tokens > 0); let completion = prepared.pending_command_completion.unwrap(); assert_eq!(completion.block_id, block_id); assert_eq!( @@ -2052,6 +2111,7 @@ fn restored_evidence_is_ignored_without_a_command_monitor() { cli_task_id: None, command: "sleep 10".to_owned(), state: BlockState::Executing, + finished: false, output: "running".to_owned(), exit_code: 0, }; @@ -2096,6 +2156,7 @@ fn restored_command_rejects_mismatched_or_invalid_terminal_evidence() { cli_task_id: Some(cli_task_id.clone()), command: "sleep 10".to_owned(), state: BlockState::Executing, + finished: false, output: String::new(), exit_code: 0, }), @@ -2112,7 +2173,8 @@ fn restored_command_rejects_mismatched_or_invalid_terminal_evidence() { requested_command_action_id: Some(action_id), cli_task_id: Some(cli_task_id), command: "sleep 10".to_owned(), - state: BlockState::Background, + state: BlockState::Static, + finished: false, output: String::new(), exit_code: 0, }), @@ -2403,6 +2465,119 @@ fn provider_monitor_prose_retry_is_bounded_then_parks() { ); } +#[test] +fn provider_command_snapshot_polling_tracks_only_unfinished_executing_states() { + for state in [ + BlockState::BeforeExecution, + BlockState::Executing, + BlockState::Background, + ] { + assert!(super::provider_command_snapshot_needs_poll( + &super::ProviderCommandOutputSnapshot { + state, + finished: false, + output: String::new(), + } + )); + assert!(!super::provider_command_snapshot_needs_poll( + &super::ProviderCommandOutputSnapshot { + state, + finished: true, + output: String::new(), + } + )); + } + for state in [BlockState::DoneWithExecution, BlockState::Static] { + assert!(!super::provider_command_snapshot_needs_poll( + &super::ProviderCommandOutputSnapshot { + state, + finished: false, + output: String::new(), + } + )); + } +} + +#[test] +fn live_command_output_tokens_only_grow() { + let mut live_tokens = 0; + assert!(super::reconcile_live_command_output_tokens( + &mut live_tokens, + &"x".repeat(100), + )); + let after_growth = live_tokens; + + assert!(!super::reconcile_live_command_output_tokens( + &mut live_tokens, + "short", + )); + assert_eq!(live_tokens, after_growth); + + assert!(super::reconcile_live_command_output_tokens( + &mut live_tokens, + &"x".repeat(1_000), + )); + assert!(live_tokens > after_growth); +} + +#[test] +fn provider_poll_monitor_identity_matches_the_full_owner() { + let action_id = AIAgentActionId::from("command-1".to_owned()); + let work_id = ExternalWorkId { + run_id: ProviderRunId::new("run-1"), + epoch: RunEpoch::new(2), + }; + let monitor = super::ProviderCommandMonitorState { + run_id: work_id.run_id.clone(), + originating_work_id: work_id.clone(), + originating_call_id: action_id.to_string(), + initial_requested_command_action_id: action_id.clone(), + block_id: BlockId::new(), + command: "sleep 10".to_owned(), + cli_task_id: TaskId::new("cli-task".to_owned()), + }; + assert!(super::provider_monitor_matches_scheduled( + &monitor, &monitor + )); + + let mut replacement_call = monitor.clone(); + replacement_call.originating_call_id = "command-2".to_owned(); + assert!(!super::provider_monitor_matches_scheduled( + &replacement_call, + &monitor, + )); + + let mut replacement_action = monitor.clone(); + replacement_action.initial_requested_command_action_id = + AIAgentActionId::from("command-2".to_owned()); + assert!(!super::provider_monitor_matches_scheduled( + &replacement_action, + &monitor, + )); + + let mut replacement_block = monitor.clone(); + replacement_block.block_id = BlockId::new(); + assert!(!super::provider_monitor_matches_scheduled( + &replacement_block, + &monitor, + )); + + let mut replacement_cli_task = monitor; + replacement_cli_task.cli_task_id = TaskId::new("replacement-cli-task".to_owned()); + assert!(!super::provider_monitor_matches_scheduled( + &replacement_cli_task, + &super::ProviderCommandMonitorState { + run_id: work_id.run_id.clone(), + originating_work_id: work_id, + originating_call_id: action_id.to_string(), + initial_requested_command_action_id: action_id, + block_id: replacement_cli_task.block_id.clone(), + command: "sleep 10".to_owned(), + cli_task_id: TaskId::new("cli-task".to_owned()), + }, + )); +} + #[test] fn provider_completion_requires_current_run_and_exact_monitor_identity() { let conversation_id = AIConversationId::new(); @@ -2547,7 +2722,7 @@ fn nonzero_provider_completion_is_continuation_evidence() { output: "one test failed".to_owned(), exit_code: 17, }; - let galaxy_agent_core::MessageContent::Text(observation) = completion.observation() else { + let galaxy_agent_core::MessageContent::Text(observation) = completion.observation(None) else { panic!("command completion must be text evidence"); }; assert!(observation.contains("exit code 17")); @@ -3783,6 +3958,46 @@ fn progressive_summary_includes_an_oversized_final_tool_batch() { ); } +#[test] +fn progressive_summary_input_is_bounded_and_discloses_omitted_history() { + let messages = vec![ + ConversationMessage { + role: MessageRole::User, + content: MessageContent::Text(format!("old-start:{}", "x".repeat(10_000))), + }, + ConversationMessage { + role: MessageRole::Assistant, + content: MessageContent::Text(format!("{}:new-end", "y".repeat(10_000))), + }, + ]; + + let content = super::progressive_summary_content(&messages, None, 1_000); + + assert_eq!(content.chars().count(), 1_000); + assert!(content.starts_with("")); + assert!(content.ends_with("new-end\n")); + assert!(content.contains("summary input truncated to fit model context")); + assert!(content.contains("omitted history remains in the conversation archive")); +} + +#[test] +fn progressive_summary_budget_is_character_safe_for_multibyte_history() { + let messages = vec![ConversationMessage { + role: MessageRole::User, + content: MessageContent::ToolResult { + tool_use_id: "large-output".to_string(), + content: format!("first:{}:last", "界".repeat(10_000)), + is_error: false, + }, + }]; + + let content = super::progressive_summary_content(&messages, None, 1_000); + + assert_eq!(content.chars().count(), 1_000); + assert!(content.contains("summary input truncated to fit model context")); + assert!(content.contains("omitted history remains in the conversation archive")); +} + #[test] fn progressive_summary_keeps_retained_tool_results_paired() { let messages = vec![ diff --git a/app/src/ai/blocklist/permissions_tests.rs b/app/src/ai/blocklist/permissions_tests.rs index fc6784e2..594bc216 100644 --- a/app/src/ai/blocklist/permissions_tests.rs +++ b/app/src/ai/blocklist/permissions_tests.rs @@ -504,6 +504,61 @@ fn test_can_write_files_mcp_config_always_denied() { }) } +#[test] +fn test_fresh_conversation_uses_active_command_only_profile() { + App::test((), |mut app| async move { + let PermissionsTestState { + permissions, + history, + profile_model, + terminal_view_id, + .. + } = initialize_permissions_test(&mut app); + + profile_model.update(&mut app, |model, ctx| { + let profile_id = model + .create_profile(ctx) + .expect("profile should be created"); + model.set_execute_commands(profile_id, &ActionPermission::AlwaysAllow, ctx); + model.set_apply_code_diffs(profile_id, &ActionPermission::AlwaysAsk, ctx); + model.set_active_profile(terminal_view_id, profile_id, ctx); + }); + + let fresh_conversation_id = history.update(&mut app, |history_model, ctx| { + history_model.start_new_conversation(terminal_view_id, false, false, false, ctx) + }); + + permissions.read(&app, |model, ctx| { + let command_permission = model.can_autoexecute_command( + &fresh_conversation_id, + "git status", + EscapeChar::Backslash, + false, + None, + Some(terminal_view_id), + ctx, + ); + assert!(matches!( + command_permission, + CommandExecutionPermission::Allowed( + CommandExecutionPermissionAllowedReason::AlwaysAllowed + ) + )); + + let file_write_permission = model.can_write_files( + &fresh_conversation_id, + &[PathBuf::from("file.txt")], + Some(terminal_view_id), + ctx, + ); + assert!(matches!( + file_write_permission, + FileWritePermission::Denied(FileWritePermissionDeniedReason::AlwaysAskEnabled) + )); + }); + }) +} + #[test] fn test_can_autoexecute_command_workspace_settings_override_profile() { App::test((), |mut app| async move { diff --git a/app/src/ai/runtime/provider_run_coordinator.rs b/app/src/ai/runtime/provider_run_coordinator.rs index 025bcac8..66eac805 100644 --- a/app/src/ai/runtime/provider_run_coordinator.rs +++ b/app/src/ai/runtime/provider_run_coordinator.rs @@ -163,11 +163,21 @@ impl From for ProviderRunCoordinatorError { pub(crate) struct ProviderRunProfile { pub(crate) runtime: Arc, pub(crate) request: TurnRequest, + pub(crate) max_context_tokens: Option, } impl ProviderRunProfile { pub(crate) fn new(runtime: Arc, request: TurnRequest) -> Self { - Self { runtime, request } + Self { + runtime, + request, + max_context_tokens: None, + } + } + + pub(crate) fn with_max_context_tokens(mut self, max_context_tokens: Option) -> Self { + self.max_context_tokens = max_context_tokens; + self } } @@ -183,23 +193,19 @@ pub(crate) struct ProviderRunCoordinator { impl ProviderRunCoordinator { pub(crate) fn from_request( run_id: impl Into, - runtime: Arc, - request: TurnRequest, + base_profile: ProviderRunProfile, tool_result_archive: Vec, limits: ProviderRunLimits, ) -> Result { let mut run = ProviderRun::new( run_id, - request.messages.clone(), + base_profile.request.messages.clone(), BASE_PROVIDER_PROFILE, limits, ); run.replace_tool_result_archive(tool_result_archive); let mut profiles = BTreeMap::new(); - profiles.insert( - BASE_PROVIDER_PROFILE.to_string(), - ProviderRunProfile::new(runtime, request), - ); + profiles.insert(BASE_PROVIDER_PROFILE.to_string(), base_profile); Self::new(run, profiles) } @@ -266,13 +272,11 @@ impl ProviderRunCoordinator { pub(crate) fn insert_profile( &mut self, profile: impl Into, - runtime: Arc, - request: TurnRequest, + config: ProviderRunProfile, ) -> Result<(), ProviderRunCoordinatorError> { let profile = profile.into(); - validate_profile_runtime(&profile, runtime.as_ref())?; - self.profiles - .insert(profile, ProviderRunProfile::new(runtime, request)); + validate_profile_runtime(&profile, config.runtime.as_ref())?; + self.profiles.insert(profile, config); Ok(()) } @@ -461,13 +465,15 @@ impl ProviderRunCoordinator { } fn request_needs_context_compaction(&self, call: &ProviderModelCall) -> bool { - let Some(max_context_tokens) = self.max_context_tokens else { - return false; - }; let Some(profile) = self.profiles.get(call.profile.as_str()) else { return false; }; - let request = request_for_model_call(profile.request.clone(), call); + let Some(max_context_tokens) = profile.max_context_tokens.or(self.max_context_tokens) + else { + return false; + }; + let request = + request_for_model_call(profile.request.clone(), call, Some(max_context_tokens)); // Provider tokenizers and request envelopes differ. Compact before the // estimate reaches the advertised limit, leaving room for that uncertainty. let safe_context_tokens = u64::from(max_context_tokens) * 9 / 10; @@ -547,7 +553,11 @@ impl ProviderRunCoordinator { { return Ok(()); } - let request = request_for_model_call(profile.request.clone(), &call); + let request = request_for_model_call( + profile.request.clone(), + &call, + profile.max_context_tokens.or(self.max_context_tokens), + ); let started_at = Instant::now(); let (attempt_sender, attempt_control) = turn_control(); let start_future = profile.runtime.start_turn(request, attempt_control).fuse(); @@ -1132,9 +1142,18 @@ fn validate_profile_runtime( Ok(()) } -fn request_for_model_call(mut template: TurnRequest, call: &ProviderModelCall) -> TurnRequest { +fn request_for_model_call( + mut template: TurnRequest, + call: &ProviderModelCall, + max_context_tokens: Option, +) -> TurnRequest { template.prompt = None; template.messages = call.messages.clone(); + for message in &mut template.messages { + message.truncate_tool_results_for_provider_request_with_context(max_context_tokens); + message + .truncate_tagged_terminal_output_for_provider_request_with_context(max_context_tokens); + } add_exploration_reminder(&mut template); template } @@ -1206,33 +1225,92 @@ fn consecutive_discovery_batches(messages: &[ConversationMessage]) -> usize { const ESTIMATED_BYTES_PER_TOKEN: u64 = 3; const ESTIMATED_REQUEST_OVERHEAD_TOKENS: u64 = 512; +const ESTIMATED_IMAGE_TOKENS: u64 = 1_600; -/// Conservative estimate, not an exact tokenizer count. UTF-8 bytes account for -/// non-ASCII content better than character counts; the caller also leaves headroom. +/// Heuristic estimate, not an exact tokenizer count or provider-specific bound. +/// UTF-8 bytes account for non-ASCII content better than character counts; raw +/// image payload bytes are excluded from text accounting and charged separately. fn estimate_turn_request_tokens(request: &TurnRequest) -> u64 { - let mut bytes = request + let mut text_bytes = request .system_prompt .as_deref() .map_or(0, |text| text.len()); - bytes += request.prompt.as_ref().map_or(0, |prompt| { - serde_json::to_string(prompt).unwrap_or_default().len() - }); - bytes += request + text_bytes += request + .prompt + .as_ref() + .map_or(0, estimate_message_content_bytes); + text_bytes += request .messages .iter() - .map(|message| serde_json::to_string(message).unwrap_or_default().len()) + .map(estimate_conversation_message_bytes) .sum::(); - bytes += request + text_bytes += request .tools .iter() .map(|tool| serde_json::to_string(tool).unwrap_or_default().len()) .sum::(); - let input_tokens = (bytes as u64).div_ceil(ESTIMATED_BYTES_PER_TOKEN); + + let image_tokens = request + .prompt + .as_ref() + .map_or(0, count_images_in_content) + .saturating_add( + request + .messages + .iter() + .map(|message| count_images_in_content(&message.content)) + .sum::(), + ) + .saturating_mul(ESTIMATED_IMAGE_TOKENS); + let input_tokens = (text_bytes as u64).div_ceil(ESTIMATED_BYTES_PER_TOKEN); input_tokens + .saturating_add(image_tokens) .saturating_add(ESTIMATED_REQUEST_OVERHEAD_TOKENS) .saturating_add(request.max_output_tokens.unwrap_or_default()) } +fn estimate_conversation_message_bytes(message: &ConversationMessage) -> usize { + let mut sanitized_message = message.clone(); + remove_image_payloads_from_content(&mut sanitized_message.content); + serde_json::to_string(&sanitized_message) + .unwrap_or_default() + .len() +} + +fn estimate_message_content_bytes(content: &MessageContent) -> usize { + let mut sanitized_content = content.clone(); + remove_image_payloads_from_content(&mut sanitized_content); + serde_json::to_string(&sanitized_content) + .unwrap_or_default() + .len() +} + +fn remove_image_payloads_from_content(content: &mut MessageContent) { + if let MessageContent::MultiPart(parts) = content { + for part in parts { + match part { + ContentPart::Image { data, .. } => data.clear(), + ContentPart::Text(_) + | ContentPart::Reasoning { .. } + | ContentPart::ToolUse { .. } + | ContentPart::ToolResult { .. } => {} + } + } + } +} + +fn count_images_in_content(content: &MessageContent) -> u64 { + match content { + MessageContent::Text(_) + | MessageContent::ToolUse { .. } + | MessageContent::ToolResult { .. } => 0, + MessageContent::MultiPart(parts) => parts + .iter() + .filter(|part| matches!(part, ContentPart::Image { .. })) + .count() as u64, + } +} + fn tool_event_call_id(event: &ToolEvent) -> Result<&str, ProviderRunCoordinatorError> { match event { ToolEvent::Proposed { call } => Ok(&call.id), diff --git a/app/src/ai/runtime/provider_run_coordinator_tests.rs b/app/src/ai/runtime/provider_run_coordinator_tests.rs index 0e07ecce..d7eb2e17 100644 --- a/app/src/ai/runtime/provider_run_coordinator_tests.rs +++ b/app/src/ai/runtime/provider_run_coordinator_tests.rs @@ -190,6 +190,34 @@ fn estimates_turn_request_with_system_tools_messages_and_output() { ); } +#[test] +fn estimates_image_payloads_without_counting_raw_bytes_as_text() { + let image_bytes = vec![255; 3_750_000]; + let serialized_text_tokens = serde_json::to_string(&image_bytes) + .unwrap() + .len() + .div_ceil(ESTIMATED_BYTES_PER_TOKEN as usize) as u64; + let request = TurnRequest::new( + "test-model", + vec![ConversationMessage { + role: MessageRole::User, + content: MessageContent::MultiPart(vec![ + ContentPart::Text("what is in this screenshot?".to_string()), + ContentPart::Image { + data: image_bytes, + mime_type: "image/jpeg".to_string(), + }, + ]), + }], + ); + + let estimate = estimate_turn_request_tokens(&request); + + assert!(estimate < 10_000); + assert!(estimate > ESTIMATED_REQUEST_OVERHEAD_TOKENS + ESTIMATED_IMAGE_TOKENS); + assert!(serialized_text_tokens > estimate); +} + #[tokio::test] async fn proactively_compacts_before_starting_an_oversized_request() { let runtime = Arc::new(ScriptedRuntime::new(vec![answer_turn()])); @@ -276,8 +304,7 @@ fn answer_turn() -> ScriptedTurn { fn coordinator(runtime: Arc) -> ProviderRunCoordinator { ProviderRunCoordinator::from_request( "run-1", - runtime, - request(), + ProviderRunProfile::new(runtime, request()), Vec::new(), ProviderRunLimits::default(), ) @@ -1357,8 +1384,7 @@ fn session_runtime_is_rejected_before_any_turn_can_start() { }); let error = ProviderRunCoordinator::from_request( "run-1", - runtime, - request(), + ProviderRunProfile::new(runtime, request()), Vec::new(), ProviderRunLimits::default(), ) @@ -1379,8 +1405,7 @@ fn context_estimate_tracks_live_transcript_and_reserves_fixed_context() { request.max_output_tokens = Some(4_000); let mut coordinator = ProviderRunCoordinator::from_request( "context-budget", - runtime, - request, + ProviderRunProfile::new(runtime, request), Vec::new(), ProviderRunLimits::default(), ) @@ -1402,6 +1427,140 @@ fn context_estimate_tracks_live_transcript_and_reserves_fixed_context() { assert_eq!(fixed_before, fixed_after); } +#[tokio::test] +async fn monitor_request_uses_its_profile_context_limit() { + let base_runtime = Arc::new(ScriptedRuntime::new(vec![answer_turn()])); + let cli_runtime = Arc::new(ScriptedRuntime::new(vec![answer_turn()])); + let original_content = format!("start:{}:end", "x".repeat(20_000)); + let large_result = ConversationMessage { + role: MessageRole::User, + content: MessageContent::ToolResult { + tool_use_id: "large-result".to_string(), + content: original_content.clone(), + is_error: false, + }, + }; + let run = ProviderRun::new( + "profile-limit-run", + vec![large_result], + CLI_MONITOR_PROVIDER_PROFILE, + ProviderRunLimits::default(), + ); + let mut profiles = BTreeMap::new(); + profiles.insert( + BASE_PROVIDER_PROFILE.to_string(), + ProviderRunProfile::new(base_runtime, request()).with_max_context_tokens(Some(100_000)), + ); + profiles.insert( + CLI_MONITOR_PROVIDER_PROFILE.to_string(), + ProviderRunProfile::new(cli_runtime.clone(), request()) + .with_max_context_tokens(Some(8_000)), + ); + let mut coordinator = ProviderRunCoordinator::new(run, profiles).unwrap(); + coordinator.set_max_context_tokens(Some(100_000)); + + let (_sender, control) = turn_control(); + let block = coordinator + .drive_until_blocked(control, |_| Ok(())) + .await + .unwrap(); + + assert!(matches!(block, ProviderRunBlock::AwaitingDriver { .. })); + let requests = cli_runtime.requests(); + let MessageContent::ToolResult { content, .. } = &requests[0].messages[0].content else { + panic!("expected bounded monitor tool result"); + }; + assert!(content.chars().count() <= galaxy_agent_core::provider_output_char_limit(Some(8_000))); + assert!(content.contains("tool result truncated")); + let MessageContent::ToolResult { content, .. } = &coordinator.run().transcript()[0].content + else { + panic!("expected original monitor tool result"); + }; + assert_eq!(content, &original_content); +} + +#[tokio::test] +async fn final_model_request_bounds_restored_tagged_terminal_text() { + let runtime = Arc::new(ScriptedRuntime::new(vec![answer_turn()])); + let original_output = format!("first:{}:last", "x".repeat(20_000)); + let restored_message = ConversationMessage { + role: MessageRole::User, + content: MessageContent::Text(format!( + "[Running command: verbose]\n[Terminal output:\n{original_output}\n]\nContinue" + )), + }; + let mut turn_request = request(); + turn_request.messages = vec![restored_message]; + let mut coordinator = ProviderRunCoordinator::from_request( + "restored-terminal-run", + ProviderRunProfile::new(runtime.clone(), turn_request).with_max_context_tokens(Some(8_000)), + Vec::new(), + ProviderRunLimits::default(), + ) + .unwrap(); + + let (_sender, control) = turn_control(); + let block = coordinator + .drive_until_blocked(control, |_| Ok(())) + .await + .unwrap(); + + assert!(matches!(block, ProviderRunBlock::AwaitingDriver { .. })); + let requests = runtime.requests(); + let MessageContent::Text(content) = &requests[0].messages[0].content else { + panic!("expected bounded restored terminal text"); + }; + assert!(content.contains("terminal output truncated for model context")); + assert!(!content.contains(&original_output)); + let MessageContent::Text(content) = &coordinator.run().transcript()[0].content else { + panic!("expected original restored terminal text"); + }; + assert!(content.contains(&original_output)); +} + +#[tokio::test] +async fn final_model_request_bounds_new_tool_results_to_the_active_context() { + let runtime = Arc::new(ScriptedRuntime::new(vec![answer_turn()])); + let original_content = format!("start:{}:end", "界".repeat(20_000)); + let large_result = ConversationMessage { + role: MessageRole::User, + content: MessageContent::ToolResult { + tool_use_id: "large-result".to_string(), + content: original_content.clone(), + is_error: false, + }, + }; + let mut turn_request = request(); + turn_request.messages = vec![large_result]; + let mut coordinator = ProviderRunCoordinator::from_request( + "large-result-run", + ProviderRunProfile::new(runtime.clone(), turn_request), + Vec::new(), + ProviderRunLimits::default(), + ) + .unwrap(); + coordinator.set_max_context_tokens(Some(8_000)); + + let (_sender, control) = turn_control(); + let block = coordinator + .drive_until_blocked(control, |_| Ok(())) + .await + .unwrap(); + + assert!(matches!(block, ProviderRunBlock::AwaitingDriver { .. })); + let requests = runtime.requests(); + let MessageContent::ToolResult { content, .. } = &requests[0].messages[0].content else { + panic!("expected bounded tool result"); + }; + assert!(content.chars().count() <= galaxy_agent_core::provider_output_char_limit(Some(8_000))); + assert!(content.contains("tool result truncated")); + let MessageContent::ToolResult { content, .. } = &coordinator.run().transcript()[0].content + else { + panic!("expected original tool result"); + }; + assert_eq!(content, &original_content); +} + #[tokio::test] async fn preflight_leaves_headroom_before_the_advertised_context_limit() { let runtime = Arc::new(ScriptedRuntime::new(vec![answer_turn()])); diff --git a/app/src/ai/runtime/rig.rs b/app/src/ai/runtime/rig.rs index c524c8bd..e875e7a4 100644 --- a/app/src/ai/runtime/rig.rs +++ b/app/src/ai/runtime/rig.rs @@ -129,11 +129,6 @@ pub(crate) async fn prepare_provider_run( let (supported_tools, supported_cli_agent_tools) = crate::ai::agent::api::prepare_direct_provider_params(&mut params); let skill_path_origin = params.session_context.skill_path_origin(); - let max_context_tokens = provider_context_window_tokens( - &base_provider_config, - params.model.as_str(), - params.context_window_limit, - ); let mut cli_params = params.clone(); let cli_model_is_placeholder = params.cli_agent_model.as_str().trim().is_empty() || params @@ -178,9 +173,13 @@ pub(crate) async fn prepare_provider_run( Some(RigRequestMode::Cli), ) .await?; - let cli_monitor_profile = Some(ProviderRunProfile::new(cli_runtime, cli_prepared.request)); + let cli_monitor_profile = Some( + ProviderRunProfile::new(cli_runtime, cli_prepared.request) + .with_max_context_tokens(cli_prepared.max_context_tokens), + ); let PreparedRigTurn { + max_context_tokens: prepared_context_tokens, task_id, needs_create_task, user_query, @@ -203,12 +202,13 @@ pub(crate) async fn prepare_provider_run( needs_create_task, user_query, model_id: request.model.as_str().to_string(), - max_context_tokens, + max_context_tokens: prepared_context_tokens, capabilities: base_runtime.descriptor().capabilities.clone(), empty_output_message: None, }; Ok(PreparedProviderRun { - base_profile: ProviderRunProfile::new(base_runtime, request), + base_profile: ProviderRunProfile::new(base_runtime, request) + .with_max_context_tokens(prepared_context_tokens), cli_monitor_profile, tool_result_archive, messages_sent, diff --git a/app/src/ai/runtime/rig_request.rs b/app/src/ai/runtime/rig_request.rs index d2ec9e74..bca4f5ef 100644 --- a/app/src/ai/runtime/rig_request.rs +++ b/app/src/ai/runtime/rig_request.rs @@ -6,9 +6,10 @@ use ai::skills::SkillReference; use base64::engine::general_purpose; use base64::Engine as _; use galaxy_agent_core::{ - ContentPart, ConversationMessage, MessageContent, MessageRole, ToolDefinition, ToolResult, - TurnRequest, + bounded_terminal_output_for_provider_request, ContentPart, ConversationMessage, MessageContent, + MessageRole, ToolDefinition, ToolResult, TurnRequest, }; +use galaxy_bedrock_model_catalog::model_metadata; use sha2::{Digest as _, Sha256}; use uuid::Uuid; use warp_multi_agent_api as api; @@ -26,6 +27,7 @@ use crate::ai::provider::request_translator::{ use crate::ai::provider::types::flatten_tool_history_for_no_tools_turn; pub(crate) struct PreparedRigTurn { + pub max_context_tokens: Option, pub task_id: String, pub needs_create_task: bool, pub user_query: Option, @@ -135,9 +137,12 @@ pub(crate) fn prepare_rig_turn( supported_tools: Vec, supported_cli_agent_tools: Vec, ) -> PreparedRigTurn { + let max_context_tokens = + effective_context_limit(config.max_input_tokens, params.context_window_limit); prepare_rig_turn_for_provider( config.model.clone(), config.max_output_tokens.map(u64::from), + max_context_tokens, RigRequestSanitizer::OpenAICompatible, params, supported_tools, @@ -153,9 +158,12 @@ pub(crate) fn prepare_rig_turn_for_mode( supported_cli_agent_tools: Vec, mode: RigRequestMode, ) -> PreparedRigTurn { + let max_context_tokens = + effective_context_limit(config.max_input_tokens, params.context_window_limit); prepare_rig_turn_for_provider( config.model.clone(), config.max_output_tokens.map(u64::from), + max_context_tokens, RigRequestSanitizer::OpenAICompatible, params, supported_tools, @@ -189,9 +197,12 @@ pub(crate) fn prepare_bedrock_rig_turn_for_mode( supported_cli_agent_tools: Vec, mode: Option, ) -> PreparedRigTurn { + let provider_limit = model_metadata(&model).and_then(|metadata| metadata.context_window_tokens); + let max_context_tokens = effective_context_limit(provider_limit, params.context_window_limit); prepare_rig_turn_for_provider( Some(model), max_output_tokens, + max_context_tokens, RigRequestSanitizer::Bedrock, params, supported_tools, @@ -206,9 +217,22 @@ enum RigRequestSanitizer { Bedrock, } +fn effective_context_limit( + provider_limit: Option, + configured_limit: Option, +) -> Option { + match (provider_limit, configured_limit) { + (Some(provider), Some(configured)) => Some(provider.min(configured)), + (provider, configured) => provider.or(configured), + } +} + +// Keep provider limits and tool sets explicit at this shared request adapter boundary. +#[allow(clippy::too_many_arguments)] fn prepare_rig_turn_for_provider( model_override: Option, max_output_tokens: Option, + max_context_tokens: Option, sanitizer: RigRequestSanitizer, params: RequestParams, supported_tools: Vec, @@ -269,11 +293,13 @@ fn prepare_rig_turn_for_provider( } let mut system_prompt = build_system_prompt(&input, &tools, &global_rules, mode); - let mut new_messages = input_messages(input, tool_results); + let mut new_messages = input_messages(input, tool_results, max_context_tokens); let mut persistent_messages = message_history; persistent_messages.append(&mut new_messages); for message in &mut persistent_messages { - message.truncate_tool_results_for_provider_request(); + message.truncate_tool_results_for_provider_request_with_context(max_context_tokens); + message + .truncate_tagged_terminal_output_for_provider_request_with_context(max_context_tokens); } match sanitizer { RigRequestSanitizer::OpenAICompatible => { @@ -311,6 +337,7 @@ fn prepare_rig_turn_for_provider( request.max_output_tokens = max_output_tokens; PreparedRigTurn { + max_context_tokens, task_id, needs_create_task, user_query, @@ -338,6 +365,7 @@ fn model_workflow_guidance(model_id: &str) -> &'static str { fn input_messages( inputs: Vec, tool_results: Vec, + max_context_tokens: Option, ) -> Vec { let mut messages = Vec::new(); if !tool_results.is_empty() { @@ -375,11 +403,18 @@ fn input_messages( }); } - messages.extend(inputs.into_iter().filter_map(input_message)); + messages.extend( + inputs + .into_iter() + .filter_map(|input| input_message(input, max_context_tokens)), + ); messages } -fn input_message(input: AIAgentInput) -> Option { +fn input_message( + input: AIAgentInput, + max_context_tokens: Option, +) -> Option { let (text, images) = match input { AIAgentInput::UserQuery { query, @@ -388,9 +423,13 @@ fn input_message(input: AIAgentInput) -> Option { .. } => { let text = if let Some(command) = running_command { + let output = bounded_terminal_output_for_provider_request( + &command.grid_contents, + max_context_tokens, + ); format!( "[Running command: {}]\n[Command ID: {}]\n[Terminal output:\n{}\n]\n{}", - command.command, command.block_id, command.grid_contents, query + command.command, command.block_id, output, query ) } else { query @@ -401,16 +440,19 @@ fn input_message(input: AIAgentInput) -> Option { prompt, context, completed_command, - } => ( - format!( - "[Completed command: {}]\n[Command ID: {}]\n[Final terminal output:\n{}\n]\n{}", - completed_command.command, - completed_command.block_id, - completed_command.grid_contents, - prompt - ), - image_parts(&context), - ), + } => { + let output = bounded_terminal_output_for_provider_request( + &completed_command.grid_contents, + max_context_tokens, + ); + ( + format!( + "[Completed command: {}]\n[Command ID: {}]\n[Final terminal output:\n{}\n]\n{}", + completed_command.command, completed_command.block_id, output, prompt + ), + image_parts(&context), + ) + } AIAgentInput::ActionResult { .. } => return None, AIAgentInput::AutoCodeDiffQuery { query, .. } => (query, Vec::new()), AIAgentInput::ResumeConversation { .. } => ( diff --git a/app/src/ai/runtime/rig_request_tests.rs b/app/src/ai/runtime/rig_request_tests.rs index 5df8dcd5..fe29c588 100644 --- a/app/src/ai/runtime/rig_request_tests.rs +++ b/app/src/ai/runtime/rig_request_tests.rs @@ -437,6 +437,38 @@ fn completed_command_assessment_uses_root_history_without_tools_or_monitor_instr assert_eq!(prepared.request.messages, prepared.persistent_messages); } +#[test] +fn completed_command_assessment_bounds_only_the_provider_payload() { + let full_output = format!("first\n{}\nlast", "x".repeat(100_000)); + let block_id: galaxy_terminal::model::BlockId = "large-completed-command".to_string().into(); + let mut params = RequestParams::new_for_test(); + params.context_window_limit = Some(8_000); + params.input = vec![AIAgentInput::CommandCompletionAssessment { + prompt: "Report the final result to the user.".to_string(), + context: Arc::from([]), + completed_command: crate::ai::agent::RunningCommand { + command: "verbose-command".to_string(), + block_id, + grid_contents: full_output.clone(), + cursor: String::new(), + requested_command_id: None, + is_alt_screen_active: false, + }, + }]; + + let prepared = prepare_rig_turn(&config(), params, Vec::new(), Vec::new()); + let MessageContent::Text(provider_text) = &prepared.request.messages[0].content else { + panic!("expected text command assessment"); + }; + + assert!(provider_text.contains("terminal output truncated for model context")); + assert!(provider_text.contains("full output remains available in the terminal")); + assert!(provider_text.contains("first")); + assert!(provider_text.contains("last")); + assert!(!provider_text.contains(&full_output)); + assert_eq!(full_output.chars().count(), 100_011); +} + #[test] fn forced_cli_profile_uses_monitor_prompt_and_tools_for_an_initial_query() { let mut params = RequestParams::new_for_test(); @@ -801,7 +833,7 @@ fn normalized_tool_outcomes_are_the_only_action_results_sent_to_rig() { }) .collect(); - let messages = input_messages(Vec::new(), tool_results); + let messages = input_messages(Vec::new(), tool_results, Some(128_000)); assert_eq!(messages.len(), 1); let MessageContent::MultiPart(parts) = &messages[0].content else { diff --git a/app/src/editor/view/mod.rs b/app/src/editor/view/mod.rs index ca2a6657..99b45d5c 100644 --- a/app/src/editor/view/mod.rs +++ b/app/src/editor/view/mod.rs @@ -125,8 +125,8 @@ use crate::util::bindings::{cmd_or_ctrl_shift, keybinding_name_to_keystroke, Cus use crate::util::clipboard::clipboard_content_with_escaped_paths; use crate::util::color::{ContrastingColor, MinimumAllowedContrast}; use crate::util::image::{ - infer_mime_type, is_supported_image_mime_type, resize_image, MAX_IMAGE_COUNT_FOR_QUERY, - MAX_IMAGE_SIZE_BYTES, MIME_SNIFF_BYTES, + infer_mime_type, is_supported_image_mime_type, prepare_image_for_agent, + MAX_IMAGE_COUNT_FOR_QUERY, MIME_SNIFF_BYTES, }; use crate::util::merge_ranges; use crate::view_components::DismissibleToast; @@ -5371,8 +5371,12 @@ impl EditorView { for image in pending_images { let is_figma = is_figma_png(&image.data); - let resized_image_bytes = match resize_image(&image.data) { - Ok(resized_image_bytes) => resized_image_bytes, + let processed_image = match prepare_image_for_agent(&image.data) { + Ok(Some(processed_image)) => processed_image, + Ok(None) => { + num_oversized_images += 1; + continue; + } Err(err) => { num_unprocessed_images += 1; log::warn!("Error resizing attached image {err:?}"); @@ -5380,26 +5384,28 @@ impl EditorView { } }; - if resized_image_bytes.len() > MAX_IMAGE_SIZE_BYTES { - num_oversized_images += 1; - continue; - } - - let sniff_len = resized_image_bytes.len().min(MIME_SNIFF_BYTES); - let mime_type = infer_mime_type( + let sniff_len = processed_image.data.len().min(MIME_SNIFF_BYTES); + let sniffed_mime_type = infer_mime_type( Path::new(&image.file_name), - &resized_image_bytes[..sniff_len], + &processed_image.data[..sniff_len], ); - if !is_supported_image_mime_type(&mime_type) { + if !is_supported_image_mime_type(&processed_image.mime_type) { num_unprocessed_images += 1; continue; } + if sniffed_mime_type != processed_image.mime_type { + log::warn!( + "Processed image MIME mismatch: declared {}, sniffed {}", + processed_image.mime_type, + sniffed_mime_type + ); + } - let base64_str = general_purpose::STANDARD.encode(&resized_image_bytes); + let base64_str = general_purpose::STANDARD.encode(&processed_image.data); processed_pending_images.push(ImageContext { data: base64_str, - mime_type, + mime_type: processed_image.mime_type, file_name: image.file_name, is_figma, }); diff --git a/app/src/terminal/view.rs b/app/src/terminal/view.rs index b8adae4d..e9d12bc6 100644 --- a/app/src/terminal/view.rs +++ b/app/src/terminal/view.rs @@ -7306,6 +7306,7 @@ impl TerminalView { if is_agent_in_control { self.redetermine_terminal_focus(ctx); } + ctx.notify(); } BlocklistAIActionEvent::ExecutingAction { .. } => { if event_matches_active_conversation() { diff --git a/app/src/util/image.rs b/app/src/util/image.rs index a31db8be..89353546 100644 --- a/app/src/util/image.rs +++ b/app/src/util/image.rs @@ -3,15 +3,24 @@ //! This module provides common functionality for processing images before they are //! sent to the AI agent, whether attached by the user or read via the read_files tool. +use std::io::Cursor; use std::path::Path; -use image::{GenericImageView, ImageError}; +use image::codecs::jpeg::JpegEncoder; +use image::{ + DynamicImage, ExtendedColorType, GenericImageView, ImageEncoder, ImageError, ImageFormat, +}; use mime_guess::from_path; /// Max image size is 3.75 MB. /// The max size of an image we will send is 5MB. However, due to the 33% inflation of Base64, this means /// the largest size a user can attach is actually ~3.75MB. pub const MAX_IMAGE_SIZE_BYTES: usize = 3750 * 1000; +const JPEG_MIME_TYPE: &str = "image/jpeg"; +const PNG_MIME_TYPE: &str = "image/png"; +const MIN_READABLE_IMAGE_DIMENSION: u32 = 512; +const JPEG_QUALITIES: &[u8] = &[85, 75, 65, 55]; +const MAX_IMAGE_ENCODE_ATTEMPTS: usize = 8; /// Cap on individual image bytes when relaying a dropped/pasted image to a /// CLI agent (Claude Code etc.) via the system clipboard. CLI agents handle @@ -62,44 +71,187 @@ pub fn is_supported_image_mime_type(mime_type: &str) -> bool { SUPPORTED_IMAGE_MIME_TYPES.contains(&mime_type) } -/// Resizes an image if it exceeds the maximum pixel count, and ensures -/// resized outputs also respect the maximum dimension (width or height). +#[derive(Debug)] +pub struct ProcessedImage { + pub data: Vec, + pub mime_type: String, +} + +/// Resizes and re-encodes an image until it fits the model-facing byte budget. /// -/// Returns the original image bytes if the image is already within the -/// pixel limit; otherwise returns the resized image bytes in the original -/// format. -pub fn resize_image(image: &[u8]) -> Result, ImageError> { +/// The source bytes are never modified. Opaque images may be converted to JPEG +/// because that is the most reliable way to fit photo-like uploads under the +/// provider cap; images with transparency stay PNG so alpha is preserved. +pub fn prepare_image_for_agent(image: &[u8]) -> Result, ImageError> { let img = image::load_from_memory(image)?; + let original_format = image::guess_format(image)?; + let original_mime_type = mime_type_for_format(original_format); - let (current_width, current_height) = img.dimensions(); - let current_pixels = (current_width * current_height) as f64; - - if current_pixels <= MAX_IMAGE_PIXELS { - return Ok(image.to_vec()); + if image.len() <= MAX_IMAGE_SIZE_BYTES && image_within_dimension_limits(&img) { + return Ok(Some(ProcessedImage { + data: image.to_vec(), + mime_type: original_mime_type.to_string(), + })); } - let original_format = image::guess_format(image)?; + let has_transparency = image_has_transparency(&img); + let mut width = target_dimension(img.width(), img.height()).0; + let mut height = target_dimension(img.width(), img.height()).1; + let mut encode_attempts = 0; - let scale = (MAX_IMAGE_PIXELS / current_pixels).sqrt(); + while encode_attempts < MAX_IMAGE_ENCODE_ATTEMPTS { + let candidate = if width == img.width() && height == img.height() { + img.clone() + } else { + img.thumbnail(width, height) + }; - let mut new_width = current_width as f64 * scale; - let mut new_height = current_height as f64 * scale; + if has_transparency { + encode_attempts += 1; + let data = encode_png(&candidate)?; + if data.len() <= MAX_IMAGE_SIZE_BYTES { + return Ok(Some(ProcessedImage { + data, + mime_type: PNG_MIME_TYPE.to_string(), + })); + } + } else { + encode_attempts += 1; + let data = encode_png(&candidate)?; + if data.len() <= MAX_IMAGE_SIZE_BYTES { + return Ok(Some(ProcessedImage { + data, + mime_type: PNG_MIME_TYPE.to_string(), + })); + } - let scale_by_width = MAX_IMAGE_DIMENSION / new_width; - let scale_by_height = MAX_IMAGE_DIMENSION / new_height; - let scale = scale_by_width.min(scale_by_height).min(1.0); + for quality in JPEG_QUALITIES { + if encode_attempts >= MAX_IMAGE_ENCODE_ATTEMPTS { + break; + } + encode_attempts += 1; + let data = encode_jpeg(&candidate, *quality)?; + if data.len() <= MAX_IMAGE_SIZE_BYTES { + return Ok(Some(ProcessedImage { + data, + mime_type: JPEG_MIME_TYPE.to_string(), + })); + } + } + } - new_width *= scale; - new_height *= scale; + let next = next_smaller_dimensions(width, height); + if next == (width, height) { + break; + } + (width, height) = next; + } - let resized_img = img.thumbnail(new_width.round() as u32, new_height.round() as u32); + Ok(None) +} - let mut output_bytes: Vec = Vec::new(); - let mut writer = std::io::Cursor::new(&mut output_bytes); +/// Backwards-compatible helper for callers that only need processed bytes. +pub fn resize_image(image: &[u8]) -> Result, ImageError> { + prepare_image_for_agent(image)?.map_or_else( + || { + Err(ImageError::Limits(image::error::LimitError::from_kind( + image::error::LimitErrorKind::InsufficientMemory, + ))) + }, + |processed| Ok(processed.data), + ) +} - resized_img.write_to(&mut writer, original_format)?; +fn image_within_dimension_limits(img: &DynamicImage) -> bool { + let (width, height) = img.dimensions(); + let pixels = (width as f64) * (height as f64); + pixels <= MAX_IMAGE_PIXELS + && (width as f64) <= MAX_IMAGE_DIMENSION + && (height as f64) <= MAX_IMAGE_DIMENSION +} - Ok(output_bytes) +fn target_dimension(width: u32, height: u32) -> (u32, u32) { + let pixels = (width as f64) * (height as f64); + let pixel_scale = if pixels > MAX_IMAGE_PIXELS { + (MAX_IMAGE_PIXELS / pixels).sqrt() + } else { + 1.0 + }; + let dimension_scale = (MAX_IMAGE_DIMENSION / width as f64) + .min(MAX_IMAGE_DIMENSION / height as f64) + .min(1.0); + let scale = pixel_scale.min(dimension_scale).min(1.0); + scaled_dimensions(width, height, scale) +} + +fn next_smaller_dimensions(width: u32, height: u32) -> (u32, u32) { + if width <= MIN_READABLE_IMAGE_DIMENSION && height <= MIN_READABLE_IMAGE_DIMENSION { + return (width, height); + } + + let next_width = ((width as f64) * 0.85).round() as u32; + let next_height = ((height as f64) * 0.85).round() as u32; + ( + next_width.max(MIN_READABLE_IMAGE_DIMENSION.min(width)), + next_height.max(MIN_READABLE_IMAGE_DIMENSION.min(height)), + ) +} + +fn scaled_dimensions(width: u32, height: u32, scale: f64) -> (u32, u32) { + ( + ((width as f64) * scale).round().max(1.0) as u32, + ((height as f64) * scale).round().max(1.0) as u32, + ) +} + +fn image_has_transparency(img: &DynamicImage) -> bool { + if !img.color().has_alpha() { + return false; + } + img.to_rgba8().pixels().any(|pixel| pixel.0[3] < 255) +} + +fn encode_image_format(img: &DynamicImage, format: ImageFormat) -> Result, ImageError> { + let mut data = Vec::new(); + img.write_to(&mut Cursor::new(&mut data), format)?; + Ok(data) +} + +fn encode_png(img: &DynamicImage) -> Result, ImageError> { + encode_image_format(img, ImageFormat::Png) +} + +fn encode_jpeg(img: &DynamicImage, quality: u8) -> Result, ImageError> { + let rgb = img.to_rgb8(); + let mut data = Vec::new(); + JpegEncoder::new_with_quality(&mut data, quality).write_image( + rgb.as_raw(), + rgb.width(), + rgb.height(), + ExtendedColorType::Rgb8, + )?; + Ok(data) +} + +fn mime_type_for_format(format: ImageFormat) -> &'static str { + match format { + ImageFormat::Png => PNG_MIME_TYPE, + ImageFormat::Jpeg => JPEG_MIME_TYPE, + ImageFormat::Gif => "image/gif", + ImageFormat::WebP => "image/webp", + ImageFormat::Pnm => "image/x-portable-anymap", + ImageFormat::Tiff => "image/tiff", + ImageFormat::Tga => "image/x-tga", + ImageFormat::Dds => "image/vnd-ms.dds", + ImageFormat::Bmp => "image/bmp", + ImageFormat::Ico => "image/x-icon", + ImageFormat::Hdr => "image/vnd.radiance", + ImageFormat::OpenExr => "image/x-exr", + ImageFormat::Farbfeld => "image/farbfeld", + ImageFormat::Avif => "image/avif", + ImageFormat::Qoi => "image/qoi", + _ => "application/octet-stream", + } } /// Result of processing an image for agent mode. @@ -107,8 +259,10 @@ pub fn resize_image(image: &[u8]) -> Result, ImageError> { pub enum ProcessImageResult { /// Image was successfully processed and is within size limits. Success { - /// The processed image bytes (resized if needed). + /// The processed image bytes (resized/re-encoded if needed). data: Vec, + /// MIME type matching the processed image bytes. + mime_type: String, }, /// Image is too large even after resizing. TooLarge, @@ -120,16 +274,12 @@ pub enum ProcessImageResult { /// /// This applies the same processing that user-attached images go through. pub fn process_image_for_agent(image_data: &[u8]) -> ProcessImageResult { - match resize_image(image_data) { - Ok(resized_bytes) => { - if resized_bytes.len() > MAX_IMAGE_SIZE_BYTES { - ProcessImageResult::TooLarge - } else { - ProcessImageResult::Success { - data: resized_bytes, - } - } - } + match prepare_image_for_agent(image_data) { + Ok(Some(processed)) => ProcessImageResult::Success { + data: processed.data, + mime_type: processed.mime_type, + }, + Ok(None) => ProcessImageResult::TooLarge, Err(err) => ProcessImageResult::Error(err), } } diff --git a/app/src/util/image_tests.rs b/app/src/util/image_tests.rs index 8d3de972..7b20cc20 100644 --- a/app/src/util/image_tests.rs +++ b/app/src/util/image_tests.rs @@ -1,4 +1,4 @@ -use image::{ImageBuffer, Rgba}; +use image::{GenericImageView, ImageBuffer, Rgba}; use super::*; @@ -32,11 +32,12 @@ fn test_process_image_for_agent_small_image() { let result = process_image_for_agent(&small_png); match result { - ProcessImageResult::Success { data } => { + ProcessImageResult::Success { data, mime_type } => { // The processed image should be non-empty assert!(!data.is_empty()); // For a small image, it should not be resized, so data should be similar assert!(data.len() <= MAX_IMAGE_SIZE_BYTES); + assert_eq!(mime_type, "image/png"); } other => panic!("Expected Success, got {:?}", other), } @@ -52,6 +53,67 @@ fn test_resize_image_small_image_unchanged() { assert_eq!(result.len(), original_len); } +fn create_noisy_png(width: u32, height: u32, transparent: bool) -> Vec { + let img: ImageBuffer, Vec> = ImageBuffer::from_fn(width, height, |x, y| { + let seed = x + .wrapping_mul(1_103_515_245) + .wrapping_add(y.wrapping_mul(12_345)); + let alpha = if transparent && (x + y) % 3 == 0 { + 128 + } else { + 255 + }; + Rgba([ + (seed & 0xff) as u8, + ((seed >> 8) & 0xff) as u8, + ((seed >> 16) & 0xff) as u8, + alpha, + ]) + }); + + let mut bytes: Vec = Vec::new(); + let mut cursor = std::io::Cursor::new(&mut bytes); + img.write_to(&mut cursor, image::ImageFormat::Png).unwrap(); + bytes +} + +#[test] +fn test_process_image_for_agent_compresses_byte_heavy_opaque_image() { + let noisy_png = create_noisy_png(1500, 1500, false); + assert!(noisy_png.len() > MAX_IMAGE_SIZE_BYTES); + + let result = process_image_for_agent(&noisy_png); + + match result { + ProcessImageResult::Success { data, mime_type } => { + assert!(data.len() <= MAX_IMAGE_SIZE_BYTES); + assert_eq!(mime_type, "image/jpeg"); + let img = image::load_from_memory(&data).unwrap(); + let (width, height) = img.dimensions(); + assert!((width as f64 * height as f64) <= MAX_IMAGE_PIXELS); + assert_eq!(width, height); + } + other => panic!("Expected Success, got {:?}", other), + } +} + +#[test] +fn test_process_image_for_agent_preserves_transparency_as_png() { + let transparent_png = create_noisy_png(1000, 1000, true); + assert!(transparent_png.len() > MAX_IMAGE_SIZE_BYTES); + + let result = process_image_for_agent(&transparent_png); + + match result { + ProcessImageResult::Success { data, mime_type } => { + assert!(data.len() <= MAX_IMAGE_SIZE_BYTES); + assert_eq!(mime_type, "image/png"); + assert!(image::load_from_memory(&data).unwrap().color().has_alpha()); + } + other => panic!("Expected Success, got {:?}", other), + } +} + #[test] fn test_process_image_for_agent_invalid_data() { let invalid_data = vec![0u8; 100]; @@ -98,6 +160,24 @@ fn test_resize_image_large_image_gets_resized() { assert!(new_width < width || new_height < height); } +#[test] +fn test_process_image_for_agent_keeps_resized_opaque_png_when_it_fits() { + let large_png = create_large_test_png(); + + let result = process_image_for_agent(&large_png); + + match result { + ProcessImageResult::Success { data, mime_type } => { + assert!(data.len() <= MAX_IMAGE_SIZE_BYTES); + assert_eq!(mime_type, "image/png"); + let img = image::load_from_memory(&data).unwrap(); + let (width, height) = img.dimensions(); + assert!((width as f64 * height as f64) <= MAX_IMAGE_PIXELS); + } + other => panic!("Expected Success, got {:?}", other), + } +} + /// Creates a very tall/narrow image to test dimension clamping. fn create_tall_test_png() -> Vec { // Create a 100x20000 image (2M pixels, but very tall) diff --git a/crates/galaxy_agent_core/src/types.rs b/crates/galaxy_agent_core/src/types.rs index ad89c348..8f5c720d 100644 --- a/crates/galaxy_agent_core/src/types.rs +++ b/crates/galaxy_agent_core/src/types.rs @@ -4,6 +4,9 @@ use serde::{Deserialize, Serialize}; use serde_json::Value as JsonValue; pub const MAX_TOOL_RESULT_CHARS_FOR_PROVIDER_REQUEST: usize = 64_000; +const MIN_PROVIDER_OUTPUT_DISCLOSURE_CHARS: usize = 27; +const ESTIMATED_CHARS_PER_TOKEN: usize = 3; +const MAX_SINGLE_OUTPUT_CONTEXT_FRACTION: usize = 4; #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] pub struct ConversationMessage { @@ -13,8 +16,51 @@ pub struct ConversationMessage { impl ConversationMessage { pub fn truncate_tool_results_for_provider_request(&mut self) { - truncate_tool_results_in_content(&mut self.content); + self.truncate_tool_results_for_provider_request_with_context(None); } + + pub fn truncate_tool_results_for_provider_request_with_context( + &mut self, + max_context_tokens: Option, + ) { + truncate_tool_results_in_content( + &mut self.content, + provider_output_char_limit(max_context_tokens), + ); + } + + pub fn truncate_tagged_terminal_output_for_provider_request_with_context( + &mut self, + max_context_tokens: Option, + ) { + truncate_tagged_terminal_output_in_content(&mut self.content, max_context_tokens); + } +} + +pub fn provider_output_char_limit(max_context_tokens: Option) -> usize { + let context_aware_limit = max_context_tokens + .and_then(|tokens| usize::try_from(tokens).ok()) + .map(|tokens| { + tokens + .saturating_mul(ESTIMATED_CHARS_PER_TOKEN) + .saturating_div(MAX_SINGLE_OUTPUT_CONTEXT_FRACTION) + }) + .unwrap_or(MAX_TOOL_RESULT_CHARS_FOR_PROVIDER_REQUEST); + context_aware_limit.clamp( + MIN_PROVIDER_OUTPUT_DISCLOSURE_CHARS, + MAX_TOOL_RESULT_CHARS_FOR_PROVIDER_REQUEST, + ) +} + +pub fn bounded_terminal_output_for_provider_request( + output: &str, + max_context_tokens: Option, +) -> String { + truncate_provider_text( + output, + provider_output_char_limit(max_context_tokens), + "terminal output truncated for model context; full output remains available in the terminal", + ) } #[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] @@ -323,15 +369,77 @@ pub enum AgentEvent { }, } -fn truncate_tool_results_in_content(content: &mut MessageContent) { +fn truncate_tagged_terminal_output_in_content( + content: &mut MessageContent, + max_context_tokens: Option, +) { + match content { + MessageContent::Text(text) => { + *text = truncate_tagged_terminal_output(text, max_context_tokens); + } + MessageContent::MultiPart(parts) => { + for part in parts { + if let ContentPart::Text(text) = part { + *text = truncate_tagged_terminal_output(text, max_context_tokens); + } + } + } + MessageContent::ToolUse { .. } | MessageContent::ToolResult { .. } => {} + } +} + +fn truncate_tagged_terminal_output(content: &str, max_context_tokens: Option) -> String { + const TAGS: [&str; 2] = ["[Terminal output:\n", "[Final terminal output:\n"]; + const END_TAG: &str = "\n]\n"; + + let mut result = String::with_capacity(content.len()); + let mut remaining = content; + loop { + let Some((tag_offset, tag)) = TAGS + .iter() + .filter_map(|tag| remaining.find(tag).map(|offset| (offset, *tag))) + .min_by_key(|(offset, _)| *offset) + else { + result.push_str(remaining); + break; + }; + let output_start = tag_offset + tag.len(); + result.push_str(&remaining[..output_start]); + let after_tag = &remaining[output_start..]; + let mut output_ends = after_tag.match_indices(END_TAG); + let Some((output_end, _)) = output_ends.next() else { + result.push_str(&bounded_terminal_output_for_provider_request( + after_tag, + max_context_tokens, + )); + break; + }; + if output_ends.next().is_some() { + return truncate_provider_text( + content, + provider_output_char_limit(max_context_tokens), + "message text truncated for model context; full content remains in conversation history", + ); + } + result.push_str(&bounded_terminal_output_for_provider_request( + &after_tag[..output_end], + max_context_tokens, + )); + result.push_str(END_TAG); + remaining = &after_tag[output_end + END_TAG.len()..]; + } + result +} + +fn truncate_tool_results_in_content(content: &mut MessageContent, max_chars: usize) { match content { MessageContent::Text(_) | MessageContent::ToolUse { .. } => {} - MessageContent::ToolResult { content, .. } => truncate_tool_result_text(content), + MessageContent::ToolResult { content, .. } => truncate_tool_result_text(content, max_chars), MessageContent::MultiPart(parts) => { for part in parts { match part { ContentPart::ToolResult { content, .. } => { - truncate_tool_result_text(content); + truncate_tool_result_text(content, max_chars); } ContentPart::Text(_) | ContentPart::Reasoning { .. } @@ -343,16 +451,31 @@ fn truncate_tool_results_in_content(content: &mut MessageContent) { } } -fn truncate_tool_result_text(content: &mut String) { +fn truncate_tool_result_text(content: &mut String, max_chars: usize) { + *content = truncate_provider_text(content, max_chars, "tool result truncated"); +} + +fn truncate_provider_text(content: &str, max_chars: usize, reason: &str) -> String { let char_count = content.chars().count(); - if char_count <= MAX_TOOL_RESULT_CHARS_FOR_PROVIDER_REQUEST { - return; + if char_count <= max_chars { + return content.to_string(); } - let omitted_chars = char_count.saturating_sub(MAX_TOOL_RESULT_CHARS_FOR_PROVIDER_REQUEST); - let marker = format!("\n... [tool result truncated; omitted {omitted_chars} chars] ...\n"); + let omitted_chars = char_count.saturating_sub(max_chars); + let marker = format!("\n... [{reason}; omitted {omitted_chars} chars] ...\n"); let marker_chars = marker.chars().count(); - let retained_chars = MAX_TOOL_RESULT_CHARS_FOR_PROVIDER_REQUEST.saturating_sub(marker_chars); + if max_chars <= marker_chars { + let short_reason = if reason.starts_with("terminal output") { + "terminal output truncated" + } else if reason.starts_with("message text") { + "message text truncated" + } else { + "tool result truncated" + }; + let short_marker = format!("[{short_reason}]"); + return short_marker.chars().take(max_chars).collect(); + } + let retained_chars = max_chars.saturating_sub(marker_chars); let head_chars = retained_chars / 2; let tail_chars = retained_chars.saturating_sub(head_chars); let head: String = content.chars().take(head_chars).collect(); @@ -364,7 +487,7 @@ fn truncate_tool_result_text(content: &mut String) { .chars() .rev() .collect(); - *content = format!("{head}{marker}{tail}"); + format!("{head}{marker}{tail}") } #[cfg(test)] diff --git a/crates/galaxy_agent_core/src/types_tests.rs b/crates/galaxy_agent_core/src/types_tests.rs index 6984c060..54e88198 100644 --- a/crates/galaxy_agent_core/src/types_tests.rs +++ b/crates/galaxy_agent_core/src/types_tests.rs @@ -19,12 +19,116 @@ fn truncates_large_tool_results_for_provider_request() { let MessageContent::ToolResult { content, .. } = message.content else { panic!("expected tool result"); }; - assert!(content.len() <= MAX_TOOL_RESULT_CHARS_FOR_PROVIDER_REQUEST + 128); + assert!(content.chars().count() <= MAX_TOOL_RESULT_CHARS_FOR_PROVIDER_REQUEST); assert!(content.starts_with(prefix)); assert!(content.ends_with(suffix)); assert!(content.contains("tool result truncated")); } +#[test] +fn context_aware_tool_result_limit_preserves_the_original_message() { + let output = format!("start:{}:end", "x".repeat(20_000)); + let mut message = ConversationMessage { + role: MessageRole::User, + content: MessageContent::ToolResult { + tool_use_id: "toolu_1".to_string(), + content: output.clone(), + is_error: false, + }, + }; + + message.truncate_tool_results_for_provider_request_with_context(Some(8_000)); + + let MessageContent::ToolResult { content, .. } = message.content else { + panic!("expected tool result"); + }; + assert!(content.chars().count() <= provider_output_char_limit(Some(8_000))); + assert!(content.contains("tool result truncated")); + assert!(content.starts_with("start:")); + assert!(content.ends_with(":end")); + assert_eq!(output.chars().count(), 20_010); +} + +#[test] +fn terminal_output_bound_discloses_that_full_output_remains_in_terminal() { + let output = format!("first-line\n{}\nlast-line", "x".repeat(20_000)); + + let bounded = bounded_terminal_output_for_provider_request(&output, Some(8_000)); + + assert!(bounded.chars().count() <= provider_output_char_limit(Some(8_000))); + assert!(bounded.contains("terminal output truncated for model context")); + assert!(bounded.contains("full output remains available in the terminal")); + assert!(bounded.starts_with("first-line")); + assert!(bounded.ends_with("last-line")); + assert_eq!(output.chars().count(), 20_021); +} + +#[test] +fn terminal_output_bound_respects_character_budget_for_multibyte_text() { + let output = format!("first:{}:last", "界".repeat(20_000)); + + let bounded = bounded_terminal_output_for_provider_request(&output, Some(8_000)); + + assert!(bounded.chars().count() <= provider_output_char_limit(Some(8_000))); + assert!(bounded.starts_with("first:")); + assert!(bounded.ends_with(":last")); + assert!(bounded.contains("terminal output truncated for model context")); +} + +#[test] +fn tagged_terminal_output_is_bounded_without_mutating_surrounding_text() { + let output = format!("first:{}:last", "x".repeat(20_000)); + let mut message = ConversationMessage { + role: MessageRole::User, + content: MessageContent::Text(format!( + "[Running command: verbose]\n[Terminal output:\n{output}\n]\nContinue the task" + )), + }; + + message.truncate_tagged_terminal_output_for_provider_request_with_context(Some(8_000)); + let once = message.clone(); + message.truncate_tagged_terminal_output_for_provider_request_with_context(Some(8_000)); + + assert_eq!(message, once); + let MessageContent::Text(content) = message.content else { + panic!("expected text message"); + }; + assert!(content.starts_with("[Running command: verbose]")); + assert!(content.ends_with("Continue the task")); + assert!(content.contains("terminal output truncated for model context")); + assert!(content.contains("first:")); + assert!(content.contains(":last")); + assert!(!content.contains(&output)); +} + +#[test] +fn ambiguous_tagged_terminal_output_bounds_the_whole_provider_message() { + let output = format!("first\n]\n{}\n]\nlast", "x".repeat(20_000)); + let mut message = ConversationMessage { + role: MessageRole::User, + content: MessageContent::Text(format!("[Terminal output:\n{output}\n]\nContinue the task")), + }; + + message.truncate_tagged_terminal_output_for_provider_request_with_context(Some(8_000)); + + let MessageContent::Text(content) = message.content else { + panic!("expected text message"); + }; + assert!(content.chars().count() <= provider_output_char_limit(Some(8_000))); + assert!(content.contains("message text truncated for model context")); + assert!(content.contains("full content remains in conversation history")); + assert!(!content.contains(&output)); +} + +#[test] +fn terminal_output_bound_keeps_a_disclosure_for_tiny_budgets() { + let bounded = bounded_terminal_output_for_provider_request(&"x".repeat(200), Some(4)); + + assert_eq!(bounded, "[terminal output truncated]"); + assert_eq!(bounded.chars().count(), provider_output_char_limit(Some(4))); + assert!(!bounded.contains(&"x".repeat(200))); +} + #[test] fn usage_total_combines_disjoint_cache_and_miss_buckets() { let usage = Usage {