Align context usage display with provider request budgets

This commit is contained in:
Ryan Ward
2026-10-30 01:49:32 +00:00
parent f399ada2a3
commit 06c25af731
10 changed files with 823 additions and 80 deletions
+27
View File
@@ -60,6 +60,7 @@ use crate::ai::blocklist::{
SerializedBlockListItem, SerializedBlockListItem,
}; };
use crate::ai::llms::LLMPreferences; use crate::ai::llms::LLMPreferences;
use crate::ai::runtime::ProviderContextUsage;
use crate::ai::skills::SkillDescriptor; use crate::ai::skills::SkillDescriptor;
use crate::code_review::CodeReviewTelemetryEvent; use crate::code_review::CodeReviewTelemetryEvent;
use crate::notebooks::NotebookId; use crate::notebooks::NotebookId;
@@ -341,6 +342,9 @@ pub struct AIConversation {
progressive_summary: Option<String>, progressive_summary: Option<String>,
messages_summarized_up_to: usize, messages_summarized_up_to: usize,
current_context_tokens: u32, current_context_tokens: u32,
/// The bounded request projection used by direct-provider context budgeting.
/// Recomputed at committed model/tool boundaries rather than restored as live usage.
provider_context_usage: Option<ProviderContextUsage>,
latest_model_call_usage: ProviderUsage, latest_model_call_usage: ProviderUsage,
has_pending_progressive_summary: bool, has_pending_progressive_summary: bool,
} }
@@ -419,6 +423,7 @@ impl AIConversation {
progressive_summary: None, progressive_summary: None,
messages_summarized_up_to: 0, messages_summarized_up_to: 0,
current_context_tokens: 0, current_context_tokens: 0,
provider_context_usage: None,
latest_model_call_usage: ProviderUsage::default(), latest_model_call_usage: ProviderUsage::default(),
has_pending_progressive_summary: false, has_pending_progressive_summary: false,
} }
@@ -732,6 +737,7 @@ impl AIConversation {
// The previous response stream no longer exists; its drain count cannot apply. // The previous response stream no longer exists; its drain count cannot apply.
messages_summarized_up_to: 0, messages_summarized_up_to: 0,
current_context_tokens: 0, current_context_tokens: 0,
provider_context_usage: None,
latest_model_call_usage: ProviderUsage::default(), latest_model_call_usage: ProviderUsage::default(),
has_pending_progressive_summary: false, has_pending_progressive_summary: false,
}) })
@@ -786,6 +792,9 @@ impl AIConversation {
if !can_change_backend && self.agent_backend != agent_backend { if !can_change_backend && self.agent_backend != agent_backend {
return false; return false;
} }
if self.agent_backend != agent_backend {
self.provider_context_usage = None;
}
self.agent_backend = agent_backend; self.agent_backend = agent_backend;
true true
} }
@@ -813,6 +822,24 @@ impl AIConversation {
self.current_context_tokens = val; self.current_context_tokens = val;
} }
pub(crate) fn provider_context_usage(&self) -> Option<&ProviderContextUsage> {
self.provider_context_usage.as_ref()
}
pub(crate) fn set_provider_context_usage(&mut self, usage: Option<ProviderContextUsage>) {
self.provider_context_usage = usage;
}
/// Provider-reported input for the last model call, separate from the next-request estimate.
pub(crate) fn latest_model_call_reported_context_tokens(&self) -> Option<u32> {
let tokens = self
.latest_model_call_usage
.input_tokens
.saturating_add(self.latest_model_call_usage.cached_input_tokens)
.saturating_add(self.latest_model_call_usage.cache_creation_input_tokens);
(tokens > 0).then(|| u32::try_from(tokens).unwrap_or(u32::MAX))
}
pub fn set_latest_model_call_usage(&mut self, usage: ProviderUsage) { pub fn set_latest_model_call_usage(&mut self, usage: ProviderUsage) {
self.current_context_tokens = u32::try_from( self.current_context_tokens = u32::try_from(
usage usage
@@ -52,6 +52,7 @@ pub(crate) use self::environment_selector::{
use crate::ai::blocklist::agent_view::is_in_cloud_context; use crate::ai::blocklist::agent_view::is_in_cloud_context;
use crate::ai::blocklist::history_model::{BlocklistAIHistoryEvent, BlocklistAIHistoryModel}; use crate::ai::blocklist::history_model::{BlocklistAIHistoryEvent, BlocklistAIHistoryModel};
use crate::ai::blocklist::prompt::prompt_alert::{PromptAlertEvent, PromptAlertView}; use crate::ai::blocklist::prompt::prompt_alert::{PromptAlertEvent, PromptAlertView};
use crate::ai::blocklist::usage::conversation_usage_view::provider_context_usage_tooltip;
use crate::ai::blocklist::usage::icon_for_context_window_usage; use crate::ai::blocklist::usage::icon_for_context_window_usage;
use crate::ai::blocklist::BlocklistAIInputModel; use crate::ai::blocklist::BlocklistAIInputModel;
use crate::ai::execution_profiles::profiles::AIExecutionProfilesModel; use crate::ai::execution_profiles::profiles::AIExecutionProfilesModel;
@@ -802,6 +803,18 @@ impl AgentInputFooter {
} }
match event { match event {
BlocklistAIHistoryEvent::ConversationUsageMetadataUpdated {
conversation_id,
} => {
if BlocklistAIHistoryModel::as_ref(ctx)
.active_conversation(me.terminal_view_id)
.is_some_and(|conversation| conversation.id() == *conversation_id)
{
me.update_context_window_button(ctx);
me.update_llm_usage_buttons(ctx);
ctx.notify();
}
}
BlocklistAIHistoryEvent::StartedNewConversation { .. } BlocklistAIHistoryEvent::StartedNewConversation { .. }
| BlocklistAIHistoryEvent::SetActiveConversation { .. } | BlocklistAIHistoryEvent::SetActiveConversation { .. }
| BlocklistAIHistoryEvent::ClearedActiveConversation { .. } | BlocklistAIHistoryEvent::ClearedActiveConversation { .. }
@@ -1985,7 +1998,11 @@ impl AgentInputFooter {
if let Some(conversation) = if let Some(conversation) =
BlocklistAIHistoryModel::as_ref(ctx).active_conversation(self.terminal_view_id) BlocklistAIHistoryModel::as_ref(ctx).active_conversation(self.terminal_view_id)
{ {
let usage = conversation.context_window_usage().clamp(0.0, 1.0); let provider_context = conversation.provider_context_usage();
let usage = provider_context
.map(|usage| usage.compaction_usage())
.unwrap_or_else(|| conversation.context_window_usage())
.clamp(0.0, 1.0);
let icon = icon_for_context_window_usage(usage); let icon = icon_for_context_window_usage(usage);
let remaining_pct = ((1.0 - usage) * 100.0).round() as i32; let remaining_pct = ((1.0 - usage) * 100.0).round() as i32;
@@ -1995,7 +2012,11 @@ impl AgentInputFooter {
}); });
let is_cache_expired = FeatureFlag::PromptCacheExpiryWarning.is_enabled() let is_cache_expired = FeatureFlag::PromptCacheExpiryWarning.is_enabled()
&& expiry.is_some_and(|expiry| expiry <= Local::now()); && expiry.is_some_and(|expiry| expiry <= Local::now());
let context_remaining_tooltip = format!("{remaining_pct}% context remaining"); let context_remaining_tooltip = if provider_context.is_some() {
format!("~{remaining_pct}% compaction history budget remaining")
} else {
format!("{remaining_pct}% context remaining")
};
let tooltip = if is_cache_expired { let tooltip = if is_cache_expired {
format!("{context_remaining_tooltip} · prompt cache expired") format!("{context_remaining_tooltip} · prompt cache expired")
} else { } else {
@@ -2022,10 +2043,20 @@ impl AgentInputFooter {
conversation.latest_model_call_cache_read_tokens(), conversation.latest_model_call_cache_read_tokens(),
conversation.latest_model_call_cache_write_tokens(), conversation.latest_model_call_cache_write_tokens(),
conversation.latest_model_call_cache_miss_tokens(), conversation.latest_model_call_cache_miss_tokens(),
conversation.provider_context_usage().cloned(),
conversation.latest_model_call_reported_context_tokens(),
) )
}); });
let Some((current_context, reported_usage, cache_read, cache_write, cache_miss)) = metrics let Some((
current_context,
reported_usage,
cache_read,
cache_write,
cache_miss,
provider_context,
last_reported_input,
)) = metrics
else { else {
self.llm_context_usage_button.update(ctx, |button, ctx| { self.llm_context_usage_button.update(ctx, |button, ctx| {
button.set_label("Ctx —", ctx); button.set_label("Ctx —", ctx);
@@ -2051,29 +2082,48 @@ impl AgentInputFooter {
.active_profile(Some(self.terminal_view_id), ctx) .active_profile(Some(self.terminal_view_id), ctx)
.data() .data()
.context_window_display_value_for_model(active_model, ctx); .context_window_display_value_for_model(active_model, ctx);
let model_max_context = active_model let model_context = [
.context_window active_model.context_window.default_max,
.default_max active_model.context_window.max,
.max(active_model.context_window.max); ]
let max_context = profile_context .into_iter()
.or((model_max_context > 0).then_some(model_max_context)) .find(|limit| *limit > 0);
.unwrap_or(200_000); let max_context = profile_context.or(model_context).unwrap_or(200_000);
let context_usage = if max_context > 0 && current_context > 0 { let context_usage = if let Some(usage) = &provider_context {
usage.compaction_usage().clamp(0.0, 1.0)
} else if max_context > 0 && current_context > 0 {
(current_context as f32 / max_context as f32).clamp(0.0, 1.0) (current_context as f32 / max_context as f32).clamp(0.0, 1.0)
} else { } else {
reported_usage.clamp(0.0, 1.0) reported_usage.clamp(0.0, 1.0)
}; };
let context_pct = context_usage * 100.0; let context_pct = provider_context
let context_label = format!( .as_ref()
.map(|usage| usage.compaction_usage())
.unwrap_or(context_usage)
* 100.0;
let (context_label, context_tooltip) = if let Some(usage) = &provider_context {
(
format!(
"Ctx budget ~{context_pct:.1}% · ~{}/{}",
format_llm_token_count(usage.compaction.history_tokens),
format_llm_token_count(usage.compaction.history_capacity_tokens),
),
provider_context_usage_tooltip(usage, last_reported_input),
)
} else {
(
format!(
"Ctx {context_pct:.1}% · {}/{}", "Ctx {context_pct:.1}% · {}/{}",
format_llm_token_count(current_context), format_llm_token_count(current_context),
format_llm_token_count(max_context), format_llm_token_count(max_context),
); ),
let context_tooltip = format!( format!(
"Current LLM context: {} of {} tokens ({context_pct:.1}%). This is not cumulative.", "Current LLM context: {} of {} tokens ({context_pct:.1}%). This is not cumulative.",
format_llm_token_count(current_context), format_llm_token_count(current_context),
format_llm_token_count(max_context), format_llm_token_count(max_context),
); ),
)
};
let context_icon = icon_for_context_window_usage(context_usage); let context_icon = icon_for_context_window_usage(context_usage);
self.llm_context_usage_button.update(ctx, |button, ctx| { self.llm_context_usage_button.update(ctx, |button, ctx| {
button.set_label(context_label, ctx); button.set_label(context_label, ctx);
+215 -20
View File
@@ -15,7 +15,7 @@ use std::collections::{BTreeMap, HashMap, HashSet, VecDeque};
#[cfg(not(target_family = "wasm"))] #[cfg(not(target_family = "wasm"))]
use std::path::PathBuf; use std::path::PathBuf;
use std::sync::Arc; use std::sync::Arc;
use std::time::Duration; use std::time::{Duration, Instant};
use ai::skills::SkillPathOrigin; use ai::skills::SkillPathOrigin;
use anyhow::anyhow; use anyhow::anyhow;
@@ -84,12 +84,13 @@ use crate::ai::provider::ProviderConfig;
#[cfg(not(target_family = "wasm"))] #[cfg(not(target_family = "wasm"))]
use crate::ai::remote_logging::{self, RemoteLogLevel, RemoteLogRecord}; use crate::ai::remote_logging::{self, RemoteLogLevel, RemoteLogRecord};
use crate::ai::runtime::{ use crate::ai::runtime::{
prepare_provider_run, provider_context_window_tokens, provider_runtime_for_request, estimated_context_message_tokens, prepare_provider_run, provider_context_window_tokens,
restore_provider_config, OrchestrationModelOption, PreparedProviderRun, ProviderActionContext, provider_runtime_for_request, restore_provider_config, OrchestrationModelOption,
ProviderRunBlock, ProviderRunCoordinator, ProviderRunProfile, ProviderRunProjection, PreparedProviderRun, ProviderActionContext, ProviderContextUsage, ProviderRunBlock,
ProviderRunCoordinator, ProviderRunProfile, ProviderRunProjection,
ProviderRunResponseProjector, ProviderToolExecutionRef, ProviderToolLifecycleOutcome, ProviderRunResponseProjector, ProviderToolExecutionRef, ProviderToolLifecycleOutcome,
RuntimeResponseConfig, BASE_PROVIDER_PROFILE, CLI_MONITOR_PROVIDER_PROFILE, RuntimeResponseConfig, BASE_PROVIDER_PROFILE, CLI_MONITOR_PROVIDER_PROFILE,
COMPUTER_USE_PROVIDER_PROFILE, COMPUTER_USE_PROVIDER_PROFILE, PROGRESSIVE_SUMMARY_TRIGGER_USAGE,
}; };
use crate::ai::AIRequestUsageModel; use crate::ai::AIRequestUsageModel;
use crate::cloud_object::model::persistence::CloudModel; use crate::cloud_object::model::persistence::CloudModel;
@@ -114,7 +115,6 @@ use crate::terminal::ShellLaunchData;
use crate::workspaces::update_manager::TeamUpdateManager; use crate::workspaces::update_manager::TeamUpdateManager;
use crate::workspaces::user_workspaces::UserWorkspaces; use crate::workspaces::user_workspaces::UserWorkspaces;
const PROGRESSIVE_SUMMARY_TRIGGER_USAGE: f32 = 0.70;
const PROGRESSIVE_SUMMARY_RETAINED_USAGE: f32 = 0.50; const PROGRESSIVE_SUMMARY_RETAINED_USAGE: f32 = 0.50;
const PROGRESSIVE_SUMMARY_MIN_RECENT_MESSAGES: usize = 1; const PROGRESSIVE_SUMMARY_MIN_RECENT_MESSAGES: usize = 1;
const PROGRESSIVE_SUMMARY_OUTPUT_TOKENS: u32 = 8_000; const PROGRESSIVE_SUMMARY_OUTPUT_TOKENS: u32 = 8_000;
@@ -148,6 +148,8 @@ struct ActiveProgressiveSummaryPlan {
candidates: Vec<ProgressiveSummaryCandidate>, candidates: Vec<ProgressiveSummaryCandidate>,
summarized_prefix: Vec<ConversationMessage>, summarized_prefix: Vec<ConversationMessage>,
summarize_messages: Vec<ConversationMessage>, summarize_messages: Vec<ConversationMessage>,
context_usage: Option<ProviderContextUsage>,
started_at: Instant,
} }
#[derive(Clone, PartialEq, Eq)] #[derive(Clone, PartialEq, Eq)]
@@ -265,6 +267,18 @@ fn content_contains_tool_result(content: &MessageContent) -> bool {
fn progressive_summary_split_point( fn progressive_summary_split_point(
messages: &[ConversationMessage], messages: &[ConversationMessage],
retained_token_budget: u32, retained_token_budget: u32,
) -> usize {
progressive_summary_split_point_with_estimator(
messages,
retained_token_budget,
estimated_message_tokens,
)
}
fn progressive_summary_split_point_with_estimator(
messages: &[ConversationMessage],
retained_token_budget: u32,
estimate_message: fn(&ConversationMessage) -> u32,
) -> usize { ) -> usize {
if messages.len() <= PROGRESSIVE_SUMMARY_MIN_RECENT_MESSAGES { if messages.len() <= PROGRESSIVE_SUMMARY_MIN_RECENT_MESSAGES {
return 0; return 0;
@@ -273,7 +287,7 @@ fn progressive_summary_split_point(
let mut retained_tokens = 0_u32; let mut retained_tokens = 0_u32;
let mut split_point = messages.len(); let mut split_point = messages.len();
for index in (0..messages.len()).rev() { for index in (0..messages.len()).rev() {
let message_tokens = estimated_message_tokens(&messages[index]); let message_tokens = estimate_message(&messages[index]);
let retained_count = messages.len() - index; let retained_count = messages.len() - index;
// Keep the latest message, then retain only as much history as fits. // Keep the latest message, then retain only as much history as fits.
// A fixed message-count floor can retain many oversized file reads. // A fixed message-count floor can retain many oversized file reads.
@@ -302,7 +316,12 @@ fn progressive_summary_split_point(
if messages if messages
.last() .last()
.is_some_and(|message| content_contains_tool_result(&message.content)) .is_some_and(|message| content_contains_tool_result(&message.content))
&& estimated_history_tokens(&messages[split_point..]) > retained_token_budget && messages[split_point..]
.iter()
.fold(0_u32, |total, message| {
total.saturating_add(estimate_message(message))
})
> retained_token_budget
{ {
return messages.len(); return messages.len();
} }
@@ -2046,6 +2065,7 @@ fn provider_llm_lifecycle(projection: &ProviderRunProjection) -> Option<Provider
runtime_id, runtime_id,
model_id, model_id,
retry_attempt, retry_attempt,
..
} => ProviderLlmLifecycle { } => ProviderLlmLifecycle {
phase: ProviderLlmLifecyclePhase::Requested, phase: ProviderLlmLifecyclePhase::Requested,
work_id: work_id.clone(), work_id: work_id.clone(),
@@ -2245,6 +2265,7 @@ fn provider_run_terminal_remote_log_record(
enum ProviderDriveMessage { enum ProviderDriveMessage {
Projection { Projection {
lifecycle: Box<Option<ProviderLlmLifecycle>>, lifecycle: Box<Option<ProviderLlmLifecycle>>,
context_usage: Option<ProviderContextUsage>,
latest_usage: Option<Usage>, latest_usage: Option<Usage>,
tool_call_progress: ProviderToolCallProgressUpdate, tool_call_progress: ProviderToolCallProgressUpdate,
events: Vec<warp_multi_agent_api::ResponseEvent>, events: Vec<warp_multi_agent_api::ResponseEvent>,
@@ -7137,6 +7158,73 @@ impl BlocklistAIController {
} }
} }
fn publish_provider_context_usage(
&self,
conversation_id: AIConversationId,
usage: Option<ProviderContextUsage>,
phase: &str,
ctx: &mut ModelContext<Self>,
) {
let changed = BlocklistAIHistoryModel::as_ref(ctx)
.conversation(&conversation_id)
.is_some_and(|conversation| conversation.provider_context_usage() != usage.as_ref());
if !changed {
return;
}
BlocklistAIHistoryModel::handle(ctx).update(ctx, |history, ctx| {
if let Some(conversation) = history.conversation_mut(&conversation_id) {
conversation.set_provider_context_usage(usage.clone());
ctx.emit(BlocklistAIHistoryEvent::ConversationUsageMetadataUpdated {
conversation_id,
});
}
});
#[cfg(not(target_family = "wasm"))]
if let Some(usage) = usage {
let slot = self.active_provider_runs.get(&conversation_id);
remote_logging::log_model_event(
ctx,
RemoteLogRecord {
level: RemoteLogLevel::Info,
message: "Provider context budget updated".to_owned(),
context: serde_json::json!({
"event": "provider_context_usage",
"conversation_id": conversation_id,
"provider_run_id": slot.map(|slot| slot.run_id.as_str()),
"stream_id": slot.map(|slot| slot.stream_id.as_str()),
"phase": phase,
"source": "projected_request_estimate",
"active": usage.active,
"compaction": usage.compaction,
"compaction_usage": usage.compaction_usage(),
"compaction_trigger_usage": PROGRESSIVE_SUMMARY_TRIGGER_USAGE,
"request_safety_usage": 0.90,
}),
},
);
}
}
fn refresh_active_provider_context_usage(
&self,
conversation_id: AIConversationId,
phase: &str,
ctx: &mut ModelContext<Self>,
) {
if let Some(run) = self
.active_provider_runs
.get(&conversation_id)
.and_then(|slot| slot.run.as_ref())
{
self.publish_provider_context_usage(
conversation_id,
run.coordinator.context_usage_snapshot(),
phase,
ctx,
);
}
}
fn active_provider_context_estimate( fn active_provider_context_estimate(
&self, &self,
conversation_id: AIConversationId, conversation_id: AIConversationId,
@@ -7166,9 +7254,7 @@ impl BlocklistAIController {
.map(Self::live_provider_command_context_tokens) .map(Self::live_provider_command_context_tokens)
.unwrap_or(0); .unwrap_or(0);
Some(( Some((
current_context_tokens estimated_context_tokens.saturating_add(live_command_output_tokens),
.max(estimated_context_tokens)
.saturating_add(live_command_output_tokens),
persistent_context_tokens, persistent_context_tokens,
)) ))
} }
@@ -7229,17 +7315,22 @@ impl BlocklistAIController {
.coordinator .coordinator
.progressive_summary_context_budget() .progressive_summary_context_budget()
.unwrap_or((active_context_limit, current_persistent_context_tokens)); .unwrap_or((active_context_limit, current_persistent_context_tokens));
let context_usage = active_run.coordinator.context_usage_snapshot();
// Trigger and retain against the space history can actually use. Fixed // Trigger and retain against the space history can actually use. Fixed
// prompts and a large output reservation cannot be reduced by summarizing. // prompts and a large output reservation cannot be reduced by summarizing.
let history_capacity = let history_capacity =
summary_context_limit.saturating_sub(summary_persistent_context_tokens); summary_context_limit.saturating_sub(summary_persistent_context_tokens);
let current_history_tokens = let current_history_tokens = context_usage.as_ref().map_or_else(
current_context_tokens.saturating_sub(current_persistent_context_tokens); || current_context_tokens.saturating_sub(current_persistent_context_tokens),
|usage| usage.compaction.history_tokens,
);
let compaction_usage = context_usage.as_ref().map_or_else(
|| current_history_tokens as f32 / history_capacity.max(1) as f32,
ProviderContextUsage::compaction_usage,
);
if summary_pending if summary_pending
|| last_failure_context_tokens == Some(current_context_tokens) || last_failure_context_tokens == Some(current_context_tokens)
|| (!force || (!force && compaction_usage < PROGRESSIVE_SUMMARY_TRIGGER_USAGE)
&& current_history_tokens as f32 / (history_capacity.max(1) as f32)
< PROGRESSIVE_SUMMARY_TRIGGER_USAGE)
{ {
return false; return false;
} }
@@ -7247,7 +7338,28 @@ impl BlocklistAIController {
(history_capacity as f32 * PROGRESSIVE_SUMMARY_RETAINED_USAGE) as u32; (history_capacity as f32 * PROGRESSIVE_SUMMARY_RETAINED_USAGE) as u32;
let retained_history_budget = let retained_history_budget =
retained_history_budget.saturating_sub(PROGRESSIVE_SUMMARY_OUTPUT_TOKENS); retained_history_budget.saturating_sub(PROGRESSIVE_SUMMARY_OUTPUT_TOKENS);
let split_point = progressive_summary_split_point(transcript, retained_history_budget); // Satisfy both profiles' own retention targets: per-profile output bounds
// can produce different history sizes even for the same durable transcript.
let split_point = [
active_run.coordinator.run().profile().as_str(),
CLI_MONITOR_PROVIDER_PROFILE,
]
.into_iter()
.filter_map(|profile| {
let usage = active_run.coordinator.context_profile_usage(profile)?;
let messages = active_run
.coordinator
.projected_transcript_for_context(profile)?;
let budget =
(usage.history_capacity_tokens as f32 * PROGRESSIVE_SUMMARY_RETAINED_USAGE) as u32;
Some(progressive_summary_split_point_with_estimator(
&messages,
budget.saturating_sub(PROGRESSIVE_SUMMARY_OUTPUT_TOKENS),
estimated_context_message_tokens,
))
})
.max()
.unwrap_or_else(|| progressive_summary_split_point(transcript, retained_history_budget));
if split_point == 0 { if split_point == 0 {
return false; return false;
} }
@@ -7307,6 +7419,8 @@ impl BlocklistAIController {
role: MessageRole::User, role: MessageRole::User,
content: MessageContent::Text(summarize_content), content: MessageContent::Text(summarize_content),
}], }],
context_usage,
started_at: Instant::now(),
}; };
let Some(slot) = self.active_provider_runs.get(&conversation_id) else { let Some(slot) = self.active_provider_runs.get(&conversation_id) else {
@@ -7349,6 +7463,26 @@ impl BlocklistAIController {
plan.retained_messages_count, plan.retained_messages_count,
candidate_names, candidate_names,
); );
#[cfg(not(target_family = "wasm"))]
remote_logging::log_model_event(
ctx,
RemoteLogRecord {
level: RemoteLogLevel::Info,
message: "Provider background compaction started".to_owned(),
context: serde_json::json!({
"event": "provider_compaction_started",
"conversation_id": conversation_id,
"provider_run_id": identity.run_id.as_str(),
"stream_id": identity.stream_id.as_str(),
"forced": force,
"context_usage": plan.context_usage,
"trigger_usage": PROGRESSIVE_SUMMARY_TRIGGER_USAGE,
"retained_usage": PROGRESSIVE_SUMMARY_RETAINED_USAGE,
"summarized_message_count": plan.split_point,
"retained_message_count": plan.retained_messages_count,
}),
},
);
ctx.spawn( ctx.spawn(
async move { async move {
let mut errors = Vec::new(); let mut errors = Vec::new();
@@ -7419,6 +7553,25 @@ impl BlocklistAIController {
return; return;
} }
#[cfg(not(target_family = "wasm"))]
remote_logging::log_model_event(
ctx,
RemoteLogRecord {
level: RemoteLogLevel::Info,
message: "Provider background compaction finished".to_owned(),
context: serde_json::json!({
"event": "provider_compaction_finished",
"conversation_id": conversation_id,
"provider_run_id": identity.run_id.as_str(),
"stream_id": identity.stream_id.as_str(),
"elapsed_ms": plan.started_at.elapsed().as_millis() as u64,
"succeeded": result.is_ok(),
"waiting_for_context": waiting_for_context,
"context_usage": plan.context_usage,
}),
},
);
self.active_provider_progressive_summaries.insert( self.active_provider_progressive_summaries.insert(
conversation_id, conversation_id,
ActiveProviderProgressiveSummaryState::Ready { ActiveProviderProgressiveSummaryState::Ready {
@@ -7630,6 +7783,7 @@ impl BlocklistAIController {
); );
return true; return true;
} }
self.refresh_active_provider_context_usage(conversation_id, "compaction_applied", ctx);
false false
} }
@@ -7807,6 +7961,7 @@ impl BlocklistAIController {
if self.apply_ready_active_provider_progressive_summary(conversation_id, ctx) { if self.apply_ready_active_provider_progressive_summary(conversation_id, ctx) {
return; return;
} }
self.refresh_active_provider_context_usage(conversation_id, "committed_boundary", ctx);
self.schedule_provider_command_output_poll(conversation_id, ctx); self.schedule_provider_command_output_poll(conversation_id, ctx);
// Ordinary background work never blocks the agent. Once a model has // Ordinary background work never blocks the agent. Once a model has
// exhausted its context, polling must not relaunch that same request // exhausted its context, polling must not relaunch that same request
@@ -7866,6 +8021,12 @@ impl BlocklistAIController {
turn_control, turn_control,
|projection| { |projection| {
let lifecycle = provider_llm_lifecycle(&projection); let lifecycle = provider_llm_lifecycle(&projection);
let context_usage = match &projection {
ProviderRunProjection::ModelTurnRequested {
context_usage, ..
} => context_usage.clone(),
_ => None,
};
let tool_call_progress = let tool_call_progress =
provider_tool_call_progress_update(&projection); provider_tool_call_progress_update(&projection);
let latest_usage = match &projection { let latest_usage = match &projection {
@@ -7888,6 +8049,7 @@ impl BlocklistAIController {
projection_sender projection_sender
.send(ProviderDriveMessage::Projection { .send(ProviderDriveMessage::Projection {
lifecycle: Box::new(lifecycle), lifecycle: Box::new(lifecycle),
context_usage,
latest_usage, latest_usage,
tool_call_progress, tool_call_progress,
events, events,
@@ -7951,20 +8113,32 @@ impl BlocklistAIController {
match message { match message {
ProviderDriveMessage::Projection { ProviderDriveMessage::Projection {
lifecycle, lifecycle,
context_usage,
latest_usage, latest_usage,
tool_call_progress, tool_call_progress,
events, events,
acknowledgement, acknowledgement,
} => { } => {
let response_stream = slot.response_stream.clone();
let did_input_contain_user_query = slot.did_input_contain_user_query;
if let Some(usage) = context_usage {
self.publish_provider_context_usage(
conversation_id,
Some(usage),
"model_requested",
ctx,
);
}
if let Some(usage) = latest_usage { if let Some(usage) = latest_usage {
BlocklistAIHistoryModel::handle(ctx).update(ctx, |history, _| { BlocklistAIHistoryModel::handle(ctx).update(ctx, |history, ctx| {
if let Some(conversation) = history.conversation_mut(&conversation_id) { if let Some(conversation) = history.conversation_mut(&conversation_id) {
conversation.set_latest_model_call_usage(usage); conversation.set_latest_model_call_usage(usage);
ctx.emit(BlocklistAIHistoryEvent::ConversationUsageMetadataUpdated {
conversation_id,
});
} }
}); });
} }
let response_stream = slot.response_stream.clone();
let did_input_contain_user_query = slot.did_input_contain_user_query;
let mut result = Ok(()); let mut result = Ok(());
for event in events { for event in events {
let event = ResponseStream::projected_event(event); let event = ResponseStream::projected_event(event);
@@ -8271,6 +8445,18 @@ impl BlocklistAIController {
result: Result<ProviderRunBlock, String>, result: Result<ProviderRunBlock, String>,
ctx: &mut ModelContext<Self>, ctx: &mut ModelContext<Self>,
) { ) {
if self
.active_provider_runs
.get(&conversation_id)
.is_some_and(|slot| slot.run_id == *run.coordinator.run().id())
{
self.publish_provider_context_usage(
conversation_id,
run.coordinator.context_usage_snapshot(),
"committed_result",
ctx,
);
}
let Some(slot) = self.active_provider_runs.get_mut(&conversation_id) else { let Some(slot) = self.active_provider_runs.get_mut(&conversation_id) else {
return; return;
}; };
@@ -9380,6 +9566,11 @@ impl BlocklistAIController {
let retained = transcript.split_off(plan.split_point); let retained = transcript.split_off(plan.split_point);
let retained_count = retained.len(); let retained_count = retained.len();
let mut projected_history = progressive_summary_messages(summary.clone());
projected_history.extend_from_slice(&retained);
let context_usage = conversation.provider_context_usage()
.or(plan.context_usage.as_ref())
.map(|usage| usage.with_transcript(&projected_history));
let new_context_tokens = plan let new_context_tokens = plan
.persistent_context_tokens .persistent_context_tokens
.saturating_add(summary_prefix_tokens) .saturating_add(summary_prefix_tokens)
@@ -9391,6 +9582,10 @@ impl BlocklistAIController {
conversation.set_progressive_summary(Some(summary), 0); conversation.set_progressive_summary(Some(summary), 0);
conversation.set_current_context_tokens(new_context_tokens); conversation.set_current_context_tokens(new_context_tokens);
conversation.set_context_window_usage(new_usage.clamp(0.0, 1.0)); conversation.set_context_window_usage(new_usage.clamp(0.0, 1.0));
conversation.set_provider_context_usage(context_usage);
ctx.emit(BlocklistAIHistoryEvent::ConversationUsageMetadataUpdated {
conversation_id,
});
conversation.write_updated_conversation_state(ctx); conversation.write_updated_conversation_state(ctx);
log::info!( log::info!(
"[progressive-summary] Applied background summary after active run completion {:?}: ~{} tokens ({:.1}% of {} context), {} tail messages preserved", "[progressive-summary] Applied background summary after active run completion {:?}: ~{} tokens ({:.1}% of {} context), {} tail messages preserved",
+3
View File
@@ -56,6 +56,7 @@ fn provider_lifecycle_logs_expose_llm_completion_and_sanitize_errors() {
retry_error.recoverable = true; retry_error.recoverable = true;
let projections = [ let projections = [
ProviderRunProjection::ModelTurnRequested { ProviderRunProjection::ModelTurnRequested {
context_usage: None,
work_id: work_id.clone(), work_id: work_id.clone(),
profile: "base".into(), profile: "base".into(),
runtime_id: "rig:openai".to_owned(), runtime_id: "rig:openai".to_owned(),
@@ -4301,6 +4302,8 @@ fn progressive_summary_test_plan(
candidates: Vec::new(), candidates: Vec::new(),
summarize_messages: summarized_prefix.clone(), summarize_messages: summarized_prefix.clone(),
summarized_prefix, summarized_prefix,
context_usage: None,
started_at: std::time::Instant::now(),
} }
} }
@@ -29,6 +29,7 @@ use crate::ai::blocklist::usage::rollup::{
}; };
use crate::ai::blocklist::view_util::{format_credits, format_token_count}; use crate::ai::blocklist::view_util::{format_credits, format_token_count};
use crate::ai::blocklist::{BlocklistAIHistoryEvent, BlocklistAIHistoryModel}; use crate::ai::blocklist::{BlocklistAIHistoryEvent, BlocklistAIHistoryModel};
use crate::ai::runtime::{ProviderContextUsage, PROGRESSIVE_SUMMARY_TRIGGER_USAGE};
use crate::appearance::Appearance; use crate::appearance::Appearance;
use crate::persistence::model::{ use crate::persistence::model::{
token_usage_category_display_name, ContextWindowSegment, ContextWindowSegmentType, token_usage_category_display_name, ContextWindowSegment, ContextWindowSegmentType,
@@ -61,6 +62,10 @@ pub struct ConversationUsageInfo {
pub commands_executed: i32, pub commands_executed: i32,
/// Live context window token count (from most recent Bedrock response). /// Live context window token count (from most recent Bedrock response).
pub current_context_tokens: u32, pub current_context_tokens: u32,
/// Local estimate for the projected direct-provider request and its compaction budget.
pub provider_context_usage: Option<ProviderContextUsage>,
/// Input tokens reported by the last completed provider call, separate from local estimates.
pub latest_model_call_reported_context_tokens: Option<u32>,
/// Cumulative cost across all requests. /// Cumulative cost across all requests.
pub estimated_cost_cents: f32, pub estimated_cost_cents: f32,
/// Cumulative cache read tokens (session total). /// Cumulative cache read tokens (session total).
@@ -71,6 +76,44 @@ pub struct ConversationUsageInfo {
pub total_input_tokens: u32, pub total_input_tokens: u32,
} }
pub(crate) fn provider_context_usage_tooltip(
usage: &ProviderContextUsage,
last_reported_input_tokens: Option<u32>,
) -> String {
let active = &usage.active;
let compaction = &usage.compaction;
let last_reported = last_reported_input_tokens
.map(|tokens| format!("{tokens} tokens"))
.unwrap_or_else(|| "not yet available".to_string());
let trigger_percent = PROGRESSIVE_SUMMARY_TRIGGER_USAGE * 100.0;
format!(
"Local estimate for the projected request; provider tokenization may differ. Updates at model/tool boundaries.\n\
Active request: ~{} input tokens of {} effective context tokens ({:.1}%).\n\
Active profile: {} ({}).\n\
Fixed prompts and tools: ~{} input tokens. Reserved output: {} tokens.\n\
Compaction budget: ~{} history tokens of {} usable history tokens ({:.1}%).\n\
Limiting profile: {} ({}), {} effective context tokens, ~{} fixed input tokens and {} reserved output tokens.\n\
Compaction starts at {trigger_percent:.0}% of usable history.\n\
Last provider-reported input: {}.",
active.estimated_input_tokens,
active.context_limit,
active.input_usage() * 100.0,
active.profile,
active.model_id,
active.fixed_input_tokens,
active.reserved_output_tokens,
compaction.history_tokens,
compaction.history_capacity_tokens,
usage.compaction_usage() * 100.0,
compaction.profile,
compaction.model_id,
compaction.context_limit,
compaction.fixed_input_tokens,
compaction.reserved_output_tokens,
last_reported,
)
}
/// Timing information for the last set of agent responses /// Timing information for the last set of agent responses
/// (all blocks since the last user input, as this is the granularity /// (all blocks since the last user input, as this is the granularity
/// at which we show the usage footer) /// at which we show the usage footer)
@@ -139,6 +182,29 @@ pub struct ConversationUsageView {
} }
impl ConversationUsageView { impl ConversationUsageView {
fn provider_context_usage(
&self,
app: &AppContext,
) -> Option<(ProviderContextUsage, Option<u32>)> {
if let Some(conversation) = self
.parent_conversation_id
.and_then(|id| BlocklistAIHistoryModel::as_ref(app).conversation(&id))
{
if let Some(usage) = conversation.provider_context_usage() {
return Some((
usage.clone(),
conversation.latest_model_call_reported_context_tokens(),
));
}
}
self.usage_info.provider_context_usage.clone().map(|usage| {
(
usage,
self.usage_info.latest_model_call_reported_context_tokens,
)
})
}
pub fn new( pub fn new(
usage_info: ConversationUsageInfo, usage_info: ConversationUsageInfo,
display_mode: DisplayMode, display_mode: DisplayMode,
@@ -316,8 +382,14 @@ impl ConversationUsageView {
let theme = appearance.theme(); let theme = appearance.theme();
let font_size = appearance.ui_font_size() + 2.; let font_size = appearance.ui_font_size() + 2.;
let text_color = blended_colors::text_main(theme, theme.surface_2()); let text_color = blended_colors::text_main(theme, theme.surface_2());
let context_window_breakdown_enabled = FeatureFlag::ContextWindowUsageBreakdown let provider_context_usage = self.provider_context_usage(app);
.is_enabled() let context_window_usage = provider_context_usage
.as_ref()
.map(|(usage, _)| usage.compaction_usage())
.unwrap_or(self.usage_info.context_window_usage)
.clamp(0.0, 1.0);
let context_window_breakdown_enabled = provider_context_usage.is_none()
&& FeatureFlag::ContextWindowUsageBreakdown.is_enabled()
&& !context_window_segment_display_rows( && !context_window_segment_display_rows(
self.usage_info.context_window_usage, self.usage_info.context_window_usage,
&self.usage_info.context_window_segments, &self.usage_info.context_window_segments,
@@ -478,8 +550,53 @@ impl ConversationUsageView {
); );
} }
// Context tokens (live state — current context window size) if let Some((usage, last_reported)) = &provider_context_usage {
if self.usage_info.current_context_tokens > 0 { let active = &usage.active;
let compaction = &usage.compaction;
for (label, value) in [
(
"Projected input (local estimate)",
format!("~{} tokens", active.estimated_input_tokens),
),
(
"Effective request context",
format!("{} tokens", active.context_limit),
),
(
"Fixed prompts and tools (estimate)",
format!("~{} tokens", active.fixed_input_tokens),
),
(
"Reserved output",
format!("{} tokens", active.reserved_output_tokens),
),
(
"Compaction history (estimate)",
format!(
"~{} / {} tokens",
compaction.history_tokens, compaction.history_capacity_tokens,
),
),
("Compaction profile", compaction.profile.clone()),
("Compaction model", compaction.model_id.clone()),
(
"Compaction trigger",
format!(
"{:.0}% of usable history",
PROGRESSIVE_SUMMARY_TRIGGER_USAGE * 100.0
),
),
(
"Last provider-reported input",
last_reported
.map(|tokens| format!("{tokens} tokens"))
.unwrap_or_else(|| "Not yet available".to_string()),
),
] {
labels.push(render_label_text(label, appearance));
values.push(render_value_text(value, appearance));
}
} else if self.usage_info.current_context_tokens > 0 {
labels.push(render_label_text("Context tokens", appearance)); labels.push(render_label_text("Context tokens", appearance));
values.push(render_value_text( values.push(render_value_text(
format_token_count(self.usage_info.current_context_tokens), format_token_count(self.usage_info.current_context_tokens),
@@ -521,11 +638,22 @@ impl ConversationUsageView {
} }
} }
labels.push(render_label_text("Context window used", appearance)); labels.push(render_label_text(
let context_window_usage = self.usage_info.context_window_usage.clamp(0.0, 1.0); if provider_context_usage.is_some() {
let context_usage_pct = context_window_usage * 100.; "History budget used (estimate)"
let context_usage_str = if context_window_breakdown_enabled && self.context_window_expanded } else {
{ "Context window used"
},
appearance,
));
let context_usage_pct = provider_context_usage
.as_ref()
.map(|(usage, _)| usage.compaction_usage())
.unwrap_or(context_window_usage)
* 100.;
let context_usage_str = if provider_context_usage.is_some() {
format!("~{context_usage_pct:.1}%")
} else if context_window_breakdown_enabled && self.context_window_expanded {
format!("{context_usage_pct:.2}%") format!("{context_usage_pct:.2}%")
} else { } else {
format!("{}%", context_usage_pct.round()) format!("{}%", context_usage_pct.round())
@@ -728,6 +728,7 @@ fn authentication_failures_identify_the_requested_transport_before_stream_start(
] { ] {
projector projector
.project(ProviderRunProjection::ModelTurnRequested { .project(ProviderRunProjection::ModelTurnRequested {
context_usage: None,
work_id: galaxy_agent_core::ExternalWorkId { work_id: galaxy_agent_core::ExternalWorkId {
run_id: galaxy_agent_core::ProviderRunId::new("run"), run_id: galaxy_agent_core::ProviderRunId::new("run"),
epoch: galaxy_agent_core::RunEpoch::new(1), epoch: galaxy_agent_core::RunEpoch::new(1),
+4 -3
View File
@@ -13,9 +13,10 @@ pub(crate) use event_translator::{
ProviderRunResponseProjector, RuntimeResponseConfig, RuntimeResponseTranslator, ProviderRunResponseProjector, RuntimeResponseConfig, RuntimeResponseTranslator,
}; };
pub(crate) use provider_run_coordinator::{ pub(crate) use provider_run_coordinator::{
ProviderRunBlock, ProviderRunCoordinator, ProviderRunProfile, ProviderRunProjection, estimated_context_message_tokens, ProviderContextUsage, ProviderRunBlock,
ProviderToolExecutionRef, ProviderToolLifecycleOutcome, BASE_PROVIDER_PROFILE, ProviderRunCoordinator, ProviderRunProfile, ProviderRunProjection, ProviderToolExecutionRef,
CLI_MONITOR_PROVIDER_PROFILE, COMPUTER_USE_PROVIDER_PROFILE, ProviderToolLifecycleOutcome, BASE_PROVIDER_PROFILE, CLI_MONITOR_PROVIDER_PROFILE,
COMPUTER_USE_PROVIDER_PROFILE, PROGRESSIVE_SUMMARY_TRIGGER_USAGE,
}; };
pub(crate) use rig::{ pub(crate) use rig::{
prepare_provider_run, provider_context_window_tokens, provider_runtime_for_request, prepare_provider_run, provider_context_window_tokens, provider_runtime_for_request,
+215 -25
View File
@@ -25,6 +25,7 @@ use crate::ai::model_output_logging;
pub(crate) const BASE_PROVIDER_PROFILE: &str = "base"; pub(crate) const BASE_PROVIDER_PROFILE: &str = "base";
pub(crate) const CLI_MONITOR_PROVIDER_PROFILE: &str = "cli-monitor"; pub(crate) const CLI_MONITOR_PROVIDER_PROFILE: &str = "cli-monitor";
pub(crate) const COMPUTER_USE_PROVIDER_PROFILE: &str = "computer-use"; pub(crate) const COMPUTER_USE_PROVIDER_PROFILE: &str = "computer-use";
pub(crate) const PROGRESSIVE_SUMMARY_TRIGGER_USAGE: f32 = 0.70;
const PROVIDER_MODEL_START_TIMEOUT: Duration = Duration::from_secs(120); const PROVIDER_MODEL_START_TIMEOUT: Duration = Duration::from_secs(120);
const PROVIDER_MODEL_EVENT_IDLE_TIMEOUT: Duration = Duration::from_secs(90); const PROVIDER_MODEL_EVENT_IDLE_TIMEOUT: Duration = Duration::from_secs(90);
const PROVIDER_MODEL_RETRY_DELAYS: [Duration; 3] = [ const PROVIDER_MODEL_RETRY_DELAYS: [Duration; 3] = [
@@ -33,6 +34,99 @@ const PROVIDER_MODEL_RETRY_DELAYS: [Duration; 3] = [
Duration::from_secs(5), Duration::from_secs(5),
]; ];
/// Estimated occupancy of the same bounded request projection sent to a provider.
/// The durable transcript remains complete; this snapshot contains no message content.
#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub(crate) struct ProviderContextProfileUsage {
pub(crate) profile: String,
pub(crate) model_id: String,
pub(crate) context_limit: u32,
pub(crate) estimated_input_tokens: u32,
pub(crate) fixed_input_tokens: u32,
pub(crate) reserved_output_tokens: u32,
pub(crate) history_tokens: u32,
pub(crate) history_capacity_tokens: u32,
}
impl ProviderContextProfileUsage {
pub(crate) fn input_usage(&self) -> f32 {
self.estimated_input_tokens as f32 / self.context_limit.max(1) as f32
}
pub(crate) fn total_reserved_tokens(&self) -> u32 {
self.estimated_input_tokens
.saturating_add(self.reserved_output_tokens)
}
fn with_transcript(&self, messages: &[ConversationMessage]) -> Self {
let messages = project_messages_for_context(messages, Some(self.context_limit));
let text_bytes = messages
.iter()
.map(estimate_conversation_message_bytes)
.sum::<usize>();
let image_tokens = messages
.iter()
.map(|message| count_images_in_content(&message.content))
.sum::<u64>()
.saturating_mul(ESTIMATED_IMAGE_TOKENS);
let history_tokens = (text_bytes as u64)
.div_ceil(ESTIMATED_BYTES_PER_TOKEN)
.saturating_add(image_tokens);
let mut usage = self.clone();
usage.history_tokens = u32::try_from(history_tokens).unwrap_or(u32::MAX);
// Without retaining the fixed prompt itself, rounding its token estimate
// separately from history can differ from the next full estimate by one token.
usage.estimated_input_tokens = usage
.fixed_input_tokens
.saturating_add(usage.history_tokens);
usage
}
}
#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub(crate) struct ProviderContextUsage {
pub(crate) active: ProviderContextProfileUsage,
pub(crate) compaction: ProviderContextProfileUsage,
#[serde(default)]
pub(crate) monitor: Option<ProviderContextProfileUsage>,
}
impl ProviderContextUsage {
pub(crate) fn compaction_usage(&self) -> f32 {
self.compaction.history_tokens as f32
/ self.compaction.history_capacity_tokens.max(1) as f32
}
pub(crate) fn with_transcript(&self, messages: &[ConversationMessage]) -> Self {
let active = self.active.with_transcript(messages);
let monitor = self
.monitor
.as_ref()
.map(|usage| usage.with_transcript(messages));
let compaction = controlling_context_profile(&active, monitor.as_ref());
Self {
active,
compaction,
monitor,
}
}
}
fn controlling_context_profile(
active: &ProviderContextProfileUsage,
monitor: Option<&ProviderContextProfileUsage>,
) -> ProviderContextProfileUsage {
monitor
.filter(|usage| {
// Profiles bound individual outputs differently. Compare actual
// projected pressure rather than only the available capacity.
u64::from(usage.history_tokens) * u64::from(active.history_capacity_tokens.max(1))
> u64::from(active.history_tokens) * u64::from(usage.history_capacity_tokens.max(1))
})
.unwrap_or(active)
.clone()
}
#[derive(Clone, Debug, PartialEq)] #[derive(Clone, Debug, PartialEq)]
pub(crate) enum ProviderRunProjection { pub(crate) enum ProviderRunProjection {
ModelTurnRequested { ModelTurnRequested {
@@ -41,6 +135,7 @@ pub(crate) enum ProviderRunProjection {
runtime_id: String, runtime_id: String,
model_id: String, model_id: String,
retry_attempt: u32, retry_attempt: u32,
context_usage: Option<ProviderContextUsage>,
}, },
ModelTurnStarted { ModelTurnStarted {
work_id: ExternalWorkId, work_id: ExternalWorkId,
@@ -267,33 +362,66 @@ impl ProviderRunCoordinator {
pub(crate) fn progressive_summary_context_budget(&self) -> Option<(u32, u32)> { pub(crate) fn progressive_summary_context_budget(&self) -> Option<(u32, u32)> {
[self.run.profile().as_str(), CLI_MONITOR_PROVIDER_PROFILE] [self.run.profile().as_str(), CLI_MONITOR_PROVIDER_PROFILE]
.into_iter() .into_iter()
.filter_map(|name| { .filter_map(|name| self.context_profile_usage(name))
.min_by_key(|usage| usage.history_capacity_tokens)
.map(|usage| {
(
usage.context_limit,
usage
.fixed_input_tokens
.saturating_add(usage.reserved_output_tokens),
)
})
}
/// Shared by compaction decisions and the context display. Each profile is
/// projected with its own output bounds before measuring its usable history.
pub(crate) fn context_usage_snapshot(&self) -> Option<ProviderContextUsage> {
let active = self.context_profile_usage(self.run.profile().as_str())?;
let monitor = self.context_profile_usage(CLI_MONITOR_PROVIDER_PROFILE);
let compaction = controlling_context_profile(&active, monitor.as_ref());
Some(ProviderContextUsage {
active,
compaction,
monitor,
})
}
/// Retention can use these projected message sizes while preserving the same
/// indices into the complete durable transcript for prefix validation.
pub(crate) fn projected_transcript_for_context(
&self,
profile: &str,
) -> Option<Vec<ConversationMessage>> {
Some(self.projected_request_for_context(profile)?.messages)
}
fn projected_request_for_context(&self, name: &str) -> Option<TurnRequest> {
let profile = self.profiles.get(name)?;
Some(project_request_for_messages(
profile.request.clone(),
self.run.transcript(),
profile.max_context_tokens.or(self.max_context_tokens),
self.run.task_list(),
))
}
pub(crate) fn context_profile_usage(&self, name: &str) -> Option<ProviderContextProfileUsage> {
let profile = self.profiles.get(name)?; let profile = self.profiles.get(name)?;
let limit = profile.max_context_tokens.or(self.max_context_tokens)?; let limit = profile.max_context_tokens.or(self.max_context_tokens)?;
let mut request = profile.request.clone(); let request = self.projected_request_for_context(name)?;
request.prompt = None; Some(context_profile_usage(name, limit, &request))
request.messages.clear();
add_task_list_context(&mut request, self.run.task_list());
let fixed = u32::try_from(estimate_turn_request_tokens(&request))
.unwrap_or(u32::MAX);
Some((limit, fixed))
})
.min_by_key(|(limit, fixed)| limit.saturating_sub(*fixed))
} }
/// Includes the current transcript, fixed request context, tools, and output reservation. /// Includes the current transcript, fixed request context, tools, and output reservation.
/// Read at the model boundary so newly committed tool output is included. /// Read at the model boundary so newly committed tool output is included.
pub(crate) fn estimated_context_tokens(&self) -> Option<(u32, u32)> { pub(crate) fn estimated_context_tokens(&self) -> Option<(u32, u32)> {
let profile = self.profiles.get(self.run.profile().as_str())?; let request = self.projected_request_for_context(self.run.profile().as_str())?;
let mut request = profile.request.clone(); let mut fixed_request = request.clone();
request.prompt = None; fixed_request.messages.clear();
request.messages.clear();
add_task_list_context(&mut request, self.run.task_list());
let persistent_tokens = estimate_turn_request_tokens(&request);
request.messages = self.run.transcript().to_vec();
Some(( Some((
u32::try_from(estimate_turn_request_tokens(&request)).unwrap_or(u32::MAX), u32::try_from(estimate_turn_request_tokens(&request)).unwrap_or(u32::MAX),
u32::try_from(persistent_tokens).unwrap_or(u32::MAX), u32::try_from(estimate_turn_request_tokens(&fixed_request)).unwrap_or(u32::MAX),
)) ))
} }
@@ -516,7 +644,10 @@ impl ProviderRunCoordinator {
// Provider tokenizers and request envelopes differ. Compact before the // Provider tokenizers and request envelopes differ. Compact before the
// estimate reaches the advertised limit, leaving room for that uncertainty. // estimate reaches the advertised limit, leaving room for that uncertainty.
let safe_context_tokens = u64::from(max_context_tokens) * 9 / 10; let safe_context_tokens = u64::from(max_context_tokens) * 9 / 10;
estimate_turn_request_tokens(&request) >= safe_context_tokens u64::from(
context_profile_usage(call.profile.as_str(), max_context_tokens, &request)
.total_reserved_tokens(),
) >= safe_context_tokens
} }
async fn checkpoint_or_fail<C>( async fn checkpoint_or_fail<C>(
@@ -585,6 +716,7 @@ impl ProviderRunCoordinator {
runtime_id: runtime_id.clone(), runtime_id: runtime_id.clone(),
model_id: model_id.clone(), model_id: model_id.clone(),
retry_attempt: call.retry_attempt, retry_attempt: call.retry_attempt,
context_usage: self.context_usage_snapshot(),
}, },
project, project,
) )
@@ -1185,21 +1317,68 @@ fn validate_profile_runtime(
} }
fn request_for_model_call( fn request_for_model_call(
mut template: TurnRequest, template: TurnRequest,
call: &ProviderModelCall, call: &ProviderModelCall,
max_context_tokens: Option<u32>, max_context_tokens: Option<u32>,
task_list: Option<&TaskListSnapshot>, task_list: Option<&TaskListSnapshot>,
) -> TurnRequest {
project_request_for_messages(template, &call.messages, max_context_tokens, task_list)
}
fn project_request_for_messages(
mut template: TurnRequest,
messages: &[ConversationMessage],
max_context_tokens: Option<u32>,
task_list: Option<&TaskListSnapshot>,
) -> TurnRequest { ) -> TurnRequest {
template.prompt = None; template.prompt = None;
template.messages = call.messages.clone(); template.messages = project_messages_for_context(messages, max_context_tokens);
for message in &mut template.messages { add_task_list_context(&mut template, task_list);
add_exploration_reminder(&mut template);
template
}
fn project_messages_for_context(
messages: &[ConversationMessage],
max_context_tokens: Option<u32>,
) -> Vec<ConversationMessage> {
let mut messages = messages.to_vec();
for message in &mut messages {
message.truncate_tool_results_for_provider_request_with_context(max_context_tokens); message.truncate_tool_results_for_provider_request_with_context(max_context_tokens);
message message
.truncate_tagged_terminal_output_for_provider_request_with_context(max_context_tokens); .truncate_tagged_terminal_output_for_provider_request_with_context(max_context_tokens);
} }
add_task_list_context(&mut template, task_list); messages
add_exploration_reminder(&mut template); }
template
fn context_profile_usage(
profile: &str,
context_limit: u32,
request: &TurnRequest,
) -> ProviderContextProfileUsage {
let reserved_output_tokens = request.max_output_tokens.unwrap_or_default();
let estimated_input_tokens =
estimate_turn_request_tokens(request).saturating_sub(reserved_output_tokens);
let mut fixed_request = request.clone();
fixed_request.messages.clear();
let fixed_input_tokens =
estimate_turn_request_tokens(&fixed_request).saturating_sub(reserved_output_tokens);
ProviderContextProfileUsage {
profile: profile.to_owned(),
model_id: request.model.as_str().to_owned(),
context_limit,
estimated_input_tokens: u32::try_from(estimated_input_tokens).unwrap_or(u32::MAX),
fixed_input_tokens: u32::try_from(fixed_input_tokens).unwrap_or(u32::MAX),
reserved_output_tokens: u32::try_from(reserved_output_tokens).unwrap_or(u32::MAX),
history_tokens: u32::try_from(estimated_input_tokens.saturating_sub(fixed_input_tokens))
.unwrap_or(u32::MAX),
history_capacity_tokens: u32::try_from(
u64::from(context_limit)
.saturating_sub(fixed_input_tokens)
.saturating_sub(reserved_output_tokens),
)
.unwrap_or(u32::MAX),
}
} }
/// A transient prompt reminder, never a limit or a synthetic tool result. It is /// A transient prompt reminder, never a limit or a synthetic tool result. It is
@@ -1271,6 +1450,17 @@ const ESTIMATED_BYTES_PER_TOKEN: u64 = 3;
const ESTIMATED_REQUEST_OVERHEAD_TOKENS: u64 = 512; const ESTIMATED_REQUEST_OVERHEAD_TOKENS: u64 = 512;
const ESTIMATED_IMAGE_TOKENS: u64 = 1_600; const ESTIMATED_IMAGE_TOKENS: u64 = 1_600;
/// Per-message history estimate from the same bounded projection used by requests.
/// Rounding each message separately is conservative by at most one token per message.
pub(crate) fn estimated_context_message_tokens(message: &ConversationMessage) -> u32 {
let tokens = (estimate_conversation_message_bytes(message) as u64)
.div_ceil(ESTIMATED_BYTES_PER_TOKEN)
.saturating_add(
count_images_in_content(&message.content).saturating_mul(ESTIMATED_IMAGE_TOKENS),
);
u32::try_from(tokens).unwrap_or(u32::MAX)
}
/// Heuristic estimate, not an exact tokenizer count or provider-specific bound. /// Heuristic estimate, not an exact tokenizer count or provider-specific bound.
/// UTF-8 bytes account for non-ASCII content better than character counts; raw /// UTF-8 bytes account for non-ASCII content better than character counts; raw
/// image payload bytes are excluded from text accounting and charged separately. /// image payload bytes are excluded from text accounting and charged separately.
@@ -1900,6 +1900,152 @@ async fn final_model_request_bounds_new_tool_results_to_the_active_context() {
assert_eq!(content, &original_content); assert_eq!(content, &original_content);
} }
#[tokio::test]
async fn context_usage_snapshot_matches_the_bounded_model_request() {
let runtime = Arc::new(ScriptedRuntime::new(vec![answer_turn()]));
let original_content = "x".repeat(20_000);
let mut turn_request = request();
turn_request.max_output_tokens = Some(500);
turn_request.messages = vec![ConversationMessage {
role: MessageRole::User,
content: MessageContent::ToolResult {
tool_use_id: "large-result".to_string(),
content: original_content.clone(),
is_error: false,
},
}];
let mut coordinator = ProviderRunCoordinator::from_request(
"projected-context-snapshot",
ProviderRunProfile::new(runtime.clone(), turn_request).with_max_context_tokens(Some(8_000)),
Vec::new(),
ProviderRunLimits::default(),
)
.unwrap();
let snapshot = coordinator.context_usage_snapshot().unwrap();
// The complete result would exceed the early-compaction threshold, but the
// bounded request fits comfortably. Occupancy must describe that request.
assert!(snapshot.compaction_usage() < 0.70);
assert_eq!(snapshot.active.reserved_output_tokens, 500);
let MessageContent::ToolResult { content, .. } = &coordinator.run().transcript()[0].content
else {
panic!("expected original tool result");
};
assert_eq!(content, &original_content);
let (_sender, control) = turn_control();
let block = coordinator
.drive_until_blocked(control, |_| Ok(()))
.await
.unwrap();
assert!(matches!(block, ProviderRunBlock::AwaitingDriver { .. }));
let requests = runtime.requests();
assert_eq!(
u64::from(snapshot.active.total_reserved_tokens()),
estimate_turn_request_tokens(&requests[0])
);
assert_eq!(
snapshot.active.estimated_input_tokens,
snapshot
.active
.total_reserved_tokens()
.saturating_sub(snapshot.active.reserved_output_tokens)
);
}
#[test]
fn context_usage_snapshot_projects_history_for_the_limiting_profile() {
let runtime = Arc::new(ScriptedRuntime::new(vec![]));
let mut turn_request = request();
turn_request.messages = vec![ConversationMessage {
role: MessageRole::User,
content: MessageContent::ToolResult {
tool_use_id: "large-result".to_string(),
content: "x".repeat(20_000),
is_error: false,
},
}];
let mut coordinator = ProviderRunCoordinator::from_request(
"limiting-profile-context-snapshot",
ProviderRunProfile::new(runtime.clone(), turn_request)
.with_max_context_tokens(Some(100_000)),
Vec::new(),
ProviderRunLimits::default(),
)
.unwrap();
coordinator
.insert_profile(
CLI_MONITOR_PROVIDER_PROFILE,
ProviderRunProfile::new(runtime, request()).with_max_context_tokens(Some(8_000)),
)
.unwrap();
let snapshot = coordinator.context_usage_snapshot().unwrap();
assert_eq!(snapshot.active.context_limit, 100_000);
assert_eq!(snapshot.compaction.profile, CLI_MONITOR_PROVIDER_PROFILE);
assert_eq!(snapshot.compaction.context_limit, 8_000);
assert!(snapshot.active.history_tokens > snapshot.compaction.history_tokens);
assert!(snapshot.compaction_usage() < 0.70);
assert_eq!(
snapshot.compaction.history_capacity_tokens,
8_000 - snapshot.compaction.fixed_input_tokens
);
}
#[test]
fn compaction_pressure_uses_projected_history_as_well_as_profile_capacity() {
let runtime = Arc::new(ScriptedRuntime::new(vec![]));
let mut turn_request = request();
turn_request.max_output_tokens = Some(70_000);
turn_request.messages = vec![ConversationMessage {
role: MessageRole::User,
content: MessageContent::ToolResult {
tool_use_id: "large-result".to_string(),
content: "x".repeat(100_000),
is_error: false,
},
}];
let mut coordinator = ProviderRunCoordinator::from_request(
"profile-pressure-order",
ProviderRunProfile::new(runtime.clone(), turn_request)
.with_max_context_tokens(Some(100_000)),
Vec::new(),
ProviderRunLimits::default(),
)
.unwrap();
coordinator
.insert_profile(
CLI_MONITOR_PROVIDER_PROFILE,
ProviderRunProfile::new(runtime, request()).with_max_context_tokens(Some(8_000)),
)
.unwrap();
let snapshot = coordinator.context_usage_snapshot().unwrap();
let monitor = coordinator
.context_profile_usage(CLI_MONITOR_PROVIDER_PROFILE)
.unwrap();
// The smaller monitor clips the result much more aggressively, so its
// smaller capacity must not hide compaction pressure in the active model.
assert!(snapshot.active.history_capacity_tokens > monitor.history_capacity_tokens);
assert_eq!(snapshot.compaction.profile, BASE_PROVIDER_PROFILE);
assert!(snapshot.compaction_usage() >= 0.70);
assert!(monitor.history_tokens as f32 / (monitor.history_capacity_tokens as f32) < 0.70);
assert_eq!(
coordinator.progressive_summary_context_budget().unwrap().0,
8_000
);
let restored: ProviderContextUsage =
serde_json::from_str(&serde_json::to_string(&snapshot).unwrap()).unwrap();
let compacted = restored.with_transcript(&[ConversationMessage {
role: MessageRole::User,
content: MessageContent::Text("Task checkpoint".to_string()),
}]);
assert_eq!(compacted.compaction.profile, CLI_MONITOR_PROVIDER_PROFILE);
}
#[tokio::test] #[tokio::test]
async fn preflight_leaves_headroom_before_the_advertised_context_limit() { async fn preflight_leaves_headroom_before_the_advertised_context_limit() {
let runtime = Arc::new(ScriptedRuntime::new(vec![answer_turn()])); let runtime = Arc::new(ScriptedRuntime::new(vec![answer_turn()]));
+2
View File
@@ -263,6 +263,8 @@ impl From<&gql_usage::ConversationUsage> for ConversationUsageInfo {
lines_removed: tool.apply_file_diff_stats.lines_removed, lines_removed: tool.apply_file_diff_stats.lines_removed,
commands_executed: tool.run_command_stats.commands_executed, commands_executed: tool.run_command_stats.commands_executed,
current_context_tokens: 0, current_context_tokens: 0,
provider_context_usage: None,
latest_model_call_reported_context_tokens: None,
estimated_cost_cents: 0.0, estimated_cost_cents: 0.0,
total_cache_read_tokens: 0, total_cache_read_tokens: 0,
total_cache_write_tokens: 0, total_cache_write_tokens: 0,