From cb830e1ba76b8da16e07ce5fdda376dbb62f9dde Mon Sep 17 00:00:00 2001 From: Ryan Ward Date: Tue, 20 Oct 2026 21:05:57 +0000 Subject: [PATCH] Run progressive compaction alongside foreground provider work --- AGENTS.md | 1 + app/src/ai/blocklist/controller.rs | 307 ++++++++++++-- app/src/ai/blocklist/controller_tests.rs | 387 +++++++++++++++++- .../ai/runtime/provider_run_coordinator.rs | 18 + .../runtime/provider_run_coordinator_tests.rs | 112 +++++ 5 files changed, 794 insertions(+), 31 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index c7bd8c01..df151ec7 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -219,6 +219,7 @@ Key invariants: - Bedrock prompt caching uses three cache points: system prompt, second-to-last history message, and tool configuration - `ensure_tool_results_paired()` enforces Bedrock's invariant that every `tool_use` has a matching `tool_result` - Progressive summaries are prepended to provider requests as a user/assistant pair; direct-provider runs compact their live transcript at a checkpointed `ReadyToCallModel` boundary before another model call, while background summarization remains independent for session-owned runtimes +- Direct-provider progressive compaction runs in a separate request alongside foreground work. Start early against the smaller usable history budget of the current and terminal-use profiles, including each profile's prompt/tools/output reservation. Apply a result only at a safe model boundary after validating its exact snapshotted prefix, preserving all newer history and recomputing live context accounting. Wait for pending compaction only after context exhaustion; cancellation and terminal cleanup must remain immediate. Cleanly completed conversations may accept a late summary while idle; stale, cancelled, or superseded generations must never replace a newer run's history. - Loop prevention in `controller.rs` detects repeated tool failures (3+ identical) and injects a corrective instruction - Direct-provider long-running shell follow-ups retain CLI tasks as UI projections while command monitoring and completion stay owned by the same provider run - A direct-provider command completion is only queued when the terminal reports it; the CLI task remains active until the provider run applies that completion at a safe boundary and deactivates it diff --git a/app/src/ai/blocklist/controller.rs b/app/src/ai/blocklist/controller.rs index bfafeb5a..b1a72b83 100644 --- a/app/src/ai/blocklist/controller.rs +++ b/app/src/ai/blocklist/controller.rs @@ -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, ) { 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, ) { + // 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, + ) { + 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( diff --git a/app/src/ai/blocklist/controller_tests.rs b/app/src/ai/blocklist/controller_tests.rs index 097e5005..e4de9412 100644 --- a/app/src/ai/blocklist/controller_tests.rs +++ b/app/src/ai/blocklist/controller_tests.rs @@ -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, +) -> 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, +) -> 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![ diff --git a/app/src/ai/runtime/provider_run_coordinator.rs b/app/src/ai/runtime/provider_run_coordinator.rs index 391e2138..5eb9e291 100644 --- a/app/src/ai/runtime/provider_run_coordinator.rs +++ b/app/src/ai/runtime/provider_run_coordinator.rs @@ -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)> { diff --git a/app/src/ai/runtime/provider_run_coordinator_tests.rs b/app/src/ai/runtime/provider_run_coordinator_tests.rs index fa685f81..a8eed5b7 100644 --- a/app/src/ai/runtime/provider_run_coordinator_tests.rs +++ b/app/src/ai/runtime/provider_run_coordinator_tests.rs @@ -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()]));