1use std::collections::BTreeMap;
26use std::pin::Pin;
27use std::sync::Arc;
28
29use pi_ai::{
30 AssistantMessage, AssistantMessageEvent, Context, Model, ProviderError, StopReason,
31 StreamOptions,
32};
33use regex::Regex;
34use tokio_util::sync::CancellationToken;
35
36use crate::core::compaction::{
37 self, BeforeCompactResult, CompactOptions, CompactionError, CompactionResult,
38 CompactionSettings, SummarizeStreamFn, preparation_none_error, prepare_compaction,
39 should_compact,
40};
41use crate::core::model_runtime::ModelRuntimeAuthOverrides;
42use crate::core::sessions::{CompactionEntry, SessionEntry, get_latest_compaction_entry};
43use crate::core::settings::ResolvedCompactionSettings;
44
45use super::AgentSession;
46use super::events::{AgentSessionEvent, CompactionReason};
47use super::extension_runner::ExtensionRunner;
48
49#[derive(Clone)]
59struct CompactionInputs {
60 stream_fn: SummarizeStreamFn,
62 api_key: Option<String>,
65 headers: Option<BTreeMap<String, Option<String>>>,
67 env: Option<BTreeMap<String, String>>,
69}
70
71#[derive(Clone)]
76pub struct CompactionStreamHandle {
77 pub stream_fn: SummarizeStreamFn,
79 pub api_key: Option<String>,
81}
82
83impl CompactionStreamHandle {
84 #[must_use]
86 pub fn new(stream_fn: SummarizeStreamFn) -> Self {
87 Self {
88 stream_fn,
89 api_key: None,
90 }
91 }
92}
93
94impl AgentSession {
99 pub async fn compact(
118 &self,
119 custom_instructions: Option<&str>,
120 ) -> Result<CompactionResult, CompactionError> {
121 self.disconnect_from_agent();
124 self.abort().await;
125
126 let abort_token = self.begin_compaction_abort();
127 self.emit_public_awaited(&AgentSessionEvent::CompactionStart {
128 reason: CompactionReason::Manual,
129 })
130 .await;
131
132 let weak_self = self.upgrade_self();
135
136 let result = self
137 .run_compaction_core(
138 CompactionReason::Manual,
139 custom_instructions,
140 false,
141 abort_token,
142 true,
143 )
144 .await;
145
146 let outcome = match result {
147 Ok(Some(compaction_result)) => {
148 self.emit_public_awaited(&AgentSessionEvent::CompactionEnd {
149 reason: CompactionReason::Manual,
150 result: Some(compaction_result.clone()),
151 aborted: false,
152 will_retry: false,
153 error_message: None,
154 })
155 .await;
156 Ok(compaction_result)
157 }
158 Ok(None) => {
159 let path_entries = self.snapshot_branch_entries().await;
162 let path_refs: Vec<&SessionEntry> = path_entries.iter().collect();
163 let err = preparation_none_error(&path_refs);
164 let message = err.to_string();
165 self.emit_public_awaited(&AgentSessionEvent::CompactionEnd {
166 reason: CompactionReason::Manual,
167 result: None,
168 aborted: false,
169 will_retry: false,
170 error_message: Some(message),
171 })
172 .await;
173 Err(err)
174 }
175 Err(err) => {
176 let aborted = matches!(err, CompactionError::Cancelled);
177 let message = if aborted {
178 None
179 } else {
180 Some(format!("Compaction failed: {err}"))
181 };
182 self.emit_public_awaited(&AgentSessionEvent::CompactionEnd {
183 reason: CompactionReason::Manual,
184 result: None,
185 aborted,
186 will_retry: false,
187 error_message: message,
188 })
189 .await;
190 Err(err)
191 }
192 };
193
194 self.clear_compaction_abort();
196 if let Some(arc) = weak_self.and_then(|w| w.upgrade()) {
197 arc.reconnect_to_agent();
198 }
199
200 outcome
201 }
202
203 pub(super) async fn check_compaction(
222 &self,
223 assistant_message: &AssistantMessage,
224 skip_aborted_check: bool,
225 ) -> bool {
226 let settings = self.compaction_settings();
227 let enabled = self.lock_inner().auto_compaction_enabled;
231 if !enabled {
232 return false;
233 }
234
235 if skip_aborted_check && assistant_message.stop_reason == StopReason::Aborted {
236 return false;
237 }
238
239 let model = self.model();
240 let context_window = model.context_window;
241
242 let same_model = same_model(&model, assistant_message);
245
246 let compaction_entry = self.latest_compaction_entry().await;
249 if let Some(latest) = &compaction_entry {
250 let boundary_ts = parse_iso_to_millis(&latest.timestamp);
251 if boundary_ts > 0 && assistant_message.timestamp <= boundary_ts {
252 return false;
253 }
254 }
255
256 if same_model && is_context_overflow(assistant_message, context_window) {
258 let will_retry = assistant_message.stop_reason != StopReason::Stop;
259
260 if !will_retry {
261 return self
264 .run_auto_compaction(CompactionReason::Overflow, false)
265 .await;
266 }
267
268 let already_attempted = {
270 let inner = self.lock_inner();
271 inner.overflow_recovery_attempted
272 };
273
274 if already_attempted {
275 self.emit_public_awaited(&AgentSessionEvent::CompactionEnd {
278 reason: CompactionReason::Overflow,
279 result: None,
280 aborted: false,
281 will_retry: false,
282 error_message: Some(
283 "Context overflow recovery failed after one compact-and-retry attempt"
284 .to_owned(),
285 ),
286 })
287 .await;
288 return false;
289 }
290
291 {
295 let mut inner = self.lock_inner();
296 inner.overflow_recovery_attempted = true;
297 }
298 let _ = self.agent.pop_last_if_assistant();
299 return self
300 .run_auto_compaction(CompactionReason::Overflow, will_retry)
301 .await;
302 }
303
304 let context_tokens =
306 self.threshold_context_tokens(assistant_message, compaction_entry.as_ref());
307 if should_compact(context_tokens, context_window, &settings_to_pure(settings)) {
308 return self
309 .run_auto_compaction(CompactionReason::Threshold, false)
310 .await;
311 }
312
313 false
314 }
315
316 async fn run_auto_compaction(&self, reason: CompactionReason, will_retry: bool) -> bool {
323 let abort_token = self.begin_auto_compaction_abort();
324
325 self.emit_public_awaited(&AgentSessionEvent::CompactionStart { reason })
326 .await;
327
328 let result = self
329 .run_compaction_core(reason, None, will_retry, abort_token, false)
330 .await;
331
332 let should_continue = match result {
333 Ok(Some(compaction_result)) => {
334 self.emit_public_awaited(&AgentSessionEvent::CompactionEnd {
335 reason,
336 result: Some(compaction_result),
337 aborted: false,
338 will_retry,
339 error_message: None,
340 })
341 .await;
342
343 if will_retry {
344 let _ = self.agent.pop_last_if_assistant();
347 true
348 } else {
349 self.agent.has_queued_messages()
353 }
354 }
355 Ok(None) => {
356 self.emit_public_awaited(&AgentSessionEvent::CompactionEnd {
361 reason,
362 result: None,
363 aborted: false,
364 will_retry: false,
365 error_message: None,
366 })
367 .await;
368 false
369 }
370 Err(CompactionError::Cancelled) => {
371 self.emit_public_awaited(&AgentSessionEvent::CompactionEnd {
372 reason,
373 result: None,
374 aborted: true,
375 will_retry: false,
376 error_message: None,
377 })
378 .await;
379 false
380 }
381 Err(err) => {
382 let message = if reason == CompactionReason::Overflow {
383 format!("Context overflow recovery failed: {err}")
384 } else {
385 format!("Auto-compaction failed: {err}")
386 };
387 self.emit_public_awaited(&AgentSessionEvent::CompactionEnd {
388 reason,
389 result: None,
390 aborted: false,
391 will_retry: false,
392 error_message: Some(message),
393 })
394 .await;
395 false
396 }
397 };
398
399 self.clear_auto_compaction_abort();
400 should_continue
401 }
402
403 async fn run_compaction_core(
416 &self,
417 reason: CompactionReason,
418 custom_instructions: Option<&str>,
419 will_retry: bool,
420 abort_token: CancellationToken,
421 is_manual: bool,
422 ) -> Result<Option<CompactionResult>, CompactionError> {
423 let model = self.model();
424
425 let (path_entries, settings) = {
428 let sm = self.session_manager.lock().await;
429 let branch: Vec<&SessionEntry> = sm.get_branch(None);
430 let entries: Vec<SessionEntry> = branch.into_iter().cloned().collect();
431 (entries, self.compaction_settings())
432 };
433 let pure_settings = settings_to_pure(settings);
434
435 let path_refs: Vec<&SessionEntry> = path_entries.iter().collect();
436 let preparation = prepare_compaction(&path_refs, pure_settings)?;
437 let Some(preparation) = preparation else {
438 return Ok(None);
439 };
440
441 let inputs = self.resolve_compaction_inputs().await?;
443
444 let runner = self.hooks.runner();
446 if runner.has_handlers("session_before_compact") {
447 let event = AgentSessionEvent::CompactionStart { reason };
448 let cancel = self.extension_before_compact(&runner, event).await?;
449 if cancel.cancel {
450 return Err(CompactionError::Cancelled);
451 }
452 if let Some(replacement) = cancel.compaction {
453 let mut replacement = replacement;
456 replacement.from_hook = Some(true);
457 return self
458 .finalize_compaction_result(replacement, true, reason, will_retry, &abort_token)
459 .await
460 .map(Some);
461 }
462 }
463
464 let thinking_level = thinking_level_str(self.thinking_level());
466
467 let result = compaction::compact(
468 &preparation,
469 CompactOptions {
470 model: &model,
471 api_key: inputs.api_key.clone(),
472 headers: inputs.headers.clone(),
473 custom_instructions: if is_manual { custom_instructions } else { None },
474 signal: Some(abort_token.clone()),
475 thinking_level: thinking_level.as_deref(),
476 stream_fn: inputs.stream_fn.clone(),
477 env: inputs.env.clone(),
478 hooks: None,
479 },
480 )
481 .await?;
482
483 self.finalize_compaction_result(result, false, reason, will_retry, &abort_token)
485 .await
486 .map(Some)
487 }
488
489 async fn finalize_compaction_result(
492 &self,
493 mut result: CompactionResult,
494 from_hook: bool,
495 reason: CompactionReason,
496 will_retry: bool,
497 abort_token: &CancellationToken,
498 ) -> Result<CompactionResult, CompactionError> {
499 if abort_token.is_cancelled() {
500 return Err(CompactionError::Cancelled);
501 }
502
503 let (_entry_id, saved_entry) = {
505 let mut sm = self.session_manager.lock().await;
506 let tokens_before_i64: i64 = i64::try_from(result.tokens_before).unwrap_or(i64::MAX);
507 let details = result.details.clone();
508 let entry_id = sm
509 .append_compaction(
510 &result.summary,
511 &result.first_kept_entry_id,
512 tokens_before_i64,
513 details,
514 if from_hook { Some(true) } else { None },
515 )
516 .map_err(|err| {
517 CompactionError::SummarizationFailed(format!(
518 "failed to persist compaction: {err}"
519 ))
520 })?;
521 let saved = sm.get_entry(&entry_id).cloned();
522 (entry_id, saved)
523 };
524
525 let new_messages = {
527 let sm = self.session_manager.lock().await;
528 sm.build_session_context()
529 .map_err(CompactionError::MessageConversion)?
530 .messages
531 };
532 self.agent.replace_messages(new_messages);
533
534 let estimated_after = compaction::estimate_context_tokens(&self.agent.transcript()).tokens;
536 result.estimated_tokens_after = Some(estimated_after);
537
538 if let Some(SessionEntry::Compaction(compaction_entry)) = &saved_entry {
540 let runner = self.hooks.runner();
541 self.extension_after_compact(&runner, compaction_entry, from_hook, reason, will_retry)
542 .await;
543 }
544
545 if let Some(entry) = saved_entry {
547 self.emit_public(AgentSessionEvent::EntryAppended { entry });
548 }
549
550 Ok(result)
551 }
552
553 async fn resolve_compaction_inputs(&self) -> Result<CompactionInputs, CompactionError> {
565 if let Some(runtime) = self.model_runtime_handle() {
566 let model = self.model();
567
568 let auth = runtime
569 .get_auth_for_model(&model, ModelRuntimeAuthOverrides::default())
570 .await
571 .map_err(|err| {
572 CompactionError::SummarizationFailed(format!("auth resolution failed: {err}"))
573 })?;
574
575 let api_key = auth.as_ref().and_then(|a| a.auth.api_key.clone());
576 let headers = auth
577 .as_ref()
578 .and_then(|a| a.auth.headers.clone())
579 .map(header_map_to_option);
580 let env = auth.and_then(|a| {
581 a.env.map(|provider_env| {
582 provider_env
583 .into_iter()
584 .collect::<BTreeMap<String, String>>()
585 })
586 });
587
588 let stream_fn: SummarizeStreamFn = {
589 let runtime = runtime.clone();
590 Arc::new(move |model: Model, ctx: Context, opts: StreamOptions| {
591 let runtime = runtime.clone();
592 Box::pin(async move {
593 runtime.stream_simple(model, ctx, opts)
594 as Pin<
595 Box<
596 dyn futures::Stream<
597 Item = Result<AssistantMessageEvent, ProviderError>,
598 > + Send,
599 >,
600 >
601 })
602 })
603 };
604
605 return Ok(CompactionInputs {
606 stream_fn,
607 api_key,
608 headers,
609 env,
610 });
611 }
612
613 if let Some(handle) = &self.compaction_stream_override {
614 return Ok(CompactionInputs {
615 stream_fn: handle.stream_fn.clone(),
616 api_key: handle.api_key.clone(),
617 headers: None,
618 env: None,
619 });
620 }
621
622 Err(CompactionError::SummarizationFailed(
623 "No model runtime configured for compaction".to_owned(),
624 ))
625 }
626
627 async fn extension_before_compact(
641 &self,
642 runner: &Arc<dyn ExtensionRunner>,
643 event: AgentSessionEvent,
644 ) -> Result<BeforeCompactResult, CompactionError> {
645 match runner.emit(event).await {
646 Ok(Some(cancel)) if cancel.cancel => Ok(BeforeCompactResult {
647 cancel: true,
648 compaction: None,
649 }),
650 Ok(_) => Ok(BeforeCompactResult::default()),
651 Err(err) => {
652 runner.emit_error(err.to_string());
653 Ok(BeforeCompactResult::default())
654 }
655 }
656 }
657
658 async fn extension_after_compact(
662 &self,
663 runner: &Arc<dyn ExtensionRunner>,
664 _entry: &CompactionEntry,
665 _from_hook: bool,
666 reason: CompactionReason,
667 will_retry: bool,
668 ) {
669 let event = AgentSessionEvent::CompactionEnd {
670 reason,
671 result: None,
672 aborted: false,
673 will_retry,
674 error_message: None,
675 };
676 if let Err(err) = runner.emit(event).await {
677 runner.emit_error(err.to_string());
678 }
679 }
680
681 fn begin_auto_compaction_abort(&self) -> CancellationToken {
685 let token = CancellationToken::new();
686 let mut inner = self.lock_inner();
687 if let Some(prev) = inner.auto_compaction_abort.take() {
688 prev.cancel();
689 }
690 inner.auto_compaction_abort = Some(token.clone());
691 token
692 }
693
694 fn clear_auto_compaction_abort(&self) {
696 self.lock_inner().auto_compaction_abort = None;
697 }
698
699 fn compaction_settings(&self) -> ResolvedCompactionSettings {
703 self.lock_settings().get_compaction_settings()
704 }
705
706 async fn snapshot_branch_entries(&self) -> Vec<SessionEntry> {
708 let sm = self.session_manager.lock().await;
709 sm.get_branch(None).into_iter().cloned().collect()
710 }
711
712 async fn latest_compaction_entry(&self) -> Option<CompactionEntry> {
714 let sm = self.session_manager.lock().await;
715 let branch = sm.get_branch(None);
716 get_latest_compaction_entry(&branch).cloned()
717 }
718
719 fn threshold_context_tokens(
725 &self,
726 assistant_message: &AssistantMessage,
727 compaction_entry: Option<&CompactionEntry>,
728 ) -> u64 {
729 let direct = if assistant_message.usage.total_tokens != 0 {
730 compaction::calculate_context_tokens(&assistant_message.usage)
731 } else {
732 0
733 };
734
735 if assistant_message.stop_reason != StopReason::Error && direct > 0 {
736 return direct;
737 }
738
739 let messages = self.agent.transcript();
741 let estimate = compaction::estimate_context_tokens(&messages);
742 if estimate.last_usage_index.is_none() {
743 return 0;
744 }
745
746 if let Some(latest) = compaction_entry
748 && let boundary_ts = parse_iso_to_millis(&latest.timestamp)
749 && boundary_ts > 0
750 && let Some(idx) = estimate.last_usage_index
751 && let Some(msg) = messages.get(idx)
752 && let Some(pi_ai::Message::Assistant(assistant)) = msg.as_llm()
753 && assistant.timestamp <= boundary_ts
754 {
755 return 0;
756 }
757
758 estimate.tokens
759 }
760}
761
762fn settings_to_pure(resolved: ResolvedCompactionSettings) -> CompactionSettings {
768 CompactionSettings {
769 enabled: resolved.enabled,
770 reserve_tokens: resolved.reserve_tokens,
771 keep_recent_tokens: resolved.keep_recent_tokens,
772 }
773}
774
775fn same_model(model: &Model, assistant: &AssistantMessage) -> bool {
777 assistant.provider == model.provider && assistant.model == model.id
778}
779
780fn thinking_level_str(level: pi_ai::ModelThinkingLevel) -> Option<String> {
782 serde_json::to_value(level)
783 .ok()
784 .and_then(|v| v.as_str().map(str::to_owned))
785}
786
787fn header_map_to_option(headers: BTreeMap<String, String>) -> BTreeMap<String, Option<String>> {
790 headers.into_iter().map(|(k, v)| (k, Some(v))).collect()
791}
792
793fn parse_iso_to_millis(timestamp: &str) -> i64 {
797 parse_rfc3339_millis(timestamp).unwrap_or(0)
802}
803
804fn parse_rfc3339_millis(ts: &str) -> Option<i64> {
806 let date_part = ts.get(..10)?; let time_part = ts.get(11..)?;
810
811 let year: i64 = date_part.get(..4)?.parse().ok()?;
812 let month: u32 = date_part.get(5..7)?.parse().ok()?;
813 let day: u32 = date_part.get(8..10)?.parse().ok()?;
814
815 let time_core = time_part.split(['+', '-']).next()?.trim_end_matches('Z');
817
818 let (hms, millis) = time_core.split_once('.').unwrap_or((time_core, "0"));
819 let hour: u32 = hms.get(..2)?.parse().ok()?;
820 let minute: u32 = hms.get(3..5)?.parse().ok()?;
821 let second: u32 = hms.get(6..8).unwrap_or("00").parse().ok().unwrap_or(0);
822 let ms: u32 = millis.get(..3).unwrap_or("000").parse().ok().unwrap_or(0);
823
824 civil_to_millis(year, month, day, hour, minute, second, ms)
825}
826
827fn civil_to_millis(
829 year: i64,
830 month: u32,
831 day: u32,
832 hour: u32,
833 min: u32,
834 sec: u32,
835 ms: u32,
836) -> Option<i64> {
837 let y = if month <= 2 { year - 1 } else { year };
839 let era = if y >= 0 { y } else { y - 399 } / 400;
840 let yoe = u32::try_from(y - era * 400).ok()?;
841 let m = if month > 2 { month - 3 } else { month + 9 };
842 let doy = (153 * m + 2) / 5 + day - 1;
843 let doe = yoe * 365 + yoe / 4 - yoe / 100 + doy;
844 let days_since_epoch = era * 146_097 + i64::from(doe) - 719_468;
845
846 let secs =
847 days_since_epoch * 86_400 + i64::from(hour) * 3_600 + i64::from(min) * 60 + i64::from(sec);
848 Some(secs * 1_000 + i64::from(ms))
849}
850
851static OVERFLOW_REGEXES: std::sync::LazyLock<Vec<Regex>> = std::sync::LazyLock::new(|| {
858 [
859 r"(?i)prompt is too long",
860 r"(?i)request_too_large",
861 r"(?i)input is too long for requested model",
862 r"(?i)exceeds the context window",
863 r"(?i)exceeds (?:the )?(?:model'?s )?maximum context length(?: of [\d,]+ tokens?|\s*\([\d,]+\))?",
864 r"(?i)input token count.*exceeds the maximum",
865 r"(?i)maximum prompt length is \d+",
866 r"(?i)reduce the length of the messages",
867 r"(?i)maximum context length is \d+ tokens",
868 r"(?i)exceeds (?:the )?maximum allowed input length of [\d,]+ tokens?",
869 r"(?i)input \(\d+ tokens\) is longer than the model'?s context length \(\d+ tokens\)",
870 r"(?i)exceeds the limit of \d+",
871 r"(?i)exceeds the available context size",
872 r"(?i)greater than the context length",
873 r"(?i)context window exceeds limit",
874 r"(?i)exceeded model token limit",
875 r"(?i)too large for model with \d+ maximum context length",
876 r"(?i)prompt has [\d,]+ tokens?, but the configured context size is [\d,]+ tokens?",
877 r"(?i)model_context_window_exceeded",
878 r"(?i)prompt too long; exceeded (?:max )?context length",
879 r"(?i)context[_ ]length[_ ]exceeded",
880 r"(?i)too many tokens",
881 r"(?i)token limit exceeded",
882 r"(?i)^4(?:00|13)\s*(?:status code)?\s*\(no body\)",
883 ]
884 .into_iter()
885 .filter_map(|pat| Regex::new(pat).ok())
886 .collect()
887});
888
889static NON_OVERFLOW_REGEXES: std::sync::LazyLock<Vec<Regex>> = std::sync::LazyLock::new(|| {
890 [
891 r"(?i)^(Throttling error|Service unavailable):",
892 r"(?i)rate limit",
893 r"(?i)too many requests",
894 ]
895 .into_iter()
896 .filter_map(|pat| Regex::new(pat).ok())
897 .collect()
898});
899
900fn overflow_patterns() -> &'static [Regex] {
901 &OVERFLOW_REGEXES
902}
903
904fn non_overflow_patterns() -> &'static [Regex] {
905 &NON_OVERFLOW_REGEXES
906}
907
908pub(super) fn is_context_overflow(message: &AssistantMessage, context_window: u64) -> bool {
915 if message.stop_reason == StopReason::Error
917 && let Some(error_message) = message.error_message.as_deref()
918 {
919 let is_non_overflow = non_overflow_patterns()
920 .iter()
921 .any(|pattern| pattern.is_match(error_message));
922 if !is_non_overflow
923 && overflow_patterns()
924 .iter()
925 .any(|pattern| pattern.is_match(error_message))
926 {
927 return true;
928 }
929 }
930
931 if context_window > 0 && message.stop_reason == StopReason::Stop {
933 let input_tokens = message.usage.input.saturating_add(message.usage.cache_read);
934 if input_tokens > context_window {
935 return true;
936 }
937 }
938
939 if context_window > 0 && message.stop_reason == StopReason::Length && message.usage.output == 0
941 {
942 let input_tokens = message.usage.input.saturating_add(message.usage.cache_read);
943 let threshold = context_window / 100 * 99 + context_window % 100 * 99 / 100;
944 if input_tokens >= threshold {
945 return true;
946 }
947 }
948
949 false
950}
951
952#[cfg(test)]
957mod tests {
958 use super::*;
959 use crate::core::agent_session::{AgentSession, AgentSessionConfig};
960 use crate::core::sessions::SessionEntry;
961 use futures::stream::{self, BoxStream};
962 use pi_agent::user_text;
963 use pi_ai::{
964 AssistantContent, AssistantMessageEvent, Context, DoneReason, Model, ModelCost, ModelInput,
965 ModelThinkingLevel, Provider, ProviderError, StreamOptions, TextContent, Usage, UsageCost,
966 };
967 use serde_json::{Map, Value};
968 use std::sync::Arc;
969 use tokio::sync::mpsc;
970 use tokio::time::{sleep, timeout};
971
972 type TestResult<T = ()> = Result<T, Box<dyn std::error::Error>>;
973
974 fn test_model(context_window: u64) -> Model {
977 Model {
978 id: "m".to_owned(),
979 name: "m".to_owned(),
980 api: "test-api".to_owned(),
981 provider: "test-provider".to_owned(),
982 base_url: String::new(),
983 reasoning: false,
984 thinking_level_map: None,
985 input: vec![ModelInput::Text],
986 cost: ModelCost::default(),
987 context_window,
988 max_tokens: 1_024,
989 headers: None,
990 compat: None,
991 extra: std::collections::BTreeMap::new(),
992 }
993 }
994
995 fn mock_provider() -> Arc<dyn Provider> {
996 struct NoopProvider;
997 impl Provider for NoopProvider {
998 fn stream(
999 &self,
1000 _model: &Model,
1001 _ctx: Context,
1002 _opts: StreamOptions,
1003 ) -> BoxStream<'static, Result<AssistantMessageEvent, ProviderError>> {
1004 Box::pin(stream::iter(Vec::new()))
1006 }
1007 }
1008 Arc::new(NoopProvider)
1009 }
1010
1011 fn summary_stream_fn(text: &str) -> SummarizeStreamFn {
1012 let text = text.to_owned();
1013 Arc::new(move |_model, _ctx, _opts| {
1014 let text = text.clone();
1015 Box::pin(async move {
1016 let mut msg = AssistantMessage::new("a", "p", "m", 1);
1017 msg.content = vec![AssistantContent::Text(TextContent::new(text.clone()))];
1018 msg.stop_reason = StopReason::Stop;
1019 let stream = stream::iter(vec![Ok(AssistantMessageEvent::Done {
1020 reason: DoneReason::Stop,
1021 message: msg,
1022 })]);
1023 Box::pin(stream)
1024 as Pin<
1025 Box<
1026 dyn futures::Stream<Item = Result<AssistantMessageEvent, ProviderError>>
1027 + Send,
1028 >,
1029 >
1030 })
1031 })
1032 }
1033
1034 fn error_stream_fn(msg: &str) -> SummarizeStreamFn {
1035 let msg = msg.to_owned();
1036 Arc::new(move |_model, _ctx, _opts| {
1037 let msg = msg.clone();
1038 Box::pin(async move {
1039 let mut message = AssistantMessage::new("a", "p", "m", 1);
1040 message.stop_reason = StopReason::Error;
1041 message.error_message = Some(msg.clone());
1042 let stream = stream::iter(vec![Ok(AssistantMessageEvent::Done {
1043 reason: DoneReason::Stop,
1044 message,
1045 })]);
1046 Box::pin(stream)
1047 as Pin<
1048 Box<
1049 dyn futures::Stream<Item = Result<AssistantMessageEvent, ProviderError>>
1050 + Send,
1051 >,
1052 >
1053 })
1054 })
1055 }
1056
1057 fn make_session(
1058 context_window: u64,
1059 stream_fn: SummarizeStreamFn,
1060 messages: Vec<pi_agent::AgentMessage>,
1061 ) -> TestResult<Arc<AgentSession>> {
1062 let provider = mock_provider();
1063 let mut config = AgentSessionConfig::test_config(provider, test_model(context_window))?;
1064 config.system_prompt = "sys".into();
1065 config.messages = messages;
1066 config.compaction_stream_override = Some(CompactionStreamHandle::new(stream_fn));
1067 Ok(AgentSession::new(config)?)
1068 }
1069
1070 fn assistant_with_usage(text: &str, usage: Usage, stop: StopReason) -> pi_agent::AgentMessage {
1071 let mut msg =
1072 AssistantMessage::new("test-api", "test-provider", "m", pi_agent::now_millis());
1073 msg.content = vec![AssistantContent::Text(TextContent::new(text))];
1074 msg.usage = usage;
1075 msg.stop_reason = stop;
1076 pi_agent::AgentMessage::Llm(Box::new(pi_ai::Message::Assistant(msg)))
1077 }
1078
1079 fn assistant_overflow_message() -> AssistantMessage {
1080 let mut msg =
1081 AssistantMessage::new("test-api", "test-provider", "m", pi_agent::now_millis());
1082 msg.stop_reason = StopReason::Error;
1083 msg.error_message = Some("prompt is too long: 999999 tokens > 8192 maximum".into());
1084 msg
1085 }
1086
1087 async fn session_with_history(
1095 context_window: u64,
1096 stream_fn: SummarizeStreamFn,
1097 ) -> TestResult<Arc<AgentSession>> {
1098 let mut messages = Vec::new();
1100 for i in 0..20 {
1101 messages.push(user_text(
1102 format!("user message number {i} with some padding text to make it larger"),
1103 std::iter::empty(),
1104 ));
1105 let usage = Usage {
1106 input: 500 + i * 100,
1107 output: 200,
1108 cache_read: 0,
1109 cache_write: 0,
1110 cache_write1h: None,
1111 reasoning: None,
1112 total_tokens: 700 + i * 100,
1113 cost: UsageCost::default(),
1114 };
1115 messages.push(assistant_with_usage(
1116 format!("assistant reply {i} with enough text to be a real message").as_str(),
1117 usage,
1118 StopReason::Stop,
1119 ));
1120 }
1121
1122 let session = make_session(context_window, stream_fn, messages)?;
1123
1124 {
1127 let mut settings = session.lock_settings();
1128 let mut compaction = Map::new();
1129 compaction.insert("keepRecentTokens".into(), Value::from(1u64));
1130 let mut overrides = Map::new();
1131 overrides.insert("compaction".into(), Value::Object(compaction));
1132 settings.apply_overrides(&overrides);
1133 }
1134
1135 {
1137 let mut sm = session.session_manager.lock().await;
1138 for msg in session.agent.transcript() {
1139 sm.append_message(&msg)?;
1140 }
1141 assert!(
1142 !sm.get_branch(None).is_empty(),
1143 "session_with_history must leave a non-empty leaf branch"
1144 );
1145 }
1146
1147 Ok(session)
1148 }
1149
1150 async fn collect_events(rx: &mut mpsc::UnboundedReceiver<String>, n: usize) -> Vec<String> {
1151 let mut out = Vec::new();
1152 while out.len() < n {
1153 match timeout(std::time::Duration::from_secs(2), rx.recv()).await {
1154 Ok(Some(v)) => out.push(v),
1155 Ok(None) | Err(_) => break,
1156 }
1157 }
1158 out
1159 }
1160
1161 #[test]
1164 fn is_context_overflow_error_pattern() {
1165 let mut msg = AssistantMessage::new("a", "anthropic", "m", 1);
1166 msg.stop_reason = StopReason::Error;
1167 msg.error_message = Some("prompt is too long: 213462 tokens > 200000 maximum".into());
1168 assert!(is_context_overflow(&msg, 0));
1169 }
1170
1171 #[test]
1172 fn is_context_overflow_non_overflow_rate_limit() {
1173 let mut msg = AssistantMessage::new("a", "bedrock", "m", 1);
1174 msg.stop_reason = StopReason::Error;
1175 msg.error_message = Some("rate limit exceeded".into());
1176 assert!(!is_context_overflow(&msg, 0));
1177 }
1178
1179 #[test]
1180 fn is_context_overflow_silent_zai_style() {
1181 let mut msg = AssistantMessage::new("a", "zai", "m", 1);
1182 msg.stop_reason = StopReason::Stop;
1183 msg.usage = Usage {
1184 input: 100_000,
1185 output: 10,
1186 cache_read: 0,
1187 cache_write: 0,
1188 cache_write1h: None,
1189 reasoning: None,
1190 total_tokens: 100_010,
1191 cost: UsageCost::default(),
1192 };
1193 assert!(is_context_overflow(&msg, 80_000));
1194 assert!(!is_context_overflow(&msg, 120_000));
1195 }
1196
1197 #[test]
1198 fn is_context_overflow_length_stop_mimo_style() {
1199 let mut msg = AssistantMessage::new("a", "mimo", "m", 1);
1200 msg.stop_reason = StopReason::Length;
1201 msg.usage = Usage {
1202 input: 8_100,
1203 output: 0,
1204 cache_read: 0,
1205 cache_write: 0,
1206 cache_write1h: None,
1207 reasoning: None,
1208 total_tokens: 8_100,
1209 cost: UsageCost::default(),
1210 };
1211 assert!(!is_context_overflow(&msg, 8_192));
1213 msg.usage.input = 8_110;
1215 assert!(is_context_overflow(&msg, 8_192));
1216 }
1217
1218 #[test]
1219 fn same_model_checks_provider_and_id() {
1220 let model = test_model(8192);
1221 let mut msg = AssistantMessage::new("test-api", "test-provider", "m", 1);
1222 assert!(same_model(&model, &msg));
1223 msg.provider = "other".into();
1224 assert!(!same_model(&model, &msg));
1225 }
1226
1227 #[test]
1228 fn parse_rfc3339_millis_basic() -> TestResult {
1229 let ts = "2025-01-01T00:00:00.000Z";
1230 let millis = parse_rfc3339_millis(ts).ok_or("valid basic RFC3339 timestamp")?;
1231 assert_eq!(millis, 1_735_689_600_000);
1232 Ok(())
1233 }
1234
1235 #[test]
1236 fn parse_rfc3339_millis_with_fractional() -> TestResult {
1237 let ts = "2025-01-01T12:30:45.123Z";
1238 let millis = parse_rfc3339_millis(ts).ok_or("valid fractional RFC3339 timestamp")?;
1239 assert_eq!(millis, 1_735_734_645_123);
1240 Ok(())
1241 }
1242
1243 #[test]
1244 fn thinking_level_str_serializes() {
1245 assert_eq!(
1246 thinking_level_str(ModelThinkingLevel::Off),
1247 Some("off".to_owned())
1248 );
1249 assert_eq!(
1250 thinking_level_str(ModelThinkingLevel::High),
1251 Some("high".to_owned())
1252 );
1253 }
1254
1255 #[tokio::test]
1258 async fn manual_compact_produces_summary_and_rebuilds_messages() -> TestResult {
1259 let session =
1260 session_with_history(8_192, summary_stream_fn("## Goal\nTest summary")).await?;
1261
1262 let (tx, mut rx) = mpsc::unbounded_channel();
1263 let _unsub = session.subscribe(move |event| {
1264 let _ = tx.send(event.type_name().to_owned());
1265 });
1266
1267 let result = session.compact(None).await?;
1268 sleep(std::time::Duration::from_millis(50)).await;
1269
1270 assert!(
1271 result.summary.contains("Test summary"),
1272 "summary should contain stream text: {}",
1273 result.summary
1274 );
1275 assert!(!result.first_kept_entry_id.is_empty());
1276 assert!(result.tokens_before > 0);
1277 assert!(result.estimated_tokens_after.is_some());
1278
1279 let messages = session.agent.transcript();
1280 assert!(
1281 messages.len() < 40,
1282 "transcript should be smaller after compaction: {}",
1283 messages.len()
1284 );
1285
1286 let events = collect_events(&mut rx, 4).await;
1287 assert!(
1288 events.iter().any(|e| e == "compaction_start"),
1289 "missing compaction_start: {events:?}"
1290 );
1291 assert!(
1292 events.iter().any(|e| e == "compaction_end"),
1293 "missing compaction_end: {events:?}"
1294 );
1295 let start_idx = events
1296 .iter()
1297 .position(|e| e == "compaction_start")
1298 .ok_or("compaction_start event position")?;
1299 let end_idx = events
1300 .iter()
1301 .position(|e| e == "compaction_end")
1302 .ok_or("compaction_end event position")?;
1303 assert!(
1304 start_idx < end_idx,
1305 "compaction_start before compaction_end"
1306 );
1307 Ok(())
1308 }
1309
1310 #[tokio::test]
1311 async fn manual_compact_already_compacted_error() -> TestResult {
1312 let session = session_with_history(8_192, summary_stream_fn("summary")).await?;
1313 session.compact(None).await?;
1314 let second = session.compact(None).await;
1315 assert!(
1316 matches!(second, Err(CompactionError::AlreadyCompacted)),
1317 "expected AlreadyCompacted, got: {second:?}"
1318 );
1319 Ok(())
1320 }
1321
1322 #[tokio::test]
1323 async fn manual_compact_too_small_error() -> TestResult {
1324 let messages = vec![
1325 user_text("hi", std::iter::empty()),
1326 assistant_with_usage("hello", Usage::default(), StopReason::Stop),
1327 ];
1328 let session = make_session(8_192, summary_stream_fn("summary"), messages)?;
1329 {
1330 let mut sm = session.session_manager.lock().await;
1331 for msg in session.agent.transcript() {
1332 sm.append_message(&msg)?;
1333 }
1334 }
1335 let result = session.compact(None).await;
1336 assert!(
1337 matches!(result, Err(CompactionError::NothingToCompact)),
1338 "expected NothingToCompact, got: {result:?}"
1339 );
1340 Ok(())
1341 }
1342
1343 #[tokio::test]
1344 async fn manual_compact_cancelled_by_abort() -> TestResult {
1345 let hang_stream: SummarizeStreamFn = Arc::new(move |_model, _ctx, opts| {
1348 let signal = opts.signal.clone();
1349 Box::pin(async move {
1350 if let Some(token) = signal {
1352 token.cancelled().await;
1353 } else {
1354 sleep(std::time::Duration::from_secs(5)).await;
1355 }
1356 let mut msg = AssistantMessage::new("a", "p", "m", 1);
1358 msg.content = vec![AssistantContent::Text(TextContent::new("late"))];
1359 msg.stop_reason = StopReason::Stop;
1360 let stream = stream::iter(vec![Ok(AssistantMessageEvent::Done {
1361 reason: DoneReason::Stop,
1362 message: msg,
1363 })]);
1364 Box::pin(stream)
1365 as Pin<
1366 Box<
1367 dyn futures::Stream<Item = Result<AssistantMessageEvent, ProviderError>>
1368 + Send,
1369 >,
1370 >
1371 })
1372 });
1373 let session = session_with_history(8_192, hang_stream).await?;
1374
1375 let session_clone = Arc::clone(&session);
1377 tokio::spawn(async move {
1378 sleep(std::time::Duration::from_millis(5)).await;
1379 session_clone.abort().await;
1380 });
1381
1382 let result = timeout(std::time::Duration::from_secs(2), session.compact(None)).await?;
1383 assert!(
1384 matches!(result, Err(CompactionError::Cancelled)),
1385 "expected Cancelled, got: {result:?}"
1386 );
1387 Ok(())
1388 }
1389
1390 #[tokio::test]
1391 async fn manual_compact_error_emits_error_message() -> TestResult {
1392 let session =
1393 session_with_history(8_192, error_stream_fn("summarization exploded")).await?;
1394
1395 let (tx, mut rx) = mpsc::unbounded_channel::<AgentSessionEvent>();
1396 let _unsub = session.subscribe(move |event| {
1397 let _ = tx.send(event.clone());
1398 });
1399
1400 let result = session.compact(None).await;
1401 assert!(result.is_err());
1402
1403 sleep(std::time::Duration::from_millis(50)).await;
1404
1405 let mut end_events = Vec::new();
1407 while let Ok(Some(event)) = timeout(std::time::Duration::from_millis(100), rx.recv()).await
1408 {
1409 if event.type_name() == "compaction_end" {
1410 end_events.push(event);
1411 }
1412 }
1413 assert!(!end_events.is_empty(), "should emit compaction_end");
1414 let Some(AgentSessionEvent::CompactionEnd { error_message, .. }) = end_events.first()
1415 else {
1416 return Err("expected CompactionEnd variant".into());
1417 };
1418 assert!(
1419 error_message
1420 .as_ref()
1421 .is_some_and(|m| m.contains("Compaction failed")),
1422 "error message should contain 'Compaction failed': {error_message:?}"
1423 );
1424 Ok(())
1425 }
1426
1427 #[tokio::test]
1428 async fn manual_compact_custom_instructions_passed_to_summary() -> TestResult {
1429 let session = session_with_history(8_192, summary_stream_fn("## Goal\nSummary")).await?;
1430 session.compact(Some("Focus on file changes")).await?;
1431 Ok(())
1432 }
1433
1434 #[tokio::test]
1435 async fn manual_compact_previous_summary_uses_update_prompt() -> TestResult {
1436 let session = session_with_history(8_192, summary_stream_fn("## Goal\nUpdated")).await?;
1437 let first = session.compact(None).await?;
1438 assert!(first.summary.contains("Updated"));
1439 {
1440 let mut sm = session.session_manager.lock().await;
1441 let new_user = user_text("after compaction question", std::iter::empty());
1442 let new_assistant = assistant_with_usage(
1443 "after compaction answer with some content",
1444 Usage {
1445 input: 1000,
1446 output: 500,
1447 ..Usage::default()
1448 },
1449 StopReason::Stop,
1450 );
1451 sm.append_message(&new_user)?;
1452 sm.append_message(&new_assistant)?;
1453 }
1454 session
1456 .agent
1457 .push_message(user_text("after compaction question", std::iter::empty()));
1458 session.agent.push_message(assistant_with_usage(
1459 "after compaction answer with some content",
1460 Usage {
1461 input: 1000,
1462 output: 500,
1463 ..Usage::default()
1464 },
1465 StopReason::Stop,
1466 ));
1467
1468 let second = session.compact(None).await;
1472 assert!(
1476 !matches!(second, Err(CompactionError::AlreadyCompacted)),
1477 "should not be AlreadyCompacted with new messages: {second:?}"
1478 );
1479 Ok(())
1480 }
1481
1482 #[tokio::test]
1483 async fn check_compaction_threshold_triggers_compaction() -> TestResult {
1484 let session = session_with_history(8_192, summary_stream_fn("threshold summary")).await?;
1485 let last_assistant = session
1486 .agent
1487 .last_assistant()
1488 .ok_or("threshold fixture has a last assistant")?;
1489 let _context_tokens = compaction::calculate_context_tokens(&last_assistant.usage);
1490 let should_continue = session.check_compaction(&last_assistant, true).await;
1491 assert!(!should_continue, "threshold should not continue");
1492 assert!(!session.is_compacting(), "should clear compaction latch");
1493 let messages = session.agent.transcript();
1494 assert!(
1495 messages.len() < 40,
1496 "transcript should be smaller after auto-compaction: {}",
1497 messages.len()
1498 );
1499 Ok(())
1500 }
1501
1502 #[tokio::test]
1503 async fn check_compaction_high_usage_aborted_message_is_skipped() -> TestResult {
1504 let session = session_with_history(8_192, summary_stream_fn("should not compact")).await?;
1505 let mut aborted_assistant = session
1506 .agent
1507 .last_assistant()
1508 .ok_or("aborted fixture has a last assistant")?;
1509 aborted_assistant.stop_reason = StopReason::Aborted;
1510 let messages_before = session.agent.transcript().len();
1511
1512 let should_continue = session.check_compaction(&aborted_assistant, true).await;
1513
1514 assert!(!should_continue);
1515 assert_eq!(
1516 session.agent.transcript().len(),
1517 messages_before,
1518 "aborted assistant should not trigger high-usage compaction"
1519 );
1520 Ok(())
1521 }
1522
1523 #[tokio::test]
1524 async fn check_compaction_below_threshold_no_compaction() -> TestResult {
1525 let session = session_with_history(1_000_000, summary_stream_fn("no compact")).await?;
1526 let last_assistant = session
1527 .agent
1528 .last_assistant()
1529 .ok_or("below-threshold fixture has a last assistant")?;
1530 let messages_before = session.agent.transcript().len();
1531 let should_continue = session.check_compaction(&last_assistant, true).await;
1532 assert!(!should_continue);
1533 assert_eq!(
1534 session.agent.transcript().len(),
1535 messages_before,
1536 "transcript should be unchanged"
1537 );
1538 Ok(())
1539 }
1540
1541 #[tokio::test]
1542 async fn check_compaction_overflow_successful_response_compacts_without_retry() -> TestResult {
1543 let session = session_with_history(8_192, summary_stream_fn("overflow recovery")).await?;
1544 let mut overflow_assistant =
1545 AssistantMessage::new("test-api", "test-provider", "m", pi_agent::now_millis());
1546 overflow_assistant.stop_reason = StopReason::Stop;
1547 overflow_assistant.usage = Usage {
1548 input: 10_000,
1549 output: 500,
1550 cache_read: 0,
1551 cache_write: 0,
1552 cache_write1h: None,
1553 reasoning: None,
1554 total_tokens: 10_500,
1555 cost: UsageCost::default(),
1556 };
1557 overflow_assistant.content = vec![AssistantContent::Text(TextContent::new("overflow"))];
1558 let should_continue = session.check_compaction(&overflow_assistant, true).await;
1559 assert!(!should_continue, "successful overflow should not retry");
1560 assert!(!session.is_compacting(), "should clear compaction latch");
1561 let messages = session.agent.transcript();
1562 assert!(
1563 messages.len() < 40,
1564 "transcript should be smaller: {}",
1565 messages.len()
1566 );
1567 Ok(())
1568 }
1569
1570 #[tokio::test]
1571 async fn check_compaction_overflow_error_compacts_and_retries_once() -> TestResult {
1572 let session = session_with_history(8_192, summary_stream_fn("overflow recovery")).await?;
1573 let overflow_assistant = assistant_overflow_message();
1574 let overflow_msg = pi_agent::AgentMessage::Llm(Box::new(pi_ai::Message::Assistant(
1575 overflow_assistant.clone(),
1576 )));
1577 session.agent.push_message(overflow_msg.clone());
1578 {
1579 let mut sm = session.session_manager.lock().await;
1580 sm.append_message(&overflow_msg)?;
1581 }
1582 let should_continue = session.check_compaction(&overflow_assistant, true).await;
1583 assert!(should_continue, "error overflow should signal continuation");
1584 Ok(())
1585 }
1586
1587 #[tokio::test]
1588 async fn check_compaction_overflow_second_attempt_emits_terminal_error() -> TestResult {
1589 let session = session_with_history(8_192, summary_stream_fn("overflow recovery")).await?;
1590 {
1591 let mut inner = session.lock_inner();
1592 inner.overflow_recovery_attempted = true;
1593 }
1594 let overflow_assistant = assistant_overflow_message();
1595 let (tx, mut rx) = mpsc::unbounded_channel::<AgentSessionEvent>();
1596 let _unsub = session.subscribe(move |event| {
1597 let _ = tx.send(event.clone());
1598 });
1599 let should_continue = session.check_compaction(&overflow_assistant, true).await;
1600 assert!(!should_continue, "second overflow should not continue");
1601 sleep(std::time::Duration::from_millis(50)).await;
1602 let mut ends = Vec::new();
1603 while let Ok(Some(event)) = timeout(std::time::Duration::from_millis(100), rx.recv()).await
1604 {
1605 if event.type_name() == "compaction_end" {
1606 ends.push(event);
1607 }
1608 }
1609 assert!(!ends.is_empty(), "should emit compaction_end");
1610 let Some(AgentSessionEvent::CompactionEnd {
1611 error_message,
1612 reason,
1613 ..
1614 }) = ends.first()
1615 else {
1616 return Err("expected terminal CompactionEnd event".into());
1617 };
1618 assert_eq!(*reason, CompactionReason::Overflow);
1619 assert!(
1620 error_message.as_ref().is_some_and(|m| m
1621 .contains("Context overflow recovery failed after one compact-and-retry attempt")),
1622 "exact terminal error: {error_message:?}"
1623 );
1624 Ok(())
1625 }
1626
1627 #[tokio::test]
1628 async fn check_compaction_different_model_skips_overflow() -> TestResult {
1629 let session = session_with_history(1_000_000, summary_stream_fn("no compact")).await?;
1632
1633 let mut other_model_msg = AssistantMessage::new(
1634 "test-api",
1635 "different-provider",
1636 "other-model",
1637 pi_agent::now_millis(),
1638 );
1639 other_model_msg.stop_reason = StopReason::Error;
1640 other_model_msg.error_message = Some("prompt is too long".into());
1641
1642 let messages_before = session.agent.transcript().len();
1643 let should_continue = session.check_compaction(&other_model_msg, true).await;
1644
1645 assert!(!should_continue, "different model should not compact");
1646 assert_eq!(
1647 session.agent.transcript().len(),
1648 messages_before,
1649 "transcript unchanged"
1650 );
1651 Ok(())
1652 }
1653
1654 #[tokio::test]
1655 async fn check_compaction_stale_pre_compaction_usage_ignored() -> TestResult {
1656 let session = session_with_history(8_192, summary_stream_fn("first compact")).await?;
1657 session.compact(None).await?;
1658 let mut old_msg = AssistantMessage::new("test-api", "test-provider", "m", 1);
1659 old_msg.stop_reason = StopReason::Stop;
1660 old_msg.usage = Usage {
1661 input: 10_000,
1662 output: 500,
1663 ..Usage::default()
1664 };
1665 let should_continue = session.check_compaction(&old_msg, true).await;
1666 assert!(
1667 !should_continue,
1668 "stale pre-compaction message should not trigger compaction"
1669 );
1670 Ok(())
1671 }
1672
1673 #[tokio::test]
1674 async fn check_compaction_disabled_returns_false() -> TestResult {
1675 let session = session_with_history(8_192, summary_stream_fn("disabled")).await?;
1676 session.set_auto_compaction_enabled(false);
1677 assert!(
1678 !session.lock_settings().get_compaction_settings().enabled,
1679 "runtime toggle must update persisted settings state"
1680 );
1681
1682 let last_assistant = session
1683 .agent
1684 .last_assistant()
1685 .ok_or("disabled fixture has a last assistant")?;
1686 let messages_before = session.agent.transcript().len();
1687
1688 let should_continue = session.check_compaction(&last_assistant, true).await;
1689
1690 assert!(!should_continue);
1691 assert_eq!(
1692 session.agent.transcript().len(),
1693 messages_before,
1694 "transcript unchanged when disabled"
1695 );
1696 Ok(())
1697 }
1698
1699 #[tokio::test]
1700 async fn manual_compact_no_model_runtime_fails() -> TestResult {
1701 let mut messages = Vec::new();
1702 for i in 0..4 {
1703 messages.push(user_text(
1704 format!("user message {i} padding for compaction"),
1705 std::iter::empty(),
1706 ));
1707 messages.push(assistant_with_usage(
1708 format!("assistant reply {i} with content").as_str(),
1709 Usage {
1710 input: 100 + i * 10,
1711 output: 50,
1712 total_tokens: 150 + i * 10,
1713 ..Usage::default()
1714 },
1715 StopReason::Stop,
1716 ));
1717 }
1718 let provider = mock_provider();
1719 let mut config = AgentSessionConfig::test_config(provider, test_model(8_192))?;
1720 config.messages = messages;
1721 let session = AgentSession::new(config)?;
1722 {
1723 let mut settings = session.lock_settings();
1724 let mut compaction = Map::new();
1725 compaction.insert("keepRecentTokens".into(), Value::from(1u64));
1726 let mut overrides = Map::new();
1727 overrides.insert("compaction".into(), Value::Object(compaction));
1728 settings.apply_overrides(&overrides);
1729 }
1730 {
1731 let mut sm = session.session_manager.lock().await;
1732 for msg in session.agent.transcript() {
1733 sm.append_message(&msg)?;
1734 }
1735 }
1736 let result = session.compact(None).await;
1737 assert!(
1738 matches!(result, Err(CompactionError::SummarizationFailed(_))),
1739 "expected SummarizationFailed without model_runtime: {result:?}"
1740 );
1741 Ok(())
1742 }
1743
1744 #[tokio::test]
1745 async fn run_auto_compaction_queued_message_continuation() -> TestResult {
1746 let session = session_with_history(8_192, summary_stream_fn("queued summary")).await?;
1747 session
1748 .agent
1749 .follow_up(user_text("queued follow-up", std::iter::empty()));
1750 session.mirror_follow_up_push("queued follow-up".into());
1751 let last_assistant = session
1752 .agent
1753 .last_assistant()
1754 .ok_or("queued fixture has a last assistant")?;
1755 let should_continue = session.check_compaction(&last_assistant, true).await;
1756 assert!(
1757 should_continue,
1758 "should continue when queued messages exist after auto-compaction"
1759 );
1760 Ok(())
1761 }
1762
1763 #[tokio::test]
1764 async fn manual_compact_persists_entry_to_session_tree() -> TestResult {
1765 let session = session_with_history(8_192, summary_stream_fn("persisted")).await?;
1766 let result = session.compact(None).await?;
1767 let entries: Vec<SessionEntry> = {
1768 let sm = session.session_manager.lock().await;
1769 sm.get_entries().into_iter().cloned().collect()
1770 };
1771 let compaction_entries: Vec<_> = entries
1772 .iter()
1773 .filter(|e| e.discriminant() == "compaction")
1774 .collect();
1775 assert_eq!(compaction_entries.len(), 1, "exactly one compaction entry");
1776 let Some(SessionEntry::Compaction(compaction)) = compaction_entries.first().copied() else {
1777 return Err("expected Compaction entry".into());
1778 };
1779 assert_eq!(compaction.summary, result.summary);
1780 assert_eq!(compaction.first_kept_entry_id, result.first_kept_entry_id);
1781 Ok(())
1782 }
1783
1784 #[tokio::test]
1785 async fn manual_compact_emits_entry_appended() -> TestResult {
1786 let session = session_with_history(8_192, summary_stream_fn("entry event")).await?;
1787
1788 let (tx, mut rx) = mpsc::unbounded_channel::<String>();
1789 let _unsub = session.subscribe(move |event| {
1790 let _ = tx.send(event.type_name().to_owned());
1791 });
1792
1793 let _ = session.compact(None).await;
1794
1795 sleep(std::time::Duration::from_millis(50)).await;
1796
1797 let events = collect_events(&mut rx, 10).await;
1798 assert!(
1799 events.iter().any(|e| e == "entry_appended"),
1800 "should emit entry_appended for the compaction entry: {events:?}"
1801 );
1802 Ok(())
1803 }
1804}