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

Run progressive compaction alongside foreground provider work

Browse Source
This commit is contained in:
Ryan Ward
2026-10-20 21:05:57 +00:00
parent 0383166f4e
commit cb830e1ba7
5 changed files with 794 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/src/ai/blocklist/controller.rs
+279 -28
View File
@@ -158,14 +158,24 @@ struct ActiveProviderRunIdentity {
enum ActiveProviderProgressiveSummaryState {
InFlight {
identity: ActiveProviderRunIdentity,
waiting_for_context: bool,
},
Ready {
identity: ActiveProviderRunIdentity,
waiting_for_context: bool,
plan: ActiveProgressiveSummaryPlan,
result: anyhow::Result<(String, Usage, String)>,
},
}
impl ActiveProviderProgressiveSummaryState {
fn identity(&self) -> &ActiveProviderRunIdentity {
match self {
Self::InFlight { identity, .. } | Self::Ready { identity, .. } => identity,
}
}
}
fn configured_llm_context_limit(info: &LLMInfo) -> u32 {
[
info.context_window.default_max,
@@ -6955,16 +6965,21 @@ impl BlocklistAIController {
conversation.has_pending_progressive_summary()
};
let transcript = active_run.coordinator.run().transcript();
let Some((current_context_tokens, persistent_context_tokens)) =
let Some((current_context_tokens, current_persistent_context_tokens)) =
self.active_provider_context_estimate(conversation_id, ctx)
else {
return false;
};
let (summary_context_limit, summary_persistent_context_tokens) = active_run
.coordinator
.progressive_summary_context_budget()
.unwrap_or((active_context_limit, current_persistent_context_tokens));
// 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);
let history_capacity =
summary_context_limit.saturating_sub(summary_persistent_context_tokens);
let current_history_tokens =
current_context_tokens.saturating_sub(persistent_context_tokens);
current_context_tokens.saturating_sub(current_persistent_context_tokens);
if summary_pending
|| last_failure_context_tokens == Some(current_context_tokens)
|| (!force
@@ -7030,7 +7045,7 @@ impl BlocklistAIController {
retained_messages_count: transcript.len().saturating_sub(split_point),
trigger_context_tokens: current_context_tokens,
active_context_limit,
persistent_context_tokens,
persistent_context_tokens: current_persistent_context_tokens,
candidates,
summarized_prefix: transcript[..split_point].to_vec(),
summarize_messages: vec![ConversationMessage {
@@ -7051,6 +7066,7 @@ impl BlocklistAIController {
conversation_id,
ActiveProviderProgressiveSummaryState::InFlight {
identity: identity.clone(),
waiting_for_context: force,
},
);
BlocklistAIHistoryModel::handle(ctx).update(ctx, |history, _| {
@@ -7121,27 +7137,23 @@ impl BlocklistAIController {
ctx: &mut ModelContext<Self>,
) {
let conversation_id = identity.conversation_id;
let owns_summary = self
let waiting_for_context = match self
.active_provider_progressive_summaries
.get(&conversation_id)
.is_some_and(|state| {
matches!(
state,
ActiveProviderProgressiveSummaryState::InFlight {
identity: active_identity,
} if active_identity == &identity
)
});
if !owns_summary {
return;
}
{
Some(ActiveProviderProgressiveSummaryState::InFlight {
identity: active_identity,
waiting_for_context,
}) if active_identity == &identity => *waiting_for_context,
_ => return,
};
let identity_is_current =
self.active_provider_runs
.get(&conversation_id)
.is_some_and(|slot| {
slot.stream_id == identity.stream_id && slot.run_id == identity.run_id
});
if !identity_is_current {
if !identity_is_current && self.active_provider_runs.contains_key(&conversation_id) {
self.active_provider_progressive_summaries
.remove(&conversation_id);
BlocklistAIHistoryModel::handle(ctx).update(ctx, |history, _| {
@@ -7156,11 +7168,29 @@ impl BlocklistAIController {
conversation_id,
ActiveProviderProgressiveSummaryState::Ready {
identity,
waiting_for_context,
plan,
result,
},
);
self.drive_active_provider_run(conversation_id, ctx);
if !self.active_provider_runs.contains_key(&conversation_id) {
self.apply_ready_idle_provider_progressive_summary(conversation_id, ctx);
return;
}
// A background result must not re-enter an active model call or dispatch
// an awaiting tool batch again. Normal work consumes it at its next boundary.
if waiting_for_context
&& self.active_provider_runs.get(&conversation_id).is_some_and(|slot| {
slot.run.as_ref().is_some_and(|run| {
matches!(
run.coordinator.run().state(),
ProviderRunState::ReadyToCallModel
)
})
})
{
self.drive_active_provider_run(conversation_id, ctx);
}
}
fn apply_ready_active_provider_progressive_summary(
@@ -7193,6 +7223,7 @@ impl BlocklistAIController {
let Some(ActiveProviderProgressiveSummaryState::Ready {
identity: _,
waiting_for_context,
plan,
result,
}) = self
@@ -7209,6 +7240,7 @@ impl BlocklistAIController {
let mut summary_applied = false;
let mut recorded_usage = None;
let mut failure = None;
match result {
Ok((summary, usage, summarizer_model_id))
if estimated_history_tokens(&progressive_summary_messages(summary.clone()))
@@ -7218,6 +7250,7 @@ impl BlocklistAIController {
"[progressive-summary] Summary did not reduce context for conversation {:?}",
conversation_id,
);
failure = Some("summary did not reduce context".to_string());
recorded_usage = Some((usage, summarizer_model_id));
}
Ok((summary, usage, summarizer_model_id)) => {
@@ -7234,6 +7267,9 @@ impl BlocklistAIController {
"[progressive-summary] Discarding stale background summary for active conversation {:?}; the snapshotted prefix no longer matches",
conversation_id,
);
failure = Some(
"summary prefix no longer matches conversation history".to_string(),
);
} else {
let drained = transcript[..plan.split_point].to_vec();
let retained = transcript[plan.split_point..].to_vec();
@@ -7241,6 +7277,14 @@ impl BlocklistAIController {
let summary_prefix = progressive_summary_messages(summary.clone());
let summary_prefix_len = summary_prefix.len();
let summary_prefix_tokens = estimated_history_tokens(&summary_prefix);
let persistent_context_tokens = run
.coordinator
.estimated_context_tokens()
.map_or(plan.persistent_context_tokens, |(_, fixed)| fixed);
let active_context_limit = run
.coordinator
.current_context_limit()
.unwrap_or(plan.active_context_limit);
if let Err(error) = run
.coordinator
.run_mut()
@@ -7255,12 +7299,11 @@ impl BlocklistAIController {
if let Ok(mut messages_sent) = run.messages_sent.lock() {
*messages_sent = retained.clone();
}
let new_context_tokens = plan
.persistent_context_tokens
let new_context_tokens = persistent_context_tokens
.saturating_add(summary_prefix_tokens)
.saturating_add(estimated_history_tokens(&retained));
let new_usage =
new_context_tokens as f32 / plan.active_context_limit.max(1) as f32;
new_context_tokens as f32 / active_context_limit.max(1) as f32;
BlocklistAIHistoryModel::handle(ctx).update(ctx, |history, _| {
let Some(conversation) = history.conversation_mut(&conversation_id)
else {
@@ -7278,7 +7321,7 @@ impl BlocklistAIController {
summarizer_model_id,
new_context_tokens,
new_usage * 100.0,
plan.active_context_limit,
active_context_limit,
retained_count,
);
recorded_usage = Some((usage, summarizer_model_id));
@@ -7291,9 +7334,13 @@ impl BlocklistAIController {
"[progressive-summary] Failed for active conversation {:?}: {error:#}",
conversation_id,
);
failure = Some(error.to_string());
}
}
let failure_context_tokens = self
.active_provider_context_estimate(conversation_id, ctx)
.map_or(plan.trigger_context_tokens, |(tokens, _)| tokens);
if let Some(run) = self
.active_provider_runs
.get_mut(&conversation_id)
@@ -7303,12 +7350,23 @@ impl BlocklistAIController {
None
} else {
// Fail open and suppress another request at this unchanged context boundary.
Some(plan.trigger_context_tokens)
Some(failure_context_tokens)
};
}
if let Some((usage, summarizer_model_id)) = recorded_usage {
self.record_progressive_summary_usage(conversation_id, usage, summarizer_model_id, ctx);
}
if waiting_for_context && !summary_applied {
self.fail_active_provider_run(
conversation_id,
format!(
"context compaction could not recover enough space: {}",
failure.unwrap_or_else(|| "could not replace the summarized history".to_string())
),
ctx,
);
return true;
}
if let Err(error) = self.persist_active_provider_run(conversation_id, ctx) {
self.fail_active_provider_run(
conversation_id,
@@ -7464,18 +7522,55 @@ impl BlocklistAIController {
conversation_id: AIConversationId,
ctx: &mut ModelContext<Self>,
) {
// A follow-up owns a new generation. Its history must never be blocked
// by, or replaced with, a detached summary from the previous generation.
// Cancellation also discards summaries before they can fail recovery.
let stale_summary = self
.active_provider_progressive_summaries
.get(&conversation_id)
.zip(self.active_provider_runs.get(&conversation_id))
.is_some_and(|(summary, slot)| {
let identity = summary.identity();
slot.stream_id != identity.stream_id || slot.run_id != identity.run_id
});
let cancelling = self
.active_provider_runs
.get(&conversation_id)
.is_some_and(|slot| slot.cancellation_reason.is_some());
if (stale_summary || cancelling)
&& self
.active_provider_progressive_summaries
.remove(&conversation_id)
.is_some()
{
BlocklistAIHistoryModel::handle(ctx).update(ctx, |history, _| {
if let Some(conversation) = history.conversation_mut(&conversation_id) {
conversation.set_has_pending_progressive_summary(false);
}
});
}
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);
// Ordinary background work never blocks the agent. Once a model has
// exhausted its context, polling must not relaunch that same request
// while its recovery summary is still running. Cancellation still drives.
if matches!(
self.active_provider_progressive_summaries
.get(&conversation_id),
Some(ActiveProviderProgressiveSummaryState::InFlight { .. })
) {
self.active_provider_progressive_summaries.get(&conversation_id),
Some(ActiveProviderProgressiveSummaryState::InFlight {
waiting_for_context: true,
..
})
) && self.active_provider_runs.get(&conversation_id).is_some_and(|slot| {
slot.cancellation_reason.is_none()
&& slot.run.as_ref().is_some_and(|run| {
matches!(run.coordinator.run().state(), ProviderRunState::ReadyToCallModel)
})
}) {
return;
}
self.begin_active_provider_progressive_summary(conversation_id, false, ctx);
let Some(slot) = self.active_provider_runs.get_mut(&conversation_id) else {
return;
};
@@ -7977,6 +8072,24 @@ impl BlocklistAIController {
);
return;
}
match self.active_provider_progressive_summaries.get_mut(&conversation_id) {
Some(ActiveProviderProgressiveSummaryState::InFlight {
waiting_for_context,
..
}) => {
*waiting_for_context = true;
return;
}
Some(ActiveProviderProgressiveSummaryState::Ready {
waiting_for_context,
..
}) => {
*waiting_for_context = true;
self.drive_active_provider_run(conversation_id, ctx);
return;
}
None => {}
}
if !self.begin_active_provider_progressive_summary(conversation_id, true, ctx) {
self.fail_active_provider_run(
conversation_id,
@@ -8826,6 +8939,11 @@ impl BlocklistAIController {
&outcome,
),
);
let completed_cleanly = matches!(
&outcome,
ProviderRunOutcome::Completed(completion)
if completion.stop_reason == StopReason::Completed
);
match outcome {
ProviderRunOutcome::Completed(completion) => match completion.stop_reason {
StopReason::Completed => {
@@ -8890,7 +9008,140 @@ impl BlocklistAIController {
}
}
}
// A background checkpoint remains useful after the answer finishes. Detach
// it while cleanup releases this generation, then retain it only if no
// successor has started or claimed its own summary in the meantime.
let pending_summary = if completed_cleanly {
self.active_provider_progressive_summaries
.remove(&conversation_id)
} else {
None
};
self.cleanup_active_provider_run(conversation_id, &stream_id, &response_stream, ctx);
if let Some(summary) = pending_summary {
let conversation_succeeded = BlocklistAIHistoryModel::as_ref(ctx)
.conversation_status(&conversation_id)
== Some(&ConversationStatus::Success);
if conversation_succeeded
&& !self.active_provider_runs.contains_key(&conversation_id)
&& !self
.active_provider_progressive_summaries
.contains_key(&conversation_id)
{
self.active_provider_progressive_summaries
.insert(conversation_id, summary);
BlocklistAIHistoryModel::handle(ctx).update(ctx, |history, _| {
if let Some(conversation) = history.conversation_mut(&conversation_id) {
conversation.set_has_pending_progressive_summary(true);
}
});
self.apply_ready_idle_provider_progressive_summary(conversation_id, ctx);
}
}
}
fn apply_ready_idle_provider_progressive_summary(
&mut self,
conversation_id: AIConversationId,
ctx: &mut ModelContext<Self>,
) {
if self.active_provider_runs.contains_key(&conversation_id)
|| !matches!(
self.active_provider_progressive_summaries
.get(&conversation_id),
Some(ActiveProviderProgressiveSummaryState::Ready { .. })
)
{
return;
}
let Some(ActiveProviderProgressiveSummaryState::Ready { plan, result, .. }) = self
.active_provider_progressive_summaries
.remove(&conversation_id)
else {
return;
};
BlocklistAIHistoryModel::handle(ctx).update(ctx, |history, _| {
if let Some(conversation) = history.conversation_mut(&conversation_id) {
conversation.set_has_pending_progressive_summary(false);
}
});
if BlocklistAIHistoryModel::as_ref(ctx).conversation_status(&conversation_id)
!= Some(&ConversationStatus::Success)
{
return;
}
match result {
Ok((summary, usage, summarizer_model_id)) => {
let summary_prefix_tokens =
estimated_history_tokens(&progressive_summary_messages(summary.clone()));
let summary_reduces_context =
summary_prefix_tokens < estimated_history_tokens(&plan.summarized_prefix);
BlocklistAIHistoryModel::handle(ctx).update(ctx, |history, ctx| {
let Some(conversation) = history.conversation_mut(&conversation_id) else {
return;
};
let mut transcript = conversation
.progressive_summary()
.map(|existing| progressive_summary_messages(existing.to_owned()))
.unwrap_or_default();
transcript.extend_from_slice(conversation.bedrock_message_history());
if transcript.get(..plan.split_point)
!= Some(plan.summarized_prefix.as_slice())
{
log::warn!(
"[progressive-summary] Discarding stale idle summary for conversation {:?}; the snapshotted prefix no longer matches",
conversation_id,
);
return;
}
if !summary_reduces_context {
log::warn!(
"[progressive-summary] Idle summary did not reduce context for conversation {:?}",
conversation_id,
);
return;
}
let retained = transcript.split_off(plan.split_point);
let retained_count = retained.len();
let new_context_tokens = plan
.persistent_context_tokens
.saturating_add(summary_prefix_tokens)
.saturating_add(estimated_history_tokens(&retained));
let new_usage =
new_context_tokens as f32 / plan.active_context_limit.max(1) as f32;
conversation.archive_tool_results(transcript);
*conversation.bedrock_message_history_mut() = retained;
conversation.set_progressive_summary(Some(summary), 0);
conversation.set_current_context_tokens(new_context_tokens);
conversation.set_context_window_usage(new_usage.clamp(0.0, 1.0));
conversation.write_updated_conversation_state(ctx);
log::info!(
"[progressive-summary] Applied background summary after active run completion {:?}: ~{} tokens ({:.1}% of {} context), {} tail messages preserved",
conversation_id,
new_context_tokens,
new_usage * 100.0,
plan.active_context_limit,
retained_count,
);
});
// This Ready result is consumed once, even when its prefix became
// stale or its response could not reduce the retained context.
self.record_progressive_summary_usage(
conversation_id,
usage,
summarizer_model_id,
ctx,
);
}
Err(error) => {
log::error!(
"[progressive-summary] Failed after active conversation {:?} completed: {error:#}",
conversation_id,
);
}
}
}
fn finalize_failed_provider_conversation(
app/src/ai/blocklist/controller_tests.rs
+384 -3
View File
@@ -4,16 +4,16 @@ use std::sync::{Arc, Mutex};
use ai::agent::action::{AskUserQuestionItem, AskUserQuestionType};
use chrono::Local;
use galaxy_agent_core::{
AgentError, AgentErrorKind, CompletedModelTurn, ContentPart, ConversationMessage,
turn_control, AgentError, AgentErrorKind, CompletedModelTurn, ContentPart, ConversationMessage,
ExternalWorkId, MessageContent, MessageRole, PermissionKind, PermissionRequest, ProviderRun,
ProviderRunFailure, ProviderRunFailureKind, ProviderRunId, ProviderRunLimits,
ProviderRunOutcome, ProviderRunState, ProviderRunStep, RunEpoch, RuntimeCapabilities,
StopReason, ToolCall, TurnRequest, Usage,
StopReason, ToolCall, TurnCommand, TurnRequest, Usage,
};
use galaxy_core::command::ExitCode;
use uuid::Uuid;
use warp_multi_agent_api::response_event;
use warpui::{App, EntityId, SingletonEntity};
use warpui::{App, EntityId, ModelHandle, SingletonEntity};
use crate::ai::agent::conversation::{AIConversationId, ConversationStatus};
use crate::ai::agent::task::TaskId;
@@ -3915,6 +3915,387 @@ fn optimistic_cli_subagent_completion_with_in_flight_stream_reports_success() {
});
}
#[test]
fn progressive_summary_completion_preserves_an_in_flight_model_turn() {
App::test((), |mut app| async move {
initialize_app_for_terminal_view(&mut app);
let terminal = add_window_with_terminal(&mut app, None);
terminal.update(&mut app, |terminal, ctx| {
let conversation_id =
BlocklistAIHistoryModel::handle(ctx).update(ctx, |history, ctx| {
let conversation_id =
history.start_new_conversation(terminal.id(), false, false, false, ctx);
history
.conversation_mut(&conversation_id)
.unwrap()
.set_has_pending_progressive_summary(true);
conversation_id
});
let stream_id = ResponseStreamId::new_for_test();
let response_stream =
ctx.add_model(|_| ResponseStream::new_for_test(stream_id.clone()));
let mut slot = progressive_summary_test_slot(conversation_id, stream_id, response_stream);
let identity = super::ActiveProviderRunIdentity {
conversation_id,
stream_id: slot.stream_id.clone(),
run_id: slot.run_id.clone(),
};
let (sender, foreground_control) = turn_control();
slot.turn_control = Some(sender);
let prefix = vec![ConversationMessage {
role: MessageRole::User,
content: MessageContent::Text("old project details".repeat(1_000)),
}];
terminal.ai_controller().update(ctx, |controller, ctx| {
controller.active_provider_runs.insert(conversation_id, slot);
controller.active_provider_progressive_summaries.insert(
conversation_id,
super::ActiveProviderProgressiveSummaryState::InFlight {
identity: identity.clone(),
waiting_for_context: true,
},
);
controller.finish_active_provider_progressive_summary(
identity.clone(),
progressive_summary_test_plan(prefix),
Ok((
"checkpoint".to_owned(),
Usage::default(),
"summary-model".to_owned(),
)),
ctx,
);
let active = &controller.active_provider_runs[&conversation_id];
assert_eq!(active.stream_id, identity.stream_id);
assert_eq!(active.run_id, identity.run_id);
assert!(active.run.is_none());
assert!(active.checkpoint.is_none());
assert!(active.cancellation_reason.is_none());
let steer = TurnCommand::Steer {
display_text: "continue".to_owned(),
model_text: "continue".to_owned(),
};
assert!(foreground_control.try_receive().is_err());
active
.turn_control
.as_ref()
.unwrap()
.try_send(steer.clone())
.unwrap();
assert_eq!(foreground_control.try_receive().unwrap(), steer);
match &controller.active_provider_progressive_summaries[&conversation_id] {
super::ActiveProviderProgressiveSummaryState::Ready {
identity: ready_identity,
waiting_for_context,
result,
..
} => {
assert_eq!(ready_identity.run_id, identity.run_id);
assert!(*waiting_for_context);
assert_eq!(result.as_ref().unwrap().0, "checkpoint");
}
super::ActiveProviderProgressiveSummaryState::InFlight { .. } => {
panic!("completed summary must remain ready until the next model boundary");
}
}
});
let conversation = BlocklistAIHistoryModel::as_ref(ctx)
.conversation(&conversation_id)
.unwrap();
assert!(conversation.has_pending_progressive_summary());
assert_eq!(conversation.progressive_summary(), None);
});
});
}
#[test]
fn stale_progressive_summary_completion_preserves_a_successor_generation() {
App::test((), |mut app| async move {
initialize_app_for_terminal_view(&mut app);
let terminal = add_window_with_terminal(&mut app, None);
terminal.update(&mut app, |terminal, ctx| {
let conversation_id =
BlocklistAIHistoryModel::handle(ctx).update(ctx, |history, ctx| {
let conversation_id =
history.start_new_conversation(terminal.id(), false, false, false, ctx);
history
.conversation_mut(&conversation_id)
.unwrap()
.set_has_pending_progressive_summary(true);
conversation_id
});
let stream_id = ResponseStreamId::new_for_test();
let response_stream =
ctx.add_model(|_| ResponseStream::new_for_test(stream_id.clone()));
let slot = progressive_summary_test_slot(conversation_id, stream_id, response_stream);
let identity = super::ActiveProviderRunIdentity {
conversation_id,
stream_id: slot.stream_id.clone(),
run_id: slot.run_id.clone(),
};
let stale_identity = super::ActiveProviderRunIdentity {
conversation_id,
stream_id: ResponseStreamId::new_for_test(),
run_id: ProviderRunId::new("previous-generation"),
};
terminal.ai_controller().update(ctx, |controller, ctx| {
controller.active_provider_runs.insert(conversation_id, slot);
controller.active_provider_progressive_summaries.insert(
conversation_id,
super::ActiveProviderProgressiveSummaryState::InFlight {
identity: identity.clone(),
waiting_for_context: false,
},
);
controller.finish_active_provider_progressive_summary(
stale_identity,
progressive_summary_test_plan(Vec::new()),
Ok((
"stale checkpoint".to_owned(),
Usage::default(),
"summary-model".to_owned(),
)),
ctx,
);
let active = &controller.active_provider_runs[&conversation_id];
assert_eq!(active.run_id, identity.run_id);
assert_eq!(active.stream_id, identity.stream_id);
match &controller.active_provider_progressive_summaries[&conversation_id] {
super::ActiveProviderProgressiveSummaryState::InFlight {
identity: pending_identity,
waiting_for_context,
} => {
assert_eq!(pending_identity.run_id, identity.run_id);
assert_eq!(pending_identity.stream_id, identity.stream_id);
assert!(!*waiting_for_context);
}
super::ActiveProviderProgressiveSummaryState::Ready { .. } => {
panic!("stale completion must not replace the successor summary");
}
}
});
let conversation = BlocklistAIHistoryModel::as_ref(ctx)
.conversation(&conversation_id)
.unwrap();
assert!(conversation.has_pending_progressive_summary());
assert_eq!(conversation.progressive_summary(), None);
});
});
}
#[test]
fn idle_progressive_summary_preserves_appended_messages_and_archives_tool_results() {
App::test((), |mut app| async move {
initialize_app_for_terminal_view(&mut app);
let terminal = add_window_with_terminal(&mut app, None);
terminal.update(&mut app, |terminal, ctx| {
let prefix = vec![
ConversationMessage {
role: MessageRole::Assistant,
content: MessageContent::ToolUse {
tool_use_id: "read-original".to_owned(),
name: "read_file".to_owned(),
input: serde_json::json!({"path": "src/original.rs"}),
},
},
ConversationMessage {
role: MessageRole::User,
content: MessageContent::ToolResult {
tool_use_id: "read-original".to_owned(),
content: "original file contents".repeat(1_000),
is_error: false,
},
},
];
let tail = vec![
ConversationMessage {
role: MessageRole::Assistant,
content: MessageContent::Text("The change is ready".to_owned()),
},
ConversationMessage {
role: MessageRole::User,
content: MessageContent::Text("Appended while the summary was running".to_owned()),
},
];
let conversation_id =
BlocklistAIHistoryModel::handle(ctx).update(ctx, |history, ctx| {
let conversation_id =
history.start_new_conversation(terminal.id(), false, false, false, ctx);
let conversation = history.conversation_mut(&conversation_id).unwrap();
conversation.append_to_bedrock_history(prefix.clone());
conversation.append_to_bedrock_history(tail.clone());
conversation.set_has_pending_progressive_summary(true);
conversation.set_status_for_test(ConversationStatus::Success);
conversation_id
});
let identity = super::ActiveProviderRunIdentity {
conversation_id,
stream_id: ResponseStreamId::new_for_test(),
run_id: ProviderRunId::new("finished-generation"),
};
terminal.ai_controller().update(ctx, |controller, ctx| {
controller.active_provider_progressive_summaries.insert(
conversation_id,
super::ActiveProviderProgressiveSummaryState::InFlight {
identity: identity.clone(),
waiting_for_context: false,
},
);
controller.finish_active_provider_progressive_summary(
identity,
progressive_summary_test_plan(prefix.clone()),
Ok((
"checkpoint".to_owned(),
Usage::default(),
"summary-model".to_owned(),
)),
ctx,
);
assert!(!controller.active_provider_runs.contains_key(&conversation_id));
assert!(!controller
.active_provider_progressive_summaries
.contains_key(&conversation_id));
});
let conversation = BlocklistAIHistoryModel::as_ref(ctx)
.conversation(&conversation_id)
.unwrap();
assert_eq!(conversation.progressive_summary(), Some("checkpoint"));
assert_eq!(conversation.bedrock_message_history(), tail.as_slice());
assert_eq!(conversation.tool_result_archive(), prefix.as_slice());
assert!(!conversation.has_pending_progressive_summary());
assert_eq!(conversation.status(), &ConversationStatus::Success);
assert_eq!(
conversation.current_context_tokens(),
100 + super::estimated_history_tokens(&super::progressive_summary_messages(
"checkpoint".to_owned()
)) + super::estimated_history_tokens(&tail)
);
});
});
}
#[test]
fn idle_progressive_summary_rejects_a_changed_prefix() {
App::test((), |mut app| async move {
initialize_app_for_terminal_view(&mut app);
let terminal = add_window_with_terminal(&mut app, None);
terminal.update(&mut app, |terminal, ctx| {
let prefix = vec![ConversationMessage {
role: MessageRole::User,
content: MessageContent::Text("original project context".repeat(1_000)),
}];
let replacement = vec![ConversationMessage {
role: MessageRole::User,
content: MessageContent::Text("replacement context that must survive".repeat(1_000)),
}];
let conversation_id =
BlocklistAIHistoryModel::handle(ctx).update(ctx, |history, ctx| {
let conversation_id =
history.start_new_conversation(terminal.id(), false, false, false, ctx);
let conversation = history.conversation_mut(&conversation_id).unwrap();
conversation.append_to_bedrock_history(replacement.clone());
conversation.set_has_pending_progressive_summary(true);
conversation.set_status_for_test(ConversationStatus::Success);
conversation_id
});
let identity = super::ActiveProviderRunIdentity {
conversation_id,
stream_id: ResponseStreamId::new_for_test(),
run_id: ProviderRunId::new("finished-generation"),
};
terminal.ai_controller().update(ctx, |controller, ctx| {
controller.active_provider_progressive_summaries.insert(
conversation_id,
super::ActiveProviderProgressiveSummaryState::InFlight {
identity: identity.clone(),
waiting_for_context: false,
},
);
controller.finish_active_provider_progressive_summary(
identity,
progressive_summary_test_plan(prefix),
Ok((
"stale checkpoint".to_owned(),
Usage::default(),
"summary-model".to_owned(),
)),
ctx,
);
assert!(!controller
.active_provider_progressive_summaries
.contains_key(&conversation_id));
});
let conversation = BlocklistAIHistoryModel::as_ref(ctx)
.conversation(&conversation_id)
.unwrap();
assert_eq!(conversation.progressive_summary(), None);
assert_eq!(conversation.bedrock_message_history(), replacement.as_slice());
assert!(conversation.tool_result_archive().is_empty());
assert!(!conversation.has_pending_progressive_summary());
});
});
}
fn progressive_summary_test_plan(
summarized_prefix: Vec<ConversationMessage>,
) -> super::ActiveProgressiveSummaryPlan {
super::ActiveProgressiveSummaryPlan {
split_point: summarized_prefix.len(),
retained_messages_count: 1,
trigger_context_tokens: 10_000,
active_context_limit: 128_000,
persistent_context_tokens: 100,
candidates: Vec::new(),
summarize_messages: summarized_prefix.clone(),
summarized_prefix,
}
}
fn progressive_summary_test_slot(
conversation_id: AIConversationId,
stream_id: ResponseStreamId,
response_stream: ModelHandle<ResponseStream>,
) -> super::ActiveProviderRunSlot {
let snapshot = provider_snapshot(conversation_id);
super::ActiveProviderRunSlot {
stream_id,
response_stream,
did_input_contain_user_query: true,
run_id: snapshot.run.id().clone(),
root_task_id: snapshot.root_task_id,
projection_target: snapshot.projection_target,
run: None,
checkpoint: None,
turn_control: None,
cancellation_reason: None,
committed_provider_batch: None,
finished_provider_batch: None,
command_action_refs: HashMap::new(),
command_monitor: None,
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(),
}
}
#[test]
fn progressive_summary_compacts_short_histories_to_the_token_budget() {
let messages = vec![
app/src/ai/runtime/provider_run_coordinator.rs
+18
View File
@@ -260,6 +260,24 @@ impl ProviderRunCoordinator {
.or(self.max_context_tokens)
}
/// Leave enough history space for a command-monitor turn before switching
/// profiles, accounting for each profile's own fixed prompt and output budget.
pub(crate) fn progressive_summary_context_budget(&self) -> Option<(u32, u32)> {
[self.run.profile().as_str(), CLI_MONITOR_PROVIDER_PROFILE]
.into_iter()
.filter_map(|name| {
let profile = self.profiles.get(name)?;
let limit = profile.max_context_tokens.or(self.max_context_tokens)?;
let mut request = profile.request.clone();
request.prompt = None;
request.messages.clear();
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.
/// Read at the model boundary so newly committed tool output is included.
pub(crate) fn estimated_context_tokens(&self) -> Option<(u32, u32)> {
app/src/ai/runtime/provider_run_coordinator_tests.rs
+112
View File
@@ -1608,6 +1608,118 @@ async fn compaction_boundary_keeps_the_monitor_context_limit() {
assert_eq!(coordinator.current_context_limit(), Some(100_000));
}
#[test]
fn progressive_summary_budget_anticipates_switching_to_a_smaller_monitor() {
let runtime = Arc::new(ScriptedRuntime::new(vec![]));
let mut coordinator = coordinator(runtime.clone());
coordinator.set_max_context_tokens(Some(500_000));
coordinator
.insert_profile(
CLI_MONITOR_PROVIDER_PROFILE,
ProviderRunProfile::new(runtime, request()).with_max_context_tokens(Some(275_000)),
)
.unwrap();
assert_eq!(coordinator.current_context_limit(), Some(500_000));
assert_eq!(
coordinator
.progressive_summary_context_budget()
.map(|(limit, _)| limit),
Some(275_000)
);
coordinator
.run_mut()
.set_ready_profile(CLI_MONITOR_PROVIDER_PROFILE)
.unwrap();
assert_eq!(coordinator.current_context_limit(), Some(275_000));
assert_eq!(
coordinator
.progressive_summary_context_budget()
.map(|(limit, _)| limit),
Some(275_000)
);
coordinator
.run_mut()
.set_ready_profile(BASE_PROVIDER_PROFILE)
.unwrap();
assert_eq!(coordinator.current_context_limit(), Some(500_000));
assert_eq!(
coordinator
.progressive_summary_context_budget()
.map(|(limit, _)| limit),
Some(275_000)
);
}
#[test]
fn progressive_summary_budget_respects_smaller_active_limits_and_profile_fallbacks() {
for (base_limit, monitor_limit, fallback_limit, expected) in [
(
Some(200_000),
Some(Some(275_000)),
Some(500_000),
Some(200_000),
),
(Some(500_000), Some(None), Some(275_000), Some(275_000)),
(Some(500_000), None, Some(275_000), Some(500_000)),
(None, None, None, None),
] {
let runtime = Arc::new(ScriptedRuntime::new(vec![]));
let mut coordinator = ProviderRunCoordinator::from_request(
"progressive-summary-budget",
ProviderRunProfile::new(runtime.clone(), request()).with_max_context_tokens(base_limit),
Vec::new(),
ProviderRunLimits::default(),
)
.unwrap();
coordinator.set_max_context_tokens(fallback_limit);
if let Some(monitor_limit) = monitor_limit {
coordinator
.insert_profile(
CLI_MONITOR_PROVIDER_PROFILE,
ProviderRunProfile::new(runtime, request())
.with_max_context_tokens(monitor_limit),
)
.unwrap();
}
assert_eq!(
coordinator.current_context_limit(),
base_limit.or(fallback_limit)
);
assert_eq!(
coordinator
.progressive_summary_context_budget()
.map(|(limit, _)| limit),
expected
);
}
}
#[test]
fn progressive_summary_budget_accounts_for_output_reservations_when_selecting_a_profile() {
let runtime = Arc::new(ScriptedRuntime::new(vec![]));
let mut coordinator = coordinator(runtime.clone());
coordinator.set_max_context_tokens(Some(275_000));
let mut monitor_request = request();
monitor_request.max_output_tokens = Some(300_000);
coordinator
.insert_profile(
CLI_MONITOR_PROVIDER_PROFILE,
ProviderRunProfile::new(runtime, monitor_request)
.with_max_context_tokens(Some(500_000)),
)
.unwrap();
let (limit, persistent_tokens) = coordinator.progressive_summary_context_budget().unwrap();
assert_eq!(coordinator.current_context_limit(), Some(275_000));
assert_eq!(limit, 500_000);
assert!(persistent_tokens > 300_000);
}
#[tokio::test]
async fn monitor_request_uses_its_profile_context_limit() {
let base_runtime = Arc::new(ScriptedRuntime::new(vec![answer_turn()]));