1use std::collections::{BTreeMap, HashMap, HashSet, VecDeque};
2use std::fmt::Write;
3use std::path::PathBuf;
4use std::str::FromStr;
5
6use serde::{Deserialize, Serialize};
7use serde_json::Value;
8use ucm_core::{Block, BlockId, Content, Document};
9
10use crate::model::{
11 CODEGRAPH_PROFILE_MARKER, META_CODEREF, META_EXPORTED, META_LANGUAGE, META_LOGICAL_KEY,
12 META_NODE_CLASS, META_SYMBOL_NAME,
13};
14
15#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize, Default)]
16#[serde(rename_all = "snake_case")]
17pub enum CodeGraphDetailLevel {
18 #[default]
19 Skeleton,
20 SymbolCard,
21 Neighborhood,
22 Source,
23}
24
25impl CodeGraphDetailLevel {
26 fn max(self, other: Self) -> Self {
27 std::cmp::max(self, other)
28 }
29
30 fn demoted(self) -> Self {
31 match self {
32 Self::Source => Self::Neighborhood,
33 Self::Neighborhood => Self::SymbolCard,
34 Self::SymbolCard => Self::Skeleton,
35 Self::Skeleton => Self::Skeleton,
36 }
37 }
38
39 fn includes_neighborhood(self) -> bool {
40 matches!(self, Self::Neighborhood | Self::Source)
41 }
42
43 fn includes_source(self) -> bool {
44 matches!(self, Self::Source)
45 }
46}
47
48fn default_true() -> bool {
49 true
50}
51
52fn default_relation_prune_priority() -> BTreeMap<String, u8> {
53 [
54 ("references", 60),
55 ("cited_by", 60),
56 ("links_to", 55),
57 ("uses_symbol", 35),
58 ("imports_symbol", 30),
59 ("reexports_symbol", 25),
60 ("calls", 20),
61 ("inherits", 15),
62 ("implements", 15),
63 ]
64 .into_iter()
65 .map(|(name, score)| (name.to_string(), score))
66 .collect()
67}
68
69fn selection_origin(
70 kind: CodeGraphSelectionOriginKind,
71 relation: Option<&str>,
72 anchor: Option<BlockId>,
73) -> Option<CodeGraphSelectionOrigin> {
74 Some(CodeGraphSelectionOrigin {
75 kind,
76 relation: relation.map(str::to_string),
77 anchor,
78 })
79}
80
81#[derive(Debug, Clone, Serialize, Deserialize)]
82pub struct HydratedSourceExcerpt {
83 pub path: String,
84 pub display: String,
85 pub start_line: usize,
86 pub end_line: usize,
87 pub snippet: String,
88}
89
90#[derive(Debug, Clone, Serialize, Deserialize)]
91pub struct CodeGraphContextNode {
92 pub block_id: BlockId,
93 pub detail_level: CodeGraphDetailLevel,
94 #[serde(default)]
95 pub pinned: bool,
96 #[serde(default, skip_serializing_if = "Option::is_none")]
97 pub origin: Option<CodeGraphSelectionOrigin>,
98 #[serde(default, skip_serializing_if = "Option::is_none")]
99 pub hydrated_source: Option<HydratedSourceExcerpt>,
100}
101
102#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
103#[serde(rename_all = "snake_case")]
104pub enum CodeGraphSelectionOriginKind {
105 Overview,
106 Manual,
107 FileSymbols,
108 Dependencies,
109 Dependents,
110}
111
112#[derive(Debug, Clone, Serialize, Deserialize)]
113pub struct CodeGraphSelectionOrigin {
114 pub kind: CodeGraphSelectionOriginKind,
115 #[serde(default, skip_serializing_if = "Option::is_none")]
116 pub relation: Option<String>,
117 #[serde(default, skip_serializing_if = "Option::is_none")]
118 pub anchor: Option<BlockId>,
119}
120
121#[derive(Debug, Clone, Serialize, Deserialize)]
122pub struct CodeGraphPrunePolicy {
123 pub max_selected: usize,
124 #[serde(default = "default_true")]
125 pub demote_before_remove: bool,
126 #[serde(default = "default_true")]
127 pub protect_focus: bool,
128 #[serde(default = "default_relation_prune_priority")]
129 pub relation_prune_priority: BTreeMap<String, u8>,
130}
131
132impl Default for CodeGraphPrunePolicy {
133 fn default() -> Self {
134 Self {
135 max_selected: 48,
136 demote_before_remove: true,
137 protect_focus: true,
138 relation_prune_priority: default_relation_prune_priority(),
139 }
140 }
141}
142
143#[derive(Debug, Clone, Default, Serialize, Deserialize)]
144pub struct CodeGraphContextSession {
145 #[serde(default)]
146 pub selected: HashMap<BlockId, CodeGraphContextNode>,
147 #[serde(default, skip_serializing_if = "Option::is_none")]
148 pub focus: Option<BlockId>,
149 #[serde(default)]
150 pub prune_policy: CodeGraphPrunePolicy,
151 #[serde(default)]
152 pub history: Vec<String>,
153}
154
155#[derive(Debug, Clone, Default, Serialize, Deserialize)]
156pub struct CodeGraphContextUpdate {
157 #[serde(default)]
158 pub added: Vec<BlockId>,
159 #[serde(default)]
160 pub removed: Vec<BlockId>,
161 #[serde(default)]
162 pub changed: Vec<BlockId>,
163 #[serde(default)]
164 pub focus: Option<BlockId>,
165 #[serde(default)]
166 pub warnings: Vec<String>,
167}
168
169#[derive(Debug, Clone, Serialize, Deserialize)]
170pub struct CodeGraphContextSummary {
171 pub selected: usize,
172 pub max_selected: usize,
173 pub repositories: usize,
174 pub directories: usize,
175 pub files: usize,
176 pub symbols: usize,
177 pub hydrated_sources: usize,
178}
179
180#[derive(Debug, Clone, Serialize, Deserialize)]
181pub struct CodeGraphRenderConfig {
182 pub max_edges_per_node: usize,
183 pub max_source_lines: usize,
184}
185
186#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
187#[serde(rename_all = "snake_case")]
188pub enum CodeGraphExportMode {
189 #[default]
190 Full,
191 Compact,
192}
193
194#[derive(Debug, Clone, Serialize, Deserialize)]
195pub struct CodeGraphExportConfig {
196 #[serde(default)]
197 pub mode: CodeGraphExportMode,
198 #[serde(default = "default_true")]
199 pub include_rendered: bool,
200 #[serde(default = "default_true")]
201 pub dedupe_edges: bool,
202 #[serde(default, skip_serializing_if = "Option::is_none")]
203 pub visible_levels: Option<usize>,
204 #[serde(default, skip_serializing_if = "Vec::is_empty")]
205 pub only_node_classes: Vec<String>,
206 #[serde(default, skip_serializing_if = "Vec::is_empty")]
207 pub exclude_node_classes: Vec<String>,
208 pub max_frontier_actions: usize,
209}
210
211#[derive(Debug, Clone, Serialize, Deserialize)]
212pub struct CodeGraphTraversalConfig {
213 #[serde(default = "default_one")]
214 pub depth: usize,
215 #[serde(default, skip_serializing_if = "Vec::is_empty")]
216 pub relation_filters: Vec<String>,
217 #[serde(default, skip_serializing_if = "Option::is_none")]
218 pub max_add: Option<usize>,
219 #[serde(default, skip_serializing_if = "Option::is_none")]
220 pub priority_threshold: Option<u16>,
221}
222
223impl Default for CodeGraphExportConfig {
224 fn default() -> Self {
225 Self {
226 mode: CodeGraphExportMode::Full,
227 include_rendered: true,
228 dedupe_edges: true,
229 visible_levels: None,
230 only_node_classes: Vec::new(),
231 exclude_node_classes: Vec::new(),
232 max_frontier_actions: 12,
233 }
234 }
235}
236
237impl Default for CodeGraphTraversalConfig {
238 fn default() -> Self {
239 Self {
240 depth: 1,
241 relation_filters: Vec::new(),
242 max_add: None,
243 priority_threshold: None,
244 }
245 }
246}
247
248impl CodeGraphExportConfig {
249 pub fn compact() -> Self {
250 Self {
251 mode: CodeGraphExportMode::Compact,
252 include_rendered: false,
253 dedupe_edges: true,
254 visible_levels: None,
255 only_node_classes: Vec::new(),
256 exclude_node_classes: Vec::new(),
257 max_frontier_actions: 6,
258 }
259 }
260}
261
262impl CodeGraphTraversalConfig {
263 fn depth(&self) -> usize {
264 self.depth.max(1)
265 }
266
267 fn relation_filter_set(&self) -> Option<HashSet<String>> {
268 if self.relation_filters.is_empty() {
269 None
270 } else {
271 Some(self.relation_filters.iter().cloned().collect())
272 }
273 }
274}
275
276#[derive(Debug, Clone, Serialize, Deserialize)]
277pub struct CodeGraphCoderef {
278 pub path: String,
279 pub display: String,
280 #[serde(default, skip_serializing_if = "Option::is_none")]
281 pub start_line: Option<usize>,
282 #[serde(default, skip_serializing_if = "Option::is_none")]
283 pub end_line: Option<usize>,
284}
285
286#[derive(Debug, Clone, Serialize, Deserialize)]
287pub struct CodeGraphContextNodeExport {
288 pub block_id: BlockId,
289 pub short_id: String,
290 pub node_class: String,
291 pub label: String,
292 pub detail_level: CodeGraphDetailLevel,
293 pub pinned: bool,
294 #[serde(default, skip_serializing_if = "Option::is_none")]
295 pub distance_from_focus: Option<usize>,
296 pub relevance_score: u16,
297 #[serde(default, skip_serializing_if = "Option::is_none")]
298 pub logical_key: Option<String>,
299 #[serde(default, skip_serializing_if = "Option::is_none")]
300 pub symbol_name: Option<String>,
301 #[serde(default, skip_serializing_if = "Option::is_none")]
302 pub path: Option<String>,
303 #[serde(default, skip_serializing_if = "Option::is_none")]
304 pub signature: Option<String>,
305 #[serde(default, skip_serializing_if = "Option::is_none")]
306 pub docs: Option<String>,
307 #[serde(default, skip_serializing_if = "Option::is_none")]
308 pub origin: Option<CodeGraphSelectionOrigin>,
309 #[serde(default, skip_serializing_if = "Option::is_none")]
310 pub coderef: Option<CodeGraphCoderef>,
311 #[serde(default, skip_serializing_if = "Option::is_none")]
312 pub hydrated_source: Option<HydratedSourceExcerpt>,
313}
314
315#[derive(Debug, Clone, Serialize, Deserialize)]
316pub struct CodeGraphContextEdgeExport {
317 pub source: BlockId,
318 pub source_short_id: String,
319 pub target: BlockId,
320 pub target_short_id: String,
321 pub relation: String,
322 #[serde(default = "default_one")]
323 pub multiplicity: usize,
324}
325
326#[derive(Debug, Clone, Serialize, Deserialize)]
327pub struct CodeGraphContextFrontierAction {
328 pub block_id: BlockId,
329 pub short_id: String,
330 pub action: String,
331 #[serde(default, skip_serializing_if = "Option::is_none")]
332 pub relation: Option<String>,
333 #[serde(default, skip_serializing_if = "Option::is_none")]
334 pub direction: Option<String>,
335 pub candidate_count: usize,
336 pub priority: u16,
337 pub description: String,
338}
339
340#[derive(Debug, Clone, Serialize, Deserialize)]
341pub struct CodeGraphContextHeuristics {
342 pub should_stop: bool,
343 #[serde(default, skip_serializing_if = "Vec::is_empty")]
344 pub reasons: Vec<String>,
345 pub hidden_candidate_count: usize,
346 pub low_value_candidate_count: usize,
347 #[serde(default, skip_serializing_if = "Option::is_none")]
348 pub recommended_next_action: Option<CodeGraphContextFrontierAction>,
349 #[serde(default, skip_serializing_if = "Vec::is_empty")]
350 pub recommended_actions: Vec<CodeGraphContextFrontierAction>,
351}
352
353#[derive(Debug, Clone, Serialize, Deserialize)]
354pub struct CodeGraphHiddenLevelSummary {
355 pub level: usize,
356 pub count: usize,
357 #[serde(default, skip_serializing_if = "Option::is_none")]
358 pub relation: Option<String>,
359 #[serde(default, skip_serializing_if = "Option::is_none")]
360 pub direction: Option<String>,
361}
362
363#[derive(Debug, Clone, Serialize, Deserialize)]
364pub struct CodeGraphContextExport {
365 pub summary: CodeGraphContextSummary,
366 #[serde(default)]
367 pub export_mode: CodeGraphExportMode,
368 #[serde(default, skip_serializing_if = "Option::is_none")]
369 pub visible_levels: Option<usize>,
370 #[serde(default, skip_serializing_if = "Option::is_none")]
371 pub focus: Option<BlockId>,
372 #[serde(default, skip_serializing_if = "Option::is_none")]
373 pub focus_short_id: Option<String>,
374 #[serde(default, skip_serializing_if = "Option::is_none")]
375 pub focus_label: Option<String>,
376 pub visible_node_count: usize,
377 pub hidden_unreachable_count: usize,
378 #[serde(default, skip_serializing_if = "Vec::is_empty")]
379 pub hidden_levels: Vec<CodeGraphHiddenLevelSummary>,
380 pub frontier: Vec<CodeGraphContextFrontierAction>,
381 pub heuristics: CodeGraphContextHeuristics,
382 pub nodes: Vec<CodeGraphContextNodeExport>,
383 pub edges: Vec<CodeGraphContextEdgeExport>,
384 pub omitted_symbol_count: usize,
385 pub total_selected_edges: usize,
386 #[serde(default, skip_serializing_if = "String::is_empty")]
387 pub rendered: String,
388}
389
390impl Default for CodeGraphRenderConfig {
391 fn default() -> Self {
392 Self {
393 max_edges_per_node: 6,
394 max_source_lines: 12,
395 }
396 }
397}
398
399impl CodeGraphRenderConfig {
400 pub fn for_max_tokens(max_tokens: usize) -> Self {
401 if max_tokens <= 512 {
402 Self {
403 max_edges_per_node: 2,
404 max_source_lines: 4,
405 }
406 } else if max_tokens <= 1024 {
407 Self {
408 max_edges_per_node: 3,
409 max_source_lines: 6,
410 }
411 } else if max_tokens <= 2048 {
412 Self {
413 max_edges_per_node: 4,
414 max_source_lines: 8,
415 }
416 } else {
417 Self::default()
418 }
419 }
420}
421
422#[derive(Debug, Clone)]
423struct IndexedEdge {
424 other: BlockId,
425 relation: String,
426}
427
428#[derive(Debug, Clone)]
429struct CodeGraphQueryIndex {
430 logical_keys: HashMap<BlockId, String>,
431 logical_key_to_id: HashMap<String, BlockId>,
432 paths_to_id: HashMap<String, BlockId>,
433 display_to_id: HashMap<String, BlockId>,
434 symbol_names_to_id: HashMap<String, Vec<BlockId>>,
435 node_classes: HashMap<BlockId, String>,
436 outgoing: HashMap<BlockId, Vec<IndexedEdge>>,
437 incoming: HashMap<BlockId, Vec<IndexedEdge>>,
438 file_symbols: HashMap<BlockId, Vec<BlockId>>,
439 symbol_children: HashMap<BlockId, Vec<BlockId>>,
440 structure_parent: HashMap<BlockId, BlockId>,
441}
442
443impl CodeGraphContextSession {
444 pub fn new() -> Self {
445 Self::default()
446 }
447
448 pub fn selected_block_ids(&self) -> Vec<BlockId> {
449 let mut ids: Vec<_> = self.selected.keys().copied().collect();
450 ids.sort_by_key(BlockId::to_string);
451 ids
452 }
453
454 pub fn summary(&self, doc: &Document) -> CodeGraphContextSummary {
455 let index = CodeGraphQueryIndex::new(doc);
456 let mut summary = CodeGraphContextSummary {
457 selected: self.selected.len(),
458 max_selected: self.prune_policy.max_selected,
459 repositories: 0,
460 directories: 0,
461 files: 0,
462 symbols: 0,
463 hydrated_sources: 0,
464 };
465
466 for node in self.selected.values() {
467 match index.node_class(&node.block_id).unwrap_or("unknown") {
468 "repository" => summary.repositories += 1,
469 "directory" => summary.directories += 1,
470 "file" => summary.files += 1,
471 "symbol" => summary.symbols += 1,
472 _ => {}
473 }
474 if node.hydrated_source.is_some() {
475 summary.hydrated_sources += 1;
476 }
477 }
478
479 summary
480 }
481
482 pub fn clear(&mut self) {
483 self.selected.clear();
484 self.focus = None;
485 self.history.push("clear".to_string());
486 }
487
488 pub fn set_prune_policy(&mut self, policy: CodeGraphPrunePolicy) {
489 self.prune_policy = policy;
490 self.history.push(format!(
491 "policy:max_selected:{}:demote:{}:protect_focus:{}",
492 self.prune_policy.max_selected,
493 self.prune_policy.demote_before_remove,
494 self.prune_policy.protect_focus
495 ));
496 }
497
498 pub fn set_focus(
499 &mut self,
500 doc: &Document,
501 block_id: Option<BlockId>,
502 ) -> CodeGraphContextUpdate {
503 let mut update = CodeGraphContextUpdate::default();
504 if let Some(block_id) = block_id {
505 if doc.get_block(&block_id).is_none() {
506 update
507 .warnings
508 .push(format!("focus block not found: {}", block_id));
509 return update;
510 }
511 self.ensure_selected_with_origin(
512 block_id,
513 CodeGraphDetailLevel::Skeleton,
514 selection_origin(CodeGraphSelectionOriginKind::Manual, None, None),
515 &mut update,
516 );
517 }
518 self.focus = block_id;
519 self.apply_prune_policy(doc, &mut update);
520 update.focus = self.focus;
521 self.history.push(match self.focus {
522 Some(id) => format!("focus:{}", id),
523 None => "focus:clear".to_string(),
524 });
525 update
526 }
527
528 pub fn select_block(
529 &mut self,
530 doc: &Document,
531 block_id: BlockId,
532 detail_level: CodeGraphDetailLevel,
533 ) -> CodeGraphContextUpdate {
534 let mut update = CodeGraphContextUpdate::default();
535 if doc.get_block(&block_id).is_none() {
536 update
537 .warnings
538 .push(format!("block not found: {}", block_id));
539 return update;
540 }
541 self.ensure_selected_with_origin(
542 block_id,
543 detail_level,
544 selection_origin(CodeGraphSelectionOriginKind::Manual, None, None),
545 &mut update,
546 );
547 self.apply_prune_policy(doc, &mut update);
548 self.history
549 .push(format!("select:{}:{:?}", block_id, detail_level));
550 update.focus = self.focus;
551 update
552 }
553
554 pub fn remove_block(&mut self, block_id: BlockId) -> CodeGraphContextUpdate {
555 let mut update = CodeGraphContextUpdate::default();
556 if self.selected.remove(&block_id).is_some() {
557 update.removed.push(block_id);
558 if self.focus == Some(block_id) {
559 self.focus = None;
560 }
561 self.history.push(format!("remove:{}", block_id));
562 }
563 update.focus = self.focus;
564 update
565 }
566
567 pub fn pin(&mut self, block_id: BlockId, pinned: bool) -> CodeGraphContextUpdate {
568 let mut update = CodeGraphContextUpdate::default();
569 if let Some(node) = self.selected.get_mut(&block_id) {
570 node.pinned = pinned;
571 update.changed.push(block_id);
572 self.history.push(format!(
573 "{}:{}",
574 if pinned { "pin" } else { "unpin" },
575 block_id
576 ));
577 }
578 update.focus = self.focus;
579 update
580 }
581
582 pub fn seed_overview(&mut self, doc: &Document) -> CodeGraphContextUpdate {
583 self.seed_overview_with_depth(doc, None)
584 }
585
586 pub fn seed_overview_with_depth(
587 &mut self,
588 doc: &Document,
589 max_depth: Option<usize>,
590 ) -> CodeGraphContextUpdate {
591 let index = CodeGraphQueryIndex::new(doc);
592 let previous: HashSet<_> = self.selected.keys().copied().collect();
593 self.selected.clear();
594 self.focus = None;
595
596 let mut update = CodeGraphContextUpdate::default();
597 let mut selected = Vec::new();
598 for block_id in index.overview_nodes(doc, max_depth) {
599 self.ensure_selected_with_origin(
600 block_id,
601 CodeGraphDetailLevel::Skeleton,
602 selection_origin(CodeGraphSelectionOriginKind::Overview, None, None),
603 &mut update,
604 );
605 selected.push(block_id);
606 }
607
608 if self.focus.is_none() {
609 self.focus = selected.first().copied().or(Some(doc.root));
610 }
611 self.apply_prune_policy(doc, &mut update);
612 update.focus = self.focus;
613 update.removed = previous
614 .into_iter()
615 .filter(|block_id| !self.selected.contains_key(block_id))
616 .collect();
617 update.removed.sort_by_key(BlockId::to_string);
618 self.history.push(match max_depth {
619 Some(depth) => format!("seed:overview:{}", depth),
620 None => "seed:overview:all".to_string(),
621 });
622 update
623 }
624
625 pub fn expand_file(&mut self, doc: &Document, file_id: BlockId) -> CodeGraphContextUpdate {
626 self.expand_file_with_config(doc, file_id, &CodeGraphTraversalConfig::default())
627 }
628
629 pub fn expand_file_with_depth(
630 &mut self,
631 doc: &Document,
632 file_id: BlockId,
633 depth: usize,
634 ) -> CodeGraphContextUpdate {
635 self.expand_file_with_config(
636 doc,
637 file_id,
638 &CodeGraphTraversalConfig {
639 depth,
640 ..CodeGraphTraversalConfig::default()
641 },
642 )
643 }
644
645 pub fn expand_file_with_config(
646 &mut self,
647 doc: &Document,
648 file_id: BlockId,
649 traversal: &CodeGraphTraversalConfig,
650 ) -> CodeGraphContextUpdate {
651 let index = CodeGraphQueryIndex::new(doc);
652 let mut update = CodeGraphContextUpdate::default();
653 self.ensure_selected_with_origin(
654 file_id,
655 CodeGraphDetailLevel::Neighborhood,
656 selection_origin(CodeGraphSelectionOriginKind::Manual, None, None),
657 &mut update,
658 );
659 let mut added_count = 0usize;
660 let mut skipped_for_threshold = 0usize;
661 let mut budget_exhausted = false;
662 if traversal.depth() > 0 {
663 let mut queue: VecDeque<(BlockId, usize)> = index
664 .file_symbols(&file_id)
665 .into_iter()
666 .map(|symbol_id| (symbol_id, 1usize))
667 .collect();
668 while let Some((symbol_id, symbol_depth)) = queue.pop_front() {
669 let candidate_priority = frontier_priority("expand_file", None, 1, false);
670 if traversal
671 .priority_threshold
672 .map(|threshold| candidate_priority < threshold)
673 .unwrap_or(false)
674 {
675 skipped_for_threshold += 1;
676 continue;
677 }
678 if traversal
679 .max_add
680 .map(|limit| added_count >= limit)
681 .unwrap_or(false)
682 {
683 budget_exhausted = true;
684 break;
685 }
686 let was_selected = self.selected.contains_key(&symbol_id);
687 self.ensure_selected_with_origin(
688 symbol_id,
689 CodeGraphDetailLevel::SymbolCard,
690 selection_origin(
691 CodeGraphSelectionOriginKind::FileSymbols,
692 None,
693 Some(file_id),
694 ),
695 &mut update,
696 );
697 if !was_selected && self.selected.contains_key(&symbol_id) {
698 added_count += 1;
699 }
700 if symbol_depth >= traversal.depth() {
701 continue;
702 }
703 for child in index.symbol_children(&symbol_id) {
704 queue.push_back((child, symbol_depth + 1));
705 }
706 }
707 }
708 if skipped_for_threshold > 0 {
709 update.warnings.push(format!(
710 "skipped {} file-symbol candidates below priority threshold",
711 skipped_for_threshold
712 ));
713 }
714 if budget_exhausted {
715 update.warnings.push(format!(
716 "stopped file expansion after adding {} nodes due to max_add budget",
717 added_count
718 ));
719 }
720 self.focus = Some(file_id);
721 self.apply_prune_policy(doc, &mut update);
722 update.focus = self.focus;
723 self.history.push(format!(
724 "expand:file:{}:{}:{}:{}",
725 file_id,
726 traversal.depth(),
727 traversal
728 .max_add
729 .map(|value| value.to_string())
730 .unwrap_or_else(|| "*".to_string()),
731 traversal
732 .priority_threshold
733 .map(|value| value.to_string())
734 .unwrap_or_else(|| "*".to_string())
735 ));
736 update
737 }
738
739 pub fn expand_dependencies(
740 &mut self,
741 doc: &Document,
742 block_id: BlockId,
743 relation_filter: Option<&str>,
744 ) -> CodeGraphContextUpdate {
745 self.expand_dependencies_with_config(
746 doc,
747 block_id,
748 &CodeGraphTraversalConfig {
749 relation_filters: relation_filter
750 .map(|relation| vec![relation.to_string()])
751 .unwrap_or_default(),
752 ..CodeGraphTraversalConfig::default()
753 },
754 )
755 }
756
757 pub fn expand_dependencies_with_filters(
758 &mut self,
759 doc: &Document,
760 block_id: BlockId,
761 relation_filters: Option<&HashSet<String>>,
762 depth: usize,
763 ) -> CodeGraphContextUpdate {
764 self.expand_dependencies_with_config(
765 doc,
766 block_id,
767 &CodeGraphTraversalConfig {
768 depth,
769 relation_filters: relation_filters
770 .map(|filters| filters.iter().cloned().collect())
771 .unwrap_or_default(),
772 ..CodeGraphTraversalConfig::default()
773 },
774 )
775 }
776
777 pub fn expand_dependencies_with_config(
778 &mut self,
779 doc: &Document,
780 block_id: BlockId,
781 traversal: &CodeGraphTraversalConfig,
782 ) -> CodeGraphContextUpdate {
783 self.expand_neighbors(doc, block_id, traversal, TraversalKind::Outgoing)
784 }
785
786 pub fn expand_dependents(
787 &mut self,
788 doc: &Document,
789 block_id: BlockId,
790 relation_filter: Option<&str>,
791 ) -> CodeGraphContextUpdate {
792 self.expand_dependents_with_config(
793 doc,
794 block_id,
795 &CodeGraphTraversalConfig {
796 relation_filters: relation_filter
797 .map(|relation| vec![relation.to_string()])
798 .unwrap_or_default(),
799 ..CodeGraphTraversalConfig::default()
800 },
801 )
802 }
803
804 pub fn expand_dependents_with_filters(
805 &mut self,
806 doc: &Document,
807 block_id: BlockId,
808 relation_filters: Option<&HashSet<String>>,
809 depth: usize,
810 ) -> CodeGraphContextUpdate {
811 self.expand_dependents_with_config(
812 doc,
813 block_id,
814 &CodeGraphTraversalConfig {
815 depth,
816 relation_filters: relation_filters
817 .map(|filters| filters.iter().cloned().collect())
818 .unwrap_or_default(),
819 ..CodeGraphTraversalConfig::default()
820 },
821 )
822 }
823
824 pub fn expand_dependents_with_config(
825 &mut self,
826 doc: &Document,
827 block_id: BlockId,
828 traversal: &CodeGraphTraversalConfig,
829 ) -> CodeGraphContextUpdate {
830 self.expand_neighbors(doc, block_id, traversal, TraversalKind::Incoming)
831 }
832
833 pub fn collapse(
834 &mut self,
835 doc: &Document,
836 block_id: BlockId,
837 include_descendants: bool,
838 ) -> CodeGraphContextUpdate {
839 let index = CodeGraphQueryIndex::new(doc);
840 let mut update = CodeGraphContextUpdate::default();
841 let mut to_remove = vec![block_id];
842 if include_descendants {
843 to_remove.extend(index.descendants(block_id));
844 }
845
846 for id in to_remove {
847 let Some(node) = self.selected.get(&id) else {
848 continue;
849 };
850 if node.pinned {
851 update
852 .warnings
853 .push(format!("{} is pinned and was not removed", id));
854 continue;
855 }
856 self.selected.remove(&id);
857 update.removed.push(id);
858 if self.focus == Some(id) {
859 self.focus = None;
860 }
861 }
862
863 if self.focus.is_none() {
864 self.focus = self.selected.keys().next().copied();
865 }
866 update.focus = self.focus;
867 self.history
868 .push(format!("collapse:{}:{}", block_id, include_descendants));
869 update
870 }
871
872 pub fn hydrate_source(
873 &mut self,
874 doc: &Document,
875 block_id: BlockId,
876 padding: usize,
877 ) -> CodeGraphContextUpdate {
878 let mut update = CodeGraphContextUpdate::default();
879 self.ensure_selected_with_origin(
880 block_id,
881 CodeGraphDetailLevel::Source,
882 selection_origin(CodeGraphSelectionOriginKind::Manual, None, None),
883 &mut update,
884 );
885 match hydrate_source_excerpt(doc, block_id, padding) {
886 Ok(Some(excerpt)) => {
887 if let Some(node) = self.selected.get_mut(&block_id) {
888 node.detail_level = CodeGraphDetailLevel::Source;
889 node.hydrated_source = Some(excerpt);
890 update.changed.push(block_id);
891 }
892 }
893 Ok(None) => update
894 .warnings
895 .push(format!("no coderef available for {}", block_id)),
896 Err(error) => update.warnings.push(error),
897 }
898 self.focus = Some(block_id);
899 self.apply_prune_policy(doc, &mut update);
900 update.focus = self.focus;
901 self.history.push(format!("hydrate:{}", block_id));
902 update
903 }
904
905 pub fn prune(&mut self, doc: &Document, max_selected: Option<usize>) -> CodeGraphContextUpdate {
906 let mut update = CodeGraphContextUpdate::default();
907 if let Some(limit) = max_selected {
908 self.prune_policy.max_selected = limit.max(1);
909 }
910 self.apply_prune_policy(doc, &mut update);
911 self.history
912 .push(format!("prune:{}", self.prune_policy.max_selected));
913 update.focus = self.focus;
914 update
915 }
916
917 pub fn render_for_prompt(&self, doc: &Document, config: &CodeGraphRenderConfig) -> String {
918 let index = CodeGraphQueryIndex::new(doc);
919 let summary = self.summary(doc);
920 let short_ids = make_short_ids(self, &index);
921 let selected_ids: HashSet<_> = self.selected.keys().copied().collect();
922
923 let mut repository_nodes = Vec::new();
924 let mut directory_nodes = Vec::new();
925 let mut file_nodes = Vec::new();
926 let mut symbol_nodes = Vec::new();
927
928 for block_id in self.selected_block_ids() {
929 match index.node_class(&block_id).unwrap_or("unknown") {
930 "repository" => repository_nodes.push(block_id),
931 "directory" => directory_nodes.push(block_id),
932 "file" => file_nodes.push(block_id),
933 "symbol" => symbol_nodes.push(block_id),
934 _ => {}
935 }
936 }
937
938 let mut out = String::new();
939 let _ = writeln!(out, "CodeGraph working set");
940 let focus = self
941 .focus
942 .and_then(|id| render_reference(doc, &index, &short_ids, id));
943 let _ = writeln!(
944 out,
945 "focus: {}",
946 focus.unwrap_or_else(|| "none".to_string())
947 );
948 let _ = writeln!(
949 out,
950 "summary: selected={}/{} repositories={} directories={} files={} symbols={} hydrated={}",
951 summary.selected,
952 summary.max_selected,
953 summary.repositories,
954 summary.directories,
955 summary.files,
956 summary.symbols,
957 summary.hydrated_sources
958 );
959
960 if !repository_nodes.is_empty() || !directory_nodes.is_empty() || !file_nodes.is_empty() {
961 let _ = writeln!(out, "\nfilesystem:");
962 for block_id in repository_nodes
963 .into_iter()
964 .chain(directory_nodes.into_iter())
965 .chain(file_nodes.into_iter())
966 {
967 let block = match doc.get_block(&block_id) {
968 Some(block) => block,
969 None => continue,
970 };
971 let short = short_ids
972 .get(&block_id)
973 .cloned()
974 .unwrap_or_else(|| block_id.to_string());
975 let label = index
976 .display_label(doc, &block_id)
977 .unwrap_or_else(|| block_id.to_string());
978 let language = block
979 .metadata
980 .custom
981 .get(META_LANGUAGE)
982 .and_then(Value::as_str)
983 .map(|value| format!(" [{}]", value))
984 .unwrap_or_default();
985 let pin = self
986 .selected
987 .get(&block_id)
988 .filter(|node| node.pinned)
989 .map(|_| " [pinned]")
990 .unwrap_or("");
991 let _ = writeln!(out, "- [{}] {}{}{}", short, label, language, pin);
992 }
993 }
994
995 if !symbol_nodes.is_empty() {
996 let _ = writeln!(out, "\nopened symbols:");
997 for block_id in symbol_nodes {
998 let Some(block) = doc.get_block(&block_id) else {
999 continue;
1000 };
1001 let Some(node) = self.selected.get(&block_id) else {
1002 continue;
1003 };
1004 let short = short_ids
1005 .get(&block_id)
1006 .cloned()
1007 .unwrap_or_else(|| block_id.to_string());
1008 let coderef = metadata_coderef_display(block)
1009 .or_else(|| content_coderef_display(block))
1010 .unwrap_or_else(|| {
1011 index
1012 .display_label(doc, &block_id)
1013 .unwrap_or_else(|| block_id.to_string())
1014 });
1015 let pin = if node.pinned { " [pinned]" } else { "" };
1016 let _ = writeln!(
1017 out,
1018 "- [{}] {}{} @ {}",
1019 short,
1020 format_symbol_signature(block),
1021 format_symbol_modifiers(block),
1022 coderef
1023 );
1024 if !pin.is_empty() {
1025 let _ = writeln!(out, " flags:{}", pin);
1026 }
1027 if let Some(description) =
1028 content_string(block, "description").or_else(|| block.metadata.summary.clone())
1029 {
1030 let _ = writeln!(out, " docs: {}", description);
1031 }
1032
1033 if node.detail_level.includes_neighborhood() {
1034 render_edge_section(
1035 &mut out,
1036 "outgoing",
1037 index.outgoing_edges(&block_id),
1038 &selected_ids,
1039 &short_ids,
1040 doc,
1041 &index,
1042 config.max_edges_per_node,
1043 );
1044 render_edge_section(
1045 &mut out,
1046 "incoming",
1047 index.incoming_edges(&block_id),
1048 &selected_ids,
1049 &short_ids,
1050 doc,
1051 &index,
1052 config.max_edges_per_node,
1053 );
1054 }
1055
1056 if node.detail_level.includes_source() {
1057 if let Some(source) = &node.hydrated_source {
1058 let _ = writeln!(
1059 out,
1060 " source: {}:{}-{}",
1061 source.path, source.start_line, source.end_line
1062 );
1063 for line in source.snippet.lines().take(config.max_source_lines) {
1064 let _ = writeln!(out, " {}", line);
1065 }
1066 }
1067 }
1068 }
1069 }
1070
1071 let total_symbols = index.total_symbols();
1072 let omitted_symbols = total_symbols.saturating_sub(summary.symbols);
1073 let _ = writeln!(out, "\nomissions:");
1074 let _ = writeln!(
1075 out,
1076 "- symbols omitted from working set: {}",
1077 omitted_symbols
1078 );
1079 let _ = writeln!(
1080 out,
1081 "- prune policy: max_selected={} demote_before_remove={} protect_focus={}",
1082 self.prune_policy.max_selected,
1083 self.prune_policy.demote_before_remove,
1084 self.prune_policy.protect_focus
1085 );
1086
1087 let _ = writeln!(out, "\nfrontier:");
1088 if let Some(focus_id) = self.focus {
1089 match index.node_class(&focus_id).unwrap_or("unknown") {
1090 "file" => {
1091 let short = short_ids
1092 .get(&focus_id)
1093 .cloned()
1094 .unwrap_or_else(|| focus_id.to_string());
1095 let _ = writeln!(out, "- [{}] expand file symbols", short);
1096 let _ = writeln!(out, "- [{}] hydrate file source", short);
1097 }
1098 "symbol" => {
1099 let short = short_ids
1100 .get(&focus_id)
1101 .cloned()
1102 .unwrap_or_else(|| focus_id.to_string());
1103 let _ = writeln!(out, "- [{}] expand dependencies", short);
1104 let _ = writeln!(out, "- [{}] expand dependents", short);
1105 let _ = writeln!(out, "- [{}] hydrate source", short);
1106 let _ = writeln!(out, "- [{}] collapse", short);
1107 }
1108 _ => {
1109 let _ = writeln!(
1110 out,
1111 "- set focus to a file or symbol to expand the working set"
1112 );
1113 }
1114 }
1115 } else {
1116 let _ = writeln!(out, "- no focus block set");
1117 }
1118
1119 out.trim_end().to_string()
1120 }
1121
1122 pub fn export(&self, doc: &Document, config: &CodeGraphRenderConfig) -> CodeGraphContextExport {
1123 self.export_with_config(doc, config, &CodeGraphExportConfig::default())
1124 }
1125
1126 pub fn export_with_config(
1127 &self,
1128 doc: &Document,
1129 config: &CodeGraphRenderConfig,
1130 export_config: &CodeGraphExportConfig,
1131 ) -> CodeGraphContextExport {
1132 let index = CodeGraphQueryIndex::new(doc);
1133 let summary = self.summary(doc);
1134 let short_ids = make_short_ids(self, &index);
1135 let selected_ids: HashSet<_> = self.selected.keys().copied().collect();
1136 let distances = focus_distances(doc, self.focus, &selected_ids, &index);
1137 let visible_selected_ids = visible_selected_ids(
1138 self.focus,
1139 &selected_ids,
1140 &distances,
1141 export_config.visible_levels,
1142 );
1143 let hidden_levels = hidden_level_summaries(
1144 self,
1145 &index,
1146 &selected_ids,
1147 &visible_selected_ids,
1148 &distances,
1149 export_config.visible_levels,
1150 );
1151 let hidden_unreachable_count = selected_ids
1152 .iter()
1153 .filter(|block_id| {
1154 !visible_selected_ids.contains(block_id) && !distances.contains_key(block_id)
1155 })
1156 .count();
1157 let filtered_selected_ids =
1158 class_filtered_selected_ids(&index, &visible_selected_ids, export_config);
1159 let rendered = if export_config.include_rendered {
1160 self.render_for_prompt(doc, config)
1161 } else {
1162 String::new()
1163 };
1164
1165 let mut nodes = Vec::new();
1166 for block_id in self.selected_block_ids() {
1167 if !filtered_selected_ids.contains(&block_id) {
1168 continue;
1169 }
1170 let Some(block) = doc.get_block(&block_id) else {
1171 continue;
1172 };
1173 let Some(node) = self.selected.get(&block_id) else {
1174 continue;
1175 };
1176 let node_class = index.node_class(&block_id).unwrap_or("unknown").to_string();
1177 let label = index
1178 .display_label(doc, &block_id)
1179 .unwrap_or_else(|| block_id.to_string());
1180 let logical_key = block_logical_key(block);
1181 let content_name = content_string(block, "name");
1182 let symbol_name = block
1183 .metadata
1184 .custom
1185 .get(META_SYMBOL_NAME)
1186 .and_then(Value::as_str)
1187 .or(content_name.as_deref())
1188 .map(str::to_string);
1189 let path = metadata_coderef_path(block).or_else(|| content_coderef_path(block));
1190 let distance_from_focus = distances.get(&block_id).copied();
1191 let relevance_score =
1192 relevance_score_for_node(self, &index, block_id, distance_from_focus);
1193 let signature = if node_class == "symbol" {
1194 Some(format!(
1195 "{}{}",
1196 format_symbol_signature(block),
1197 format_symbol_modifiers(block)
1198 ))
1199 } else {
1200 None
1201 };
1202 let docs = if should_include_docs(
1203 export_config,
1204 self.focus,
1205 block_id,
1206 node,
1207 distance_from_focus,
1208 ) {
1209 content_string(block, "description").or_else(|| block.metadata.summary.clone())
1210 } else {
1211 None
1212 };
1213 let coderef = block_coderef(block).map(|coderef| CodeGraphCoderef {
1214 path: coderef.path,
1215 display: coderef.display,
1216 start_line: coderef.start_line,
1217 end_line: coderef.end_line,
1218 });
1219 let hydrated_source = if should_include_hydrated_source(
1220 export_config,
1221 self.focus,
1222 block_id,
1223 node,
1224 distance_from_focus,
1225 ) {
1226 node.hydrated_source.clone()
1227 } else {
1228 None
1229 };
1230
1231 nodes.push(CodeGraphContextNodeExport {
1232 block_id,
1233 short_id: short_ids
1234 .get(&block_id)
1235 .cloned()
1236 .unwrap_or_else(|| block_id.to_string()),
1237 node_class,
1238 label,
1239 detail_level: node.detail_level,
1240 pinned: node.pinned,
1241 distance_from_focus,
1242 relevance_score,
1243 logical_key,
1244 symbol_name,
1245 path,
1246 signature,
1247 docs,
1248 origin: node.origin.clone(),
1249 coderef,
1250 hydrated_source,
1251 });
1252 }
1253 nodes.sort_by_key(|node| {
1254 (
1255 std::cmp::Reverse(node.relevance_score),
1256 node.distance_from_focus.unwrap_or(usize::MAX),
1257 node.short_id.clone(),
1258 )
1259 });
1260
1261 let (edges, total_selected_edges) =
1262 export_edges(&index, &filtered_selected_ids, &short_ids, export_config);
1263
1264 let frontier = self.export_frontier(doc, &index, &short_ids, &selected_ids);
1265 let heuristics = self.compute_heuristics(&index, &frontier);
1266 let omitted_symbol_count = index.total_symbols().saturating_sub(summary.symbols);
1267
1268 CodeGraphContextExport {
1269 summary,
1270 export_mode: export_config.mode,
1271 visible_levels: export_config.visible_levels,
1272 focus: self.focus,
1273 focus_short_id: self.focus.and_then(|id| short_ids.get(&id).cloned()),
1274 focus_label: self.focus.and_then(|id| index.display_label(doc, &id)),
1275 visible_node_count: nodes.len(),
1276 hidden_unreachable_count,
1277 hidden_levels,
1278 nodes,
1279 edges,
1280 frontier: frontier
1281 .into_iter()
1282 .take(export_config.max_frontier_actions.max(1))
1283 .collect(),
1284 heuristics,
1285 omitted_symbol_count,
1286 total_selected_edges,
1287 rendered,
1288 }
1289 }
1290
1291 fn export_frontier(
1292 &self,
1293 doc: &Document,
1294 index: &CodeGraphQueryIndex,
1295 short_ids: &HashMap<BlockId, String>,
1296 selected_ids: &HashSet<BlockId>,
1297 ) -> Vec<CodeGraphContextFrontierAction> {
1298 let Some(focus_id) = self.focus else {
1299 return Vec::new();
1300 };
1301 let short_id = short_ids
1302 .get(&focus_id)
1303 .cloned()
1304 .unwrap_or_else(|| focus_id.to_string());
1305 let label = index
1306 .display_label(doc, &focus_id)
1307 .unwrap_or_else(|| focus_id.to_string());
1308 match index.node_class(&focus_id).unwrap_or("unknown") {
1309 "file" => {
1310 let hidden = index
1311 .file_symbols(&focus_id)
1312 .into_iter()
1313 .filter(|id| !selected_ids.contains(id))
1314 .count();
1315 let mut actions = vec![CodeGraphContextFrontierAction {
1316 block_id: focus_id,
1317 short_id,
1318 action: "expand_file".to_string(),
1319 relation: None,
1320 direction: None,
1321 candidate_count: hidden,
1322 priority: frontier_priority("expand_file", None, hidden, false),
1323 description: format!("Expand file symbols for {}", label),
1324 }];
1325 actions.push(CodeGraphContextFrontierAction {
1326 block_id: focus_id,
1327 short_id: actions[0].short_id.clone(),
1328 action: "hydrate_source".to_string(),
1329 relation: None,
1330 direction: None,
1331 candidate_count: usize::from(
1332 self.selected
1333 .get(&focus_id)
1334 .and_then(|node| node.hydrated_source.as_ref())
1335 .is_none(),
1336 ),
1337 priority: frontier_priority(
1338 "hydrate_source",
1339 None,
1340 usize::from(
1341 self.selected
1342 .get(&focus_id)
1343 .and_then(|node| node.hydrated_source.as_ref())
1344 .is_none(),
1345 ),
1346 false,
1347 ),
1348 description: format!("Hydrate source for file {}", label),
1349 });
1350 actions.sort_by_key(|action| {
1351 (
1352 std::cmp::Reverse(action.priority),
1353 action.action.clone(),
1354 action.relation.clone(),
1355 )
1356 });
1357 actions
1358 }
1359 "symbol" => {
1360 let mut actions = Vec::new();
1361 append_relation_frontier(
1362 &mut actions,
1363 focus_id,
1364 &short_id,
1365 &label,
1366 index.outgoing_edges(&focus_id),
1367 selected_ids,
1368 "expand_dependencies",
1369 "outgoing",
1370 );
1371 append_relation_frontier(
1372 &mut actions,
1373 focus_id,
1374 &short_id,
1375 &label,
1376 index.incoming_edges(&focus_id),
1377 selected_ids,
1378 "expand_dependents",
1379 "incoming",
1380 );
1381 actions.push(CodeGraphContextFrontierAction {
1382 block_id: focus_id,
1383 short_id: short_id.clone(),
1384 action: "hydrate_source".to_string(),
1385 relation: None,
1386 direction: None,
1387 candidate_count: usize::from(
1388 self.selected
1389 .get(&focus_id)
1390 .and_then(|node| node.hydrated_source.as_ref())
1391 .is_none(),
1392 ),
1393 priority: frontier_priority(
1394 "hydrate_source",
1395 None,
1396 usize::from(
1397 self.selected
1398 .get(&focus_id)
1399 .and_then(|node| node.hydrated_source.as_ref())
1400 .is_none(),
1401 ),
1402 false,
1403 ),
1404 description: format!("Hydrate source for {}", label),
1405 });
1406 actions.push(CodeGraphContextFrontierAction {
1407 block_id: focus_id,
1408 short_id,
1409 action: "collapse".to_string(),
1410 relation: None,
1411 direction: None,
1412 candidate_count: 1,
1413 priority: frontier_priority("collapse", None, 1, false),
1414 description: format!("Collapse {} from working set", label),
1415 });
1416 actions.sort_by_key(|action| {
1417 (
1418 std::cmp::Reverse(action.priority),
1419 action.action.clone(),
1420 action.relation.clone(),
1421 )
1422 });
1423 actions
1424 }
1425 _ => Vec::new(),
1426 }
1427 }
1428
1429 fn compute_heuristics(
1430 &self,
1431 index: &CodeGraphQueryIndex,
1432 frontier: &[CodeGraphContextFrontierAction],
1433 ) -> CodeGraphContextHeuristics {
1434 let focus_node = self.focus.and_then(|id| self.selected.get(&id));
1435 let focus_hydrated = focus_node
1436 .and_then(|node| node.hydrated_source.as_ref())
1437 .is_some();
1438 let hidden_candidate_count = frontier
1439 .iter()
1440 .filter(|action| action.action.starts_with("expand_") && action.candidate_count > 0)
1441 .map(|action| action.candidate_count)
1442 .sum();
1443 let low_value_candidate_count = frontier
1444 .iter()
1445 .filter(|action| action.action.starts_with("expand_") && action.priority <= 30)
1446 .map(|action| action.candidate_count)
1447 .sum();
1448 let recommended_actions: Vec<_> = frontier
1449 .iter()
1450 .filter(|action| action.candidate_count > 0)
1451 .take(3)
1452 .cloned()
1453 .collect();
1454 let recommended_next_action = recommended_actions.first().cloned();
1455
1456 let mut reasons = Vec::new();
1457 let should_stop = match self.focus {
1458 None => {
1459 reasons
1460 .push("set focus to a file or symbol before continuing expansion".to_string());
1461 false
1462 }
1463 Some(focus_id) => match index.node_class(&focus_id).unwrap_or("unknown") {
1464 "file" => {
1465 if hidden_candidate_count == 0 && focus_hydrated {
1466 reasons.push(
1467 "focus file is hydrated and no unselected file symbols remain"
1468 .to_string(),
1469 );
1470 true
1471 } else if hidden_candidate_count == 0 {
1472 reasons.push(
1473 "all file symbols for the focused file are already selected"
1474 .to_string(),
1475 );
1476 false
1477 } else {
1478 false
1479 }
1480 }
1481 "symbol" => {
1482 if hidden_candidate_count == 0 && focus_hydrated {
1483 reasons.push(
1484 "focus symbol is hydrated and no unselected dependency frontier remains"
1485 .to_string(),
1486 );
1487 true
1488 } else if focus_hydrated
1489 && hidden_candidate_count > 0
1490 && hidden_candidate_count == low_value_candidate_count
1491 {
1492 reasons.push(
1493 "remaining frontier is low-value compared to the hydrated focus symbol"
1494 .to_string(),
1495 );
1496 true
1497 } else {
1498 false
1499 }
1500 }
1501 _ => frontier.iter().all(|action| action.candidate_count == 0),
1502 },
1503 };
1504
1505 CodeGraphContextHeuristics {
1506 should_stop,
1507 reasons,
1508 hidden_candidate_count,
1509 low_value_candidate_count,
1510 recommended_next_action,
1511 recommended_actions,
1512 }
1513 }
1514
1515 fn ensure_selected_with_origin(
1516 &mut self,
1517 block_id: BlockId,
1518 detail_level: CodeGraphDetailLevel,
1519 origin: Option<CodeGraphSelectionOrigin>,
1520 update: &mut CodeGraphContextUpdate,
1521 ) {
1522 match self.selected.get_mut(&block_id) {
1523 Some(node) => {
1524 let next = node.detail_level.max(detail_level);
1525 if next != node.detail_level {
1526 node.detail_level = next;
1527 update.changed.push(block_id);
1528 }
1529 if origin_is_more_protective(origin.as_ref(), node.origin.as_ref()) {
1530 node.origin = origin;
1531 push_unique(&mut update.changed, block_id);
1532 }
1533 }
1534 None => {
1535 self.selected.insert(
1536 block_id,
1537 CodeGraphContextNode {
1538 block_id,
1539 detail_level,
1540 pinned: false,
1541 origin,
1542 hydrated_source: None,
1543 },
1544 );
1545 update.added.push(block_id);
1546 }
1547 }
1548 }
1549
1550 fn expand_neighbors(
1551 &mut self,
1552 doc: &Document,
1553 block_id: BlockId,
1554 traversal_config: &CodeGraphTraversalConfig,
1555 traversal_kind: TraversalKind,
1556 ) -> CodeGraphContextUpdate {
1557 let index = CodeGraphQueryIndex::new(doc);
1558 let mut update = CodeGraphContextUpdate::default();
1559 let relation_filters = traversal_config.relation_filter_set();
1560 self.ensure_selected_with_origin(
1561 block_id,
1562 CodeGraphDetailLevel::Neighborhood,
1563 selection_origin(
1564 CodeGraphSelectionOriginKind::Manual,
1565 relation_filters
1566 .as_ref()
1567 .and_then(|filters| join_relation_filters(filters)),
1568 None,
1569 ),
1570 &mut update,
1571 );
1572
1573 let mut queue = VecDeque::from([(block_id, 0usize)]);
1574 let mut visited = HashSet::from([block_id]);
1575 let mut added_count = 0usize;
1576 let mut skipped_for_threshold = 0usize;
1577 let mut budget_exhausted = false;
1578 while let Some((current, current_depth)) = queue.pop_front() {
1579 if current_depth >= traversal_config.depth() {
1580 continue;
1581 }
1582 let edges = match traversal_kind {
1583 TraversalKind::Outgoing => index.outgoing_edges(¤t),
1584 TraversalKind::Incoming => index.incoming_edges(¤t),
1585 };
1586
1587 for edge in edges {
1588 if !relation_matches(relation_filters.as_ref(), edge.relation.as_str()) {
1589 continue;
1590 }
1591 let action_name = match traversal_kind {
1592 TraversalKind::Outgoing => "expand_dependencies",
1593 TraversalKind::Incoming => "expand_dependents",
1594 };
1595 let candidate_priority = frontier_priority(
1596 action_name,
1597 Some(edge.relation.as_str()),
1598 1,
1599 is_low_value_relation(action_name, edge.relation.as_str()),
1600 );
1601 if traversal_config
1602 .priority_threshold
1603 .map(|threshold| candidate_priority < threshold)
1604 .unwrap_or(false)
1605 {
1606 skipped_for_threshold += 1;
1607 continue;
1608 }
1609 if traversal_config
1610 .max_add
1611 .map(|limit| added_count >= limit)
1612 .unwrap_or(false)
1613 {
1614 budget_exhausted = true;
1615 break;
1616 }
1617 let class = index.node_class(&edge.other).unwrap_or("unknown");
1618 let level = if class == "symbol" {
1619 CodeGraphDetailLevel::SymbolCard
1620 } else {
1621 CodeGraphDetailLevel::Skeleton
1622 };
1623 let was_selected = self.selected.contains_key(&edge.other);
1624 self.ensure_selected_with_origin(
1625 edge.other,
1626 level,
1627 selection_origin(
1628 match traversal_kind {
1629 TraversalKind::Outgoing => CodeGraphSelectionOriginKind::Dependencies,
1630 TraversalKind::Incoming => CodeGraphSelectionOriginKind::Dependents,
1631 },
1632 Some(edge.relation.as_str()),
1633 Some(current),
1634 ),
1635 &mut update,
1636 );
1637 if !was_selected && self.selected.contains_key(&edge.other) {
1638 added_count += 1;
1639 }
1640 if visited.insert(edge.other) {
1641 queue.push_back((edge.other, current_depth + 1));
1642 }
1643 }
1644 if budget_exhausted {
1645 break;
1646 }
1647 }
1648
1649 if skipped_for_threshold > 0 {
1650 update.warnings.push(format!(
1651 "skipped {} candidates below priority threshold",
1652 skipped_for_threshold
1653 ));
1654 }
1655 if budget_exhausted {
1656 update.warnings.push(format!(
1657 "stopped expansion after adding {} nodes due to max_add budget",
1658 added_count
1659 ));
1660 }
1661
1662 self.focus = Some(block_id);
1663 self.apply_prune_policy(doc, &mut update);
1664 update.focus = self.focus;
1665 self.history.push(format!(
1666 "expand:{}:{}:{}:{}:{}:{}",
1667 match traversal_kind {
1668 TraversalKind::Outgoing => "dependencies",
1669 TraversalKind::Incoming => "dependents",
1670 },
1671 block_id,
1672 relation_filters
1673 .as_ref()
1674 .map(join_relation_filter_string)
1675 .unwrap_or_else(|| "*".to_string()),
1676 traversal_config.depth(),
1677 traversal_config
1678 .max_add
1679 .map(|value| value.to_string())
1680 .unwrap_or_else(|| "*".to_string()),
1681 traversal_config
1682 .priority_threshold
1683 .map(|value| value.to_string())
1684 .unwrap_or_else(|| "*".to_string())
1685 ));
1686 update
1687 }
1688
1689 fn apply_prune_policy(&mut self, doc: &Document, update: &mut CodeGraphContextUpdate) {
1690 if self.selected.len() <= self.prune_policy.max_selected.max(1) {
1691 return;
1692 }
1693
1694 let index = CodeGraphQueryIndex::new(doc);
1695 let protected_focus = if self.prune_policy.protect_focus {
1696 self.focus
1697 } else {
1698 None
1699 };
1700
1701 if self.prune_policy.demote_before_remove {
1702 while self.selected.len() > self.prune_policy.max_selected.max(1) {
1703 let Some(block_id) = self.next_demotable_block(&index, protected_focus) else {
1704 break;
1705 };
1706 let Some(node) = self.selected.get_mut(&block_id) else {
1707 continue;
1708 };
1709 let next_level = node.detail_level.demoted();
1710 if next_level == node.detail_level {
1711 break;
1712 }
1713 node.detail_level = next_level;
1714 if !node.detail_level.includes_source() {
1715 node.hydrated_source = None;
1716 }
1717 push_unique(&mut update.changed, block_id);
1718 }
1719 }
1720
1721 while self.selected.len() > self.prune_policy.max_selected.max(1) {
1722 let Some(block_id) = self.next_removable_block(&index, protected_focus) else {
1723 update.warnings.push(format!(
1724 "working set has {} nodes but no removable nodes remain under current prune policy",
1725 self.selected.len()
1726 ));
1727 break;
1728 };
1729 self.selected.remove(&block_id);
1730 push_unique(&mut update.removed, block_id);
1731 update.added.retain(|id| id != &block_id);
1732 update.changed.retain(|id| id != &block_id);
1733 if self.focus == Some(block_id) {
1734 self.focus = None;
1735 }
1736 }
1737
1738 if self.focus.is_none() {
1739 self.focus = self.next_focus_candidate(&index);
1740 }
1741 update.focus = self.focus;
1742 }
1743
1744 fn next_demotable_block(
1745 &self,
1746 index: &CodeGraphQueryIndex,
1747 protected_focus: Option<BlockId>,
1748 ) -> Option<BlockId> {
1749 self.selected
1750 .values()
1751 .filter(|node| Some(node.block_id) != protected_focus && !node.pinned)
1752 .filter(|node| node.detail_level.demoted() != node.detail_level)
1753 .max_by_key(|node| {
1754 (
1755 origin_prune_rank(node.origin.as_ref(), &self.prune_policy),
1756 relation_prune_rank(node.origin.as_ref(), &self.prune_policy),
1757 node.detail_level as u8,
1758 prune_removal_rank(index.node_class(&node.block_id).unwrap_or("unknown")),
1759 node.block_id.to_string(),
1760 )
1761 })
1762 .map(|node| node.block_id)
1763 }
1764
1765 fn next_removable_block(
1766 &self,
1767 index: &CodeGraphQueryIndex,
1768 protected_focus: Option<BlockId>,
1769 ) -> Option<BlockId> {
1770 self.selected
1771 .values()
1772 .filter(|node| Some(node.block_id) != protected_focus && !node.pinned)
1773 .max_by_key(|node| {
1774 (
1775 origin_prune_rank(node.origin.as_ref(), &self.prune_policy),
1776 relation_prune_rank(node.origin.as_ref(), &self.prune_policy),
1777 prune_removal_rank(index.node_class(&node.block_id).unwrap_or("unknown")),
1778 node.detail_level as u8,
1779 node.block_id.to_string(),
1780 )
1781 })
1782 .map(|node| node.block_id)
1783 }
1784
1785 fn next_focus_candidate(&self, index: &CodeGraphQueryIndex) -> Option<BlockId> {
1786 self.selected
1787 .values()
1788 .min_by_key(|node| {
1789 (
1790 focus_preference_rank(index.node_class(&node.block_id).unwrap_or("unknown")),
1791 node.block_id.to_string(),
1792 )
1793 })
1794 .map(|node| node.block_id)
1795 }
1796}
1797
1798#[derive(Debug, Clone, Copy)]
1799enum TraversalKind {
1800 Outgoing,
1801 Incoming,
1802}
1803
1804impl CodeGraphQueryIndex {
1805 fn new(doc: &Document) -> Self {
1806 let mut logical_keys = HashMap::new();
1807 let mut logical_key_to_id = HashMap::new();
1808 let mut paths_to_id = HashMap::new();
1809 let mut display_to_id = HashMap::new();
1810 let mut symbol_names_to_id: HashMap<String, Vec<BlockId>> = HashMap::new();
1811 let mut node_classes = HashMap::new();
1812 let mut outgoing: HashMap<BlockId, Vec<IndexedEdge>> = HashMap::new();
1813 let mut incoming: HashMap<BlockId, Vec<IndexedEdge>> = HashMap::new();
1814 let mut file_symbols: HashMap<BlockId, Vec<BlockId>> = HashMap::new();
1815 let mut symbol_children: HashMap<BlockId, Vec<BlockId>> = HashMap::new();
1816 let mut structure_parent: HashMap<BlockId, BlockId> = HashMap::new();
1817
1818 for (block_id, block) in &doc.blocks {
1819 if let Some(key) = block_logical_key(block) {
1820 logical_keys.insert(*block_id, key.clone());
1821 logical_key_to_id.insert(key, *block_id);
1822 }
1823 if let Some(class) = node_class(block) {
1824 node_classes.insert(*block_id, class.clone());
1825 }
1826 if let Some(path) = metadata_coderef_path(block).or_else(|| content_coderef_path(block))
1827 {
1828 let should_replace = match paths_to_id.get(&path) {
1829 Some(existing_id) => {
1830 let existing_rank = path_selector_rank(
1831 node_classes
1832 .get(existing_id)
1833 .map(String::as_str)
1834 .unwrap_or("unknown"),
1835 );
1836 let next_rank = path_selector_rank(
1837 node_classes
1838 .get(block_id)
1839 .map(String::as_str)
1840 .unwrap_or("unknown"),
1841 );
1842 next_rank < existing_rank
1843 }
1844 None => true,
1845 };
1846 if should_replace {
1847 paths_to_id.insert(path, *block_id);
1848 }
1849 }
1850 if let Some(display) =
1851 metadata_coderef_display(block).or_else(|| content_coderef_display(block))
1852 {
1853 display_to_id.insert(display, *block_id);
1854 }
1855 let content_name = content_string(block, "name");
1856 if let Some(symbol_name) = block
1857 .metadata
1858 .custom
1859 .get(META_SYMBOL_NAME)
1860 .and_then(Value::as_str)
1861 .or(content_name.as_deref())
1862 {
1863 symbol_names_to_id
1864 .entry(symbol_name.to_string())
1865 .or_default()
1866 .push(*block_id);
1867 }
1868 }
1869
1870 for (source, block) in &doc.blocks {
1871 for edge in &block.edges {
1872 let relation = edge_type_to_string(&edge.edge_type);
1873 outgoing.entry(*source).or_default().push(IndexedEdge {
1874 other: edge.target,
1875 relation: relation.clone(),
1876 });
1877 incoming.entry(edge.target).or_default().push(IndexedEdge {
1878 other: *source,
1879 relation,
1880 });
1881 }
1882 }
1883
1884 for (parent, children) in &doc.structure {
1885 let parent_class = node_classes
1886 .get(parent)
1887 .map(String::as_str)
1888 .unwrap_or("unknown");
1889 for child in children {
1890 let child_class = node_classes
1891 .get(child)
1892 .map(String::as_str)
1893 .unwrap_or("unknown");
1894 if parent_class == "file" && child_class == "symbol" {
1895 file_symbols.entry(*parent).or_default().push(*child);
1896 }
1897 if parent_class == "symbol" && child_class == "symbol" {
1898 symbol_children.entry(*parent).or_default().push(*child);
1899 }
1900 structure_parent.insert(*child, *parent);
1901 }
1902 }
1903
1904 Self {
1905 logical_keys,
1906 logical_key_to_id,
1907 paths_to_id,
1908 display_to_id,
1909 symbol_names_to_id,
1910 node_classes,
1911 outgoing,
1912 incoming,
1913 file_symbols,
1914 symbol_children,
1915 structure_parent,
1916 }
1917 }
1918
1919 fn resolve_selector(&self, selector: &str) -> Option<BlockId> {
1920 BlockId::from_str(selector)
1921 .ok()
1922 .or_else(|| self.logical_key_to_id.get(selector).copied())
1923 .or_else(|| self.paths_to_id.get(selector).copied())
1924 .or_else(|| self.display_to_id.get(selector).copied())
1925 .or_else(|| {
1926 self.symbol_names_to_id.get(selector).and_then(|ids| {
1927 if ids.len() == 1 {
1928 ids.first().copied()
1929 } else {
1930 None
1931 }
1932 })
1933 })
1934 }
1935
1936 fn overview_nodes(&self, doc: &Document, max_depth: Option<usize>) -> Vec<BlockId> {
1937 let mut nodes = Vec::new();
1938 let limit = max_depth.unwrap_or(usize::MAX);
1939 let mut queue = VecDeque::from([(doc.root, 0usize)]);
1940 let mut visited = HashSet::new();
1941 while let Some((block_id, depth)) = queue.pop_front() {
1942 if !visited.insert(block_id) {
1943 continue;
1944 }
1945 let class = self
1946 .node_classes
1947 .get(&block_id)
1948 .map(String::as_str)
1949 .unwrap_or("unknown");
1950 if matches!(class, "repository" | "directory" | "file") {
1951 nodes.push(block_id);
1952 }
1953 if depth >= limit {
1954 continue;
1955 }
1956 for child in doc.children(&block_id) {
1957 let child_class = self
1958 .node_classes
1959 .get(child)
1960 .map(String::as_str)
1961 .unwrap_or("unknown");
1962 if matches!(child_class, "repository" | "directory" | "file") {
1963 queue.push_back((*child, depth + 1));
1964 }
1965 }
1966 }
1967 nodes.sort_by_key(|block_id| {
1968 self.logical_keys
1969 .get(block_id)
1970 .cloned()
1971 .unwrap_or_else(|| block_id.to_string())
1972 });
1973 nodes
1974 }
1975
1976 fn outgoing_edges(&self, block_id: &BlockId) -> Vec<IndexedEdge> {
1977 self.outgoing.get(block_id).cloned().unwrap_or_default()
1978 }
1979
1980 fn incoming_edges(&self, block_id: &BlockId) -> Vec<IndexedEdge> {
1981 self.incoming.get(block_id).cloned().unwrap_or_default()
1982 }
1983
1984 fn file_symbols(&self, block_id: &BlockId) -> Vec<BlockId> {
1985 let mut symbols = self.file_symbols.get(block_id).cloned().unwrap_or_default();
1986 symbols.sort_by_key(|id| {
1987 self.logical_keys
1988 .get(id)
1989 .cloned()
1990 .unwrap_or_else(|| id.to_string())
1991 });
1992 symbols
1993 }
1994
1995 fn symbol_children(&self, block_id: &BlockId) -> Vec<BlockId> {
1996 let mut children = self
1997 .symbol_children
1998 .get(block_id)
1999 .cloned()
2000 .unwrap_or_default();
2001 children.sort_by_key(|id| {
2002 self.logical_keys
2003 .get(id)
2004 .cloned()
2005 .unwrap_or_else(|| id.to_string())
2006 });
2007 children
2008 }
2009
2010 fn descendants(&self, block_id: BlockId) -> Vec<BlockId> {
2011 let mut out = Vec::new();
2012 let mut queue: VecDeque<BlockId> = self
2013 .symbol_children
2014 .get(&block_id)
2015 .cloned()
2016 .unwrap_or_default()
2017 .into();
2018 while let Some(next) = queue.pop_front() {
2019 out.push(next);
2020 if let Some(children) = self.symbol_children.get(&next) {
2021 for child in children {
2022 queue.push_back(*child);
2023 }
2024 }
2025 }
2026 out
2027 }
2028
2029 fn node_class(&self, block_id: &BlockId) -> Option<&str> {
2030 self.node_classes.get(block_id).map(String::as_str)
2031 }
2032
2033 fn structure_parent(&self, block_id: &BlockId) -> Option<BlockId> {
2034 self.structure_parent.get(block_id).copied()
2035 }
2036
2037 fn total_symbols(&self) -> usize {
2038 self.node_classes
2039 .values()
2040 .filter(|class| class.as_str() == "symbol")
2041 .count()
2042 }
2043
2044 fn display_label(&self, doc: &Document, block_id: &BlockId) -> Option<String> {
2045 let block = doc.get_block(block_id)?;
2046 match self.node_class(block_id) {
2047 Some("file") | Some("directory") | Some("repository") => metadata_coderef_path(block)
2048 .or_else(|| content_coderef_path(block))
2049 .or_else(|| block_logical_key(block)),
2050 Some("symbol") => block_logical_key(block)
2051 .or_else(|| metadata_coderef_display(block))
2052 .or_else(|| content_coderef_display(block)),
2053 _ => block_logical_key(block),
2054 }
2055 }
2056}
2057
2058pub fn is_codegraph_document(doc: &Document) -> bool {
2059 let profile = doc.metadata.custom.get("profile").and_then(Value::as_str);
2060 let marker = doc
2061 .metadata
2062 .custom
2063 .get("profile_marker")
2064 .and_then(Value::as_str);
2065
2066 profile == Some("codegraph") || marker == Some(CODEGRAPH_PROFILE_MARKER)
2067}
2068
2069pub fn resolve_codegraph_selector(doc: &Document, selector: &str) -> Option<BlockId> {
2070 CodeGraphQueryIndex::new(doc).resolve_selector(selector)
2071}
2072
2073pub fn render_codegraph_context_prompt(
2074 doc: &Document,
2075 session: &CodeGraphContextSession,
2076 config: &CodeGraphRenderConfig,
2077) -> String {
2078 session.render_for_prompt(doc, config)
2079}
2080
2081pub fn export_codegraph_context(
2082 doc: &Document,
2083 session: &CodeGraphContextSession,
2084 config: &CodeGraphRenderConfig,
2085) -> CodeGraphContextExport {
2086 session.export(doc, config)
2087}
2088
2089pub fn export_codegraph_context_with_config(
2090 doc: &Document,
2091 session: &CodeGraphContextSession,
2092 render_config: &CodeGraphRenderConfig,
2093 export_config: &CodeGraphExportConfig,
2094) -> CodeGraphContextExport {
2095 session.export_with_config(doc, render_config, export_config)
2096}
2097
2098pub fn approximate_prompt_tokens(rendered: &str) -> u32 {
2099 ((rendered.len() as f32) / 4.0).ceil() as u32
2100}
2101
2102fn origin_is_more_protective(
2103 next: Option<&CodeGraphSelectionOrigin>,
2104 current: Option<&CodeGraphSelectionOrigin>,
2105) -> bool {
2106 match (next, current) {
2107 (Some(next), Some(current)) => {
2108 selection_origin_protection_rank(next) < selection_origin_protection_rank(current)
2109 }
2110 (Some(_), None) => true,
2111 _ => false,
2112 }
2113}
2114
2115fn selection_origin_protection_rank(origin: &CodeGraphSelectionOrigin) -> u8 {
2116 match origin.kind {
2117 CodeGraphSelectionOriginKind::Manual => 0,
2118 CodeGraphSelectionOriginKind::Overview => 1,
2119 CodeGraphSelectionOriginKind::FileSymbols => 2,
2120 CodeGraphSelectionOriginKind::Dependencies => 3,
2121 CodeGraphSelectionOriginKind::Dependents => 4,
2122 }
2123}
2124
2125fn origin_prune_rank(
2126 origin: Option<&CodeGraphSelectionOrigin>,
2127 policy: &CodeGraphPrunePolicy,
2128) -> u8 {
2129 let _ = policy;
2130 match origin.map(|origin| origin.kind) {
2131 Some(CodeGraphSelectionOriginKind::Dependents) => 5,
2132 Some(CodeGraphSelectionOriginKind::Dependencies) => 4,
2133 Some(CodeGraphSelectionOriginKind::FileSymbols) => 2,
2134 Some(CodeGraphSelectionOriginKind::Overview) => 1,
2135 Some(CodeGraphSelectionOriginKind::Manual) => 0,
2136 None => 0,
2137 }
2138}
2139
2140fn relation_prune_rank(
2141 origin: Option<&CodeGraphSelectionOrigin>,
2142 policy: &CodeGraphPrunePolicy,
2143) -> u8 {
2144 origin
2145 .and_then(|origin| origin.relation.as_ref())
2146 .and_then(|relation| policy.relation_prune_priority.get(relation).copied())
2147 .unwrap_or(0)
2148}
2149
2150fn push_unique(ids: &mut Vec<BlockId>, block_id: BlockId) {
2151 if !ids.contains(&block_id) {
2152 ids.push(block_id);
2153 }
2154}
2155
2156fn default_one() -> usize {
2157 1
2158}
2159
2160fn prune_removal_rank(node_class: &str) -> u8 {
2161 match node_class {
2162 "symbol" => 4,
2163 "file" => 3,
2164 "directory" => 2,
2165 "repository" => 1,
2166 _ => 0,
2167 }
2168}
2169
2170fn focus_preference_rank(node_class: &str) -> u8 {
2171 match node_class {
2172 "symbol" => 0,
2173 "file" => 1,
2174 "directory" => 2,
2175 "repository" => 3,
2176 _ => 4,
2177 }
2178}
2179
2180fn path_selector_rank(node_class: &str) -> u8 {
2181 match node_class {
2182 "file" => 0,
2183 "directory" => 1,
2184 "repository" => 2,
2185 "symbol" => 3,
2186 _ => 4,
2187 }
2188}
2189
2190#[allow(clippy::too_many_arguments)]
2191fn render_edge_section(
2192 out: &mut String,
2193 label: &str,
2194 edges: Vec<IndexedEdge>,
2195 selected_ids: &HashSet<BlockId>,
2196 short_ids: &HashMap<BlockId, String>,
2197 doc: &Document,
2198 index: &CodeGraphQueryIndex,
2199 limit: usize,
2200) {
2201 let visible = dedupe_visible_edges(edges, selected_ids);
2202
2203 if visible.is_empty() {
2204 return;
2205 }
2206
2207 let _ = writeln!(out, " {}:", label);
2208 for (edge, multiplicity) in visible.iter().take(limit) {
2209 let short = short_ids
2210 .get(&edge.other)
2211 .cloned()
2212 .unwrap_or_else(|| edge.other.to_string());
2213 let target = index
2214 .display_label(doc, &edge.other)
2215 .unwrap_or_else(|| edge.other.to_string());
2216 let suffix = if *multiplicity > 1 {
2217 format!(" (x{})", multiplicity)
2218 } else {
2219 String::new()
2220 };
2221 let _ = writeln!(
2222 out,
2223 " - {} -> [{}] {}{}",
2224 edge.relation, short, target, suffix
2225 );
2226 }
2227
2228 if visible.len() > limit {
2229 let _ = writeln!(out, " - ... {} more", visible.len() - limit);
2230 }
2231}
2232
2233#[allow(clippy::too_many_arguments)]
2234fn append_relation_frontier(
2235 out: &mut Vec<CodeGraphContextFrontierAction>,
2236 block_id: BlockId,
2237 short_id: &str,
2238 label: &str,
2239 edges: Vec<IndexedEdge>,
2240 selected_ids: &HashSet<BlockId>,
2241 action: &str,
2242 direction: &str,
2243) {
2244 let mut counts: BTreeMap<String, usize> = BTreeMap::new();
2245 for edge in edges {
2246 if selected_ids.contains(&edge.other) {
2247 continue;
2248 }
2249 *counts.entry(edge.relation).or_default() += 1;
2250 }
2251 for (relation, candidate_count) in counts {
2252 let low_value = is_low_value_relation(action, relation.as_str());
2253 out.push(CodeGraphContextFrontierAction {
2254 block_id,
2255 short_id: short_id.to_string(),
2256 action: action.to_string(),
2257 relation: Some(relation.clone()),
2258 direction: Some(direction.to_string()),
2259 candidate_count,
2260 priority: frontier_priority(
2261 action,
2262 Some(relation.as_str()),
2263 candidate_count,
2264 low_value,
2265 ),
2266 description: format!(
2267 "{} {} neighbors via {} for {}",
2268 action, direction, relation, label
2269 ),
2270 });
2271 }
2272}
2273
2274fn is_low_value_relation(action: &str, relation: &str) -> bool {
2275 matches!(action, "expand_dependents")
2276 || relation == "references"
2277 || relation == "cited_by"
2278 || relation == "links_to"
2279}
2280
2281fn dedupe_visible_edges(
2282 edges: Vec<IndexedEdge>,
2283 selected_ids: &HashSet<BlockId>,
2284) -> Vec<(IndexedEdge, usize)> {
2285 let mut counts: HashMap<(BlockId, String), usize> = HashMap::new();
2286 for edge in edges {
2287 if !selected_ids.contains(&edge.other) {
2288 continue;
2289 }
2290 *counts.entry((edge.other, edge.relation)).or_default() += 1;
2291 }
2292 let mut deduped: Vec<_> = counts
2293 .into_iter()
2294 .map(|((other, relation), multiplicity)| (IndexedEdge { other, relation }, multiplicity))
2295 .collect();
2296 deduped.sort_by_key(|(edge, _)| (edge.relation.clone(), edge.other.to_string()));
2297 deduped
2298}
2299
2300fn export_edges(
2301 index: &CodeGraphQueryIndex,
2302 selected_ids: &HashSet<BlockId>,
2303 short_ids: &HashMap<BlockId, String>,
2304 export_config: &CodeGraphExportConfig,
2305) -> (Vec<CodeGraphContextEdgeExport>, usize) {
2306 let mut edges = Vec::new();
2307 let mut total_selected_edges = 0;
2308
2309 if export_config.dedupe_edges {
2310 let mut counts: HashMap<(BlockId, String, BlockId), usize> = HashMap::new();
2311 for source in selected_ids.iter().copied() {
2312 for edge in index.outgoing_edges(&source) {
2313 if !selected_ids.contains(&edge.other) {
2314 continue;
2315 }
2316 total_selected_edges += 1;
2317 *counts
2318 .entry((source, edge.relation, edge.other))
2319 .or_default() += 1;
2320 }
2321 }
2322 for ((source, relation, target), multiplicity) in counts {
2323 edges.push(CodeGraphContextEdgeExport {
2324 source,
2325 source_short_id: short_ids
2326 .get(&source)
2327 .cloned()
2328 .unwrap_or_else(|| source.to_string()),
2329 target,
2330 target_short_id: short_ids
2331 .get(&target)
2332 .cloned()
2333 .unwrap_or_else(|| target.to_string()),
2334 relation,
2335 multiplicity,
2336 });
2337 }
2338 } else {
2339 for source in selected_ids.iter().copied() {
2340 for edge in index.outgoing_edges(&source) {
2341 if !selected_ids.contains(&edge.other) {
2342 continue;
2343 }
2344 total_selected_edges += 1;
2345 edges.push(CodeGraphContextEdgeExport {
2346 source,
2347 source_short_id: short_ids
2348 .get(&source)
2349 .cloned()
2350 .unwrap_or_else(|| source.to_string()),
2351 target: edge.other,
2352 target_short_id: short_ids
2353 .get(&edge.other)
2354 .cloned()
2355 .unwrap_or_else(|| edge.other.to_string()),
2356 relation: edge.relation,
2357 multiplicity: 1,
2358 });
2359 }
2360 }
2361 }
2362
2363 edges.sort_by_key(|edge| {
2364 (
2365 edge.source_short_id.clone(),
2366 edge.relation.clone(),
2367 edge.target_short_id.clone(),
2368 )
2369 });
2370 (edges, total_selected_edges)
2371}
2372
2373fn focus_distances(
2374 doc: &Document,
2375 focus: Option<BlockId>,
2376 selected_ids: &HashSet<BlockId>,
2377 index: &CodeGraphQueryIndex,
2378) -> HashMap<BlockId, usize> {
2379 let mut distances = HashMap::new();
2380 let Some(focus) = focus else {
2381 return distances;
2382 };
2383 if !selected_ids.contains(&focus) {
2384 return distances;
2385 }
2386
2387 let mut queue = VecDeque::from([(focus, 0usize)]);
2388 distances.insert(focus, 0);
2389 while let Some((block_id, distance)) = queue.pop_front() {
2390 let mut neighbors: Vec<BlockId> = index
2391 .outgoing_edges(&block_id)
2392 .into_iter()
2393 .chain(index.incoming_edges(&block_id).into_iter())
2394 .map(|edge| edge.other)
2395 .collect();
2396 neighbors.extend(doc.children(&block_id));
2397 if let Some(parent) = index.structure_parent(&block_id) {
2398 neighbors.push(parent);
2399 }
2400 for neighbor in neighbors {
2401 if !selected_ids.contains(&neighbor) || distances.contains_key(&neighbor) {
2402 continue;
2403 }
2404 distances.insert(neighbor, distance + 1);
2405 queue.push_back((neighbor, distance + 1));
2406 }
2407 }
2408 distances
2409}
2410
2411fn visible_selected_ids(
2412 focus: Option<BlockId>,
2413 selected_ids: &HashSet<BlockId>,
2414 distances: &HashMap<BlockId, usize>,
2415 visible_levels: Option<usize>,
2416) -> HashSet<BlockId> {
2417 match (focus, visible_levels) {
2418 (Some(_), Some(levels)) => selected_ids
2419 .iter()
2420 .copied()
2421 .filter(|block_id| distances.get(block_id).copied().unwrap_or(usize::MAX) <= levels)
2422 .collect(),
2423 _ => selected_ids.clone(),
2424 }
2425}
2426
2427fn class_filtered_selected_ids(
2428 index: &CodeGraphQueryIndex,
2429 selected_ids: &HashSet<BlockId>,
2430 export_config: &CodeGraphExportConfig,
2431) -> HashSet<BlockId> {
2432 selected_ids
2433 .iter()
2434 .copied()
2435 .filter(|block_id| {
2436 node_class_visible(
2437 index.node_class(block_id).unwrap_or("unknown"),
2438 export_config,
2439 )
2440 })
2441 .collect()
2442}
2443
2444fn hidden_level_summaries(
2445 session: &CodeGraphContextSession,
2446 index: &CodeGraphQueryIndex,
2447 selected_ids: &HashSet<BlockId>,
2448 visible_selected_ids: &HashSet<BlockId>,
2449 distances: &HashMap<BlockId, usize>,
2450 visible_levels: Option<usize>,
2451) -> Vec<CodeGraphHiddenLevelSummary> {
2452 let Some(levels) = visible_levels else {
2453 return Vec::new();
2454 };
2455 let mut counts: BTreeMap<(usize, Option<String>, Option<String>), usize> = BTreeMap::new();
2456 for block_id in selected_ids {
2457 if visible_selected_ids.contains(block_id) {
2458 continue;
2459 }
2460 let Some(distance) = distances.get(block_id).copied() else {
2461 continue;
2462 };
2463 if distance > levels {
2464 let (relation, direction) = hidden_summary_metadata(session, index, *block_id);
2465 *counts.entry((distance, relation, direction)).or_default() += 1;
2466 }
2467 }
2468 counts
2469 .into_iter()
2470 .map(
2471 |((level, relation, direction), count)| CodeGraphHiddenLevelSummary {
2472 level,
2473 count,
2474 relation,
2475 direction,
2476 },
2477 )
2478 .collect()
2479}
2480
2481fn hidden_summary_metadata(
2482 session: &CodeGraphContextSession,
2483 index: &CodeGraphQueryIndex,
2484 block_id: BlockId,
2485) -> (Option<String>, Option<String>) {
2486 let Some(node) = session.selected.get(&block_id) else {
2487 return (None, None);
2488 };
2489 match node.origin.as_ref() {
2490 Some(origin) if origin.kind == CodeGraphSelectionOriginKind::Dependencies => {
2491 (origin.relation.clone(), Some("outgoing".to_string()))
2492 }
2493 Some(origin) if origin.kind == CodeGraphSelectionOriginKind::Dependents => {
2494 (origin.relation.clone(), Some("incoming".to_string()))
2495 }
2496 Some(origin) if origin.kind == CodeGraphSelectionOriginKind::FileSymbols => (
2497 Some("contains_symbol".to_string()),
2498 Some("structural".to_string()),
2499 ),
2500 Some(origin) if origin.kind == CodeGraphSelectionOriginKind::Overview => (
2501 Some("structure".to_string()),
2502 Some("structural".to_string()),
2503 ),
2504 Some(origin) if origin.kind == CodeGraphSelectionOriginKind::Manual => {
2505 (origin.relation.clone(), Some("manual".to_string()))
2506 }
2507 _ => match index.node_class(&block_id).unwrap_or("unknown") {
2508 "repository" | "directory" | "file" => (
2509 Some("structure".to_string()),
2510 Some("structural".to_string()),
2511 ),
2512 _ => (None, None),
2513 },
2514 }
2515}
2516
2517fn node_class_visible(node_class: &str, export_config: &CodeGraphExportConfig) -> bool {
2518 let only_matches = export_config.only_node_classes.is_empty()
2519 || export_config
2520 .only_node_classes
2521 .iter()
2522 .any(|allowed| allowed == node_class);
2523 let excluded = export_config
2524 .exclude_node_classes
2525 .iter()
2526 .any(|excluded| excluded == node_class);
2527 only_matches && !excluded
2528}
2529
2530fn relation_matches(relation_filters: Option<&HashSet<String>>, relation: &str) -> bool {
2531 relation_filters
2532 .map(|filters| filters.contains(relation))
2533 .unwrap_or(true)
2534}
2535
2536fn join_relation_filters(relation_filters: &HashSet<String>) -> Option<&str> {
2537 if relation_filters.len() == 1 {
2538 relation_filters.iter().next().map(String::as_str)
2539 } else {
2540 None
2541 }
2542}
2543
2544fn join_relation_filter_string(relation_filters: &HashSet<String>) -> String {
2545 let mut filters: Vec<_> = relation_filters.iter().cloned().collect();
2546 filters.sort();
2547 filters.join(",")
2548}
2549
2550fn relevance_score_for_node(
2551 session: &CodeGraphContextSession,
2552 index: &CodeGraphQueryIndex,
2553 block_id: BlockId,
2554 distance_from_focus: Option<usize>,
2555) -> u16 {
2556 let Some(node) = session.selected.get(&block_id) else {
2557 return 0;
2558 };
2559 let mut score = 0u16;
2560 if session.focus == Some(block_id) {
2561 score += 100;
2562 }
2563 if node.pinned {
2564 score += 40;
2565 }
2566 score += match index.node_class(&block_id).unwrap_or("unknown") {
2567 "symbol" => 40,
2568 "file" => 28,
2569 "directory" => 16,
2570 "repository" => 10,
2571 _ => 6,
2572 };
2573 score += match node.detail_level {
2574 CodeGraphDetailLevel::Source => 30,
2575 CodeGraphDetailLevel::Neighborhood => 20,
2576 CodeGraphDetailLevel::SymbolCard => 12,
2577 CodeGraphDetailLevel::Skeleton => 4,
2578 };
2579 score += match node.origin.as_ref().map(|origin| origin.kind) {
2580 Some(CodeGraphSelectionOriginKind::Manual) => 24,
2581 Some(CodeGraphSelectionOriginKind::Overview) => 18,
2582 Some(CodeGraphSelectionOriginKind::FileSymbols) => 16,
2583 Some(CodeGraphSelectionOriginKind::Dependencies) => 12,
2584 Some(CodeGraphSelectionOriginKind::Dependents) => 8,
2585 None => 0,
2586 };
2587 score += match distance_from_focus {
2588 Some(0) => 30,
2589 Some(1) => 20,
2590 Some(2) => 10,
2591 Some(3) => 4,
2592 Some(_) => 1,
2593 None => 0,
2594 };
2595 score
2596}
2597
2598fn should_include_docs(
2599 export_config: &CodeGraphExportConfig,
2600 focus: Option<BlockId>,
2601 block_id: BlockId,
2602 node: &CodeGraphContextNode,
2603 distance_from_focus: Option<usize>,
2604) -> bool {
2605 match export_config.mode {
2606 CodeGraphExportMode::Full => true,
2607 CodeGraphExportMode::Compact => {
2608 focus == Some(block_id) || node.pinned || distance_from_focus.unwrap_or(usize::MAX) <= 1
2609 }
2610 }
2611}
2612
2613fn should_include_hydrated_source(
2614 export_config: &CodeGraphExportConfig,
2615 focus: Option<BlockId>,
2616 block_id: BlockId,
2617 node: &CodeGraphContextNode,
2618 distance_from_focus: Option<usize>,
2619) -> bool {
2620 if node.hydrated_source.is_none() {
2621 return false;
2622 }
2623 match export_config.mode {
2624 CodeGraphExportMode::Full => true,
2625 CodeGraphExportMode::Compact => {
2626 focus == Some(block_id)
2627 || (node.pinned && distance_from_focus.unwrap_or(usize::MAX) <= 1)
2628 }
2629 }
2630}
2631
2632fn frontier_priority(
2633 action: &str,
2634 relation: Option<&str>,
2635 candidate_count: usize,
2636 low_value: bool,
2637) -> u16 {
2638 let base = match action {
2639 "hydrate_source" => 120,
2640 "expand_file" => 100,
2641 "expand_dependencies" => 85,
2642 "expand_dependents" => 70,
2643 "collapse" => 5,
2644 _ => 20,
2645 };
2646 let relation_adjust = match relation {
2647 Some("references") | Some("cited_by") => -20,
2648 Some("links_to") => -12,
2649 Some("uses_symbol") => 8,
2650 Some("imports_symbol") => 6,
2651 Some("reexports_symbol") => 4,
2652 Some("calls") => 6,
2653 _ => 0,
2654 };
2655 let low_value_adjust = if low_value { -10 } else { 0 };
2656 let count_bonus = candidate_count.min(12) as i32;
2657 (base + relation_adjust + low_value_adjust + count_bonus).max(0) as u16
2658}
2659
2660fn make_short_ids(
2661 session: &CodeGraphContextSession,
2662 index: &CodeGraphQueryIndex,
2663) -> HashMap<BlockId, String> {
2664 let mut by_class: BTreeMap<&str, Vec<BlockId>> = BTreeMap::new();
2665 for block_id in session.selected.keys().copied() {
2666 by_class
2667 .entry(index.node_class(&block_id).unwrap_or("node"))
2668 .or_default()
2669 .push(block_id);
2670 }
2671
2672 let mut result = HashMap::new();
2673 for (class, ids) in by_class {
2674 let mut ids = ids;
2675 ids.sort_by_key(|block_id| {
2676 index
2677 .logical_keys
2678 .get(block_id)
2679 .cloned()
2680 .unwrap_or_else(|| block_id.to_string())
2681 });
2682 for (idx, block_id) in ids.into_iter().enumerate() {
2683 let prefix = match class {
2684 "repository" => "R",
2685 "directory" => "D",
2686 "file" => "F",
2687 "symbol" => "S",
2688 _ => "N",
2689 };
2690 result.insert(block_id, format!("{}{}", prefix, idx + 1));
2691 }
2692 }
2693 result
2694}
2695
2696fn render_reference(
2697 doc: &Document,
2698 index: &CodeGraphQueryIndex,
2699 short_ids: &HashMap<BlockId, String>,
2700 block_id: BlockId,
2701) -> Option<String> {
2702 Some(format!(
2703 "[{}] {}",
2704 short_ids.get(&block_id)?.clone(),
2705 index.display_label(doc, &block_id)?
2706 ))
2707}
2708
2709fn edge_type_to_string(edge_type: &ucm_core::EdgeType) -> String {
2710 match edge_type {
2711 ucm_core::EdgeType::DerivedFrom => "derived_from".to_string(),
2712 ucm_core::EdgeType::Supersedes => "supersedes".to_string(),
2713 ucm_core::EdgeType::TransformedFrom => "transformed_from".to_string(),
2714 ucm_core::EdgeType::References => "references".to_string(),
2715 ucm_core::EdgeType::CitedBy => "cited_by".to_string(),
2716 ucm_core::EdgeType::LinksTo => "links_to".to_string(),
2717 ucm_core::EdgeType::Supports => "supports".to_string(),
2718 ucm_core::EdgeType::Contradicts => "contradicts".to_string(),
2719 ucm_core::EdgeType::Elaborates => "elaborates".to_string(),
2720 ucm_core::EdgeType::Summarizes => "summarizes".to_string(),
2721 ucm_core::EdgeType::ParentOf => "parent_of".to_string(),
2722 ucm_core::EdgeType::SiblingOf => "sibling_of".to_string(),
2723 ucm_core::EdgeType::PreviousSibling => "previous_sibling".to_string(),
2724 ucm_core::EdgeType::NextSibling => "next_sibling".to_string(),
2725 ucm_core::EdgeType::VersionOf => "version_of".to_string(),
2726 ucm_core::EdgeType::AlternativeOf => "alternative_of".to_string(),
2727 ucm_core::EdgeType::TranslationOf => "translation_of".to_string(),
2728 ucm_core::EdgeType::ChildOf => "child_of".to_string(),
2729 ucm_core::EdgeType::Custom(name) => name.clone(),
2730 }
2731}
2732
2733fn hydrate_source_excerpt(
2734 doc: &Document,
2735 block_id: BlockId,
2736 padding: usize,
2737) -> Result<Option<HydratedSourceExcerpt>, String> {
2738 let Some(block) = doc.get_block(&block_id) else {
2739 return Err(format!("block not found: {}", block_id));
2740 };
2741 let coderef =
2742 block_coderef(block).ok_or_else(|| format!("missing coderef for {}", block_id))?;
2743 let repo =
2744 repository_root(doc).ok_or_else(|| "missing repository_path metadata".to_string())?;
2745 #[cfg(target_arch = "wasm32")]
2746 {
2747 let _ = (repo, coderef, padding);
2748 Err("source hydration is not available on wasm32".to_string())
2749 }
2750 #[cfg(not(target_arch = "wasm32"))]
2751 {
2752 let path = repo.join(&coderef.path);
2753 let source = std::fs::read_to_string(&path)
2754 .map_err(|error| format!("failed to read {}: {}", path.display(), error))?;
2755 let lines: Vec<_> = source.lines().collect();
2756 let line_count = lines.len().max(1);
2757 let start_line = coderef.start_line.unwrap_or(1).max(1);
2758 let end_line = coderef
2759 .end_line
2760 .unwrap_or(start_line)
2761 .max(start_line)
2762 .min(line_count);
2763 let slice_start = start_line.saturating_sub(padding + 1);
2764 let slice_end = (end_line + padding).min(line_count);
2765
2766 let mut snippet = String::new();
2767 for (idx, line) in lines[slice_start..slice_end].iter().enumerate() {
2768 let number = slice_start + idx + 1;
2769 let _ = writeln!(snippet, "{:>4} | {}", number, line);
2770 }
2771
2772 Ok(Some(HydratedSourceExcerpt {
2773 path: coderef.path,
2774 display: coderef.display,
2775 start_line,
2776 end_line,
2777 snippet: snippet.trim_end().to_string(),
2778 }))
2779 }
2780}
2781
2782fn repository_root(doc: &Document) -> Option<PathBuf> {
2783 doc.metadata
2784 .custom
2785 .get("repository_path")
2786 .and_then(Value::as_str)
2787 .map(PathBuf::from)
2788}
2789
2790#[derive(Debug, Clone)]
2791struct BlockCoderef {
2792 path: String,
2793 display: String,
2794 start_line: Option<usize>,
2795 end_line: Option<usize>,
2796}
2797
2798fn block_coderef(block: &Block) -> Option<BlockCoderef> {
2799 let value = block
2800 .metadata
2801 .custom
2802 .get(META_CODEREF)
2803 .or_else(|| match &block.content {
2804 Content::Json { value, .. } => value.get("coderef"),
2805 _ => None,
2806 })?;
2807
2808 Some(BlockCoderef {
2809 path: value.get("path")?.as_str()?.to_string(),
2810 display: value
2811 .get("display")
2812 .and_then(Value::as_str)
2813 .unwrap_or_else(|| {
2814 value
2815 .get("path")
2816 .and_then(Value::as_str)
2817 .unwrap_or("unknown")
2818 })
2819 .to_string(),
2820 start_line: value
2821 .get("start_line")
2822 .and_then(Value::as_u64)
2823 .map(|v| v as usize),
2824 end_line: value
2825 .get("end_line")
2826 .and_then(Value::as_u64)
2827 .map(|v| v as usize),
2828 })
2829}
2830
2831fn format_symbol_signature(block: &Block) -> String {
2832 let kind = content_string(block, "kind").unwrap_or_else(|| "symbol".to_string());
2833 let name = content_string(block, "name").unwrap_or_else(|| "unknown".to_string());
2834 let inputs = content_array(block, "inputs")
2835 .into_iter()
2836 .map(|value| {
2837 let name = value.get("name").and_then(Value::as_str).unwrap_or("_");
2838 match value.get("type").and_then(Value::as_str) {
2839 Some(type_name) => format!("{}: {}", name, type_name),
2840 None => name.to_string(),
2841 }
2842 })
2843 .collect::<Vec<_>>();
2844 let output = content_string(block, "output");
2845 let type_info = content_string(block, "type");
2846 match kind.as_str() {
2847 "function" | "method" => {
2848 let mut rendered = format!("{} {}({})", kind, name, inputs.join(", "));
2849 if let Some(output) = output {
2850 let _ = write!(rendered, " -> {}", output);
2851 }
2852 rendered
2853 }
2854 _ => {
2855 let mut rendered = format!("{} {}", kind, name);
2856 if let Some(type_info) = type_info {
2857 let _ = write!(rendered, " : {}", type_info);
2858 }
2859 if block
2860 .metadata
2861 .custom
2862 .get(META_EXPORTED)
2863 .and_then(Value::as_bool)
2864 .unwrap_or(false)
2865 {
2866 let _ = write!(rendered, " [exported]");
2867 }
2868 rendered
2869 }
2870 }
2871}
2872
2873fn format_symbol_modifiers(block: &Block) -> String {
2874 let Content::Json { value, .. } = &block.content else {
2875 return String::new();
2876 };
2877 let Some(modifiers) = value.get("modifiers").and_then(Value::as_object) else {
2878 return String::new();
2879 };
2880
2881 let mut parts = Vec::new();
2882 if modifiers.get("async").and_then(Value::as_bool) == Some(true) {
2883 parts.push("async".to_string());
2884 }
2885 if modifiers.get("static").and_then(Value::as_bool) == Some(true) {
2886 parts.push("static".to_string());
2887 }
2888 if modifiers.get("generator").and_then(Value::as_bool) == Some(true) {
2889 parts.push("generator".to_string());
2890 }
2891 if let Some(visibility) = modifiers.get("visibility").and_then(Value::as_str) {
2892 parts.push(visibility.to_string());
2893 }
2894
2895 if parts.is_empty() {
2896 String::new()
2897 } else {
2898 format!(" [{}]", parts.join(", "))
2899 }
2900}
2901
2902fn content_string(block: &Block, field: &str) -> Option<String> {
2903 let Content::Json { value, .. } = &block.content else {
2904 return None;
2905 };
2906 value.get(field)?.as_str().map(|value| value.to_string())
2907}
2908
2909fn content_array(block: &Block, field: &str) -> Vec<Value> {
2910 let Content::Json { value, .. } = &block.content else {
2911 return Vec::new();
2912 };
2913 value
2914 .get(field)
2915 .and_then(Value::as_array)
2916 .cloned()
2917 .unwrap_or_default()
2918}
2919
2920fn node_class(block: &Block) -> Option<String> {
2921 block
2922 .metadata
2923 .custom
2924 .get(META_NODE_CLASS)
2925 .and_then(Value::as_str)
2926 .map(|value| value.to_string())
2927}
2928
2929fn block_logical_key(block: &Block) -> Option<String> {
2930 block
2931 .metadata
2932 .custom
2933 .get(META_LOGICAL_KEY)
2934 .and_then(Value::as_str)
2935 .map(|value| value.to_string())
2936}
2937
2938fn metadata_coderef_path(block: &Block) -> Option<String> {
2939 block
2940 .metadata
2941 .custom
2942 .get(META_CODEREF)
2943 .and_then(|value| value.get("path"))
2944 .and_then(Value::as_str)
2945 .map(|value| value.to_string())
2946}
2947
2948fn metadata_coderef_display(block: &Block) -> Option<String> {
2949 block
2950 .metadata
2951 .custom
2952 .get(META_CODEREF)
2953 .and_then(|value| value.get("display"))
2954 .and_then(Value::as_str)
2955 .map(|value| value.to_string())
2956}
2957
2958fn content_coderef_path(block: &Block) -> Option<String> {
2959 let Content::Json { value, .. } = &block.content else {
2960 return None;
2961 };
2962 value
2963 .get("coderef")
2964 .and_then(|value| value.get("path"))
2965 .and_then(Value::as_str)
2966 .map(|value| value.to_string())
2967}
2968
2969fn content_coderef_display(block: &Block) -> Option<String> {
2970 let Content::Json { value, .. } = &block.content else {
2971 return None;
2972 };
2973 value
2974 .get("coderef")
2975 .and_then(|value| value.get("display"))
2976 .and_then(Value::as_str)
2977 .map(|value| value.to_string())
2978}
2979
2980#[cfg(test)]
2981mod tests {
2982 use std::fs;
2983
2984 use tempfile::tempdir;
2985
2986 use super::*;
2987 use crate::{build_code_graph, CodeGraphBuildInput, CodeGraphExtractorConfig};
2988
2989 fn build_test_graph() -> Document {
2990 let dir = tempdir().unwrap();
2991 fs::create_dir_all(dir.path().join("src")).unwrap();
2992 fs::write(
2993 dir.path().join("src/util.rs"),
2994 "pub fn util() -> i32 { 1 }\n",
2995 )
2996 .unwrap();
2997 fs::write(
2998 dir.path().join("src/lib.rs"),
2999 "mod util;\n/// Add values.\npub async fn add(a: i32, b: i32) -> i32 { util::util() + a + b }\n\npub fn sub(a: i32, b: i32) -> i32 { util::util() + a - b }\n",
3000 )
3001 .unwrap();
3002
3003 let repository_path = dir.path().to_path_buf();
3004 std::mem::forget(dir);
3005
3006 build_code_graph(&CodeGraphBuildInput {
3007 repository_path,
3008 commit_hash: "context-tests".to_string(),
3009 config: CodeGraphExtractorConfig::default(),
3010 })
3011 .unwrap()
3012 .document
3013 }
3014
3015 #[test]
3016 fn overview_expand_dependents_and_hydrate_source_work() {
3017 let doc = build_test_graph();
3018 let mut session = CodeGraphContextSession::new();
3019 let update = session.seed_overview(&doc);
3020 assert!(!update.added.is_empty());
3021 assert_eq!(session.summary(&doc).symbols, 0);
3022
3023 let file_id = resolve_codegraph_selector(&doc, "src/lib.rs").unwrap();
3024 session.expand_file(&doc, file_id);
3025 assert!(session.summary(&doc).symbols >= 1);
3026
3027 let add_id = resolve_codegraph_selector(&doc, "symbol:src/lib.rs::add").unwrap();
3028 let util_id = resolve_codegraph_selector(&doc, "symbol:src/util.rs::util").unwrap();
3029 let deps = session.expand_dependencies(&doc, add_id, Some("uses_symbol"));
3030 assert!(deps.added.contains(&util_id) || session.selected.contains_key(&util_id));
3031
3032 let dependents = session.expand_dependents(&doc, util_id, Some("uses_symbol"));
3033 assert!(dependents.added.contains(&add_id) || session.selected.contains_key(&add_id));
3034
3035 let hydrated = session.hydrate_source(&doc, add_id, 1);
3036 assert!(hydrated.changed.contains(&add_id));
3037 assert!(session
3038 .selected
3039 .get(&add_id)
3040 .and_then(|node| node.hydrated_source.as_ref())
3041 .is_some());
3042
3043 let rendered = session.render_for_prompt(&doc, &CodeGraphRenderConfig::default());
3044 assert!(rendered.contains("CodeGraph working set"));
3045 assert!(rendered.contains("expand dependents"));
3046 assert!(rendered.contains("uses_symbol"));
3047 assert!(rendered.contains("source:"));
3048 }
3049
3050 #[test]
3051 fn resolve_selector_prefers_logical_key_and_path() {
3052 let doc = build_test_graph();
3053 let file_id = resolve_codegraph_selector(&doc, "src/lib.rs").unwrap();
3054 let logical_id = resolve_codegraph_selector(&doc, "symbol:src/lib.rs::add").unwrap();
3055 let display_id = resolve_codegraph_selector(&doc, "src/lib.rs:2-2").unwrap_or(logical_id);
3056 assert!(doc.get_block(&file_id).is_some());
3057 assert!(doc.get_block(&logical_id).is_some());
3058 assert_eq!(logical_id, display_id);
3059 }
3060
3061 #[test]
3062 fn prune_policy_demotes_before_removing() {
3063 let doc = build_test_graph();
3064 let mut session = CodeGraphContextSession::new();
3065 session.set_prune_policy(CodeGraphPrunePolicy {
3066 max_selected: 10,
3067 ..CodeGraphPrunePolicy::default()
3068 });
3069
3070 let file_id = resolve_codegraph_selector(&doc, "src/lib.rs").unwrap();
3071 let add_id = resolve_codegraph_selector(&doc, "symbol:src/lib.rs::add").unwrap();
3072 session.seed_overview(&doc);
3073 session.expand_file(&doc, file_id);
3074 session.expand_dependencies(&doc, add_id, Some("uses_symbol"));
3075 session.hydrate_source(&doc, add_id, 1);
3076 assert!(session
3077 .selected
3078 .get(&add_id)
3079 .and_then(|node| node.hydrated_source.as_ref())
3080 .is_some());
3081
3082 session.set_focus(&doc, Some(file_id));
3083 let update = session.prune(&doc, Some(4));
3084 assert!(session.selected.len() <= 4);
3085 assert!(!update.changed.is_empty() || !update.removed.is_empty());
3086
3087 let rendered = session.render_for_prompt(&doc, &CodeGraphRenderConfig::default());
3088 assert!(rendered.contains("selected=4/4"));
3089 assert!(!rendered.contains("source:"));
3090 }
3091
3092 #[test]
3093 fn prune_prefers_dependents_before_file_skeletons() {
3094 let doc = build_test_graph();
3095 let mut session = CodeGraphContextSession::new();
3096 session.set_prune_policy(CodeGraphPrunePolicy {
3097 max_selected: 20,
3098 ..CodeGraphPrunePolicy::default()
3099 });
3100
3101 let util_file_id = resolve_codegraph_selector(&doc, "src/util.rs").unwrap();
3102 let util_symbol_id = resolve_codegraph_selector(&doc, "symbol:src/util.rs::util").unwrap();
3103 let lib_file_id = resolve_codegraph_selector(&doc, "src/lib.rs").unwrap();
3104 let add_id = resolve_codegraph_selector(&doc, "symbol:src/lib.rs::add").unwrap();
3105 let sub_id = resolve_codegraph_selector(&doc, "symbol:src/lib.rs::sub").unwrap();
3106
3107 session.seed_overview(&doc);
3108 session.expand_file(&doc, util_file_id);
3109 session.expand_dependents(&doc, util_symbol_id, Some("uses_symbol"));
3110 assert!(session.selected.contains_key(&add_id));
3111 assert!(session.selected.contains_key(&sub_id));
3112
3113 session.set_focus(&doc, Some(util_file_id));
3114 session.prune(&doc, Some(5));
3115 assert!(session.selected.contains_key(&lib_file_id));
3116 assert!(!session.selected.contains_key(&add_id));
3117 assert!(!session.selected.contains_key(&sub_id));
3118 }
3119
3120 #[test]
3121 fn export_includes_frontier_and_origin_metadata() {
3122 let doc = build_test_graph();
3123 let mut session = CodeGraphContextSession::new();
3124 let file_id = resolve_codegraph_selector(&doc, "src/lib.rs").unwrap();
3125 let add_id = resolve_codegraph_selector(&doc, "symbol:src/lib.rs::add").unwrap();
3126
3127 session.seed_overview(&doc);
3128 session.expand_file(&doc, file_id);
3129 session.focus = Some(add_id);
3130 let export = session.export(&doc, &CodeGraphRenderConfig::default());
3131
3132 assert_eq!(export.focus, Some(add_id));
3133 assert!(export.nodes.iter().any(|node| {
3134 node.block_id == add_id
3135 && node.symbol_name.as_deref() == Some("add")
3136 && node.path.as_deref() == Some("src/lib.rs")
3137 && node
3138 .origin
3139 .as_ref()
3140 .map(|origin| origin.kind == CodeGraphSelectionOriginKind::FileSymbols)
3141 .unwrap_or(false)
3142 }));
3143 assert!(export
3144 .frontier
3145 .iter()
3146 .any(|action| action.action == "hydrate_source"));
3147 assert!(export.heuristics.recommended_next_action.is_some());
3148 assert!(!export.heuristics.should_stop);
3149 assert!(export
3150 .frontier
3151 .iter()
3152 .any(|action| action.action == "expand_dependencies"
3153 && action.relation.as_deref() == Some("uses_symbol")));
3154 }
3155
3156 #[test]
3157 fn overview_seed_depth_limits_structural_selection() {
3158 let doc = build_test_graph();
3159 let mut shallow = CodeGraphContextSession::new();
3160 shallow.seed_overview_with_depth(&doc, Some(1));
3161 let shallow_summary = shallow.summary(&doc);
3162 assert!(shallow_summary.repositories + shallow_summary.directories >= 1);
3163 assert_eq!(shallow_summary.files, 0);
3164
3165 let mut deeper = CodeGraphContextSession::new();
3166 deeper.seed_overview_with_depth(&doc, Some(3));
3167 let deeper_summary = deeper.summary(&doc);
3168 assert!(deeper_summary.files >= 2);
3169 }
3170
3171 #[test]
3172 fn export_with_visible_levels_summarizes_hidden_nodes() {
3173 let doc = build_test_graph();
3174 let mut session = CodeGraphContextSession::new();
3175 let file_id = resolve_codegraph_selector(&doc, "src/lib.rs").unwrap();
3176 let add_id = resolve_codegraph_selector(&doc, "symbol:src/lib.rs::add").unwrap();
3177
3178 session.seed_overview(&doc);
3179 session.expand_file(&doc, file_id);
3180 session.expand_dependencies(&doc, add_id, Some("uses_symbol"));
3181 session.focus = Some(add_id);
3182
3183 let mut export_config = CodeGraphExportConfig::compact();
3184 export_config.visible_levels = Some(1);
3185 let export =
3186 session.export_with_config(&doc, &CodeGraphRenderConfig::default(), &export_config);
3187
3188 assert_eq!(export.visible_levels, Some(1));
3189 assert!(export.visible_node_count < export.summary.selected);
3190 assert!(export.hidden_levels.iter().any(|hidden| hidden.level >= 2));
3191 assert!(export
3192 .nodes
3193 .iter()
3194 .all(|node| node.distance_from_focus.unwrap_or(usize::MAX) <= 1));
3195 }
3196
3197 #[test]
3198 fn selective_multi_hop_expansion_follows_only_requested_relations() {
3199 let mut doc = build_test_graph();
3200 let mut session = CodeGraphContextSession::new();
3201 let file_id = resolve_codegraph_selector(&doc, "src/lib.rs").unwrap();
3202 let add_id = resolve_codegraph_selector(&doc, "symbol:src/lib.rs::add").unwrap();
3203 let sub_id = resolve_codegraph_selector(&doc, "symbol:src/lib.rs::sub").unwrap();
3204 let util_id = resolve_codegraph_selector(&doc, "symbol:src/util.rs::util").unwrap();
3205
3206 doc.add_edge(&add_id, ucm_core::EdgeType::References, sub_id);
3207
3208 session.seed_overview(&doc);
3209 session.expand_file(&doc, file_id);
3210
3211 let relation_filters = HashSet::from(["references".to_string()]);
3212 session.expand_dependencies_with_filters(&doc, add_id, Some(&relation_filters), 2);
3213 assert!(session.selected.contains_key(&sub_id));
3214 assert!(!session.selected.contains_key(&util_id));
3215
3216 let mut session = CodeGraphContextSession::new();
3217 session.seed_overview(&doc);
3218 session.expand_file(&doc, file_id);
3219 let relation_filters = HashSet::from(["references".to_string(), "uses_symbol".to_string()]);
3220 session.expand_dependencies_with_filters(&doc, add_id, Some(&relation_filters), 2);
3221 assert!(session.selected.contains_key(&sub_id));
3222 assert!(session.selected.contains_key(&util_id));
3223 }
3224
3225 #[test]
3226 fn traversal_budget_caps_additions_and_reports_warning() {
3227 let doc = build_test_graph();
3228 let mut session = CodeGraphContextSession::new();
3229 let file_id = resolve_codegraph_selector(&doc, "src/lib.rs").unwrap();
3230
3231 session.seed_overview(&doc);
3232 let update = session.expand_file_with_config(
3233 &doc,
3234 file_id,
3235 &CodeGraphTraversalConfig {
3236 depth: 2,
3237 max_add: Some(1),
3238 ..CodeGraphTraversalConfig::default()
3239 },
3240 );
3241
3242 assert_eq!(update.added.len(), 1);
3243 assert!(update
3244 .warnings
3245 .iter()
3246 .any(|warning| warning.contains("max_add")));
3247 }
3248
3249 #[test]
3250 fn priority_threshold_skips_low_value_relations() {
3251 let mut doc = build_test_graph();
3252 let mut session = CodeGraphContextSession::new();
3253 let file_id = resolve_codegraph_selector(&doc, "src/lib.rs").unwrap();
3254 let add_id = resolve_codegraph_selector(&doc, "symbol:src/lib.rs::add").unwrap();
3255 let util_id = resolve_codegraph_selector(&doc, "symbol:src/util.rs::util").unwrap();
3256
3257 doc.add_edge(&add_id, ucm_core::EdgeType::References, util_id);
3258
3259 session.seed_overview(&doc);
3260 session.expand_file(&doc, file_id);
3261 let update = session.expand_dependencies_with_config(
3262 &doc,
3263 add_id,
3264 &CodeGraphTraversalConfig {
3265 depth: 1,
3266 relation_filters: vec!["references".to_string()],
3267 priority_threshold: Some(80),
3268 ..CodeGraphTraversalConfig::default()
3269 },
3270 );
3271
3272 assert!(!session.selected.contains_key(&util_id));
3273 assert!(update.added.is_empty());
3274 }
3275
3276 #[test]
3277 fn export_filters_node_classes_and_includes_hidden_relation_metadata() {
3278 let doc = build_test_graph();
3279 let mut session = CodeGraphContextSession::new();
3280 let file_id = resolve_codegraph_selector(&doc, "src/lib.rs").unwrap();
3281 let add_id = resolve_codegraph_selector(&doc, "symbol:src/lib.rs::add").unwrap();
3282
3283 session.seed_overview(&doc);
3284 session.expand_file(&doc, file_id);
3285 session.expand_dependencies(&doc, add_id, Some("uses_symbol"));
3286 session.focus = Some(add_id);
3287
3288 let mut export_config = CodeGraphExportConfig::compact();
3289 export_config.visible_levels = Some(0);
3290 export_config.only_node_classes = vec!["symbol".to_string()];
3291 let export =
3292 session.export_with_config(&doc, &CodeGraphRenderConfig::default(), &export_config);
3293
3294 assert!(export.nodes.iter().all(|node| node.node_class == "symbol"));
3295 assert!(export.edges.iter().all(|edge| {
3296 export
3297 .nodes
3298 .iter()
3299 .any(|node| node.block_id == edge.source && node.node_class == "symbol")
3300 && export
3301 .nodes
3302 .iter()
3303 .any(|node| node.block_id == edge.target && node.node_class == "symbol")
3304 }));
3305 assert!(export.hidden_levels.iter().any(|hidden| {
3306 hidden.relation.as_deref() == Some("uses_symbol")
3307 && hidden.direction.as_deref() == Some("outgoing")
3308 }));
3309 }
3310
3311 #[test]
3312 fn compact_export_dedupes_edges_and_omits_rendered_text() {
3313 let mut doc = build_test_graph();
3314 let mut session = CodeGraphContextSession::new();
3315 let file_id = resolve_codegraph_selector(&doc, "src/lib.rs").unwrap();
3316 let add_id = resolve_codegraph_selector(&doc, "symbol:src/lib.rs::add").unwrap();
3317 let util_id = resolve_codegraph_selector(&doc, "symbol:src/util.rs::util").unwrap();
3318
3319 doc.add_edge(&add_id, ucm_core::EdgeType::References, util_id);
3320
3321 session.seed_overview(&doc);
3322 session.expand_file(&doc, file_id);
3323 session.expand_dependencies(&doc, add_id, Some("uses_symbol"));
3324 session.hydrate_source(&doc, add_id, 1);
3325 session.focus = Some(add_id);
3326
3327 let export = session.export_with_config(
3328 &doc,
3329 &CodeGraphRenderConfig::default(),
3330 &CodeGraphExportConfig::compact(),
3331 );
3332
3333 assert_eq!(export.export_mode, CodeGraphExportMode::Compact);
3334 assert!(export.rendered.is_empty());
3335 assert!(export.total_selected_edges >= export.edges.len());
3336 assert!(export.edges.iter().all(|edge| edge.multiplicity >= 1));
3337 assert!(export
3338 .nodes
3339 .iter()
3340 .find(|node| node.block_id == add_id)
3341 .and_then(|node| node.hydrated_source.as_ref())
3342 .is_some());
3343 }
3344
3345 #[test]
3346 fn heuristics_stop_when_focus_is_hydrated_and_frontier_is_exhausted() {
3347 let doc = build_test_graph();
3348 let mut session = CodeGraphContextSession::new();
3349 let file_id = resolve_codegraph_selector(&doc, "src/lib.rs").unwrap();
3350 let add_id = resolve_codegraph_selector(&doc, "symbol:src/lib.rs::add").unwrap();
3351
3352 session.seed_overview(&doc);
3353 session.expand_file(&doc, file_id);
3354 session.expand_dependencies(&doc, add_id, Some("uses_symbol"));
3355 session.focus = Some(add_id);
3356 let pre_hydrate = session.export(&doc, &CodeGraphRenderConfig::default());
3357 assert!(!pre_hydrate.heuristics.should_stop);
3358 assert_eq!(
3359 pre_hydrate
3360 .heuristics
3361 .recommended_next_action
3362 .as_ref()
3363 .map(|action| action.action.as_str()),
3364 Some("hydrate_source")
3365 );
3366
3367 session.hydrate_source(&doc, add_id, 1);
3368 let exhausted = session.export(&doc, &CodeGraphRenderConfig::default());
3369 assert!(exhausted.heuristics.should_stop);
3370 assert!(exhausted
3371 .heuristics
3372 .reasons
3373 .iter()
3374 .any(|reason| reason.contains("hydrated")));
3375 }
3376}