ucp_llm/
context.rs

1//! Context management infrastructure for UCM documents.
2//!
3//! This module provides APIs for intelligent context window management,
4//! allowing external orchestration layers to load documents, traverse
5//! the knowledge graph, and curate context windows while preserving UCM invariants.
6
7use serde::{Deserialize, Serialize};
8use std::collections::{HashMap, HashSet, VecDeque};
9use ucm_core::{BlockId, Content, Document};
10
11#[cfg(test)]
12use ucm_core::Block;
13
14/// Reason why a block was included in context
15#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
16pub enum InclusionReason {
17    /// Directly referenced by task or query
18    DirectReference,
19    /// Part of navigation path
20    NavigationPath,
21    /// Structural context (parent/sibling)
22    StructuralContext,
23    /// Semantic relevance
24    SemanticRelevance,
25    /// External decision (from orchestrator)
26    ExternalDecision,
27    /// Required for understanding
28    RequiredContext,
29}
30
31/// A block in the context window with metadata
32#[derive(Debug, Clone, Serialize, Deserialize)]
33pub struct ContextBlock {
34    pub block_id: BlockId,
35    pub inclusion_reason: InclusionReason,
36    pub relevance_score: f32,
37    pub token_estimate: usize,
38    pub access_count: usize,
39    pub last_accessed: chrono::DateTime<chrono::Utc>,
40    pub compressed: bool,
41    pub original_content: Option<String>,
42}
43
44/// Relationship between context blocks
45#[derive(Debug, Clone, Serialize, Deserialize)]
46pub struct ContextRelation {
47    pub source: BlockId,
48    pub target: BlockId,
49    pub relation_type: String,
50}
51
52/// Context window metadata
53#[derive(Debug, Clone, Default, Serialize, Deserialize)]
54pub struct ContextMetadata {
55    pub focus_area: Option<BlockId>,
56    pub task_description: Option<String>,
57    pub created_at: Option<chrono::DateTime<chrono::Utc>>,
58    pub last_modified: Option<chrono::DateTime<chrono::Utc>>,
59}
60
61/// Constraints for the context window
62#[derive(Debug, Clone, Serialize, Deserialize)]
63pub struct ContextConstraints {
64    pub max_tokens: usize,
65    pub max_blocks: usize,
66    pub max_depth: usize,
67    pub min_relevance: f32,
68    pub required_roles: Vec<String>,
69    pub excluded_tags: Vec<String>,
70    pub preserve_structure: bool,
71    pub allow_compression: bool,
72}
73
74impl Default for ContextConstraints {
75    fn default() -> Self {
76        Self {
77            max_tokens: 4000,
78            max_blocks: 100,
79            max_depth: 10,
80            min_relevance: 0.0,
81            required_roles: Vec::new(),
82            excluded_tags: Vec::new(),
83            preserve_structure: true,
84            allow_compression: true,
85        }
86    }
87}
88
89/// Direction for context expansion
90#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
91pub enum ExpandDirection {
92    Up,
93    Down,
94    Both,
95    Semantic,
96}
97
98/// Policy for context expansion
99#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
100pub enum ExpansionPolicy {
101    /// Only add highly relevant blocks
102    Conservative,
103    /// Balance relevance and diversity
104    #[default]
105    Balanced,
106    /// Add potentially useful blocks
107    Aggressive,
108}
109
110/// Policy for context pruning
111#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
112pub enum PruningPolicy {
113    /// Remove lowest relevance first
114    #[default]
115    RelevanceFirst,
116    /// Remove least recently accessed
117    RecencyFirst,
118    /// Remove redundant content
119    RedundancyFirst,
120}
121
122/// Method for content compression
123#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
124pub enum CompressionMethod {
125    /// Truncate to length limit
126    #[default]
127    Truncate,
128    /// Summarize content (requires external summarizer)
129    Summarize,
130    /// Show only structure, not content
131    StructureOnly,
132}
133
134/// Result of a context operation
135#[derive(Debug, Clone, Default, Serialize, Deserialize)]
136pub struct ContextUpdateResult {
137    pub blocks_added: Vec<BlockId>,
138    pub blocks_removed: Vec<BlockId>,
139    pub blocks_compressed: Vec<BlockId>,
140    pub total_tokens: usize,
141    pub total_blocks: usize,
142    pub warnings: Vec<String>,
143}
144
145/// Statistics about the context window
146#[derive(Debug, Clone, Default, Serialize, Deserialize)]
147pub struct ContextStatistics {
148    pub total_tokens: usize,
149    pub total_blocks: usize,
150    pub blocks_by_reason: HashMap<String, usize>,
151    pub average_relevance: f32,
152    pub depth_distribution: HashMap<usize, usize>,
153    pub compressed_count: usize,
154}
155
156/// Context window with intelligent management
157#[derive(Debug, Clone, Serialize, Deserialize)]
158pub struct ContextWindow {
159    pub id: String,
160    pub blocks: HashMap<BlockId, ContextBlock>,
161    pub relationships: Vec<ContextRelation>,
162    pub metadata: ContextMetadata,
163    pub constraints: ContextConstraints,
164}
165
166impl ContextWindow {
167    /// Create a new empty context window
168    pub fn new(id: impl Into<String>, constraints: ContextConstraints) -> Self {
169        Self {
170            id: id.into(),
171            blocks: HashMap::new(),
172            relationships: Vec::new(),
173            metadata: ContextMetadata {
174                created_at: Some(chrono::Utc::now()),
175                ..Default::default()
176            },
177            constraints,
178        }
179    }
180
181    /// Get the number of blocks in the context
182    pub fn block_count(&self) -> usize {
183        self.blocks.len()
184    }
185
186    /// Get estimated total tokens
187    pub fn total_tokens(&self) -> usize {
188        self.blocks.values().map(|b| b.token_estimate).sum()
189    }
190
191    /// Check if context has room for more blocks
192    pub fn has_capacity(&self) -> bool {
193        self.blocks.len() < self.constraints.max_blocks
194            && self.total_tokens() < self.constraints.max_tokens
195    }
196
197    /// Check if a block is in the context
198    pub fn contains(&self, block_id: &BlockId) -> bool {
199        self.blocks.contains_key(block_id)
200    }
201
202    /// Get a block from the context
203    pub fn get(&self, block_id: &BlockId) -> Option<&ContextBlock> {
204        self.blocks.get(block_id)
205    }
206
207    /// Get all block IDs in the context
208    pub fn block_ids(&self) -> Vec<BlockId> {
209        self.blocks.keys().copied().collect()
210    }
211}
212
213/// Context Management Infrastructure
214///
215/// Provides APIs for external orchestration layers to manage context windows.
216pub struct ContextManager {
217    window: ContextWindow,
218    expansion_policy: ExpansionPolicy,
219    pruning_policy: PruningPolicy,
220}
221
222impl ContextManager {
223    /// Create a new context manager with default constraints
224    pub fn new(id: impl Into<String>) -> Self {
225        Self {
226            window: ContextWindow::new(id, ContextConstraints::default()),
227            expansion_policy: ExpansionPolicy::default(),
228            pruning_policy: PruningPolicy::default(),
229        }
230    }
231
232    /// Create a context manager with custom constraints
233    pub fn with_constraints(id: impl Into<String>, constraints: ContextConstraints) -> Self {
234        Self {
235            window: ContextWindow::new(id, constraints),
236            expansion_policy: ExpansionPolicy::default(),
237            pruning_policy: PruningPolicy::default(),
238        }
239    }
240
241    /// Set the expansion policy
242    pub fn with_expansion_policy(mut self, policy: ExpansionPolicy) -> Self {
243        self.expansion_policy = policy;
244        self
245    }
246
247    /// Set the pruning policy
248    pub fn with_pruning_policy(mut self, policy: PruningPolicy) -> Self {
249        self.pruning_policy = policy;
250        self
251    }
252
253    /// Get a reference to the context window
254    pub fn window(&self) -> &ContextWindow {
255        &self.window
256    }
257
258    /// Initialize context with a focus block
259    pub fn initialize_focus(
260        &mut self,
261        doc: &Document,
262        focus_id: BlockId,
263        task_description: &str,
264    ) -> ContextUpdateResult {
265        self.window.metadata.focus_area = Some(focus_id);
266        self.window.metadata.task_description = Some(task_description.to_string());
267        self.window.metadata.last_modified = Some(chrono::Utc::now());
268
269        let mut result = ContextUpdateResult::default();
270
271        // Add focus block
272        if let Some(_block) = doc.get_block(&focus_id) {
273            self.add_block_internal(doc, focus_id, InclusionReason::DirectReference, 1.0);
274            result.blocks_added.push(focus_id);
275        }
276
277        // Add structural context (ancestors)
278        let mut current = focus_id;
279        let mut depth = 0;
280        while let Some(parent) = doc.parent(&current) {
281            if *parent == doc.root || depth >= 3 {
282                break;
283            }
284            self.add_block_internal(
285                doc,
286                *parent,
287                InclusionReason::StructuralContext,
288                0.8 - depth as f32 * 0.1,
289            );
290            result.blocks_added.push(*parent);
291            current = *parent;
292            depth += 1;
293        }
294
295        result.total_tokens = self.window.total_tokens();
296        result.total_blocks = self.window.block_count();
297        result
298    }
299
300    /// Navigate to a new focus area
301    pub fn navigate_to(
302        &mut self,
303        doc: &Document,
304        target_id: BlockId,
305        task_description: &str,
306    ) -> ContextUpdateResult {
307        self.window.metadata.focus_area = Some(target_id);
308        self.window.metadata.task_description = Some(task_description.to_string());
309        self.window.metadata.last_modified = Some(chrono::Utc::now());
310
311        let mut result = ContextUpdateResult::default();
312
313        // Add target block
314        if doc.get_block(&target_id).is_some() {
315            self.add_block_internal(doc, target_id, InclusionReason::NavigationPath, 1.0);
316            result.blocks_added.push(target_id);
317        }
318
319        // Prune if needed
320        let pruned = self.prune_if_needed();
321        result.blocks_removed = pruned;
322
323        result.total_tokens = self.window.total_tokens();
324        result.total_blocks = self.window.block_count();
325        result
326    }
327
328    /// Add a block to the context
329    pub fn add_block(
330        &mut self,
331        doc: &Document,
332        block_id: BlockId,
333        reason: InclusionReason,
334    ) -> ContextUpdateResult {
335        let mut result = ContextUpdateResult::default();
336
337        if doc.get_block(&block_id).is_some() {
338            self.add_block_internal(doc, block_id, reason, 0.7);
339            result.blocks_added.push(block_id);
340        }
341
342        // Prune if needed
343        let pruned = self.prune_if_needed();
344        result.blocks_removed = pruned;
345
346        result.total_tokens = self.window.total_tokens();
347        result.total_blocks = self.window.block_count();
348        result
349    }
350
351    /// Remove a block from the context
352    pub fn remove_block(&mut self, block_id: BlockId) -> ContextUpdateResult {
353        let mut result = ContextUpdateResult::default();
354
355        if self.window.blocks.remove(&block_id).is_some() {
356            result.blocks_removed.push(block_id);
357        }
358
359        self.window.metadata.last_modified = Some(chrono::Utc::now());
360        result.total_tokens = self.window.total_tokens();
361        result.total_blocks = self.window.block_count();
362        result
363    }
364
365    /// Expand context in a direction
366    pub fn expand_context(
367        &mut self,
368        doc: &Document,
369        direction: ExpandDirection,
370        depth: usize,
371    ) -> ContextUpdateResult {
372        let mut result = ContextUpdateResult::default();
373
374        let focus_id = match self.window.metadata.focus_area {
375            Some(id) => id,
376            None => return result,
377        };
378
379        match direction {
380            ExpandDirection::Down => {
381                let added = self.expand_downward(doc, focus_id, depth);
382                result.blocks_added = added;
383            }
384            ExpandDirection::Up => {
385                let added = self.expand_upward(doc, focus_id, depth);
386                result.blocks_added = added;
387            }
388            ExpandDirection::Both => {
389                let added_down = self.expand_downward(doc, focus_id, depth);
390                let added_up = self.expand_upward(doc, focus_id, depth);
391                result.blocks_added = added_down.into_iter().chain(added_up).collect();
392            }
393            ExpandDirection::Semantic => {
394                // Semantic expansion would use edges
395                let added = self.expand_semantic(doc, focus_id, depth);
396                result.blocks_added = added;
397            }
398        }
399
400        // Prune if needed
401        let pruned = self.prune_if_needed();
402        result.blocks_removed = pruned;
403
404        self.window.metadata.last_modified = Some(chrono::Utc::now());
405        result.total_tokens = self.window.total_tokens();
406        result.total_blocks = self.window.block_count();
407        result
408    }
409
410    /// Compress blocks to fit within constraints
411    pub fn compress(&mut self, doc: &Document, method: CompressionMethod) -> ContextUpdateResult {
412        let mut result = ContextUpdateResult::default();
413
414        if !self.window.constraints.allow_compression {
415            result
416                .warnings
417                .push("Compression not allowed by constraints".to_string());
418            return result;
419        }
420
421        // Find blocks to compress (lowest relevance, not already compressed)
422        let mut blocks_to_compress: Vec<(BlockId, f32)> = self
423            .window
424            .blocks
425            .iter()
426            .filter(|(_, cb)| !cb.compressed)
427            .map(|(id, cb)| (*id, cb.relevance_score))
428            .collect();
429
430        blocks_to_compress
431            .sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
432
433        for (block_id, _) in blocks_to_compress.iter().take(10) {
434            // Extract content text before mutable borrow
435            let original = doc
436                .get_block(block_id)
437                .map(|block| self.extract_content_text(&block.content));
438
439            if let Some(context_block) = self.window.blocks.get_mut(block_id) {
440                if let Some(original_text) = original {
441                    context_block.original_content = Some(original_text);
442
443                    match method {
444                        CompressionMethod::Truncate => {
445                            // Reduce token estimate
446                            context_block.token_estimate /= 2;
447                        }
448                        CompressionMethod::StructureOnly => {
449                            // Minimal token estimate
450                            context_block.token_estimate = 10;
451                        }
452                        CompressionMethod::Summarize => {
453                            // Would need external summarizer
454                            context_block.token_estimate /= 3;
455                        }
456                    }
457
458                    context_block.compressed = true;
459                    result.blocks_compressed.push(*block_id);
460                }
461            }
462
463            // Check if we're within constraints
464            if self.window.total_tokens() <= self.window.constraints.max_tokens {
465                break;
466            }
467        }
468
469        result.total_tokens = self.window.total_tokens();
470        result.total_blocks = self.window.block_count();
471        result
472    }
473
474    /// Get statistics about the context
475    pub fn get_statistics(&self) -> ContextStatistics {
476        let mut blocks_by_reason: HashMap<String, usize> = HashMap::new();
477        let mut total_relevance = 0.0;
478        let mut compressed_count = 0;
479
480        for cb in self.window.blocks.values() {
481            let reason = format!("{:?}", cb.inclusion_reason);
482            *blocks_by_reason.entry(reason).or_insert(0) += 1;
483            total_relevance += cb.relevance_score;
484            if cb.compressed {
485                compressed_count += 1;
486            }
487        }
488
489        let average_relevance = if self.window.blocks.is_empty() {
490            0.0
491        } else {
492            total_relevance / self.window.blocks.len() as f32
493        };
494
495        ContextStatistics {
496            total_tokens: self.window.total_tokens(),
497            total_blocks: self.window.block_count(),
498            blocks_by_reason,
499            average_relevance,
500            depth_distribution: HashMap::new(), // Would need document access
501            compressed_count,
502        }
503    }
504
505    /// Render context to a format suitable for LLM prompts
506    pub fn render_for_prompt(&self, doc: &Document) -> String {
507        let mut output = String::new();
508
509        // Sort blocks by relevance for output
510        let mut blocks: Vec<(&BlockId, &ContextBlock)> = self.window.blocks.iter().collect();
511        blocks.sort_by(|a, b| {
512            b.1.relevance_score
513                .partial_cmp(&a.1.relevance_score)
514                .unwrap_or(std::cmp::Ordering::Equal)
515        });
516
517        for (block_id, context_block) in blocks {
518            if let Some(block) = doc.get_block(block_id) {
519                let content = if context_block.compressed {
520                    if let Some(ref original) = context_block.original_content {
521                        format!("[compressed] {}...", &original[..original.len().min(50)])
522                    } else {
523                        "[compressed]".to_string()
524                    }
525                } else {
526                    self.extract_content_text(&block.content)
527                };
528
529                let role = block
530                    .metadata
531                    .semantic_role
532                    .as_ref()
533                    .map(|r| r.category.as_str())
534                    .unwrap_or("block");
535
536                output.push_str(&format!("[{}] {}: {}\n", block_id, role, content));
537            }
538        }
539
540        output
541    }
542
543    // Internal helper methods
544
545    fn add_block_internal(
546        &mut self,
547        doc: &Document,
548        block_id: BlockId,
549        reason: InclusionReason,
550        relevance: f32,
551    ) {
552        if self.window.blocks.contains_key(&block_id) {
553            // Update access count
554            if let Some(cb) = self.window.blocks.get_mut(&block_id) {
555                cb.access_count += 1;
556                cb.last_accessed = chrono::Utc::now();
557            }
558            return;
559        }
560
561        if let Some(block) = doc.get_block(&block_id) {
562            let token_estimate = self.estimate_tokens(&block.content);
563
564            let context_block = ContextBlock {
565                block_id,
566                inclusion_reason: reason,
567                relevance_score: relevance,
568                token_estimate,
569                access_count: 1,
570                last_accessed: chrono::Utc::now(),
571                compressed: false,
572                original_content: None,
573            };
574
575            self.window.blocks.insert(block_id, context_block);
576        }
577    }
578
579    fn expand_downward(
580        &mut self,
581        doc: &Document,
582        start: BlockId,
583        max_depth: usize,
584    ) -> Vec<BlockId> {
585        let mut added = Vec::new();
586        let mut queue = VecDeque::new();
587        queue.push_back((start, 0usize));
588
589        while let Some((node_id, depth)) = queue.pop_front() {
590            if depth > max_depth || !self.window.has_capacity() {
591                break;
592            }
593
594            for child in doc.children(&node_id) {
595                if !self.window.contains(child) {
596                    let relevance = 0.6 - depth as f32 * 0.1;
597                    self.add_block_internal(
598                        doc,
599                        *child,
600                        InclusionReason::StructuralContext,
601                        relevance.max(0.1),
602                    );
603                    added.push(*child);
604                    queue.push_back((*child, depth + 1));
605                }
606            }
607        }
608
609        added
610    }
611
612    fn expand_upward(&mut self, doc: &Document, start: BlockId, max_depth: usize) -> Vec<BlockId> {
613        let mut added = Vec::new();
614        let mut current = start;
615        let mut depth = 0;
616
617        while let Some(parent) = doc.parent(&current) {
618            if *parent == doc.root || depth >= max_depth || !self.window.has_capacity() {
619                break;
620            }
621
622            if !self.window.contains(parent) {
623                let relevance = 0.7 - depth as f32 * 0.1;
624                self.add_block_internal(
625                    doc,
626                    *parent,
627                    InclusionReason::StructuralContext,
628                    relevance.max(0.1),
629                );
630                added.push(*parent);
631            }
632
633            current = *parent;
634            depth += 1;
635        }
636
637        added
638    }
639
640    fn expand_semantic(
641        &mut self,
642        doc: &Document,
643        start: BlockId,
644        max_depth: usize,
645    ) -> Vec<BlockId> {
646        let mut added = Vec::new();
647        let mut visited = HashSet::new();
648        let mut queue = VecDeque::new();
649        queue.push_back((start, 0usize));
650
651        while let Some((node_id, depth)) = queue.pop_front() {
652            if depth > max_depth || !self.window.has_capacity() {
653                break;
654            }
655
656            if visited.contains(&node_id) {
657                continue;
658            }
659            visited.insert(node_id);
660
661            if let Some(block) = doc.get_block(&node_id) {
662                for edge in &block.edges {
663                    if !self.window.contains(&edge.target) && !visited.contains(&edge.target) {
664                        let relevance = 0.5 - depth as f32 * 0.1;
665                        self.add_block_internal(
666                            doc,
667                            edge.target,
668                            InclusionReason::SemanticRelevance,
669                            relevance.max(0.1),
670                        );
671                        added.push(edge.target);
672                        queue.push_back((edge.target, depth + 1));
673                    }
674                }
675            }
676        }
677
678        added
679    }
680
681    fn prune_if_needed(&mut self) -> Vec<BlockId> {
682        let mut removed = Vec::new();
683
684        while self.window.block_count() > self.window.constraints.max_blocks
685            || self.window.total_tokens() > self.window.constraints.max_tokens
686        {
687            // Find block to remove based on policy
688            let to_remove = match self.pruning_policy {
689                PruningPolicy::RelevanceFirst => self.find_lowest_relevance(),
690                PruningPolicy::RecencyFirst => self.find_least_recent(),
691                PruningPolicy::RedundancyFirst => self.find_lowest_relevance(), // Simplified
692            };
693
694            if let Some(block_id) = to_remove {
695                self.window.blocks.remove(&block_id);
696                removed.push(block_id);
697            } else {
698                break;
699            }
700        }
701
702        removed
703    }
704
705    fn find_lowest_relevance(&self) -> Option<BlockId> {
706        self.window
707            .blocks
708            .iter()
709            .filter(|(id, _)| Some(**id) != self.window.metadata.focus_area)
710            .min_by(|a, b| {
711                a.1.relevance_score
712                    .partial_cmp(&b.1.relevance_score)
713                    .unwrap_or(std::cmp::Ordering::Equal)
714            })
715            .map(|(id, _)| *id)
716    }
717
718    fn find_least_recent(&self) -> Option<BlockId> {
719        self.window
720            .blocks
721            .iter()
722            .filter(|(id, _)| Some(**id) != self.window.metadata.focus_area)
723            .min_by(|a, b| a.1.last_accessed.cmp(&b.1.last_accessed))
724            .map(|(id, _)| *id)
725    }
726
727    fn estimate_tokens(&self, content: &Content) -> usize {
728        let text = self.extract_content_text(content);
729        // Rough estimate: ~4 characters per token
730        (text.len() / 4).max(1)
731    }
732
733    fn extract_content_text(&self, content: &Content) -> String {
734        match content {
735            Content::Text(t) => t.text.clone(),
736            Content::Code(c) => c.source.clone(),
737            Content::Table(t) => format!("Table: {} rows", t.rows.len()),
738            Content::Math(m) => m.expression.clone(),
739            Content::Media(m) => m.alt_text.clone().unwrap_or_else(|| "Media".to_string()),
740            Content::Json { .. } => "JSON data".to_string(),
741            Content::Binary { .. } => "Binary data".to_string(),
742            Content::Composite { children, .. } => {
743                format!("Composite: {} children", children.len())
744            }
745        }
746    }
747}
748
749#[cfg(test)]
750mod tests {
751    use super::*;
752    use ucm_core::DocumentId;
753
754    fn create_test_document() -> Document {
755        let mut doc = Document::new(DocumentId::new("test"));
756        let root = doc.root;
757
758        let h1 = Block::new(Content::text("Chapter 1"), Some("heading1"));
759        let h1_id = doc.add_block(h1, &root).unwrap();
760
761        let p1 = Block::new(
762            Content::text("Introduction paragraph with some content"),
763            Some("paragraph"),
764        );
765        doc.add_block(p1, &h1_id).unwrap();
766
767        let h2 = Block::new(Content::text("Section 1.1"), Some("heading2"));
768        let h2_id = doc.add_block(h2, &h1_id).unwrap();
769
770        let p2 = Block::new(Content::text("Section content here"), Some("paragraph"));
771        doc.add_block(p2, &h2_id).unwrap();
772
773        doc
774    }
775
776    #[test]
777    fn test_context_manager_creation() {
778        let manager = ContextManager::new("test-context");
779        assert_eq!(manager.window().block_count(), 0);
780        assert!(manager.window().has_capacity());
781    }
782
783    #[test]
784    fn test_initialize_focus() {
785        let doc = create_test_document();
786        let mut manager = ContextManager::new("test-context");
787
788        let root_children = doc.children(&doc.root);
789        let h1_id = root_children[0];
790
791        let result = manager.initialize_focus(&doc, h1_id, "Test task");
792
793        assert!(!result.blocks_added.is_empty());
794        assert!(manager.window().contains(&h1_id));
795    }
796
797    #[test]
798    fn test_expand_context() {
799        let doc = create_test_document();
800        let mut manager = ContextManager::new("test-context");
801
802        let root_children = doc.children(&doc.root);
803        let h1_id = root_children[0];
804
805        manager.initialize_focus(&doc, h1_id, "Test task");
806
807        let result = manager.expand_context(&doc, ExpandDirection::Down, 2);
808
809        assert!(result.total_blocks > 1);
810    }
811
812    #[test]
813    fn test_add_and_remove_block() {
814        let doc = create_test_document();
815        let mut manager = ContextManager::new("test-context");
816
817        let root_children = doc.children(&doc.root);
818        let h1_id = root_children[0];
819
820        let result = manager.add_block(&doc, h1_id, InclusionReason::DirectReference);
821        assert!(result.blocks_added.contains(&h1_id));
822        assert!(manager.window().contains(&h1_id));
823
824        let result = manager.remove_block(h1_id);
825        assert!(result.blocks_removed.contains(&h1_id));
826        assert!(!manager.window().contains(&h1_id));
827    }
828
829    #[test]
830    fn test_constraints() {
831        let constraints = ContextConstraints {
832            max_blocks: 5,
833            max_tokens: 100,
834            ..Default::default()
835        };
836
837        let doc = create_test_document();
838        let mut manager = ContextManager::with_constraints("test-context", constraints);
839
840        let root_children = doc.children(&doc.root);
841        let h1_id = root_children[0];
842
843        manager.initialize_focus(&doc, h1_id, "Test task");
844        manager.expand_context(&doc, ExpandDirection::Down, 10);
845
846        // Should be limited by constraints
847        assert!(manager.window().block_count() <= 5);
848    }
849
850    #[test]
851    fn test_statistics() {
852        let doc = create_test_document();
853        let mut manager = ContextManager::new("test-context");
854
855        let root_children = doc.children(&doc.root);
856        let h1_id = root_children[0];
857
858        manager.initialize_focus(&doc, h1_id, "Test task");
859
860        let stats = manager.get_statistics();
861        assert!(stats.total_blocks > 0);
862        assert!(stats.total_tokens > 0);
863    }
864
865    #[test]
866    fn test_render_for_prompt() {
867        let doc = create_test_document();
868        let mut manager = ContextManager::new("test-context");
869
870        let root_children = doc.children(&doc.root);
871        let h1_id = root_children[0];
872
873        manager.initialize_focus(&doc, h1_id, "Test task");
874
875        let prompt = manager.render_for_prompt(&doc);
876        assert!(!prompt.is_empty());
877        assert!(prompt.contains("Chapter 1"));
878    }
879}
ucp_llm/context.rs

ucp_llm/
context.rs
