11use std:: collections:: HashSet ;
22use std:: path:: PathBuf ;
3+ use std:: time:: Instant ;
34
45use anyhow:: Result ;
56
67use crate :: config;
78use crate :: core;
89
9- use super :: super :: super :: context:: extract_symbols_from_diff;
1010use super :: super :: super :: services:: PipelineServices ;
1111use super :: super :: super :: session:: ReviewSession ;
12+ use super :: traces:: { trace_record, MAX_GRAPH_TRACE_DETAILS } ;
1213
1314pub ( in super :: super ) async fn add_semantic_context (
1415 services : & PipelineServices ,
1516 session : & ReviewSession ,
1617 diff : & core:: UnifiedDiff ,
18+ symbols : & [ String ] ,
1719 context_chunks : & mut Vec < core:: LLMContextChunk > ,
20+ graph_query_records : & mut Vec < core:: dag:: DagExecutionRecord > ,
1821) {
1922 let Some ( index) = session. semantic_index . as_ref ( ) else {
2023 return ;
2124 } ;
22- let preferred_files = graph_ranked_semantic_files ( services, session, diff) ;
25+ let should_trace_graph_ranking = session. symbol_index . is_some ( ) && !symbols. is_empty ( ) ;
26+ let ranking_started = Instant :: now ( ) ;
27+ let preferred_files = graph_ranked_semantic_files ( services, session, diff, symbols) ;
28+ if should_trace_graph_ranking {
29+ append_semantic_preference_trace_records (
30+ graph_query_records,
31+ & preferred_files,
32+ ranking_started. elapsed ( ) . as_millis ( ) as u64 ,
33+ ) ;
34+ }
2335
2436 let semantic_chunks = core:: semantic_context_for_diff (
2537 index,
@@ -41,12 +53,12 @@ fn graph_ranked_semantic_files(
4153 services : & PipelineServices ,
4254 session : & ReviewSession ,
4355 diff : & core:: UnifiedDiff ,
56+ symbols : & [ String ] ,
4457) -> Vec < PathBuf > {
4558 let Some ( index) = session. symbol_index . as_ref ( ) else {
4659 return Vec :: new ( ) ;
4760 } ;
4861
49- let symbols = extract_symbols_from_diff ( diff) ;
5062 if symbols. is_empty ( ) {
5163 return Vec :: new ( ) ;
5264 }
@@ -59,7 +71,7 @@ fn graph_ranked_semantic_files(
5971 services. config . symbol_index_graph_max_files ,
6072 ) ,
6173 ) ;
62- let related_locations = retriever. related_symbol_locations ( & diff. file_path , & symbols) ;
74+ let related_locations = retriever. related_symbol_locations ( & diff. file_path , symbols) ;
6375
6476 let mut preferred_files = Vec :: new ( ) ;
6577 let mut seen = HashSet :: new ( ) ;
@@ -84,6 +96,27 @@ fn graph_ranked_semantic_files(
8496 preferred_files
8597}
8698
99+ fn append_semantic_preference_trace_records (
100+ graph_query_records : & mut Vec < core:: dag:: DagExecutionRecord > ,
101+ preferred_files : & [ PathBuf ] ,
102+ duration_ms : u64 ,
103+ ) {
104+ graph_query_records. push ( trace_record (
105+ format ! ( "semantic_preferred_files={}" , preferred_files. len( ) ) ,
106+ duration_ms,
107+ ) ) ;
108+ for ( index, file_path) in preferred_files
109+ . iter ( )
110+ . take ( MAX_GRAPH_TRACE_DETAILS )
111+ . enumerate ( )
112+ {
113+ graph_query_records. push ( trace_record (
114+ format ! ( "semantic_preferred_file[{index}]={}" , file_path. display( ) ) ,
115+ 0 ,
116+ ) ) ;
117+ }
118+ }
119+
87120pub ( in super :: super ) async fn add_path_context (
88121 services : & PipelineServices ,
89122 diff : & core:: UnifiedDiff ,
@@ -117,3 +150,22 @@ pub(in super::super) async fn add_path_context(
117150
118151 Ok ( ( ) )
119152}
153+
154+ #[ cfg( test) ]
155+ mod tests {
156+ use super :: append_semantic_preference_trace_records;
157+ use std:: path:: PathBuf ;
158+
159+ #[ test]
160+ fn semantic_preference_trace_records_keep_ranked_order ( ) {
161+ let preferred_files = vec ! [ PathBuf :: from( "src/auth.rs" ) , PathBuf :: from( "src/db.rs" ) ] ;
162+
163+ let mut records = Vec :: new ( ) ;
164+ append_semantic_preference_trace_records ( & mut records, & preferred_files, 7 ) ;
165+
166+ assert_eq ! ( records[ 0 ] . name, "semantic_preferred_files=2" ) ;
167+ assert_eq ! ( records[ 0 ] . duration_ms, 7 ) ;
168+ assert_eq ! ( records[ 1 ] . name, "semantic_preferred_file[0]=src/auth.rs" ) ;
169+ assert_eq ! ( records[ 2 ] . name, "semantic_preferred_file[1]=src/db.rs" ) ;
170+ }
171+ }
0 commit comments