ucp_codegraph/programmatic/
session.rs

1use std::collections::HashSet;
2
3use anyhow::Result;
4use ucm_core::BlockId;
5
6use crate::{
7    export_codegraph_context_with_config, render_codegraph_context_prompt, CodeGraphContextExport,
8    CodeGraphContextFrontierAction, CodeGraphContextSession, CodeGraphContextSummary,
9    CodeGraphContextUpdate, CodeGraphDetailLevel, CodeGraphExportConfig, CodeGraphRenderConfig,
10    CodeGraphSelectionOriginKind, CodeGraphTraversalConfig,
11};
12
13use super::{
14    query,
15    types::{
16        CodeGraphExpandMode, CodeGraphFindQuery, CodeGraphNodeSummary,
17        CodeGraphRecommendedActionsResult, CodeGraphSelectionExplanation, CodeGraphSessionDiff,
18    },
19    CodeGraphNavigator,
20};
21
22#[derive(Debug, Clone)]
23pub struct CodeGraphNavigatorSession {
24    graph: CodeGraphNavigator,
25    context: CodeGraphContextSession,
26}
27
28impl CodeGraphNavigatorSession {
29    pub fn new(graph: CodeGraphNavigator) -> Self {
30        Self {
31            graph,
32            context: CodeGraphContextSession::new(),
33        }
34    }
35
36    pub fn context(&self) -> &CodeGraphContextSession {
37        &self.context
38    }
39
40    pub fn selected_block_ids(&self) -> Vec<BlockId> {
41        let mut ids = self.context.selected.keys().copied().collect::<Vec<_>>();
42        ids.sort_by_key(|value| value.to_string());
43        ids
44    }
45
46    pub fn summary(&self) -> CodeGraphContextSummary {
47        self.context.summary(self.graph.document())
48    }
49
50    pub fn fork(&self) -> Self {
51        self.clone()
52    }
53
54    pub fn seed_overview(&mut self, max_depth: Option<usize>) -> CodeGraphContextUpdate {
55        self.context
56            .seed_overview_with_depth(self.graph.document(), max_depth)
57    }
58
59    pub fn focus(&mut self, selector: Option<&str>) -> Result<CodeGraphContextUpdate> {
60        let block_id = selector
61            .map(|value| self.graph.resolve_required(value))
62            .transpose()?;
63        Ok(self.context.set_focus(self.graph.document(), block_id))
64    }
65
66    pub fn select(
67        &mut self,
68        selector: &str,
69        detail_level: CodeGraphDetailLevel,
70    ) -> Result<CodeGraphContextUpdate> {
71        let block_id = self.graph.resolve_required(selector)?;
72        Ok(self
73            .context
74            .select_block(self.graph.document(), block_id, detail_level))
75    }
76
77    pub fn expand(
78        &mut self,
79        selector: &str,
80        mode: CodeGraphExpandMode,
81        traversal: &CodeGraphTraversalConfig,
82    ) -> Result<CodeGraphContextUpdate> {
83        let block_id = self.graph.resolve_required(selector)?;
84        Ok(match mode {
85            CodeGraphExpandMode::File => {
86                self.context
87                    .expand_file_with_config(self.graph.document(), block_id, traversal)
88            }
89            CodeGraphExpandMode::Dependencies => self.context.expand_dependencies_with_config(
90                self.graph.document(),
91                block_id,
92                traversal,
93            ),
94            CodeGraphExpandMode::Dependents => self.context.expand_dependents_with_config(
95                self.graph.document(),
96                block_id,
97                traversal,
98            ),
99        })
100    }
101
102    pub fn hydrate_source(
103        &mut self,
104        selector: &str,
105        padding: usize,
106    ) -> Result<CodeGraphContextUpdate> {
107        let block_id = self.graph.resolve_required(selector)?;
108        Ok(self
109            .context
110            .hydrate_source(self.graph.document(), block_id, padding))
111    }
112
113    pub fn collapse(
114        &mut self,
115        selector: &str,
116        include_descendants: bool,
117    ) -> Result<CodeGraphContextUpdate> {
118        let block_id = self.graph.resolve_required(selector)?;
119        Ok(self
120            .context
121            .collapse(self.graph.document(), block_id, include_descendants))
122    }
123
124    pub fn pin(&mut self, selector: &str, pinned: bool) -> Result<CodeGraphContextUpdate> {
125        let block_id = self.graph.resolve_required(selector)?;
126        Ok(self.context.pin(block_id, pinned))
127    }
128
129    pub fn prune(&mut self, max_selected: Option<usize>) -> CodeGraphContextUpdate {
130        self.context.prune(self.graph.document(), max_selected)
131    }
132
133    pub fn export(
134        &self,
135        render: &CodeGraphRenderConfig,
136        export: &CodeGraphExportConfig,
137    ) -> CodeGraphContextExport {
138        export_codegraph_context_with_config(self.graph.document(), &self.context, render, export)
139    }
140
141    pub fn render_prompt(&self, render: &CodeGraphRenderConfig) -> String {
142        render_codegraph_context_prompt(self.graph.document(), &self.context, render)
143    }
144
145    pub fn find_nodes(&self, query: &CodeGraphFindQuery) -> Result<Vec<CodeGraphNodeSummary>> {
146        self.graph.find_nodes(query)
147    }
148
149    pub fn why_selected(&self, selector: &str) -> Result<CodeGraphSelectionExplanation> {
150        let block_id = self.graph.resolve_required(selector)?;
151        let node = self.graph.describe_node(block_id);
152        let Some(selected) = self.context.selected.get(&block_id) else {
153            return Ok(CodeGraphSelectionExplanation {
154                block_id,
155                selected: false,
156                focus: self.context.focus == Some(block_id),
157                pinned: false,
158                detail_level: None,
159                origin: None,
160                explanation: "Node is not currently selected in the session.".to_string(),
161                node,
162                anchor: None,
163            });
164        };
165
166        let anchor = selected
167            .origin
168            .as_ref()
169            .and_then(|origin| origin.anchor)
170            .and_then(|id| self.graph.describe_node(id));
171        let explanation = match selected.origin.as_ref().map(|origin| origin.kind) {
172            Some(CodeGraphSelectionOriginKind::Manual) => {
173                "Node was selected directly by the agent.".to_string()
174            }
175            Some(CodeGraphSelectionOriginKind::Overview) => {
176                "Node was selected as part of the overview scaffold.".to_string()
177            }
178            Some(CodeGraphSelectionOriginKind::FileSymbols) => {
179                "Node was selected while expanding file symbols.".to_string()
180            }
181            Some(CodeGraphSelectionOriginKind::Dependencies) => format!(
182                "Node was selected while following dependency edges{}.",
183                relation_suffix(selected.origin.as_ref())
184            ),
185            Some(CodeGraphSelectionOriginKind::Dependents) => format!(
186                "Node was selected while following dependent edges{}.",
187                relation_suffix(selected.origin.as_ref())
188            ),
189            None => "Node is selected in the session.".to_string(),
190        };
191
192        Ok(CodeGraphSelectionExplanation {
193            block_id,
194            selected: true,
195            focus: self.context.focus == Some(block_id),
196            pinned: selected.pinned,
197            detail_level: Some(selected.detail_level),
198            origin: selected.origin.clone(),
199            explanation,
200            node,
201            anchor,
202        })
203    }
204
205    pub fn diff(&self, other: &Self) -> CodeGraphSessionDiff {
206        let before = self
207            .context
208            .selected
209            .keys()
210            .copied()
211            .collect::<HashSet<_>>();
212        let after = other
213            .context
214            .selected
215            .keys()
216            .copied()
217            .collect::<HashSet<_>>();
218        let mut added = after
219            .difference(&before)
220            .copied()
221            .filter_map(|id| other.graph.describe_node(id))
222            .collect::<Vec<_>>();
223        let mut removed = before
224            .difference(&after)
225            .copied()
226            .filter_map(|id| self.graph.describe_node(id))
227            .collect::<Vec<_>>();
228        added.sort_by(|left, right| left.label.cmp(&right.label));
229        removed.sort_by(|left, right| left.label.cmp(&right.label));
230        CodeGraphSessionDiff {
231            added,
232            removed,
233            focus_before: self.context.focus,
234            focus_after: other.context.focus,
235            changed_focus: self.context.focus != other.context.focus,
236        }
237    }
238
239    pub fn apply_recommended_actions(
240        &mut self,
241        top: usize,
242        padding: usize,
243        depth: Option<usize>,
244        max_add: Option<usize>,
245        priority_threshold: Option<u16>,
246    ) -> Result<CodeGraphRecommendedActionsResult> {
247        let export_config = CodeGraphExportConfig {
248            max_frontier_actions: top.max(1).max(8),
249            ..Default::default()
250        };
251        let actions = self
252            .export(&CodeGraphRenderConfig::default(), &export_config)
253            .frontier
254            .into_iter()
255            .filter(|action| action.candidate_count > 0)
256            .filter(|action| {
257                priority_threshold
258                    .map(|threshold| action.priority >= threshold)
259                    .unwrap_or(true)
260            })
261            .take(top.max(1))
262            .collect::<Vec<_>>();
263        if actions.is_empty() {
264            return Err(anyhow::anyhow!(
265                "No recommended actions available for the current focus"
266            ));
267        }
268
269        let mut update = CodeGraphContextUpdate::default();
270        let mut applied_actions = Vec::new();
271        for action in actions {
272            let traversal = CodeGraphTraversalConfig {
273                depth: depth.unwrap_or(1),
274                relation_filters: action.relation.clone().into_iter().collect(),
275                max_add,
276                priority_threshold,
277            };
278            applied_actions.push(action_summary(&action));
279            merge_update(
280                &mut update,
281                match action.action.as_str() {
282                    "hydrate_source" => {
283                        self.context
284                            .hydrate_source(self.graph.document(), action.block_id, padding)
285                    }
286                    "expand_file" => self.context.expand_file_with_config(
287                        self.graph.document(),
288                        action.block_id,
289                        &traversal,
290                    ),
291                    "expand_dependencies" => self.context.expand_dependencies_with_config(
292                        self.graph.document(),
293                        action.block_id,
294                        &traversal,
295                    ),
296                    "expand_dependents" => self.context.expand_dependents_with_config(
297                        self.graph.document(),
298                        action.block_id,
299                        &traversal,
300                    ),
301                    "collapse" => {
302                        self.context
303                            .collapse(self.graph.document(), action.block_id, false)
304                    }
305                    _ => CodeGraphContextUpdate::default(),
306                },
307            );
308        }
309
310        Ok(CodeGraphRecommendedActionsResult {
311            applied_actions,
312            update,
313        })
314    }
315
316    pub fn path_between(
317        &self,
318        start_selector: &str,
319        end_selector: &str,
320        max_hops: usize,
321    ) -> Result<Option<crate::programmatic::types::CodeGraphPathResult>> {
322        let start = self.graph.resolve_required(start_selector)?;
323        let end = self.graph.resolve_required(end_selector)?;
324        Ok(query::path_between(
325            self.graph.document(),
326            start,
327            end,
328            max_hops,
329        ))
330    }
331}
332
333fn merge_update(into: &mut CodeGraphContextUpdate, next: CodeGraphContextUpdate) {
334    into.added.extend(next.added);
335    into.removed.extend(next.removed);
336    into.changed.extend(next.changed);
337    into.warnings.extend(next.warnings);
338    if next.focus.is_some() {
339        into.focus = next.focus;
340    }
341}
342
343fn action_summary(action: &CodeGraphContextFrontierAction) -> String {
344    match action.relation.as_deref() {
345        Some(relation) => format!("{} {} via {}", action.action, action.short_id, relation),
346        None => format!("{} {}", action.action, action.short_id),
347    }
348}
349
350fn relation_suffix(origin: Option<&crate::CodeGraphSelectionOrigin>) -> String {
351    origin
352        .and_then(|value| value.relation.as_deref())
353        .map(|relation| format!(" via `{}`", relation))
354        .unwrap_or_default()
355}
ucp_codegraph/programmatic/session.rs

ucp_codegraph/programmatic/
session.rs
