use std::collections::{BTreeMap, HashMap, HashSet, VecDeque};
use std::fmt::Write;
use std::path::PathBuf;
use std::str::FromStr;
use std::time::Instant;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use ucm_core::{Block, BlockId, Content, Document};
use crate::model::{
CODEGRAPH_PROFILE_MARKER, META_CODEREF, META_EXPORTED, META_LANGUAGE, META_LOGICAL_KEY,
META_NODE_CLASS, META_SYMBOL_NAME,
};
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize, Default)]
#[serde(rename_all = "snake_case")]
pub enum CodeGraphDetailLevel {
#[default]
Skeleton,
SymbolCard,
Neighborhood,
Source,
}
impl CodeGraphDetailLevel {
fn max(self, other: Self) -> Self {
std::cmp::max(self, other)
}
fn demoted(self) -> Self {
match self {
Self::Source => Self::Neighborhood,
Self::Neighborhood => Self::SymbolCard,
Self::SymbolCard => Self::Skeleton,
Self::Skeleton => Self::Skeleton,
}
}
fn includes_neighborhood(self) -> bool {
matches!(self, Self::Neighborhood | Self::Source)
}
fn includes_source(self) -> bool {
matches!(self, Self::Source)
}
}
fn default_true() -> bool {
true
}
fn default_relation_prune_priority() -> BTreeMap<String, u8> {
[
("references", 60),
("cited_by", 60),
("links_to", 55),
("uses_symbol", 35),
("imports_symbol", 30),
("reexports_symbol", 25),
("calls", 20),
("inherits", 15),
("implements", 15),
]
.into_iter()
.map(|(name, score)| (name.to_string(), score))
.collect()
}
fn selection_origin(
kind: CodeGraphSelectionOriginKind,
relation: Option<&str>,
anchor: Option<BlockId>,
) -> Option<CodeGraphSelectionOrigin> {
Some(CodeGraphSelectionOrigin {
kind,
relation: relation.map(str::to_string),
anchor,
})
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HydratedSourceExcerpt {
pub path: String,
pub display: String,
pub start_line: usize,
pub end_line: usize,
pub snippet: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CodeGraphContextNode {
pub block_id: BlockId,
pub detail_level: CodeGraphDetailLevel,
#[serde(default)]
pub pinned: bool,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub origin: Option<CodeGraphSelectionOrigin>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub hydrated_source: Option<HydratedSourceExcerpt>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum CodeGraphSelectionOriginKind {
Overview,
Manual,
FileSymbols,
Dependencies,
Dependents,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CodeGraphSelectionOrigin {
pub kind: CodeGraphSelectionOriginKind,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub relation: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub anchor: Option<BlockId>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CodeGraphPrunePolicy {
pub max_selected: usize,
#[serde(default = "default_true")]
pub demote_before_remove: bool,
#[serde(default = "default_true")]
pub protect_focus: bool,
#[serde(default = "default_relation_prune_priority")]
pub relation_prune_priority: BTreeMap<String, u8>,
}
impl Default for CodeGraphPrunePolicy {
fn default() -> Self {
Self {
max_selected: 48,
demote_before_remove: true,
protect_focus: true,
relation_prune_priority: default_relation_prune_priority(),
}
}
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct CodeGraphContextSession {
#[serde(default)]
pub selected: HashMap<BlockId, CodeGraphContextNode>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub focus: Option<BlockId>,
#[serde(default)]
pub prune_policy: CodeGraphPrunePolicy,
#[serde(default)]
pub history: Vec<String>,
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct CodeGraphContextUpdate {
#[serde(default)]
pub added: Vec<BlockId>,
#[serde(default)]
pub removed: Vec<BlockId>,
#[serde(default)]
pub changed: Vec<BlockId>,
#[serde(default)]
pub focus: Option<BlockId>,
#[serde(default)]
pub warnings: Vec<String>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub telemetry: Vec<CodeGraphSessionMutation>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CodeGraphContextSummary {
pub selected: usize,
pub max_selected: usize,
pub repositories: usize,
pub directories: usize,
pub files: usize,
pub symbols: usize,
pub hydrated_sources: usize,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CodeGraphRenderConfig {
pub max_edges_per_node: usize,
pub max_source_lines: usize,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub max_rendered_bytes: Option<usize>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub max_rendered_tokens: Option<u32>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
#[serde(rename_all = "snake_case")]
pub enum CodeGraphExportMode {
#[default]
Full,
Compact,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CodeGraphExportConfig {
#[serde(default)]
pub mode: CodeGraphExportMode,
#[serde(default = "default_true")]
pub include_rendered: bool,
#[serde(default = "default_true")]
pub dedupe_edges: bool,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub visible_levels: Option<usize>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub only_node_classes: Vec<String>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub exclude_node_classes: Vec<String>,
pub max_frontier_actions: usize,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CodeGraphTraversalConfig {
#[serde(default = "default_one")]
pub depth: usize,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub relation_filters: Vec<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub max_add: Option<usize>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub priority_threshold: Option<u16>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub budget: Option<CodeGraphOperationBudget>,
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct CodeGraphOperationBudget {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub max_depth: Option<usize>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub max_nodes_visited: Option<usize>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub max_nodes_added: Option<usize>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub max_hydrated_bytes: Option<usize>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub max_elapsed_ms: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub max_emitted_telemetry_events: Option<usize>,
}
impl Default for CodeGraphExportConfig {
fn default() -> Self {
Self {
mode: CodeGraphExportMode::Full,
include_rendered: true,
dedupe_edges: true,
visible_levels: None,
only_node_classes: Vec::new(),
exclude_node_classes: Vec::new(),
max_frontier_actions: 12,
}
}
}
impl Default for CodeGraphTraversalConfig {
fn default() -> Self {
Self {
depth: 1,
relation_filters: Vec::new(),
max_add: None,
priority_threshold: None,
budget: None,
}
}
}
impl CodeGraphExportConfig {
pub fn compact() -> Self {
Self {
mode: CodeGraphExportMode::Compact,
include_rendered: false,
dedupe_edges: true,
visible_levels: None,
only_node_classes: Vec::new(),
exclude_node_classes: Vec::new(),
max_frontier_actions: 6,
}
}
}
impl CodeGraphTraversalConfig {
fn depth(&self) -> usize {
self.depth.max(1)
}
fn relation_filter_set(&self) -> Option<HashSet<String>> {
if self.relation_filters.is_empty() {
None
} else {
Some(self.relation_filters.iter().cloned().collect())
}
}
}
fn min_optional_usize(left: Option<usize>, right: Option<usize>) -> Option<usize> {
match (left, right) {
(Some(a), Some(b)) => Some(a.min(b)),
(Some(a), None) => Some(a),
(None, Some(b)) => Some(b),
(None, None) => None,
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum CodeGraphSessionMutationKind {
Select,
Focus,
ExpandFile,
ExpandDependencies,
ExpandDependents,
Hydrate,
Collapse,
Pin,
Unpin,
Prune,
SeedOverview,
ApplyRecommendedActions,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CodeGraphSessionMutation {
pub sequence: usize,
pub kind: CodeGraphSessionMutationKind,
pub operation: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub selector: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub target_block_id: Option<BlockId>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub resolved_block_ids: Vec<BlockId>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub traversal: Option<CodeGraphTraversalConfig>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub budget: Option<CodeGraphOperationBudget>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub nodes_added: Vec<BlockId>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub nodes_removed: Vec<BlockId>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub nodes_changed: Vec<BlockId>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub focus_before: Option<BlockId>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub focus_after: Option<BlockId>,
pub elapsed_ms: u64,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub reason: Option<String>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub warnings: Vec<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CodeGraphRecommendation {
pub action_kind: String,
pub target_block_id: BlockId,
pub target_short_id: String,
pub target_label: String,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub relation_set: Vec<String>,
pub priority: u16,
pub candidate_count: usize,
pub estimated_evidence_gain: usize,
pub estimated_token_cost: u32,
pub estimated_hydration_bytes: usize,
pub explanation: String,
pub rationale: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum CodeGraphExportOmissionReason {
VisibleLevelLimit,
ClassFilter,
RenderBudget,
HydratedExcerptSupersededSummary,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CodeGraphExportOmissionDetail {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub block_id: Option<BlockId>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub short_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub label: Option<String>,
pub reason: CodeGraphExportOmissionReason,
pub explanation: String,
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct CodeGraphExportOmissionReport {
pub hidden_by_visible_levels: usize,
pub excluded_by_class_filters: usize,
pub dropped_by_render_budget: usize,
pub suppressed_by_hydrated_excerpt: usize,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub details: Vec<CodeGraphExportOmissionDetail>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "event", rename_all = "snake_case")]
pub enum CodeGraphSessionEvent {
Mutation {
mutation: Box<CodeGraphSessionMutation>,
},
Recommendation {
recommendation: Box<CodeGraphRecommendation>,
},
SessionSaved {
metadata: CodeGraphSessionPersistenceMetadata,
},
SessionLoaded {
metadata: CodeGraphSessionPersistenceMetadata,
},
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CodeGraphSessionPersistenceMetadata {
pub schema_version: String,
pub session_id: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub parent_session_id: Option<String>,
pub graph_snapshot_hash: String,
pub session_snapshot_hash: String,
pub mutation_count: usize,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CodeGraphPersistedSession {
pub metadata: CodeGraphSessionPersistenceMetadata,
pub context: CodeGraphContextSession,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub mutation_log: Vec<CodeGraphSessionMutation>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub event_log: Vec<CodeGraphSessionEvent>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CodeGraphCoderef {
pub path: String,
pub display: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub start_line: Option<usize>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub end_line: Option<usize>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CodeGraphContextNodeExport {
pub block_id: BlockId,
pub short_id: String,
pub node_class: String,
pub label: String,
pub detail_level: CodeGraphDetailLevel,
pub pinned: bool,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub distance_from_focus: Option<usize>,
pub relevance_score: u16,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub logical_key: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub symbol_name: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub path: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub signature: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub docs: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub origin: Option<CodeGraphSelectionOrigin>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub coderef: Option<CodeGraphCoderef>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub hydrated_source: Option<HydratedSourceExcerpt>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CodeGraphContextEdgeExport {
pub source: BlockId,
pub source_short_id: String,
pub target: BlockId,
pub target_short_id: String,
pub relation: String,
#[serde(default = "default_one")]
pub multiplicity: usize,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CodeGraphContextFrontierAction {
pub block_id: BlockId,
pub short_id: String,
pub action: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub relation: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub direction: Option<String>,
pub candidate_count: usize,
pub priority: u16,
pub description: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub explanation: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CodeGraphContextHeuristics {
pub should_stop: bool,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub reasons: Vec<String>,
pub hidden_candidate_count: usize,
pub low_value_candidate_count: usize,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub recommended_next_action: Option<CodeGraphContextFrontierAction>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub recommended_actions: Vec<CodeGraphContextFrontierAction>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub recommendations: Vec<CodeGraphRecommendation>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CodeGraphHiddenLevelSummary {
pub level: usize,
pub count: usize,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub relation: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub direction: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CodeGraphContextExport {
pub summary: CodeGraphContextSummary,
#[serde(default)]
pub export_mode: CodeGraphExportMode,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub visible_levels: Option<usize>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub focus: Option<BlockId>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub focus_short_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub focus_label: Option<String>,
pub visible_node_count: usize,
pub hidden_unreachable_count: usize,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub hidden_levels: Vec<CodeGraphHiddenLevelSummary>,
pub frontier: Vec<CodeGraphContextFrontierAction>,
pub heuristics: CodeGraphContextHeuristics,
pub nodes: Vec<CodeGraphContextNodeExport>,
pub edges: Vec<CodeGraphContextEdgeExport>,
pub omitted_symbol_count: usize,
pub total_selected_edges: usize,
#[serde(default)]
pub omissions: CodeGraphExportOmissionReport,
#[serde(default, skip_serializing_if = "String::is_empty")]
pub rendered: String,
}
impl Default for CodeGraphRenderConfig {
fn default() -> Self {
Self {
max_edges_per_node: 6,
max_source_lines: 12,
max_rendered_bytes: None,
max_rendered_tokens: None,
}
}
}
impl CodeGraphRenderConfig {
pub fn for_max_tokens(max_tokens: usize) -> Self {
if max_tokens <= 512 {
Self {
max_edges_per_node: 2,
max_source_lines: 4,
max_rendered_bytes: None,
max_rendered_tokens: Some(max_tokens as u32),
}
} else if max_tokens <= 1024 {
Self {
max_edges_per_node: 3,
max_source_lines: 6,
max_rendered_bytes: None,
max_rendered_tokens: Some(max_tokens as u32),
}
} else if max_tokens <= 2048 {
Self {
max_edges_per_node: 4,
max_source_lines: 8,
max_rendered_bytes: None,
max_rendered_tokens: Some(max_tokens as u32),
}
} else {
Self {
max_rendered_tokens: Some(max_tokens as u32),
..Self::default()
}
}
}
}
#[derive(Debug, Clone)]
struct IndexedEdge {
other: BlockId,
relation: String,
}
#[derive(Debug, Clone)]
struct CodeGraphQueryIndex {
logical_keys: HashMap<BlockId, String>,
logical_key_to_id: HashMap<String, BlockId>,
paths_to_id: HashMap<String, BlockId>,
display_to_id: HashMap<String, BlockId>,
symbol_names_to_id: HashMap<String, Vec<BlockId>>,
node_classes: HashMap<BlockId, String>,
outgoing: HashMap<BlockId, Vec<IndexedEdge>>,
incoming: HashMap<BlockId, Vec<IndexedEdge>>,
file_symbols: HashMap<BlockId, Vec<BlockId>>,
symbol_children: HashMap<BlockId, Vec<BlockId>>,
structure_parent: HashMap<BlockId, BlockId>,
}
impl CodeGraphContextSession {
pub fn new() -> Self {
Self::default()
}
pub fn selected_block_ids(&self) -> Vec<BlockId> {
let mut ids: Vec<_> = self.selected.keys().copied().collect();
ids.sort_by_key(BlockId::to_string);
ids
}
pub fn summary(&self, doc: &Document) -> CodeGraphContextSummary {
let index = CodeGraphQueryIndex::new(doc);
let mut summary = CodeGraphContextSummary {
selected: self.selected.len(),
max_selected: self.prune_policy.max_selected,
repositories: 0,
directories: 0,
files: 0,
symbols: 0,
hydrated_sources: 0,
};
for node in self.selected.values() {
match index.node_class(&node.block_id).unwrap_or("unknown") {
"repository" => summary.repositories += 1,
"directory" => summary.directories += 1,
"file" => summary.files += 1,
"symbol" => summary.symbols += 1,
_ => {}
}
if node.hydrated_source.is_some() {
summary.hydrated_sources += 1;
}
}
summary
}
pub fn clear(&mut self) {
self.selected.clear();
self.focus = None;
self.history.push("clear".to_string());
}
pub fn set_prune_policy(&mut self, policy: CodeGraphPrunePolicy) {
self.prune_policy = policy;
self.history.push(format!(
"policy:max_selected:{}:demote:{}:protect_focus:{}",
self.prune_policy.max_selected,
self.prune_policy.demote_before_remove,
self.prune_policy.protect_focus
));
}
pub fn set_focus(
&mut self,
doc: &Document,
block_id: Option<BlockId>,
) -> CodeGraphContextUpdate {
let mut update = CodeGraphContextUpdate::default();
if let Some(block_id) = block_id {
if doc.get_block(&block_id).is_none() {
update
.warnings
.push(format!("focus block not found: {}", block_id));
return update;
}
self.ensure_selected_with_origin(
block_id,
CodeGraphDetailLevel::Skeleton,
selection_origin(CodeGraphSelectionOriginKind::Manual, None, None),
&mut update,
);
}
self.focus = block_id;
self.apply_prune_policy(doc, &mut update);
update.focus = self.focus;
self.history.push(match self.focus {
Some(id) => format!("focus:{}", id),
None => "focus:clear".to_string(),
});
update
}
pub fn select_block(
&mut self,
doc: &Document,
block_id: BlockId,
detail_level: CodeGraphDetailLevel,
) -> CodeGraphContextUpdate {
let mut update = CodeGraphContextUpdate::default();
if doc.get_block(&block_id).is_none() {
update
.warnings
.push(format!("block not found: {}", block_id));
return update;
}
self.ensure_selected_with_origin(
block_id,
detail_level,
selection_origin(CodeGraphSelectionOriginKind::Manual, None, None),
&mut update,
);
self.apply_prune_policy(doc, &mut update);
self.history
.push(format!("select:{}:{:?}", block_id, detail_level));
update.focus = self.focus;
update
}
pub fn remove_block(&mut self, block_id: BlockId) -> CodeGraphContextUpdate {
let mut update = CodeGraphContextUpdate::default();
if self.selected.remove(&block_id).is_some() {
update.removed.push(block_id);
if self.focus == Some(block_id) {
self.focus = None;
}
self.history.push(format!("remove:{}", block_id));
}
update.focus = self.focus;
update
}
pub fn pin(&mut self, block_id: BlockId, pinned: bool) -> CodeGraphContextUpdate {
let mut update = CodeGraphContextUpdate::default();
if let Some(node) = self.selected.get_mut(&block_id) {
node.pinned = pinned;
update.changed.push(block_id);
self.history.push(format!(
"{}:{}",
if pinned { "pin" } else { "unpin" },
block_id
));
}
update.focus = self.focus;
update
}
pub fn seed_overview(&mut self, doc: &Document) -> CodeGraphContextUpdate {
self.seed_overview_with_depth(doc, None)
}
pub fn seed_overview_with_depth(
&mut self,
doc: &Document,
max_depth: Option<usize>,
) -> CodeGraphContextUpdate {
let index = CodeGraphQueryIndex::new(doc);
let previous: HashSet<_> = self.selected.keys().copied().collect();
self.selected.clear();
self.focus = None;
let mut update = CodeGraphContextUpdate::default();
let mut selected = Vec::new();
for block_id in index.overview_nodes(doc, max_depth) {
self.ensure_selected_with_origin(
block_id,
CodeGraphDetailLevel::Skeleton,
selection_origin(CodeGraphSelectionOriginKind::Overview, None, None),
&mut update,
);
selected.push(block_id);
}
if self.focus.is_none() {
self.focus = selected.first().copied().or(Some(doc.root));
}
self.apply_prune_policy(doc, &mut update);
update.focus = self.focus;
update.removed = previous
.into_iter()
.filter(|block_id| !self.selected.contains_key(block_id))
.collect();
update.removed.sort_by_key(BlockId::to_string);
self.history.push(match max_depth {
Some(depth) => format!("seed:overview:{}", depth),
None => "seed:overview:all".to_string(),
});
update
}
pub fn expand_file(&mut self, doc: &Document, file_id: BlockId) -> CodeGraphContextUpdate {
self.expand_file_with_config(doc, file_id, &CodeGraphTraversalConfig::default())
}
pub fn expand_file_with_depth(
&mut self,
doc: &Document,
file_id: BlockId,
depth: usize,
) -> CodeGraphContextUpdate {
self.expand_file_with_config(
doc,
file_id,
&CodeGraphTraversalConfig {
depth,
..CodeGraphTraversalConfig::default()
},
)
}
pub fn expand_file_with_config(
&mut self,
doc: &Document,
file_id: BlockId,
traversal: &CodeGraphTraversalConfig,
) -> CodeGraphContextUpdate {
let index = CodeGraphQueryIndex::new(doc);
let mut update = CodeGraphContextUpdate::default();
let budget = traversal.budget.as_ref();
let start = Instant::now();
let effective_depth = traversal.depth().min(
budget
.and_then(|value| value.max_depth)
.unwrap_or(usize::MAX),
);
let max_add = min_optional_usize(
traversal.max_add,
budget.and_then(|value| value.max_nodes_added),
);
self.ensure_selected_with_origin(
file_id,
CodeGraphDetailLevel::Neighborhood,
selection_origin(CodeGraphSelectionOriginKind::Manual, None, None),
&mut update,
);
let mut added_count = 0usize;
let mut visited_count = 0usize;
let mut skipped_for_threshold = 0usize;
let mut budget_exhausted = false;
if effective_depth > 0 {
let mut queue: VecDeque<(BlockId, usize)> = index
.file_symbols(&file_id)
.into_iter()
.map(|symbol_id| (symbol_id, 1usize))
.collect();
while let Some((symbol_id, symbol_depth)) = queue.pop_front() {
visited_count += 1;
if budget
.and_then(|value| value.max_nodes_visited)
.map(|limit| visited_count > limit)
.unwrap_or(false)
{
update.warnings.push(format!(
"stopped file expansion after visiting {} nodes due to max_nodes_visited budget",
visited_count - 1
));
budget_exhausted = true;
break;
}
if budget
.and_then(|value| value.max_elapsed_ms)
.map(|limit| start.elapsed().as_millis() as u64 >= limit)
.unwrap_or(false)
{
update.warnings.push(format!(
"stopped file expansion after {} ms due to max_elapsed_ms budget",
start.elapsed().as_millis()
));
budget_exhausted = true;
break;
}
let candidate_priority = frontier_priority("expand_file", None, 1, false);
if traversal
.priority_threshold
.map(|threshold| candidate_priority < threshold)
.unwrap_or(false)
{
skipped_for_threshold += 1;
continue;
}
if max_add.map(|limit| added_count >= limit).unwrap_or(false) {
budget_exhausted = true;
break;
}
let was_selected = self.selected.contains_key(&symbol_id);
self.ensure_selected_with_origin(
symbol_id,
CodeGraphDetailLevel::SymbolCard,
selection_origin(
CodeGraphSelectionOriginKind::FileSymbols,
None,
Some(file_id),
),
&mut update,
);
if !was_selected && self.selected.contains_key(&symbol_id) {
added_count += 1;
}
if symbol_depth >= effective_depth {
continue;
}
for child in index.symbol_children(&symbol_id) {
queue.push_back((child, symbol_depth + 1));
}
}
}
if skipped_for_threshold > 0 {
update.warnings.push(format!(
"skipped {} file-symbol candidates below priority threshold",
skipped_for_threshold
));
}
if budget_exhausted {
update.warnings.push(format!(
"stopped file expansion after adding {} nodes due to max_add budget",
added_count
));
}
self.focus = Some(file_id);
self.apply_prune_policy(doc, &mut update);
update.focus = self.focus;
self.history.push(format!(
"expand:file:{}:{}:{}:{}",
file_id,
effective_depth,
max_add
.map(|value| value.to_string())
.unwrap_or_else(|| "*".to_string()),
traversal
.priority_threshold
.map(|value| value.to_string())
.unwrap_or_else(|| "*".to_string())
));
update
}
pub fn expand_dependencies(
&mut self,
doc: &Document,
block_id: BlockId,
relation_filter: Option<&str>,
) -> CodeGraphContextUpdate {
self.expand_dependencies_with_config(
doc,
block_id,
&CodeGraphTraversalConfig {
relation_filters: relation_filter
.map(|relation| vec![relation.to_string()])
.unwrap_or_default(),
..CodeGraphTraversalConfig::default()
},
)
}
pub fn expand_dependencies_with_filters(
&mut self,
doc: &Document,
block_id: BlockId,
relation_filters: Option<&HashSet<String>>,
depth: usize,
) -> CodeGraphContextUpdate {
self.expand_dependencies_with_config(
doc,
block_id,
&CodeGraphTraversalConfig {
depth,
relation_filters: relation_filters
.map(|filters| filters.iter().cloned().collect())
.unwrap_or_default(),
..CodeGraphTraversalConfig::default()
},
)
}
pub fn expand_dependencies_with_config(
&mut self,
doc: &Document,
block_id: BlockId,
traversal: &CodeGraphTraversalConfig,
) -> CodeGraphContextUpdate {
self.expand_neighbors(doc, block_id, traversal, TraversalKind::Outgoing)
}
pub fn expand_dependents(
&mut self,
doc: &Document,
block_id: BlockId,
relation_filter: Option<&str>,
) -> CodeGraphContextUpdate {
self.expand_dependents_with_config(
doc,
block_id,
&CodeGraphTraversalConfig {
relation_filters: relation_filter
.map(|relation| vec![relation.to_string()])
.unwrap_or_default(),
..CodeGraphTraversalConfig::default()
},
)
}
pub fn expand_dependents_with_filters(
&mut self,
doc: &Document,
block_id: BlockId,
relation_filters: Option<&HashSet<String>>,
depth: usize,
) -> CodeGraphContextUpdate {
self.expand_dependents_with_config(
doc,
block_id,
&CodeGraphTraversalConfig {
depth,
relation_filters: relation_filters
.map(|filters| filters.iter().cloned().collect())
.unwrap_or_default(),
..CodeGraphTraversalConfig::default()
},
)
}
pub fn expand_dependents_with_config(
&mut self,
doc: &Document,
block_id: BlockId,
traversal: &CodeGraphTraversalConfig,
) -> CodeGraphContextUpdate {
self.expand_neighbors(doc, block_id, traversal, TraversalKind::Incoming)
}
pub fn collapse(
&mut self,
doc: &Document,
block_id: BlockId,
include_descendants: bool,
) -> CodeGraphContextUpdate {
let index = CodeGraphQueryIndex::new(doc);
let mut update = CodeGraphContextUpdate::default();
let mut to_remove = vec![block_id];
if include_descendants {
to_remove.extend(index.descendants(block_id));
}
for id in to_remove {
let Some(node) = self.selected.get(&id) else {
continue;
};
if node.pinned {
update
.warnings
.push(format!("{} is pinned and was not removed", id));
continue;
}
self.selected.remove(&id);
update.removed.push(id);
if self.focus == Some(id) {
self.focus = None;
}
}
if self.focus.is_none() {
self.focus = self.selected.keys().next().copied();
}
update.focus = self.focus;
self.history
.push(format!("collapse:{}:{}", block_id, include_descendants));
update
}
pub fn hydrate_source(
&mut self,
doc: &Document,
block_id: BlockId,
padding: usize,
) -> CodeGraphContextUpdate {
self.hydrate_source_with_budget(doc, block_id, padding, None)
}
pub fn hydrate_source_with_budget(
&mut self,
doc: &Document,
block_id: BlockId,
padding: usize,
budget: Option<&CodeGraphOperationBudget>,
) -> CodeGraphContextUpdate {
let mut update = CodeGraphContextUpdate::default();
self.ensure_selected_with_origin(
block_id,
CodeGraphDetailLevel::Source,
selection_origin(CodeGraphSelectionOriginKind::Manual, None, None),
&mut update,
);
match hydrate_source_excerpt(doc, block_id, padding) {
Ok(Some(mut excerpt)) => {
if let Some(max_bytes) = budget.and_then(|value| value.max_hydrated_bytes) {
if excerpt.snippet.len() > max_bytes {
excerpt.snippet = truncate_utf8(&excerpt.snippet, max_bytes);
update.warnings.push(format!(
"hydrated source for {} was truncated to {} bytes due to max_hydrated_bytes budget",
block_id, max_bytes
));
}
}
if let Some(node) = self.selected.get_mut(&block_id) {
node.detail_level = CodeGraphDetailLevel::Source;
node.hydrated_source = Some(excerpt);
update.changed.push(block_id);
}
}
Ok(None) => update
.warnings
.push(format!("no coderef available for {}", block_id)),
Err(error) => update.warnings.push(error),
}
self.focus = Some(block_id);
self.apply_prune_policy(doc, &mut update);
update.focus = self.focus;
self.history.push(format!("hydrate:{}", block_id));
update
}
pub fn prune(&mut self, doc: &Document, max_selected: Option<usize>) -> CodeGraphContextUpdate {
let mut update = CodeGraphContextUpdate::default();
if let Some(limit) = max_selected {
self.prune_policy.max_selected = limit.max(1);
}
self.apply_prune_policy(doc, &mut update);
self.history
.push(format!("prune:{}", self.prune_policy.max_selected));
update.focus = self.focus;
update
}
pub fn render_for_prompt(&self, doc: &Document, config: &CodeGraphRenderConfig) -> String {
let index = CodeGraphQueryIndex::new(doc);
let summary = self.summary(doc);
let short_ids = make_short_ids(self, &index);
let selected_ids: HashSet<_> = self.selected.keys().copied().collect();
let mut repository_nodes = Vec::new();
let mut directory_nodes = Vec::new();
let mut file_nodes = Vec::new();
let mut symbol_nodes = Vec::new();
for block_id in self.selected_block_ids() {
match index.node_class(&block_id).unwrap_or("unknown") {
"repository" => repository_nodes.push(block_id),
"directory" => directory_nodes.push(block_id),
"file" => file_nodes.push(block_id),
"symbol" => symbol_nodes.push(block_id),
_ => {}
}
}
let mut out = String::new();
let _ = writeln!(out, "CodeGraph working set");
let focus = self
.focus
.and_then(|id| render_reference(doc, &index, &short_ids, id));
let _ = writeln!(
out,
"focus: {}",
focus.unwrap_or_else(|| "none".to_string())
);
let _ = writeln!(
out,
"summary: selected={}/{} repositories={} directories={} files={} symbols={} hydrated={}",
summary.selected,
summary.max_selected,
summary.repositories,
summary.directories,
summary.files,
summary.symbols,
summary.hydrated_sources
);
if !repository_nodes.is_empty() || !directory_nodes.is_empty() || !file_nodes.is_empty() {
let _ = writeln!(out, "\nfilesystem:");
for block_id in repository_nodes
.into_iter()
.chain(directory_nodes.into_iter())
.chain(file_nodes.into_iter())
{
let block = match doc.get_block(&block_id) {
Some(block) => block,
None => continue,
};
let short = short_ids
.get(&block_id)
.cloned()
.unwrap_or_else(|| block_id.to_string());
let label = index
.display_label(doc, &block_id)
.unwrap_or_else(|| block_id.to_string());
let language = block
.metadata
.custom
.get(META_LANGUAGE)
.and_then(Value::as_str)
.map(|value| format!(" [{}]", value))
.unwrap_or_default();
let pin = self
.selected
.get(&block_id)
.filter(|node| node.pinned)
.map(|_| " [pinned]")
.unwrap_or("");
let _ = writeln!(out, "- [{}] {}{}{}", short, label, language, pin);
}
}
if !symbol_nodes.is_empty() {
let _ = writeln!(out, "\nopened symbols:");
for block_id in symbol_nodes {
let Some(block) = doc.get_block(&block_id) else {
continue;
};
let Some(node) = self.selected.get(&block_id) else {
continue;
};
let short = short_ids
.get(&block_id)
.cloned()
.unwrap_or_else(|| block_id.to_string());
let coderef = metadata_coderef_display(block)
.or_else(|| content_coderef_display(block))
.unwrap_or_else(|| {
index
.display_label(doc, &block_id)
.unwrap_or_else(|| block_id.to_string())
});
let pin = if node.pinned { " [pinned]" } else { "" };
let _ = writeln!(
out,
"- [{}] {}{} @ {}",
short,
format_symbol_signature(block),
format_symbol_modifiers(block),
coderef
);
if !pin.is_empty() {
let _ = writeln!(out, " flags:{}", pin);
}
if let Some(description) =
content_string(block, "description").or_else(|| block.metadata.summary.clone())
{
let _ = writeln!(out, " docs: {}", description);
}
if node.detail_level.includes_neighborhood() {
render_edge_section(
&mut out,
"outgoing",
index.outgoing_edges(&block_id),
&selected_ids,
&short_ids,
doc,
&index,
config.max_edges_per_node,
);
render_edge_section(
&mut out,
"incoming",
index.incoming_edges(&block_id),
&selected_ids,
&short_ids,
doc,
&index,
config.max_edges_per_node,
);
}
if node.detail_level.includes_source() {
if let Some(source) = &node.hydrated_source {
let _ = writeln!(
out,
" source: {}:{}-{}",
source.path, source.start_line, source.end_line
);
for line in source.snippet.lines().take(config.max_source_lines) {
let _ = writeln!(out, " {}", line);
}
}
}
}
}
let total_symbols = index.total_symbols();
let omitted_symbols = total_symbols.saturating_sub(summary.symbols);
let _ = writeln!(out, "\nomissions:");
let _ = writeln!(
out,
"- symbols omitted from working set: {}",
omitted_symbols
);
let _ = writeln!(
out,
"- prune policy: max_selected={} demote_before_remove={} protect_focus={}",
self.prune_policy.max_selected,
self.prune_policy.demote_before_remove,
self.prune_policy.protect_focus
);
let _ = writeln!(out, "\nfrontier:");
if let Some(focus_id) = self.focus {
match index.node_class(&focus_id).unwrap_or("unknown") {
"file" => {
let short = short_ids
.get(&focus_id)
.cloned()
.unwrap_or_else(|| focus_id.to_string());
let _ = writeln!(out, "- [{}] expand file symbols", short);
let _ = writeln!(out, "- [{}] hydrate file source", short);
}
"symbol" => {
let short = short_ids
.get(&focus_id)
.cloned()
.unwrap_or_else(|| focus_id.to_string());
let _ = writeln!(out, "- [{}] expand dependencies", short);
let _ = writeln!(out, "- [{}] expand dependents", short);
let _ = writeln!(out, "- [{}] hydrate source", short);
let _ = writeln!(out, "- [{}] collapse", short);
}
_ => {
let _ = writeln!(
out,
"- set focus to a file or symbol to expand the working set"
);
}
}
} else {
let _ = writeln!(out, "- no focus block set");
}
out.trim_end().to_string()
}
pub fn export(&self, doc: &Document, config: &CodeGraphRenderConfig) -> CodeGraphContextExport {
self.export_with_config(doc, config, &CodeGraphExportConfig::default())
}
pub fn export_with_config(
&self,
doc: &Document,
config: &CodeGraphRenderConfig,
export_config: &CodeGraphExportConfig,
) -> CodeGraphContextExport {
let index = CodeGraphQueryIndex::new(doc);
let summary = self.summary(doc);
let short_ids = make_short_ids(self, &index);
let selected_ids: HashSet<_> = self.selected.keys().copied().collect();
let mut omissions = CodeGraphExportOmissionReport::default();
let distances = focus_distances(doc, self.focus, &selected_ids, &index);
let visible_selected_ids = visible_selected_ids(
self.focus,
&selected_ids,
&distances,
export_config.visible_levels,
);
let hidden_levels = hidden_level_summaries(
self,
&index,
&selected_ids,
&visible_selected_ids,
&distances,
export_config.visible_levels,
);
let hidden_unreachable_count = selected_ids
.iter()
.filter(|block_id| {
!visible_selected_ids.contains(block_id) && !distances.contains_key(block_id)
})
.count();
for block_id in selected_ids
.difference(&visible_selected_ids)
.copied()
.collect::<Vec<_>>()
{
omissions.hidden_by_visible_levels += 1;
omissions.details.push(CodeGraphExportOmissionDetail {
block_id: Some(block_id),
short_id: short_ids.get(&block_id).cloned(),
label: index.display_label(doc, &block_id),
reason: CodeGraphExportOmissionReason::VisibleLevelLimit,
explanation: format!(
"Node is outside the visible level budget from the current focus (visible_levels={}).",
export_config
.visible_levels
.map(|value| value.to_string())
.unwrap_or_else(|| "none".to_string())
),
});
}
let filtered_selected_ids =
class_filtered_selected_ids(&index, &visible_selected_ids, export_config);
for block_id in visible_selected_ids
.difference(&filtered_selected_ids)
.copied()
.collect::<Vec<_>>()
{
omissions.excluded_by_class_filters += 1;
omissions.details.push(CodeGraphExportOmissionDetail {
block_id: Some(block_id),
short_id: short_ids.get(&block_id).cloned(),
label: index.display_label(doc, &block_id),
reason: CodeGraphExportOmissionReason::ClassFilter,
explanation: format!(
"Node was excluded by export node-class filters (only={:?}, exclude={:?}).",
export_config.only_node_classes, export_config.exclude_node_classes
),
});
}
let mut nodes = Vec::new();
for block_id in self.selected_block_ids() {
if !filtered_selected_ids.contains(&block_id) {
continue;
}
let Some(block) = doc.get_block(&block_id) else {
continue;
};
let Some(node) = self.selected.get(&block_id) else {
continue;
};
let node_class = index.node_class(&block_id).unwrap_or("unknown").to_string();
let label = index
.display_label(doc, &block_id)
.unwrap_or_else(|| block_id.to_string());
let logical_key = block_logical_key(block);
let content_name = content_string(block, "name");
let symbol_name = block
.metadata
.custom
.get(META_SYMBOL_NAME)
.and_then(Value::as_str)
.or(content_name.as_deref())
.map(str::to_string);
let path = metadata_coderef_path(block).or_else(|| content_coderef_path(block));
let distance_from_focus = distances.get(&block_id).copied();
let relevance_score =
relevance_score_for_node(self, &index, block_id, distance_from_focus);
let signature = if node_class == "symbol" {
Some(format!(
"{}{}",
format_symbol_signature(block),
format_symbol_modifiers(block)
))
} else {
None
};
let docs = if should_include_docs(
export_config,
self.focus,
block_id,
node,
distance_from_focus,
) {
content_string(block, "description").or_else(|| block.metadata.summary.clone())
} else {
None
};
let coderef = block_coderef(block).map(|coderef| CodeGraphCoderef {
path: coderef.path,
display: coderef.display,
start_line: coderef.start_line,
end_line: coderef.end_line,
});
let hydrated_source = if should_include_hydrated_source(
export_config,
self.focus,
block_id,
node,
distance_from_focus,
) {
node.hydrated_source.clone()
} else {
None
};
if hydrated_source.is_some() && docs.is_none() {
omissions.suppressed_by_hydrated_excerpt += 1;
omissions.details.push(CodeGraphExportOmissionDetail {
block_id: Some(block_id),
short_id: short_ids.get(&block_id).cloned(),
label: Some(label.clone()),
reason: CodeGraphExportOmissionReason::HydratedExcerptSupersededSummary,
explanation:
"Hydrated source excerpt superseded the shorter summary/docs view for this node."
.to_string(),
});
}
nodes.push(CodeGraphContextNodeExport {
block_id,
short_id: short_ids
.get(&block_id)
.cloned()
.unwrap_or_else(|| block_id.to_string()),
node_class,
label,
detail_level: node.detail_level,
pinned: node.pinned,
distance_from_focus,
relevance_score,
logical_key,
symbol_name,
path,
signature,
docs,
origin: node.origin.clone(),
coderef,
hydrated_source,
});
}
nodes.sort_by_key(|node| {
(
std::cmp::Reverse(node.relevance_score),
node.distance_from_focus.unwrap_or(usize::MAX),
node.short_id.clone(),
)
});
apply_render_budget_to_nodes(&mut nodes, &mut omissions, config);
let budget_node_ids = nodes
.iter()
.map(|node| node.block_id)
.collect::<HashSet<_>>();
let (edges, total_selected_edges) =
export_edges(&index, &budget_node_ids, &short_ids, export_config);
let frontier = self.export_frontier(doc, &index, &short_ids, &selected_ids);
let heuristics = self.compute_heuristics(doc, &index, &short_ids, &frontier);
let omitted_symbol_count = index.total_symbols().saturating_sub(summary.symbols);
let rendered = if export_config.include_rendered {
apply_rendered_text_budget(self.render_for_prompt(doc, config), config, &mut omissions)
} else {
String::new()
};
CodeGraphContextExport {
summary,
export_mode: export_config.mode,
visible_levels: export_config.visible_levels,
focus: self.focus,
focus_short_id: self.focus.and_then(|id| short_ids.get(&id).cloned()),
focus_label: self.focus.and_then(|id| index.display_label(doc, &id)),
visible_node_count: nodes.len(),
hidden_unreachable_count,
hidden_levels,
nodes,
edges,
frontier: frontier
.into_iter()
.take(export_config.max_frontier_actions.max(1))
.collect(),
heuristics,
omitted_symbol_count,
total_selected_edges,
omissions,
rendered,
}
}
fn export_frontier(
&self,
doc: &Document,
index: &CodeGraphQueryIndex,
short_ids: &HashMap<BlockId, String>,
selected_ids: &HashSet<BlockId>,
) -> Vec<CodeGraphContextFrontierAction> {
let Some(focus_id) = self.focus else {
return Vec::new();
};
let short_id = short_ids
.get(&focus_id)
.cloned()
.unwrap_or_else(|| focus_id.to_string());
let label = index
.display_label(doc, &focus_id)
.unwrap_or_else(|| focus_id.to_string());
match index.node_class(&focus_id).unwrap_or("unknown") {
"file" => {
let hidden = index
.file_symbols(&focus_id)
.into_iter()
.filter(|id| !selected_ids.contains(id))
.count();
let mut actions = vec![CodeGraphContextFrontierAction {
block_id: focus_id,
short_id,
action: "expand_file".to_string(),
relation: None,
direction: None,
candidate_count: hidden,
priority: frontier_priority("expand_file", None, hidden, false),
description: format!("Expand file symbols for {}", label),
explanation: Some(format!(
"{} hidden symbol candidates remain under the focused file",
hidden
)),
}];
actions.push(CodeGraphContextFrontierAction {
block_id: focus_id,
short_id: actions[0].short_id.clone(),
action: "hydrate_source".to_string(),
relation: None,
direction: None,
candidate_count: usize::from(
self.selected
.get(&focus_id)
.and_then(|node| node.hydrated_source.as_ref())
.is_none(),
),
priority: frontier_priority(
"hydrate_source",
None,
usize::from(
self.selected
.get(&focus_id)
.and_then(|node| node.hydrated_source.as_ref())
.is_none(),
),
false,
),
description: format!("Hydrate source for file {}", label),
explanation: Some(format!(
"Hydrating {} will add source lines to the working set export",
label
)),
});
actions.sort_by_key(|action| {
(
std::cmp::Reverse(action.priority),
action.action.clone(),
action.relation.clone(),
)
});
actions
}
"symbol" => {
let mut actions = Vec::new();
append_relation_frontier(
&mut actions,
focus_id,
&short_id,
&label,
index.outgoing_edges(&focus_id),
selected_ids,
"expand_dependencies",
"outgoing",
);
append_relation_frontier(
&mut actions,
focus_id,
&short_id,
&label,
index.incoming_edges(&focus_id),
selected_ids,
"expand_dependents",
"incoming",
);
actions.push(CodeGraphContextFrontierAction {
block_id: focus_id,
short_id: short_id.clone(),
action: "hydrate_source".to_string(),
relation: None,
direction: None,
candidate_count: usize::from(
self.selected
.get(&focus_id)
.and_then(|node| node.hydrated_source.as_ref())
.is_none(),
),
priority: frontier_priority(
"hydrate_source",
None,
usize::from(
self.selected
.get(&focus_id)
.and_then(|node| node.hydrated_source.as_ref())
.is_none(),
),
false,
),
description: format!("Hydrate source for {}", label),
explanation: Some(format!(
"Hydrating {} will surface an anchored source excerpt for the focused symbol",
label
)),
});
actions.push(CodeGraphContextFrontierAction {
block_id: focus_id,
short_id,
action: "collapse".to_string(),
relation: None,
direction: None,
candidate_count: 1,
priority: frontier_priority("collapse", None, 1, false),
description: format!("Collapse {} from working set", label),
explanation: Some(format!(
"Collapse removes {} from the active working set when the current branch is no longer useful",
label
)),
});
actions.sort_by_key(|action| {
(
std::cmp::Reverse(action.priority),
action.action.clone(),
action.relation.clone(),
)
});
actions
}
_ => Vec::new(),
}
}
fn compute_heuristics(
&self,
doc: &Document,
index: &CodeGraphQueryIndex,
short_ids: &HashMap<BlockId, String>,
frontier: &[CodeGraphContextFrontierAction],
) -> CodeGraphContextHeuristics {
let focus_node = self.focus.and_then(|id| self.selected.get(&id));
let focus_hydrated = focus_node
.and_then(|node| node.hydrated_source.as_ref())
.is_some();
let hidden_candidate_count = frontier
.iter()
.filter(|action| action.action.starts_with("expand_") && action.candidate_count > 0)
.map(|action| action.candidate_count)
.sum();
let low_value_candidate_count = frontier
.iter()
.filter(|action| action.action.starts_with("expand_") && action.priority <= 30)
.map(|action| action.candidate_count)
.sum();
let recommended_actions: Vec<_> = frontier
.iter()
.filter(|action| action.candidate_count > 0)
.take(3)
.cloned()
.collect();
let recommendations = recommended_actions
.iter()
.map(|action| recommendation_from_frontier(doc, index, short_ids, action))
.collect::<Vec<_>>();
let recommended_next_action = recommended_actions.first().cloned();
let mut reasons = Vec::new();
let should_stop = match self.focus {
None => {
reasons
.push("set focus to a file or symbol before continuing expansion".to_string());
false
}
Some(focus_id) => match index.node_class(&focus_id).unwrap_or("unknown") {
"file" => {
if hidden_candidate_count == 0 && focus_hydrated {
reasons.push(
"focus file is hydrated and no unselected file symbols remain"
.to_string(),
);
true
} else if hidden_candidate_count == 0 {
reasons.push(
"all file symbols for the focused file are already selected"
.to_string(),
);
false
} else {
false
}
}
"symbol" => {
if hidden_candidate_count == 0 && focus_hydrated {
reasons.push(
"focus symbol is hydrated and no unselected dependency frontier remains"
.to_string(),
);
true
} else if focus_hydrated
&& hidden_candidate_count > 0
&& hidden_candidate_count == low_value_candidate_count
{
reasons.push(
"remaining frontier is low-value compared to the hydrated focus symbol"
.to_string(),
);
true
} else {
false
}
}
_ => frontier.iter().all(|action| action.candidate_count == 0),
},
};
CodeGraphContextHeuristics {
should_stop,
reasons,
hidden_candidate_count,
low_value_candidate_count,
recommended_next_action,
recommended_actions,
recommendations,
}
}
fn ensure_selected_with_origin(
&mut self,
block_id: BlockId,
detail_level: CodeGraphDetailLevel,
origin: Option<CodeGraphSelectionOrigin>,
update: &mut CodeGraphContextUpdate,
) {
match self.selected.get_mut(&block_id) {
Some(node) => {
let next = node.detail_level.max(detail_level);
if next != node.detail_level {
node.detail_level = next;
update.changed.push(block_id);
}
if origin_is_more_protective(origin.as_ref(), node.origin.as_ref()) {
node.origin = origin;
push_unique(&mut update.changed, block_id);
}
}
None => {
self.selected.insert(
block_id,
CodeGraphContextNode {
block_id,
detail_level,
pinned: false,
origin,
hydrated_source: None,
},
);
update.added.push(block_id);
}
}
}
fn expand_neighbors(
&mut self,
doc: &Document,
block_id: BlockId,
traversal_config: &CodeGraphTraversalConfig,
traversal_kind: TraversalKind,
) -> CodeGraphContextUpdate {
let index = CodeGraphQueryIndex::new(doc);
let mut update = CodeGraphContextUpdate::default();
let relation_filters = traversal_config.relation_filter_set();
let budget = traversal_config.budget.as_ref();
let start = Instant::now();
let effective_depth = traversal_config.depth().min(
budget
.and_then(|value| value.max_depth)
.unwrap_or(usize::MAX),
);
let max_add = min_optional_usize(
traversal_config.max_add,
budget.and_then(|value| value.max_nodes_added),
);
self.ensure_selected_with_origin(
block_id,
CodeGraphDetailLevel::Neighborhood,
selection_origin(
CodeGraphSelectionOriginKind::Manual,
relation_filters
.as_ref()
.and_then(|filters| join_relation_filters(filters)),
None,
),
&mut update,
);
let mut queue = VecDeque::from([(block_id, 0usize)]);
let mut visited = HashSet::from([block_id]);
let mut added_count = 0usize;
let mut visited_count = 0usize;
let mut skipped_for_threshold = 0usize;
let mut budget_exhausted = false;
while let Some((current, current_depth)) = queue.pop_front() {
visited_count += 1;
if budget
.and_then(|value| value.max_nodes_visited)
.map(|limit| visited_count > limit)
.unwrap_or(false)
{
update.warnings.push(format!(
"stopped expansion after visiting {} nodes due to max_nodes_visited budget",
visited_count - 1
));
budget_exhausted = true;
break;
}
if budget
.and_then(|value| value.max_elapsed_ms)
.map(|limit| start.elapsed().as_millis() as u64 >= limit)
.unwrap_or(false)
{
update.warnings.push(format!(
"stopped expansion after {} ms due to max_elapsed_ms budget",
start.elapsed().as_millis()
));
budget_exhausted = true;
break;
}
if current_depth >= effective_depth {
continue;
}
let edges = match traversal_kind {
TraversalKind::Outgoing => index.outgoing_edges(¤t),
TraversalKind::Incoming => index.incoming_edges(¤t),
};
for edge in edges {
if !relation_matches(relation_filters.as_ref(), edge.relation.as_str()) {
continue;
}
let action_name = match traversal_kind {
TraversalKind::Outgoing => "expand_dependencies",
TraversalKind::Incoming => "expand_dependents",
};
let candidate_priority = frontier_priority(
action_name,
Some(edge.relation.as_str()),
1,
is_low_value_relation(action_name, edge.relation.as_str()),
);
if traversal_config
.priority_threshold
.map(|threshold| candidate_priority < threshold)
.unwrap_or(false)
{
skipped_for_threshold += 1;
continue;
}
if max_add.map(|limit| added_count >= limit).unwrap_or(false) {
budget_exhausted = true;
break;
}
let class = index.node_class(&edge.other).unwrap_or("unknown");
let level = if class == "symbol" {
CodeGraphDetailLevel::SymbolCard
} else {
CodeGraphDetailLevel::Skeleton
};
let was_selected = self.selected.contains_key(&edge.other);
self.ensure_selected_with_origin(
edge.other,
level,
selection_origin(
match traversal_kind {
TraversalKind::Outgoing => CodeGraphSelectionOriginKind::Dependencies,
TraversalKind::Incoming => CodeGraphSelectionOriginKind::Dependents,
},
Some(edge.relation.as_str()),
Some(current),
),
&mut update,
);
if !was_selected && self.selected.contains_key(&edge.other) {
added_count += 1;
}
if visited.insert(edge.other) {
queue.push_back((edge.other, current_depth + 1));
}
}
if budget_exhausted {
break;
}
}
if skipped_for_threshold > 0 {
update.warnings.push(format!(
"skipped {} candidates below priority threshold",
skipped_for_threshold
));
}
if budget_exhausted {
update.warnings.push(format!(
"stopped expansion after adding {} nodes due to max_add budget",
added_count
));
}
self.focus = Some(block_id);
self.apply_prune_policy(doc, &mut update);
update.focus = self.focus;
self.history.push(format!(
"expand:{}:{}:{}:{}:{}:{}",
match traversal_kind {
TraversalKind::Outgoing => "dependencies",
TraversalKind::Incoming => "dependents",
},
block_id,
relation_filters
.as_ref()
.map(join_relation_filter_string)
.unwrap_or_else(|| "*".to_string()),
effective_depth,
max_add
.map(|value| value.to_string())
.unwrap_or_else(|| "*".to_string()),
traversal_config
.priority_threshold
.map(|value| value.to_string())
.unwrap_or_else(|| "*".to_string())
));
update
}
fn apply_prune_policy(&mut self, doc: &Document, update: &mut CodeGraphContextUpdate) {
if self.selected.len() <= self.prune_policy.max_selected.max(1) {
return;
}
let index = CodeGraphQueryIndex::new(doc);
let protected_focus = if self.prune_policy.protect_focus {
self.focus
} else {
None
};
if self.prune_policy.demote_before_remove {
while self.selected.len() > self.prune_policy.max_selected.max(1) {
let Some(block_id) = self.next_demotable_block(&index, protected_focus) else {
break;
};
let Some(node) = self.selected.get_mut(&block_id) else {
continue;
};
let next_level = node.detail_level.demoted();
if next_level == node.detail_level {
break;
}
node.detail_level = next_level;
if !node.detail_level.includes_source() {
node.hydrated_source = None;
}
push_unique(&mut update.changed, block_id);
}
}
while self.selected.len() > self.prune_policy.max_selected.max(1) {
let Some(block_id) = self.next_removable_block(&index, protected_focus) else {
update.warnings.push(format!(
"working set has {} nodes but no removable nodes remain under current prune policy",
self.selected.len()
));
break;
};
self.selected.remove(&block_id);
push_unique(&mut update.removed, block_id);
update.added.retain(|id| id != &block_id);
update.changed.retain(|id| id != &block_id);
if self.focus == Some(block_id) {
self.focus = None;
}
}
if self.focus.is_none() {
self.focus = self.next_focus_candidate(&index);
}
update.focus = self.focus;
}
fn next_demotable_block(
&self,
index: &CodeGraphQueryIndex,
protected_focus: Option<BlockId>,
) -> Option<BlockId> {
self.selected
.values()
.filter(|node| Some(node.block_id) != protected_focus && !node.pinned)
.filter(|node| node.detail_level.demoted() != node.detail_level)
.max_by_key(|node| {
(
origin_prune_rank(node.origin.as_ref(), &self.prune_policy),
relation_prune_rank(node.origin.as_ref(), &self.prune_policy),
node.detail_level as u8,
prune_removal_rank(index.node_class(&node.block_id).unwrap_or("unknown")),
node.block_id.to_string(),
)
})
.map(|node| node.block_id)
}
fn next_removable_block(
&self,
index: &CodeGraphQueryIndex,
protected_focus: Option<BlockId>,
) -> Option<BlockId> {
self.selected
.values()
.filter(|node| Some(node.block_id) != protected_focus && !node.pinned)
.max_by_key(|node| {
(
origin_prune_rank(node.origin.as_ref(), &self.prune_policy),
relation_prune_rank(node.origin.as_ref(), &self.prune_policy),
prune_removal_rank(index.node_class(&node.block_id).unwrap_or("unknown")),
node.detail_level as u8,
node.block_id.to_string(),
)
})
.map(|node| node.block_id)
}
fn next_focus_candidate(&self, index: &CodeGraphQueryIndex) -> Option<BlockId> {
self.selected
.values()
.min_by_key(|node| {
(
focus_preference_rank(index.node_class(&node.block_id).unwrap_or("unknown")),
node.block_id.to_string(),
)
})
.map(|node| node.block_id)
}
}
#[derive(Debug, Clone, Copy)]
enum TraversalKind {
Outgoing,
Incoming,
}
impl CodeGraphQueryIndex {
fn new(doc: &Document) -> Self {
let mut logical_keys = HashMap::new();
let mut logical_key_to_id = HashMap::new();
let mut paths_to_id = HashMap::new();
let mut display_to_id = HashMap::new();
let mut symbol_names_to_id: HashMap<String, Vec<BlockId>> = HashMap::new();
let mut node_classes = HashMap::new();
let mut outgoing: HashMap<BlockId, Vec<IndexedEdge>> = HashMap::new();
let mut incoming: HashMap<BlockId, Vec<IndexedEdge>> = HashMap::new();
let mut file_symbols: HashMap<BlockId, Vec<BlockId>> = HashMap::new();
let mut symbol_children: HashMap<BlockId, Vec<BlockId>> = HashMap::new();
let mut structure_parent: HashMap<BlockId, BlockId> = HashMap::new();
for (block_id, block) in &doc.blocks {
if let Some(key) = block_logical_key(block) {
logical_keys.insert(*block_id, key.clone());
logical_key_to_id.insert(key, *block_id);
}
if let Some(class) = node_class(block) {
node_classes.insert(*block_id, class.clone());
}
if let Some(path) = metadata_coderef_path(block).or_else(|| content_coderef_path(block))
{
let should_replace = match paths_to_id.get(&path) {
Some(existing_id) => {
let existing_rank = path_selector_rank(
node_classes
.get(existing_id)
.map(String::as_str)
.unwrap_or("unknown"),
);
let next_rank = path_selector_rank(
node_classes
.get(block_id)
.map(String::as_str)
.unwrap_or("unknown"),
);
next_rank < existing_rank
}
None => true,
};
if should_replace {
paths_to_id.insert(path, *block_id);
}
}
if let Some(display) =
metadata_coderef_display(block).or_else(|| content_coderef_display(block))
{
display_to_id.insert(display, *block_id);
}
let content_name = content_string(block, "name");
if let Some(symbol_name) = block
.metadata
.custom
.get(META_SYMBOL_NAME)
.and_then(Value::as_str)
.or(content_name.as_deref())
{
symbol_names_to_id
.entry(symbol_name.to_string())
.or_default()
.push(*block_id);
}
}
for (source, block) in &doc.blocks {
for edge in &block.edges {
let relation = edge_type_to_string(&edge.edge_type);
outgoing.entry(*source).or_default().push(IndexedEdge {
other: edge.target,
relation: relation.clone(),
});
incoming.entry(edge.target).or_default().push(IndexedEdge {
other: *source,
relation,
});
}
}
for (parent, children) in &doc.structure {
let parent_class = node_classes
.get(parent)
.map(String::as_str)
.unwrap_or("unknown");
for child in children {
let child_class = node_classes
.get(child)
.map(String::as_str)
.unwrap_or("unknown");
if parent_class == "file" && child_class == "symbol" {
file_symbols.entry(*parent).or_default().push(*child);
}
if parent_class == "symbol" && child_class == "symbol" {
symbol_children.entry(*parent).or_default().push(*child);
}
structure_parent.insert(*child, *parent);
}
}
Self {
logical_keys,
logical_key_to_id,
paths_to_id,
display_to_id,
symbol_names_to_id,
node_classes,
outgoing,
incoming,
file_symbols,
symbol_children,
structure_parent,
}
}
fn resolve_selector(&self, selector: &str) -> Option<BlockId> {
BlockId::from_str(selector)
.ok()
.or_else(|| self.logical_key_to_id.get(selector).copied())
.or_else(|| self.paths_to_id.get(selector).copied())
.or_else(|| self.display_to_id.get(selector).copied())
.or_else(|| {
self.symbol_names_to_id.get(selector).and_then(|ids| {
if ids.len() == 1 {
ids.first().copied()
} else {
None
}
})
})
}
fn overview_nodes(&self, doc: &Document, max_depth: Option<usize>) -> Vec<BlockId> {
let mut nodes = Vec::new();
let limit = max_depth.unwrap_or(usize::MAX);
let mut queue = VecDeque::from([(doc.root, 0usize)]);
let mut visited = HashSet::new();
while let Some((block_id, depth)) = queue.pop_front() {
if !visited.insert(block_id) {
continue;
}
let class = self
.node_classes
.get(&block_id)
.map(String::as_str)
.unwrap_or("unknown");
if matches!(class, "repository" | "directory" | "file") {
nodes.push(block_id);
}
if depth >= limit {
continue;
}
for child in doc.children(&block_id) {
let child_class = self
.node_classes
.get(child)
.map(String::as_str)
.unwrap_or("unknown");
if matches!(child_class, "repository" | "directory" | "file") {
queue.push_back((*child, depth + 1));
}
}
}
nodes.sort_by_key(|block_id| {
self.logical_keys
.get(block_id)
.cloned()
.unwrap_or_else(|| block_id.to_string())
});
nodes
}
fn outgoing_edges(&self, block_id: &BlockId) -> Vec<IndexedEdge> {
self.outgoing.get(block_id).cloned().unwrap_or_default()
}
fn incoming_edges(&self, block_id: &BlockId) -> Vec<IndexedEdge> {
self.incoming.get(block_id).cloned().unwrap_or_default()
}
fn file_symbols(&self, block_id: &BlockId) -> Vec<BlockId> {
let mut symbols = self.file_symbols.get(block_id).cloned().unwrap_or_default();
symbols.sort_by_key(|id| {
self.logical_keys
.get(id)
.cloned()
.unwrap_or_else(|| id.to_string())
});
symbols
}
fn symbol_children(&self, block_id: &BlockId) -> Vec<BlockId> {
let mut children = self
.symbol_children
.get(block_id)
.cloned()
.unwrap_or_default();
children.sort_by_key(|id| {
self.logical_keys
.get(id)
.cloned()
.unwrap_or_else(|| id.to_string())
});
children
}
fn descendants(&self, block_id: BlockId) -> Vec<BlockId> {
let mut out = Vec::new();
let mut queue: VecDeque<BlockId> = self
.symbol_children
.get(&block_id)
.cloned()
.unwrap_or_default()
.into();
while let Some(next) = queue.pop_front() {
out.push(next);
if let Some(children) = self.symbol_children.get(&next) {
for child in children {
queue.push_back(*child);
}
}
}
out
}
fn node_class(&self, block_id: &BlockId) -> Option<&str> {
self.node_classes.get(block_id).map(String::as_str)
}
fn structure_parent(&self, block_id: &BlockId) -> Option<BlockId> {
self.structure_parent.get(block_id).copied()
}
fn total_symbols(&self) -> usize {
self.node_classes
.values()
.filter(|class| class.as_str() == "symbol")
.count()
}
fn display_label(&self, doc: &Document, block_id: &BlockId) -> Option<String> {
let block = doc.get_block(block_id)?;
match self.node_class(block_id) {
Some("file") | Some("directory") | Some("repository") => metadata_coderef_path(block)
.or_else(|| content_coderef_path(block))
.or_else(|| block_logical_key(block)),
Some("symbol") => block_logical_key(block)
.or_else(|| metadata_coderef_display(block))
.or_else(|| content_coderef_display(block)),
_ => block_logical_key(block),
}
}
}
pub fn is_codegraph_document(doc: &Document) -> bool {
let profile = doc.metadata.custom.get("profile").and_then(Value::as_str);
let marker = doc
.metadata
.custom
.get("profile_marker")
.and_then(Value::as_str);
profile == Some("codegraph") || marker == Some(CODEGRAPH_PROFILE_MARKER)
}
pub fn resolve_codegraph_selector(doc: &Document, selector: &str) -> Option<BlockId> {
CodeGraphQueryIndex::new(doc).resolve_selector(selector)
}
pub fn render_codegraph_context_prompt(
doc: &Document,
session: &CodeGraphContextSession,
config: &CodeGraphRenderConfig,
) -> String {
session.render_for_prompt(doc, config)
}
pub fn export_codegraph_context(
doc: &Document,
session: &CodeGraphContextSession,
config: &CodeGraphRenderConfig,
) -> CodeGraphContextExport {
session.export(doc, config)
}
pub fn export_codegraph_context_with_config(
doc: &Document,
session: &CodeGraphContextSession,
render_config: &CodeGraphRenderConfig,
export_config: &CodeGraphExportConfig,
) -> CodeGraphContextExport {
session.export_with_config(doc, render_config, export_config)
}
pub fn approximate_prompt_tokens(rendered: &str) -> u32 {
((rendered.len() as f32) / 4.0).ceil() as u32
}
fn approximate_tokens_for_bytes(bytes: usize) -> u32 {
((bytes as f32) / 4.0).ceil() as u32
}
fn truncate_utf8(value: &str, max_bytes: usize) -> String {
if value.len() <= max_bytes {
return value.to_string();
}
let mut end = 0usize;
for (index, _) in value.char_indices() {
if index > max_bytes {
break;
}
end = index;
}
if end == 0 {
String::new()
} else {
value[..end].to_string()
}
}
fn estimated_export_node_bytes(node: &CodeGraphContextNodeExport) -> usize {
node.label.len()
+ node.short_id.len()
+ node.logical_key.as_ref().map(String::len).unwrap_or(0)
+ node.symbol_name.as_ref().map(String::len).unwrap_or(0)
+ node.path.as_ref().map(String::len).unwrap_or(0)
+ node.signature.as_ref().map(String::len).unwrap_or(0)
+ node.docs.as_ref().map(String::len).unwrap_or(0)
+ node
.hydrated_source
.as_ref()
.map(|source| source.snippet.len())
.unwrap_or(0)
}
fn content_source_bytes(block: &Block) -> Option<usize> {
match &block.content {
Content::Code(code) => Some(code.source.len()),
Content::Text(text) => Some(text.text.len()),
_ => Some(block.content.size_bytes()),
}
}
fn origin_is_more_protective(
next: Option<&CodeGraphSelectionOrigin>,
current: Option<&CodeGraphSelectionOrigin>,
) -> bool {
match (next, current) {
(Some(next), Some(current)) => {
selection_origin_protection_rank(next) < selection_origin_protection_rank(current)
}
(Some(_), None) => true,
_ => false,
}
}
fn selection_origin_protection_rank(origin: &CodeGraphSelectionOrigin) -> u8 {
match origin.kind {
CodeGraphSelectionOriginKind::Manual => 0,
CodeGraphSelectionOriginKind::Overview => 1,
CodeGraphSelectionOriginKind::FileSymbols => 2,
CodeGraphSelectionOriginKind::Dependencies => 3,
CodeGraphSelectionOriginKind::Dependents => 4,
}
}
fn origin_prune_rank(
origin: Option<&CodeGraphSelectionOrigin>,
policy: &CodeGraphPrunePolicy,
) -> u8 {
let _ = policy;
match origin.map(|origin| origin.kind) {
Some(CodeGraphSelectionOriginKind::Dependents) => 5,
Some(CodeGraphSelectionOriginKind::Dependencies) => 4,
Some(CodeGraphSelectionOriginKind::FileSymbols) => 2,
Some(CodeGraphSelectionOriginKind::Overview) => 1,
Some(CodeGraphSelectionOriginKind::Manual) => 0,
None => 0,
}
}
fn relation_prune_rank(
origin: Option<&CodeGraphSelectionOrigin>,
policy: &CodeGraphPrunePolicy,
) -> u8 {
origin
.and_then(|origin| origin.relation.as_ref())
.and_then(|relation| policy.relation_prune_priority.get(relation).copied())
.unwrap_or(0)
}
fn push_unique(ids: &mut Vec<BlockId>, block_id: BlockId) {
if !ids.contains(&block_id) {
ids.push(block_id);
}
}
fn default_one() -> usize {
1
}
fn prune_removal_rank(node_class: &str) -> u8 {
match node_class {
"symbol" => 4,
"file" => 3,
"directory" => 2,
"repository" => 1,
_ => 0,
}
}
fn focus_preference_rank(node_class: &str) -> u8 {
match node_class {
"symbol" => 0,
"file" => 1,
"directory" => 2,
"repository" => 3,
_ => 4,
}
}
fn path_selector_rank(node_class: &str) -> u8 {
match node_class {
"file" => 0,
"directory" => 1,
"repository" => 2,
"symbol" => 3,
_ => 4,
}
}
#[allow(clippy::too_many_arguments)]
fn render_edge_section(
out: &mut String,
label: &str,
edges: Vec<IndexedEdge>,
selected_ids: &HashSet<BlockId>,
short_ids: &HashMap<BlockId, String>,
doc: &Document,
index: &CodeGraphQueryIndex,
limit: usize,
) {
let visible = dedupe_visible_edges(edges, selected_ids);
if visible.is_empty() {
return;
}
let _ = writeln!(out, " {}:", label);
for (edge, multiplicity) in visible.iter().take(limit) {
let short = short_ids
.get(&edge.other)
.cloned()
.unwrap_or_else(|| edge.other.to_string());
let target = index
.display_label(doc, &edge.other)
.unwrap_or_else(|| edge.other.to_string());
let suffix = if *multiplicity > 1 {
format!(" (x{})", multiplicity)
} else {
String::new()
};
let _ = writeln!(
out,
" - {} -> [{}] {}{}",
edge.relation, short, target, suffix
);
}
if visible.len() > limit {
let _ = writeln!(out, " - ... {} more", visible.len() - limit);
}
}
#[allow(clippy::too_many_arguments)]
fn append_relation_frontier(
out: &mut Vec<CodeGraphContextFrontierAction>,
block_id: BlockId,
short_id: &str,
label: &str,
edges: Vec<IndexedEdge>,
selected_ids: &HashSet<BlockId>,
action: &str,
direction: &str,
) {
let mut counts: BTreeMap<String, usize> = BTreeMap::new();
for edge in edges {
if selected_ids.contains(&edge.other) {
continue;
}
*counts.entry(edge.relation).or_default() += 1;
}
for (relation, candidate_count) in counts {
let low_value = is_low_value_relation(action, relation.as_str());
out.push(CodeGraphContextFrontierAction {
block_id,
short_id: short_id.to_string(),
action: action.to_string(),
relation: Some(relation.clone()),
direction: Some(direction.to_string()),
candidate_count,
priority: frontier_priority(
action,
Some(relation.as_str()),
candidate_count,
low_value,
),
description: format!(
"{} {} neighbors via {} for {}",
action, direction, relation, label
),
explanation: Some(format!(
"{} hidden {} candidate{} remain for {} via {}",
candidate_count,
direction,
if candidate_count == 1 { "" } else { "s" },
label,
relation
)),
});
}
}
fn recommendation_from_frontier(
doc: &Document,
index: &CodeGraphQueryIndex,
short_ids: &HashMap<BlockId, String>,
action: &CodeGraphContextFrontierAction,
) -> CodeGraphRecommendation {
let target_label = index
.display_label(doc, &action.block_id)
.unwrap_or_else(|| action.block_id.to_string());
let estimated_evidence_gain = match action.action.as_str() {
"hydrate_source" => 4,
"collapse" => 1,
_ => action.candidate_count.max(1),
};
let estimated_hydration_bytes = if action.action == "hydrate_source" {
doc.get_block(&action.block_id)
.map(|block| content_source_bytes(block).unwrap_or(0))
.unwrap_or(0)
} else {
0
};
let estimated_token_cost = if estimated_hydration_bytes > 0 {
approximate_tokens_for_bytes(estimated_hydration_bytes)
} else {
(action.candidate_count as u32).saturating_mul(24)
};
let rationale = action
.explanation
.clone()
.unwrap_or_else(|| action.description.clone());
CodeGraphRecommendation {
action_kind: action.action.clone(),
target_block_id: action.block_id,
target_short_id: short_ids
.get(&action.block_id)
.cloned()
.unwrap_or_else(|| action.block_id.to_string()),
target_label,
relation_set: action.relation.clone().into_iter().collect(),
priority: action.priority,
candidate_count: action.candidate_count,
estimated_evidence_gain,
estimated_token_cost,
estimated_hydration_bytes,
explanation: action.description.clone(),
rationale,
}
}
fn apply_render_budget_to_nodes(
nodes: &mut Vec<CodeGraphContextNodeExport>,
omissions: &mut CodeGraphExportOmissionReport,
config: &CodeGraphRenderConfig,
) {
let max_bytes = config.max_rendered_bytes;
let max_tokens = config.max_rendered_tokens;
if max_bytes.is_none() && max_tokens.is_none() {
return;
}
let mut kept = Vec::new();
let mut used_bytes = 0usize;
let mut used_tokens = 0u32;
for node in nodes.drain(..) {
let node_bytes = estimated_export_node_bytes(&node);
let node_tokens = approximate_tokens_for_bytes(node_bytes);
let within_bytes = max_bytes
.map(|limit| used_bytes + node_bytes <= limit)
.unwrap_or(true);
let within_tokens = max_tokens
.map(|limit| used_tokens + node_tokens <= limit)
.unwrap_or(true);
if within_bytes && within_tokens {
used_bytes += node_bytes;
used_tokens = used_tokens.saturating_add(node_tokens);
kept.push(node);
continue;
}
omissions.dropped_by_render_budget += 1;
omissions.details.push(CodeGraphExportOmissionDetail {
block_id: Some(node.block_id),
short_id: Some(node.short_id.clone()),
label: Some(node.label.clone()),
reason: CodeGraphExportOmissionReason::RenderBudget,
explanation: format!(
"Node exceeded export/render budget (bytes_used={} tokens_used={}).",
used_bytes, used_tokens
),
});
}
*nodes = kept;
}
fn apply_rendered_text_budget(
rendered: String,
config: &CodeGraphRenderConfig,
omissions: &mut CodeGraphExportOmissionReport,
) -> String {
let mut limited = rendered;
if let Some(max_bytes) = config.max_rendered_bytes {
if limited.len() > max_bytes {
limited = truncate_utf8(&limited, max_bytes);
omissions.dropped_by_render_budget += 1;
omissions.details.push(CodeGraphExportOmissionDetail {
block_id: None,
short_id: None,
label: Some("rendered".to_string()),
reason: CodeGraphExportOmissionReason::RenderBudget,
explanation: format!(
"Rendered prompt text was truncated to {} bytes by max_rendered_bytes.",
max_bytes
),
});
}
}
if let Some(max_tokens) = config.max_rendered_tokens {
if approximate_prompt_tokens(&limited) > max_tokens {
let max_bytes = (max_tokens as usize).saturating_mul(4);
limited = truncate_utf8(&limited, max_bytes);
omissions.dropped_by_render_budget += 1;
omissions.details.push(CodeGraphExportOmissionDetail {
block_id: None,
short_id: None,
label: Some("rendered".to_string()),
reason: CodeGraphExportOmissionReason::RenderBudget,
explanation: format!(
"Rendered prompt text was truncated to approximately {} tokens.",
max_tokens
),
});
}
}
limited
}
fn is_low_value_relation(action: &str, relation: &str) -> bool {
matches!(action, "expand_dependents")
|| relation == "references"
|| relation == "cited_by"
|| relation == "links_to"
}
fn dedupe_visible_edges(
edges: Vec<IndexedEdge>,
selected_ids: &HashSet<BlockId>,
) -> Vec<(IndexedEdge, usize)> {
let mut counts: HashMap<(BlockId, String), usize> = HashMap::new();
for edge in edges {
if !selected_ids.contains(&edge.other) {
continue;
}
*counts.entry((edge.other, edge.relation)).or_default() += 1;
}
let mut deduped: Vec<_> = counts
.into_iter()
.map(|((other, relation), multiplicity)| (IndexedEdge { other, relation }, multiplicity))
.collect();
deduped.sort_by_key(|(edge, _)| (edge.relation.clone(), edge.other.to_string()));
deduped
}
fn export_edges(
index: &CodeGraphQueryIndex,
selected_ids: &HashSet<BlockId>,
short_ids: &HashMap<BlockId, String>,
export_config: &CodeGraphExportConfig,
) -> (Vec<CodeGraphContextEdgeExport>, usize) {
let mut edges = Vec::new();
let mut total_selected_edges = 0;
if export_config.dedupe_edges {
let mut counts: HashMap<(BlockId, String, BlockId), usize> = HashMap::new();
for source in selected_ids.iter().copied() {
for edge in index.outgoing_edges(&source) {
if !selected_ids.contains(&edge.other) {
continue;
}
total_selected_edges += 1;
*counts
.entry((source, edge.relation, edge.other))
.or_default() += 1;
}
}
for ((source, relation, target), multiplicity) in counts {
edges.push(CodeGraphContextEdgeExport {
source,
source_short_id: short_ids
.get(&source)
.cloned()
.unwrap_or_else(|| source.to_string()),
target,
target_short_id: short_ids
.get(&target)
.cloned()
.unwrap_or_else(|| target.to_string()),
relation,
multiplicity,
});
}
} else {
for source in selected_ids.iter().copied() {
for edge in index.outgoing_edges(&source) {
if !selected_ids.contains(&edge.other) {
continue;
}
total_selected_edges += 1;
edges.push(CodeGraphContextEdgeExport {
source,
source_short_id: short_ids
.get(&source)
.cloned()
.unwrap_or_else(|| source.to_string()),
target: edge.other,
target_short_id: short_ids
.get(&edge.other)
.cloned()
.unwrap_or_else(|| edge.other.to_string()),
relation: edge.relation,
multiplicity: 1,
});
}
}
}
edges.sort_by_key(|edge| {
(
edge.source_short_id.clone(),
edge.relation.clone(),
edge.target_short_id.clone(),
)
});
(edges, total_selected_edges)
}
fn focus_distances(
doc: &Document,
focus: Option<BlockId>,
selected_ids: &HashSet<BlockId>,
index: &CodeGraphQueryIndex,
) -> HashMap<BlockId, usize> {
let mut distances = HashMap::new();
let Some(focus) = focus else {
return distances;
};
if !selected_ids.contains(&focus) {
return distances;
}
let mut queue = VecDeque::from([(focus, 0usize)]);
distances.insert(focus, 0);
while let Some((block_id, distance)) = queue.pop_front() {
let mut neighbors: Vec<BlockId> = index
.outgoing_edges(&block_id)
.into_iter()
.chain(index.incoming_edges(&block_id).into_iter())
.map(|edge| edge.other)
.collect();
neighbors.extend(doc.children(&block_id));
if let Some(parent) = index.structure_parent(&block_id) {
neighbors.push(parent);
}
for neighbor in neighbors {
if !selected_ids.contains(&neighbor) || distances.contains_key(&neighbor) {
continue;
}
distances.insert(neighbor, distance + 1);
queue.push_back((neighbor, distance + 1));
}
}
distances
}
fn visible_selected_ids(
focus: Option<BlockId>,
selected_ids: &HashSet<BlockId>,
distances: &HashMap<BlockId, usize>,
visible_levels: Option<usize>,
) -> HashSet<BlockId> {
match (focus, visible_levels) {
(Some(_), Some(levels)) => selected_ids
.iter()
.copied()
.filter(|block_id| distances.get(block_id).copied().unwrap_or(usize::MAX) <= levels)
.collect(),
_ => selected_ids.clone(),
}
}
fn class_filtered_selected_ids(
index: &CodeGraphQueryIndex,
selected_ids: &HashSet<BlockId>,
export_config: &CodeGraphExportConfig,
) -> HashSet<BlockId> {
selected_ids
.iter()
.copied()
.filter(|block_id| {
node_class_visible(
index.node_class(block_id).unwrap_or("unknown"),
export_config,
)
})
.collect()
}
fn hidden_level_summaries(
session: &CodeGraphContextSession,
index: &CodeGraphQueryIndex,
selected_ids: &HashSet<BlockId>,
visible_selected_ids: &HashSet<BlockId>,
distances: &HashMap<BlockId, usize>,
visible_levels: Option<usize>,
) -> Vec<CodeGraphHiddenLevelSummary> {
let Some(levels) = visible_levels else {
return Vec::new();
};
let mut counts: BTreeMap<(usize, Option<String>, Option<String>), usize> = BTreeMap::new();
for block_id in selected_ids {
if visible_selected_ids.contains(block_id) {
continue;
}
let Some(distance) = distances.get(block_id).copied() else {
continue;
};
if distance > levels {
let (relation, direction) = hidden_summary_metadata(session, index, *block_id);
*counts.entry((distance, relation, direction)).or_default() += 1;
}
}
counts
.into_iter()
.map(
|((level, relation, direction), count)| CodeGraphHiddenLevelSummary {
level,
count,
relation,
direction,
},
)
.collect()
}
fn hidden_summary_metadata(
session: &CodeGraphContextSession,
index: &CodeGraphQueryIndex,
block_id: BlockId,
) -> (Option<String>, Option<String>) {
let Some(node) = session.selected.get(&block_id) else {
return (None, None);
};
match node.origin.as_ref() {
Some(origin) if origin.kind == CodeGraphSelectionOriginKind::Dependencies => {
(origin.relation.clone(), Some("outgoing".to_string()))
}
Some(origin) if origin.kind == CodeGraphSelectionOriginKind::Dependents => {
(origin.relation.clone(), Some("incoming".to_string()))
}
Some(origin) if origin.kind == CodeGraphSelectionOriginKind::FileSymbols => (
Some("contains_symbol".to_string()),
Some("structural".to_string()),
),
Some(origin) if origin.kind == CodeGraphSelectionOriginKind::Overview => (
Some("structure".to_string()),
Some("structural".to_string()),
),
Some(origin) if origin.kind == CodeGraphSelectionOriginKind::Manual => {
(origin.relation.clone(), Some("manual".to_string()))
}
_ => match index.node_class(&block_id).unwrap_or("unknown") {
"repository" | "directory" | "file" => (
Some("structure".to_string()),
Some("structural".to_string()),
),
_ => (None, None),
},
}
}
fn node_class_visible(node_class: &str, export_config: &CodeGraphExportConfig) -> bool {
let only_matches = export_config.only_node_classes.is_empty()
|| export_config
.only_node_classes
.iter()
.any(|allowed| allowed == node_class);
let excluded = export_config
.exclude_node_classes
.iter()
.any(|excluded| excluded == node_class);
only_matches && !excluded
}
fn relation_matches(relation_filters: Option<&HashSet<String>>, relation: &str) -> bool {
relation_filters
.map(|filters| filters.contains(relation))
.unwrap_or(true)
}
fn join_relation_filters(relation_filters: &HashSet<String>) -> Option<&str> {
if relation_filters.len() == 1 {
relation_filters.iter().next().map(String::as_str)
} else {
None
}
}
fn join_relation_filter_string(relation_filters: &HashSet<String>) -> String {
let mut filters: Vec<_> = relation_filters.iter().cloned().collect();
filters.sort();
filters.join(",")
}
fn relevance_score_for_node(
session: &CodeGraphContextSession,
index: &CodeGraphQueryIndex,
block_id: BlockId,
distance_from_focus: Option<usize>,
) -> u16 {
let Some(node) = session.selected.get(&block_id) else {
return 0;
};
let mut score = 0u16;
if session.focus == Some(block_id) {
score += 100;
}
if node.pinned {
score += 40;
}
score += match index.node_class(&block_id).unwrap_or("unknown") {
"symbol" => 40,
"file" => 28,
"directory" => 16,
"repository" => 10,
_ => 6,
};
score += match node.detail_level {
CodeGraphDetailLevel::Source => 30,
CodeGraphDetailLevel::Neighborhood => 20,
CodeGraphDetailLevel::SymbolCard => 12,
CodeGraphDetailLevel::Skeleton => 4,
};
score += match node.origin.as_ref().map(|origin| origin.kind) {
Some(CodeGraphSelectionOriginKind::Manual) => 24,
Some(CodeGraphSelectionOriginKind::Overview) => 18,
Some(CodeGraphSelectionOriginKind::FileSymbols) => 16,
Some(CodeGraphSelectionOriginKind::Dependencies) => 12,
Some(CodeGraphSelectionOriginKind::Dependents) => 8,
None => 0,
};
score += match distance_from_focus {
Some(0) => 30,
Some(1) => 20,
Some(2) => 10,
Some(3) => 4,
Some(_) => 1,
None => 0,
};
score
}
fn should_include_docs(
export_config: &CodeGraphExportConfig,
focus: Option<BlockId>,
block_id: BlockId,
node: &CodeGraphContextNode,
distance_from_focus: Option<usize>,
) -> bool {
match export_config.mode {
CodeGraphExportMode::Full => true,
CodeGraphExportMode::Compact => {
focus == Some(block_id) || node.pinned || distance_from_focus.unwrap_or(usize::MAX) <= 1
}
}
}
fn should_include_hydrated_source(
export_config: &CodeGraphExportConfig,
focus: Option<BlockId>,
block_id: BlockId,
node: &CodeGraphContextNode,
distance_from_focus: Option<usize>,
) -> bool {
if node.hydrated_source.is_none() {
return false;
}
match export_config.mode {
CodeGraphExportMode::Full => true,
CodeGraphExportMode::Compact => {
focus == Some(block_id)
|| (node.pinned && distance_from_focus.unwrap_or(usize::MAX) <= 1)
}
}
}
fn frontier_priority(
action: &str,
relation: Option<&str>,
candidate_count: usize,
low_value: bool,
) -> u16 {
let base = match action {
"hydrate_source" => 120,
"expand_file" => 100,
"expand_dependencies" => 85,
"expand_dependents" => 70,
"collapse" => 5,
_ => 20,
};
let relation_adjust = match relation {
Some("references") | Some("cited_by") => -20,
Some("links_to") => -12,
Some("uses_symbol") => 8,
Some("imports_symbol") => 6,
Some("reexports_symbol") => 4,
Some("calls") => 6,
_ => 0,
};
let low_value_adjust = if low_value { -10 } else { 0 };
let count_bonus = candidate_count.min(12) as i32;
(base + relation_adjust + low_value_adjust + count_bonus).max(0) as u16
}
fn make_short_ids(
session: &CodeGraphContextSession,
index: &CodeGraphQueryIndex,
) -> HashMap<BlockId, String> {
let mut by_class: BTreeMap<&str, Vec<BlockId>> = BTreeMap::new();
for block_id in session.selected.keys().copied() {
by_class
.entry(index.node_class(&block_id).unwrap_or("node"))
.or_default()
.push(block_id);
}
let mut result = HashMap::new();
for (class, ids) in by_class {
let mut ids = ids;
ids.sort_by_key(|block_id| {
index
.logical_keys
.get(block_id)
.cloned()
.unwrap_or_else(|| block_id.to_string())
});
for (idx, block_id) in ids.into_iter().enumerate() {
let prefix = match class {
"repository" => "R",
"directory" => "D",
"file" => "F",
"symbol" => "S",
_ => "N",
};
result.insert(block_id, format!("{}{}", prefix, idx + 1));
}
}
result
}
fn render_reference(
doc: &Document,
index: &CodeGraphQueryIndex,
short_ids: &HashMap<BlockId, String>,
block_id: BlockId,
) -> Option<String> {
Some(format!(
"[{}] {}",
short_ids.get(&block_id)?.clone(),
index.display_label(doc, &block_id)?
))
}
fn edge_type_to_string(edge_type: &ucm_core::EdgeType) -> String {
match edge_type {
ucm_core::EdgeType::DerivedFrom => "derived_from".to_string(),
ucm_core::EdgeType::Supersedes => "supersedes".to_string(),
ucm_core::EdgeType::TransformedFrom => "transformed_from".to_string(),
ucm_core::EdgeType::References => "references".to_string(),
ucm_core::EdgeType::CitedBy => "cited_by".to_string(),
ucm_core::EdgeType::LinksTo => "links_to".to_string(),
ucm_core::EdgeType::Supports => "supports".to_string(),
ucm_core::EdgeType::Contradicts => "contradicts".to_string(),
ucm_core::EdgeType::Elaborates => "elaborates".to_string(),
ucm_core::EdgeType::Summarizes => "summarizes".to_string(),
ucm_core::EdgeType::ParentOf => "parent_of".to_string(),
ucm_core::EdgeType::SiblingOf => "sibling_of".to_string(),
ucm_core::EdgeType::PreviousSibling => "previous_sibling".to_string(),
ucm_core::EdgeType::NextSibling => "next_sibling".to_string(),
ucm_core::EdgeType::VersionOf => "version_of".to_string(),
ucm_core::EdgeType::AlternativeOf => "alternative_of".to_string(),
ucm_core::EdgeType::TranslationOf => "translation_of".to_string(),
ucm_core::EdgeType::ChildOf => "child_of".to_string(),
ucm_core::EdgeType::Custom(name) => name.clone(),
}
}
fn hydrate_source_excerpt(
doc: &Document,
block_id: BlockId,
padding: usize,
) -> Result<Option<HydratedSourceExcerpt>, String> {
let Some(block) = doc.get_block(&block_id) else {
return Err(format!("block not found: {}", block_id));
};
let coderef =
block_coderef(block).ok_or_else(|| format!("missing coderef for {}", block_id))?;
let repo =
repository_root(doc).ok_or_else(|| "missing repository_path metadata".to_string())?;
#[cfg(target_arch = "wasm32")]
{
let _ = (repo, coderef, padding);
Err("source hydration is not available on wasm32".to_string())
}
#[cfg(not(target_arch = "wasm32"))]
{
let path = repo.join(&coderef.path);
let source = std::fs::read_to_string(&path)
.map_err(|error| format!("failed to read {}: {}", path.display(), error))?;
let lines: Vec<_> = source.lines().collect();
let line_count = lines.len().max(1);
let start_line = coderef.start_line.unwrap_or(1).max(1);
let end_line = coderef
.end_line
.unwrap_or(start_line)
.max(start_line)
.min(line_count);
let slice_start = start_line.saturating_sub(padding + 1);
let slice_end = (end_line + padding).min(line_count);
let mut snippet = String::new();
for (idx, line) in lines[slice_start..slice_end].iter().enumerate() {
let number = slice_start + idx + 1;
let _ = writeln!(snippet, "{:>4} | {}", number, line);
}
Ok(Some(HydratedSourceExcerpt {
path: coderef.path,
display: coderef.display,
start_line,
end_line,
snippet: snippet.trim_end().to_string(),
}))
}
}
fn repository_root(doc: &Document) -> Option<PathBuf> {
doc.metadata
.custom
.get("repository_path")
.and_then(Value::as_str)
.map(PathBuf::from)
}
#[derive(Debug, Clone)]
struct BlockCoderef {
path: String,
display: String,
start_line: Option<usize>,
end_line: Option<usize>,
}
fn block_coderef(block: &Block) -> Option<BlockCoderef> {
let value = block
.metadata
.custom
.get(META_CODEREF)
.or_else(|| match &block.content {
Content::Json { value, .. } => value.get("coderef"),
_ => None,
})?;
Some(BlockCoderef {
path: value.get("path")?.as_str()?.to_string(),
display: value
.get("display")
.and_then(Value::as_str)
.unwrap_or_else(|| {
value
.get("path")
.and_then(Value::as_str)
.unwrap_or("unknown")
})
.to_string(),
start_line: value
.get("start_line")
.and_then(Value::as_u64)
.map(|v| v as usize),
end_line: value
.get("end_line")
.and_then(Value::as_u64)
.map(|v| v as usize),
})
}
fn format_symbol_signature(block: &Block) -> String {
let kind = content_string(block, "kind").unwrap_or_else(|| "symbol".to_string());
let name = content_string(block, "name").unwrap_or_else(|| "unknown".to_string());
let inputs = content_array(block, "inputs")
.into_iter()
.map(|value| {
let name = value.get("name").and_then(Value::as_str).unwrap_or("_");
match value.get("type").and_then(Value::as_str) {
Some(type_name) => format!("{}: {}", name, type_name),
None => name.to_string(),
}
})
.collect::<Vec<_>>();
let output = content_string(block, "output");
let type_info = content_string(block, "type");
match kind.as_str() {
"function" | "method" => {
let mut rendered = format!("{} {}({})", kind, name, inputs.join(", "));
if let Some(output) = output {
let _ = write!(rendered, " -> {}", output);
}
rendered
}
_ => {
let mut rendered = format!("{} {}", kind, name);
if let Some(type_info) = type_info {
let _ = write!(rendered, " : {}", type_info);
}
if block
.metadata
.custom
.get(META_EXPORTED)
.and_then(Value::as_bool)
.unwrap_or(false)
{
let _ = write!(rendered, " [exported]");
}
rendered
}
}
}
fn format_symbol_modifiers(block: &Block) -> String {
let Content::Json { value, .. } = &block.content else {
return String::new();
};
let Some(modifiers) = value.get("modifiers").and_then(Value::as_object) else {
return String::new();
};
let mut parts = Vec::new();
if modifiers.get("async").and_then(Value::as_bool) == Some(true) {
parts.push("async".to_string());
}
if modifiers.get("static").and_then(Value::as_bool) == Some(true) {
parts.push("static".to_string());
}
if modifiers.get("generator").and_then(Value::as_bool) == Some(true) {
parts.push("generator".to_string());
}
if let Some(visibility) = modifiers.get("visibility").and_then(Value::as_str) {
parts.push(visibility.to_string());
}
if parts.is_empty() {
String::new()
} else {
format!(" [{}]", parts.join(", "))
}
}
fn content_string(block: &Block, field: &str) -> Option<String> {
let Content::Json { value, .. } = &block.content else {
return None;
};
value.get(field)?.as_str().map(|value| value.to_string())
}
fn content_array(block: &Block, field: &str) -> Vec<Value> {
let Content::Json { value, .. } = &block.content else {
return Vec::new();
};
value
.get(field)
.and_then(Value::as_array)
.cloned()
.unwrap_or_default()
}
fn node_class(block: &Block) -> Option<String> {
block
.metadata
.custom
.get(META_NODE_CLASS)
.and_then(Value::as_str)
.map(|value| value.to_string())
}
fn block_logical_key(block: &Block) -> Option<String> {
block
.metadata
.custom
.get(META_LOGICAL_KEY)
.and_then(Value::as_str)
.map(|value| value.to_string())
}
fn metadata_coderef_path(block: &Block) -> Option<String> {
block
.metadata
.custom
.get(META_CODEREF)
.and_then(|value| value.get("path"))
.and_then(Value::as_str)
.map(|value| value.to_string())
}
fn metadata_coderef_display(block: &Block) -> Option<String> {
block
.metadata
.custom
.get(META_CODEREF)
.and_then(|value| value.get("display"))
.and_then(Value::as_str)
.map(|value| value.to_string())
}
fn content_coderef_path(block: &Block) -> Option<String> {
let Content::Json { value, .. } = &block.content else {
return None;
};
value
.get("coderef")
.and_then(|value| value.get("path"))
.and_then(Value::as_str)
.map(|value| value.to_string())
}
fn content_coderef_display(block: &Block) -> Option<String> {
let Content::Json { value, .. } = &block.content else {
return None;
};
value
.get("coderef")
.and_then(|value| value.get("display"))
.and_then(Value::as_str)
.map(|value| value.to_string())
}
#[cfg(test)]
mod tests {
use std::fs;
use tempfile::tempdir;
use super::*;
use crate::{build_code_graph, CodeGraphBuildInput, CodeGraphExtractorConfig};
fn build_test_graph() -> Document {
let dir = tempdir().unwrap();
fs::create_dir_all(dir.path().join("src")).unwrap();
fs::write(
dir.path().join("src/util.rs"),
"pub fn util() -> i32 { 1 }\n",
)
.unwrap();
fs::write(
dir.path().join("src/lib.rs"),
"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",
)
.unwrap();
let repository_path = dir.path().to_path_buf();
std::mem::forget(dir);
build_code_graph(&CodeGraphBuildInput {
repository_path,
commit_hash: "context-tests".to_string(),
config: CodeGraphExtractorConfig::default(),
})
.unwrap()
.document
}
#[test]
fn overview_expand_dependents_and_hydrate_source_work() {
let doc = build_test_graph();
let mut session = CodeGraphContextSession::new();
let update = session.seed_overview(&doc);
assert!(!update.added.is_empty());
assert_eq!(session.summary(&doc).symbols, 0);
let file_id = resolve_codegraph_selector(&doc, "src/lib.rs").unwrap();
session.expand_file(&doc, file_id);
assert!(session.summary(&doc).symbols >= 1);
let add_id = resolve_codegraph_selector(&doc, "symbol:src/lib.rs::add").unwrap();
let util_id = resolve_codegraph_selector(&doc, "symbol:src/util.rs::util").unwrap();
let deps = session.expand_dependencies(&doc, add_id, Some("uses_symbol"));
assert!(deps.added.contains(&util_id) || session.selected.contains_key(&util_id));
let dependents = session.expand_dependents(&doc, util_id, Some("uses_symbol"));
assert!(dependents.added.contains(&add_id) || session.selected.contains_key(&add_id));
let hydrated = session.hydrate_source(&doc, add_id, 1);
assert!(hydrated.changed.contains(&add_id));
assert!(session
.selected
.get(&add_id)
.and_then(|node| node.hydrated_source.as_ref())
.is_some());
let rendered = session.render_for_prompt(&doc, &CodeGraphRenderConfig::default());
assert!(rendered.contains("CodeGraph working set"));
assert!(rendered.contains("expand dependents"));
assert!(rendered.contains("uses_symbol"));
assert!(rendered.contains("source:"));
}
#[test]
fn resolve_selector_prefers_logical_key_and_path() {
let doc = build_test_graph();
let file_id = resolve_codegraph_selector(&doc, "src/lib.rs").unwrap();
let logical_id = resolve_codegraph_selector(&doc, "symbol:src/lib.rs::add").unwrap();
let display_id = resolve_codegraph_selector(&doc, "src/lib.rs:2-2").unwrap_or(logical_id);
assert!(doc.get_block(&file_id).is_some());
assert!(doc.get_block(&logical_id).is_some());
assert_eq!(logical_id, display_id);
}
#[test]
fn prune_policy_demotes_before_removing() {
let doc = build_test_graph();
let mut session = CodeGraphContextSession::new();
session.set_prune_policy(CodeGraphPrunePolicy {
max_selected: 10,
..CodeGraphPrunePolicy::default()
});
let file_id = resolve_codegraph_selector(&doc, "src/lib.rs").unwrap();
let add_id = resolve_codegraph_selector(&doc, "symbol:src/lib.rs::add").unwrap();
session.seed_overview(&doc);
session.expand_file(&doc, file_id);
session.expand_dependencies(&doc, add_id, Some("uses_symbol"));
session.hydrate_source(&doc, add_id, 1);
assert!(session
.selected
.get(&add_id)
.and_then(|node| node.hydrated_source.as_ref())
.is_some());
session.set_focus(&doc, Some(file_id));
let update = session.prune(&doc, Some(4));
assert!(session.selected.len() <= 4);
assert!(!update.changed.is_empty() || !update.removed.is_empty());
let rendered = session.render_for_prompt(&doc, &CodeGraphRenderConfig::default());
assert!(rendered.contains("selected=4/4"));
assert!(!rendered.contains("source:"));
}
#[test]
fn prune_prefers_dependents_before_file_skeletons() {
let doc = build_test_graph();
let mut session = CodeGraphContextSession::new();
session.set_prune_policy(CodeGraphPrunePolicy {
max_selected: 20,
..CodeGraphPrunePolicy::default()
});
let util_file_id = resolve_codegraph_selector(&doc, "src/util.rs").unwrap();
let util_symbol_id = resolve_codegraph_selector(&doc, "symbol:src/util.rs::util").unwrap();
let lib_file_id = resolve_codegraph_selector(&doc, "src/lib.rs").unwrap();
let add_id = resolve_codegraph_selector(&doc, "symbol:src/lib.rs::add").unwrap();
let sub_id = resolve_codegraph_selector(&doc, "symbol:src/lib.rs::sub").unwrap();
session.seed_overview(&doc);
session.expand_file(&doc, util_file_id);
session.expand_dependents(&doc, util_symbol_id, Some("uses_symbol"));
assert!(session.selected.contains_key(&add_id));
assert!(session.selected.contains_key(&sub_id));
session.set_focus(&doc, Some(util_file_id));
session.prune(&doc, Some(5));
assert!(session.selected.contains_key(&lib_file_id));
assert!(!session.selected.contains_key(&add_id));
assert!(!session.selected.contains_key(&sub_id));
}
#[test]
fn export_includes_frontier_and_origin_metadata() {
let doc = build_test_graph();
let mut session = CodeGraphContextSession::new();
let file_id = resolve_codegraph_selector(&doc, "src/lib.rs").unwrap();
let add_id = resolve_codegraph_selector(&doc, "symbol:src/lib.rs::add").unwrap();
session.seed_overview(&doc);
session.expand_file(&doc, file_id);
session.focus = Some(add_id);
let export = session.export(&doc, &CodeGraphRenderConfig::default());
assert_eq!(export.focus, Some(add_id));
assert!(export.nodes.iter().any(|node| {
node.block_id == add_id
&& node.symbol_name.as_deref() == Some("add")
&& node.path.as_deref() == Some("src/lib.rs")
&& node
.origin
.as_ref()
.map(|origin| origin.kind == CodeGraphSelectionOriginKind::FileSymbols)
.unwrap_or(false)
}));
assert!(export
.frontier
.iter()
.any(|action| action.action == "hydrate_source"));
assert!(export.heuristics.recommended_next_action.is_some());
assert!(!export.heuristics.should_stop);
assert!(export
.frontier
.iter()
.any(|action| action.action == "expand_dependencies"
&& action.relation.as_deref() == Some("uses_symbol")));
}
#[test]
fn overview_seed_depth_limits_structural_selection() {
let doc = build_test_graph();
let mut shallow = CodeGraphContextSession::new();
shallow.seed_overview_with_depth(&doc, Some(1));
let shallow_summary = shallow.summary(&doc);
assert!(shallow_summary.repositories + shallow_summary.directories >= 1);
assert_eq!(shallow_summary.files, 0);
let mut deeper = CodeGraphContextSession::new();
deeper.seed_overview_with_depth(&doc, Some(3));
let deeper_summary = deeper.summary(&doc);
assert!(deeper_summary.files >= 2);
}
#[test]
fn export_with_visible_levels_summarizes_hidden_nodes() {
let doc = build_test_graph();
let mut session = CodeGraphContextSession::new();
let file_id = resolve_codegraph_selector(&doc, "src/lib.rs").unwrap();
let add_id = resolve_codegraph_selector(&doc, "symbol:src/lib.rs::add").unwrap();
session.seed_overview(&doc);
session.expand_file(&doc, file_id);
session.expand_dependencies(&doc, add_id, Some("uses_symbol"));
session.focus = Some(add_id);
let mut export_config = CodeGraphExportConfig::compact();
export_config.visible_levels = Some(1);
let export =
session.export_with_config(&doc, &CodeGraphRenderConfig::default(), &export_config);
assert_eq!(export.visible_levels, Some(1));
assert!(export.visible_node_count < export.summary.selected);
assert!(export.hidden_levels.iter().any(|hidden| hidden.level >= 2));
assert!(export
.nodes
.iter()
.all(|node| node.distance_from_focus.unwrap_or(usize::MAX) <= 1));
}
#[test]
fn selective_multi_hop_expansion_follows_only_requested_relations() {
let mut doc = build_test_graph();
let mut session = CodeGraphContextSession::new();
let file_id = resolve_codegraph_selector(&doc, "src/lib.rs").unwrap();
let add_id = resolve_codegraph_selector(&doc, "symbol:src/lib.rs::add").unwrap();
let sub_id = resolve_codegraph_selector(&doc, "symbol:src/lib.rs::sub").unwrap();
let util_id = resolve_codegraph_selector(&doc, "symbol:src/util.rs::util").unwrap();
doc.add_edge(&add_id, ucm_core::EdgeType::References, sub_id);
session.seed_overview(&doc);
session.expand_file(&doc, file_id);
let relation_filters = HashSet::from(["references".to_string()]);
session.expand_dependencies_with_filters(&doc, add_id, Some(&relation_filters), 2);
assert!(session.selected.contains_key(&sub_id));
assert!(!session.selected.contains_key(&util_id));
let mut session = CodeGraphContextSession::new();
session.seed_overview(&doc);
session.expand_file(&doc, file_id);
let relation_filters = HashSet::from(["references".to_string(), "uses_symbol".to_string()]);
session.expand_dependencies_with_filters(&doc, add_id, Some(&relation_filters), 2);
assert!(session.selected.contains_key(&sub_id));
assert!(session.selected.contains_key(&util_id));
}
#[test]
fn traversal_budget_caps_additions_and_reports_warning() {
let doc = build_test_graph();
let mut session = CodeGraphContextSession::new();
let file_id = resolve_codegraph_selector(&doc, "src/lib.rs").unwrap();
session.seed_overview(&doc);
let update = session.expand_file_with_config(
&doc,
file_id,
&CodeGraphTraversalConfig {
depth: 2,
max_add: Some(1),
..CodeGraphTraversalConfig::default()
},
);
assert_eq!(update.added.len(), 1);
assert!(update
.warnings
.iter()
.any(|warning| warning.contains("max_add")));
}
#[test]
fn priority_threshold_skips_low_value_relations() {
let mut doc = build_test_graph();
let mut session = CodeGraphContextSession::new();
let file_id = resolve_codegraph_selector(&doc, "src/lib.rs").unwrap();
let add_id = resolve_codegraph_selector(&doc, "symbol:src/lib.rs::add").unwrap();
let util_id = resolve_codegraph_selector(&doc, "symbol:src/util.rs::util").unwrap();
doc.add_edge(&add_id, ucm_core::EdgeType::References, util_id);
session.seed_overview(&doc);
session.expand_file(&doc, file_id);
let update = session.expand_dependencies_with_config(
&doc,
add_id,
&CodeGraphTraversalConfig {
depth: 1,
relation_filters: vec!["references".to_string()],
priority_threshold: Some(80),
..CodeGraphTraversalConfig::default()
},
);
assert!(!session.selected.contains_key(&util_id));
assert!(update.added.is_empty());
}
#[test]
fn export_filters_node_classes_and_includes_hidden_relation_metadata() {
let doc = build_test_graph();
let mut session = CodeGraphContextSession::new();
let file_id = resolve_codegraph_selector(&doc, "src/lib.rs").unwrap();
let add_id = resolve_codegraph_selector(&doc, "symbol:src/lib.rs::add").unwrap();
session.seed_overview(&doc);
session.expand_file(&doc, file_id);
session.expand_dependencies(&doc, add_id, Some("uses_symbol"));
session.focus = Some(add_id);
let mut export_config = CodeGraphExportConfig::compact();
export_config.visible_levels = Some(0);
export_config.only_node_classes = vec!["symbol".to_string()];
let export =
session.export_with_config(&doc, &CodeGraphRenderConfig::default(), &export_config);
assert!(export.nodes.iter().all(|node| node.node_class == "symbol"));
assert!(export.edges.iter().all(|edge| {
export
.nodes
.iter()
.any(|node| node.block_id == edge.source && node.node_class == "symbol")
&& export
.nodes
.iter()
.any(|node| node.block_id == edge.target && node.node_class == "symbol")
}));
assert!(export.hidden_levels.iter().any(|hidden| {
hidden.relation.as_deref() == Some("uses_symbol")
&& hidden.direction.as_deref() == Some("outgoing")
}));
}
#[test]
fn compact_export_dedupes_edges_and_omits_rendered_text() {
let mut doc = build_test_graph();
let mut session = CodeGraphContextSession::new();
let file_id = resolve_codegraph_selector(&doc, "src/lib.rs").unwrap();
let add_id = resolve_codegraph_selector(&doc, "symbol:src/lib.rs::add").unwrap();
let util_id = resolve_codegraph_selector(&doc, "symbol:src/util.rs::util").unwrap();
doc.add_edge(&add_id, ucm_core::EdgeType::References, util_id);
session.seed_overview(&doc);
session.expand_file(&doc, file_id);
session.expand_dependencies(&doc, add_id, Some("uses_symbol"));
session.hydrate_source(&doc, add_id, 1);
session.focus = Some(add_id);
let export = session.export_with_config(
&doc,
&CodeGraphRenderConfig::default(),
&CodeGraphExportConfig::compact(),
);
assert_eq!(export.export_mode, CodeGraphExportMode::Compact);
assert!(export.rendered.is_empty());
assert!(export.total_selected_edges >= export.edges.len());
assert!(export.edges.iter().all(|edge| edge.multiplicity >= 1));
assert!(export
.nodes
.iter()
.find(|node| node.block_id == add_id)
.and_then(|node| node.hydrated_source.as_ref())
.is_some());
}
#[test]
fn heuristics_stop_when_focus_is_hydrated_and_frontier_is_exhausted() {
let doc = build_test_graph();
let mut session = CodeGraphContextSession::new();
let file_id = resolve_codegraph_selector(&doc, "src/lib.rs").unwrap();
let add_id = resolve_codegraph_selector(&doc, "symbol:src/lib.rs::add").unwrap();
session.seed_overview(&doc);
session.expand_file(&doc, file_id);
session.expand_dependencies(&doc, add_id, Some("uses_symbol"));
session.focus = Some(add_id);
let pre_hydrate = session.export(&doc, &CodeGraphRenderConfig::default());
assert!(!pre_hydrate.heuristics.should_stop);
assert_eq!(
pre_hydrate
.heuristics
.recommended_next_action
.as_ref()
.map(|action| action.action.as_str()),
Some("hydrate_source")
);
session.hydrate_source(&doc, add_id, 1);
let exhausted = session.export(&doc, &CodeGraphRenderConfig::default());
assert!(exhausted.heuristics.should_stop);
assert!(exhausted
.heuristics
.reasons
.iter()
.any(|reason| reason.contains("hydrated")));
}
}