1use std::{
2 collections::{BTreeSet, HashMap, HashSet},
3 time::Instant,
4};
5
6use crate::{
7 api::{
8 ApiError, CodeGraphContextResponse, CodeRepositoryFreshnessDiagnostics,
9 CodeRepositoryFreshnessState, CodeRepositoryQueryResponse, RequestContext,
10 },
11 application::RelayKnowledgeService,
12 domain::{
13 CodeGraphCodeExcerpt, CodeGraphContextBudget, CodeGraphContextPack,
14 CodeGraphContextProvenance, CodeGraphContextRequest, CodeGraphImpactHint, CodeQueryKind,
15 CodeRetrievalHit, CodeRetrievalLayer, CodeRetrievalRequest,
16 },
17};
18
19const ENTRY_QUERY_KINDS: [CodeQueryKind; 3] = [
20 CodeQueryKind::Hybrid,
21 CodeQueryKind::Definition,
22 CodeQueryKind::Symbol,
23];
24const EXPANSION_QUERY_KINDS: [CodeQueryKind; 4] = [
25 CodeQueryKind::References,
26 CodeQueryKind::Callers,
27 CodeQueryKind::Callees,
28 CodeQueryKind::Imports,
29];
30const MAX_CONTEXT_SEEDS: usize = 3;
31const MAX_EXPANSION_LIMIT: usize = 4;
32
33impl RelayKnowledgeService {
34 pub async fn codegraph_context(
36 &self,
37 request: CodeGraphContextRequest,
38 context: RequestContext,
39 ) -> Result<CodeGraphContextResponse, ApiError> {
40 let started = Instant::now();
41 let mut candidate_count = 0usize;
42 let mut entry_points = Vec::new();
43 let mut freshness_parts = Vec::new();
44 let mut context_request = request.clone();
45 let mut primary = None;
46
47 for kind in ENTRY_QUERY_KINDS {
48 let response = self
49 .run_context_query(
50 &context_request,
51 kind,
52 request.query.clone(),
53 request.limit,
54 &context,
55 None,
56 )
57 .await?;
58 candidate_count = candidate_count.saturating_add(response.results.len());
59 push_unique_hits(&mut entry_points, response.results.clone());
60 if kind == CodeQueryKind::Hybrid {
61 context_request = pinned_context_request(&request, &response.scope);
62 primary = Some(response.clone());
63 } else {
64 freshness_parts.push(response.freshness);
65 }
66 }
67
68 let primary = primary.expect("hybrid entry query always runs");
69 let expansion_constraints = context_query_constraints(&context_request)
70 .map_err(|error| ApiError::invalid_argument(error.to_string()))?;
71 let seeds = context_seeds(&entry_points);
72 let mut related_symbols = Vec::new();
73 let mut graph_paths = Vec::new();
74 let mut context_roles = HashMap::new();
75
76 for seed in &seeds {
77 let seed_query = seed_query(seed, &request.query);
78 for kind in EXPANSION_QUERY_KINDS {
79 let response = self
80 .run_context_query(
81 &context_request,
82 kind,
83 seed_query.clone(),
84 request.limit.min(MAX_EXPANSION_LIMIT),
85 &context,
86 Some(&expansion_constraints),
87 )
88 .await?;
89 candidate_count = candidate_count.saturating_add(response.results.len());
90 remember_context_roles(&mut context_roles, kind, &response.results);
91 let results = response.results;
92 if kind == CodeQueryKind::References {
93 push_unique_hits(&mut related_symbols, results);
94 } else {
95 push_unique_hits(&mut graph_paths, results);
96 }
97 freshness_parts.push(response.freshness);
98 }
99 }
100
101 let count_truncated = truncate_hits(&mut entry_points, request.limit)
102 | truncate_hits(&mut related_symbols, request.limit)
103 | truncate_hits(&mut graph_paths, request.limit);
104 apply_code_visibility(&mut entry_points, request.include_code);
105 apply_code_visibility(&mut related_symbols, request.include_code);
106 apply_code_visibility(&mut graph_paths, request.include_code);
107
108 let mut pack = CodeGraphContextPack {
109 code_excerpts: code_excerpts(
110 request.include_code,
111 &entry_points,
112 &related_symbols,
113 &graph_paths,
114 &context_roles,
115 ),
116 impact_hints: impact_hints(&graph_paths, &context_roles),
117 entry_points,
118 related_symbols,
119 graph_paths,
120 };
121 let byte_truncated = pack_to_budget(&mut pack, request.max_context_bytes, &context_roles);
122 let mut truncated = count_truncated | byte_truncated;
123 let context_bytes = serialized_context_bytes(&pack);
124 truncated |= primary.results.len() > request.limit;
125 let retrieval_layers = retrieval_layers(&pack);
126 let mut freshness = merge_context_freshness(primary.freshness, freshness_parts);
127 freshness.merge_direct_source_read_paths(context_paths(&pack));
128 let returned_count = pack.entry_points.len()
129 + pack.related_symbols.len()
130 + pack.graph_paths.len()
131 + pack.code_excerpts.len();
132 let mut diagnostics = vec![format!(
133 "Expanded {} seed(s) through references, callers, callees, and imports with bounded limits.",
134 seeds.len()
135 )];
136 if count_truncated {
137 diagnostics.push("Context pack was truncated to fit the requested limit.".to_owned());
138 }
139 if byte_truncated {
140 diagnostics.push("Context pack was truncated to fit max_context_bytes.".to_owned());
141 }
142
143 Ok(CodeGraphContextResponse {
144 metadata: primary.metadata,
145 query: request.query.clone(),
146 repository_scope: primary.scope,
147 freshness,
148 budget: CodeGraphContextBudget {
149 limit: request.limit,
150 max_context_bytes: request.max_context_bytes,
151 candidate_count,
152 returned_count,
153 context_bytes,
154 elapsed_ms: started.elapsed().as_millis().try_into().unwrap_or(u64::MAX),
155 },
156 truncated,
157 retrieval_layers,
158 request,
159 pack,
160 diagnostics,
161 })
162 }
163
164 async fn run_context_query(
165 &self,
166 request: &CodeGraphContextRequest,
167 kind: CodeQueryKind,
168 query: String,
169 limit: usize,
170 context: &RequestContext,
171 constraints: Option<&CodeRetrievalRequest>,
172 ) -> Result<CodeRepositoryQueryResponse, ApiError> {
173 let mut retrieval = CodeRetrievalRequest::new(
174 query,
175 request.repository.clone(),
176 kind,
177 limit,
178 request.freshness_policy,
179 )
180 .map_err(|error| ApiError::invalid_argument(error.to_string()))?;
181 retrieval.exclude_generated = request.exclude_generated;
182 if let Some(constraints) = constraints {
183 carry_context_filters(&mut retrieval, constraints);
184 }
185
186 self.query_code_repository(retrieval, context.clone()).await
187 }
188}
189
190fn context_query_constraints(
191 request: &CodeGraphContextRequest,
192) -> Result<CodeRetrievalRequest, crate::domain::DomainError> {
193 CodeRetrievalRequest::new(
194 request.query.clone(),
195 request.repository.clone(),
196 CodeQueryKind::Hybrid,
197 request.limit,
198 request.freshness_policy,
199 )
200}
201
202fn carry_context_filters(target: &mut CodeRetrievalRequest, source: &CodeRetrievalRequest) {
203 target.query_language_filters = source.query_language_filters.clone();
204 target.query_path_substrings = source.query_path_substrings.clone();
205}
206
207fn pinned_context_request(
208 original: &CodeGraphContextRequest,
209 primary_scope: &crate::api::CodeRepositoryScopeMetadata,
210) -> CodeGraphContextRequest {
211 let mut pinned = original.clone();
212 if !primary_scope.resolved_commit_sha.is_empty() {
213 pinned.repository.ref_selector = primary_scope.resolved_commit_sha.clone();
214 }
215 pinned
216}
217
218fn context_seeds(entry_points: &[CodeRetrievalHit]) -> Vec<CodeRetrievalHit> {
219 let mut seeds = entry_points.to_vec();
220 seeds.sort_by(|left, right| right.score.total_cmp(&left.score));
221 seeds.truncate(MAX_CONTEXT_SEEDS);
222 seeds
223}
224
225fn push_unique_hits(target: &mut Vec<CodeRetrievalHit>, hits: Vec<CodeRetrievalHit>) {
226 let mut keys = target.iter().map(hit_key).collect::<HashSet<_>>();
227 for hit in hits {
228 if keys.insert(hit_key(&hit)) {
229 target.push(hit);
230 }
231 }
232 target.sort_by(|left, right| right.score.total_cmp(&left.score));
233}
234
235fn truncate_hits(hits: &mut Vec<CodeRetrievalHit>, limit: usize) -> bool {
236 if hits.len() > limit {
237 hits.truncate(limit);
238 true
239 } else {
240 false
241 }
242}
243
244fn remember_context_roles(
245 roles: &mut HashMap<String, CodeQueryKind>,
246 kind: CodeQueryKind,
247 hits: &[CodeRetrievalHit],
248) {
249 for hit in hits {
250 roles.entry(hit_key(hit)).or_insert(kind);
251 }
252}
253
254fn hit_key(hit: &CodeRetrievalHit) -> String {
255 format!(
256 "{}:{}:{}:{}:{}",
257 hit.path,
258 hit.line_range.start,
259 hit.line_range.end,
260 hit.symbol_snapshot_id.as_deref().unwrap_or(""),
261 hit.edge_kind.as_deref().unwrap_or("")
262 )
263}
264
265fn seed_query(hit: &CodeRetrievalHit, fallback: &str) -> String {
266 hit.canonical_symbol_id
267 .as_deref()
268 .and_then(extract_searchable_symbol_tail)
269 .unwrap_or(fallback)
270 .to_owned()
271}
272
273fn extract_searchable_symbol_tail(value: &str) -> Option<&str> {
274 value
275 .rsplit([':', '/', '#', '.', '@', '[', ']'])
276 .find(|part| part.chars().any(|character| character.is_alphanumeric()))
277}
278
279fn apply_code_visibility(hits: &mut [CodeRetrievalHit], include_code: bool) {
280 if include_code {
281 return;
282 }
283 for hit in hits {
284 hit.excerpt.clear();
285 }
286}
287
288fn code_excerpts(
289 include_code: bool,
290 entry_points: &[CodeRetrievalHit],
291 related_symbols: &[CodeRetrievalHit],
292 graph_paths: &[CodeRetrievalHit],
293 context_roles: &HashMap<String, CodeQueryKind>,
294) -> Vec<CodeGraphCodeExcerpt> {
295 if !include_code {
296 return Vec::new();
297 }
298 entry_points
299 .iter()
300 .chain(related_symbols)
301 .chain(graph_paths)
302 .filter(|hit| !hit.excerpt.trim().is_empty())
303 .map(|hit| CodeGraphCodeExcerpt {
304 path: hit.path.clone(),
305 language_id: hit.language_id.clone(),
306 line_range: hit.line_range.clone(),
307 symbol_snapshot_id: hit.symbol_snapshot_id.clone(),
308 provenance: CodeGraphContextProvenance {
309 query_kind: provenance_kind(hit, context_roles),
310 retrieval_layers: hit.retrieval_layers.clone(),
311 score: hit.score,
312 },
313 excerpt: hit.excerpt.clone(),
314 })
315 .collect()
316}
317
318fn impact_hints(
319 graph_paths: &[CodeRetrievalHit],
320 context_roles: &HashMap<String, CodeQueryKind>,
321) -> Vec<CodeGraphImpactHint> {
322 graph_paths
323 .iter()
324 .map(|hit| CodeGraphImpactHint {
325 path: hit.path.clone(),
326 line_range: hit.line_range.clone(),
327 relationship: context_roles
328 .get(&hit_key(hit))
329 .map(context_role_relationship)
330 .or(hit.edge_kind.as_deref())
331 .unwrap_or_else(|| relationship_from_layers(&hit.retrieval_layers))
332 .to_owned(),
333 symbol_snapshot_id: hit.symbol_snapshot_id.clone(),
334 retrieval_layers: hit.retrieval_layers.clone(),
335 score: hit.score,
336 })
337 .collect()
338}
339
340fn provenance_kind(
341 hit: &CodeRetrievalHit,
342 context_roles: &HashMap<String, CodeQueryKind>,
343) -> CodeQueryKind {
344 if let Some(kind) = context_roles.get(&hit_key(hit)) {
345 *kind
346 } else if hit
347 .retrieval_layers
348 .contains(&CodeRetrievalLayer::Reference)
349 {
350 CodeQueryKind::References
351 } else if hit
352 .retrieval_layers
353 .contains(&CodeRetrievalLayer::CallGraph)
354 {
355 CodeQueryKind::Callers
356 } else if hit
357 .retrieval_layers
358 .contains(&CodeRetrievalLayer::ImportGraph)
359 {
360 CodeQueryKind::Imports
361 } else if hit
362 .retrieval_layers
363 .contains(&CodeRetrievalLayer::Definition)
364 {
365 CodeQueryKind::Definition
366 } else if hit.retrieval_layers.contains(&CodeRetrievalLayer::Symbol) {
367 CodeQueryKind::Symbol
368 } else {
369 CodeQueryKind::Hybrid
370 }
371}
372
373fn context_role_relationship(kind: &CodeQueryKind) -> &'static str {
374 match kind {
375 CodeQueryKind::References => "reference",
376 CodeQueryKind::Callers => "caller",
377 CodeQueryKind::Callees => "callee",
378 CodeQueryKind::Imports => "import",
379 _ => "context",
380 }
381}
382
383fn relationship_from_layers(layers: &[CodeRetrievalLayer]) -> &'static str {
384 if layers.contains(&CodeRetrievalLayer::CallGraph) {
385 "call_graph"
386 } else if layers.contains(&CodeRetrievalLayer::ImportGraph) {
387 "import_graph"
388 } else if layers.contains(&CodeRetrievalLayer::Reference) {
389 "reference"
390 } else {
391 "related"
392 }
393}
394
395fn merge_context_freshness(
396 mut primary: CodeRepositoryFreshnessDiagnostics,
397 parts: Vec<CodeRepositoryFreshnessDiagnostics>,
398) -> CodeRepositoryFreshnessDiagnostics {
399 for freshness in parts {
400 primary.state = worse_freshness_state(primary.state, freshness.state);
401 primary.scope_stale |= freshness.scope_stale;
402 primary.direct_source_read_required |= freshness.direct_source_read_required;
403 primary.index_lag.requested_ref_indexed &= freshness.index_lag.requested_ref_indexed;
404 primary.index_lag.pending_task_count = primary
405 .index_lag
406 .pending_task_count
407 .max(freshness.index_lag.pending_task_count);
408 primary.index_lag.pending_file_count = max_optional_usize(
409 primary.index_lag.pending_file_count,
410 freshness.index_lag.pending_file_count,
411 );
412 primary.pending.active_for_repository |= freshness.pending.active_for_repository;
413 primary.pending.active_matches_request |= freshness.pending.active_matches_request;
414 primary.pending.queue_depth = primary
415 .pending
416 .queue_depth
417 .max(freshness.pending.queue_depth);
418 primary.pending.queued_task_count = primary
419 .pending
420 .queued_task_count
421 .max(freshness.pending.queued_task_count);
422 primary.pending.running_task_count = primary
423 .pending
424 .running_task_count
425 .max(freshness.pending.running_task_count);
426 primary.pending.retrying_task_count = primary
427 .pending
428 .retrying_task_count
429 .max(freshness.pending.retrying_task_count);
430 primary.pending.dead_letter_task_count = primary
431 .pending
432 .dead_letter_task_count
433 .max(freshness.pending.dead_letter_task_count);
434 primary.pending.running_lease_count = primary
435 .pending
436 .running_lease_count
437 .max(freshness.pending.running_lease_count);
438 primary.stale_reason = merge_reason(primary.stale_reason.take(), freshness.stale_reason);
439 primary.degraded_reason =
440 merge_reason(primary.degraded_reason.take(), freshness.degraded_reason);
441 primary.merge_direct_source_read_paths(freshness.direct_source_read_paths);
442 }
443 primary
444}
445
446fn worse_freshness_state(
447 left: CodeRepositoryFreshnessState,
448 right: CodeRepositoryFreshnessState,
449) -> CodeRepositoryFreshnessState {
450 if freshness_state_rank(left) >= freshness_state_rank(right) {
451 left
452 } else {
453 right
454 }
455}
456
457fn freshness_state_rank(state: CodeRepositoryFreshnessState) -> u8 {
458 match state {
459 CodeRepositoryFreshnessState::Fresh => 0,
460 CodeRepositoryFreshnessState::Degraded => 1,
461 CodeRepositoryFreshnessState::Stale => 2,
462 CodeRepositoryFreshnessState::Pending => 3,
463 }
464}
465
466fn max_optional_usize(left: Option<usize>, right: Option<usize>) -> Option<usize> {
467 match (left, right) {
468 (Some(left), Some(right)) => Some(left.max(right)),
469 (Some(value), None) | (None, Some(value)) => Some(value),
470 (None, None) => None,
471 }
472}
473
474fn merge_reason(left: Option<String>, right: Option<String>) -> Option<String> {
475 match (left, right) {
476 (Some(left), Some(right)) if left == right => Some(left),
477 (Some(left), Some(right)) => Some(format!("{left}; {right}")),
478 (Some(reason), None) | (None, Some(reason)) => Some(reason),
479 (None, None) => None,
480 }
481}
482
483fn pack_to_budget(
484 pack: &mut CodeGraphContextPack,
485 max_context_bytes: usize,
486 context_roles: &HashMap<String, CodeQueryKind>,
487) -> bool {
488 let mut truncated = false;
489 while serialized_context_bytes(pack) > max_context_bytes {
490 let removed_code_excerpt = pack.code_excerpts.pop().is_some();
491 let cleared_expansion_excerpts =
492 !removed_code_excerpt && clear_expansion_hit_excerpts(pack);
493 if removed_code_excerpt || cleared_expansion_excerpts {
494 truncated = true;
495 } else if pack.graph_paths.pop().is_some() {
496 pack.impact_hints = impact_hints(&pack.graph_paths, context_roles);
497 truncated = true;
498 } else if pack.related_symbols.pop().is_some()
499 || clear_entry_hit_excerpts(pack)
500 || pack.entry_points.pop().is_some()
501 {
502 truncated = true;
503 } else {
504 clear_hit_excerpts(pack);
505 return true;
506 }
507 }
508
509 truncated
510}
511
512fn clear_hit_excerpts(pack: &mut CodeGraphContextPack) {
513 clear_expansion_hit_excerpts(pack);
514 clear_entry_hit_excerpts(pack);
515 pack.code_excerpts.clear();
516 pack.impact_hints.clear();
517}
518
519fn clear_expansion_hit_excerpts(pack: &mut CodeGraphContextPack) -> bool {
520 let had_excerpts = pack
521 .related_symbols
522 .iter()
523 .chain(&pack.graph_paths)
524 .any(|hit| !hit.excerpt.is_empty());
525 apply_code_visibility(&mut pack.related_symbols, false);
526 apply_code_visibility(&mut pack.graph_paths, false);
527
528 had_excerpts
529}
530
531fn clear_entry_hit_excerpts(pack: &mut CodeGraphContextPack) -> bool {
532 let had_excerpts = pack.entry_points.iter().any(|hit| !hit.excerpt.is_empty());
533 apply_code_visibility(&mut pack.entry_points, false);
534
535 had_excerpts
536}
537
538fn retrieval_layers(pack: &CodeGraphContextPack) -> Vec<CodeRetrievalLayer> {
539 let mut layers = BTreeSet::new();
540 for hit in pack
541 .entry_points
542 .iter()
543 .chain(&pack.related_symbols)
544 .chain(&pack.graph_paths)
545 {
546 for layer in &hit.retrieval_layers {
547 layers.insert(layer.as_str());
548 }
549 }
550
551 layers
552 .into_iter()
553 .filter_map(layer_from_str)
554 .collect::<Vec<_>>()
555}
556
557fn layer_from_str(value: &str) -> Option<CodeRetrievalLayer> {
558 match value {
559 "lexical" => Some(CodeRetrievalLayer::Lexical),
560 "symbol" => Some(CodeRetrievalLayer::Symbol),
561 "definition" => Some(CodeRetrievalLayer::Definition),
562 "reference" => Some(CodeRetrievalLayer::Reference),
563 "call_graph" => Some(CodeRetrievalLayer::CallGraph),
564 "import_graph" => Some(CodeRetrievalLayer::ImportGraph),
565 "sbom" => Some(CodeRetrievalLayer::Sbom),
566 "impact" => Some(CodeRetrievalLayer::Impact),
567 "text_fallback" => Some(CodeRetrievalLayer::TextFallback),
568 _ => None,
569 }
570}
571
572fn serialized_context_bytes<T: serde::Serialize>(value: &T) -> usize {
573 serde_json::to_vec(value)
574 .map(|bytes| bytes.len())
575 .unwrap_or(usize::MAX / 4)
576}
577
578fn context_paths(pack: &CodeGraphContextPack) -> Vec<String> {
579 pack.entry_points
580 .iter()
581 .chain(&pack.related_symbols)
582 .chain(&pack.graph_paths)
583 .map(|hit| hit.path.clone())
584 .collect::<BTreeSet<_>>()
585 .into_iter()
586 .collect()
587}
588
589#[cfg(test)]
590mod tests {
591 use super::*;
592 use crate::{
593 api::{
594 CodeRepositoryIndexLag, CodeRepositoryPendingIndexWork, CodeRepositoryScopeMetadata,
595 },
596 domain::{CodeRepositorySelector, FreshnessPolicy, RepositoryCodeRange},
597 };
598
599 #[test]
600 fn context_roles_preserve_edge_kind_and_drive_provenance_and_hints() {
601 let mut hit = call_graph_hit();
602 hit.edge_kind = Some("call".to_owned());
603 let mut roles = HashMap::new();
604 remember_context_roles(&mut roles, CodeQueryKind::Callees, &[hit.clone()]);
605 let excerpts = code_excerpts(true, &[], &[], &[hit.clone()], &roles);
606 let hints = impact_hints(&[hit.clone()], &roles);
607
608 assert_eq!(hit.edge_kind.as_deref(), Some("call"));
609 assert_eq!(provenance_kind(&hit, &roles), CodeQueryKind::Callees);
610 assert_eq!(excerpts[0].provenance.query_kind, CodeQueryKind::Callees);
611 assert_eq!(hints[0].relationship, "callee");
612 }
613
614 #[test]
615 fn count_truncation_reports_when_unique_hits_exceed_limit() {
616 let mut hits = vec![call_graph_hit(), call_graph_hit_at("src/main.rs")];
617
618 assert!(truncate_hits(&mut hits, 1));
619 assert_eq!(hits.len(), 1);
620 }
621
622 #[test]
623 fn pinned_context_request_uses_primary_served_commit_for_followups() {
624 let request = CodeGraphContextRequest::new(
625 CodeRepositorySelector::new("repo", "HEAD", Vec::new(), Vec::new()).unwrap(),
626 "retry policy",
627 3,
628 FreshnessPolicy::AllowStale,
629 1024,
630 true,
631 false,
632 )
633 .unwrap();
634
635 let pinned = pinned_context_request(&request, &scope_metadata("commit-a"));
636
637 assert_eq!(pinned.repository.ref_selector, "commit-a");
638 assert_eq!(request.repository.ref_selector, "HEAD");
639 }
640
641 #[test]
642 fn context_freshness_merges_degraded_expansion_state_and_reason() {
643 let primary = freshness(CodeRepositoryFreshnessState::Fresh, None, Vec::new());
644 let degraded = freshness(
645 CodeRepositoryFreshnessState::Degraded,
646 Some("parser degraded"),
647 vec!["src/lib.rs".to_owned()],
648 );
649
650 let merged = merge_context_freshness(primary, vec![degraded]);
651
652 assert_eq!(merged.state, CodeRepositoryFreshnessState::Degraded);
653 assert_eq!(merged.degraded_reason.as_deref(), Some("parser degraded"));
654 assert_eq!(merged.direct_source_read_paths, ["src/lib.rs"]);
655 }
656
657 #[test]
658 fn budget_truncation_keeps_impact_hints_aligned_with_graph_paths() {
659 let graph_paths = (0..12)
660 .map(|index| call_graph_hit_at(&format!("src/path_{index}.rs")))
661 .collect::<Vec<_>>();
662 let mut pack = CodeGraphContextPack {
663 entry_points: Vec::new(),
664 related_symbols: Vec::new(),
665 impact_hints: impact_hints(&graph_paths, &HashMap::new()),
666 code_excerpts: Vec::new(),
667 graph_paths,
668 };
669
670 assert!(pack_to_budget(&mut pack, 1024, &HashMap::new()));
671 assert!(serialized_context_bytes(&pack) <= 1024);
672 assert_eq!(pack.impact_hints.len(), pack.graph_paths.len());
673 for hint in &pack.impact_hints {
674 assert!(pack.graph_paths.iter().any(|hit| hit.path == hint.path));
675 }
676 }
677
678 #[test]
679 fn budget_truncation_clears_hit_excerpts_before_dropping_evidence() {
680 let mut hit = call_graph_hit();
681 hit.excerpt = "x".repeat(5000);
682 let mut pack = CodeGraphContextPack {
683 entry_points: vec![hit],
684 related_symbols: Vec::new(),
685 graph_paths: Vec::new(),
686 impact_hints: Vec::new(),
687 code_excerpts: code_excerpts(true, &[call_graph_hit()], &[], &[], &HashMap::new()),
688 };
689
690 assert!(pack_to_budget(&mut pack, 1024, &HashMap::new()));
691 assert_eq!(pack.entry_points.len(), 1);
692 assert!(pack.entry_points[0].excerpt.is_empty());
693 assert!(serialized_context_bytes(&pack) <= 1024);
694 }
695
696 #[test]
697 fn budget_truncation_preserves_entry_excerpt_before_expansion_evidence() {
698 let mut entry = call_graph_hit_at("src/context.rs");
699 entry.excerpt = "pub struct AgentContextPackBuilder;".to_owned();
700 let graph_paths = (0..8)
701 .map(|index| {
702 let mut hit = call_graph_hit_at(&format!("src/expansion_{index}.rs"));
703 hit.excerpt = "x".repeat(3000);
704 hit
705 })
706 .collect::<Vec<_>>();
707 let mut pack = CodeGraphContextPack {
708 entry_points: vec![entry],
709 related_symbols: Vec::new(),
710 impact_hints: impact_hints(&graph_paths, &HashMap::new()),
711 code_excerpts: Vec::new(),
712 graph_paths,
713 };
714
715 assert!(pack_to_budget(&mut pack, 2048, &HashMap::new()));
716 assert_eq!(pack.entry_points.len(), 1);
717 assert!(
718 pack.entry_points[0]
719 .excerpt
720 .contains("AgentContextPackBuilder")
721 );
722 assert!(serialized_context_bytes(&pack) <= 2048);
723 }
724
725 #[test]
726 fn expansion_queries_carry_inline_scope_filters_without_kind_or_name_terms() {
727 let source = CodeRetrievalRequest::new(
728 "path:src lang:rust name:Retry kind:function retry policy",
729 CodeRepositorySelector::new("repo", "HEAD", Vec::new(), Vec::new()).unwrap(),
730 CodeQueryKind::Hybrid,
731 3,
732 crate::domain::FreshnessPolicy::AllowStale,
733 )
734 .unwrap();
735 let mut target = CodeRetrievalRequest::new(
736 "RetryPolicy",
737 CodeRepositorySelector::new("repo", "HEAD", Vec::new(), Vec::new()).unwrap(),
738 CodeQueryKind::References,
739 3,
740 crate::domain::FreshnessPolicy::AllowStale,
741 )
742 .unwrap();
743
744 carry_context_filters(&mut target, &source);
745
746 assert_eq!(target.query_path_substrings, ["src"]);
747 assert_eq!(target.query_language_filters, ["rust"]);
748 assert!(target.query_kind_filters.is_empty());
749 assert!(target.query_name_substrings.is_empty());
750 }
751
752 fn call_graph_hit() -> CodeRetrievalHit {
753 call_graph_hit_at("src/lib.rs")
754 }
755
756 fn call_graph_hit_at(path: &str) -> CodeRetrievalHit {
757 CodeRetrievalHit {
758 repository_id: "repo".to_owned(),
759 scope_id: "scope".to_owned(),
760 resolved_commit_sha: "commit".to_owned(),
761 tree_hash: "tree".to_owned(),
762 path: path.to_owned(),
763 language_id: "rust".to_owned(),
764 byte_range: RepositoryCodeRange { start: 0, end: 1 },
765 line_range: RepositoryCodeRange { start: 1, end: 1 },
766 symbol_snapshot_id: None,
767 canonical_symbol_id: None,
768 file_id: None,
769 retrieval_layers: vec![CodeRetrievalLayer::CallGraph],
770 index_versions: Vec::new(),
771 stale: false,
772 staleness_hint: None,
773 degraded_reason: None,
774 edge_kind: None,
775 edge_resolution_state: None,
776 edge_target_hint: None,
777 edge_confidence_basis_points: None,
778 edge_confidence_tier: None,
779 score: 1.0,
780 excerpt: "call();".to_owned(),
781 }
782 }
783
784 fn scope_metadata(resolved_commit_sha: &str) -> CodeRepositoryScopeMetadata {
785 CodeRepositoryScopeMetadata {
786 scope_id: "scope".to_owned(),
787 repository_id: "repo".to_owned(),
788 alias: "repo".to_owned(),
789 requested_ref: "HEAD".to_owned(),
790 resolved_commit_sha: resolved_commit_sha.to_owned(),
791 tree_hash: "tree".to_owned(),
792 path_filters: Vec::new(),
793 language_filters: Vec::new(),
794 indexed_file_count: 1,
795 index_versions: Vec::new(),
796 stale: false,
797 }
798 }
799
800 fn freshness(
801 state: CodeRepositoryFreshnessState,
802 degraded_reason: Option<&str>,
803 direct_source_read_paths: Vec<String>,
804 ) -> CodeRepositoryFreshnessDiagnostics {
805 CodeRepositoryFreshnessDiagnostics {
806 state,
807 freshness_policy: FreshnessPolicy::AllowStale,
808 graph_version: 1,
809 source_scope: Some("scope".to_owned()),
810 scope_stale: matches!(
811 state,
812 CodeRepositoryFreshnessState::Stale | CodeRepositoryFreshnessState::Pending
813 ),
814 stale_reason: None,
815 degraded_reason: degraded_reason.map(str::to_owned),
816 index_lag: CodeRepositoryIndexLag {
817 requested_ref: "HEAD".to_owned(),
818 requested_resolved_ref: "commit".to_owned(),
819 served_ref: "commit".to_owned(),
820 requested_ref_indexed: true,
821 pending_file_count: None,
822 pending_task_count: 0,
823 },
824 pending: CodeRepositoryPendingIndexWork::default(),
825 cursor: None,
826 direct_source_read_required: false,
827 direct_source_read_paths,
828 agent_instructions: Vec::new(),
829 }
830 }
831}