use anyhow::Result;
use unfault_core::graph::CodeGraph;
use crate::integration::gcp::trace::RemoteCallPattern;
#[derive(Debug, Default)]
pub struct TraceEnrichmentResult {
pub remote_services_added: usize,
pub edges_added: usize,
pub linked_services: Vec<String>,
}
pub fn enrich_graph(
graph: &mut CodeGraph,
patterns: &[RemoteCallPattern],
http_caller_files: &[String],
verbose: bool,
) -> Result<TraceEnrichmentResult> {
let mut result = TraceEnrichmentResult::default();
for pattern in patterns {
if pattern.remote_service_name.is_empty() {
continue;
}
let svc_idx = graph.get_or_create_remote_service(
&pattern.remote_service_name,
&pattern.remote_endpoint,
pattern.observed_count,
pattern.p99_latency_ms,
);
result.remote_services_added += 1;
result
.linked_services
.push(pattern.remote_service_name.clone());
if verbose {
eprintln!(
" trace: remote service '{}' ({} calls{})",
pattern.remote_service_name,
pattern.observed_count,
pattern
.p99_latency_ms
.map(|p| format!(", p99={:.0}ms", p))
.unwrap_or_default()
);
}
let matched_files = match_caller_files(pattern, http_caller_files, graph);
let files_to_link: Vec<&str> = if matched_files.is_empty() {
http_caller_files.iter().map(|s| s.as_str()).collect()
} else {
matched_files.iter().copied().collect()
};
for file_path in files_to_link {
if graph.add_remote_call_edge(file_path, svc_idx) {
result.edges_added += 1;
if verbose {
eprintln!(
" → {} --[RemoteCall]--> {}",
file_path.rsplit('/').next().unwrap_or(file_path),
pattern.remote_service_name
);
}
}
}
}
Ok(result)
}
fn match_caller_files<'a>(
pattern: &RemoteCallPattern,
http_caller_files: &'a [String],
graph: &CodeGraph,
) -> Vec<&'a str> {
use unfault_core::graph::GraphNode;
if pattern.local_callers.is_empty() {
return vec![];
}
let mut matched: Vec<&str> = Vec::new();
for caller_name in &pattern.local_callers {
for idx in graph.graph.node_indices() {
let node = &graph.graph[idx];
if let GraphNode::Function {
http_path: Some(path),
name,
..
} = node
{
let caller_lower = caller_name.to_lowercase();
let path_fragment = path.split('/').find(|s| s.len() > 2).unwrap_or(path);
if caller_lower.contains(path_fragment)
|| caller_lower.contains(&name.to_lowercase())
{
if let Some(file_path) = find_file_for_function(idx, graph, http_caller_files) {
if !matched.contains(&file_path) {
matched.push(file_path);
}
}
}
}
}
}
matched
}
fn find_file_for_function<'a>(
func_idx: unfault_core::graph::GraphNodeIndex,
graph: &CodeGraph,
candidates: &'a [String],
) -> Option<&'a str> {
use petgraph::visit::EdgeRef;
use petgraph::Direction;
use unfault_core::graph::{GraphEdgeKind, GraphNode};
for edge in graph.graph.edges_directed(func_idx, Direction::Incoming) {
if !matches!(edge.weight(), GraphEdgeKind::Contains) {
continue;
}
if let GraphNode::File { path, .. } = &graph.graph[edge.source()] {
if let Some(candidate) = candidates.iter().find(|c| *c == path) {
return Some(candidate.as_str());
}
}
}
None
}
#[cfg(test)]
mod tests {
use super::*;
use unfault_core::graph::{CodeGraph, GraphEdgeKind, GraphNode, GraphNodeIndex};
use unfault_core::parse::ast::FileId;
use unfault_core::types::context::Language;
fn make_graph_with_handler() -> (CodeGraph, GraphNodeIndex) {
let mut graph = CodeGraph::new();
let file_idx = graph.graph.add_node(GraphNode::File {
file_id: FileId(1),
path: "src/checkout.py".to_string(),
language: Language::Python,
});
graph.file_nodes.insert(FileId(1), file_idx);
graph
.path_to_file
.insert("src/checkout.py".to_string(), file_idx);
let fn_idx = graph.graph.add_node(GraphNode::Function {
file_id: FileId(1),
name: "process_checkout".to_string(),
qualified_name: "process_checkout".to_string(),
is_async: true,
is_handler: true,
http_method: Some("POST".to_string()),
http_path: Some("/checkout".to_string()),
});
graph
.graph
.add_edge(file_idx, fn_idx, GraphEdgeKind::Contains);
(graph, file_idx)
}
#[test]
fn enrich_creates_remote_service_and_edge() {
let (mut graph, _) = make_graph_with_handler();
let patterns = vec![RemoteCallPattern {
remote_service_name: "inventory-service".to_string(),
remote_endpoint: "https://inventory-svc:8080".to_string(),
observed_count: 42,
p99_latency_ms: Some(12.5),
local_callers: vec![],
}];
let http_callers = vec!["src/checkout.py".to_string()];
let result = enrich_graph(&mut graph, &patterns, &http_callers, false).unwrap();
assert_eq!(result.remote_services_added, 1);
assert_eq!(result.edges_added, 1);
assert_eq!(result.linked_services[0], "inventory-service");
let remote_nodes = graph.get_remote_services();
assert_eq!(remote_nodes.len(), 1);
}
#[test]
fn enrich_no_callers_no_edges() {
let (mut graph, _) = make_graph_with_handler();
let patterns = vec![RemoteCallPattern {
remote_service_name: "payments".to_string(),
remote_endpoint: String::new(),
observed_count: 5,
p99_latency_ms: None,
local_callers: vec![],
}];
let result = enrich_graph(&mut graph, &patterns, &[], false).unwrap();
assert_eq!(result.edges_added, 0);
assert_eq!(result.remote_services_added, 1);
}
#[test]
fn enrich_deduplicates_edges() {
let (mut graph, _) = make_graph_with_handler();
let pattern = RemoteCallPattern {
remote_service_name: "inventory-service".to_string(),
remote_endpoint: String::new(),
observed_count: 1,
p99_latency_ms: None,
local_callers: vec![],
};
let http_callers = vec!["src/checkout.py".to_string()];
enrich_graph(&mut graph, &[pattern.clone()], &http_callers, false).unwrap();
let result2 = enrich_graph(&mut graph, &[pattern], &http_callers, false).unwrap();
assert_eq!(result2.edges_added, 0);
}
}