Skip to main content

ucp_agent/
operations.rs

1//! Core operations for agent graph traversal.
2
3use crate::cursor::{CursorNeighborhood, ViewMode};
4use crate::error::{AgentError, AgentSessionId, Result};
5use crate::rag::{RagProvider, RagSearchOptions, RagSearchResults};
6use crate::safety::{CircuitBreaker, DepthGuard, GlobalLimits};
7use crate::session::{AgentSession, SessionConfig};
8use std::collections::HashMap;
9use std::sync::{Arc, RwLock};
10use std::time::Instant;
11use ucm_core::{BlockId, Document, EdgeType};
12use ucm_engine::traversal::{NavigateDirection, TraversalEngine, TraversalFilter, TraversalOutput};
13
14/// Result of a navigation operation.
15#[derive(Debug, Clone)]
16pub struct NavigationResult {
17    /// New position after navigation.
18    pub position: BlockId,
19    /// Whether the neighborhood was refreshed.
20    pub refreshed: bool,
21    /// Current neighborhood.
22    pub neighborhood: CursorNeighborhood,
23}
24
25/// Options for expansion operations.
26#[derive(Debug, Clone, Default)]
27pub struct ExpandOptions {
28    /// Maximum depth to expand.
29    pub depth: usize,
30    /// View mode for results.
31    pub view_mode: ViewMode,
32    /// Filter by semantic roles.
33    pub roles: Option<Vec<String>>,
34    /// Filter by tags.
35    pub tags: Option<Vec<String>>,
36}
37
38impl ExpandOptions {
39    pub fn new() -> Self {
40        Self {
41            depth: 3,
42            ..Default::default()
43        }
44    }
45
46    pub fn with_depth(mut self, depth: usize) -> Self {
47        self.depth = depth;
48        self
49    }
50
51    pub fn with_view_mode(mut self, mode: ViewMode) -> Self {
52        self.view_mode = mode;
53        self
54    }
55
56    pub fn with_roles(mut self, roles: Vec<String>) -> Self {
57        self.roles = Some(roles);
58        self
59    }
60
61    pub fn with_tags(mut self, tags: Vec<String>) -> Self {
62        self.tags = Some(tags);
63        self
64    }
65}
66
67/// Result of an expansion operation.
68#[derive(Debug, Clone)]
69pub struct ExpansionResult {
70    /// Root of the expansion.
71    pub root: BlockId,
72    /// Expanded blocks by depth level.
73    pub levels: Vec<Vec<BlockId>>,
74    /// Total blocks expanded.
75    pub total_blocks: usize,
76}
77
78/// Direction for expansion.
79#[derive(Debug, Clone, Copy, PartialEq, Eq)]
80pub enum ExpandDirection {
81    /// Expand to children (descendants).
82    Down,
83    /// Expand to parents (ancestors).
84    Up,
85    /// Expand in both directions.
86    Both,
87    /// Expand via semantic edges only.
88    Semantic,
89}
90
91impl From<ucl_parser::ast::ExpandDirection> for ExpandDirection {
92    fn from(d: ucl_parser::ast::ExpandDirection) -> Self {
93        match d {
94            ucl_parser::ast::ExpandDirection::Down => ExpandDirection::Down,
95            ucl_parser::ast::ExpandDirection::Up => ExpandDirection::Up,
96            ucl_parser::ast::ExpandDirection::Both => ExpandDirection::Both,
97            ucl_parser::ast::ExpandDirection::Semantic => ExpandDirection::Semantic,
98        }
99    }
100}
101
102/// Options for search operations.
103#[derive(Debug, Clone, Default)]
104pub struct SearchOptions {
105    /// Maximum results to return.
106    pub limit: usize,
107    /// Minimum similarity threshold.
108    pub min_similarity: f32,
109    /// Filter by semantic roles.
110    pub roles: Option<Vec<String>>,
111    /// Filter by tags.
112    pub tags: Option<Vec<String>>,
113}
114
115impl SearchOptions {
116    pub fn new() -> Self {
117        Self {
118            limit: 10,
119            min_similarity: 0.0,
120            roles: None,
121            tags: None,
122        }
123    }
124
125    pub fn with_limit(mut self, limit: usize) -> Self {
126        self.limit = limit;
127        self
128    }
129
130    pub fn with_min_similarity(mut self, threshold: f32) -> Self {
131        self.min_similarity = threshold;
132        self
133    }
134}
135
136/// Result of a find operation.
137#[derive(Debug, Clone)]
138pub struct FindResult {
139    /// Matching block IDs.
140    pub matches: Vec<BlockId>,
141    /// Total blocks searched.
142    pub total_searched: usize,
143}
144
145/// View of a block's content.
146#[derive(Debug, Clone)]
147pub struct BlockView {
148    /// Block ID.
149    pub block_id: BlockId,
150    /// Content (based on view mode).
151    pub content: Option<String>,
152    /// Semantic role.
153    pub role: Option<String>,
154    /// Tags.
155    pub tags: Vec<String>,
156    /// Children count.
157    pub children_count: usize,
158    /// Incoming edges count.
159    pub incoming_edges: usize,
160    /// Outgoing edges count.
161    pub outgoing_edges: usize,
162}
163
164/// View of the cursor neighborhood.
165#[derive(Debug, Clone)]
166pub struct NeighborhoodView {
167    /// Current position.
168    pub position: BlockId,
169    /// Parent blocks.
170    pub ancestors: Vec<BlockView>,
171    /// Child blocks.
172    pub children: Vec<BlockView>,
173    /// Sibling blocks.
174    pub siblings: Vec<BlockView>,
175    /// Connected blocks via semantic edges.
176    pub connections: Vec<(BlockView, EdgeType)>,
177}
178
179/// Main interface for agent graph traversal operations.
180pub struct AgentTraversal {
181    /// Active sessions.
182    sessions: RwLock<HashMap<AgentSessionId, AgentSession>>,
183    /// The document being traversed.
184    document: Arc<RwLock<Document>>,
185    /// Optional RAG provider for semantic search.
186    rag_provider: Option<Arc<dyn RagProvider>>,
187    /// Global limits for all sessions.
188    global_limits: GlobalLimits,
189    /// Circuit breaker for fault tolerance.
190    circuit_breaker: CircuitBreaker,
191    /// Depth guard for recursion protection.
192    depth_guard: DepthGuard,
193}
194
195impl AgentTraversal {
196    /// Create a new agent traversal system.
197    pub fn new(document: Document) -> Self {
198        Self {
199            sessions: RwLock::new(HashMap::new()),
200            document: Arc::new(RwLock::new(document)),
201            rag_provider: None,
202            global_limits: GlobalLimits::default(),
203            circuit_breaker: CircuitBreaker::new(5, std::time::Duration::from_secs(30)),
204            depth_guard: DepthGuard::new(100),
205        }
206    }
207
208    /// Create with a RAG provider.
209    pub fn with_rag_provider(mut self, provider: Arc<dyn RagProvider>) -> Self {
210        self.rag_provider = Some(provider);
211        self
212    }
213
214    /// Create with custom global limits.
215    pub fn with_global_limits(mut self, limits: GlobalLimits) -> Self {
216        self.global_limits = limits;
217        self
218    }
219
220    /// Update the internal document with a new copy.
221    ///
222    /// Use this when you've added blocks to the original document
223    /// after creating the AgentTraversal.
224    pub fn update_document(&self, document: Document) -> Result<()> {
225        let mut doc = self
226            .document
227            .write()
228            .map_err(|_| AgentError::Internal("Failed to acquire document lock".to_string()))?;
229        *doc = document;
230        Ok(())
231    }
232
233    /// Get a clone of the internal document.
234    pub fn get_document(&self) -> Result<Document> {
235        let doc = self
236            .document
237            .read()
238            .map_err(|_| AgentError::Internal("Failed to acquire document lock".to_string()))?;
239        Ok(doc.clone())
240    }
241
242    // ==================== Session Management ====================
243
244    /// Create a new agent session.
245    pub fn create_session(&self, config: SessionConfig) -> Result<AgentSessionId> {
246        self.circuit_breaker.can_proceed()?;
247
248        let mut sessions = self
249            .sessions
250            .write()
251            .map_err(|_| AgentError::Internal("Failed to acquire sessions lock".to_string()))?;
252
253        if sessions.len() >= self.global_limits.max_sessions {
254            return Err(AgentError::MaxSessionsReached {
255                max: self.global_limits.max_sessions,
256            });
257        }
258
259        let doc = self
260            .document
261            .read()
262            .map_err(|_| AgentError::Internal("Failed to acquire document lock".to_string()))?;
263
264        // Get start block or default to document root
265        let start_block = config.start_block.unwrap_or(doc.root);
266
267        let session = AgentSession::new(start_block, config);
268        let session_id = session.id.clone();
269        sessions.insert(session_id.clone(), session);
270
271        Ok(session_id)
272    }
273
274    /// Get a reference to a session.
275    pub fn get_session(
276        &self,
277        id: &AgentSessionId,
278    ) -> Result<std::sync::RwLockReadGuard<'_, HashMap<AgentSessionId, AgentSession>>> {
279        let sessions = self
280            .sessions
281            .read()
282            .map_err(|_| AgentError::Internal("Failed to acquire sessions lock".to_string()))?;
283        if !sessions.contains_key(id) {
284            return Err(AgentError::SessionNotFound(id.clone()));
285        }
286        Ok(sessions)
287    }
288
289    /// Close a session.
290    pub fn close_session(&self, id: &AgentSessionId) -> Result<()> {
291        let mut sessions = self
292            .sessions
293            .write()
294            .map_err(|_| AgentError::Internal("Failed to acquire sessions lock".to_string()))?;
295
296        let session = sessions
297            .get_mut(id)
298            .ok_or_else(|| AgentError::SessionNotFound(id.clone()))?;
299
300        session.complete();
301        sessions.remove(id);
302        Ok(())
303    }
304
305    // ==================== Navigation ====================
306
307    /// Navigate to a specific block.
308    pub fn navigate_to(
309        &self,
310        session_id: &AgentSessionId,
311        target: BlockId,
312    ) -> Result<NavigationResult> {
313        self.circuit_breaker.can_proceed()?;
314        let _guard = self.depth_guard.try_enter()?;
315        let start = Instant::now();
316
317        let mut sessions = self
318            .sessions
319            .write()
320            .map_err(|_| AgentError::Internal("Failed to acquire sessions lock".to_string()))?;
321
322        let session = sessions
323            .get_mut(session_id)
324            .ok_or_else(|| AgentError::SessionNotFound(session_id.clone()))?;
325
326        session.check_can_traverse()?;
327
328        // Verify block exists
329        let doc = self
330            .document
331            .read()
332            .map_err(|_| AgentError::Internal("Failed to acquire document lock".to_string()))?;
333
334        if doc.get_block(&target).is_none() {
335            return Err(AgentError::BlockNotFound(target));
336        }
337
338        // Move cursor
339        session.cursor.move_to(target);
340        session.touch();
341        session.metrics.record_navigation();
342        session.budget.record_traversal();
343
344        // Refresh neighborhood
345        let neighborhood = self.compute_neighborhood(&doc, &target)?;
346        session.cursor.update_neighborhood(neighborhood.clone());
347
348        session.metrics.record_execution_time(start.elapsed());
349
350        Ok(NavigationResult {
351            position: target,
352            refreshed: true,
353            neighborhood,
354        })
355    }
356
357    /// Go back in navigation history.
358    pub fn go_back(&self, session_id: &AgentSessionId, steps: usize) -> Result<NavigationResult> {
359        self.circuit_breaker.can_proceed()?;
360        let start = Instant::now();
361
362        let mut sessions = self
363            .sessions
364            .write()
365            .map_err(|_| AgentError::Internal("Failed to acquire sessions lock".to_string()))?;
366
367        let session = sessions
368            .get_mut(session_id)
369            .ok_or_else(|| AgentError::SessionNotFound(session_id.clone()))?;
370
371        session.check_can_traverse()?;
372
373        let position = session
374            .cursor
375            .go_back(steps)
376            .ok_or(AgentError::EmptyHistory)?;
377
378        session.touch();
379        session.metrics.record_navigation();
380
381        // Refresh neighborhood
382        let doc = self
383            .document
384            .read()
385            .map_err(|_| AgentError::Internal("Failed to acquire document lock".to_string()))?;
386
387        let neighborhood = self.compute_neighborhood(&doc, &position)?;
388        session.cursor.update_neighborhood(neighborhood.clone());
389
390        session.metrics.record_execution_time(start.elapsed());
391
392        Ok(NavigationResult {
393            position,
394            refreshed: true,
395            neighborhood,
396        })
397    }
398
399    // ==================== Expansion ====================
400
401    /// Expand from a block in a given direction.
402    pub fn expand(
403        &self,
404        session_id: &AgentSessionId,
405        block_id: BlockId,
406        direction: ExpandDirection,
407        options: ExpandOptions,
408    ) -> Result<ExpansionResult> {
409        self.circuit_breaker.can_proceed()?;
410        let _guard = self.depth_guard.try_enter()?;
411        let start = Instant::now();
412
413        let sessions = self
414            .sessions
415            .read()
416            .map_err(|_| AgentError::Internal("Failed to acquire sessions lock".to_string()))?;
417
418        let session = sessions
419            .get(session_id)
420            .ok_or_else(|| AgentError::SessionNotFound(session_id.clone()))?;
421
422        session.check_can_traverse()?;
423
424        // Check depth limit
425        if options.depth > session.limits.max_expand_depth {
426            return Err(AgentError::DepthLimitExceeded {
427                current: options.depth,
428                max: session.limits.max_expand_depth,
429            });
430        }
431
432        let doc = self
433            .document
434            .read()
435            .map_err(|_| AgentError::Internal("Failed to acquire document lock".to_string()))?;
436
437        // Build traversal filter
438        let filter = self.build_traversal_filter(&options);
439
440        let levels = match direction {
441            ExpandDirection::Down => self.expand_down(&doc, &block_id, options.depth, &filter)?,
442            ExpandDirection::Up => self.expand_up(&doc, &block_id, options.depth)?,
443            ExpandDirection::Both => {
444                let mut down = self.expand_down(&doc, &block_id, options.depth, &filter)?;
445                let up = self.expand_up(&doc, &block_id, options.depth)?;
446                down.extend(up);
447                down
448            }
449            ExpandDirection::Semantic => self.expand_semantic(&doc, &block_id, options.depth)?,
450        };
451
452        let total_blocks: usize = levels.iter().map(|l| l.len()).sum();
453
454        // Update metrics
455        drop(sessions);
456        let mut sessions_mut = self
457            .sessions
458            .write()
459            .map_err(|_| AgentError::Internal("Failed to acquire sessions lock".to_string()))?;
460
461        if let Some(session) = sessions_mut.get_mut(session_id) {
462            session.metrics.record_expansion(total_blocks);
463            session.budget.record_traversal();
464            session.metrics.record_execution_time(start.elapsed());
465            session.touch();
466        }
467
468        Ok(ExpansionResult {
469            root: block_id,
470            levels,
471            total_blocks,
472        })
473    }
474
475    // ==================== Search ====================
476
477    /// Perform semantic search (requires RAG provider).
478    pub async fn search(
479        &self,
480        session_id: &AgentSessionId,
481        query: &str,
482        options: SearchOptions,
483    ) -> Result<RagSearchResults> {
484        self.circuit_breaker.can_proceed()?;
485        let start = Instant::now();
486
487        let rag = self
488            .rag_provider
489            .as_ref()
490            .ok_or(AgentError::RagNotConfigured)?;
491
492        {
493            let sessions = self
494                .sessions
495                .read()
496                .map_err(|_| AgentError::Internal("Failed to acquire sessions lock".to_string()))?;
497
498            let session = sessions
499                .get(session_id)
500                .ok_or_else(|| AgentError::SessionNotFound(session_id.clone()))?;
501
502            session.check_can_search()?;
503        }
504
505        let rag_options = RagSearchOptions::new()
506            .with_limit(options.limit)
507            .with_min_similarity(options.min_similarity);
508
509        let results = rag.search(query, rag_options).await?;
510
511        // Store results for CTX ADD RESULTS
512        let mut sessions = self
513            .sessions
514            .write()
515            .map_err(|_| AgentError::Internal("Failed to acquire sessions lock".to_string()))?;
516
517        if let Some(session) = sessions.get_mut(session_id) {
518            session.store_results(results.block_ids());
519            session.metrics.record_search();
520            session.metrics.record_execution_time(start.elapsed());
521            session.touch();
522        }
523
524        Ok(results)
525    }
526
527    /// Find blocks by pattern (no RAG required).
528    pub fn find_by_pattern(
529        &self,
530        session_id: &AgentSessionId,
531        role: Option<&str>,
532        tag: Option<&str>,
533        label: Option<&str>,
534        pattern: Option<&str>,
535    ) -> Result<FindResult> {
536        self.circuit_breaker.can_proceed()?;
537        let start = Instant::now();
538
539        let sessions = self
540            .sessions
541            .read()
542            .map_err(|_| AgentError::Internal("Failed to acquire sessions lock".to_string()))?;
543
544        let session = sessions
545            .get(session_id)
546            .ok_or_else(|| AgentError::SessionNotFound(session_id.clone()))?;
547
548        session.check_can_search()?;
549
550        let doc = self
551            .document
552            .read()
553            .map_err(|_| AgentError::Internal("Failed to acquire document lock".to_string()))?;
554
555        let mut matches = Vec::new();
556        let mut total_searched = 0;
557
558        // Build regex if pattern provided
559        let regex = pattern
560            .map(regex::Regex::new)
561            .transpose()
562            .map_err(|e| AgentError::Internal(format!("Invalid regex pattern: {}", e)))?;
563
564        for block in doc.blocks.values() {
565            total_searched += 1;
566
567            // Filter by role
568            if let Some(r) = role {
569                let block_role = block
570                    .metadata
571                    .semantic_role
572                    .as_ref()
573                    .map(|sr| sr.category.as_str())
574                    .unwrap_or("");
575                if block_role != r {
576                    continue;
577                }
578            }
579
580            // Filter by tag
581            if let Some(t) = tag {
582                if !block.metadata.tags.contains(&t.to_string()) {
583                    continue;
584                }
585            }
586
587            // Filter by label
588            if let Some(l) = label {
589                if block.metadata.label.as_deref() != Some(l) {
590                    continue;
591                }
592            }
593
594            // Filter by content pattern
595            if let Some(ref re) = regex {
596                let content = self.extract_content_text(&block.content);
597                if !re.is_match(&content) {
598                    continue;
599                }
600            }
601
602            matches.push(block.id);
603        }
604
605        // Store results for CTX ADD RESULTS
606        drop(sessions);
607        let mut sessions_mut = self
608            .sessions
609            .write()
610            .map_err(|_| AgentError::Internal("Failed to acquire sessions lock".to_string()))?;
611
612        if let Some(session) = sessions_mut.get_mut(session_id) {
613            session.store_results(matches.clone());
614            session.metrics.record_search();
615            session.metrics.record_execution_time(start.elapsed());
616            session.touch();
617        }
618
619        Ok(FindResult {
620            matches,
621            total_searched,
622        })
623    }
624
625    // ==================== View ====================
626
627    /// View a specific block.
628    pub fn view_block(
629        &self,
630        session_id: &AgentSessionId,
631        block_id: BlockId,
632        mode: ViewMode,
633    ) -> Result<BlockView> {
634        self.circuit_breaker.can_proceed()?;
635
636        let sessions = self
637            .sessions
638            .read()
639            .map_err(|_| AgentError::Internal("Failed to acquire sessions lock".to_string()))?;
640
641        let session = sessions
642            .get(session_id)
643            .ok_or_else(|| AgentError::SessionNotFound(session_id.clone()))?;
644
645        session.check_active()?;
646
647        let doc = self
648            .document
649            .read()
650            .map_err(|_| AgentError::Internal("Failed to acquire document lock".to_string()))?;
651
652        self.view_block_internal(&doc, &block_id, &mode)
653    }
654
655    /// View the neighborhood around the current cursor position.
656    pub fn view_neighborhood(&self, session_id: &AgentSessionId) -> Result<NeighborhoodView> {
657        self.circuit_breaker.can_proceed()?;
658
659        let sessions = self
660            .sessions
661            .read()
662            .map_err(|_| AgentError::Internal("Failed to acquire sessions lock".to_string()))?;
663
664        let session = sessions
665            .get(session_id)
666            .ok_or_else(|| AgentError::SessionNotFound(session_id.clone()))?;
667
668        session.check_active()?;
669
670        let position = session.cursor.position;
671        let view_mode = session.cursor.view_mode.clone();
672
673        // Release session lock before calling view_block
674        drop(sessions);
675
676        let doc = self
677            .document
678            .read()
679            .map_err(|_| AgentError::Internal("Failed to acquire document lock".to_string()))?;
680
681        let neighborhood = self.compute_neighborhood(&doc, &position)?;
682
683        // Build views for each block in neighborhood
684        let ancestors: Vec<BlockView> = neighborhood
685            .ancestors
686            .iter()
687            .filter_map(|id| self.view_block_internal(&doc, id, &view_mode).ok())
688            .collect();
689
690        let children: Vec<BlockView> = neighborhood
691            .children
692            .iter()
693            .filter_map(|id| self.view_block_internal(&doc, id, &view_mode).ok())
694            .collect();
695
696        let siblings: Vec<BlockView> = neighborhood
697            .siblings
698            .iter()
699            .filter_map(|id| self.view_block_internal(&doc, id, &view_mode).ok())
700            .collect();
701
702        let connections: Vec<(BlockView, EdgeType)> = neighborhood
703            .connections
704            .iter()
705            .filter_map(|(id, edge_type)| {
706                self.view_block_internal(&doc, id, &view_mode)
707                    .ok()
708                    .map(|view| (view, edge_type.clone()))
709            })
710            .collect();
711
712        Ok(NeighborhoodView {
713            position,
714            ancestors,
715            children,
716            siblings,
717            connections,
718        })
719    }
720
721    // ==================== Path Finding ====================
722
723    /// Find a path between two blocks.
724    pub fn find_path(
725        &self,
726        session_id: &AgentSessionId,
727        from: BlockId,
728        to: BlockId,
729        max_length: Option<usize>,
730    ) -> Result<Vec<BlockId>> {
731        self.circuit_breaker.can_proceed()?;
732        let _guard = self.depth_guard.try_enter()?;
733        let start = Instant::now();
734
735        let sessions = self
736            .sessions
737            .read()
738            .map_err(|_| AgentError::Internal("Failed to acquire sessions lock".to_string()))?;
739
740        let session = sessions
741            .get(session_id)
742            .ok_or_else(|| AgentError::SessionNotFound(session_id.clone()))?;
743
744        session.check_can_traverse()?;
745
746        let doc = self
747            .document
748            .read()
749            .map_err(|_| AgentError::Internal("Failed to acquire document lock".to_string()))?;
750
751        // Simple BFS path finding
752        let max_depth = max_length.unwrap_or(10);
753        let path = self.bfs_path(&doc, &from, &to, max_depth)?;
754
755        // Update metrics
756        drop(sessions);
757        let mut sessions_mut = self
758            .sessions
759            .write()
760            .map_err(|_| AgentError::Internal("Failed to acquire sessions lock".to_string()))?;
761
762        if let Some(session) = sessions_mut.get_mut(session_id) {
763            session.metrics.record_traversal();
764            session.budget.record_traversal();
765            session.metrics.record_execution_time(start.elapsed());
766            session.touch();
767        }
768
769        Ok(path)
770    }
771
772    // ==================== Context Operations ====================
773
774    /// Add a block to the context window.
775    pub fn context_add(
776        &self,
777        session_id: &AgentSessionId,
778        block_id: BlockId,
779        _reason: Option<String>,
780        _relevance: Option<f32>,
781    ) -> Result<()> {
782        let mut sessions = self
783            .sessions
784            .write()
785            .map_err(|_| AgentError::Internal("Failed to acquire sessions lock".to_string()))?;
786
787        let session = sessions
788            .get_mut(session_id)
789            .ok_or_else(|| AgentError::SessionNotFound(session_id.clone()))?;
790
791        session.check_can_modify_context()?;
792
793        // For now, we just verify the block exists
794        let doc = self
795            .document
796            .read()
797            .map_err(|_| AgentError::Internal("Failed to acquire document lock".to_string()))?;
798
799        if doc.get_block(&block_id).is_none() {
800            return Err(AgentError::BlockNotFound(block_id));
801        }
802
803        session.metrics.record_context_add(1);
804        session.touch();
805
806        // Context management is handled by the external context manager
807        // This is a placeholder for integration
808        Ok(())
809    }
810
811    /// Add all last results to context.
812    pub fn context_add_results(&self, session_id: &AgentSessionId) -> Result<Vec<BlockId>> {
813        let mut sessions = self
814            .sessions
815            .write()
816            .map_err(|_| AgentError::Internal("Failed to acquire sessions lock".to_string()))?;
817
818        let session = sessions
819            .get_mut(session_id)
820            .ok_or_else(|| AgentError::SessionNotFound(session_id.clone()))?;
821
822        session.check_can_modify_context()?;
823
824        let results = session.get_last_results()?.to_vec();
825        session.metrics.record_context_add(results.len());
826        session.touch();
827
828        Ok(results)
829    }
830
831    /// Remove a block from context.
832    pub fn context_remove(&self, session_id: &AgentSessionId, _block_id: BlockId) -> Result<()> {
833        let mut sessions = self
834            .sessions
835            .write()
836            .map_err(|_| AgentError::Internal("Failed to acquire sessions lock".to_string()))?;
837
838        let session = sessions
839            .get_mut(session_id)
840            .ok_or_else(|| AgentError::SessionNotFound(session_id.clone()))?;
841
842        session.check_can_modify_context()?;
843        session.metrics.record_context_remove();
844        session.touch();
845
846        Ok(())
847    }
848
849    /// Clear the context window.
850    pub fn context_clear(&self, session_id: &AgentSessionId) -> Result<()> {
851        let mut sessions = self
852            .sessions
853            .write()
854            .map_err(|_| AgentError::Internal("Failed to acquire sessions lock".to_string()))?;
855
856        let session = sessions
857            .get_mut(session_id)
858            .ok_or_else(|| AgentError::SessionNotFound(session_id.clone()))?;
859
860        session.check_can_modify_context()?;
861        session.touch();
862
863        Ok(())
864    }
865
866    /// Set focus block.
867    pub fn context_focus(
868        &self,
869        session_id: &AgentSessionId,
870        block_id: Option<BlockId>,
871    ) -> Result<()> {
872        let mut sessions = self
873            .sessions
874            .write()
875            .map_err(|_| AgentError::Internal("Failed to acquire sessions lock".to_string()))?;
876
877        let session = sessions
878            .get_mut(session_id)
879            .ok_or_else(|| AgentError::SessionNotFound(session_id.clone()))?;
880
881        session.set_focus(block_id);
882        session.touch();
883
884        Ok(())
885    }
886
887    // ==================== Internal Helpers ====================
888
889    fn compute_neighborhood(
890        &self,
891        doc: &Document,
892        position: &BlockId,
893    ) -> Result<CursorNeighborhood> {
894        let mut neighborhood = CursorNeighborhood::new();
895
896        // Get ancestors
897        let mut current = *position;
898        for _ in 0..5 {
899            if let Some(parent) = doc.parent(&current) {
900                neighborhood.ancestors.push(*parent);
901                current = *parent;
902            } else {
903                break;
904            }
905        }
906
907        // Get children
908        neighborhood.children = doc.children(position).to_vec();
909
910        // Get siblings
911        if let Some(parent) = doc.parent(position) {
912            neighborhood.siblings = doc
913                .children(parent)
914                .iter()
915                .filter(|id| *id != position)
916                .copied()
917                .collect();
918        }
919
920        // Get semantic connections via edge index
921        for (edge_type, target) in doc.edge_index.outgoing_from(position) {
922            neighborhood.connections.push((*target, edge_type.clone()));
923        }
924
925        neighborhood.stale = false;
926        Ok(neighborhood)
927    }
928
929    fn view_block_internal(
930        &self,
931        doc: &Document,
932        block_id: &BlockId,
933        mode: &ViewMode,
934    ) -> Result<BlockView> {
935        let block = doc
936            .get_block(block_id)
937            .ok_or(AgentError::BlockNotFound(*block_id))?;
938
939        let content = match mode {
940            ViewMode::IdsOnly => None,
941            ViewMode::Preview { length } => {
942                let text = self.extract_content_text(&block.content);
943                Some(text.chars().take(*length).collect())
944            }
945            ViewMode::Full => Some(self.extract_content_text(&block.content)),
946            ViewMode::Metadata => None,
947            ViewMode::Adaptive { .. } => Some(self.extract_content_text(&block.content)),
948        };
949
950        let outgoing_edges = doc.edge_index.outgoing_from(block_id).len();
951        let incoming_edges = doc.edge_index.incoming_to(block_id).len();
952
953        Ok(BlockView {
954            block_id: *block_id,
955            content,
956            role: block
957                .metadata
958                .semantic_role
959                .as_ref()
960                .map(|r| r.category.as_str().to_string()),
961            tags: block.metadata.tags.clone(),
962            children_count: doc.children(block_id).len(),
963            incoming_edges,
964            outgoing_edges,
965        })
966    }
967
968    fn extract_content_text(&self, content: &ucm_core::Content) -> String {
969        match content {
970            ucm_core::Content::Text(t) => t.text.clone(),
971            ucm_core::Content::Code(c) => c.source.clone(),
972            ucm_core::Content::Table(t) => format!("Table: {} rows", t.rows.len()),
973            ucm_core::Content::Math(m) => m.expression.clone(),
974            ucm_core::Content::Media(m) => {
975                m.alt_text.clone().unwrap_or_else(|| "Media".to_string())
976            }
977            ucm_core::Content::Json { .. } => "JSON data".to_string(),
978            ucm_core::Content::Binary { .. } => "Binary data".to_string(),
979            ucm_core::Content::Composite { children, .. } => {
980                format!("Composite: {} children", children.len())
981            }
982        }
983    }
984
985    fn build_traversal_filter(&self, options: &ExpandOptions) -> TraversalFilter {
986        let mut filter = TraversalFilter::default();
987
988        if let Some(ref roles) = options.roles {
989            filter.include_roles = roles.clone();
990        }
991
992        if let Some(ref tags) = options.tags {
993            filter.include_tags = tags.clone();
994        }
995
996        filter
997    }
998
999    fn expand_down(
1000        &self,
1001        doc: &Document,
1002        block_id: &BlockId,
1003        depth: usize,
1004        filter: &TraversalFilter,
1005    ) -> Result<Vec<Vec<BlockId>>> {
1006        let engine = TraversalEngine::new();
1007        let result = engine
1008            .navigate(
1009                doc,
1010                Some(*block_id),
1011                NavigateDirection::BreadthFirst,
1012                Some(depth),
1013                Some(filter.clone()),
1014                TraversalOutput::StructureOnly,
1015            )
1016            .map_err(|e| AgentError::EngineError(e.to_string()))?;
1017
1018        // Group by depth level
1019        let mut levels: Vec<Vec<BlockId>> = vec![Vec::new(); depth + 1];
1020        for node in result.nodes {
1021            if node.depth <= depth {
1022                levels[node.depth].push(node.id);
1023            }
1024        }
1025
1026        // Remove empty trailing levels
1027        while levels.last().map(|l| l.is_empty()).unwrap_or(false) {
1028            levels.pop();
1029        }
1030
1031        Ok(levels)
1032    }
1033
1034    fn expand_up(
1035        &self,
1036        doc: &Document,
1037        block_id: &BlockId,
1038        depth: usize,
1039    ) -> Result<Vec<Vec<BlockId>>> {
1040        let mut levels = Vec::new();
1041        let mut current = *block_id;
1042
1043        for _ in 0..depth {
1044            if let Some(parent) = doc.parent(&current) {
1045                levels.push(vec![*parent]);
1046                current = *parent;
1047            } else {
1048                break;
1049            }
1050        }
1051
1052        Ok(levels)
1053    }
1054
1055    fn expand_semantic(
1056        &self,
1057        doc: &Document,
1058        block_id: &BlockId,
1059        depth: usize,
1060    ) -> Result<Vec<Vec<BlockId>>> {
1061        let mut levels = Vec::new();
1062        let mut visited = std::collections::HashSet::new();
1063        let mut current_level = vec![*block_id];
1064        visited.insert(*block_id);
1065
1066        for _ in 0..depth {
1067            let mut next_level = Vec::new();
1068
1069            for id in &current_level {
1070                for (_, target) in doc.edge_index.outgoing_from(id) {
1071                    if !visited.contains(target) {
1072                        visited.insert(*target);
1073                        next_level.push(*target);
1074                    }
1075                }
1076            }
1077
1078            if next_level.is_empty() {
1079                break;
1080            }
1081
1082            levels.push(next_level.clone());
1083            current_level = next_level;
1084        }
1085
1086        Ok(levels)
1087    }
1088
1089    fn bfs_path(
1090        &self,
1091        doc: &Document,
1092        from: &BlockId,
1093        to: &BlockId,
1094        max_depth: usize,
1095    ) -> Result<Vec<BlockId>> {
1096        use std::collections::{HashSet, VecDeque};
1097
1098        if from == to {
1099            return Ok(vec![*from]);
1100        }
1101
1102        let mut visited = HashSet::new();
1103        let mut queue = VecDeque::new();
1104        let mut parent_map: HashMap<BlockId, BlockId> = HashMap::new();
1105
1106        queue.push_back((*from, 0));
1107        visited.insert(*from);
1108
1109        while let Some((current, depth)) = queue.pop_front() {
1110            if depth >= max_depth {
1111                continue;
1112            }
1113
1114            // Check children
1115            for child in doc.children(&current) {
1116                if !visited.contains(child) {
1117                    visited.insert(*child);
1118                    parent_map.insert(*child, current);
1119
1120                    if child == to {
1121                        // Reconstruct path
1122                        let mut path = vec![*to];
1123                        let mut c = *to;
1124                        while let Some(p) = parent_map.get(&c) {
1125                            path.push(*p);
1126                            c = *p;
1127                        }
1128                        path.reverse();
1129                        return Ok(path);
1130                    }
1131
1132                    queue.push_back((*child, depth + 1));
1133                }
1134            }
1135
1136            // Check parent
1137            if let Some(parent) = doc.parent(&current) {
1138                if !visited.contains(parent) {
1139                    visited.insert(*parent);
1140                    parent_map.insert(*parent, current);
1141
1142                    if parent == to {
1143                        let mut path = vec![*to];
1144                        let mut c = *to;
1145                        while let Some(p) = parent_map.get(&c) {
1146                            path.push(*p);
1147                            c = *p;
1148                        }
1149                        path.reverse();
1150                        return Ok(path);
1151                    }
1152
1153                    queue.push_back((*parent, depth + 1));
1154                }
1155            }
1156
1157            // Check semantic edges
1158            for (_, target) in doc.edge_index.outgoing_from(&current) {
1159                if !visited.contains(target) {
1160                    visited.insert(*target);
1161                    parent_map.insert(*target, current);
1162
1163                    if target == to {
1164                        let mut path = vec![*to];
1165                        let mut c = *to;
1166                        while let Some(p) = parent_map.get(&c) {
1167                            path.push(*p);
1168                            c = *p;
1169                        }
1170                        path.reverse();
1171                        return Ok(path);
1172                    }
1173
1174                    queue.push_back((*target, depth + 1));
1175                }
1176            }
1177        }
1178
1179        Err(AgentError::NoPathExists {
1180            from: *from,
1181            to: *to,
1182        })
1183    }
1184}
1185
1186#[cfg(test)]
1187mod tests {
1188    use super::*;
1189
1190    fn create_test_document() -> Document {
1191        Document::create()
1192    }
1193
1194    #[test]
1195    fn test_create_session() {
1196        let doc = create_test_document();
1197        let traversal = AgentTraversal::new(doc);
1198
1199        let session_id = traversal.create_session(SessionConfig::default()).unwrap();
1200        assert!(!session_id.0.is_nil());
1201    }
1202
1203    #[test]
1204    fn test_close_session() {
1205        let doc = create_test_document();
1206        let traversal = AgentTraversal::new(doc);
1207
1208        let session_id = traversal.create_session(SessionConfig::default()).unwrap();
1209        assert!(traversal.close_session(&session_id).is_ok());
1210
1211        // Session should be removed
1212        assert!(traversal.close_session(&session_id).is_err());
1213    }
1214
1215    #[test]
1216    fn test_max_sessions_limit() {
1217        let doc = create_test_document();
1218        let traversal = AgentTraversal::new(doc).with_global_limits(GlobalLimits {
1219            max_sessions: 2,
1220            ..Default::default()
1221        });
1222
1223        // Create 2 sessions
1224        let _ = traversal.create_session(SessionConfig::default()).unwrap();
1225        let _ = traversal.create_session(SessionConfig::default()).unwrap();
1226
1227        // Third should fail
1228        let result = traversal.create_session(SessionConfig::default());
1229        assert!(matches!(
1230            result,
1231            Err(AgentError::MaxSessionsReached { max: 2 })
1232        ));
1233    }
1234}