From 200bfde21c20eb7d040fc1f1ff537ad70e47771c Mon Sep 17 00:00:00 2001 From: ScriptedAlchemy Date: Sat, 26 Sep 2026 16:23:03 +0000 Subject: [PATCH] perf(code-index): resolve trait dispatch from relation keys With the MCP defaults (depth 3, trait dispatch on) every callees page read each direct callee's summary and incoming edges to find its container, so a page of a 309-row listing read 1,066 entities over 315 queries. The dispatch check now finds every callee's container in one batched relation-key step and reads only the distinct containers, and the callees, implementors' children, and names under a trait. Callers and the other symbol-graph primitive reads (symbol and signature search, implementations, type hierarchy) open a metered reader too, so their envelopes carry the request cost receipt. Callable-code reads report the files of the rows they returned as touched files, so their responses carry the tracedecay_metrics footer. Fixes #2226 --- .../src/primitives/runtime.rs | 10 +- .../src/primitives/symbol_graph.rs | 55 +++- .../src/code_index_scheduler/queries.rs | 296 +++++++++--------- .../src/retrieval/callable_code.rs | 50 +++ .../src/retrieval/callable_code_service.rs | 10 +- .../tracedecay-contracts/src/retrieval/mod.rs | 12 +- .../src/retrieval/symbol_graph.rs | 23 +- .../graph_query_test/relation_page_cost.rs | 69 +++- .../typed_evidence_trailers.rs | 121 ++++++- 9 files changed, 459 insertions(+), 187 deletions(-) diff --git a/crates/tracedecay-application/src/primitives/runtime.rs b/crates/tracedecay-application/src/primitives/runtime.rs index 6692932968..49edaea18b 100644 --- a/crates/tracedecay-application/src/primitives/runtime.rs +++ b/crates/tracedecay-application/src/primitives/runtime.rs @@ -35,8 +35,8 @@ use tracedecay_contracts::{ Deadline, DisclosureClass, EvidenceCoverage, EvidenceDomain, EvidencePacket, FreshnessState, LegalAction, Omission, OmissionReason, OpaqueCursor, OperationBudgetUsage, OperationReceipt, OperationTermination, PageCursor, PageRequest, PageState, PolicyDecisionRef, RequestAdmission, - RequestContext, RequestId, ResolvedScope, RetrievalEvidence, RetryDirective, SafeDiagnostic, - TemporalState, + RequestContext, RequestCostReceiptV1, RequestId, ResolvedScope, RetrievalEvidence, + RetryDirective, SafeDiagnostic, TemporalState, }; use tracedecay_domain::text::forward_slash_path; use tracedecay_domain::{CodeGenerationId, CommitId, ComponentVersion, UtcMicros}; @@ -1312,6 +1312,7 @@ fn symbol_outcome( page, finished_at, budget, + cost, } => symbol_page( access, context, @@ -1320,12 +1321,14 @@ fn symbol_outcome( page, finished_at, budget, + cost, false, ), SymbolGraphPortOutcome::Partial { page, finished_at, budget, + cost, } => symbol_page( access, context, @@ -1334,6 +1337,7 @@ fn symbol_outcome( page, finished_at, budget, + cost, true, ), SymbolGraphPortOutcome::Failed { failure, .. } => { @@ -1351,6 +1355,7 @@ fn symbol_page( page: SymbolGraphPage, finished_at: UtcMicros, budget: OperationBudgetUsage, + cost: Option, partial: bool, ) -> Result, ApplicationContractError> { let returned = page.items.len() as u64; @@ -1385,6 +1390,7 @@ fn symbol_page( )?; if let Ok(envelope) = &mut result { envelope.touched_files = touched_files; + envelope.cost = cost; } if unsupported && let Ok(envelope) = &mut result diff --git a/crates/tracedecay-application/src/primitives/symbol_graph.rs b/crates/tracedecay-application/src/primitives/symbol_graph.rs index b1195ca28c..bb54374666 100644 --- a/crates/tracedecay-application/src/primitives/symbol_graph.rs +++ b/crates/tracedecay-application/src/primitives/symbol_graph.rs @@ -5,7 +5,8 @@ use std::pin::Pin; use std::sync::Arc; use tracedecay_code_index::graph_projection::{ - CodeGraphInteractiveReader, CodeGraphSymbolBindingV1, CodeGraphSymbolSummaryV1, + CodeGraphInteractiveReader, CodeGraphReadCostMeter, CodeGraphSymbolBindingV1, + CodeGraphSymbolSummaryV1, }; use tracedecay_code_index::lineage::LineageSymbolRecordV1; use tracedecay_contracts::retrieval::{ @@ -191,7 +192,7 @@ where &request.meta.page, "search", &claim, - graph.freshness, + &graph, records, Vec::new(), None, @@ -262,7 +263,7 @@ where &request.meta.page, "exact", &claim, - graph.freshness, + &graph, records, Vec::new(), None, @@ -325,7 +326,7 @@ where &request.meta.page, "signature", &claim, - graph.freshness, + &graph, records, Vec::new(), None, @@ -414,7 +415,7 @@ where &request.meta.page, "implementations", &claim, - graph.freshness, + &graph, records, Vec::new(), None, @@ -468,7 +469,7 @@ where &request.meta.page, "hierarchy", &claim, - graph.freshness, + &graph, Vec::new(), Vec::new(), None, @@ -533,7 +534,7 @@ where &request.meta.page, "hierarchy", &claim, - graph.freshness, + &graph, records, Vec::new(), None, @@ -591,7 +592,7 @@ where &request.meta.page, "callers", &claim, - graph.freshness, + &graph, records, gaps, None, @@ -689,7 +690,7 @@ where &request.meta.page, "callees", &claim, - graph.freshness, + &graph, records, Vec::new(), None, @@ -748,7 +749,7 @@ where &request.meta.page, "impact", &claim, - graph.freshness, + &graph, records, Vec::new(), Some(edge_count), @@ -761,7 +762,9 @@ where } struct OpenSymbolGraph { + /// Counts every store read on [`Self::cost`]. reader: CodeGraphInteractiveReader, + cost: CodeGraphReadCostMeter, cancellation: Arc, freshness: tracedecay_graph_query::CodeGraphReadFreshnessV1, } @@ -781,15 +784,18 @@ async fn open_graph( .await .map_err(|_| ())?; let freshness = verified.freshness(); + let cost = CodeGraphReadCostMeter::start(); let reader = verified .reader_with_cancellation( context.request, context.observed_at, Arc::clone(&cancellation), ) - .map_err(|_| ())?; + .map_err(|_| ())? + .metered(&cost); Ok(OpenSymbolGraph { reader, + cost, cancellation, freshness, }) @@ -1137,17 +1143,19 @@ async fn with_implementation_bodies( context: SymbolGraphPortContext<'_>, outcome: SymbolGraphPortOutcome, ) -> SymbolGraphPortOutcome { - let (page, partial, finished_at, budget) = match outcome { + let (page, partial, finished_at, budget, cost) = match outcome { SymbolGraphPortOutcome::Completed { page, finished_at, budget, - } => (page, false, finished_at, budget), + cost, + } => (page, false, finished_at, budget, cost), SymbolGraphPortOutcome::Partial { page, finished_at, budget, - } => (page, true, finished_at, budget), + cost, + } => (page, true, finished_at, budget, cost), SymbolGraphPortOutcome::Failed { failure, finished_at, @@ -1204,12 +1212,14 @@ async fn with_implementation_bodies( page, finished_at, budget, + cost, } } else { SymbolGraphPortOutcome::Completed { page, finished_at, budget, + cost, } } } @@ -1265,12 +1275,22 @@ async fn complete_or_failed( request: &PageRequest, lane: &str, claim: &SymbolGraphPageClaim, - freshness: tracedecay_graph_query::CodeGraphReadFreshnessV1, + graph: &OpenSymbolGraph, items: Vec, gaps: Vec, related_edge_count: Option, ) -> SymbolGraphPortOutcome { - let mut page = match paginate(cursors, context, request, lane, claim, freshness, items).await { + let mut page = match paginate( + cursors, + context, + request, + lane, + claim, + graph.freshness, + items, + ) + .await + { Ok(page) => page, Err(failure) => { return SymbolGraphPortOutcome::Failed { @@ -1282,17 +1302,20 @@ async fn complete_or_failed( }; page.related_edge_count = related_edge_count; page.support_gaps = gaps; + let cost = Some(graph.cost.receipt()); if page.support_gaps.is_empty() { SymbolGraphPortOutcome::Completed { page, finished_at: context.observed_at, budget: OperationBudgetUsage::default(), + cost, } } else { SymbolGraphPortOutcome::Partial { page, finished_at: context.observed_at, budget: OperationBudgetUsage::default(), + cost, } } } diff --git a/crates/tracedecay-code-index-runtime/src/code_index_scheduler/queries.rs b/crates/tracedecay-code-index-runtime/src/code_index_scheduler/queries.rs index d4de835ed9..199fa59780 100644 --- a/crates/tracedecay-code-index-runtime/src/code_index_scheduler/queries.rs +++ b/crates/tracedecay-code-index-runtime/src/code_index_scheduler/queries.rs @@ -4,7 +4,7 @@ //! It selects one already-mounted worktree generation and translates the //! generic lane evidence into the typed application-operation records. -use std::collections::{BTreeSet, HashMap}; +use std::collections::{BTreeMap, BTreeSet, HashMap}; use std::future::Future; use std::pin::Pin; use std::sync::{Arc, LazyLock}; @@ -1310,136 +1310,63 @@ fn check_dispatch_control( Ok(()) } -/// Visits canonical same-generation adjacency and stops as soon as `visit` -/// reports that the shared candidate capacity is exhausted. -/// -/// The Boolean return is false when adjacency remains unvisited, allowing the -/// public coverage receipt to expose that remainder as unknown rather than -/// claiming complete trait dispatch. -fn visit_trait_dispatch_targets( +/// One trait-dispatch expansion step over relation keys: `kinds` edges from +/// every `seed` at once, `reverse` for incoming. `None` means the step hit the +/// key ceiling, so dispatch is reported incomplete rather than guessed. +fn dispatch_step( reader: &CodeGraphInteractiveReader, - callee: &CodeGraphSymbolRefV1, - scope: &tracedecay_contracts::CodeQueryScope, + seeds: &[CodeGraphSymbolRefV1], + kinds: &[RelationEdgeKindV1], + reverse: bool, budget: RetrievalBudget, control: &Arc, - mut visit: impl FnMut( - &tracedecay_code_index::graph_projection::CodeGraphSymbolSummaryV1, - ) -> Result, -) -> Result { +) -> Result>>, DispatchExpansionStop> { check_dispatch_control(control.as_ref(), budget)?; - let Some(callee_summary) = reader + let step = reader + .relation_keys( + seeds, + kinds, + reverse, + MAX_RELATION_CANDIDATE_KEYS.saturating_add(1), + graph_read_cancellation(Arc::clone(control), budget.deadline_micros), + ) + .map_err(|_| dispatch_read_stop(control.as_ref(), budget))?; + if step.truncated { + return Ok(None); + } + Ok(Some( + step.per_seed + .into_iter() + .map(|keys| keys.into_iter().map(|key| key.neighbor).collect()) + .collect(), + )) +} + +fn dispatch_summary( + reader: &CodeGraphInteractiveReader, + symbol: &CodeGraphSymbolRefV1, + budget: RetrievalBudget, + control: &Arc, +) -> Result +{ + check_dispatch_control(control.as_ref(), budget)?; + reader .symbol_summary_for( - callee, + symbol, graph_read_cancellation(Arc::clone(control), budget.deadline_micros), ) .map_err(|_| dispatch_read_stop(control.as_ref(), budget))? - else { - return Ok(true); - }; - let callee_name = callee_summary + .ok_or(DispatchExpansionStop::Unavailable) +} + +fn summary_kind( + summary: &tracedecay_code_index::graph_projection::CodeGraphSymbolSummaryV1, +) -> Result, DispatchExpansionStop> { + summary .metadata .as_ref() - .ok_or(DispatchExpansionStop::Unavailable)? - .simple_name - .clone(); - let relation_limit = MAX_RELATION_CANDIDATE_KEYS; - check_dispatch_control(control.as_ref(), budget)?; - let parent_batches = match reader.callers( - std::slice::from_ref(&callee_summary.occurrence), - &[RelationEdgeKindV1::Contains], - relation_limit, - graph_read_cancellation(Arc::clone(control), budget.deadline_micros), - ) { - Ok(batches) => batches, - Err( - tracedecay_code_index::graph_projection::CodeGraphProjectionError::BudgetExhausted { - .. - }, - ) => return Ok(false), - Err(_) => return Err(dispatch_read_stop(control.as_ref(), budget)), - }; - let mut traits = Vec::new(); - for edge in parent_batches.into_iter().flatten() { - let metadata = edge - .neighbor - .metadata - .as_ref() - .ok_or(DispatchExpansionStop::Unavailable)?; - if matches!( - NodeKind::from_str(&metadata.kind), - Some(NodeKind::Trait | NodeKind::Interface | NodeKind::InterfaceType) - ) { - traits.push(edge.neighbor.occurrence); - } - } - if traits.is_empty() { - return Ok(true); - } - check_dispatch_control(control.as_ref(), budget)?; - let implementor_batches = match reader.callers( - &traits, - &[RelationEdgeKindV1::Implements], - relation_limit, - graph_read_cancellation(Arc::clone(control), budget.deadline_micros), - ) { - Ok(batches) => batches, - Err( - tracedecay_code_index::graph_projection::CodeGraphProjectionError::BudgetExhausted { - .. - }, - ) => return Ok(false), - Err(_) => return Err(dispatch_read_stop(control.as_ref(), budget)), - }; - let implementors = implementor_batches - .into_iter() - .flatten() - .map(|edge| edge.neighbor.occurrence) - .collect::>(); - if implementors.is_empty() { - return Ok(true); - } - check_dispatch_control(control.as_ref(), budget)?; - let child_batches = match reader.callees( - &implementors, - &[RelationEdgeKindV1::Contains], - relation_limit, - graph_read_cancellation(Arc::clone(control), budget.deadline_micros), - ) { - Ok(batches) => batches, - Err( - tracedecay_code_index::graph_projection::CodeGraphProjectionError::BudgetExhausted { - .. - }, - ) => return Ok(false), - Err(_) => return Err(dispatch_read_stop(control.as_ref(), budget)), - }; - for child in child_batches - .into_iter() - .flatten() - .map(|edge| edge.neighbor) - { - check_dispatch_control(control.as_ref(), budget)?; - let Some(metadata) = child.metadata.as_ref() else { - return Err(DispatchExpansionStop::Unavailable); - }; - if !matches!( - NodeKind::from_str(&metadata.kind), - Some(NodeKind::Function | NodeKind::Method) - ) || metadata.simple_name != callee_name - || !child - .binding - .as_ref() - .and_then(|binding| binding.logical_path.as_deref()) - .is_some_and(|path| path_is_in_code_query_scope(path, scope)) - { - continue; - } - if !visit(&child)? { - return Ok(false); - } - } - check_dispatch_control(control.as_ref(), budget)?; - Ok(true) + .map(|metadata| NodeKind::from_str(&metadata.kind)) + .ok_or(DispatchExpansionStop::Unavailable) } /// Appends the concrete impl methods reachable through each direct callee's @@ -1447,6 +1374,10 @@ fn visit_trait_dispatch_targets( /// page and hydrate like every other relation. Dispatch keys follow the /// direct keys in canonical order; a trait fan-out past the key ceiling /// leaves the listing incomplete rather than claiming full dispatch. +/// +/// Every callee's container comes from one batched relation-key step; only +/// the distinct containers, and the callees and impl children under a trait, +/// are read. A callee with no trait container costs no read at all. fn augment_callee_dispatch_keys( reader: &CodeGraphInteractiveReader, found: &mut GraphRelationKeysV1, @@ -1454,38 +1385,119 @@ fn augment_callee_dispatch_keys( budget: RetrievalBudget, control: &Arc, ) -> Result<(), DispatchExpansionStop> { - let mut seen = found + if found.keys.is_empty() { + return Ok(()); + } + let callees = found .keys .iter() .map(|key| key.symbol.clone()) - .collect::>(); + .collect::>(); + let Some(containers) = dispatch_step( + reader, + &callees, + &[RelationEdgeKindV1::Contains], + true, + budget, + control, + )? + else { + found.complete = false; + return Ok(()); + }; + let mut is_trait = BTreeMap::::new(); + for container in containers.iter().flatten() { + if is_trait.contains_key(container) { + continue; + } + let kind = summary_kind(&dispatch_summary(reader, container, budget, control)?)?; + is_trait.insert( + container.clone(), + matches!( + kind, + Some(NodeKind::Trait | NodeKind::Interface | NodeKind::InterfaceType) + ), + ); + } + let mut seen = callees.iter().cloned().collect::>(); let mut dispatch = Vec::new(); - for callee in &found.keys { - let exhausted = visit_trait_dispatch_targets( + 'callees: for (callee, containers) in found.keys.iter().zip(containers) { + let traits = containers + .into_iter() + .filter(|container| is_trait.get(container).copied().unwrap_or(false)) + .collect::>(); + if traits.is_empty() { + continue; + } + let callee_summary = dispatch_summary(reader, &callee.symbol, budget, control)?; + let callee_name = callee_summary + .metadata + .as_ref() + .ok_or(DispatchExpansionStop::Unavailable)? + .simple_name + .clone(); + let Some(implementors) = dispatch_step( reader, - &callee.symbol, - scope, + &traits, + &[RelationEdgeKindV1::Implements], + true, budget, control, - |target| { - let symbol = CodeGraphSymbolRefV1::for_occurrence(&target.occurrence) - .map_err(|_| DispatchExpansionStop::Unavailable)?; - if seen.insert(symbol.clone()) { - dispatch.push(RelationKeyV1 { - symbol, - edge_kind: RelationEdgeKindV1::Calls, - dispatch_from: Some(callee.symbol.clone()), - depth: callee.depth, - }); - } - Ok(found.keys.len() + dispatch.len() < MAX_RELATION_CANDIDATE_KEYS) - }, - )?; - if !exhausted { + )? + else { + found.complete = false; + break; + }; + let implementors = implementors.into_iter().flatten().collect::>(); + if implementors.is_empty() { + continue; + } + let Some(children) = dispatch_step( + reader, + &implementors, + &[RelationEdgeKindV1::Contains], + false, + budget, + control, + )? + else { found.complete = false; break; + }; + for child in children.into_iter().flatten() { + if seen.contains(&child) { + continue; + } + let summary = dispatch_summary(reader, &child, budget, control)?; + let Some(metadata) = summary.metadata.as_ref() else { + return Err(DispatchExpansionStop::Unavailable); + }; + if !matches!( + NodeKind::from_str(&metadata.kind), + Some(NodeKind::Function | NodeKind::Method) + ) || metadata.simple_name != callee_name + || !summary + .binding + .as_ref() + .and_then(|binding| binding.logical_path.as_deref()) + .is_some_and(|path| path_is_in_code_query_scope(path, scope)) + { + continue; + } + seen.insert(child.clone()); + dispatch.push(RelationKeyV1 { + symbol: child, + edge_kind: RelationEdgeKindV1::Calls, + dispatch_from: Some(callee.symbol.clone()), + depth: callee.depth, + }); + if found.keys.len() + dispatch.len() >= MAX_RELATION_CANDIDATE_KEYS { + found.complete = false; + break 'callees; + } } } + check_dispatch_control(control.as_ref(), budget)?; dispatch.sort_by(|left, right| { left.depth .cmp(&right.depth) diff --git a/crates/tracedecay-contracts/src/retrieval/callable_code.rs b/crates/tracedecay-contracts/src/retrieval/callable_code.rs index f094a161fc..644e4cba82 100644 --- a/crates/tracedecay-contracts/src/retrieval/callable_code.rs +++ b/crates/tracedecay-contracts/src/retrieval/callable_code.rs @@ -75,6 +75,56 @@ pub struct CodeQueryPage { pub query_fallback: Option, } +/// The project file a callable-code row was read from, when it names one. +pub trait CodeQueryRow { + fn source_path(&self) -> Option<&str>; +} + +impl CodeQueryPage { + /// Project-relative files of the rows this page returned, once each. + pub fn touched_files(&self) -> Vec { + let mut files = self + .items + .iter() + .filter_map(CodeQueryRow::source_path) + .map(str::to_owned) + .collect::>(); + files.sort(); + files.dedup(); + files + } +} + +impl CodeQueryRow for ExactOccurrenceRecord { + fn source_path(&self) -> Option<&str> { + Some(&self.occurrence.path) + } +} + +impl CodeQueryRow for LexicalOccurrenceRecord { + fn source_path(&self) -> Option<&str> { + Some(&self.occurrence.path) + } +} + +impl CodeQueryRow for SourceMetadataRecord { + fn source_path(&self) -> Option<&str> { + Some(&self.path) + } +} + +impl CodeQueryRow for CodeFacetRecord { + fn source_path(&self) -> Option<&str> { + None + } +} + +impl CodeQueryRow for CodeTimelineRecord { + fn source_path(&self) -> Option<&str> { + None + } +} + impl CodeQueryPage { pub fn new( generation: CodeGenerationId, diff --git a/crates/tracedecay-contracts/src/retrieval/callable_code_service.rs b/crates/tracedecay-contracts/src/retrieval/callable_code_service.rs index a552885860..fa85cc1f7e 100644 --- a/crates/tracedecay-contracts/src/retrieval/callable_code_service.rs +++ b/crates/tracedecay-contracts/src/retrieval/callable_code_service.rs @@ -236,7 +236,7 @@ macro_rules! callable_code_service_method { ) { return problem_envelope(context, operation, problem); } - evidence_envelope_with_async_publication_recheck( + let mut result = evidence_envelope_with_async_publication_recheck( context, operation, &admission, @@ -251,7 +251,13 @@ macro_rules! callable_code_service_method { ) }, ) - .await + .await?; + if let Ok(envelope) = &mut result + && let Some(page) = envelope.outcome.payload() + { + envelope.touched_files = page.touched_files(); + } + Ok(result) } }; } diff --git a/crates/tracedecay-contracts/src/retrieval/mod.rs b/crates/tracedecay-contracts/src/retrieval/mod.rs index 552c908333..9781d063c4 100644 --- a/crates/tracedecay-contracts/src/retrieval/mod.rs +++ b/crates/tracedecay-contracts/src/retrieval/mod.rs @@ -86,12 +86,12 @@ pub use callable_code::{ CALLABLE_CODE_OPERATION_COUNT, CallableCodeOperationKind, CallableCodeOperations, CodeFacetDimension, CodeFacetRecord, CodeFacetRequest, CodeHierarchyRequest, CodeImpactRequest, CodeImplementationsRequest, CodeLexicalField, CodeLexicalFieldFilter, CodeNavigationRequest, - CodeOccurrenceRecord, CodeQueryPage, CodeQueryScope, CodeRelationRequest, CodeSignatureRequest, - CodeSymbolSearchRequest, CodeTimelineRecord, CodeTimelineRequest, ExactOccurrenceRecord, - ExactOccurrenceRequest, LexicalOccurrenceRecord, MAX_CALLABLE_CODE_DEPTH, - MAX_CALLABLE_CODE_FILTERS, MAX_CALLABLE_CODE_FUZZY_EXPANSIONS, MAX_CALLABLE_CODE_QUERY_BYTES, - MAX_SOURCE_METADATA_FILES, ModuleApiRequest, PhraseSearchRequest, QualifiedNameRequest, - SourceMetadataRecord, SourceMetadataRequest, + CodeOccurrenceRecord, CodeQueryPage, CodeQueryRow, CodeQueryScope, CodeRelationRequest, + CodeSignatureRequest, CodeSymbolSearchRequest, CodeTimelineRecord, CodeTimelineRequest, + ExactOccurrenceRecord, ExactOccurrenceRequest, LexicalOccurrenceRecord, + MAX_CALLABLE_CODE_DEPTH, MAX_CALLABLE_CODE_FILTERS, MAX_CALLABLE_CODE_FUZZY_EXPANSIONS, + MAX_CALLABLE_CODE_QUERY_BYTES, MAX_SOURCE_METADATA_FILES, ModuleApiRequest, + PhraseSearchRequest, QualifiedNameRequest, SourceMetadataRecord, SourceMetadataRequest, }; pub use callable_code_catalog::{ callable_code_catalog_contribution, callable_code_handler_descriptors, callable_code_operation, diff --git a/crates/tracedecay-contracts/src/retrieval/symbol_graph.rs b/crates/tracedecay-contracts/src/retrieval/symbol_graph.rs index 2f01a68e89..c0ec990e9d 100644 --- a/crates/tracedecay-contracts/src/retrieval/symbol_graph.rs +++ b/crates/tracedecay-contracts/src/retrieval/symbol_graph.rs @@ -10,7 +10,7 @@ use crate::error::ApplicationContractError; use crate::handlers::ApplicationOperation; use crate::result::{OpaqueCursor, OperationBudgetUsage}; -use super::RetrievalRequestMeta; +use super::{CodeQueryRow, RetrievalRequestMeta}; pub const MAX_SYMBOL_GRAPH_DEPTH: u32 = 10; pub const MAX_SYMBOL_GRAPH_QUERY_BYTES: usize = 4_096; @@ -497,11 +497,14 @@ pub enum SymbolGraphPortOutcome { page: SymbolGraphPage, finished_at: UtcMicros, budget: OperationBudgetUsage, + /// What the read cost its graph lease. + cost: Option, }, Partial { page: SymbolGraphPage, finished_at: UtcMicros, budget: OperationBudgetUsage, + cost: Option, }, Failed { failure: PrimitiveFailure, @@ -571,3 +574,21 @@ pub trait SymbolGraphPrimitivePort { request: &'a GraphImpactPrimitiveRequest, ) -> SymbolGraphPortFuture<'a, SymbolPrimitiveRecord>; } + +impl CodeQueryRow for SymbolPrimitiveRecord { + fn source_path(&self) -> Option<&str> { + Some(&self.symbol().file) + } +} + +impl CodeQueryRow for SymbolRelationRecord { + fn source_path(&self) -> Option<&str> { + Some(&self.symbol().file) + } +} + +impl CodeQueryRow for TypeHierarchyRecord { + fn source_path(&self) -> Option<&str> { + Some(&self.symbol().file) + } +} diff --git a/crates/tracedecay/tests/mcp_suite/mcp_handler_test/graph_query_test/relation_page_cost.rs b/crates/tracedecay/tests/mcp_suite/mcp_handler_test/graph_query_test/relation_page_cost.rs index 0ec821a5e3..51dffd497a 100644 --- a/crates/tracedecay/tests/mcp_suite/mcp_handler_test/graph_query_test/relation_page_cost.rs +++ b/crates/tracedecay/tests/mcp_suite/mcp_handler_test/graph_query_test/relation_page_cost.rs @@ -22,13 +22,22 @@ fn hub_source() -> String { format!("pub fn hub() {{\n{calls}}}\n\n{leaves}") } -async fn callees_page(fixture: &GraphQueryFixture, node_id: &str, cursor: Option<&Value>) -> Value { - let mut arguments = json!({ - "node_id": node_id, - "maximum_depth": 1, - "resolve_trait_dispatch": false, - "format": "json", - }); +async fn callees_page( + fixture: &GraphQueryFixture, + node_id: &str, + defaults: bool, + cursor: Option<&Value>, +) -> Value { + let mut arguments = if defaults { + json!({"node_id": node_id, "format": "json"}) + } else { + json!({ + "node_id": node_id, + "maximum_depth": 1, + "resolve_trait_dispatch": false, + "format": "json", + }) + }; if let Some(cursor) = cursor { arguments["meta"] = json!({ "projection": "evidence", @@ -77,7 +86,7 @@ async fn a_callees_page_reads_its_own_rows_not_every_relation() { let target: Value = serde_json::from_str(extract_text(&target.value)).unwrap(); let node_id = target[0]["node_id"].as_str().expect("node id").to_owned(); - let first = callees_page(&fixture, &node_id, None).await; + let first = callees_page(&fixture, &node_id, false, None).await; let payload = &first["outcome"]["value"]["payload"]; assert_eq!( ( @@ -92,7 +101,7 @@ async fn a_callees_page_reads_its_own_rows_not_every_relation() { // each call's target) that decode no entity. assert_eq!(page_cost(&first), (11, 2, 210), "{first:#}"); - let second = callees_page(&fixture, &node_id, Some(&payload["next_cursor"])).await; + let second = callees_page(&fixture, &node_id, false, Some(&payload["next_cursor"])).await; assert_eq!( second["outcome"]["value"]["payload"]["items"] .as_array() @@ -103,3 +112,45 @@ async fn a_callees_page_reads_its_own_rows_not_every_relation() { assert_eq!(page_cost(&second), (11, 2, 210), "{second:#}"); shutdown_graph_fixture(fixture).await; } + +/// With the MCP defaults (depth 3, trait dispatch on) a page still reads in +/// proportion to its own rows: the trait-dispatch check resolves containers +/// from relation keys rather than reading every callee. +#[tokio::test] +async fn a_default_callees_page_resolves_dispatch_without_reading_every_callee() { + let fixture = graph_query_fixture_with_sources(|project| { + fs::create_dir_all(project.join("src")).unwrap(); + fs::write(project.join("src/lib.rs"), hub_source()).unwrap(); + }) + .await; + let target = call_production_tool( + &fixture, + "tracedecay_by_qualified_name", + json!({"qualified_name": "src/lib.rs::hub", "format": "json"}), + None, + None, + ) + .await + .expect("exact function"); + let target: Value = serde_json::from_str(extract_text(&target.value)).unwrap(); + let node_id = target[0]["node_id"].as_str().expect("node id").to_owned(); + + let first = callees_page(&fixture, &node_id, true, None).await; + let payload = &first["outcome"]["value"]["payload"]; + assert_eq!( + ( + payload["total"].as_u64(), + payload["items"].as_array().map(Vec::len) + ), + (Some(105), Some(10)), + "{first:#}" + ); + // Eleven reads, as with dispatch off: the dispatch check finds every + // callee's container in one batched step (the 105 incoming calls, none a + // `Contains`), so no callee is read. Five fan-outs: the two walk hops, the + // empty second level, and the two container hops. + assert_eq!(page_cost(&first), (11, 5, 315), "{first:#}"); + let second = callees_page(&fixture, &node_id, true, Some(&payload["next_cursor"])).await; + assert_eq!(page_cost(&second), (11, 5, 315), "{second:#}"); + shutdown_graph_fixture(fixture).await; +} diff --git a/crates/tracedecay/tests/mcp_suite/mcp_handler_test/graph_query_test/typed_evidence_trailers.rs b/crates/tracedecay/tests/mcp_suite/mcp_handler_test/graph_query_test/typed_evidence_trailers.rs index f51909ea5d..ce866669fb 100644 --- a/crates/tracedecay/tests/mcp_suite/mcp_handler_test/graph_query_test/typed_evidence_trailers.rs +++ b/crates/tracedecay/tests/mcp_suite/mcp_handler_test/graph_query_test/typed_evidence_trailers.rs @@ -254,11 +254,11 @@ async fn typed_callees_carry_their_read_cost_on_the_envelope_and_the_trailer() { "src/walk.rs::Walk::read", "{payload:#}" ); - // Six point reads: the seed `known`; the trait-dispatch check on its one - // callee (the callee's summary, its two incoming edges, and the impl that - // contains it); and the callee's summary for the page. Three fan-outs: + // Three point reads: the seed `known`; the container of its one callee, + // which the dispatch check finds in one batched step and reads to learn + // it is no trait; and the callee's summary for the page. Four fan-outs: // `known`'s call relations and their targets (one row each), then the - // callee's incoming edges (two rows). + // callee's incoming edges (two rows) and their `Contains` source (one). let cost = &payload["cost"]; assert_eq!( ( @@ -267,17 +267,17 @@ async fn typed_callees_carry_their_read_cost_on_the_envelope_and_the_trailer() { &cost["adjacency_rows"], ), ( - &json!({"graph_sealed": 6, "graph_staging": 0}), - &json!(3), + &json!({"graph_sealed": 3, "graph_staging": 0}), &json!(4), + &json!(5), ), "{payload:#}" ); assert_eq!( cost_trailer(&texts), format!( - "\ntracedecay_cost: wall_us={} graph_sealed_reads=6 graph_staging_reads=0 \ - adjacency_queries=3 adjacency_rows=4 bytes_hydrated={}", + "\ntracedecay_cost: wall_us={} graph_sealed_reads=3 graph_staging_reads=0 \ + adjacency_queries=4 adjacency_rows=5 bytes_hydrated={}", cost["wall_micros"], cost["bytes_hydrated"] ), "the trailer renders the envelope's receipt" @@ -290,8 +290,111 @@ async fn typed_callees_carry_their_read_cost_on_the_envelope_and_the_trailer() { ) .await; assert!( - cost_trailer(&markdown).contains(" graph_sealed_reads=6 graph_staging_reads=0 "), + cost_trailer(&markdown).contains(" graph_sealed_reads=3 graph_staging_reads=0 "), "{markdown:?}" ); shutdown_graph_fixture(fixture).await; } + +/// A callers read is metered on the same lease as every other graph read: +/// one incoming fan-out over `Walk::read` (the call from `known` and the impl +/// that contains it), both edges read to learn their kind, and the one +/// caller they reach. +#[tokio::test] +async fn typed_callers_carry_their_read_cost() { + let fixture = trailer_fixture().await; + let target = call_production_tool( + &fixture, + "tracedecay_by_qualified_name", + json!({"qualified_name": "src/walk.rs::Walk::read", "format": "json"}), + None, + None, + ) + .await + .expect("exact method"); + let target: Value = serde_json::from_str(extract_text(&target.value)).unwrap(); + let node_id = target[0]["node_id"].as_str().expect("node id").to_owned(); + + let texts = call( + &fixture, + "tracedecay_callers", + json!({"node_id": node_id, "maximum_depth": 1, "format": "json"}), + ) + .await; + let payload: Value = serde_json::from_str(&texts[0]).unwrap(); + let cost = &payload["cost"]; + assert_eq!( + ( + &cost["point_reads"], + &cost["adjacency_queries"], + &cost["adjacency_rows"], + ), + ( + &json!({"graph_sealed": 3, "graph_staging": 0}), + &json!(1), + &json!(2), + ), + "{payload:#}" + ); + assert_eq!( + cost_trailer(&texts), + format!( + "\ntracedecay_cost: wall_us={} graph_sealed_reads=3 graph_staging_reads=0 \ + adjacency_queries=1 adjacency_rows=2 bytes_hydrated={}", + cost["wall_micros"], cost["bytes_hydrated"] + ) + ); + shutdown_graph_fixture(fixture).await; +} + +/// Callable-code reads report the files of the rows they returned, so their +/// responses end with the token-accounting footer like every other code read. +#[tokio::test] +async fn callable_code_reads_carry_their_files_and_the_accounting_footer() { + let fixture = trailer_fixture().await; + let search = call( + &fixture, + "tracedecay_code_symbol_search", + json!({ + "query": "Walk", + "scope": {"path_prefix": null}, + "lazy_index_ignored_dependencies": false, + "meta": {"projection": "evidence", "order": "source_position"}, + "format": "json", + }), + ) + .await; + let body = body_before_footer(&search, 672); + let payload: Value = serde_json::from_str(&body[0]).unwrap(); + assert_eq!( + payload["touched_files"], + json!(["src/walk.rs", "src/lib.rs"]), + "{payload:#}" + ); + + let target = call_production_tool( + &fixture, + "tracedecay_by_qualified_name", + json!({"qualified_name": "src/lib.rs::known", "format": "json"}), + None, + None, + ) + .await + .expect("exact function"); + let target: Value = serde_json::from_str(extract_text(&target.value)).unwrap(); + let node_id = target[0]["node_id"].as_str().expect("node id").to_owned(); + let callees = call( + &fixture, + "tracedecay_callees", + json!({"node_id": node_id, "maximum_depth": 1, "format": "json"}), + ) + .await; + let body = body_before_footer(&callees, 35); + let payload: Value = serde_json::from_str(&body[0]).unwrap(); + assert_eq!( + payload["touched_files"], + json!(["src/walk.rs"]), + "{payload:#}" + ); + shutdown_graph_fixture(fixture).await; +}