1use serde::{Deserialize, Serialize};
4
5use crate::plugins::{Plugin, PluginDescriptor, PluginRegistrar};
6use crate::state::{MutationBatch, StateKey, StateKeyOptions};
7
8pub const CONTEXT_COMPACTION_PLUGIN_ID: &str = "context_compaction";
10
11#[derive(Debug, Clone, Serialize, Deserialize, schemars::JsonSchema)]
20pub struct CompactionConfig {
21 pub summarizer_system_prompt: String,
23 pub summarizer_user_prompt: String,
25 #[serde(default, skip_serializing_if = "Option::is_none")]
27 pub summary_max_tokens: Option<u32>,
28 #[serde(default, skip_serializing_if = "Option::is_none")]
30 pub summary_model: Option<String>,
31 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
47pub struct CompactionConfigKey;
49
50impl awaken_contract::registry_spec::PluginConfigKey for CompactionConfigKey {
51 const KEY: &'static str = "compaction";
52 type Config = CompactionConfig;
53}
54
55#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
61pub struct CompactionBoundary {
62 pub summary: String,
64 pub pre_tokens: usize,
66 pub post_tokens: usize,
68 pub timestamp_ms: u64,
70}
71
72#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
77pub struct CompactionState {
78 pub boundaries: Vec<CompactionBoundary>,
80 pub total_compactions: u64,
82}
83
84#[derive(Debug, Clone, Serialize, Deserialize)]
86#[serde(tag = "type", rename_all = "snake_case")]
87pub enum CompactionAction {
88 RecordBoundary(CompactionBoundary),
90 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 pub fn latest_boundary(&self) -> Option<&CompactionBoundary> {
110 self.boundaries.last()
111 }
112}
113
114pub 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#[derive(Debug, Clone, Default)]
137pub struct CompactionPlugin {
138 pub config: CompactionConfig,
140}
141
142impl CompactionPlugin {
143 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
170pub const CONTEXT_TRANSFORM_PLUGIN_ID: &str = "context_transform";
176
177pub 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 #[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 #[test]
534 fn compaction_fires_at_threshold() {
535 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, 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, 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 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 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 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 let store = StateStore::new();
616 store.install_plugin(CompactionPlugin::default()).unwrap();
617
618 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 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 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 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 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 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 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 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 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 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}