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