Skip to main content

awaken_runtime/context/
plugin.rs

1//! CompactionPlugin, CompactionConfig, and compaction state tracking.
2
3use serde::{Deserialize, Serialize};
4
5use crate::plugins::{Plugin, PluginDescriptor, PluginRegistrar};
6use crate::state::{MutationBatch, StateKey, StateKeyOptions};
7
8/// Plugin ID for context compaction.
9pub const CONTEXT_COMPACTION_PLUGIN_ID: &str = "context_compaction";
10
11// ---------------------------------------------------------------------------
12// CompactionConfig — configurable prompts and thresholds
13// ---------------------------------------------------------------------------
14
15/// Configuration for the compaction subsystem.
16///
17/// Controls summarizer prompts, model selection, and savings thresholds.
18/// Stored in `AgentSpec.sections["compaction"]` and read via `PluginConfigKey`.
19#[derive(Debug, Clone, Serialize, Deserialize, schemars::JsonSchema)]
20pub struct CompactionConfig {
21    /// System prompt for the summarizer LLM call.
22    pub summarizer_system_prompt: String,
23    /// User prompt template. `{messages}` is replaced with the conversation transcript.
24    pub summarizer_user_prompt: String,
25    /// Maximum tokens for the summary response.
26    #[serde(default, skip_serializing_if = "Option::is_none")]
27    pub summary_max_tokens: Option<u32>,
28    /// Model to use for summarization (if different from the agent's model).
29    #[serde(default, skip_serializing_if = "Option::is_none")]
30    pub summary_model: Option<String>,
31    /// Minimum token savings ratio to accept a compaction (0.0-1.0).
32    pub min_savings_ratio: f64,
33}
34
35impl Default for CompactionConfig {
36    fn default() -> Self {
37        Self {
38            summarizer_system_prompt: "You are a conversation summarizer. Preserve all key facts, decisions, tool results, and action items. Be concise but complete.".into(),
39            summarizer_user_prompt: "Summarize the following conversation:\n\n{messages}".into(),
40            summary_max_tokens: None,
41            summary_model: None,
42            min_savings_ratio: 0.3,
43        }
44    }
45}
46
47/// Plugin config key for [`CompactionConfig`].
48pub struct CompactionConfigKey;
49
50impl awaken_contract::registry_spec::PluginConfigKey for CompactionConfigKey {
51    const KEY: &'static str = "compaction";
52    type Config = CompactionConfig;
53}
54
55// ---------------------------------------------------------------------------
56// Compaction boundary tracking
57// ---------------------------------------------------------------------------
58
59/// A recorded compaction boundary — snapshot of a single compaction event.
60#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
61pub struct CompactionBoundary {
62    /// Summary text produced by the compaction pass.
63    pub summary: String,
64    /// Estimated tokens before compaction (in the compacted range).
65    pub pre_tokens: usize,
66    /// Estimated tokens after compaction (summary message tokens).
67    pub post_tokens: usize,
68    /// Timestamp of the compaction event (millis since UNIX epoch).
69    pub timestamp_ms: u64,
70}
71
72/// Durable state for context compaction tracking.
73///
74/// Stores a history of compaction boundaries so that load-time trimming
75/// and plugin queries can identify already-summarized ranges.
76#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
77pub struct CompactionState {
78    /// Ordered list of compaction boundaries (most recent last).
79    pub boundaries: Vec<CompactionBoundary>,
80    /// Total number of compaction passes performed.
81    pub total_compactions: u64,
82}
83
84/// Reducer actions for [`CompactionState`].
85#[derive(Debug, Clone, Serialize, Deserialize)]
86#[serde(tag = "type", rename_all = "snake_case")]
87pub enum CompactionAction {
88    /// Record a new compaction boundary.
89    RecordBoundary(CompactionBoundary),
90    /// Clear all tracked boundaries (e.g. on thread reset).
91    Clear,
92}
93
94impl CompactionState {
95    fn reduce(&mut self, action: CompactionAction) {
96        match action {
97            CompactionAction::RecordBoundary(boundary) => {
98                self.boundaries.push(boundary);
99                self.total_compactions += 1;
100            }
101            CompactionAction::Clear => {
102                self.boundaries.clear();
103                self.total_compactions = 0;
104            }
105        }
106    }
107
108    /// Latest compaction boundary, if any.
109    pub fn latest_boundary(&self) -> Option<&CompactionBoundary> {
110        self.boundaries.last()
111    }
112}
113
114/// State key for context compaction state.
115pub struct CompactionStateKey;
116
117impl StateKey for CompactionStateKey {
118    const KEY: &'static str = "__context_compaction";
119    type Value = CompactionState;
120    type Update = CompactionAction;
121
122    fn apply(value: &mut Self::Value, update: Self::Update) {
123        value.reduce(update);
124    }
125}
126
127// ---------------------------------------------------------------------------
128// CompactionPlugin
129// ---------------------------------------------------------------------------
130
131/// Plugin that integrates context compaction state into the plugin system.
132///
133/// Registers the [`CompactionStateKey`] state key so that compaction boundaries
134/// are tracked durably and available to other plugins and external observers.
135/// Accepts an optional [`CompactionConfig`] for configurable prompts and thresholds.
136#[derive(Debug, Clone, Default)]
137pub struct CompactionPlugin {
138    /// Compaction configuration (prompts, model, thresholds).
139    pub config: CompactionConfig,
140}
141
142impl CompactionPlugin {
143    /// Create with explicit config.
144    pub fn new(config: CompactionConfig) -> Self {
145        Self { config }
146    }
147}
148
149impl Plugin for CompactionPlugin {
150    fn descriptor(&self) -> PluginDescriptor {
151        PluginDescriptor {
152            name: CONTEXT_COMPACTION_PLUGIN_ID,
153        }
154    }
155
156    fn register(&self, registrar: &mut PluginRegistrar) -> Result<(), awaken_contract::StateError> {
157        registrar.register_key::<CompactionStateKey>(StateKeyOptions::default())?;
158        Ok(())
159    }
160
161    fn on_activate(
162        &self,
163        _agent_spec: &awaken_contract::registry_spec::AgentSpec,
164        _patch: &mut MutationBatch,
165    ) -> Result<(), awaken_contract::StateError> {
166        Ok(())
167    }
168}
169
170// ---------------------------------------------------------------------------
171// ContextTransformPlugin — registers the context truncation request transform
172// ---------------------------------------------------------------------------
173
174/// Plugin ID for context truncation transform.
175pub const CONTEXT_TRANSFORM_PLUGIN_ID: &str = "context_transform";
176
177/// Plugin that registers the built-in context truncation request transform.
178///
179/// Wraps a `ContextWindowPolicy` and registers a `ContextTransform` via
180/// `register_request_transform()` during plugin registration. This ensures
181/// the transform flows through the standard plugin mechanism (ADR-0001)
182/// instead of being manually appended post-hoc.
183pub struct ContextTransformPlugin {
184    policy: awaken_contract::contract::inference::ContextWindowPolicy,
185}
186
187impl ContextTransformPlugin {
188    pub fn new(policy: awaken_contract::contract::inference::ContextWindowPolicy) -> Self {
189        Self { policy }
190    }
191}
192
193impl Plugin for ContextTransformPlugin {
194    fn descriptor(&self) -> PluginDescriptor {
195        PluginDescriptor {
196            name: CONTEXT_TRANSFORM_PLUGIN_ID,
197        }
198    }
199
200    fn register(&self, registrar: &mut PluginRegistrar) -> Result<(), awaken_contract::StateError> {
201        registrar.register_request_transform(
202            CONTEXT_TRANSFORM_PLUGIN_ID,
203            super::ContextTransform::new(self.policy.clone()),
204        );
205        Ok(())
206    }
207}
208
209#[cfg(test)]
210mod tests {
211    use super::*;
212    use crate::state::StateStore;
213    use awaken_contract::contract::message::Message;
214
215    #[test]
216    fn compaction_state_record_boundary() {
217        let mut state = CompactionState::default();
218        assert_eq!(state.total_compactions, 0);
219        assert!(state.boundaries.is_empty());
220
221        state.reduce(CompactionAction::RecordBoundary(CompactionBoundary {
222            summary: "User asked to implement feature X.".into(),
223            pre_tokens: 5000,
224            post_tokens: 200,
225            timestamp_ms: 1234567890,
226        }));
227
228        assert_eq!(state.total_compactions, 1);
229        assert_eq!(state.boundaries.len(), 1);
230        assert_eq!(
231            state.latest_boundary().unwrap().summary,
232            "User asked to implement feature X."
233        );
234    }
235
236    #[test]
237    fn compaction_state_multiple_boundaries() {
238        let mut state = CompactionState::default();
239
240        for i in 0..3 {
241            state.reduce(CompactionAction::RecordBoundary(CompactionBoundary {
242                summary: format!("summary {i}"),
243                pre_tokens: 1000 * (i + 1),
244                post_tokens: 100 * (i + 1),
245                timestamp_ms: 1000 + i as u64,
246            }));
247        }
248
249        assert_eq!(state.total_compactions, 3);
250        assert_eq!(state.boundaries.len(), 3);
251        assert_eq!(state.latest_boundary().unwrap().summary, "summary 2");
252    }
253
254    #[test]
255    fn compaction_state_clear() {
256        let mut state = CompactionState {
257            boundaries: vec![CompactionBoundary {
258                summary: "old".into(),
259                pre_tokens: 100,
260                post_tokens: 10,
261                timestamp_ms: 1,
262            }],
263            total_compactions: 1,
264        };
265
266        state.reduce(CompactionAction::Clear);
267        assert!(state.boundaries.is_empty());
268        assert_eq!(state.total_compactions, 0);
269    }
270
271    #[test]
272    fn compaction_state_latest_boundary_empty() {
273        let state = CompactionState::default();
274        assert!(state.latest_boundary().is_none());
275    }
276
277    #[test]
278    fn compaction_state_serde_roundtrip() {
279        let state = CompactionState {
280            boundaries: vec![
281                CompactionBoundary {
282                    summary: "first".into(),
283                    pre_tokens: 5000,
284                    post_tokens: 200,
285                    timestamp_ms: 1000,
286                },
287                CompactionBoundary {
288                    summary: "second".into(),
289                    pre_tokens: 3000,
290                    post_tokens: 150,
291                    timestamp_ms: 2000,
292                },
293            ],
294            total_compactions: 2,
295        };
296
297        let json = serde_json::to_string(&state).unwrap();
298        let parsed: CompactionState = serde_json::from_str(&json).unwrap();
299        assert_eq!(parsed, state);
300    }
301
302    #[test]
303    fn compaction_plugin_registers_key() {
304        let store = StateStore::new();
305        store.install_plugin(CompactionPlugin::default()).unwrap();
306        let registry = store.registry.lock();
307        assert!(registry.keys_by_name.contains_key("__context_compaction"));
308    }
309
310    #[test]
311    fn compaction_plugin_state_via_store() {
312        let store = StateStore::new();
313        store.install_plugin(CompactionPlugin::default()).unwrap();
314
315        let mut patch = store.begin_mutation();
316        patch.update::<CompactionStateKey>(super::super::record_compaction_boundary(
317            CompactionBoundary {
318                summary: "test summary".into(),
319                pre_tokens: 4000,
320                post_tokens: 180,
321                timestamp_ms: 9999,
322            },
323        ));
324        store.commit(patch).unwrap();
325
326        let state = store.read::<CompactionStateKey>().unwrap();
327        assert_eq!(state.total_compactions, 1);
328        assert_eq!(state.boundaries[0].summary, "test summary");
329    }
330
331    #[test]
332    fn record_compaction_boundary_constructor() {
333        let action = super::super::record_compaction_boundary(CompactionBoundary {
334            summary: "s".into(),
335            pre_tokens: 100,
336            post_tokens: 10,
337            timestamp_ms: 0,
338        });
339        assert!(matches!(action, CompactionAction::RecordBoundary(_)));
340    }
341
342    #[test]
343    fn compaction_state_record_then_clear_then_record() {
344        let mut state = CompactionState::default();
345
346        state.reduce(CompactionAction::RecordBoundary(CompactionBoundary {
347            summary: "first".into(),
348            pre_tokens: 1000,
349            post_tokens: 100,
350            timestamp_ms: 1,
351        }));
352        assert_eq!(state.total_compactions, 1);
353
354        state.reduce(CompactionAction::Clear);
355        assert_eq!(state.total_compactions, 0);
356        assert!(state.boundaries.is_empty());
357        assert!(state.latest_boundary().is_none());
358
359        state.reduce(CompactionAction::RecordBoundary(CompactionBoundary {
360            summary: "after clear".into(),
361            pre_tokens: 2000,
362            post_tokens: 150,
363            timestamp_ms: 2,
364        }));
365        assert_eq!(state.total_compactions, 1);
366        assert_eq!(state.latest_boundary().unwrap().summary, "after clear");
367    }
368
369    #[test]
370    fn compaction_state_key_properties() {
371        assert_eq!(CompactionStateKey::KEY, "__context_compaction");
372    }
373
374    #[test]
375    fn compaction_state_key_apply() {
376        let mut state = CompactionState::default();
377        CompactionStateKey::apply(
378            &mut state,
379            CompactionAction::RecordBoundary(CompactionBoundary {
380                summary: "via apply".into(),
381                pre_tokens: 500,
382                post_tokens: 50,
383                timestamp_ms: 42,
384            }),
385        );
386        assert_eq!(state.total_compactions, 1);
387        assert_eq!(state.boundaries[0].summary, "via apply");
388    }
389
390    #[test]
391    fn compaction_plugin_descriptor_name() {
392        let plugin = CompactionPlugin::default();
393        assert_eq!(plugin.descriptor().name, CONTEXT_COMPACTION_PLUGIN_ID);
394    }
395
396    #[test]
397    fn compaction_plugin_new_with_config() {
398        let config = CompactionConfig {
399            min_savings_ratio: 0.8,
400            ..Default::default()
401        };
402        let plugin = CompactionPlugin::new(config);
403        assert!((plugin.config.min_savings_ratio - 0.8).abs() < f64::EPSILON);
404    }
405
406    #[test]
407    fn compaction_boundary_equality() {
408        let a = CompactionBoundary {
409            summary: "s".into(),
410            pre_tokens: 100,
411            post_tokens: 10,
412            timestamp_ms: 0,
413        };
414        let b = a.clone();
415        assert_eq!(a, b);
416    }
417
418    #[test]
419    fn compaction_boundary_serde_roundtrip() {
420        let boundary = CompactionBoundary {
421            summary: "test summary".into(),
422            pre_tokens: 3000,
423            post_tokens: 200,
424            timestamp_ms: 1234567890,
425        };
426        let json = serde_json::to_string(&boundary).unwrap();
427        let parsed: CompactionBoundary = serde_json::from_str(&json).unwrap();
428        assert_eq!(parsed, boundary);
429    }
430
431    // -----------------------------------------------------------------------
432    // Migrated from uncarve: additional compaction state tests
433    // -----------------------------------------------------------------------
434
435    #[test]
436    fn compaction_state_default_is_empty() {
437        let state = CompactionState::default();
438        assert!(state.boundaries.is_empty());
439        assert_eq!(state.total_compactions, 0);
440        assert!(state.latest_boundary().is_none());
441    }
442
443    #[test]
444    fn compaction_state_boundary_ordering_preserved() {
445        let mut state = CompactionState::default();
446        for i in 0..5 {
447            state.reduce(CompactionAction::RecordBoundary(CompactionBoundary {
448                summary: format!("boundary_{i}"),
449                pre_tokens: 1000,
450                post_tokens: 100,
451                timestamp_ms: i as u64,
452            }));
453        }
454        assert_eq!(state.boundaries.len(), 5);
455        assert_eq!(state.total_compactions, 5);
456        for (i, b) in state.boundaries.iter().enumerate() {
457            assert_eq!(b.summary, format!("boundary_{i}"));
458            assert_eq!(b.timestamp_ms, i as u64);
459        }
460    }
461
462    #[test]
463    fn compaction_state_clear_twice_is_idempotent() {
464        let mut state = CompactionState::default();
465        state.reduce(CompactionAction::RecordBoundary(CompactionBoundary {
466            summary: "s".into(),
467            pre_tokens: 1,
468            post_tokens: 1,
469            timestamp_ms: 0,
470        }));
471        state.reduce(CompactionAction::Clear);
472        state.reduce(CompactionAction::Clear);
473        assert!(state.boundaries.is_empty());
474        assert_eq!(state.total_compactions, 0);
475    }
476
477    #[test]
478    fn compaction_config_default_has_sane_values() {
479        let config = CompactionConfig::default();
480        assert!(!config.summarizer_system_prompt.is_empty());
481        assert!(config.summarizer_user_prompt.contains("{messages}"));
482        assert!(config.min_savings_ratio > 0.0);
483        assert!(config.min_savings_ratio < 1.0);
484        assert!(config.summary_max_tokens.is_none());
485        assert!(config.summary_model.is_none());
486    }
487
488    #[test]
489    fn compaction_config_serde_roundtrip() {
490        let config = CompactionConfig {
491            summarizer_system_prompt: "custom system".into(),
492            summarizer_user_prompt: "custom user: {messages}".into(),
493            summary_max_tokens: Some(512),
494            summary_model: Some("claude-3-haiku".into()),
495            min_savings_ratio: 0.5,
496        };
497        let json = serde_json::to_string(&config).unwrap();
498        let parsed: CompactionConfig = serde_json::from_str(&json).unwrap();
499        assert_eq!(
500            parsed.summarizer_system_prompt,
501            config.summarizer_system_prompt
502        );
503        assert_eq!(parsed.summary_max_tokens, Some(512));
504        assert_eq!(parsed.summary_model.as_deref(), Some("claude-3-haiku"));
505    }
506
507    #[test]
508    fn compaction_state_pre_post_tokens_preserved() {
509        let mut state = CompactionState::default();
510        state.reduce(CompactionAction::RecordBoundary(CompactionBoundary {
511            summary: "test".into(),
512            pre_tokens: 10_000,
513            post_tokens: 500,
514            timestamp_ms: 99,
515        }));
516        let b = state.latest_boundary().unwrap();
517        assert_eq!(b.pre_tokens, 10_000);
518        assert_eq!(b.post_tokens, 500);
519        assert_eq!(b.timestamp_ms, 99);
520    }
521
522    #[test]
523    fn context_transform_plugin_descriptor_name() {
524        let policy = awaken_contract::contract::inference::ContextWindowPolicy::default();
525        let plugin = ContextTransformPlugin::new(policy);
526        assert_eq!(plugin.descriptor().name, CONTEXT_TRANSFORM_PLUGIN_ID);
527    }
528
529    // -----------------------------------------------------------------------
530    // Additional compaction tests
531    // -----------------------------------------------------------------------
532
533    #[test]
534    fn compaction_fires_at_threshold() {
535        // Verify savings ratio check: only accept compaction when savings >= min_savings_ratio
536        let config = CompactionConfig {
537            min_savings_ratio: 0.5,
538            ..Default::default()
539        };
540        let boundary_good = CompactionBoundary {
541            summary: "good".into(),
542            pre_tokens: 1000,
543            post_tokens: 400, // 60% savings > 50% threshold
544            timestamp_ms: 1,
545        };
546        let savings_good =
547            1.0 - (boundary_good.post_tokens as f64 / boundary_good.pre_tokens as f64);
548        assert!(
549            savings_good >= config.min_savings_ratio,
550            "60% savings should meet 50% threshold"
551        );
552
553        let boundary_bad = CompactionBoundary {
554            summary: "bad".into(),
555            pre_tokens: 1000,
556            post_tokens: 600, // 40% savings < 50% threshold
557            timestamp_ms: 2,
558        };
559        let savings_bad = 1.0 - (boundary_bad.post_tokens as f64 / boundary_bad.pre_tokens as f64);
560        assert!(
561            savings_bad < config.min_savings_ratio,
562            "40% savings should not meet 50% threshold"
563        );
564    }
565
566    #[test]
567    fn compaction_state_tracks_across_multiple_rounds() {
568        let mut state = CompactionState::default();
569        // Simulate 5 compaction rounds with increasing pre-token counts
570        for round in 1..=5u64 {
571            state.reduce(CompactionAction::RecordBoundary(CompactionBoundary {
572                summary: format!("round {round}"),
573                pre_tokens: 1000 * round as usize,
574                post_tokens: 100 * round as usize,
575                timestamp_ms: round * 1000,
576            }));
577            assert_eq!(state.total_compactions, round);
578            assert_eq!(state.boundaries.len(), round as usize);
579        }
580        // Latest boundary should be the last round
581        assert_eq!(state.latest_boundary().unwrap().summary, "round 5");
582        assert_eq!(state.latest_boundary().unwrap().pre_tokens, 5000);
583    }
584
585    #[test]
586    fn compaction_config_serialization_omits_none_fields() {
587        let config = CompactionConfig::default();
588        let json = serde_json::to_value(&config).unwrap();
589        // summary_max_tokens and summary_model are None, should be omitted via skip_serializing_if
590        assert!(
591            !json.as_object().unwrap().contains_key("summary_max_tokens"),
592            "None fields should be omitted"
593        );
594        assert!(
595            !json.as_object().unwrap().contains_key("summary_model"),
596            "None fields should be omitted"
597        );
598    }
599
600    #[test]
601    fn compaction_config_serialization_includes_some_fields() {
602        let config = CompactionConfig {
603            summary_max_tokens: Some(1024),
604            summary_model: Some("claude-3-sonnet".into()),
605            ..Default::default()
606        };
607        let json = serde_json::to_value(&config).unwrap();
608        assert_eq!(json["summary_max_tokens"], 1024);
609        assert_eq!(json["summary_model"], "claude-3-sonnet");
610    }
611
612    #[test]
613    fn compaction_with_tool_messages_records_correctly() {
614        // Simulate compaction that includes tool messages in the summarized range
615        let store = StateStore::new();
616        store.install_plugin(CompactionPlugin::default()).unwrap();
617
618        // Record a boundary representing a range that included tool messages
619        let mut patch = store.begin_mutation();
620        patch.update::<CompactionStateKey>(super::super::record_compaction_boundary(
621            CompactionBoundary {
622                summary: "User asked to search files. Tool search returned 3 results. Assistant presented findings.".into(),
623                pre_tokens: 8000,
624                post_tokens: 200,
625                timestamp_ms: 1000,
626            },
627        ));
628        store.commit(patch).unwrap();
629
630        let state = store.read::<CompactionStateKey>().unwrap();
631        assert_eq!(state.total_compactions, 1);
632        assert!(state.boundaries[0].summary.contains("Tool search"));
633        assert_eq!(state.boundaries[0].pre_tokens, 8000);
634    }
635
636    #[test]
637    fn context_transform_plugin_registers_transform() {
638        use crate::plugins::PluginRegistrar;
639        let policy = awaken_contract::contract::inference::ContextWindowPolicy::default();
640        let plugin = ContextTransformPlugin::new(policy);
641        let mut registrar = PluginRegistrar::new();
642        plugin.register(&mut registrar).unwrap();
643        assert_eq!(
644            registrar.request_transforms.len(),
645            1,
646            "should have registered one transform"
647        );
648        assert_eq!(
649            registrar.request_transforms[0].plugin_id,
650            CONTEXT_TRANSFORM_PLUGIN_ID
651        );
652    }
653
654    #[test]
655    fn transform_ordering_compaction_then_context() {
656        use crate::plugins::PluginRegistrar;
657        // Compaction plugin should register no transforms
658        let mut reg_compaction = PluginRegistrar::new();
659        CompactionPlugin::default()
660            .register(&mut reg_compaction)
661            .unwrap();
662        assert!(
663            reg_compaction.request_transforms.is_empty(),
664            "CompactionPlugin should not register request transforms"
665        );
666        // ContextTransformPlugin should register exactly one transform
667        let policy = awaken_contract::contract::inference::ContextWindowPolicy::default();
668        let mut reg_transform = PluginRegistrar::new();
669        ContextTransformPlugin::new(policy)
670            .register(&mut reg_transform)
671            .unwrap();
672        assert_eq!(reg_transform.request_transforms.len(), 1);
673    }
674
675    #[test]
676    fn token_count_estimation_for_various_content_types() {
677        use awaken_contract::contract::transform::estimate_message_tokens;
678
679        // Text message
680        let text_msg = Message::user("Hello, this is a test message with some content.");
681        let text_tokens = estimate_message_tokens(&text_msg);
682        assert!(
683            text_tokens > 4,
684            "text message should have tokens beyond overhead"
685        );
686
687        // Empty content message
688        let empty_msg = Message::user("");
689        let empty_tokens = estimate_message_tokens(&empty_msg);
690        assert_eq!(
691            empty_tokens, 4,
692            "empty message should have only overhead tokens"
693        );
694
695        // Very long message
696        let long_msg = Message::user("x".repeat(4000));
697        let long_tokens = estimate_message_tokens(&long_msg);
698        assert!(
699            long_tokens >= 1000,
700            "4000-char message should estimate >= 1000 tokens, got {long_tokens}"
701        );
702    }
703
704    #[test]
705    fn enable_prompt_cache_flag_in_policy() {
706        let policy_cached = awaken_contract::contract::inference::ContextWindowPolicy {
707            enable_prompt_cache: true,
708            ..Default::default()
709        };
710        assert!(policy_cached.enable_prompt_cache);
711
712        let policy_uncached = awaken_contract::contract::inference::ContextWindowPolicy {
713            enable_prompt_cache: false,
714            ..Default::default()
715        };
716        assert!(!policy_uncached.enable_prompt_cache);
717
718        // Both should create valid transform plugins
719        let _ = ContextTransformPlugin::new(policy_cached);
720        let _ = ContextTransformPlugin::new(policy_uncached);
721    }
722
723    #[test]
724    fn autocompact_threshold_check() {
725        use awaken_contract::contract::transform::estimate_tokens;
726
727        let policy_with_threshold = awaken_contract::contract::inference::ContextWindowPolicy {
728            autocompact_threshold: Some(500),
729            ..Default::default()
730        };
731
732        // Simulate checking if messages exceed autocompact threshold
733        let messages = vec![Message::user("short"), Message::assistant("reply")];
734        let total = estimate_tokens(&messages);
735        assert!(
736            total < policy_with_threshold.autocompact_threshold.unwrap(),
737            "short conversation should be under threshold"
738        );
739
740        // Longer conversation should exceed threshold
741        let long_messages: Vec<Message> = (0..100)
742            .map(|i| Message::user(format!("message {i} with some filler text to add tokens")))
743            .collect();
744        let long_total = estimate_tokens(&long_messages);
745        assert!(
746            long_total > policy_with_threshold.autocompact_threshold.unwrap(),
747            "100-message conversation should exceed threshold of 500, got {long_total}"
748        );
749    }
750
751    #[test]
752    fn compaction_action_serde_roundtrip() {
753        let actions = vec![
754            CompactionAction::RecordBoundary(CompactionBoundary {
755                summary: "s".into(),
756                pre_tokens: 1,
757                post_tokens: 1,
758                timestamp_ms: 0,
759            }),
760            CompactionAction::Clear,
761        ];
762        for action in actions {
763            let json = serde_json::to_string(&action).unwrap();
764            let parsed: CompactionAction = serde_json::from_str(&json).unwrap();
765            // Verify the action type roundtrips
766            match (&action, &parsed) {
767                (CompactionAction::Clear, CompactionAction::Clear) => {}
768                (CompactionAction::RecordBoundary(a), CompactionAction::RecordBoundary(b)) => {
769                    assert_eq!(a.summary, b.summary);
770                }
771                _ => panic!("action type mismatch after serde roundtrip"),
772            }
773        }
774    }
775}