1use super::{Capability, CapabilityLocalization, CapabilityStatus};
7use crate::message::{ContentPart, Message, MessageRole};
8use crate::message_filter::{
9 ExcludedNoticeTransform, MessageFilterProvider, MessageQuery, anchored_window,
10};
11use crate::tool_types::ToolHints;
12use crate::tools::{Tool, ToolExecutionResult};
13use async_trait::async_trait;
14use everruns_core::{tool_context::ToolContext, tool_context::ToolContextService};
15use serde::{Deserialize, Serialize};
16use serde_json::{Value, json};
17use std::cmp::Ordering;
18use std::io::{self, Write};
19use std::sync::Arc;
20
21pub const INFINITY_CONTEXT_CAPABILITY_ID: &str = "infinity_context";
23
24pub struct InfinityContextCapability;
26
27pub struct InfinityContextFilterOnlyCapability;
32
33impl Capability for InfinityContextFilterOnlyCapability {
34 fn id(&self) -> &str {
35 INFINITY_CONTEXT_CAPABILITY_ID
36 }
37
38 fn name(&self) -> &str {
39 "Infinity Context"
40 }
41
42 fn description(&self) -> &str {
43 "Trims older conversation history out of the live provider prompt."
44 }
45
46 fn status(&self) -> CapabilityStatus {
47 CapabilityStatus::Available
48 }
49
50 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
51 Some(Arc::new(InfinityContextFilterProvider))
52 }
53
54 fn message_filter_config(&self, config: &Value, compaction_enabled: bool) -> Value {
55 message_filter_config(config, compaction_enabled)
56 }
57}
58
59impl Capability for InfinityContextCapability {
60 fn id(&self) -> &str {
61 INFINITY_CONTEXT_CAPABILITY_ID
62 }
63
64 fn name(&self) -> &str {
65 "Infinity Context"
66 }
67
68 fn description(&self) -> &str {
69 r#"Trims older conversation history out of the live prompt while keeping it queryable with `query_history`.
70
71> [!TIP]
72> Use this for long-running sessions where earlier discussion still matters but should not consume prompt budget every turn."#
73 }
74
75 fn status(&self) -> CapabilityStatus {
76 CapabilityStatus::Available
77 }
78
79 fn icon(&self) -> Option<&str> {
80 Some("infinity")
81 }
82
83 fn category(&self) -> Option<&str> {
84 Some("Optimization")
85 }
86
87 fn system_prompt_addition(&self) -> Option<&str> {
88 Some(INFINITY_CONTEXT_SYSTEM_PROMPT)
89 }
90
91 fn tools(&self) -> Vec<Box<dyn Tool>> {
92 vec![Box::new(QueryHistoryTool)]
93 }
94
95 fn message_filter_provider(&self) -> Option<Arc<dyn MessageFilterProvider>> {
96 Some(Arc::new(InfinityContextFilterProvider))
97 }
98
99 fn message_filter_config(&self, config: &Value, compaction_enabled: bool) -> Value {
100 message_filter_config(config, compaction_enabled)
101 }
102
103 fn config_schema(&self) -> Option<Value> {
108 Some(json!({
109 "type": "object",
110 "properties": {
111 "context_budget_tokens": {
112 "type": "integer",
113 "title": "Context budget (tokens)",
114 "description": "Maximum prompt budget reserved for message history.",
115 "minimum": 1,
116 "default": default_context_budget_tokens()
117 },
118 "min_recent_messages": {
119 "type": "integer",
120 "title": "Minimum recent messages",
121 "description": "Number of recent messages always kept, even when the token budget is tight.",
122 "minimum": 1,
123 "default": default_min_recent_messages()
124 },
125 "max_recent_messages": {
126 "type": "integer",
127 "title": "Maximum recent messages",
128 "description": "Optional hard cap on recent messages kept in the live prompt.",
129 "minimum": 1
130 },
131 "keep_first_messages": {
132 "type": "integer",
133 "title": "Anchored first messages",
134 "description": "Optional leading messages kept as an anchor (the original task), even under a tight budget. Additional to the maximum recent messages. The anchor is fetched as a bounded head+tail load (capped at 16), so it is guaranteed even for histories far longer than the candidate load window. Defaults to 0 so untrusted first messages cannot bypass the configured token budget or recent-message cap; raise it only for trusted sessions where anchoring leading context is worth the extra prompt cost.",
135 "minimum": 0,
136 "maximum": MAX_KEEP_FIRST_MESSAGES,
137 "default": default_keep_first_messages()
138 }
139 }
140 }))
141 }
142
143 fn validate_config(&self, config: &Value) -> Result<(), String> {
144 if config.is_null() {
145 return Ok(());
146 }
147 let typed: InfinityContextConfig = serde_json::from_value(config.clone())
148 .map_err(|e| format!("invalid infinity_context config: {e}"))?;
149 if typed.context_budget_tokens == 0 {
150 return Err("context_budget_tokens must be >= 1".to_string());
151 }
152 if typed.min_recent_messages == 0 {
153 return Err("min_recent_messages must be >= 1".to_string());
154 }
155 if typed.max_recent_messages == Some(0) {
156 return Err("max_recent_messages must be >= 1".to_string());
157 }
158 if typed.keep_first_messages > MAX_KEEP_FIRST_MESSAGES {
159 return Err(format!(
160 "keep_first_messages must be <= {MAX_KEEP_FIRST_MESSAGES}"
161 ));
162 }
163 Ok(())
164 }
165
166 fn localizations(&self) -> Vec<CapabilityLocalization> {
167 vec![
168 CapabilityLocalization {
169 locale: "en",
170 name: None,
171 description: None,
172 config_description: Some(
173 "Controls the token budget for history and the minimum/maximum number \
174 of recent messages kept in the prompt.",
175 ),
176 config_overlay: None,
177 },
178 CapabilityLocalization {
179 locale: "uk",
180 name: Some("Нескінченний контекст"),
181 description: Some(
182 "Прибирає старішу історію розмови з активного запиту, зберігаючи її \
183 доступною через інструмент query_history.",
184 ),
185 config_description: Some(
186 "Визначає бюджет токенів для історії та мінімальну й максимальну \
187 кількість останніх повідомлень у запиті.",
188 ),
189 config_overlay: Some(json!({
190 "properties": {
191 "context_budget_tokens": {
192 "title": "Бюджет контексту (токени)",
193 "description": "Максимальний бюджет запиту, зарезервований для історії повідомлень."
194 },
195 "min_recent_messages": {
196 "title": "Мінімум останніх повідомлень",
197 "description": "Кількість останніх повідомлень, які зберігаються завжди, навіть коли бюджет токенів обмежений."
198 },
199 "max_recent_messages": {
200 "title": "Максимум останніх повідомлень",
201 "description": "Необов'язкове жорстке обмеження кількості останніх повідомлень в активному запиті."
202 }
203 }
204 })),
205 },
206 ]
207 }
208}
209
210fn message_filter_config(base: &Value, compaction_enabled: bool) -> Value {
211 if !compaction_enabled {
212 return base.clone();
213 }
214 let mut config = base.clone();
215 match config.as_object_mut() {
216 Some(map) => {
217 map.insert("compaction_active".to_string(), Value::Bool(true));
218 }
219 None => config = json!({ "compaction_active": true }),
220 }
221 config
222}
223
224const INFINITY_CONTEXT_SYSTEM_PROMPT: &str = r#"## Conversation history
225
226Earlier messages may be trimmed from the live prompt. Use `query_history`
227to retrieve them when needed. The window is trimmed automatically; do not
228abandon tasks for token reasons — persist important state via file or
229memory tools when available."#;
230
231#[derive(Debug, Clone, Serialize, Deserialize)]
232struct InfinityContextConfig {
233 #[serde(default = "default_context_budget_tokens")]
235 context_budget_tokens: usize,
236
237 #[serde(default = "default_min_recent_messages")]
239 min_recent_messages: usize,
240
241 #[serde(default)]
246 max_recent_messages: Option<usize>,
247
248 #[serde(default = "default_keep_first_messages")]
253 keep_first_messages: usize,
254
255 #[serde(default)]
260 compaction_active: bool,
261}
262
263fn default_context_budget_tokens() -> usize {
264 100_000
265}
266
267fn default_min_recent_messages() -> usize {
268 10
269}
270
271fn default_keep_first_messages() -> usize {
272 0
273}
274
275impl Default for InfinityContextConfig {
276 fn default() -> Self {
277 Self {
278 context_budget_tokens: default_context_budget_tokens(),
279 min_recent_messages: default_min_recent_messages(),
280 max_recent_messages: None,
281 keep_first_messages: default_keep_first_messages(),
282 compaction_active: false,
283 }
284 }
285}
286
287const CANDIDATE_AVG_TOKENS_PER_MESSAGE: usize = 250;
288const CANDIDATE_OVERFETCH_FACTOR: usize = 4;
289const CANDIDATE_MAX_MESSAGES: usize = 2_000;
290const MAX_KEEP_FIRST_MESSAGES: usize = 16;
291
292struct InfinityContextFilterProvider;
293
294impl MessageFilterProvider for InfinityContextFilterProvider {
295 fn apply_filters(&self, query: &mut MessageQuery, config: &Value) {
296 let config: InfinityContextConfig =
297 serde_json::from_value(config.clone()).unwrap_or_default();
298
299 query.limit = Some(resolve_candidate_load_limit(&config) as i64);
300 let keep_first_messages = resolve_keep_first_messages(&config);
308 if keep_first_messages > 0 {
309 query.keep_head = Some(keep_first_messages);
310 }
311 query.prepend_transform = Some(Arc::new(ExcludedNoticeTransform::infinity_context()));
312 }
313
314 fn post_load(&self, messages: &mut Vec<Message>, config: &Value) {
315 let config: InfinityContextConfig =
316 serde_json::from_value(config.clone()).unwrap_or_default();
317 let existing_notice_count = take_existing_excluded_notice(messages);
318
319 if config.compaction_active {
323 if existing_notice_count > 0 {
324 insert_excluded_notice(
325 messages,
326 resolve_keep_first_messages(&config),
327 existing_notice_count,
328 );
329 }
330 return;
331 }
332
333 let outcome = trim_messages_to_token_budget(messages, &config);
334 let mut total_excluded_count = existing_notice_count.saturating_add(outcome.hidden_count);
335 if total_excluded_count > 0 {
336 let head_ids: std::collections::HashSet<String> = messages
342 .iter()
343 .take(outcome.head_len)
344 .map(|message| message.id.to_string())
345 .collect();
346 let count_before_integrity = messages.len();
347 *messages = everruns_core::retain_complete_message_tool_exchanges(messages, true);
348 total_excluded_count = total_excluded_count
349 .saturating_add(count_before_integrity.saturating_sub(messages.len()));
350 let retained_head_len = messages
351 .iter()
352 .take_while(|message| head_ids.contains(&message.id.to_string()))
353 .count();
354 insert_excluded_notice(messages, retained_head_len, total_excluded_count);
357 }
358 }
359
360 fn priority(&self) -> i32 {
361 100
362 }
363}
364
365fn resolve_keep_first_messages(config: &InfinityContextConfig) -> usize {
366 config.keep_first_messages.min(MAX_KEEP_FIRST_MESSAGES)
369}
370
371fn resolve_candidate_load_limit(config: &InfinityContextConfig) -> usize {
372 let budget_derived_limit = (config.context_budget_tokens / CANDIDATE_AVG_TOKENS_PER_MESSAGE)
376 .saturating_mul(CANDIDATE_OVERFETCH_FACTOR)
377 .max(config.min_recent_messages)
378 .clamp(1, CANDIDATE_MAX_MESSAGES);
379
380 if let Some(max_recent_messages) = config.max_recent_messages {
381 return budget_derived_limit.min(max_recent_messages.max(1));
382 }
383
384 budget_derived_limit
385}
386
387fn estimate_message_tokens(message: &Message) -> usize {
388 const TOKEN_CHARS: usize = 4;
389 let role_overhead = message.role.to_string().len() + 8;
390 let content_len: usize = message
391 .content
392 .iter()
393 .map(|part| match part {
394 ContentPart::Text(text) => text.text.len(),
395 ContentPart::Image(image) => {
396 image.url.as_ref().map_or(0, String::len)
397 + image.base64.as_ref().map_or(50, String::len)
398 + image.media_type.as_ref().map_or(0, String::len)
399 }
400 ContentPart::ImageFile(file) => {
401 file.image_id.to_string().len() + file.filename.as_ref().map_or(0, String::len)
402 }
403 ContentPart::File(file) => {
404 file.file_id.to_string().len() + file.filename.as_ref().map_or(0, String::len)
405 }
406 ContentPart::ToolCall(call) => {
407 call.id.len() + call.name.len() + estimate_json_value_len(&call.arguments) + 20
408 }
409 ContentPart::ToolResult(result) => {
410 result.tool_call_id.len()
411 + result.result.as_ref().map_or(0, estimate_json_value_len)
412 + result.error.as_ref().map_or(0, String::len)
413 + 20
414 }
415 ContentPart::Reasoning(reasoning) => {
418 reasoning.provider.len()
419 + reasoning.item_id.as_ref().map_or(0, String::len)
420 + reasoning.signature.as_ref().map_or(0, String::len)
421 + reasoning.encrypted.as_ref().map_or(0, String::len)
422 + reasoning.display_text().map_or(0, |text| text.len())
423 }
424 _ => 0,
428 })
429 .sum();
430 (role_overhead + content_len) / TOKEN_CHARS
431}
432
433struct CountingWriter {
434 len: usize,
435}
436
437impl Write for CountingWriter {
438 fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
439 self.len = self.len.saturating_add(buf.len());
440 Ok(buf.len())
441 }
442
443 fn flush(&mut self) -> io::Result<()> {
444 Ok(())
445 }
446}
447
448fn estimate_json_value_len(value: &Value) -> usize {
449 let mut writer = CountingWriter { len: 0 };
450 serde_json::to_writer(&mut writer, value)
451 .map(|_| writer.len)
452 .unwrap_or(0)
453}
454
455fn take_existing_excluded_notice(messages: &mut Vec<Message>) -> usize {
456 let Some(first) = messages.first() else {
457 return 0;
458 };
459 let Some(count) = parse_excluded_notice_count(first) else {
460 return 0;
461 };
462
463 messages.remove(0);
464 count
465}
466
467fn parse_excluded_notice_count(message: &Message) -> Option<usize> {
468 let text = message.text()?;
469 let rest = text.strip_prefix("[IMPORTANT: ")?;
470 let (count, rest) = rest.split_once(' ')?;
471 if !rest.starts_with("earlier messages are NOT visible in this context.") {
472 return None;
473 }
474 count.parse().ok()
475}
476
477#[derive(Default)]
479struct TrimOutcome {
480 hidden_count: usize,
482 head_len: usize,
484}
485
486fn insert_excluded_notice(messages: &mut Vec<Message>, position: usize, count: usize) {
488 let text = ExcludedNoticeTransform::infinity_context()
489 .format
490 .replace("{}", &count.to_string());
491 messages.insert(position.min(messages.len()), Message::system(text));
492}
493
494fn trim_messages_to_token_budget(
498 messages: &mut Vec<Message>,
499 config: &InfinityContextConfig,
500) -> TrimOutcome {
501 if messages.is_empty() {
502 return TrimOutcome::default();
503 }
504
505 let costs: Vec<usize> = messages.iter().map(estimate_message_tokens).collect();
506 let window = anchored_window(
507 &costs,
508 resolve_keep_first_messages(config),
509 config.min_recent_messages,
510 config.max_recent_messages,
511 config.context_budget_tokens,
512 );
513
514 let hidden_count = window.hidden();
515 if hidden_count > 0 {
516 let tail = messages.split_off(window.recent_start);
518 messages.truncate(window.head_len);
519 messages.extend(tail);
520 }
521
522 TrimOutcome {
523 hidden_count,
524 head_len: window.head_len,
525 }
526}
527
528pub struct QueryHistoryTool;
530
531#[derive(Debug, Deserialize)]
532struct QueryHistoryParams {
533 #[serde(default)]
534 query: Option<String>,
535 #[serde(default)]
536 message_range: Option<MessageRange>,
537 #[serde(default = "default_query_limit")]
538 limit: usize,
539}
540
541#[derive(Debug, Deserialize)]
542struct MessageRange {
543 from: usize,
544 to: usize,
545}
546
547fn default_query_limit() -> usize {
548 20
549}
550
551#[async_trait]
552impl Tool for QueryHistoryTool {
553 fn narrate(
554 &self,
555 tool_call: &crate::tool_types::ToolCall,
556 phase: crate::tool_narration::ToolNarrationPhase,
557 locale: Option<&str>,
558 _ctx: crate::tool_narration::ToolNarrationContext<'_>,
559 ) -> Option<String> {
560 Some(crate::tool_narration::narrate_query_history(
561 &tool_call.arguments,
562 phase,
563 locale,
564 ))
565 }
566
567 fn name(&self) -> &str {
568 "query_history"
569 }
570
571 fn display_name(&self) -> Option<&str> {
572 Some("Query History")
573 }
574
575 fn description(&self) -> &str {
576 "Search or retrieve earlier messages from this conversation that may not be visible in the current prompt."
577 }
578
579 fn parameters_schema(&self) -> Value {
580 json!({
581 "type": "object",
582 "properties": {
583 "query": {
584 "type": "string",
585 "description": "Keyword search over earlier messages"
586 },
587 "message_range": {
588 "type": "object",
589 "properties": {
590 "from": { "type": "integer", "minimum": 0, "description": "Start index (0-based, inclusive)" },
591 "to": { "type": "integer", "minimum": 0, "description": "End index (0-based, exclusive)" }
592 },
593 "required": ["from", "to"],
594 "additionalProperties": false,
595 "description": "Retrieve messages by absolute position in the conversation"
596 },
597 "limit": {
598 "type": "integer",
599 "minimum": 1,
600 "default": 20,
601 "description": "Maximum number of messages to return"
602 }
603 },
604 "additionalProperties": false
605 })
606 }
607
608 fn hints(&self) -> ToolHints {
609 ToolHints::default()
610 .with_readonly(true)
611 .with_idempotent(true)
612 }
613
614 async fn execute(&self, _arguments: Value) -> ToolExecutionResult {
615 ToolExecutionResult::tool_error(
616 "query_history requires session context. Execute it with ToolContext.",
617 )
618 }
619
620 fn requires_context(&self) -> bool {
621 true
622 }
623
624 fn required_context_services(&self) -> &'static [ToolContextService] {
625 &[ToolContextService::MessageRetriever]
626 }
627
628 async fn execute_with_context(
629 &self,
630 arguments: Value,
631 context: &ToolContext,
632 ) -> ToolExecutionResult {
633 let params: QueryHistoryParams = match serde_json::from_value(arguments) {
634 Ok(params) => params,
635 Err(error) => {
636 return ToolExecutionResult::tool_error(format!("Invalid parameters: {error}"));
637 }
638 };
639
640 let Some(retriever) = &context.message_retriever else {
641 return ToolExecutionResult::tool_error("No message retriever available");
642 };
643
644 let messages = match retriever.load(context.session_id).await {
645 Ok(messages) => messages,
646 Err(error) => {
647 return ToolExecutionResult::internal_error(error);
648 }
649 };
650
651 if messages.is_empty() {
652 return ToolExecutionResult::success(json!({
653 "count": 0,
654 "message": "No history available."
655 }));
656 }
657
658 let limit = params.limit.min(50);
659 let total = messages.len();
660
661 if let Some(range) = params.message_range {
662 let from = range.from.min(total);
663 let to = range.to.min(total).max(from);
664 let range_messages: Vec<_> = messages[from..to].iter().take(limit).collect();
665 return format_range_result(&range_messages, from, total);
666 }
667
668 if let Some(query) = params.query.as_deref() {
669 let results = search_messages(&messages, query, limit);
670 return format_search_result(&results, total);
671 }
672
673 let recent: Vec<_> = messages.iter().rev().take(limit).collect();
674 format_recent_result(&recent, total)
675 }
676}
677
678struct SearchResult<'a> {
679 index: usize,
680 message: &'a Message,
681 score: f64,
682}
683
684fn search_messages<'a>(
685 messages: &'a [Message],
686 query: &str,
687 limit: usize,
688) -> Vec<SearchResult<'a>> {
689 let query_lower = query.to_lowercase();
690 let mut results = Vec::new();
691
692 for (index, message) in messages.iter().enumerate() {
693 let content = extract_text_content(message).to_lowercase();
694 if !content.contains(&query_lower) {
695 continue;
696 }
697
698 let mut score = 1.0;
699
700 if content.split_whitespace().any(|word| word == query_lower) {
701 score += 0.5;
702 }
703
704 if !messages.is_empty() {
705 score += (index as f64 / messages.len() as f64) * 0.3;
706 }
707
708 match message.role {
709 MessageRole::User | MessageRole::Agent => score += 0.2,
710 MessageRole::System => score += 0.1,
711 MessageRole::ToolResult => {}
712 }
713
714 results.push(SearchResult {
715 index,
716 message,
717 score,
718 });
719 }
720
721 results.sort_by(|left, right| {
722 right
723 .score
724 .partial_cmp(&left.score)
725 .unwrap_or(Ordering::Equal)
726 });
727 results.truncate(limit);
728 results
729}
730
731fn extract_text_content(message: &Message) -> String {
732 message
733 .content
734 .iter()
735 .filter_map(|part| match part {
736 ContentPart::Text(text) => Some(text.text.clone()),
737 ContentPart::ToolResult(result) => result.result.as_ref().map(ToString::to_string),
738 _ => None,
739 })
740 .collect::<Vec<_>>()
741 .join(" ")
742}
743
744fn truncate_content(content: &str, max_len: usize) -> String {
745 let char_count = content.chars().count();
746 if char_count <= max_len {
747 return content.to_string();
748 }
749
750 format!("{}...", content.chars().take(max_len).collect::<String>())
751}
752
753fn format_message(message: &Message, index: usize, total: usize) -> Value {
754 json!({
755 "index": index,
756 "position": format!("{}/{}", index + 1, total),
757 "role": message.role.to_string(),
758 "created_at": message.created_at.to_rfc3339(),
759 "content": truncate_content(&extract_text_content(message), 500)
760 })
761}
762
763fn format_range_result(
764 messages: &[&Message],
765 start_index: usize,
766 total: usize,
767) -> ToolExecutionResult {
768 if messages.is_empty() {
769 return ToolExecutionResult::success(json!({
770 "count": 0,
771 "message": "No messages in the requested range."
772 }));
773 }
774
775 let formatted: Vec<Value> = messages
776 .iter()
777 .enumerate()
778 .map(|(offset, message)| format_message(message, start_index + offset, total))
779 .collect();
780
781 ToolExecutionResult::success(json!({
782 "messages": formatted,
783 "count": messages.len(),
784 "total_in_history": total,
785 "range": format!("{}-{}", start_index + 1, start_index + messages.len())
786 }))
787}
788
789fn format_search_result(results: &[SearchResult<'_>], total: usize) -> ToolExecutionResult {
790 if results.is_empty() {
791 return ToolExecutionResult::success(json!({
792 "count": 0,
793 "message": "No matching messages found."
794 }));
795 }
796
797 let formatted: Vec<Value> = results
798 .iter()
799 .map(|result| {
800 let mut message = format_message(result.message, result.index, total);
801 message["relevance_score"] = json!(format!("{:.2}", result.score));
802 message
803 })
804 .collect();
805
806 ToolExecutionResult::success(json!({
807 "messages": formatted,
808 "count": results.len(),
809 "total_in_history": total
810 }))
811}
812
813fn format_recent_result(messages: &[&Message], total: usize) -> ToolExecutionResult {
814 let formatted: Vec<Value> = messages
815 .iter()
816 .enumerate()
817 .map(|(offset, message)| format_message(message, total - messages.len() + offset, total))
818 .collect();
819
820 ToolExecutionResult::success(json!({
821 "messages": formatted,
822 "count": messages.len(),
823 "total_in_history": total,
824 "note": "Showing most recent history. Use `query` to search or `message_range` to fetch older messages."
825 }))
826}
827
828#[cfg(test)]
829mod tests {
830 use super::*;
831 use crate::test_fixtures::TestMessageRetriever;
832 use crate::typed_id::SessionId;
833
834 #[test]
837 fn test_provides_message_filter() {
838 let capability = InfinityContextCapability;
839 assert!(capability.message_filter_provider().is_some());
840 }
841
842 #[test]
843 fn filter_only_view_preserves_filter_without_model_contributions() {
844 let capability = InfinityContextFilterOnlyCapability;
845
846 assert!(capability.message_filter_provider().is_some());
847 assert!(capability.tools().is_empty());
848 assert!(capability.system_prompt_addition().is_none());
849 }
850
851 #[test]
852 fn test_config_schema_and_validate_config() {
853 let capability = InfinityContextCapability;
854
855 let schema = capability.config_schema().expect("config schema");
856 assert_eq!(schema["type"], "object");
857 assert!(schema["properties"]["context_budget_tokens"].is_object());
858 assert!(schema["properties"]["min_recent_messages"].is_object());
859 assert!(schema["properties"]["max_recent_messages"].is_object());
860 assert_eq!(
861 schema["properties"]["keep_first_messages"]["maximum"],
862 MAX_KEEP_FIRST_MESSAGES
863 );
864
865 assert!(capability.validate_config(&Value::Null).is_ok());
867 assert!(capability.validate_config(&json!({})).is_ok());
868 assert!(
869 capability
870 .validate_config(&json!({
871 "context_budget_tokens": 50_000,
872 "min_recent_messages": 5,
873 "max_recent_messages": 100
874 }))
875 .is_ok()
876 );
877
878 assert!(
880 capability
881 .validate_config(&json!({"context_budget_tokens": "lots"}))
882 .is_err()
883 );
884 assert!(
885 capability
886 .validate_config(&json!({"context_budget_tokens": 0}))
887 .is_err()
888 );
889 assert!(
890 capability
891 .validate_config(&json!({"max_recent_messages": 0}))
892 .is_err()
893 );
894 assert!(
895 capability
896 .validate_config(&json!({
897 "keep_first_messages": MAX_KEEP_FIRST_MESSAGES + 1
898 }))
899 .is_err()
900 );
901 }
902
903 #[test]
904 fn test_localizations_resolve_uk() {
905 let capability = InfinityContextCapability;
906 assert_eq!(
907 capability.localized_name(Some("uk-UA")),
908 "Нескінченний контекст"
909 );
910 assert!(capability.describe_schema(None).is_some());
911 }
912
913 #[test]
914 fn test_filter_provider_sets_bounded_candidate_load_limit_without_hard_cap() {
915 let mut query = MessageQuery::new(SessionId::new());
916 let provider = InfinityContextFilterProvider;
917 provider.apply_filters(
918 &mut query,
919 &json!({"context_budget_tokens": 1_000, "min_recent_messages": 3}),
920 );
921
922 assert_eq!(query.limit, Some(16));
923 assert!(query.prepend_transform.is_some());
924 assert_eq!(query.keep_head, None);
927 }
928
929 #[test]
930 fn test_filter_provider_sets_keep_head_from_keep_first_messages() {
931 let mut query = MessageQuery::new(SessionId::new());
932 let provider = InfinityContextFilterProvider;
933 provider.apply_filters(
934 &mut query,
935 &json!({"context_budget_tokens": 1_000, "keep_first_messages": 3}),
936 );
937 assert_eq!(query.keep_head, Some(3));
938 }
939
940 #[test]
941 fn test_filter_provider_omits_keep_head_when_zero() {
942 let mut query = MessageQuery::new(SessionId::new());
943 let provider = InfinityContextFilterProvider;
944 provider.apply_filters(
945 &mut query,
946 &json!({"context_budget_tokens": 1_000, "keep_first_messages": 0}),
947 );
948 assert_eq!(query.keep_head, None);
949 }
950
951 #[test]
952 fn test_filter_provider_caps_keep_head_for_unvalidated_config() {
953 let mut query = MessageQuery::new(SessionId::new());
954 let provider = InfinityContextFilterProvider;
955 provider.apply_filters(
956 &mut query,
957 &json!({
958 "context_budget_tokens": 1_000,
959 "keep_first_messages": usize::MAX
960 }),
961 );
962
963 assert_eq!(query.keep_head, Some(MAX_KEEP_FIRST_MESSAGES));
964 }
965
966 #[test]
967 fn test_filter_provider_caps_explicit_max_to_bounded_candidate_window() {
968 let mut query = MessageQuery::new(SessionId::new());
969 let provider = InfinityContextFilterProvider;
970 provider.apply_filters(
971 &mut query,
972 &json!({
973 "context_budget_tokens": 500_000,
974 "min_recent_messages": 10,
975 "max_recent_messages": 1_000_000
976 }),
977 );
978
979 assert_eq!(query.limit, Some(CANDIDATE_MAX_MESSAGES as i64));
980 assert!(query.prepend_transform.is_some());
981 }
982
983 #[test]
984 fn test_filter_provider_caps_large_min_recent_messages() {
985 let mut query = MessageQuery::new(SessionId::new());
986 let provider = InfinityContextFilterProvider;
987 provider.apply_filters(
988 &mut query,
989 &json!({
990 "context_budget_tokens": 1_000,
991 "min_recent_messages": 1_000_000,
992 }),
993 );
994
995 assert_eq!(query.limit, Some(CANDIDATE_MAX_MESSAGES as i64));
996 assert!(query.prepend_transform.is_some());
997 }
998
999 #[test]
1000 fn test_filter_provider_allows_small_public_chat_window() {
1001 let mut query = MessageQuery::new(SessionId::new());
1002 let provider = InfinityContextFilterProvider;
1003 provider.apply_filters(
1004 &mut query,
1005 &json!({
1006 "context_budget_tokens": 10_000,
1007 "min_recent_messages": 10,
1008 "max_recent_messages": 30
1009 }),
1010 );
1011
1012 assert_eq!(query.limit, Some(30));
1013 assert!(query.prepend_transform.is_some());
1014 }
1015
1016 #[test]
1017 fn test_filter_provider_falls_back_to_defaults_for_invalid_config() {
1018 let mut query = MessageQuery::new(SessionId::new());
1019 let provider = InfinityContextFilterProvider;
1020 provider.apply_filters(
1021 &mut query,
1022 &json!({"context_budget_tokens": "not-a-number"}),
1023 );
1024
1025 assert_eq!(query.limit, Some(1_600));
1026 assert!(query.prepend_transform.is_some());
1027 }
1028
1029 #[test]
1030 fn test_filter_provider_trims_loaded_messages_by_token_budget() {
1031 let provider = InfinityContextFilterProvider;
1032 let mut messages = vec![
1033 Message::user("the original task"),
1034 Message::assistant("old ".repeat(400)),
1035 Message::user("recent one"),
1036 Message::assistant("recent two"),
1037 ];
1038
1039 provider.post_load(
1040 &mut messages,
1041 &json!({"context_budget_tokens": 1, "min_recent_messages": 2}),
1042 );
1043
1044 assert_eq!(messages.len(), 3);
1047 assert!(
1048 extract_text_content(&messages[0])
1049 .contains("2 earlier messages are NOT visible in this context")
1050 );
1051 assert_eq!(extract_text_content(&messages[1]), "recent one");
1052 assert_eq!(extract_text_content(&messages[2]), "recent two");
1053 assert!(
1054 !messages
1055 .iter()
1056 .any(|m| extract_text_content(m) == "the original task")
1057 );
1058 }
1059
1060 #[test]
1061 fn test_filter_provider_applies_hard_cap_after_loading() {
1062 let provider = InfinityContextFilterProvider;
1063 let mut messages = vec![
1064 Message::user("one"),
1065 Message::assistant("two"),
1066 Message::user("three"),
1067 Message::assistant("four"),
1068 Message::user("five"),
1069 ];
1070
1071 provider.post_load(
1072 &mut messages,
1073 &json!({
1074 "context_budget_tokens": 10_000,
1075 "min_recent_messages": 10,
1076 "max_recent_messages": 2
1077 }),
1078 );
1079
1080 assert_eq!(messages.len(), 3);
1082 assert!(
1083 extract_text_content(&messages[0])
1084 .contains("3 earlier messages are NOT visible in this context")
1085 );
1086 assert_eq!(extract_text_content(&messages[1]), "four");
1087 assert_eq!(extract_text_content(&messages[2]), "five");
1088 }
1089
1090 #[test]
1091 fn test_filter_provider_anchors_task_through_full_flow() {
1092 let provider = InfinityContextFilterProvider;
1093 let config = json!({
1096 "context_budget_tokens": 600,
1097 "min_recent_messages": 2,
1098 "keep_first_messages": 1
1099 });
1100 let mut query = MessageQuery::new(SessionId::new());
1101 provider.apply_filters(&mut query, &config);
1102 let mut messages = vec![
1103 Message::user("TASK: build the widget"),
1104 Message::assistant("X".repeat(2000)),
1105 Message::assistant("Y".repeat(2000)),
1106 Message::user("recent a"),
1107 Message::assistant("recent b"),
1108 ];
1109
1110 query.apply_windowing(&mut messages);
1111 provider.post_load(&mut messages, &config);
1112
1113 assert_eq!(extract_text_content(&messages[0]), "TASK: build the widget");
1116 assert!(
1117 extract_text_content(&messages[1])
1118 .contains("earlier messages are NOT visible in this context")
1119 );
1120 assert_eq!(extract_text_content(messages.last().unwrap()), "recent b");
1121 assert!(
1123 !messages
1124 .iter()
1125 .any(|m| extract_text_content(m).starts_with("XXX"))
1126 );
1127 }
1128
1129 #[test]
1130 fn message_filter_hook_coordinates_with_compaction_without_changing_user_config() {
1131 let base = json!({ "context_budget_tokens": 1000 });
1132
1133 let coordinated = message_filter_config(&base, true);
1134 assert_eq!(coordinated["compaction_active"], json!(true));
1135 assert_eq!(coordinated["context_budget_tokens"], json!(1000));
1136 assert!(base.get("compaction_active").is_none());
1137
1138 assert_eq!(message_filter_config(&base, false), base);
1139 assert_eq!(
1140 message_filter_config(&Value::Null, true),
1141 json!({ "compaction_active": true })
1142 );
1143 }
1144
1145 #[test]
1146 fn test_filter_provider_defers_eviction_to_compaction() {
1147 let provider = InfinityContextFilterProvider;
1148 let mut messages = vec![
1149 Message::user("task"),
1150 Message::assistant("old ".repeat(400)),
1151 Message::user("recent one"),
1152 Message::assistant("recent two"),
1153 ];
1154
1155 provider.post_load(
1156 &mut messages,
1157 &json!({
1158 "context_budget_tokens": 1,
1159 "min_recent_messages": 2,
1160 "compaction_active": true
1161 }),
1162 );
1163
1164 assert_eq!(messages.len(), 4);
1167 assert!(
1168 messages
1169 .iter()
1170 .all(|m| !extract_text_content(m).contains("NOT visible"))
1171 );
1172 }
1173
1174 #[test]
1175 fn test_filter_provider_caps_keep_first_messages_during_post_load() {
1176 let provider = InfinityContextFilterProvider;
1177 let mut messages: Vec<Message> = (0..20)
1178 .map(|idx| Message::user(format!("message {idx}")))
1179 .collect();
1180
1181 provider.post_load(
1182 &mut messages,
1183 &json!({
1184 "context_budget_tokens": 1,
1185 "min_recent_messages": 1,
1186 "keep_first_messages": usize::MAX
1187 }),
1188 );
1189
1190 assert_eq!(extract_text_content(&messages[0]), "message 0");
1191 assert_eq!(
1192 extract_text_content(&messages[MAX_KEEP_FIRST_MESSAGES - 1]),
1193 format!("message {}", MAX_KEEP_FIRST_MESSAGES - 1)
1194 );
1195 assert!(extract_text_content(&messages[MAX_KEEP_FIRST_MESSAGES]).contains("NOT visible"));
1196 assert_eq!(extract_text_content(messages.last().unwrap()), "message 19");
1197 }
1198
1199 #[test]
1200 fn test_filter_provider_default_drops_oversized_first_message() {
1201 let provider = InfinityContextFilterProvider;
1202 let mut messages = vec![
1203 Message::user("attacker ".repeat(20_000)),
1204 Message::assistant("middle"),
1205 Message::user("recent one"),
1206 Message::assistant("recent two"),
1207 ];
1208
1209 provider.post_load(
1210 &mut messages,
1211 &json!({
1212 "context_budget_tokens": 10,
1213 "min_recent_messages": 2,
1214 "max_recent_messages": 2
1215 }),
1216 );
1217
1218 assert_eq!(messages.len(), 3);
1219 assert!(
1220 extract_text_content(&messages[0])
1221 .contains("2 earlier messages are NOT visible in this context")
1222 );
1223 assert_eq!(extract_text_content(&messages[1]), "recent one");
1224 assert_eq!(extract_text_content(&messages[2]), "recent two");
1225 assert!(
1226 !messages
1227 .iter()
1228 .any(|m| extract_text_content(m).starts_with("attacker"))
1229 );
1230 }
1231
1232 #[test]
1233 fn test_filter_provider_keep_first_messages_anchors_multiple() {
1234 let provider = InfinityContextFilterProvider;
1235 let mut messages = vec![
1236 Message::user("anchor one"),
1237 Message::user("anchor two"),
1238 Message::assistant("mid ".repeat(400)),
1239 Message::user("recent"),
1240 ];
1241
1242 provider.post_load(
1243 &mut messages,
1244 &json!({
1245 "context_budget_tokens": 1,
1246 "min_recent_messages": 1,
1247 "keep_first_messages": 2
1248 }),
1249 );
1250
1251 assert_eq!(extract_text_content(&messages[0]), "anchor one");
1252 assert_eq!(extract_text_content(&messages[1]), "anchor two");
1253 assert!(extract_text_content(&messages[2]).contains("NOT visible"));
1254 assert_eq!(extract_text_content(messages.last().unwrap()), "recent");
1255 }
1256
1257 #[test]
1258 fn test_estimate_json_value_len_matches_serialized_length() {
1259 let value = json!({
1260 "stdout": ["alpha", "beta"],
1261 "ok": true,
1262 "count": 2
1263 });
1264
1265 assert_eq!(
1266 estimate_json_value_len(&value),
1267 serde_json::to_string(&value).unwrap().len()
1268 );
1269 }
1270
1271 #[test]
1272 fn test_query_history_requires_context() {
1273 let tool = QueryHistoryTool;
1274 assert!(tool.requires_context());
1275 }
1276
1277 #[tokio::test]
1278 async fn test_query_history_tool_errors_without_retriever() {
1279 let tool = QueryHistoryTool;
1280 let result = tool
1281 .execute_with_context(json!({"query": "api"}), &ToolContext::new(SessionId::new()))
1282 .await;
1283
1284 match result {
1285 ToolExecutionResult::ToolError(message) => {
1286 assert!(message.contains("No message retriever available"));
1287 }
1288 other => panic!("expected tool error, got {other:?}"),
1289 }
1290 }
1291
1292 #[tokio::test]
1293 async fn test_query_history_tool_rejects_invalid_params() {
1294 let result = QueryHistoryTool.execute(json!({"limit": "oops"})).await;
1295
1296 match result {
1297 ToolExecutionResult::ToolError(message) => {
1298 assert!(message.contains("requires session context"));
1299 }
1300 other => panic!("expected tool error, got {other:?}"),
1301 }
1302
1303 let session_id = SessionId::new();
1304 let retriever = TestMessageRetriever::new();
1305 let result = QueryHistoryTool
1306 .execute_with_context(
1307 json!({"message_range": {"from": "bad", "to": 1}}),
1308 &ToolContext::new(session_id).with_message_retriever(Arc::new(retriever)),
1309 )
1310 .await;
1311
1312 match result {
1313 ToolExecutionResult::ToolError(message) => {
1314 assert!(message.contains("Invalid parameters"));
1315 }
1316 other => panic!("expected tool error, got {other:?}"),
1317 }
1318 }
1319
1320 #[tokio::test]
1321 async fn test_query_history_tool_empty_history() {
1322 let session_id = SessionId::new();
1323 let retriever = TestMessageRetriever::new();
1324
1325 let result = QueryHistoryTool
1326 .execute_with_context(
1327 json!({}),
1328 &ToolContext::new(session_id).with_message_retriever(Arc::new(retriever)),
1329 )
1330 .await;
1331
1332 match result {
1333 ToolExecutionResult::Success(value) => {
1334 assert_eq!(value["count"], 0);
1335 assert_eq!(value["message"], "No history available.");
1336 }
1337 other => panic!("expected success, got {other:?}"),
1338 }
1339 }
1340
1341 #[tokio::test]
1342 async fn test_query_history_tool_searches_history() {
1343 let session_id = SessionId::new();
1344 let retriever = TestMessageRetriever::new();
1345 retriever
1346 .seed(
1347 session_id,
1348 vec![
1349 Message::user("First topic"),
1350 Message::assistant("The API key is abc123"),
1351 Message::user("We should keep discussing logging"),
1352 ],
1353 )
1354 .await;
1355
1356 let result = QueryHistoryTool
1357 .execute_with_context(
1358 json!({"query": "api key"}),
1359 &ToolContext::new(session_id).with_message_retriever(Arc::new(retriever)),
1360 )
1361 .await;
1362
1363 match result {
1364 ToolExecutionResult::Success(value) => {
1365 assert_eq!(value["count"], 1);
1366 assert_eq!(value["messages"][0]["content"], "The API key is abc123");
1367 }
1368 other => panic!("expected success, got {other:?}"),
1369 }
1370 }
1371
1372 #[tokio::test]
1373 async fn test_query_history_tool_search_no_match() {
1374 let session_id = SessionId::new();
1375 let retriever = TestMessageRetriever::new();
1376 retriever
1377 .seed(
1378 session_id,
1379 vec![Message::user("one"), Message::assistant("two")],
1380 )
1381 .await;
1382
1383 let result = QueryHistoryTool
1384 .execute_with_context(
1385 json!({"query": "missing"}),
1386 &ToolContext::new(session_id).with_message_retriever(Arc::new(retriever)),
1387 )
1388 .await;
1389
1390 match result {
1391 ToolExecutionResult::Success(value) => {
1392 assert_eq!(value["count"], 0);
1393 assert_eq!(value["message"], "No matching messages found.");
1394 }
1395 other => panic!("expected success, got {other:?}"),
1396 }
1397 }
1398
1399 #[tokio::test]
1400 async fn test_query_history_tool_reads_range() {
1401 let session_id = SessionId::new();
1402 let retriever = TestMessageRetriever::new();
1403 retriever
1404 .seed(
1405 session_id,
1406 vec![
1407 Message::user("one"),
1408 Message::assistant("two"),
1409 Message::user("three"),
1410 ],
1411 )
1412 .await;
1413
1414 let result = QueryHistoryTool
1415 .execute_with_context(
1416 json!({"message_range": {"from": 1, "to": 3}, "limit": 10}),
1417 &ToolContext::new(session_id).with_message_retriever(Arc::new(retriever)),
1418 )
1419 .await;
1420
1421 match result {
1422 ToolExecutionResult::Success(value) => {
1423 assert_eq!(value["count"], 2);
1424 assert_eq!(value["messages"][0]["content"], "two");
1425 assert_eq!(value["messages"][1]["content"], "three");
1426 }
1427 other => panic!("expected success, got {other:?}"),
1428 }
1429 }
1430
1431 #[tokio::test]
1432 async fn test_query_history_tool_clamps_out_of_bounds_range() {
1433 let session_id = SessionId::new();
1434 let retriever = TestMessageRetriever::new();
1435 retriever
1436 .seed(
1437 session_id,
1438 vec![
1439 Message::user("one"),
1440 Message::assistant("two"),
1441 Message::user("three"),
1442 ],
1443 )
1444 .await;
1445
1446 let result = QueryHistoryTool
1447 .execute_with_context(
1448 json!({"message_range": {"from": 99, "to": 100}}),
1449 &ToolContext::new(session_id).with_message_retriever(Arc::new(retriever)),
1450 )
1451 .await;
1452
1453 match result {
1454 ToolExecutionResult::Success(value) => {
1455 assert_eq!(value["count"], 0);
1456 assert_eq!(value["message"], "No messages in the requested range.");
1457 }
1458 other => panic!("expected success, got {other:?}"),
1459 }
1460 }
1461
1462 #[test]
1463 fn test_truncate_content_is_utf8_safe() {
1464 let truncated = truncate_content("hello🙂world", 6);
1465 assert_eq!(truncated, "hello🙂...");
1466 }
1467
1468 #[test]
1469 fn trim_preserves_locally_unmatched_tool_result_for_stateful_responses() {
1470 use crate::tool_types::ToolCall;
1471
1472 let provider = InfinityContextFilterProvider;
1473 let mut messages = vec![
1479 Message::user("old question"),
1480 Message::assistant_with_tools(
1481 "calling tool",
1482 vec![ToolCall {
1483 id: "call_old".to_string(),
1484 name: "edit_file".to_string(),
1485 arguments: serde_json::json!({}),
1486 }],
1487 ),
1488 Message::tool_result("call_old", Some(serde_json::json!("done")), None),
1490 Message::user("new question"),
1491 Message::assistant("answer"),
1492 ];
1493
1494 provider.post_load(
1495 &mut messages,
1496 &serde_json::json!({"context_budget_tokens": 1, "min_recent_messages": 3}),
1497 );
1498
1499 assert!(
1500 messages.iter().any(|m| m.role == MessageRole::ToolResult),
1501 "locally unmatched tool result must be preserved until provider serialization"
1502 );
1503
1504 let llm_messages = messages
1505 .iter()
1506 .map(crate::llm_conversions::llm_message_from_message)
1507 .collect();
1508 let stateless_view =
1509 everruns_core::retain_complete_llm_tool_exchanges_for_request(llm_messages, false);
1510 assert!(
1511 stateless_view
1512 .iter()
1513 .all(|message| message.tool_call_id.is_none()),
1514 "a stateless runtime view must drop a result whose call was trimmed"
1515 );
1516 }
1517
1518 #[test]
1519 fn trim_removes_a_visible_call_when_its_result_is_evicted() {
1520 use crate::tool_types::ToolCall;
1521
1522 let provider = InfinityContextFilterProvider;
1523 let mut messages = vec![
1524 Message::user("original task"),
1525 Message::assistant_with_tools(
1526 "calling tool",
1527 vec![ToolCall {
1528 id: "call_old".to_string(),
1529 name: "bash".to_string(),
1530 arguments: serde_json::json!({}),
1531 }],
1532 ),
1533 Message::tool_result(
1534 "call_old",
1535 Some(serde_json::json!("large result ".repeat(500))),
1536 None,
1537 ),
1538 Message::user("recent question"),
1539 Message::assistant("recent answer"),
1540 ];
1541
1542 provider.post_load(
1543 &mut messages,
1544 &serde_json::json!({
1545 "context_budget_tokens": 1,
1546 "min_recent_messages": 2,
1547 "keep_first_messages": 2
1548 }),
1549 );
1550
1551 assert!(
1552 messages
1553 .iter()
1554 .flat_map(Message::tool_calls)
1555 .next()
1556 .is_none(),
1557 "the anchored assistant message must not retain a call after its result is hidden"
1558 );
1559 }
1560
1561 #[test]
1562 fn trim_keeps_tool_result_when_tool_call_is_visible() {
1563 use crate::tool_types::ToolCall;
1564
1565 let provider = InfinityContextFilterProvider;
1566 let mut messages = vec![
1568 Message::assistant_with_tools(
1569 "calling tool",
1570 vec![ToolCall {
1571 id: "call_1".to_string(),
1572 name: "read_file".to_string(),
1573 arguments: serde_json::json!({}),
1574 }],
1575 ),
1576 Message::tool_result("call_1", Some(serde_json::json!("content")), None),
1577 Message::user("thanks"),
1578 ];
1579
1580 provider.post_load(
1581 &mut messages,
1582 &serde_json::json!({"context_budget_tokens": 100_000, "min_recent_messages": 10}),
1583 );
1584
1585 assert!(
1586 messages.iter().any(|m| m.role == MessageRole::ToolResult),
1587 "tool result must be kept when its tool call is visible"
1588 );
1589 }
1590}