Skip to main content

ucp_graph/
session.rs

1use std::collections::{HashMap, HashSet, VecDeque};
2
3use serde::{Deserialize, Serialize};
4use ucm_core::BlockId;
5
6use crate::{
7    navigator::GraphNavigator,
8    query::GraphNeighborMode,
9    store::GraphStoreError,
10    types::{GraphDetailLevel, GraphEdgeSummary, GraphNodeSummary},
11};
12
13#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
14#[serde(rename_all = "snake_case")]
15pub enum GraphSelectionOriginKind {
16    Overview,
17    Manual,
18    Children,
19    Parents,
20    Outgoing,
21    Incoming,
22}
23
24#[derive(Debug, Clone, Serialize, Deserialize)]
25pub struct GraphSelectionOrigin {
26    pub kind: GraphSelectionOriginKind,
27    #[serde(default, skip_serializing_if = "Option::is_none")]
28    pub relation: Option<String>,
29    #[serde(default, skip_serializing_if = "Option::is_none")]
30    pub anchor: Option<BlockId>,
31}
32
33#[derive(Debug, Clone, Serialize, Deserialize)]
34pub struct GraphSessionNode {
35    pub detail_level: GraphDetailLevel,
36    #[serde(default)]
37    pub pinned: bool,
38    #[serde(default, skip_serializing_if = "Option::is_none")]
39    pub origin: Option<GraphSelectionOrigin>,
40}
41
42#[derive(Debug, Clone, Default, Serialize, Deserialize)]
43pub struct GraphSessionUpdate {
44    #[serde(default)]
45    pub added: Vec<BlockId>,
46    #[serde(default)]
47    pub removed: Vec<BlockId>,
48    #[serde(default)]
49    pub changed: Vec<BlockId>,
50    #[serde(default)]
51    pub focus: Option<BlockId>,
52    #[serde(default)]
53    pub warnings: Vec<String>,
54}
55
56#[derive(Debug, Clone, Serialize, Deserialize)]
57pub struct GraphSessionSummary {
58    pub selected: usize,
59    pub pinned: usize,
60    pub focused: bool,
61    pub roots: usize,
62    pub leaves: usize,
63}
64
65#[derive(Debug, Clone, Serialize, Deserialize)]
66pub struct GraphSelectionExplanation {
67    pub block_id: BlockId,
68    pub selected: bool,
69    pub focus: bool,
70    pub pinned: bool,
71    #[serde(default, skip_serializing_if = "Option::is_none")]
72    pub detail_level: Option<GraphDetailLevel>,
73    #[serde(default, skip_serializing_if = "Option::is_none")]
74    pub origin: Option<GraphSelectionOrigin>,
75    pub explanation: String,
76    #[serde(default, skip_serializing_if = "Option::is_none")]
77    pub node: Option<GraphNodeSummary>,
78    #[serde(default, skip_serializing_if = "Option::is_none")]
79    pub anchor: Option<GraphNodeSummary>,
80}
81
82#[derive(Debug, Clone, Serialize, Deserialize)]
83pub struct GraphSessionDiff {
84    #[serde(default)]
85    pub added: Vec<GraphNodeSummary>,
86    #[serde(default)]
87    pub removed: Vec<GraphNodeSummary>,
88    #[serde(default, skip_serializing_if = "Option::is_none")]
89    pub focus_before: Option<BlockId>,
90    #[serde(default, skip_serializing_if = "Option::is_none")]
91    pub focus_after: Option<BlockId>,
92    pub changed_focus: bool,
93}
94
95#[derive(Debug, Clone, Serialize, Deserialize)]
96pub struct GraphExportNode {
97    pub block_id: BlockId,
98    pub label: String,
99    pub content_type: String,
100    #[serde(default, skip_serializing_if = "Option::is_none")]
101    pub semantic_role: Option<String>,
102    #[serde(default)]
103    pub tags: Vec<String>,
104    pub detail_level: GraphDetailLevel,
105    pub pinned: bool,
106    #[serde(default, skip_serializing_if = "Option::is_none")]
107    pub parent: Option<BlockId>,
108    pub children: usize,
109    pub outgoing_edges: usize,
110    pub incoming_edges: usize,
111}
112
113#[derive(Debug, Clone, Serialize, Deserialize)]
114pub struct GraphExportEdge {
115    pub source: BlockId,
116    pub target: BlockId,
117    pub relation: String,
118    pub direction: String,
119}
120
121#[derive(Debug, Clone, Serialize, Deserialize)]
122pub struct GraphExport {
123    pub summary: GraphSessionSummary,
124    #[serde(default)]
125    pub nodes: Vec<GraphExportNode>,
126    #[serde(default)]
127    pub edges: Vec<GraphExportEdge>,
128}
129
130#[derive(Debug, Clone)]
131pub struct GraphSession {
132    graph: GraphNavigator,
133    selected: HashMap<BlockId, GraphSessionNode>,
134    focus: Option<BlockId>,
135    history: Vec<String>,
136}
137
138impl GraphSession {
139    pub fn new(graph: GraphNavigator) -> Self {
140        Self {
141            graph,
142            selected: HashMap::new(),
143            focus: None,
144            history: Vec::new(),
145        }
146    }
147
148    pub fn selected_block_ids(&self) -> Vec<BlockId> {
149        let mut ids = self.selected.keys().copied().collect::<Vec<_>>();
150        ids.sort_by_key(|id| id.to_string());
151        ids
152    }
153
154    pub fn summary(&self) -> GraphSessionSummary {
155        let roots = self
156            .selected
157            .keys()
158            .filter(|id| {
159                self.graph
160                    .describe_node(**id)
161                    .map(|node| node.parent.is_none())
162                    .unwrap_or(false)
163            })
164            .count();
165        let leaves = self
166            .selected
167            .keys()
168            .filter(|id| {
169                self.graph
170                    .describe_node(**id)
171                    .map(|node| node.children == 0)
172                    .unwrap_or(false)
173            })
174            .count();
175        GraphSessionSummary {
176            selected: self.selected.len(),
177            pinned: self.selected.values().filter(|node| node.pinned).count(),
178            focused: self.focus.is_some(),
179            roots,
180            leaves,
181        }
182    }
183
184    pub fn fork(&self) -> Self {
185        self.clone()
186    }
187
188    pub fn seed_overview(&mut self, max_depth: Option<usize>) -> GraphSessionUpdate {
189        let max_depth = max_depth.unwrap_or(2).max(1);
190        let mut update = GraphSessionUpdate::default();
191        let root = self.graph.root_id();
192        let mut queue = VecDeque::from([(root, 0usize)]);
193        while let Some((current, depth)) = queue.pop_front() {
194            self.select_id(
195                current,
196                GraphDetailLevel::Summary,
197                Some(GraphSelectionOrigin {
198                    kind: GraphSelectionOriginKind::Overview,
199                    relation: None,
200                    anchor: None,
201                }),
202                &mut update,
203            );
204            if depth < max_depth {
205                for child in self.graph.neighbors(current, GraphNeighborMode::Children) {
206                    queue.push_back((child.to, depth + 1));
207                }
208            }
209        }
210        update.focus = self.focus;
211        self.history
212            .push(format!("seed_overview depth={max_depth}"));
213        update
214    }
215
216    pub fn focus(&mut self, selector: Option<&str>) -> Result<GraphSessionUpdate, GraphStoreError> {
217        self.focus = selector.and_then(|value| self.graph.resolve_selector(value));
218        Ok(GraphSessionUpdate {
219            focus: self.focus,
220            ..GraphSessionUpdate::default()
221        })
222    }
223
224    pub fn select(
225        &mut self,
226        selector: &str,
227        detail_level: GraphDetailLevel,
228    ) -> Result<GraphSessionUpdate, GraphStoreError> {
229        let block_id = self
230            .graph
231            .resolve_selector(selector)
232            .ok_or_else(|| GraphStoreError::GraphNotFound(selector.to_string()))?;
233        let mut update = GraphSessionUpdate::default();
234        self.select_id(
235            block_id,
236            detail_level,
237            Some(GraphSelectionOrigin {
238                kind: GraphSelectionOriginKind::Manual,
239                relation: None,
240                anchor: None,
241            }),
242            &mut update,
243        );
244        self.history.push(format!("select {selector}"));
245        Ok(update)
246    }
247
248    pub fn expand(
249        &mut self,
250        selector: &str,
251        mode: GraphNeighborMode,
252        depth: usize,
253        max_add: Option<usize>,
254    ) -> Result<GraphSessionUpdate, GraphStoreError> {
255        let start = self
256            .graph
257            .resolve_selector(selector)
258            .ok_or_else(|| GraphStoreError::GraphNotFound(selector.to_string()))?;
259        let mut update = GraphSessionUpdate::default();
260        let mut queue = VecDeque::from([(start, 0usize)]);
261        let mut seen = HashSet::from([start]);
262        let mut added = 0usize;
263
264        while let Some((current, current_depth)) = queue.pop_front() {
265            if current_depth >= depth.max(1) {
266                continue;
267            }
268            for neighbor in self.graph.neighbors(current, mode) {
269                if !seen.insert(neighbor.to) {
270                    continue;
271                }
272                if max_add.map(|limit| added >= limit).unwrap_or(false) {
273                    update.warnings.push(format!(
274                        "Stopped expansion after reaching max_add={}",
275                        max_add.unwrap_or_default()
276                    ));
277                    update.focus = self.focus;
278                    return Ok(update);
279                }
280                self.select_id(
281                    neighbor.to,
282                    GraphDetailLevel::Summary,
283                    Some(origin_for(mode, &neighbor, current)),
284                    &mut update,
285                );
286                queue.push_back((neighbor.to, current_depth + 1));
287                added += 1;
288            }
289        }
290
291        self.history
292            .push(format!("expand {selector} mode={mode:?} depth={depth}"));
293        update.focus = self.focus;
294        Ok(update)
295    }
296
297    pub fn collapse(
298        &mut self,
299        selector: &str,
300        include_descendants: bool,
301    ) -> Result<GraphSessionUpdate, GraphStoreError> {
302        let start = self
303            .graph
304            .resolve_selector(selector)
305            .ok_or_else(|| GraphStoreError::GraphNotFound(selector.to_string()))?;
306        let mut update = GraphSessionUpdate::default();
307        let mut remove = vec![start];
308        if include_descendants {
309            let mut queue = VecDeque::from([start]);
310            while let Some(current) = queue.pop_front() {
311                for child in self.graph.neighbors(current, GraphNeighborMode::Children) {
312                    remove.push(child.to);
313                    queue.push_back(child.to);
314                }
315            }
316        }
317        for block_id in remove {
318            if self.selected.remove(&block_id).is_some() {
319                update.removed.push(block_id);
320            }
321        }
322        if self
323            .focus
324            .map(|id| update.removed.contains(&id))
325            .unwrap_or(false)
326        {
327            self.focus = None;
328        }
329        update.focus = self.focus;
330        Ok(update)
331    }
332
333    pub fn pin(
334        &mut self,
335        selector: &str,
336        pinned: bool,
337    ) -> Result<GraphSessionUpdate, GraphStoreError> {
338        let block_id = self
339            .graph
340            .resolve_selector(selector)
341            .ok_or_else(|| GraphStoreError::GraphNotFound(selector.to_string()))?;
342        let mut update = GraphSessionUpdate::default();
343        if let Some(node) = self.selected.get_mut(&block_id) {
344            node.pinned = pinned;
345            update.changed.push(block_id);
346        }
347        update.focus = self.focus;
348        Ok(update)
349    }
350
351    pub fn prune(&mut self, max_selected: Option<usize>) -> GraphSessionUpdate {
352        let limit = max_selected.unwrap_or(32).max(1);
353        let mut update = GraphSessionUpdate::default();
354        if self.selected.len() <= limit {
355            update.focus = self.focus;
356            return update;
357        }
358
359        let mut candidates = self
360            .selected
361            .iter()
362            .filter(|(id, node)| !node.pinned && Some(**id) != self.focus)
363            .map(|(id, _)| *id)
364            .collect::<Vec<_>>();
365        candidates.sort_by_key(|id| id.to_string());
366
367        while self.selected.len() > limit {
368            let Some(block_id) = candidates.pop() else {
369                break;
370            };
371            if self.selected.remove(&block_id).is_some() {
372                update.removed.push(block_id);
373            }
374        }
375        update.focus = self.focus;
376        update
377    }
378
379    pub fn export(&self) -> GraphExport {
380        let mut nodes = self
381            .selected
382            .iter()
383            .filter_map(|(block_id, node)| {
384                self.graph
385                    .describe_node(*block_id)
386                    .map(|summary| GraphExportNode {
387                        block_id: *block_id,
388                        label: summary.label,
389                        content_type: summary.content_type,
390                        semantic_role: summary.semantic_role,
391                        tags: summary.tags,
392                        detail_level: node.detail_level,
393                        pinned: node.pinned,
394                        parent: summary.parent,
395                        children: summary.children,
396                        outgoing_edges: summary.outgoing_edges,
397                        incoming_edges: summary.incoming_edges,
398                    })
399            })
400            .collect::<Vec<_>>();
401        nodes.sort_by(|left, right| left.label.cmp(&right.label));
402
403        let selected_ids = self.selected.keys().copied().collect::<HashSet<_>>();
404        let mut seen_edges = HashSet::new();
405        let mut edges = Vec::new();
406        for block_id in &selected_ids {
407            for edge in self
408                .graph
409                .neighbors(*block_id, GraphNeighborMode::Neighborhood)
410            {
411                if !selected_ids.contains(&edge.to) {
412                    continue;
413                }
414                let key = (
415                    edge.from,
416                    edge.to,
417                    edge.relation.clone(),
418                    edge.direction.clone(),
419                );
420                if seen_edges.insert(key.clone()) {
421                    edges.push(GraphExportEdge {
422                        source: edge.from,
423                        target: edge.to,
424                        relation: edge.relation,
425                        direction: edge.direction,
426                    });
427                }
428            }
429        }
430        edges.sort_by(|left, right| {
431            left.relation
432                .cmp(&right.relation)
433                .then(left.source.to_string().cmp(&right.source.to_string()))
434                .then(left.target.to_string().cmp(&right.target.to_string()))
435        });
436
437        GraphExport {
438            summary: self.summary(),
439            nodes,
440            edges,
441        }
442    }
443
444    pub fn why_selected(
445        &self,
446        selector: &str,
447    ) -> Result<GraphSelectionExplanation, GraphStoreError> {
448        let block_id = self
449            .graph
450            .resolve_selector(selector)
451            .ok_or_else(|| GraphStoreError::GraphNotFound(selector.to_string()))?;
452        let node = self.graph.describe_node(block_id);
453        let Some(selected) = self.selected.get(&block_id) else {
454            return Ok(GraphSelectionExplanation {
455                block_id,
456                selected: false,
457                focus: self.focus == Some(block_id),
458                pinned: false,
459                detail_level: None,
460                origin: None,
461                explanation: "Node is not currently selected in the session.".to_string(),
462                node,
463                anchor: None,
464            });
465        };
466        let anchor = selected
467            .origin
468            .as_ref()
469            .and_then(|origin| origin.anchor)
470            .and_then(|id| self.graph.describe_node(id));
471        let explanation = match selected.origin.as_ref().map(|origin| origin.kind) {
472            Some(GraphSelectionOriginKind::Overview) => {
473                "Node was selected as part of the overview scaffold.".to_string()
474            }
475            Some(GraphSelectionOriginKind::Manual) => {
476                "Node was selected directly by the agent.".to_string()
477            }
478            Some(GraphSelectionOriginKind::Children) => {
479                "Node was selected while expanding child relationships.".to_string()
480            }
481            Some(GraphSelectionOriginKind::Parents) => {
482                "Node was selected while traversing toward parent relationships.".to_string()
483            }
484            Some(GraphSelectionOriginKind::Outgoing) => {
485                "Node was selected while following outgoing semantic edges.".to_string()
486            }
487            Some(GraphSelectionOriginKind::Incoming) => {
488                "Node was selected while following incoming semantic edges.".to_string()
489            }
490            None => "Node is selected in the session.".to_string(),
491        };
492        Ok(GraphSelectionExplanation {
493            block_id,
494            selected: true,
495            focus: self.focus == Some(block_id),
496            pinned: selected.pinned,
497            detail_level: Some(selected.detail_level),
498            origin: selected.origin.clone(),
499            explanation,
500            node,
501            anchor,
502        })
503    }
504
505    pub fn diff(&self, other: &Self) -> GraphSessionDiff {
506        let before = self.selected.keys().copied().collect::<HashSet<_>>();
507        let after = other.selected.keys().copied().collect::<HashSet<_>>();
508        let mut added = after
509            .difference(&before)
510            .copied()
511            .filter_map(|id| other.graph.describe_node(id))
512            .collect::<Vec<_>>();
513        let mut removed = before
514            .difference(&after)
515            .copied()
516            .filter_map(|id| self.graph.describe_node(id))
517            .collect::<Vec<_>>();
518        added.sort_by(|left, right| left.label.cmp(&right.label));
519        removed.sort_by(|left, right| left.label.cmp(&right.label));
520        GraphSessionDiff {
521            added,
522            removed,
523            focus_before: self.focus,
524            focus_after: other.focus,
525            changed_focus: self.focus != other.focus,
526        }
527    }
528
529    fn select_id(
530        &mut self,
531        block_id: BlockId,
532        detail_level: GraphDetailLevel,
533        origin: Option<GraphSelectionOrigin>,
534        update: &mut GraphSessionUpdate,
535    ) {
536        match self.selected.get_mut(&block_id) {
537            Some(node) => {
538                if node.detail_level < detail_level {
539                    node.detail_level = detail_level;
540                    update.changed.push(block_id);
541                }
542                if node.origin.is_none() {
543                    node.origin = origin;
544                }
545            }
546            None => {
547                self.selected.insert(
548                    block_id,
549                    GraphSessionNode {
550                        detail_level,
551                        pinned: false,
552                        origin,
553                    },
554                );
555                update.added.push(block_id);
556            }
557        }
558    }
559}
560
561fn origin_for(
562    mode: GraphNeighborMode,
563    hop: &crate::types::GraphPathHop,
564    anchor: BlockId,
565) -> GraphSelectionOrigin {
566    GraphSelectionOrigin {
567        kind: match mode {
568            GraphNeighborMode::Children => GraphSelectionOriginKind::Children,
569            GraphNeighborMode::Parents => GraphSelectionOriginKind::Parents,
570            GraphNeighborMode::Outgoing => GraphSelectionOriginKind::Outgoing,
571            GraphNeighborMode::Incoming => GraphSelectionOriginKind::Incoming,
572            GraphNeighborMode::Neighborhood => {
573                if hop.direction == "incoming" {
574                    GraphSelectionOriginKind::Incoming
575                } else if hop.relation == "parent" {
576                    GraphSelectionOriginKind::Parents
577                } else if hop.relation == "contains" {
578                    GraphSelectionOriginKind::Children
579                } else {
580                    GraphSelectionOriginKind::Outgoing
581                }
582            }
583        },
584        relation: Some(hop.relation.clone()),
585        anchor: Some(anchor),
586    }
587}
588
589#[allow(dead_code)]
590fn _edge_to_export(edge: GraphEdgeSummary) -> GraphExportEdge {
591    GraphExportEdge {
592        source: edge.source,
593        target: edge.target,
594        relation: edge.relation,
595        direction: edge.direction,
596    }
597}