Run progressive compaction alongside foreground provider work
This commit is contained in:
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user