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

Bound agent context and preprocess oversized image attachments

Browse Source
This commit is contained in:
Ryan Ward
2026-09-24 21:54:20 -05:00
parent 6bf81519dc
commit 2a89d03866
15 changed files with 1606 additions and 283 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/src/ai/blocklist/action_model/execute.rs
+15 -2
View File
@@ -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);
app/src/ai/blocklist/controller.rs
+384 -119
View File
@@ -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::<String>();
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("</messages-to-summarize>");
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::<String>();
let tail = content
.chars()
.rev()
.take(tail_chars)
.collect::<String>()
.chars()
.rev()
.collect::<String>();
format!("{head}{marker}{tail}")
}
fn progressive_summary_messages(summary: String) -> Vec<ConversationMessage> {
@@ -896,10 +917,46 @@ struct RestoredProviderCommandEvidence {
cli_task_id: Option<TaskId>,
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<u32>) -> 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<PendingProviderMonitorObservation>,
pending_command_completion: Option<PendingProviderCommandCompletion>,
monitor_prose_continuations: usize,
live_command_output_tokens: u32,
command_output_poll_scheduled: bool,
retry_status: Option<ProviderRetryStatus>,
tool_call_progress: Vec<ProviderToolCallProgressStatus>,
}
@@ -1215,6 +1274,8 @@ struct ActiveProviderRunSnapshot {
pending_command_completion: Option<PendingProviderCommandCompletion>,
monitor_prose_continuations: usize,
#[serde(default)]
live_command_output_tokens: u32,
#[serde(default)]
queued_follow_ups: Vec<QueuedProviderRunSnapshot>,
}
@@ -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<ProgressiveSummaryCandidate> {
@@ -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<Self>,
) -> 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::<Vec<_>>();
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<ProviderCommandOutputSnapshot> {
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<Self>,
) {
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<Self>,
) {
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("<prior-summary>\n");
summarize_content.push_str(prior);
summarize_content.push_str("\n</prior-summary>\n\n");
}
summarize_content.push_str("<messages-to-summarize>\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::<String>();
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::<Vec<_>>()
.join("\n"),
};
summarize_content.push_str(&format!("[{role_str}]: {content_str}\n"));
}
summarize_content.push_str("</messages-to-summarize>");
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)
app/src/ai/blocklist/controller_tests.rs
+250 -35
View File
@@ -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("<messages-to-summarize>"));
assert!(content.ends_with("new-end\n</messages-to-summarize>"));
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![
app/src/ai/blocklist/permissions_tests.rs
+55
View File
@@ -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 {
app/src/ai/runtime/provider_run_coordinator.rs
+107 -29
View File
@@ -163,11 +163,21 @@ impl From<ProviderRunProtocolError> for ProviderRunCoordinatorError {
pub(crate) struct ProviderRunProfile {
pub(crate) runtime: Arc<dyn AgentRuntime>,
pub(crate) request: TurnRequest,
pub(crate) max_context_tokens: Option<u32>,
}
impl ProviderRunProfile {
pub(crate) fn new(runtime: Arc<dyn AgentRuntime>, 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<u32>) -> 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<galaxy_agent_core::ProviderRunId>,
runtime: Arc<dyn AgentRuntime>,
request: TurnRequest,
base_profile: ProviderRunProfile,
tool_result_archive: Vec<galaxy_agent_core::ConversationMessage>,
limits: ProviderRunLimits,
) -> Result<Self, ProviderRunCoordinatorError> {
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<String>,
runtime: Arc<dyn AgentRuntime>,
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<u32>,
) -> 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::<usize>();
bytes += request
text_bytes += request
.tools
.iter()
.map(|tool| serde_json::to_string(tool).unwrap_or_default().len())
.sum::<usize>();
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::<u64>(),
)
.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),
app/src/ai/runtime/provider_run_coordinator_tests.rs
+165 -6
View File
@@ -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<dyn AgentRuntime>) -> 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()]));
app/src/ai/runtime/rig.rs
+8 -8
View File
@@ -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,
app/src/ai/runtime/rig_request.rs
+59 -17
View File
@@ -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<u32>,
pub task_id: String,
pub needs_create_task: bool,
pub user_query: Option<String>,
@@ -135,9 +137,12 @@ pub(crate) fn prepare_rig_turn(
supported_tools: Vec<ToolType>,
supported_cli_agent_tools: Vec<ToolType>,
) -> 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<ToolType>,
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<ToolType>,
mode: Option<RigRequestMode>,
) -> 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<u32>,
configured_limit: Option<u32>,
) -> Option<u32> {
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<String>,
max_output_tokens: Option<u64>,
max_context_tokens: Option<u32>,
sanitizer: RigRequestSanitizer,
params: RequestParams,
supported_tools: Vec<ToolType>,
@@ -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<AIAgentInput>,
tool_results: Vec<ToolResult>,
max_context_tokens: Option<u32>,
) -> Vec<ConversationMessage> {
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<ConversationMessage> {
fn input_message(
input: AIAgentInput,
max_context_tokens: Option<u32>,
) -> Option<ConversationMessage> {
let (text, images) = match input {
AIAgentInput::UserQuery {
query,
@@ -388,9 +423,13 @@ fn input_message(input: AIAgentInput) -> Option<ConversationMessage> {
..
} => {
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<ConversationMessage> {
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 { .. } => (
app/src/ai/runtime/rig_request_tests.rs
+33 -1
View File
@@ -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 {