1use crate::skill_history::{is_skill_tool_name, SkillNameMatch};
8use crate::{TiktokenTokenCounter, TokenBudget, TokenCounter};
9use bamboo_domain::{
10 sha256_hex, CompressionEvent, CompressionEventKind, CompressionTriggerType, Message,
11 MessagePart, ModelContextResetReason, Role, Session, TokenBudgetUsage,
12};
13use chrono::Utc;
14use std::collections::{BTreeMap, HashMap, HashSet};
15use thiserror::Error;
16
17#[derive(Debug, Clone, Copy, PartialEq, Eq)]
22pub struct RetrievalWindowPolicy {
23 pub min_recent_user_turns: usize,
25 pub target_usage_percent: u8,
27}
28
29#[derive(Debug, Clone, Default, PartialEq, Eq)]
36pub struct RetrievalWindowTokenAccounting {
37 pub fixed_prompt_tokens: u32,
40 pub boundary_reclaimable_tokens: u32,
44 pub provider_message_tokens: BTreeMap<String, u32>,
46}
47
48#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize)]
50pub struct RetrievalWindowCandidatePlan {
51 pub message_ids_to_archive: Vec<String>,
53 pub archive_message_count: usize,
55 pub archive_group_count: usize,
57 pub archive_user_turn_count: usize,
59 pub archive_message_tokens: u32,
61 pub active_message_count: usize,
63 pub active_state_sha256: String,
66 pub token_accounting_sha256: String,
69 pub active_tokens_before: u32,
71 pub projected_active_tokens_after: u32,
73 pub context_window_tokens: u32,
75 pub request_input_limit_tokens: u32,
77 pub target_tokens: u32,
79 pub target_usage_percent: u8,
81 pub min_recent_user_turns: usize,
83 pub fixed_prompt_tokens: u32,
85 pub boundary_reclaimable_tokens: u32,
88 pub system_message_tokens: u32,
90 pub provider_message_token_override_count: usize,
92 pub protected_active_tokens: u32,
94 pub retained_recent_user_turn_count: usize,
96 pub retained_user_turn_count: usize,
98 pub oldest_retained_message_id: Option<String>,
100 pub oldest_retained_user_message_id: Option<String>,
102 pub incomplete_protocol_group_count: usize,
104 #[serde(skip)]
111 evidence_sha256: String,
112}
113
114#[derive(Debug, Clone, PartialEq, Eq, Error)]
116pub enum RetrievalWindowPlanError {
117 #[error("retrieval-window requires at least one recent user turn")]
118 InvalidRecentUserTurnFloor,
119 #[error("retrieval-window target usage percent must be between 1 and 100 (got {target_usage_percent})")]
120 InvalidTargetUsagePercent { target_usage_percent: u8 },
121 #[error(
122 "invalid token budget for retrieval-window planning (context={context_window_tokens}, request_input={max_request_input_tokens})"
123 )]
124 InvalidTokenBudget {
125 context_window_tokens: u32,
126 max_request_input_tokens: u32,
127 },
128 #[error("token accounting overflow while planning retrieval-window archive candidates")]
129 TokenAccountingOverflow,
130 #[error("provider-aware token accounting is required for image message {message_id}")]
131 MissingProviderMessageTokenEstimate { message_id: String },
132 #[error("no active messages are available for retrieval-window planning")]
133 NoActiveMessages,
134 #[error(
135 "retrieval-window target is already satisfied (active={active_tokens}, target={target_tokens})"
136 )]
137 TargetAlreadySatisfied {
138 active_tokens: u32,
139 target_tokens: u32,
140 },
141 #[error(
142 "no eligible active logical group can be archived (active={active_tokens}, target={target_tokens}, protected={protected_active_tokens}, incomplete_protocol_groups={incomplete_protocol_group_count})"
143 )]
144 NothingToArchive {
145 active_tokens: u32,
146 target_tokens: u32,
147 protected_active_tokens: u32,
148 incomplete_protocol_group_count: usize,
149 },
150 #[error(
151 "protected active content prevents the retrieval-window target (projected={projected_tokens}, target={target_tokens}, protected={protected_active_tokens}, fixed={fixed_prompt_tokens}, incomplete_protocol_groups={incomplete_protocol_group_count})"
152 )]
153 ProtectedContentExceedsTarget {
154 projected_tokens: u32,
155 target_tokens: u32,
156 protected_active_tokens: u32,
157 fixed_prompt_tokens: u32,
158 incomplete_protocol_group_count: usize,
159 },
160}
161
162#[derive(Debug, Clone, PartialEq, Eq)]
164pub struct RetrievalWindowApplyResult {
165 pub event_id: String,
166 pub newly_archived_message_count: usize,
167 pub idempotent_replay: bool,
168}
169
170#[derive(Debug, Clone, PartialEq, Eq, Error)]
172pub enum RetrievalWindowApplyError {
173 #[error("retrieval-window plan has no candidate messages")]
174 EmptyCandidateSet,
175 #[error("retrieval-window plan candidate at index {index} has an empty message ID")]
176 EmptyCandidateId { index: usize },
177 #[error("retrieval-window plan repeats candidate message ID {message_id}")]
178 DuplicateCandidateId { message_id: String },
179 #[error("retrieval-window candidate message {message_id} is missing from the session")]
180 MissingCandidateMessage { message_id: String },
181 #[error("retrieval-window candidate message ID {message_id} is not unique in the session")]
182 DuplicateSessionMessageId { message_id: String },
183 #[error("retrieval-window candidate message {message_id} is out of session order")]
184 CandidateOrderMismatch { message_id: String },
185 #[error("retrieval-window plan evidence is inconsistent: {field}")]
186 InconsistentPlan { field: &'static str },
187 #[error("session already has a conversation summary")]
188 PreExistingConversationSummary,
189 #[error("retrieval-window plan was only partially applied")]
190 PartialApplication,
191 #[error("active candidate message {message_id} is already correlated to an archive event")]
192 ActiveCandidateAlreadyCorrelated { message_id: String },
193 #[error("retrieval-window candidate message {message_id} is a system message")]
194 SystemMessageCandidate { message_id: String },
195 #[error("retrieval-window candidate message {message_id} is marked never_compress")]
196 NeverCompressCandidate { message_id: String },
197 #[error("retrieval-window candidates do not cover a complete logical group")]
198 IncompleteLogicalGroup,
199 #[error("retrieval-window candidates contain an unsafe tool call/result chain")]
200 UnsafeToolProtocol,
201 #[error("retrieval-window candidates contain a protected Skill-loading chain")]
202 ProtectedSkillChain,
203 #[error("retrieval-window candidates include a retained recent user turn")]
204 RecentUserTurnProtected,
205 #[error("retrieval-window plan is stale: {invariant}")]
206 StalePlan { invariant: &'static str },
207 #[error("retrieval-window replay candidates refer to different archive events")]
208 MixedArchiveEvents,
209 #[error("retrieval-window replay event {event_id} is missing or duplicated")]
210 MissingOrDuplicateArchiveEvent { event_id: String },
211 #[error("archive event {event_id} was produced by a different compression strategy")]
212 ArchiveEventKindConflict { event_id: String },
213 #[error("retrieval-window replay evidence does not match event {event_id}: {field}")]
214 ArchiveEventEvidenceMismatch {
215 event_id: String,
216 field: &'static str,
217 },
218 #[error("token accounting overflow while applying retrieval-window plan")]
219 TokenAccountingOverflow,
220}
221
222#[derive(Debug, Clone, Copy)]
223struct IndexedMessage<'a> {
224 session_index: usize,
225 message: &'a Message,
226}
227
228#[derive(Debug)]
229struct LogicalGroup<'a> {
230 messages: Vec<IndexedMessage<'a>>,
231 user_message_id: Option<String>,
232 protocol_safe: bool,
233 protected: bool,
234 token_count: u32,
235}
236
237impl<'a> LogicalGroup<'a> {
238 fn preamble(message: IndexedMessage<'a>) -> Self {
239 Self {
240 messages: vec![message],
241 user_message_id: None,
242 protocol_safe: true,
243 protected: false,
244 token_count: 0,
245 }
246 }
247
248 fn user_turn(message: IndexedMessage<'a>) -> Self {
249 Self {
250 user_message_id: Some(message.message.id.clone()),
251 messages: vec![message],
252 protocol_safe: true,
253 protected: false,
254 token_count: 0,
255 }
256 }
257}
258
259pub fn build_retrieval_window_candidate_plan(
261 session: &Session,
262 budget: &TokenBudget,
263 policy: RetrievalWindowPolicy,
264) -> Result<RetrievalWindowCandidatePlan, RetrievalWindowPlanError> {
265 build_retrieval_window_candidate_plan_with_token_accounting(
266 session,
267 budget,
268 policy,
269 &RetrievalWindowTokenAccounting::default(),
270 )
271}
272
273pub fn build_retrieval_window_candidate_plan_with_fixed_tokens(
276 session: &Session,
277 budget: &TokenBudget,
278 policy: RetrievalWindowPolicy,
279 fixed_prompt_tokens: u32,
280) -> Result<RetrievalWindowCandidatePlan, RetrievalWindowPlanError> {
281 build_retrieval_window_candidate_plan_with_token_accounting(
282 session,
283 budget,
284 policy,
285 &RetrievalWindowTokenAccounting {
286 fixed_prompt_tokens,
287 ..RetrievalWindowTokenAccounting::default()
288 },
289 )
290}
291
292pub fn build_retrieval_window_candidate_plan_with_token_accounting(
299 session: &Session,
300 budget: &TokenBudget,
301 policy: RetrievalWindowPolicy,
302 accounting: &RetrievalWindowTokenAccounting,
303) -> Result<RetrievalWindowCandidatePlan, RetrievalWindowPlanError> {
304 build_retrieval_window_candidate_plan_with_counter(
305 session,
306 budget,
307 policy,
308 accounting,
309 false,
310 &TiktokenTokenCounter::default(),
311 )
312}
313
314pub fn build_retrieval_window_critical_overflow_plan_with_token_accounting(
326 session: &Session,
327 budget: &TokenBudget,
328 policy: RetrievalWindowPolicy,
329 accounting: &RetrievalWindowTokenAccounting,
330) -> Result<RetrievalWindowCandidatePlan, RetrievalWindowPlanError> {
331 build_retrieval_window_candidate_plan_with_counter(
332 session,
333 budget,
334 policy,
335 accounting,
336 true,
337 &TiktokenTokenCounter::default(),
338 )
339}
340
341fn build_retrieval_window_candidate_plan_with_counter(
342 session: &Session,
343 budget: &TokenBudget,
344 policy: RetrievalWindowPolicy,
345 accounting: &RetrievalWindowTokenAccounting,
346 force_archive_after_provider_overflow: bool,
347 counter: &impl TokenCounter,
348) -> Result<RetrievalWindowCandidatePlan, RetrievalWindowPlanError> {
349 validate_inputs(budget, policy)?;
350
351 let (system_messages, mut groups, active_message_count) = build_logical_groups(session);
352 if active_message_count == 0 {
353 return Err(RetrievalWindowPlanError::NoActiveMessages);
354 }
355
356 let system_tokens = count_indexed_messages(&system_messages, accounting, counter)?;
357 for group in &mut groups {
358 group.token_count = count_indexed_messages(&group.messages, accounting, counter)?;
359 }
360 let group_tokens = checked_sum(groups.iter().map(|group| group.token_count))?;
361 let post_boundary_active_tokens = accounting
362 .fixed_prompt_tokens
363 .checked_add(system_tokens)
364 .and_then(|tokens| tokens.checked_add(group_tokens))
365 .ok_or(RetrievalWindowPlanError::TokenAccountingOverflow)?;
366 let active_tokens_before = post_boundary_active_tokens
367 .checked_add(accounting.boundary_reclaimable_tokens)
368 .ok_or(RetrievalWindowPlanError::TokenAccountingOverflow)?;
369 mark_protocol_safety(&mut groups);
370 mark_protected_groups(&mut groups, policy.min_recent_user_turns);
371
372 let protected_group_tokens = checked_sum(
373 groups
374 .iter()
375 .filter(|group| group.protected)
376 .map(|group| group.token_count),
377 )?;
378 let protected_active_tokens = system_tokens
379 .checked_add(protected_group_tokens)
380 .ok_or(RetrievalWindowPlanError::TokenAccountingOverflow)?;
381 let incomplete_protocol_group_count =
382 groups.iter().filter(|group| !group.protocol_safe).count();
383 let configured_target_tokens =
384 effective_retrieval_window_target_tokens(budget, policy.target_usage_percent);
385 let minimum_projected_active_tokens = accounting
386 .fixed_prompt_tokens
387 .checked_add(protected_active_tokens)
388 .ok_or(RetrievalWindowPlanError::TokenAccountingOverflow)?;
389 if force_archive_after_provider_overflow
390 && minimum_projected_active_tokens > budget.max_request_input_tokens()
391 {
392 return Err(RetrievalWindowPlanError::ProtectedContentExceedsTarget {
393 projected_tokens: minimum_projected_active_tokens,
394 target_tokens: budget.max_request_input_tokens(),
395 protected_active_tokens,
396 fixed_prompt_tokens: accounting.fixed_prompt_tokens,
397 incomplete_protocol_group_count,
398 });
399 }
400 let target_tokens = if force_archive_after_provider_overflow {
401 minimum_projected_active_tokens
402 } else {
403 configured_target_tokens
404 };
405
406 if active_tokens_before <= target_tokens {
407 return Err(RetrievalWindowPlanError::TargetAlreadySatisfied {
408 active_tokens: active_tokens_before,
409 target_tokens,
410 });
411 }
412
413 let mut selected_group_indexes = HashSet::new();
414 let mut message_ids_to_archive = Vec::new();
415 let mut archive_message_tokens = 0u32;
416 let mut projected_active_tokens_after = post_boundary_active_tokens;
421 let mut archive_user_turn_count = 0usize;
422
423 for (group_index, group) in groups.iter().enumerate() {
424 if projected_active_tokens_after <= target_tokens && !message_ids_to_archive.is_empty() {
425 break;
426 }
427 if group.protected {
428 continue;
429 }
430
431 selected_group_indexes.insert(group_index);
432 archive_message_tokens = archive_message_tokens
433 .checked_add(group.token_count)
434 .ok_or(RetrievalWindowPlanError::TokenAccountingOverflow)?;
435 projected_active_tokens_after = projected_active_tokens_after
436 .checked_sub(group.token_count)
437 .ok_or(RetrievalWindowPlanError::TokenAccountingOverflow)?;
438 if group.user_message_id.is_some() {
439 archive_user_turn_count += 1;
440 }
441 message_ids_to_archive.extend(
442 group
443 .messages
444 .iter()
445 .map(|indexed| indexed.message.id.clone()),
446 );
447 }
448
449 if message_ids_to_archive.is_empty() {
450 return Err(RetrievalWindowPlanError::NothingToArchive {
451 active_tokens: active_tokens_before,
452 target_tokens,
453 protected_active_tokens,
454 incomplete_protocol_group_count,
455 });
456 }
457 if projected_active_tokens_after > target_tokens {
458 return Err(RetrievalWindowPlanError::ProtectedContentExceedsTarget {
459 projected_tokens: projected_active_tokens_after,
460 target_tokens,
461 protected_active_tokens,
462 fixed_prompt_tokens: accounting.fixed_prompt_tokens,
463 incomplete_protocol_group_count,
464 });
465 }
466
467 let oldest_retained_message_id = groups
468 .iter()
469 .enumerate()
470 .filter(|(index, _)| !selected_group_indexes.contains(index))
471 .flat_map(|(_, group)| group.messages.iter())
472 .min_by_key(|indexed| indexed.session_index)
473 .map(|indexed| indexed.message.id.clone());
474 let oldest_retained_user_message_id = groups
475 .iter()
476 .enumerate()
477 .find(|(index, group)| {
478 !selected_group_indexes.contains(index) && group.user_message_id.is_some()
479 })
480 .and_then(|(_, group)| group.user_message_id.clone());
481 let retained_user_turn_count = groups
482 .iter()
483 .enumerate()
484 .filter(|(index, group)| {
485 !selected_group_indexes.contains(index) && group.user_message_id.is_some()
486 })
487 .count();
488 let total_user_turn_count = groups
489 .iter()
490 .filter(|group| group.user_message_id.is_some())
491 .count();
492
493 let mut plan = RetrievalWindowCandidatePlan {
494 archive_message_count: message_ids_to_archive.len(),
495 archive_group_count: selected_group_indexes.len(),
496 archive_user_turn_count,
497 archive_message_tokens,
498 active_message_count,
499 active_state_sha256: active_state_sha256(session),
500 token_accounting_sha256: token_accounting_sha256(accounting),
501 active_tokens_before,
502 projected_active_tokens_after,
503 context_window_tokens: budget.max_context_tokens,
504 request_input_limit_tokens: budget.max_request_input_tokens(),
505 target_tokens,
506 target_usage_percent: policy.target_usage_percent,
507 min_recent_user_turns: policy.min_recent_user_turns,
508 fixed_prompt_tokens: accounting.fixed_prompt_tokens,
509 boundary_reclaimable_tokens: accounting.boundary_reclaimable_tokens,
510 system_message_tokens: system_tokens,
511 provider_message_token_override_count: system_messages
512 .iter()
513 .chain(groups.iter().flat_map(|group| group.messages.iter()))
514 .filter(|indexed| {
515 accounting
516 .provider_message_tokens
517 .contains_key(&indexed.message.id)
518 })
519 .count(),
520 protected_active_tokens,
521 retained_recent_user_turn_count: total_user_turn_count.min(policy.min_recent_user_turns),
522 retained_user_turn_count,
523 oldest_retained_message_id,
524 oldest_retained_user_message_id,
525 incomplete_protocol_group_count,
526 message_ids_to_archive,
527 evidence_sha256: String::new(),
528 };
529 plan.evidence_sha256 = plan_evidence_sha256(&plan);
530 Ok(plan)
531}
532
533#[derive(Debug, Clone, Copy)]
534struct ValidatedRetrievalWindowUsage {
535 system_tokens: u32,
536 window_tokens: u32,
537}
538
539pub fn apply_retrieval_window_plan(
549 session: &mut Session,
550 plan: &RetrievalWindowCandidatePlan,
551 current_policy: RetrievalWindowPolicy,
552 current_budget: &TokenBudget,
553 current_accounting: &RetrievalWindowTokenAccounting,
554) -> Result<RetrievalWindowApplyResult, RetrievalWindowApplyError> {
555 apply_retrieval_window_plan_with_trigger(
556 session,
557 plan,
558 current_policy,
559 current_budget,
560 current_accounting,
561 CompressionTriggerType::Auto,
562 )
563}
564
565pub fn apply_retrieval_window_plan_with_trigger(
569 session: &mut Session,
570 plan: &RetrievalWindowCandidatePlan,
571 current_policy: RetrievalWindowPolicy,
572 current_budget: &TokenBudget,
573 current_accounting: &RetrievalWindowTokenAccounting,
574 trigger_type: CompressionTriggerType,
575) -> Result<RetrievalWindowApplyResult, RetrievalWindowApplyError> {
576 let usage = validate_apply_plan_arithmetic(plan)?;
577 if session.conversation_summary.is_some() {
578 return Err(RetrievalWindowApplyError::PreExistingConversationSummary);
579 }
580 validate_current_policy(plan, current_policy)?;
581 validate_current_token_budget(plan, current_budget)?;
582 validate_current_token_accounting(plan, current_accounting)?;
583
584 let candidate_indexes = resolve_candidate_indexes(session, plan)?;
585 let active_count = candidate_indexes
586 .iter()
587 .filter(|index| !session.messages[**index].compressed)
588 .count();
589 let archived_count = candidate_indexes.len().saturating_sub(active_count);
590
591 if active_count > 0 && archived_count > 0 {
592 return Err(RetrievalWindowApplyError::PartialApplication);
593 }
594 if archived_count == candidate_indexes.len() {
595 validate_plan_evidence(plan)?;
596 return validate_idempotent_replay(session, plan, &candidate_indexes, &trigger_type);
597 }
598
599 for index in &candidate_indexes {
600 let message = &session.messages[*index];
601 if message.compressed_by_event_id.is_some() {
602 return Err(
603 RetrievalWindowApplyError::ActiveCandidateAlreadyCorrelated {
604 message_id: message.id.clone(),
605 },
606 );
607 }
608 if matches!(message.role, Role::System) {
609 return Err(RetrievalWindowApplyError::SystemMessageCandidate {
610 message_id: message.id.clone(),
611 });
612 }
613 if message.never_compress {
614 return Err(RetrievalWindowApplyError::NeverCompressCandidate {
615 message_id: message.id.clone(),
616 });
617 }
618 }
619
620 validate_active_plan_structure(session, plan)?;
621 validate_plan_evidence(plan)?;
622
623 let mut event = CompressionEvent::new(
624 plan.archive_message_count,
625 plan.archive_group_count,
626 usage_percentage(plan.active_tokens_before, plan.context_window_tokens),
627 usage_percentage(
628 plan.projected_active_tokens_after,
629 plan.context_window_tokens,
630 ),
631 0,
632 trigger_type,
633 0.0,
634 None,
635 0,
636 );
637 event.kind = CompressionEventKind::RetrievalWindow;
638 event.source_tokens = plan.archive_message_tokens;
639 event.fixed_prompt_tokens = plan.fixed_prompt_tokens;
640 event.retrieval_active_tokens_before = plan.active_tokens_before;
641 event.retrieval_active_message_count_before = plan.active_message_count;
642 event.retrieval_active_state_sha256 = Some(plan.active_state_sha256.clone());
643 event.retrieval_token_accounting_sha256 = Some(plan.token_accounting_sha256.clone());
644 event.retrieval_active_tokens_after = plan.projected_active_tokens_after;
645 event.retrieval_target_tokens = plan.target_tokens;
646 event.retrieval_target_usage_percent = plan.target_usage_percent;
647 event.retrieval_archived_group_count = plan.archive_group_count;
648 event.retrieval_archived_user_turn_count = plan.archive_user_turn_count;
649 event.retrieval_archived_message_tokens = plan.archive_message_tokens;
650 event.retrieval_boundary_reclaimed_tokens = plan.boundary_reclaimable_tokens;
651 event.retrieval_system_message_tokens = plan.system_message_tokens;
652 event.retrieval_context_window_tokens = plan.context_window_tokens;
653 event.retrieval_request_input_limit_tokens = plan.request_input_limit_tokens;
654 event.retrieval_min_recent_user_turns = plan.min_recent_user_turns;
655 event.retrieval_retained_recent_user_turn_count = plan.retained_recent_user_turn_count;
656 event.retrieval_retained_user_turn_count = plan.retained_user_turn_count;
657 event
658 .retrieval_oldest_retained_message_id
659 .clone_from(&plan.oldest_retained_message_id);
660 event
661 .retrieval_oldest_retained_user_message_id
662 .clone_from(&plan.oldest_retained_user_message_id);
663 event.retrieval_provider_message_token_override_count =
664 plan.provider_message_token_override_count;
665 event.retrieval_protected_active_tokens = plan.protected_active_tokens;
666 event.retrieval_incomplete_protocol_group_count = plan.incomplete_protocol_group_count;
667 let event_id = event.id.clone();
668
669 for index in candidate_indexes {
670 session.messages[index].compressed = true;
671 session.messages[index].compressed_by_event_id = Some(event_id.clone());
672 }
673 session.compression_events.push(event);
674 let previous_cache_read_input_tokens = session
675 .token_usage
676 .as_ref()
677 .map(|usage| usage.cache_read_input_tokens)
678 .unwrap_or(0);
679 let previous_provider_prompt_usage = session
680 .token_usage
681 .as_ref()
682 .and_then(|usage| usage.provider_prompt_usage)
683 .map(|mut usage| {
684 usage.retained_from_previous_call = true;
685 usage
686 });
687 session.token_usage = Some(TokenBudgetUsage {
688 system_tokens: usage.system_tokens,
689 summary_tokens: 0,
690 window_tokens: usage.window_tokens,
691 total_tokens: plan.projected_active_tokens_after,
692 max_context_tokens: plan.context_window_tokens,
693 budget_limit: plan.request_input_limit_tokens,
694 truncation_occurred: false,
695 segments_removed: plan.archive_group_count,
696 prompt_cached_tool_outputs: 0,
697 prompt_cached_tool_tokens_saved: 0,
698 thinking_tokens: 0,
699 cache_read_input_tokens: previous_cache_read_input_tokens,
700 provider_prompt_usage: previous_provider_prompt_usage,
701 });
702 session.reset_model_context_epoch(ModelContextResetReason::Compression);
703 session.updated_at = Utc::now();
704
705 Ok(RetrievalWindowApplyResult {
706 event_id,
707 newly_archived_message_count: plan.archive_message_count,
708 idempotent_replay: false,
709 })
710}
711
712fn validate_apply_plan_arithmetic(
713 plan: &RetrievalWindowCandidatePlan,
714) -> Result<ValidatedRetrievalWindowUsage, RetrievalWindowApplyError> {
715 if plan.message_ids_to_archive.is_empty() {
716 return Err(RetrievalWindowApplyError::EmptyCandidateSet);
717 }
718 if plan.archive_message_count != plan.message_ids_to_archive.len() {
719 return Err(RetrievalWindowApplyError::InconsistentPlan {
720 field: "archive_message_count",
721 });
722 }
723 if plan.archive_group_count == 0 || plan.archive_group_count > plan.archive_message_count {
724 return Err(RetrievalWindowApplyError::InconsistentPlan {
725 field: "archive_group_count",
726 });
727 }
728 if plan.archive_user_turn_count > plan.archive_group_count {
729 return Err(RetrievalWindowApplyError::InconsistentPlan {
730 field: "archive_user_turn_count",
731 });
732 }
733 if plan.active_message_count < plan.archive_message_count {
734 return Err(RetrievalWindowApplyError::InconsistentPlan {
735 field: "active_message_count",
736 });
737 }
738 if !(1..=100).contains(&plan.target_usage_percent) {
739 return Err(RetrievalWindowApplyError::InconsistentPlan {
740 field: "target_usage_percent",
741 });
742 }
743 if plan.min_recent_user_turns == 0 {
744 return Err(RetrievalWindowApplyError::InconsistentPlan {
745 field: "min_recent_user_turns",
746 });
747 }
748 if plan.context_window_tokens == 0
749 || plan.request_input_limit_tokens == 0
750 || plan.request_input_limit_tokens > plan.context_window_tokens
751 {
752 return Err(RetrievalWindowApplyError::InconsistentPlan {
753 field: "request_input_limit_tokens",
754 });
755 }
756 let expected_target = effective_target_tokens(
757 plan.context_window_tokens,
758 plan.request_input_limit_tokens,
759 plan.target_usage_percent,
760 );
761 let protected_with_fixed = plan
767 .fixed_prompt_tokens
768 .checked_add(plan.protected_active_tokens)
769 .ok_or(RetrievalWindowApplyError::TokenAccountingOverflow)?;
770 let sealed_provider_overflow_plan = protected_with_fixed <= plan.request_input_limit_tokens
771 && plan.target_tokens == protected_with_fixed
772 && plan.projected_active_tokens_after == protected_with_fixed;
773 if plan.target_tokens > expected_target && !sealed_provider_overflow_plan {
774 return Err(RetrievalWindowApplyError::InconsistentPlan {
775 field: "target_tokens",
776 });
777 }
778 if plan.active_tokens_before <= plan.target_tokens
779 || plan.projected_active_tokens_after > plan.target_tokens
780 || plan.archive_message_tokens == 0
781 || plan
782 .active_tokens_before
783 .checked_sub(plan.boundary_reclaimable_tokens)
784 .and_then(|tokens| tokens.checked_sub(plan.archive_message_tokens))
785 != Some(plan.projected_active_tokens_after)
786 {
787 return Err(RetrievalWindowApplyError::InconsistentPlan {
788 field: "active_token_arithmetic",
789 });
790 }
791 if plan.retained_recent_user_turn_count > plan.retained_user_turn_count {
792 return Err(RetrievalWindowApplyError::InconsistentPlan {
793 field: "retained_recent_user_turn_count",
794 });
795 }
796 if plan.protected_active_tokens < plan.system_message_tokens {
797 return Err(RetrievalWindowApplyError::InconsistentPlan {
798 field: "protected_active_tokens",
799 });
800 }
801
802 let system_tokens = plan
803 .fixed_prompt_tokens
804 .checked_add(plan.system_message_tokens)
805 .ok_or(RetrievalWindowApplyError::TokenAccountingOverflow)?;
806 if system_tokens > plan.projected_active_tokens_after
807 || protected_with_fixed > plan.projected_active_tokens_after
808 {
809 return Err(RetrievalWindowApplyError::InconsistentPlan {
810 field: "projected_active_tokens_after",
811 });
812 }
813 let window_tokens = plan
814 .projected_active_tokens_after
815 .checked_sub(system_tokens)
816 .ok_or(RetrievalWindowApplyError::TokenAccountingOverflow)?;
817
818 Ok(ValidatedRetrievalWindowUsage {
819 system_tokens,
820 window_tokens,
821 })
822}
823
824fn resolve_candidate_indexes(
825 session: &Session,
826 plan: &RetrievalWindowCandidatePlan,
827) -> Result<Vec<usize>, RetrievalWindowApplyError> {
828 let mut candidate_ids = HashSet::new();
829 for (index, message_id) in plan.message_ids_to_archive.iter().enumerate() {
830 if message_id.is_empty() {
831 return Err(RetrievalWindowApplyError::EmptyCandidateId { index });
832 }
833 if !candidate_ids.insert(message_id.as_str()) {
834 return Err(RetrievalWindowApplyError::DuplicateCandidateId {
835 message_id: message_id.clone(),
836 });
837 }
838 }
839
840 let mut indexes_by_id: HashMap<&str, Vec<usize>> = HashMap::new();
841 for (index, message) in session.messages.iter().enumerate() {
842 if candidate_ids.contains(message.id.as_str()) {
843 indexes_by_id
844 .entry(message.id.as_str())
845 .or_default()
846 .push(index);
847 }
848 }
849
850 let mut indexes = Vec::with_capacity(plan.archive_message_count);
851 let mut previous_index = None;
852 for message_id in &plan.message_ids_to_archive {
853 let matches = indexes_by_id
854 .get(message_id.as_str())
855 .map(Vec::as_slice)
856 .unwrap_or(&[]);
857 let index = match matches {
858 [] => {
859 return Err(RetrievalWindowApplyError::MissingCandidateMessage {
860 message_id: message_id.clone(),
861 });
862 }
863 [index] => *index,
864 _ => {
865 return Err(RetrievalWindowApplyError::DuplicateSessionMessageId {
866 message_id: message_id.clone(),
867 });
868 }
869 };
870 if previous_index.is_some_and(|previous| previous >= index) {
871 return Err(RetrievalWindowApplyError::CandidateOrderMismatch {
872 message_id: message_id.clone(),
873 });
874 }
875 previous_index = Some(index);
876 indexes.push(index);
877 }
878 Ok(indexes)
879}
880
881fn usage_percentage(tokens: u32, limit: u32) -> f64 {
882 if limit == 0 {
883 0.0
884 } else {
885 (f64::from(tokens) / f64::from(limit)) * 100.0
886 }
887}
888
889const ACTIVE_STATE_DIGEST_DOMAIN: &[u8] = b"bamboo.retrieval-window.active-state.v1\0";
890const TOKEN_ACCOUNTING_DIGEST_DOMAIN: &[u8] = b"bamboo.retrieval-window.token-accounting.v2\0";
891const PLAN_EVIDENCE_DIGEST_DOMAIN: &[u8] = b"bamboo.retrieval-window.plan-evidence.v2\0";
892
893fn active_state_sha256(session: &Session) -> String {
894 let active_state = session
895 .messages
896 .iter()
897 .filter(|message| !message.compressed)
898 .map(|message| {
899 serde_json::json!([
900 message.id,
901 message.role,
902 message.content,
903 message.reasoning,
904 message.reasoning_signature,
905 message.content_parts,
906 message.image_ocr,
907 message.phase,
908 message.tool_calls,
909 message.tool_call_id,
910 message.tool_success,
911 message.never_compress,
912 message.compression_level,
913 ])
914 })
915 .collect::<Vec<_>>();
916 let encoded = serde_json::to_vec(&active_state)
917 .expect("token-relevant retrieval-window message state must serialize");
918 let mut payload = Vec::with_capacity(ACTIVE_STATE_DIGEST_DOMAIN.len() + encoded.len());
919 payload.extend_from_slice(ACTIVE_STATE_DIGEST_DOMAIN);
920 payload.extend_from_slice(&encoded);
921 sha256_hex(&payload)
922}
923
924fn token_accounting_sha256(accounting: &RetrievalWindowTokenAccounting) -> String {
925 let encoded = serde_json::to_vec(&(
926 accounting.fixed_prompt_tokens,
927 accounting.boundary_reclaimable_tokens,
928 &accounting.provider_message_tokens,
929 ))
930 .expect("retrieval-window token accounting must serialize");
931 let mut payload = Vec::with_capacity(TOKEN_ACCOUNTING_DIGEST_DOMAIN.len() + encoded.len());
932 payload.extend_from_slice(TOKEN_ACCOUNTING_DIGEST_DOMAIN);
933 payload.extend_from_slice(&encoded);
934 sha256_hex(&payload)
935}
936
937fn plan_evidence_sha256(plan: &RetrievalWindowCandidatePlan) -> String {
938 let encoded =
939 serde_json::to_vec(plan).expect("retrieval-window candidate plan evidence must serialize");
940 let mut payload = Vec::with_capacity(PLAN_EVIDENCE_DIGEST_DOMAIN.len() + encoded.len());
941 payload.extend_from_slice(PLAN_EVIDENCE_DIGEST_DOMAIN);
942 payload.extend_from_slice(&encoded);
943 sha256_hex(&payload)
944}
945
946fn validate_plan_evidence(
947 plan: &RetrievalWindowCandidatePlan,
948) -> Result<(), RetrievalWindowApplyError> {
949 if plan_evidence_sha256(plan) != plan.evidence_sha256 {
950 return Err(RetrievalWindowApplyError::StalePlan {
951 invariant: "plan_evidence_sha256",
952 });
953 }
954 Ok(())
955}
956
957fn validate_current_token_accounting(
958 plan: &RetrievalWindowCandidatePlan,
959 current_accounting: &RetrievalWindowTokenAccounting,
960) -> Result<(), RetrievalWindowApplyError> {
961 if current_accounting.fixed_prompt_tokens != plan.fixed_prompt_tokens
962 || token_accounting_sha256(current_accounting) != plan.token_accounting_sha256
963 {
964 return Err(RetrievalWindowApplyError::StalePlan {
965 invariant: "token_accounting_sha256",
966 });
967 }
968 Ok(())
969}
970
971fn validate_current_policy(
972 plan: &RetrievalWindowCandidatePlan,
973 current_policy: RetrievalWindowPolicy,
974) -> Result<(), RetrievalWindowApplyError> {
975 if current_policy.min_recent_user_turns != plan.min_recent_user_turns
976 || current_policy.target_usage_percent != plan.target_usage_percent
977 {
978 return Err(RetrievalWindowApplyError::StalePlan {
979 invariant: "retrieval_window_policy",
980 });
981 }
982 Ok(())
983}
984
985fn validate_current_token_budget(
986 plan: &RetrievalWindowCandidatePlan,
987 current_budget: &TokenBudget,
988) -> Result<(), RetrievalWindowApplyError> {
989 if current_budget.max_context_tokens != plan.context_window_tokens
990 || current_budget.max_request_input_tokens() != plan.request_input_limit_tokens
991 {
992 return Err(RetrievalWindowApplyError::StalePlan {
993 invariant: "token_budget",
994 });
995 }
996 Ok(())
997}
998
999fn validate_active_plan_structure(
1000 session: &Session,
1001 plan: &RetrievalWindowCandidatePlan,
1002) -> Result<(), RetrievalWindowApplyError> {
1003 let candidate_ids = plan
1004 .message_ids_to_archive
1005 .iter()
1006 .map(String::as_str)
1007 .collect::<HashSet<_>>();
1008 let (_, mut groups, active_message_count) = build_logical_groups(session);
1009 if active_message_count != plan.active_message_count {
1010 return Err(RetrievalWindowApplyError::StalePlan {
1011 invariant: "active_message_count",
1012 });
1013 }
1014 if plan.provider_message_token_override_count > active_message_count {
1015 return Err(RetrievalWindowApplyError::InconsistentPlan {
1016 field: "provider_message_token_override_count",
1017 });
1018 }
1019
1020 mark_protocol_safety(&mut groups);
1021 let incomplete_protocol_group_count =
1022 groups.iter().filter(|group| !group.protocol_safe).count();
1023
1024 let total_user_turn_count = groups
1025 .iter()
1026 .filter(|group| group.user_message_id.is_some())
1027 .count();
1028 let expected_retained_recent_user_turn_count =
1029 total_user_turn_count.min(plan.min_recent_user_turns);
1030 if plan.retained_recent_user_turn_count != expected_retained_recent_user_turn_count {
1031 return Err(RetrievalWindowApplyError::InconsistentPlan {
1032 field: "retained_recent_user_turn_count",
1033 });
1034 }
1035 mark_protected_groups(&mut groups, plan.min_recent_user_turns);
1036
1037 let mut selected_group_indexes = Vec::new();
1038 for (group_index, group) in groups.iter().enumerate() {
1039 let selected_message_count = group
1040 .messages
1041 .iter()
1042 .filter(|indexed| candidate_ids.contains(indexed.message.id.as_str()))
1043 .count();
1044 if selected_message_count == 0 {
1045 continue;
1046 }
1047 if selected_message_count != group.messages.len() {
1048 return Err(RetrievalWindowApplyError::IncompleteLogicalGroup);
1049 }
1050 if !group.protocol_safe {
1051 return Err(RetrievalWindowApplyError::UnsafeToolProtocol);
1052 }
1053 if logical_group_has_protected_skill(group) {
1054 return Err(RetrievalWindowApplyError::ProtectedSkillChain);
1055 }
1056 if group.protected {
1057 return Err(RetrievalWindowApplyError::RecentUserTurnProtected);
1058 }
1059 selected_group_indexes.push(group_index);
1060 }
1061
1062 if incomplete_protocol_group_count != plan.incomplete_protocol_group_count {
1063 return Err(RetrievalWindowApplyError::StalePlan {
1064 invariant: "incomplete_protocol_group_count",
1065 });
1066 }
1067
1068 if selected_group_indexes.len() != plan.archive_group_count {
1069 return Err(RetrievalWindowApplyError::StalePlan {
1070 invariant: "archive_group_count",
1071 });
1072 }
1073 let selected_group_index_set = selected_group_indexes
1074 .iter()
1075 .copied()
1076 .collect::<HashSet<_>>();
1077 let selected_ids_in_session_order = groups
1078 .iter()
1079 .enumerate()
1080 .filter(|(index, _)| selected_group_index_set.contains(index))
1081 .flat_map(|(_, group)| group.messages.iter())
1082 .map(|indexed| indexed.message.id.as_str())
1083 .collect::<Vec<_>>();
1084 if selected_ids_in_session_order
1085 != plan
1086 .message_ids_to_archive
1087 .iter()
1088 .map(String::as_str)
1089 .collect::<Vec<_>>()
1090 {
1091 return Err(RetrievalWindowApplyError::StalePlan {
1092 invariant: "candidate_group_order",
1093 });
1094 }
1095
1096 let Some(last_selected_group_index) = selected_group_indexes.last().copied() else {
1097 return Err(RetrievalWindowApplyError::IncompleteLogicalGroup);
1098 };
1099 if groups
1100 .iter()
1101 .enumerate()
1102 .take(last_selected_group_index + 1)
1103 .any(|(index, group)| !group.protected && !selected_group_index_set.contains(&index))
1104 {
1105 return Err(RetrievalWindowApplyError::StalePlan {
1106 invariant: "oldest_first_selection",
1107 });
1108 }
1109
1110 let archive_user_turn_count = selected_group_indexes
1111 .iter()
1112 .filter(|index| groups[**index].user_message_id.is_some())
1113 .count();
1114 if archive_user_turn_count != plan.archive_user_turn_count {
1115 return Err(RetrievalWindowApplyError::StalePlan {
1116 invariant: "archive_user_turn_count",
1117 });
1118 }
1119 let retained_user_turn_count = groups
1120 .iter()
1121 .enumerate()
1122 .filter(|(index, group)| {
1123 !selected_group_index_set.contains(index) && group.user_message_id.is_some()
1124 })
1125 .count();
1126 if retained_user_turn_count != plan.retained_user_turn_count {
1127 return Err(RetrievalWindowApplyError::StalePlan {
1128 invariant: "retained_user_turn_count",
1129 });
1130 }
1131
1132 let oldest_retained_message_id = groups
1133 .iter()
1134 .enumerate()
1135 .filter(|(index, _)| !selected_group_index_set.contains(index))
1136 .flat_map(|(_, group)| group.messages.iter())
1137 .min_by_key(|indexed| indexed.session_index)
1138 .map(|indexed| indexed.message.id.as_str());
1139 if oldest_retained_message_id != plan.oldest_retained_message_id.as_deref() {
1140 return Err(RetrievalWindowApplyError::StalePlan {
1141 invariant: "oldest_retained_message_id",
1142 });
1143 }
1144 let oldest_retained_user_message_id = groups
1145 .iter()
1146 .enumerate()
1147 .find(|(index, group)| {
1148 !selected_group_index_set.contains(index) && group.user_message_id.is_some()
1149 })
1150 .and_then(|(_, group)| group.user_message_id.as_deref());
1151 if oldest_retained_user_message_id != plan.oldest_retained_user_message_id.as_deref() {
1152 return Err(RetrievalWindowApplyError::StalePlan {
1153 invariant: "oldest_retained_user_message_id",
1154 });
1155 }
1156 if active_state_sha256(session) != plan.active_state_sha256 {
1157 return Err(RetrievalWindowApplyError::StalePlan {
1158 invariant: "active_state_sha256",
1159 });
1160 }
1161 Ok(())
1162}
1163
1164fn validate_idempotent_replay(
1165 session: &Session,
1166 plan: &RetrievalWindowCandidatePlan,
1167 candidate_indexes: &[usize],
1168 trigger_type: &CompressionTriggerType,
1169) -> Result<RetrievalWindowApplyResult, RetrievalWindowApplyError> {
1170 let Some(event_id) = candidate_indexes
1171 .first()
1172 .and_then(|index| session.messages[*index].compressed_by_event_id.as_deref())
1173 else {
1174 return Err(RetrievalWindowApplyError::MixedArchiveEvents);
1175 };
1176 if candidate_indexes
1177 .iter()
1178 .any(|index| session.messages[*index].compressed_by_event_id.as_deref() != Some(event_id))
1179 {
1180 return Err(RetrievalWindowApplyError::MixedArchiveEvents);
1181 }
1182
1183 let matching_events = session
1184 .compression_events
1185 .iter()
1186 .filter(|event| event.id == event_id)
1187 .collect::<Vec<_>>();
1188 let [event] = matching_events.as_slice() else {
1189 return Err(RetrievalWindowApplyError::MissingOrDuplicateArchiveEvent {
1190 event_id: event_id.to_string(),
1191 });
1192 };
1193 if event.kind != CompressionEventKind::RetrievalWindow {
1194 return Err(RetrievalWindowApplyError::ArchiveEventKindConflict {
1195 event_id: event_id.to_string(),
1196 });
1197 }
1198
1199 let correlated_message_ids = session
1200 .messages
1201 .iter()
1202 .filter(|message| message.compressed_by_event_id.as_deref() == Some(event_id))
1203 .map(|message| message.id.as_str())
1204 .collect::<Vec<_>>();
1205 if correlated_message_ids
1206 != plan
1207 .message_ids_to_archive
1208 .iter()
1209 .map(String::as_str)
1210 .collect::<Vec<_>>()
1211 {
1212 return Err(replay_evidence_mismatch(event_id, "message_correlation"));
1213 }
1214
1215 macro_rules! require_event_evidence {
1216 ($actual:expr, $expected:expr, $field:literal) => {
1217 if $actual != $expected {
1218 return Err(replay_evidence_mismatch(event_id, $field));
1219 }
1220 };
1221 }
1222
1223 require_event_evidence!(
1224 event.messages_compressed,
1225 plan.archive_message_count,
1226 "messages_compressed"
1227 );
1228 require_event_evidence!(
1229 event.segments_removed,
1230 plan.archive_group_count,
1231 "segments_removed"
1232 );
1233 require_event_evidence!(event.summary_tokens, 0, "summary_tokens");
1234 require_event_evidence!(event.actual_summary_tokens, 0, "actual_summary_tokens");
1235 require_event_evidence!(&event.trigger_type, trigger_type, "trigger_type");
1236 require_event_evidence!(
1237 event.compression_ratio.to_bits(),
1238 0.0f64.to_bits(),
1239 "compression_ratio"
1240 );
1241 require_event_evidence!(event.model_used.as_ref(), None, "model_used");
1242 require_event_evidence!(event.latency_ms, 0, "latency_ms");
1243 require_event_evidence!(
1244 event.source_tokens,
1245 plan.archive_message_tokens,
1246 "source_tokens"
1247 );
1248 require_event_evidence!(
1249 event.fixed_prompt_tokens,
1250 plan.fixed_prompt_tokens,
1251 "fixed_prompt_tokens"
1252 );
1253 require_event_evidence!(event.target_summary_tokens, 0, "target_summary_tokens");
1254 require_event_evidence!(
1255 event.summary_target_ratio.to_bits(),
1256 0.0f64.to_bits(),
1257 "summary_target_ratio"
1258 );
1259 require_event_evidence!(
1260 event.actual_summary_ratio.to_bits(),
1261 0.0f64.to_bits(),
1262 "actual_summary_ratio"
1263 );
1264 require_event_evidence!(
1265 event.summary_budget_clamped,
1266 false,
1267 "summary_budget_clamped"
1268 );
1269 require_event_evidence!(
1270 event.summary_budget_clamp_reason.as_ref(),
1271 None,
1272 "summary_budget_clamp_reason"
1273 );
1274 require_event_evidence!(event.summarization_map_calls, 0, "summarization_map_calls");
1275 require_event_evidence!(
1276 event.summarization_reduce_calls,
1277 0,
1278 "summarization_reduce_calls"
1279 );
1280 require_event_evidence!(
1281 event.summarization_fallback_used,
1282 false,
1283 "summarization_fallback_used"
1284 );
1285 require_event_evidence!(
1286 event.retrieval_active_tokens_before,
1287 plan.active_tokens_before,
1288 "retrieval_active_tokens_before"
1289 );
1290 require_event_evidence!(
1291 event.retrieval_active_message_count_before,
1292 plan.active_message_count,
1293 "retrieval_active_message_count_before"
1294 );
1295 require_event_evidence!(
1296 event.retrieval_active_state_sha256.as_deref(),
1297 Some(plan.active_state_sha256.as_str()),
1298 "retrieval_active_state_sha256"
1299 );
1300 require_event_evidence!(
1301 event.retrieval_token_accounting_sha256.as_deref(),
1302 Some(plan.token_accounting_sha256.as_str()),
1303 "retrieval_token_accounting_sha256"
1304 );
1305 require_event_evidence!(
1306 event.retrieval_active_tokens_after,
1307 plan.projected_active_tokens_after,
1308 "retrieval_active_tokens_after"
1309 );
1310 require_event_evidence!(
1311 event.retrieval_target_tokens,
1312 plan.target_tokens,
1313 "retrieval_target_tokens"
1314 );
1315 require_event_evidence!(
1316 event.retrieval_target_usage_percent,
1317 plan.target_usage_percent,
1318 "retrieval_target_usage_percent"
1319 );
1320 require_event_evidence!(
1321 event.retrieval_archived_group_count,
1322 plan.archive_group_count,
1323 "retrieval_archived_group_count"
1324 );
1325 require_event_evidence!(
1326 event.retrieval_archived_user_turn_count,
1327 plan.archive_user_turn_count,
1328 "retrieval_archived_user_turn_count"
1329 );
1330 require_event_evidence!(
1331 event.retrieval_archived_message_tokens,
1332 plan.archive_message_tokens,
1333 "retrieval_archived_message_tokens"
1334 );
1335 require_event_evidence!(
1336 event.retrieval_boundary_reclaimed_tokens,
1337 plan.boundary_reclaimable_tokens,
1338 "retrieval_boundary_reclaimed_tokens"
1339 );
1340 require_event_evidence!(
1341 event.retrieval_system_message_tokens,
1342 plan.system_message_tokens,
1343 "retrieval_system_message_tokens"
1344 );
1345 require_event_evidence!(
1346 event.retrieval_context_window_tokens,
1347 plan.context_window_tokens,
1348 "retrieval_context_window_tokens"
1349 );
1350 require_event_evidence!(
1351 event.retrieval_request_input_limit_tokens,
1352 plan.request_input_limit_tokens,
1353 "retrieval_request_input_limit_tokens"
1354 );
1355 require_event_evidence!(
1356 event.retrieval_min_recent_user_turns,
1357 plan.min_recent_user_turns,
1358 "retrieval_min_recent_user_turns"
1359 );
1360 require_event_evidence!(
1361 event.retrieval_retained_recent_user_turn_count,
1362 plan.retained_recent_user_turn_count,
1363 "retrieval_retained_recent_user_turn_count"
1364 );
1365 require_event_evidence!(
1366 event.retrieval_retained_user_turn_count,
1367 plan.retained_user_turn_count,
1368 "retrieval_retained_user_turn_count"
1369 );
1370 require_event_evidence!(
1371 event.retrieval_oldest_retained_message_id.as_ref(),
1372 plan.oldest_retained_message_id.as_ref(),
1373 "retrieval_oldest_retained_message_id"
1374 );
1375 require_event_evidence!(
1376 event.retrieval_oldest_retained_user_message_id.as_ref(),
1377 plan.oldest_retained_user_message_id.as_ref(),
1378 "retrieval_oldest_retained_user_message_id"
1379 );
1380 require_event_evidence!(
1381 event.retrieval_provider_message_token_override_count,
1382 plan.provider_message_token_override_count,
1383 "retrieval_provider_message_token_override_count"
1384 );
1385 require_event_evidence!(
1386 event.retrieval_protected_active_tokens,
1387 plan.protected_active_tokens,
1388 "retrieval_protected_active_tokens"
1389 );
1390 require_event_evidence!(
1391 event.retrieval_incomplete_protocol_group_count,
1392 plan.incomplete_protocol_group_count,
1393 "retrieval_incomplete_protocol_group_count"
1394 );
1395 require_event_evidence!(
1396 event.usage_before_percent.to_bits(),
1397 usage_percentage(plan.active_tokens_before, plan.context_window_tokens).to_bits(),
1398 "usage_before_percent"
1399 );
1400 require_event_evidence!(
1401 event.usage_after_percent.to_bits(),
1402 usage_percentage(
1403 plan.projected_active_tokens_after,
1404 plan.context_window_tokens
1405 )
1406 .to_bits(),
1407 "usage_after_percent"
1408 );
1409
1410 Ok(RetrievalWindowApplyResult {
1411 event_id: event_id.to_string(),
1412 newly_archived_message_count: 0,
1413 idempotent_replay: true,
1414 })
1415}
1416
1417fn replay_evidence_mismatch(event_id: &str, field: &'static str) -> RetrievalWindowApplyError {
1418 RetrievalWindowApplyError::ArchiveEventEvidenceMismatch {
1419 event_id: event_id.to_string(),
1420 field,
1421 }
1422}
1423
1424fn validate_inputs(
1425 budget: &TokenBudget,
1426 policy: RetrievalWindowPolicy,
1427) -> Result<(), RetrievalWindowPlanError> {
1428 if policy.min_recent_user_turns == 0 {
1429 return Err(RetrievalWindowPlanError::InvalidRecentUserTurnFloor);
1430 }
1431 if !(1..=100).contains(&policy.target_usage_percent) {
1432 return Err(RetrievalWindowPlanError::InvalidTargetUsagePercent {
1433 target_usage_percent: policy.target_usage_percent,
1434 });
1435 }
1436 let max_request_input_tokens = budget.max_request_input_tokens();
1437 if budget.max_context_tokens == 0 || max_request_input_tokens == 0 {
1438 return Err(RetrievalWindowPlanError::InvalidTokenBudget {
1439 context_window_tokens: budget.max_context_tokens,
1440 max_request_input_tokens,
1441 });
1442 }
1443 Ok(())
1444}
1445
1446pub fn effective_retrieval_window_target_tokens(
1453 budget: &TokenBudget,
1454 target_usage_percent: u8,
1455) -> u32 {
1456 effective_target_tokens(
1457 budget.max_context_tokens,
1458 budget.max_request_input_tokens(),
1459 target_usage_percent,
1460 )
1461}
1462
1463fn effective_target_tokens(
1464 context_window_tokens: u32,
1465 request_input_limit_tokens: u32,
1466 target_usage_percent: u8,
1467) -> u32 {
1468 let percentage_target =
1469 ((u64::from(context_window_tokens) * u64::from(target_usage_percent)) / 100).max(1) as u32;
1470 percentage_target.min(request_input_limit_tokens)
1471}
1472
1473fn checked_sum(values: impl IntoIterator<Item = u32>) -> Result<u32, RetrievalWindowPlanError> {
1474 values.into_iter().try_fold(0u32, |total, value| {
1475 total
1476 .checked_add(value)
1477 .ok_or(RetrievalWindowPlanError::TokenAccountingOverflow)
1478 })
1479}
1480
1481fn count_indexed_messages(
1482 messages: &[IndexedMessage<'_>],
1483 accounting: &RetrievalWindowTokenAccounting,
1484 counter: &impl TokenCounter,
1485) -> Result<u32, RetrievalWindowPlanError> {
1486 messages.iter().try_fold(0u32, |total, indexed| {
1487 total
1488 .checked_add(count_provider_visible_message_tokens(
1489 indexed.message,
1490 accounting,
1491 counter,
1492 )?)
1493 .ok_or(RetrievalWindowPlanError::TokenAccountingOverflow)
1494 })
1495}
1496
1497fn count_provider_visible_message_tokens(
1498 message: &Message,
1499 accounting: &RetrievalWindowTokenAccounting,
1500 counter: &impl TokenCounter,
1501) -> Result<u32, RetrievalWindowPlanError> {
1502 if let Some(provider_tokens) = accounting.provider_message_tokens.get(&message.id) {
1503 return Ok(*provider_tokens);
1504 }
1505
1506 let mut tokens = counter.count_message(message);
1507
1508 if matches!(message.role, Role::Assistant) {
1512 if let Some(reasoning) = message
1513 .reasoning
1514 .as_deref()
1515 .filter(|reasoning| !reasoning.is_empty())
1516 {
1517 tokens = tokens
1518 .checked_add(counter.count_text(reasoning))
1519 .ok_or(RetrievalWindowPlanError::TokenAccountingOverflow)?;
1520 }
1521 }
1522
1523 if let Some(parts) = message.content_parts.as_deref() {
1529 for part in parts {
1530 let part_tokens = match part {
1531 MessagePart::Text { text } => counter.count_text(text),
1532 MessagePart::ImageUrl { .. } => {
1533 return Err(
1534 RetrievalWindowPlanError::MissingProviderMessageTokenEstimate {
1535 message_id: message.id.clone(),
1536 },
1537 );
1538 }
1539 };
1540 tokens = tokens
1541 .checked_add(part_tokens)
1542 .ok_or(RetrievalWindowPlanError::TokenAccountingOverflow)?;
1543 }
1544 }
1545
1546 Ok(tokens)
1547}
1548
1549fn build_logical_groups(
1550 session: &Session,
1551) -> (Vec<IndexedMessage<'_>>, Vec<LogicalGroup<'_>>, usize) {
1552 let mut system_messages = Vec::new();
1553 let mut groups = Vec::new();
1554 let mut active_message_count = 0usize;
1555
1556 for (session_index, message) in session.messages.iter().enumerate() {
1557 if message.compressed {
1558 continue;
1559 }
1560 active_message_count += 1;
1561 let indexed = IndexedMessage {
1562 session_index,
1563 message,
1564 };
1565
1566 match message.role {
1567 Role::System => system_messages.push(indexed),
1568 Role::User => groups.push(LogicalGroup::user_turn(indexed)),
1569 Role::Assistant | Role::Tool => match groups.last_mut() {
1570 Some(group) => group.messages.push(indexed),
1571 None => groups.push(LogicalGroup::preamble(indexed)),
1572 },
1573 }
1574 }
1575
1576 (system_messages, groups, active_message_count)
1577}
1578
1579fn mark_protocol_safety(groups: &mut [LogicalGroup<'_>]) {
1580 let mut calls_by_id: HashMap<String, Vec<(usize, usize)>> = HashMap::new();
1581 let mut results_by_id: HashMap<String, Vec<(usize, usize)>> = HashMap::new();
1582
1583 for (group_index, group) in groups.iter_mut().enumerate() {
1584 let mut pending_call_ids = HashSet::new();
1585
1586 for indexed in &group.messages {
1587 let message = indexed.message;
1588 match message.role {
1589 Role::Assistant => {
1590 if message.tool_call_id.is_some() {
1591 group.protocol_safe = false;
1592 }
1593 let calls = message
1594 .tool_calls
1595 .as_ref()
1596 .filter(|calls| !calls.is_empty());
1597 match calls {
1598 Some(calls) => {
1599 if !pending_call_ids.is_empty() {
1600 group.protocol_safe = false;
1601 pending_call_ids.clear();
1602 }
1603 let mut local_call_ids = HashSet::new();
1604 for call in calls {
1605 if call.id.is_empty() || !local_call_ids.insert(call.id.clone()) {
1606 group.protocol_safe = false;
1607 }
1608 pending_call_ids.insert(call.id.clone());
1609 calls_by_id
1610 .entry(call.id.clone())
1611 .or_default()
1612 .push((group_index, indexed.session_index));
1613 }
1614 }
1615 None if !pending_call_ids.is_empty() => {
1616 group.protocol_safe = false;
1617 pending_call_ids.clear();
1618 }
1619 None => {}
1620 }
1621 }
1622 Role::Tool => {
1623 if message
1624 .tool_calls
1625 .as_ref()
1626 .is_some_and(|calls| !calls.is_empty())
1627 {
1628 group.protocol_safe = false;
1629 }
1630 match message.tool_call_id.as_deref() {
1631 Some(call_id) if !call_id.is_empty() => {
1632 results_by_id
1633 .entry(call_id.to_string())
1634 .or_default()
1635 .push((group_index, indexed.session_index));
1636 if !pending_call_ids.remove(call_id) {
1637 group.protocol_safe = false;
1638 }
1639 }
1640 _ => group.protocol_safe = false,
1641 }
1642 }
1643 Role::User => {
1644 if !pending_call_ids.is_empty()
1645 || message.tool_call_id.is_some()
1646 || message
1647 .tool_calls
1648 .as_ref()
1649 .is_some_and(|calls| !calls.is_empty())
1650 {
1651 group.protocol_safe = false;
1652 pending_call_ids.clear();
1653 }
1654 }
1655 Role::System => unreachable!("system messages are not placed in logical groups"),
1656 }
1657 }
1658
1659 if !pending_call_ids.is_empty() {
1660 group.protocol_safe = false;
1661 }
1662 }
1663
1664 let call_ids = calls_by_id
1665 .keys()
1666 .chain(results_by_id.keys())
1667 .cloned()
1668 .collect::<HashSet<_>>();
1669 for call_id in call_ids {
1670 let calls = calls_by_id.get(&call_id).map(Vec::as_slice).unwrap_or(&[]);
1671 let results = results_by_id
1672 .get(&call_id)
1673 .map(Vec::as_slice)
1674 .unwrap_or(&[]);
1675 let valid = calls.len() == 1
1676 && results.len() == 1
1677 && calls[0].0 == results[0].0
1678 && calls[0].1 < results[0].1;
1679 if valid {
1680 continue;
1681 }
1682 for (group_index, _) in calls.iter().chain(results.iter()) {
1683 if let Some(group) = groups.get_mut(*group_index) {
1684 group.protocol_safe = false;
1685 }
1686 }
1687 }
1688}
1689
1690fn logical_group_has_protected_skill(group: &LogicalGroup<'_>) -> bool {
1691 group.messages.iter().any(|indexed| {
1692 indexed.message.tool_calls.as_ref().is_some_and(|calls| {
1693 calls
1694 .iter()
1695 .any(|call| is_skill_tool_name(&call.function.name, SkillNameMatch::Canonical))
1696 })
1697 })
1698}
1699
1700fn mark_protected_groups(groups: &mut [LogicalGroup<'_>], min_recent_user_turns: usize) {
1701 for group in groups.iter_mut() {
1702 group.protected = !group.protocol_safe
1703 || group
1704 .messages
1705 .iter()
1706 .any(|indexed| indexed.message.never_compress)
1707 || logical_group_has_protected_skill(group);
1708 }
1709
1710 let user_group_indexes = groups
1711 .iter()
1712 .enumerate()
1713 .filter_map(|(index, group)| group.user_message_id.as_ref().map(|_| index))
1714 .collect::<Vec<_>>();
1715 let retained_start = user_group_indexes
1716 .len()
1717 .saturating_sub(min_recent_user_turns);
1718 for group_index in &user_group_indexes[retained_start..] {
1719 groups[*group_index].protected = true;
1720 }
1721}
1722
1723#[cfg(test)]
1724mod tests {
1725 use super::*;
1726 use bamboo_domain::{
1727 provider_transcript_boundary_sha256, ConversationSummary, FunctionCall, ModelContextState,
1728 ProviderFamily, ProviderProtocol, ProviderTranscriptAuthor, ProviderTranscriptItem,
1729 ProviderTranscriptOrigin, ProviderTranscriptResetReason, ToolCall,
1730 };
1731
1732 #[derive(Debug)]
1733 struct CharacterTokenCounter;
1734
1735 impl TokenCounter for CharacterTokenCounter {
1736 fn count_message(&self, message: &Message) -> u32 {
1737 message.content.chars().count() as u32
1738 }
1739
1740 fn count_text(&self, text: &str) -> u32 {
1741 text.chars().count() as u32
1742 }
1743 }
1744
1745 fn budget(context_window: u32) -> TokenBudget {
1746 let mut budget = TokenBudget::for_model(context_window);
1747 budget.max_output_tokens = 0;
1748 budget.safety_margin = 0;
1749 budget
1750 }
1751
1752 fn policy(min_recent_user_turns: usize, target_usage_percent: u8) -> RetrievalWindowPolicy {
1753 RetrievalWindowPolicy {
1754 min_recent_user_turns,
1755 target_usage_percent,
1756 }
1757 }
1758
1759 fn identified(mut message: Message, id: &str, tokens: usize) -> Message {
1760 message.id = id.to_string();
1761 message.content = "x".repeat(tokens);
1762 message
1763 }
1764
1765 fn system(id: &str, tokens: usize) -> Message {
1766 identified(Message::system(""), id, tokens)
1767 }
1768
1769 fn user(id: &str, tokens: usize) -> Message {
1770 identified(Message::user(""), id, tokens)
1771 }
1772
1773 fn assistant(id: &str, tokens: usize) -> Message {
1774 identified(Message::assistant("", None), id, tokens)
1775 }
1776
1777 fn tool_call(message_id: &str, call_id: &str, tokens: usize, name: &str) -> Message {
1778 let mut message = assistant(message_id, tokens);
1779 message.tool_calls = Some(vec![ToolCall {
1780 id: call_id.to_string(),
1781 tool_type: "function".to_string(),
1782 function: FunctionCall {
1783 name: name.to_string(),
1784 arguments: "{}".to_string(),
1785 },
1786 }]);
1787 message
1788 }
1789
1790 fn tool_result(message_id: &str, call_id: &str, tokens: usize) -> Message {
1791 identified(Message::tool_result(call_id, ""), message_id, tokens)
1792 }
1793
1794 fn add_turn(session: &mut Session, prefix: &str, tokens_per_message: usize) {
1795 session.add_message(user(&format!("{prefix}-u"), tokens_per_message));
1796 session.add_message(assistant(&format!("{prefix}-a"), tokens_per_message));
1797 }
1798
1799 fn basic_session_and_plan() -> (Session, RetrievalWindowCandidatePlan) {
1800 let mut session = Session::new("retrieval-window-apply", "test-model");
1801 session.add_message(system("system", 5));
1802 for turn in 1..=4 {
1803 add_turn(&mut session, &format!("t{turn}"), 10);
1804 }
1805 let plan = plan_with_counter(&session, 50, 2, 5).expect("plan should build");
1806 (session, plan)
1807 }
1808
1809 fn safe_tool_session_and_plan() -> (Session, RetrievalWindowCandidatePlan) {
1810 let mut session = Session::new("retrieval-window-apply-tool", "test-model");
1811 session.add_message(system("system", 5));
1812 session.add_message(user("old-u", 5));
1813 session.add_message(tool_call("old-call-1", "call-1", 5, "Read"));
1814 session.add_message(tool_result("old-result-1", "call-1", 5));
1815 session.add_message(tool_call("old-call-2", "call-2", 5, "Grep"));
1816 session.add_message(tool_result("old-result-2", "call-2", 5));
1817 session.add_message(assistant("old-final", 5));
1818 add_turn(&mut session, "recent", 10);
1819 let plan = plan_with_counter(&session, 50, 1, 0).expect("tool plan should build");
1820 (session, plan)
1821 }
1822
1823 fn accounting_for_plan(plan: &RetrievalWindowCandidatePlan) -> RetrievalWindowTokenAccounting {
1824 RetrievalWindowTokenAccounting {
1825 fixed_prompt_tokens: plan.fixed_prompt_tokens,
1826 boundary_reclaimable_tokens: plan.boundary_reclaimable_tokens,
1827 ..RetrievalWindowTokenAccounting::default()
1828 }
1829 }
1830
1831 fn policy_for_plan(plan: &RetrievalWindowCandidatePlan) -> RetrievalWindowPolicy {
1832 RetrievalWindowPolicy {
1833 min_recent_user_turns: plan.min_recent_user_turns,
1834 target_usage_percent: plan.target_usage_percent,
1835 }
1836 }
1837
1838 fn apply_test_plan(
1839 session: &mut Session,
1840 plan: &RetrievalWindowCandidatePlan,
1841 ) -> Result<RetrievalWindowApplyResult, RetrievalWindowApplyError> {
1842 apply_retrieval_window_plan(
1843 session,
1844 plan,
1845 policy_for_plan(plan),
1846 &budget(100),
1847 &accounting_for_plan(plan),
1848 )
1849 }
1850
1851 fn assert_apply_error_without_mutation(
1852 session: &mut Session,
1853 plan: &RetrievalWindowCandidatePlan,
1854 expected: RetrievalWindowApplyError,
1855 ) {
1856 let before = serde_json::to_vec(session).expect("session should serialize");
1857 assert_eq!(apply_test_plan(session, plan), Err(expected));
1858 assert_eq!(
1859 serde_json::to_vec(session).expect("session should serialize"),
1860 before,
1861 "failed application must not mutate durable session state"
1862 );
1863 }
1864
1865 fn plan_with_counter(
1866 session: &Session,
1867 target_usage_percent: u8,
1868 min_recent_user_turns: usize,
1869 fixed_prompt_tokens: u32,
1870 ) -> Result<RetrievalWindowCandidatePlan, RetrievalWindowPlanError> {
1871 plan_with_counter_and_accounting(
1872 session,
1873 target_usage_percent,
1874 min_recent_user_turns,
1875 &RetrievalWindowTokenAccounting {
1876 fixed_prompt_tokens,
1877 ..RetrievalWindowTokenAccounting::default()
1878 },
1879 )
1880 }
1881
1882 fn plan_with_counter_and_accounting(
1883 session: &Session,
1884 target_usage_percent: u8,
1885 min_recent_user_turns: usize,
1886 accounting: &RetrievalWindowTokenAccounting,
1887 ) -> Result<RetrievalWindowCandidatePlan, RetrievalWindowPlanError> {
1888 build_retrieval_window_candidate_plan_with_counter(
1889 session,
1890 &budget(100),
1891 policy(min_recent_user_turns, target_usage_percent),
1892 accounting,
1893 false,
1894 &CharacterTokenCounter,
1895 )
1896 }
1897
1898 #[test]
1899 fn selects_oldest_complete_turns_deterministically_without_mutation() {
1900 let mut session = Session::new("retrieval-window-basic", "test-model");
1901 session.add_message(system("system", 5));
1902 for turn in 1..=4 {
1903 add_turn(&mut session, &format!("t{turn}"), 10);
1904 }
1905 let before = serde_json::to_value(&session).expect("session should serialize");
1906
1907 let first = plan_with_counter(&session, 50, 2, 5).expect("plan should build");
1908 let second = plan_with_counter(&session, 50, 2, 5).expect("plan should repeat");
1909
1910 assert_eq!(first, second);
1911 assert_eq!(
1912 first.message_ids_to_archive,
1913 vec!["t1-u", "t1-a", "t2-u", "t2-a"]
1914 );
1915 assert_eq!(first.archive_message_count, 4);
1916 assert_eq!(first.archive_group_count, 2);
1917 assert_eq!(first.archive_user_turn_count, 2);
1918 assert_eq!(first.archive_message_tokens, 40);
1919 assert_eq!(first.active_message_count, 9);
1920 assert_eq!(first.active_state_sha256.len(), 64);
1921 assert_eq!(first.token_accounting_sha256.len(), 64);
1922 assert_eq!(first.active_tokens_before, 90);
1923 assert_eq!(first.projected_active_tokens_after, 50);
1924 assert_eq!(first.context_window_tokens, 100);
1925 assert_eq!(first.request_input_limit_tokens, 100);
1926 assert_eq!(first.min_recent_user_turns, 2);
1927 assert_eq!(first.fixed_prompt_tokens, 5);
1928 assert_eq!(first.system_message_tokens, 5);
1929 assert_eq!(first.protected_active_tokens, 45);
1930 assert_eq!(first.retained_recent_user_turn_count, 2);
1931 assert_eq!(first.retained_user_turn_count, 2);
1932 assert_eq!(first.oldest_retained_message_id.as_deref(), Some("t3-u"));
1933 assert_eq!(
1934 first.oldest_retained_user_message_id.as_deref(),
1935 Some("t3-u")
1936 );
1937 assert_eq!(
1938 serde_json::to_value(&session).expect("session should serialize"),
1939 before,
1940 "planning must not mutate the authoritative session"
1941 );
1942 }
1943
1944 #[test]
1945 fn keeps_multiple_tool_rounds_atomic_with_their_user_turn() {
1946 let mut session = Session::new("retrieval-window-tools", "test-model");
1947 session.add_message(system("system", 5));
1948 session.add_message(user("old-u", 5));
1949 session.add_message(tool_call("old-call-1", "call-1", 5, "Read"));
1950 session.add_message(tool_result("old-result-1", "call-1", 5));
1951 session.add_message(tool_call("old-call-2", "call-2", 5, "Grep"));
1952 session.add_message(tool_result("old-result-2", "call-2", 5));
1953 session.add_message(assistant("old-final", 5));
1954 add_turn(&mut session, "recent", 10);
1955
1956 let plan = plan_with_counter(&session, 50, 1, 0).expect("tool turn is safe");
1957
1958 assert_eq!(
1959 plan.message_ids_to_archive,
1960 vec![
1961 "old-u",
1962 "old-call-1",
1963 "old-result-1",
1964 "old-call-2",
1965 "old-result-2",
1966 "old-final"
1967 ]
1968 );
1969 assert_eq!(plan.archive_group_count, 1);
1970 assert_eq!(plan.archive_user_turn_count, 1);
1971 assert_eq!(plan.incomplete_protocol_group_count, 0);
1972 }
1973
1974 #[test]
1975 fn incomplete_tool_chain_is_protected_and_reported() {
1976 let mut session = Session::new("retrieval-window-incomplete", "test-model");
1977 session.add_message(system("system", 5));
1978 session.add_message(user("old-u", 10));
1979 session.add_message(tool_call("old-call", "missing-result", 10, "Read"));
1980 add_turn(&mut session, "recent", 15);
1981
1982 let error = plan_with_counter(&session, 50, 1, 0).expect_err("chain is incomplete");
1983
1984 assert_eq!(
1985 error,
1986 RetrievalWindowPlanError::NothingToArchive {
1987 active_tokens: 55,
1988 target_tokens: 50,
1989 protected_active_tokens: 55,
1990 incomplete_protocol_group_count: 1,
1991 }
1992 );
1993 }
1994
1995 #[test]
1996 fn duplicate_tool_results_make_the_whole_turn_ambiguous() {
1997 let mut session = Session::new("retrieval-window-ambiguous", "test-model");
1998 session.add_message(system("system", 5));
1999 session.add_message(user("old-u", 5));
2000 session.add_message(tool_call("old-call", "duplicate-result", 5, "Read"));
2001 session.add_message(tool_result("old-result-1", "duplicate-result", 5));
2002 session.add_message(tool_result("old-result-2", "duplicate-result", 5));
2003 add_turn(&mut session, "recent", 15);
2004
2005 let error = plan_with_counter(&session, 50, 1, 0).expect_err("chain is ambiguous");
2006
2007 assert_eq!(
2008 error,
2009 RetrievalWindowPlanError::NothingToArchive {
2010 active_tokens: 55,
2011 target_tokens: 50,
2012 protected_active_tokens: 55,
2013 incomplete_protocol_group_count: 1,
2014 }
2015 );
2016 }
2017
2018 #[test]
2019 fn never_compress_turn_is_retained_while_another_old_turn_is_selected() {
2020 let mut session = Session::new("retrieval-window-protected", "test-model");
2021 session.add_message(system("system", 5));
2022 let mut protected = user("protected-u", 10);
2023 protected.never_compress = true;
2024 session.add_message(protected);
2025 session.add_message(assistant("protected-a", 10));
2026 add_turn(&mut session, "eligible", 10);
2027 add_turn(&mut session, "recent", 10);
2028
2029 let plan = plan_with_counter(&session, 50, 1, 0).expect("eligible turn should fit");
2030
2031 assert_eq!(
2032 plan.message_ids_to_archive,
2033 vec!["eligible-u", "eligible-a"]
2034 );
2035 assert_eq!(plan.protected_active_tokens, 45);
2036 assert_eq!(
2037 plan.oldest_retained_user_message_id.as_deref(),
2038 Some("protected-u")
2039 );
2040 }
2041
2042 #[test]
2043 fn leading_non_system_preamble_is_an_explicit_candidate_group() {
2044 let mut session = Session::new("retrieval-window-preamble", "test-model");
2045 session.add_message(system("system", 5));
2046 session.add_message(assistant("preamble", 20));
2047 add_turn(&mut session, "old", 10);
2048 add_turn(&mut session, "recent", 10);
2049
2050 let plan = plan_with_counter(&session, 50, 1, 0).expect("preamble should archive");
2051
2052 assert_eq!(plan.message_ids_to_archive, vec!["preamble"]);
2053 assert_eq!(plan.archive_group_count, 1);
2054 assert_eq!(plan.archive_user_turn_count, 0);
2055 assert_eq!(plan.oldest_retained_message_id.as_deref(), Some("old-u"));
2056 }
2057
2058 #[test]
2059 fn already_archived_messages_are_ignored_on_repeated_planning() {
2060 let mut session = Session::new("retrieval-window-repeat", "test-model");
2061 session.add_message(system("system", 5));
2062 for turn in 1..=4 {
2063 add_turn(&mut session, &format!("t{turn}"), 10);
2064 }
2065
2066 let first = plan_with_counter(&session, 50, 2, 5).expect("first plan should build");
2067 let archived_ids = first
2068 .message_ids_to_archive
2069 .iter()
2070 .map(String::as_str)
2071 .collect::<HashSet<_>>();
2072 for message in &mut session.messages {
2073 if archived_ids.contains(message.id.as_str()) {
2074 message.compressed = true;
2075 }
2076 }
2077
2078 assert_eq!(
2079 plan_with_counter(&session, 50, 2, 5),
2080 Err(RetrievalWindowPlanError::TargetAlreadySatisfied {
2081 active_tokens: 50,
2082 target_tokens: 50,
2083 })
2084 );
2085 }
2086
2087 #[test]
2088 fn fixed_and_protected_content_over_target_returns_typed_error() {
2089 let mut session = Session::new("retrieval-window-fixed-overflow", "test-model");
2090 session.add_message(system("system", 5));
2091 add_turn(&mut session, "old", 10);
2092 add_turn(&mut session, "recent", 10);
2093
2094 let error =
2095 plan_with_counter(&session, 50, 1, 60).expect_err("fixed prompt cost prevents target");
2096
2097 assert_eq!(
2098 error,
2099 RetrievalWindowPlanError::ProtectedContentExceedsTarget {
2100 projected_tokens: 85,
2101 target_tokens: 50,
2102 protected_active_tokens: 25,
2103 fixed_prompt_tokens: 60,
2104 incomplete_protocol_group_count: 0,
2105 }
2106 );
2107 }
2108
2109 #[test]
2110 fn boundary_reclaim_is_subtracted_once_and_still_commits_an_archive_group() {
2111 let mut session = Session::new("retrieval-window-boundary-reclaim", "test-model");
2112 session.add_message(system("system", 5));
2113 add_turn(&mut session, "old", 10);
2114 add_turn(&mut session, "recent", 10);
2115 let accounting = RetrievalWindowTokenAccounting {
2116 fixed_prompt_tokens: 5,
2117 boundary_reclaimable_tokens: 40,
2118 ..RetrievalWindowTokenAccounting::default()
2119 };
2120
2121 let plan = plan_with_counter_and_accounting(&session, 50, 1, &accounting)
2122 .expect("one archive group should commit the boundary reclaim");
2123
2124 assert_eq!(plan.active_tokens_before, 90);
2125 assert_eq!(plan.boundary_reclaimable_tokens, 40);
2126 assert_eq!(plan.message_ids_to_archive, vec!["old-u", "old-a"]);
2127 assert_eq!(plan.archive_message_tokens, 20);
2128 assert_eq!(plan.projected_active_tokens_after, 30);
2129
2130 apply_retrieval_window_plan(
2131 &mut session,
2132 &plan,
2133 policy_for_plan(&plan),
2134 &budget(100),
2135 &accounting,
2136 )
2137 .expect("boundary-aware plan should apply");
2138 let event = session
2139 .compression_events
2140 .last()
2141 .expect("retrieval event should be recorded");
2142 assert_eq!(event.retrieval_boundary_reclaimed_tokens, 40);
2143 assert_eq!(event.retrieval_active_tokens_after, 30);
2144 assert_eq!(session.token_usage.as_ref().unwrap().total_tokens, 30);
2145 }
2146
2147 #[test]
2148 fn stored_assistant_reasoning_counts_with_its_logical_group() {
2149 let mut session = Session::new("retrieval-window-reasoning", "test-model");
2150 session.add_message(system("system", 5));
2151 session.add_message(user("old-u", 5));
2152 let mut reasoned = assistant("old-a", 5);
2153 reasoned.reasoning = Some("r".repeat(40));
2154 reasoned.reasoning_signature = Some("provider-signature".to_string());
2155 session.add_message(reasoned);
2156 add_turn(&mut session, "recent", 10);
2157
2158 let plan = plan_with_counter(&session, 50, 1, 0)
2159 .expect("provider-visible reasoning should put the session over target");
2160
2161 assert_eq!(plan.active_tokens_before, 75);
2162 assert_eq!(plan.message_ids_to_archive, vec!["old-u", "old-a"]);
2163 assert_eq!(plan.archive_message_tokens, 50);
2164 assert_eq!(plan.projected_active_tokens_after, 25);
2165 }
2166
2167 #[test]
2168 fn provider_prepared_multimodal_cost_counts_with_its_logical_group() {
2169 let mut session = Session::new("retrieval-window-multimodal", "test-model");
2170 session.add_message(system("system", 5));
2171 let mut multimodal = user("old-u", 0);
2172 multimodal.content_parts = Some(vec![
2173 MessagePart::Text {
2174 text: "t".repeat(20),
2175 },
2176 MessagePart::ImageUrl {
2177 image_url: bamboo_domain::ImageUrlRef {
2178 url: "i".repeat(20),
2179 detail: Some("high".to_string()),
2180 },
2181 },
2182 ]);
2183 session.add_message(multimodal);
2184 session.add_message(assistant("old-a", 5));
2185 add_turn(&mut session, "recent", 10);
2186
2187 let mut accounting = RetrievalWindowTokenAccounting::default();
2188 accounting
2189 .provider_message_tokens
2190 .insert("old-u".to_string(), 44);
2191
2192 let plan = plan_with_counter_and_accounting(&session, 50, 1, &accounting)
2193 .expect("provider-visible multimodal parts should exceed the target");
2194
2195 assert_eq!(plan.active_tokens_before, 74);
2196 assert_eq!(plan.message_ids_to_archive, vec!["old-u", "old-a"]);
2197 assert_eq!(plan.archive_message_tokens, 49);
2198 assert_eq!(plan.projected_active_tokens_after, 25);
2199 assert_eq!(plan.provider_message_token_override_count, 1);
2200 }
2201
2202 #[test]
2203 fn image_without_provider_prepared_cost_fails_closed() {
2204 let mut session = Session::new("retrieval-window-attachment", "test-model");
2205 session.add_message(system("system", 5));
2206 let mut image = user("old-u", 0);
2207 image.content_parts = Some(vec![MessagePart::ImageUrl {
2208 image_url: bamboo_domain::ImageUrlRef {
2209 url: "bamboo-attachment://retrieval-window-attachment/image-1".to_string(),
2210 detail: Some("high".to_string()),
2211 },
2212 }]);
2213 session.add_message(image);
2214 session.add_message(assistant("old-a", 5));
2215 add_turn(&mut session, "recent", 10);
2216
2217 assert_eq!(
2218 plan_with_counter(&session, 50, 1, 0),
2219 Err(
2220 RetrievalWindowPlanError::MissingProviderMessageTokenEstimate {
2221 message_id: "old-u".to_string(),
2222 }
2223 )
2224 );
2225 }
2226
2227 #[test]
2228 fn oversized_latest_turn_returns_protected_content_error() {
2229 let mut session = Session::new("retrieval-window-latest-overflow", "test-model");
2230 session.add_message(system("system", 5));
2231 add_turn(&mut session, "old", 10);
2232 add_turn(&mut session, "recent", 30);
2233
2234 let error = plan_with_counter(&session, 50, 1, 0)
2235 .expect_err("latest protected turn cannot meet target");
2236
2237 assert_eq!(
2238 error,
2239 RetrievalWindowPlanError::ProtectedContentExceedsTarget {
2240 projected_tokens: 65,
2241 target_tokens: 50,
2242 protected_active_tokens: 65,
2243 fixed_prompt_tokens: 0,
2244 incomplete_protocol_group_count: 0,
2245 }
2246 );
2247 }
2248
2249 #[test]
2250 fn invalid_policy_and_already_satisfied_target_are_explicit() {
2251 let mut session = Session::new("retrieval-window-inputs", "test-model");
2252 session.add_message(system("system", 5));
2253 add_turn(&mut session, "one", 10);
2254 add_turn(&mut session, "two", 10);
2255
2256 assert_eq!(
2257 build_retrieval_window_candidate_plan_with_counter(
2258 &session,
2259 &budget(100),
2260 policy(0, 50),
2261 &RetrievalWindowTokenAccounting::default(),
2262 false,
2263 &CharacterTokenCounter,
2264 ),
2265 Err(RetrievalWindowPlanError::InvalidRecentUserTurnFloor)
2266 );
2267 assert_eq!(
2268 build_retrieval_window_candidate_plan_with_counter(
2269 &session,
2270 &budget(100),
2271 policy(1, 0),
2272 &RetrievalWindowTokenAccounting::default(),
2273 false,
2274 &CharacterTokenCounter,
2275 ),
2276 Err(RetrievalWindowPlanError::InvalidTargetUsagePercent {
2277 target_usage_percent: 0,
2278 })
2279 );
2280 assert_eq!(
2281 plan_with_counter(&session, 50, 1, 0),
2282 Err(RetrievalWindowPlanError::TargetAlreadySatisfied {
2283 active_tokens: 45,
2284 target_tokens: 50,
2285 })
2286 );
2287 }
2288
2289 #[test]
2290 fn critical_provider_overflow_archives_every_eligible_group_below_configured_target() {
2291 let mut session = Session::new("retrieval-window-provider-overflow", "test-model");
2292 session.add_message(system("system", 5));
2293 add_turn(&mut session, "one", 10);
2294 add_turn(&mut session, "two", 10);
2295 add_turn(&mut session, "three", 10);
2296 add_turn(&mut session, "four", 10);
2297
2298 let plan = build_retrieval_window_candidate_plan_with_counter(
2299 &session,
2300 &budget(200),
2301 policy(1, 50),
2302 &RetrievalWindowTokenAccounting::default(),
2303 true,
2304 &CharacterTokenCounter,
2305 )
2306 .expect("provider overflow must override a locally satisfied target");
2307
2308 assert_eq!(plan.active_tokens_before, 85);
2309 assert_eq!(plan.target_tokens, 25);
2310 assert_eq!(plan.archive_group_count, 3);
2311 assert_eq!(plan.retained_user_turn_count, 1);
2312 assert!(plan.projected_active_tokens_after <= plan.target_tokens);
2313 let mut committed = session;
2314 apply_retrieval_window_plan(
2315 &mut committed,
2316 &plan,
2317 policy(1, 50),
2318 &budget(200),
2319 &RetrievalWindowTokenAccounting::default(),
2320 )
2321 .expect("the sealed emergency target must remain valid at application");
2322 assert!(committed.messages.iter().any(|message| message.compressed));
2323 }
2324
2325 #[test]
2326 fn critical_provider_overflow_archives_every_eligible_group_above_configured_target() {
2327 let mut session = Session::new("retrieval-window-provider-overflow-above", "test-model");
2328 session.add_message(system("system", 5));
2329 add_turn(&mut session, "one", 10);
2330 add_turn(&mut session, "two", 10);
2331 add_turn(&mut session, "three", 10);
2332 add_turn(&mut session, "four", 10);
2333
2334 let plan = build_retrieval_window_candidate_plan_with_counter(
2338 &session,
2339 &budget(100),
2340 policy(1, 20),
2341 &RetrievalWindowTokenAccounting::default(),
2342 true,
2343 &CharacterTokenCounter,
2344 )
2345 .expect("provider overflow must maximize safe headroom above the local target");
2346
2347 assert_eq!(plan.active_tokens_before, 85);
2348 assert_eq!(plan.target_tokens, 25);
2349 assert_eq!(plan.projected_active_tokens_after, 25);
2350 assert_eq!(plan.archive_group_count, 3);
2351 assert_eq!(plan.retained_user_turn_count, 1);
2352 let mut committed = session;
2353 apply_retrieval_window_plan_with_trigger(
2354 &mut committed,
2355 &plan,
2356 policy(1, 20),
2357 &budget(100),
2358 &RetrievalWindowTokenAccounting::default(),
2359 CompressionTriggerType::CriticalOverflow,
2360 )
2361 .expect("the sealed protected-floor target must remain valid at application");
2362 assert_eq!(
2363 committed
2364 .compression_events
2365 .last()
2366 .expect("critical overflow event")
2367 .retrieval_archived_group_count,
2368 3
2369 );
2370 }
2371
2372 #[test]
2373 fn critical_provider_overflow_rejects_a_protected_floor_above_the_hard_input_limit() {
2374 let mut session = Session::new(
2375 "retrieval-window-provider-overflow-hard-limit",
2376 "test-model",
2377 );
2378 session.add_message(system("system", 5));
2379 add_turn(&mut session, "old", 10);
2380 add_turn(&mut session, "recent", 30);
2381
2382 let error = build_retrieval_window_candidate_plan_with_counter(
2383 &session,
2384 &budget(50),
2385 policy(1, 20),
2386 &RetrievalWindowTokenAccounting::default(),
2387 true,
2388 &CharacterTokenCounter,
2389 )
2390 .expect_err("a provider retry cannot fit protected content beyond the hard input limit");
2391
2392 assert_eq!(
2393 error,
2394 RetrievalWindowPlanError::ProtectedContentExceedsTarget {
2395 projected_tokens: 65,
2396 target_tokens: 50,
2397 protected_active_tokens: 65,
2398 fixed_prompt_tokens: 0,
2399 incomplete_protocol_group_count: 0,
2400 }
2401 );
2402 }
2403
2404 #[test]
2405 fn applies_summary_free_boundary_and_exact_replay_is_a_noop() {
2406 let (mut session, plan) = basic_session_and_plan();
2407 let transcript_before = session
2408 .messages
2409 .iter()
2410 .map(|message| {
2411 (
2412 message.id.clone(),
2413 message.role.clone(),
2414 message.content.clone(),
2415 )
2416 })
2417 .collect::<Vec<_>>();
2418
2419 session.model_context_state = Some(ModelContextState {
2420 prefix_epoch: 7,
2421 cache_scope_sha256: Some("prepared-scope".to_string()),
2422 ..ModelContextState::default()
2423 });
2424 let provider_boundary =
2425 provider_transcript_boundary_sha256(Some("provider-a"), Some("openai"))
2426 .expect("provider boundary");
2427 session
2428 .activate_provider_transcript_route(
2429 ProviderFamily::OpenAi,
2430 ProviderProtocol::OpenAiResponsesV1,
2431 &provider_boundary,
2432 )
2433 .expect("route should activate");
2434 let model_epoch_before = session
2435 .model_context_state
2436 .as_ref()
2437 .expect("model context")
2438 .prefix_epoch;
2439 let provider_epoch_before = session.provider_transcript.epoch();
2440
2441 let result = apply_test_plan(&mut session, &plan).expect("plan should apply");
2442
2443 assert!(!result.idempotent_replay);
2444 assert_eq!(result.newly_archived_message_count, 4);
2445 assert_eq!(session.messages.len(), transcript_before.len());
2446 assert_eq!(
2447 session
2448 .messages
2449 .iter()
2450 .map(|message| {
2451 (
2452 message.id.clone(),
2453 message.role.clone(),
2454 message.content.clone(),
2455 )
2456 })
2457 .collect::<Vec<_>>(),
2458 transcript_before,
2459 "archive flags must not rewrite authoritative transcript content or order"
2460 );
2461 assert!(session.conversation_summary.is_none());
2462 assert_eq!(session.compression_events.len(), 1);
2463
2464 let event = &session.compression_events[0];
2465 assert_eq!(event.id, result.event_id);
2466 assert_eq!(event.kind, CompressionEventKind::RetrievalWindow);
2467 assert_eq!(event.messages_compressed, 4);
2468 assert_eq!(event.segments_removed, 2);
2469 assert_eq!(event.summary_tokens, 0);
2470 assert_eq!(event.actual_summary_tokens, 0);
2471 assert!(event.model_used.is_none());
2472 assert_eq!(event.retrieval_active_tokens_before, 90);
2473 assert_eq!(event.retrieval_active_message_count_before, 9);
2474 assert_eq!(
2475 event.retrieval_active_state_sha256.as_deref(),
2476 Some(plan.active_state_sha256.as_str())
2477 );
2478 assert_eq!(
2479 event.retrieval_token_accounting_sha256.as_deref(),
2480 Some(plan.token_accounting_sha256.as_str())
2481 );
2482 assert_eq!(event.retrieval_active_tokens_after, 50);
2483 assert_eq!(event.retrieval_target_tokens, 50);
2484 assert_eq!(event.retrieval_archived_message_tokens, 40);
2485 assert_eq!(event.retrieval_boundary_reclaimed_tokens, 0);
2486 assert_eq!(event.retrieval_system_message_tokens, 5);
2487 assert_eq!(event.retrieval_context_window_tokens, 100);
2488 assert_eq!(event.retrieval_request_input_limit_tokens, 100);
2489 assert_eq!(event.retrieval_min_recent_user_turns, 2);
2490
2491 let archived_ids = plan
2492 .message_ids_to_archive
2493 .iter()
2494 .map(String::as_str)
2495 .collect::<HashSet<_>>();
2496 for message in &session.messages {
2497 if archived_ids.contains(message.id.as_str()) {
2498 assert!(message.compressed);
2499 assert_eq!(
2500 message.compressed_by_event_id.as_deref(),
2501 Some(result.event_id.as_str())
2502 );
2503 } else {
2504 assert!(!message.compressed);
2505 assert!(message.compressed_by_event_id.is_none());
2506 }
2507 }
2508
2509 let usage = session.token_usage.as_ref().expect("token usage snapshot");
2510 assert_eq!(usage.system_tokens, 10);
2511 assert_eq!(usage.summary_tokens, 0);
2512 assert_eq!(usage.window_tokens, 40);
2513 assert_eq!(usage.total_tokens, 50);
2514 assert_eq!(usage.max_context_tokens, 100);
2515 assert_eq!(usage.budget_limit, 100);
2516 assert_eq!(usage.segments_removed, 2);
2517
2518 assert_eq!(
2519 session
2520 .model_context_state
2521 .as_ref()
2522 .expect("model context")
2523 .prefix_epoch,
2524 model_epoch_before + 1
2525 );
2526 assert_eq!(
2527 session.provider_transcript.epoch(),
2528 provider_epoch_before + 1
2529 );
2530 assert_eq!(
2531 session.provider_transcript.last_reset_reason(),
2532 Some(ProviderTranscriptResetReason::Compression)
2533 );
2534
2535 let round_trip: Session = serde_json::from_slice(
2536 &serde_json::to_vec(&session).expect("session should serialize"),
2537 )
2538 .expect("session should deserialize");
2539 assert_eq!(
2540 round_trip.compression_events[0].kind,
2541 CompressionEventKind::RetrievalWindow
2542 );
2543 assert_eq!(
2544 round_trip.compression_events[0].retrieval_active_tokens_after,
2545 50
2546 );
2547
2548 let serialized_after_first =
2549 serde_json::to_vec(&session).expect("session should serialize");
2550 let model_epoch_after_first = session
2551 .model_context_state
2552 .as_ref()
2553 .expect("model context")
2554 .prefix_epoch;
2555 let provider_epoch_after_first = session.provider_transcript.epoch();
2556 let replay =
2557 apply_test_plan(&mut session, &plan).expect("exact replay should be idempotent");
2558 assert!(replay.idempotent_replay);
2559 assert_eq!(replay.event_id, result.event_id);
2560 assert_eq!(replay.newly_archived_message_count, 0);
2561 assert_eq!(
2562 serde_json::to_vec(&session).expect("session should serialize"),
2563 serialized_after_first
2564 );
2565 assert_eq!(
2566 session
2567 .model_context_state
2568 .as_ref()
2569 .expect("model context")
2570 .prefix_epoch,
2571 model_epoch_after_first
2572 );
2573 assert_eq!(
2574 session.provider_transcript.epoch(),
2575 provider_epoch_after_first
2576 );
2577 }
2578
2579 #[test]
2580 fn trigger_aware_boundary_resets_provider_replay_across_family_restart_matrix() {
2581 let routes = [
2582 (
2583 ProviderFamily::OpenAi,
2584 ProviderProtocol::OpenAiResponsesV1,
2585 "openai",
2586 ),
2587 (
2588 ProviderFamily::Anthropic,
2589 ProviderProtocol::AnthropicMessages2023_06_01,
2590 "anthropic",
2591 ),
2592 (
2593 ProviderFamily::Copilot,
2594 ProviderProtocol::OpenAiResponsesV1,
2595 "copilot",
2596 ),
2597 ];
2598
2599 for (family, protocol, provider_type) in routes {
2600 let (mut session, plan) = basic_session_and_plan();
2601 let boundary = provider_transcript_boundary_sha256(
2602 Some(&format!("{provider_type}-retrieval-test")),
2603 Some(provider_type),
2604 )
2605 .expect("provider boundary");
2606 session
2607 .activate_provider_transcript_route(family, protocol, &boundary)
2608 .expect("provider route should activate");
2609 let items = if family == ProviderFamily::Anthropic {
2610 vec![
2611 ProviderTranscriptItem::try_from_payload(
2612 family,
2613 protocol,
2614 ProviderTranscriptOrigin::Provider,
2615 ProviderTranscriptAuthor::Model,
2616 serde_json::json!({
2617 "type":"server_tool_use",
2618 "id":"srvtoolu_archive_matrix",
2619 "name":"tool_search_tool_regex",
2620 "input":{"pattern":"history"}
2621 }),
2622 )
2623 .expect("anthropic search call"),
2624 ProviderTranscriptItem::try_from_payload(
2625 family,
2626 protocol,
2627 ProviderTranscriptOrigin::Provider,
2628 ProviderTranscriptAuthor::ToolResult,
2629 serde_json::json!({
2630 "type":"tool_search_tool_result",
2631 "tool_use_id":"srvtoolu_archive_matrix",
2632 "content":{
2633 "type":"tool_search_tool_search_result",
2634 "tool_references":[{
2635 "type":"tool_reference",
2636 "tool_name":"session_history_current"
2637 }]
2638 }
2639 }),
2640 )
2641 .expect("anthropic search result"),
2642 ]
2643 } else {
2644 vec![ProviderTranscriptItem::try_from_payload(
2645 family,
2646 protocol,
2647 ProviderTranscriptOrigin::Provider,
2648 ProviderTranscriptAuthor::Model,
2649 serde_json::json!({
2650 "type":"tool_search_call",
2651 "id":"tsc_archive_matrix",
2652 "execution":"client",
2653 "call_id":"search_archive_matrix",
2654 "status":"completed",
2655 "arguments":{"query":"history"}
2656 }),
2657 )
2658 .expect("OpenAI-family search call")]
2659 };
2660 session
2661 .append_provider_transcript_group("t4-a", None, items)
2662 .expect("provider-native group should append");
2663 assert_eq!(
2664 session
2665 .provider_transcript
2666 .replayable_groups(family, protocol, &boundary)
2667 .len(),
2668 1
2669 );
2670 let provider_epoch_before = session.provider_transcript.epoch();
2671
2672 apply_retrieval_window_plan_with_trigger(
2673 &mut session,
2674 &plan,
2675 policy_for_plan(&plan),
2676 &budget(100),
2677 &accounting_for_plan(&plan),
2678 CompressionTriggerType::CriticalOverflow,
2679 )
2680 .expect("trigger-aware boundary should apply");
2681
2682 assert_eq!(
2683 session
2684 .compression_events
2685 .last()
2686 .expect("retrieval event")
2687 .trigger_type,
2688 CompressionTriggerType::CriticalOverflow
2689 );
2690 assert_eq!(
2691 session.provider_transcript.epoch(),
2692 provider_epoch_before + 1
2693 );
2694 assert!(session
2695 .provider_transcript
2696 .replayable_groups(family, protocol, &boundary)
2697 .is_empty());
2698
2699 let restarted: Session = serde_json::from_slice(
2700 &serde_json::to_vec(&session).expect("session should serialize"),
2701 )
2702 .expect("session should reload");
2703 assert_eq!(
2704 restarted.provider_transcript.last_reset_reason(),
2705 Some(ProviderTranscriptResetReason::Compression)
2706 );
2707 assert!(restarted
2708 .provider_transcript
2709 .replayable_groups(family, protocol, &boundary)
2710 .is_empty());
2711 }
2712 }
2713
2714 #[test]
2715 fn trigger_aware_idempotent_replay_requires_the_original_trigger() {
2716 let (mut session, plan) = basic_session_and_plan();
2717 let first = apply_retrieval_window_plan_with_trigger(
2718 &mut session,
2719 &plan,
2720 policy_for_plan(&plan),
2721 &budget(100),
2722 &accounting_for_plan(&plan),
2723 CompressionTriggerType::Manual,
2724 )
2725 .expect("manual boundary should apply");
2726 assert_eq!(
2727 session.compression_events[0].trigger_type,
2728 CompressionTriggerType::Manual
2729 );
2730
2731 let replay = apply_retrieval_window_plan_with_trigger(
2732 &mut session,
2733 &plan,
2734 policy_for_plan(&plan),
2735 &budget(100),
2736 &accounting_for_plan(&plan),
2737 CompressionTriggerType::Manual,
2738 )
2739 .expect("the same manual request should be idempotent");
2740 assert!(replay.idempotent_replay);
2741 assert_eq!(replay.event_id, first.event_id);
2742
2743 assert!(matches!(
2744 apply_retrieval_window_plan(
2745 &mut session,
2746 &plan,
2747 policy_for_plan(&plan),
2748 &budget(100),
2749 &accounting_for_plan(&plan),
2750 ),
2751 Err(RetrievalWindowApplyError::ArchiveEventEvidenceMismatch {
2752 field: "trigger_type",
2753 ..
2754 })
2755 ));
2756 }
2757
2758 #[test]
2759 fn later_plan_archives_additional_groups_without_reassigning_old_messages() {
2760 let (mut session, first_plan) = basic_session_and_plan();
2761 let first = apply_test_plan(&mut session, &first_plan).expect("first plan should apply");
2762 add_turn(&mut session, "t5", 10);
2763 add_turn(&mut session, "t6", 10);
2764
2765 let second_plan = plan_with_counter(&session, 50, 2, 5).expect("second plan should build");
2766 assert_eq!(
2767 second_plan.message_ids_to_archive,
2768 vec!["t3-u", "t3-a", "t4-u", "t4-a"]
2769 );
2770 let second = apply_test_plan(&mut session, &second_plan).expect("second plan should apply");
2771
2772 assert_ne!(second.event_id, first.event_id);
2773 assert_eq!(session.compression_events.len(), 2);
2774 for message_id in &first_plan.message_ids_to_archive {
2775 let message = session
2776 .messages
2777 .iter()
2778 .find(|message| &message.id == message_id)
2779 .expect("first message remains exact");
2780 assert_eq!(
2781 message.compressed_by_event_id.as_deref(),
2782 Some(first.event_id.as_str())
2783 );
2784 }
2785 for message_id in &second_plan.message_ids_to_archive {
2786 let message = session
2787 .messages
2788 .iter()
2789 .find(|message| &message.id == message_id)
2790 .expect("second message remains exact");
2791 assert_eq!(
2792 message.compressed_by_event_id.as_deref(),
2793 Some(second.event_id.as_str())
2794 );
2795 }
2796 }
2797
2798 #[test]
2799 fn missing_duplicate_system_and_never_compress_candidates_fail_transactionally() {
2800 let (mut session, mut plan) = basic_session_and_plan();
2801 plan.message_ids_to_archive[0] = "missing".to_string();
2802 assert_apply_error_without_mutation(
2803 &mut session,
2804 &plan,
2805 RetrievalWindowApplyError::MissingCandidateMessage {
2806 message_id: "missing".to_string(),
2807 },
2808 );
2809
2810 let (mut session, mut plan) = basic_session_and_plan();
2811 plan.message_ids_to_archive[1] = plan.message_ids_to_archive[0].clone();
2812 assert_apply_error_without_mutation(
2813 &mut session,
2814 &plan,
2815 RetrievalWindowApplyError::DuplicateCandidateId {
2816 message_id: "t1-u".to_string(),
2817 },
2818 );
2819
2820 let (mut session, plan) = basic_session_and_plan();
2821 session
2822 .messages
2823 .iter_mut()
2824 .find(|message| message.id == "t4-a")
2825 .expect("retained message")
2826 .id = "t1-u".to_string();
2827 assert_apply_error_without_mutation(
2828 &mut session,
2829 &plan,
2830 RetrievalWindowApplyError::DuplicateSessionMessageId {
2831 message_id: "t1-u".to_string(),
2832 },
2833 );
2834
2835 let (mut session, mut plan) = basic_session_and_plan();
2836 plan.message_ids_to_archive[0] = "system".to_string();
2837 assert_apply_error_without_mutation(
2838 &mut session,
2839 &plan,
2840 RetrievalWindowApplyError::SystemMessageCandidate {
2841 message_id: "system".to_string(),
2842 },
2843 );
2844
2845 let (mut session, plan) = basic_session_and_plan();
2846 session
2847 .messages
2848 .iter_mut()
2849 .find(|message| message.id == "t1-u")
2850 .expect("candidate message")
2851 .never_compress = true;
2852 assert_apply_error_without_mutation(
2853 &mut session,
2854 &plan,
2855 RetrievalWindowApplyError::NeverCompressCandidate {
2856 message_id: "t1-u".to_string(),
2857 },
2858 );
2859 }
2860
2861 #[test]
2862 fn inconsistent_summary_and_stale_plans_fail_before_mutation() {
2863 let (mut session, mut plan) = basic_session_and_plan();
2864 plan.archive_message_tokens += 1;
2865 assert_apply_error_without_mutation(
2866 &mut session,
2867 &plan,
2868 RetrievalWindowApplyError::InconsistentPlan {
2869 field: "active_token_arithmetic",
2870 },
2871 );
2872
2873 let (mut session, plan) = basic_session_and_plan();
2874 session.conversation_summary = Some(ConversationSummary::new("existing", 1, 1));
2875 assert_apply_error_without_mutation(
2876 &mut session,
2877 &plan,
2878 RetrievalWindowApplyError::PreExistingConversationSummary,
2879 );
2880
2881 let (mut session, plan) = basic_session_and_plan();
2882 add_turn(&mut session, "newer", 10);
2883 assert_apply_error_without_mutation(
2884 &mut session,
2885 &plan,
2886 RetrievalWindowApplyError::StalePlan {
2887 invariant: "active_message_count",
2888 },
2889 );
2890
2891 let (mut session, plan) = basic_session_and_plan();
2892 session
2893 .messages
2894 .iter_mut()
2895 .find(|message| message.id == "t1-u")
2896 .expect("candidate message")
2897 .content
2898 .push_str("changed-after-planning");
2899 assert_apply_error_without_mutation(
2900 &mut session,
2901 &plan,
2902 RetrievalWindowApplyError::StalePlan {
2903 invariant: "active_state_sha256",
2904 },
2905 );
2906 }
2907
2908 #[test]
2909 fn authenticated_plan_rejects_self_consistent_candidate_token_tampering() {
2910 let (mut session, mut plan) = basic_session_and_plan();
2911
2912 plan.message_ids_to_archive.truncate(2);
2917 plan.archive_message_count = 2;
2918 plan.archive_group_count = 1;
2919 plan.archive_user_turn_count = 1;
2920 plan.retained_user_turn_count = 3;
2921 plan.oldest_retained_message_id = Some("t2-u".to_string());
2922 plan.oldest_retained_user_message_id = Some("t2-u".to_string());
2923
2924 assert_apply_error_without_mutation(
2925 &mut session,
2926 &plan,
2927 RetrievalWindowApplyError::StalePlan {
2928 invariant: "plan_evidence_sha256",
2929 },
2930 );
2931 }
2932
2933 #[test]
2934 fn active_state_digest_covers_provider_visible_message_fields() {
2935 let (session, _) = safe_tool_session_and_plan();
2936 let original = active_state_sha256(&session);
2937
2938 let mut reasoning_changed = session.clone();
2939 reasoning_changed.messages[1].reasoning = Some("new reasoning".to_string());
2940 assert_ne!(active_state_sha256(&reasoning_changed), original);
2941
2942 let mut parts_changed = session.clone();
2943 parts_changed.messages[1].content_parts = Some(vec![MessagePart::Text {
2944 text: "provider-visible part".to_string(),
2945 }]);
2946 assert_ne!(active_state_sha256(&parts_changed), original);
2947
2948 let mut tool_arguments_changed = session;
2949 tool_arguments_changed.messages[2]
2950 .tool_calls
2951 .as_mut()
2952 .and_then(|calls| calls.first_mut())
2953 .expect("tool call")
2954 .function
2955 .arguments = "{\"changed\":true}".to_string();
2956 assert_ne!(active_state_sha256(&tool_arguments_changed), original);
2957 }
2958
2959 #[test]
2960 fn token_accounting_digest_covers_fixed_boundary_and_provider_message_costs() {
2961 let mut accounting = RetrievalWindowTokenAccounting {
2962 fixed_prompt_tokens: 5,
2963 ..RetrievalWindowTokenAccounting::default()
2964 };
2965 accounting
2966 .provider_message_tokens
2967 .insert("t1-u".to_string(), 10);
2968 let original = token_accounting_sha256(&accounting);
2969
2970 let mut fixed_changed = accounting.clone();
2971 fixed_changed.fixed_prompt_tokens += 1;
2972 assert_ne!(token_accounting_sha256(&fixed_changed), original);
2973
2974 let mut boundary_changed = accounting.clone();
2975 boundary_changed.boundary_reclaimable_tokens += 1;
2976 assert_ne!(token_accounting_sha256(&boundary_changed), original);
2977
2978 let mut override_changed = accounting.clone();
2979 override_changed
2980 .provider_message_tokens
2981 .insert("t1-u".to_string(), 11);
2982 assert_ne!(token_accounting_sha256(&override_changed), original);
2983
2984 let mut override_added = accounting;
2985 override_added
2986 .provider_message_tokens
2987 .insert("t2-u".to_string(), 10);
2988 assert_ne!(token_accounting_sha256(&override_added), original);
2989 }
2990
2991 #[test]
2992 fn changed_token_accounting_fails_before_mutation() {
2993 let mut session = Session::new("retrieval-window-accounting-stale", "test-model");
2994 session.add_message(system("system", 5));
2995 for turn in 1..=4 {
2996 add_turn(&mut session, &format!("t{turn}"), 10);
2997 }
2998 let mut accounting = RetrievalWindowTokenAccounting {
2999 fixed_prompt_tokens: 5,
3000 ..RetrievalWindowTokenAccounting::default()
3001 };
3002 accounting
3003 .provider_message_tokens
3004 .insert("t1-u".to_string(), 10);
3005 let plan = plan_with_counter_and_accounting(&session, 50, 2, &accounting)
3006 .expect("provider-prepared plan should build");
3007 let before = serde_json::to_vec(&session).expect("session should serialize");
3008
3009 let mut changed_override = accounting.clone();
3010 changed_override
3011 .provider_message_tokens
3012 .insert("t1-u".to_string(), 11);
3013 assert_eq!(
3014 apply_retrieval_window_plan(
3015 &mut session,
3016 &plan,
3017 policy_for_plan(&plan),
3018 &budget(100),
3019 &changed_override,
3020 ),
3021 Err(RetrievalWindowApplyError::StalePlan {
3022 invariant: "token_accounting_sha256",
3023 })
3024 );
3025 assert_eq!(
3026 serde_json::to_vec(&session).expect("session should serialize"),
3027 before
3028 );
3029
3030 let mut changed_boundary = accounting.clone();
3031 changed_boundary.boundary_reclaimable_tokens += 1;
3032 assert_eq!(
3033 apply_retrieval_window_plan(
3034 &mut session,
3035 &plan,
3036 policy_for_plan(&plan),
3037 &budget(100),
3038 &changed_boundary,
3039 ),
3040 Err(RetrievalWindowApplyError::StalePlan {
3041 invariant: "token_accounting_sha256",
3042 })
3043 );
3044 assert_eq!(
3045 serde_json::to_vec(&session).expect("session should serialize"),
3046 before
3047 );
3048
3049 let mut changed_fixed = accounting;
3050 changed_fixed.fixed_prompt_tokens += 1;
3051 assert_eq!(
3052 apply_retrieval_window_plan(
3053 &mut session,
3054 &plan,
3055 policy_for_plan(&plan),
3056 &budget(100),
3057 &changed_fixed,
3058 ),
3059 Err(RetrievalWindowApplyError::StalePlan {
3060 invariant: "token_accounting_sha256",
3061 })
3062 );
3063 assert_eq!(
3064 serde_json::to_vec(&session).expect("session should serialize"),
3065 before
3066 );
3067 }
3068
3069 #[test]
3070 fn changed_retrieval_window_policy_fails_before_mutation() {
3071 let (mut session, plan) = basic_session_and_plan();
3072 let accounting = accounting_for_plan(&plan);
3073 let before = serde_json::to_vec(&session).expect("session should serialize");
3074 let changed_policies = [
3075 RetrievalWindowPolicy {
3076 min_recent_user_turns: plan.min_recent_user_turns + 1,
3077 target_usage_percent: plan.target_usage_percent,
3078 },
3079 RetrievalWindowPolicy {
3080 min_recent_user_turns: plan.min_recent_user_turns,
3081 target_usage_percent: plan.target_usage_percent - 1,
3082 },
3083 ];
3084
3085 for current_policy in changed_policies {
3086 assert_eq!(
3087 apply_retrieval_window_plan(
3088 &mut session,
3089 &plan,
3090 current_policy,
3091 &budget(100),
3092 &accounting,
3093 ),
3094 Err(RetrievalWindowApplyError::StalePlan {
3095 invariant: "retrieval_window_policy",
3096 })
3097 );
3098 assert_eq!(
3099 serde_json::to_vec(&session).expect("session should serialize"),
3100 before
3101 );
3102 }
3103 }
3104
3105 #[test]
3106 fn changed_token_budget_fails_before_mutation() {
3107 let (mut session, plan) = basic_session_and_plan();
3108 let accounting = accounting_for_plan(&plan);
3109 let before = serde_json::to_vec(&session).expect("session should serialize");
3110
3111 let smaller_context = budget(64);
3112 let mut larger_output_reserve = budget(100);
3113 larger_output_reserve.max_output_tokens = 10;
3114 let mut larger_safety_margin = budget(100);
3115 larger_safety_margin.safety_margin = 5;
3116
3117 for current_budget in [smaller_context, larger_output_reserve, larger_safety_margin] {
3118 assert_eq!(
3119 apply_retrieval_window_plan(
3120 &mut session,
3121 &plan,
3122 policy_for_plan(&plan),
3123 ¤t_budget,
3124 &accounting,
3125 ),
3126 Err(RetrievalWindowApplyError::StalePlan {
3127 invariant: "token_budget",
3128 })
3129 );
3130 assert_eq!(
3131 serde_json::to_vec(&session).expect("session should serialize"),
3132 before
3133 );
3134 }
3135 }
3136
3137 #[test]
3138 fn legacy_skill_retrieval_canonical_names_do_not_narrow_roles() {
3139 for role in [Role::Assistant, Role::User, Role::Tool, Role::System] {
3140 for (name, expected) in [
3141 ("load_skill", true),
3142 ("read_skill_resource", true),
3143 ("default::LoAd_SkIlL", true),
3144 ("outer::inner::READ_SKILL_RESOURCE", true),
3145 ("mcp__server__load_skill", false),
3146 ("load_skill_extra", false),
3147 ] {
3148 let mut message = tool_call("candidate", "candidate-call", 1, name);
3149 message.role = role.clone();
3150 let group = LogicalGroup::preamble(IndexedMessage {
3151 session_index: 0,
3152 message: &message,
3153 });
3154 assert_eq!(
3155 logical_group_has_protected_skill(&group),
3156 expected,
3157 "{role:?}/{name}"
3158 );
3159 }
3160 }
3161 }
3162
3163 #[test]
3164 fn legacy_skill_mixed_retrieval_rejects_protected_and_incomplete_groups_without_mutation() {
3165 let mut session = Session::new("mixed-retrieval-skill", "test-model");
3166 session.add_message(system("system", 5));
3167 session.add_message(user("old-u", 10));
3168 let mut mixed = tool_call("mixed-call", "skill-call", 10, "Read");
3169 mixed.tool_calls.as_mut().unwrap().push(ToolCall {
3170 id: "other-call".into(),
3171 tool_type: "function".into(),
3172 function: FunctionCall {
3173 name: "Grep".into(),
3174 arguments: "{}".into(),
3175 },
3176 });
3177 session.add_message(mixed);
3178 session.add_message(tool_result("skill-result", "skill-call", 10));
3179 session.add_message(tool_result("other-result", "other-call", 10));
3180 session.add_message(assistant("old-final", 10));
3181 add_turn(&mut session, "recent", 10);
3182 let plan = plan_with_counter(&session, 50, 1, 0).unwrap();
3183 assert_eq!(
3184 plan.message_ids_to_archive,
3185 [
3186 "old-u",
3187 "mixed-call",
3188 "skill-result",
3189 "other-result",
3190 "old-final"
3191 ]
3192 );
3193 for name in ["namespace::LoAd_SkIlL", "READ_SKILL_RESOURCE"] {
3194 let mut candidate = session.clone();
3195 candidate.messages[2].tool_calls.as_mut().unwrap()[0]
3196 .function
3197 .name = name.into();
3198 assert_apply_error_without_mutation(
3199 &mut candidate,
3200 &plan,
3201 RetrievalWindowApplyError::ProtectedSkillChain,
3202 );
3203 let mut partial = plan.clone();
3204 partial
3205 .message_ids_to_archive
3206 .retain(|id| id != "other-result");
3207 partial.archive_message_count -= 1;
3208 assert_apply_error_without_mutation(
3209 &mut candidate,
3210 &partial,
3211 RetrievalWindowApplyError::IncompleteLogicalGroup,
3212 );
3213 candidate.messages[4].tool_call_id = Some("missing-call".into());
3214 assert_apply_error_without_mutation(
3215 &mut candidate,
3216 &plan,
3217 RetrievalWindowApplyError::UnsafeToolProtocol,
3218 );
3219 }
3220 }
3221
3222 #[test]
3223 fn incomplete_tool_and_skill_groups_fail_before_mutation() {
3224 let (mut session, mut plan) = basic_session_and_plan();
3225 plan.message_ids_to_archive.remove(1);
3226 plan.archive_message_count -= 1;
3227 assert_apply_error_without_mutation(
3228 &mut session,
3229 &plan,
3230 RetrievalWindowApplyError::IncompleteLogicalGroup,
3231 );
3232
3233 let (mut session, plan) = safe_tool_session_and_plan();
3234 session
3235 .messages
3236 .iter_mut()
3237 .find(|message| message.id == "old-result-1")
3238 .expect("tool result")
3239 .tool_call_id = Some("wrong-call".to_string());
3240 assert_apply_error_without_mutation(
3241 &mut session,
3242 &plan,
3243 RetrievalWindowApplyError::UnsafeToolProtocol,
3244 );
3245
3246 let (mut session, plan) = safe_tool_session_and_plan();
3247 session
3248 .messages
3249 .iter_mut()
3250 .find(|message| message.id == "old-call-1")
3251 .and_then(|message| message.tool_calls.as_mut())
3252 .and_then(|calls| calls.first_mut())
3253 .expect("tool call")
3254 .function
3255 .name = "namespace::load_skill".to_string();
3256 assert_apply_error_without_mutation(
3257 &mut session,
3258 &plan,
3259 RetrievalWindowApplyError::ProtectedSkillChain,
3260 );
3261 }
3262
3263 #[test]
3264 fn partial_mixed_and_wrong_strategy_replays_are_typed_conflicts() {
3265 let (mut session, plan) = basic_session_and_plan();
3266 session
3267 .messages
3268 .iter_mut()
3269 .find(|message| message.id == "t1-u")
3270 .expect("candidate")
3271 .compressed = true;
3272 assert_apply_error_without_mutation(
3273 &mut session,
3274 &plan,
3275 RetrievalWindowApplyError::PartialApplication,
3276 );
3277
3278 let (mut session, plan) = basic_session_and_plan();
3279 apply_test_plan(&mut session, &plan).expect("plan should apply");
3280 session
3281 .messages
3282 .iter_mut()
3283 .find(|message| message.id == "t1-a")
3284 .expect("candidate")
3285 .compressed_by_event_id = Some("different-event".to_string());
3286 assert_apply_error_without_mutation(
3287 &mut session,
3288 &plan,
3289 RetrievalWindowApplyError::MixedArchiveEvents,
3290 );
3291
3292 let (mut session, plan) = basic_session_and_plan();
3293 let result = apply_test_plan(&mut session, &plan).expect("plan should apply");
3294 session.compression_events[0].kind = CompressionEventKind::Summary;
3295 assert_apply_error_without_mutation(
3296 &mut session,
3297 &plan,
3298 RetrievalWindowApplyError::ArchiveEventKindConflict {
3299 event_id: result.event_id,
3300 },
3301 );
3302
3303 let (mut session, plan) = basic_session_and_plan();
3304 let result = apply_test_plan(&mut session, &plan).expect("plan should apply");
3305 session.compression_events.clear();
3306 assert_apply_error_without_mutation(
3307 &mut session,
3308 &plan,
3309 RetrievalWindowApplyError::MissingOrDuplicateArchiveEvent {
3310 event_id: result.event_id,
3311 },
3312 );
3313
3314 let (mut session, plan) = basic_session_and_plan();
3315 let result = apply_test_plan(&mut session, &plan).expect("plan should apply");
3316 session.compression_events[0].retrieval_target_tokens += 1;
3317 assert_apply_error_without_mutation(
3318 &mut session,
3319 &plan,
3320 RetrievalWindowApplyError::ArchiveEventEvidenceMismatch {
3321 event_id: result.event_id,
3322 field: "retrieval_target_tokens",
3323 },
3324 );
3325 }
3326
3327 #[test]
3328 fn skill_loading_turn_remains_protected() {
3329 let mut session = Session::new("retrieval-window-skill", "test-model");
3330 session.add_message(system("system", 5));
3331 session.add_message(user("skill-u", 5));
3332 session.add_message(tool_call("skill-call", "load", 5, "default::LoAd_SkIlL"));
3333 session.add_message(tool_result("skill-result", "load", 5));
3334 session.add_message(assistant("skill-final", 5));
3335 add_turn(&mut session, "eligible", 10);
3336 add_turn(&mut session, "recent", 10);
3337
3338 let plan = plan_with_counter(&session, 50, 1, 0).expect("eligible turn should archive");
3339
3340 assert_eq!(
3341 plan.message_ids_to_archive,
3342 vec!["eligible-u", "eligible-a"]
3343 );
3344 assert_eq!(
3345 plan.oldest_retained_user_message_id.as_deref(),
3346 Some("skill-u")
3347 );
3348 }
3349}