Skip to content

Commit 4f6949f

Browse files
jruokolajruokola
committed
feat: enhance project support in graph functions and cache key generation
1 parent 8c1d916 commit 4f6949f

File tree

5 files changed

+181
-64
lines changed

5 files changed

+181
-64
lines changed

‎Cargo.lock‎

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎crates/codegraph-graph/Cargo.toml‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,10 +20,11 @@ sha2 = { workspace = true }
2020
# SurrealDB support (optional, default on via feature flag)
2121
surrealdb = { version = "2.2", optional = true }
2222
dotenvy = { version = "0.15", optional = true }
23+
tokio = { workspace = true, optional = true, features = ["macros", "rt-multi-thread"] }
2324

2425
[features]
2526
default = []
26-
surrealdb = ["dep:surrealdb", "dep:dotenvy"]
27+
surrealdb = ["dep:surrealdb", "dep:dotenvy", "dep:tokio"]
2728

2829
[[bin]]
2930
name = "surreal_smoke_test"

‎crates/codegraph-graph/src/graph_functions.rs‎

Lines changed: 62 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -12,11 +12,38 @@ use tracing::{debug, error};
1212
#[derive(Clone)]
1313
pub struct GraphFunctions {
1414
db: Arc<Surreal<Any>>,
15+
project_id: String,
1516
}
1617

1718
impl GraphFunctions {
1819
pub fn new(db: Arc<Surreal<Any>>) -> Self {
19-
Self { db }
20+
Self {
21+
db,
22+
project_id: Self::default_project_id(),
23+
}
24+
}
25+
26+
pub fn new_with_project_id(db: Arc<Surreal<Any>>, project_id: impl Into<String>) -> Self {
27+
Self {
28+
db,
29+
project_id: project_id.into(),
30+
}
31+
}
32+
33+
pub fn project_id(&self) -> &str {
34+
&self.project_id
35+
}
36+
37+
fn default_project_id() -> String {
38+
std::env::var("CODEGRAPH_PROJECT_ID")
39+
.ok()
40+
.filter(|v| !v.trim().is_empty())
41+
.or_else(|| {
42+
std::env::current_dir()
43+
.ok()
44+
.map(|p| p.display().to_string())
45+
})
46+
.unwrap_or_else(|| "default-project".to_string())
2047
}
2148

2249
/// Get transitive dependencies of a node up to specified depth
@@ -35,13 +62,16 @@ impl GraphFunctions {
3562
depth: i32,
3663
) -> Result<Vec<DependencyNode>> {
3764
debug!(
38-
"Calling fn::get_transitive_dependencies({}, {}, {})",
39-
node_id, edge_type, depth
65+
"Calling fn::get_transitive_dependencies({}, {}, {}, project={})",
66+
node_id, edge_type, depth, self.project_id
4067
);
4168

4269
let result: Vec<DependencyNode> = self
4370
.db
44-
.query("RETURN fn::get_transitive_dependencies($node_id, $edge_type, $depth)")
71+
.query(
72+
"RETURN fn::get_transitive_dependencies($project_id, $node_id, $edge_type, $depth)",
73+
)
74+
.bind(("project_id", self.project_id.clone()))
4575
.bind(("node_id", node_id.to_string()))
4676
.bind(("edge_type", edge_type.to_string()))
4777
.bind(("depth", depth))
@@ -70,11 +100,15 @@ impl GraphFunctions {
70100
&self,
71101
edge_type: &str,
72102
) -> Result<Vec<CircularDependency>> {
73-
debug!("Calling fn::detect_circular_dependencies({})", edge_type);
103+
debug!(
104+
"Calling fn::detect_circular_dependencies({}, project={})",
105+
edge_type, self.project_id
106+
);
74107

75108
let result: Vec<CircularDependency> = self
76109
.db
77-
.query("RETURN fn::detect_circular_dependencies($edge_type)")
110+
.query("RETURN fn::detect_circular_dependencies($project_id, $edge_type)")
111+
.bind(("project_id", self.project_id.clone()))
78112
.bind(("edge_type", edge_type.to_string()))
79113
.await
80114
.map_err(|e| {
@@ -103,11 +137,15 @@ impl GraphFunctions {
103137
from_node: &str,
104138
max_depth: i32,
105139
) -> Result<Vec<CallChainNode>> {
106-
debug!("Calling fn::trace_call_chain({}, {})", from_node, max_depth);
140+
debug!(
141+
"Calling fn::trace_call_chain({}, {}, project={})",
142+
from_node, max_depth, self.project_id
143+
);
107144

108145
let result: Vec<CallChainNode> = self
109146
.db
110-
.query("RETURN fn::trace_call_chain($from_node, $max_depth)")
147+
.query("RETURN fn::trace_call_chain($project_id, $from_node, $max_depth)")
148+
.bind(("project_id", self.project_id.clone()))
111149
.bind(("from_node", from_node.to_string()))
112150
.bind(("max_depth", max_depth))
113151
.await
@@ -132,11 +170,15 @@ impl GraphFunctions {
132170
/// # Returns
133171
/// Coupling metrics including afferent, efferent, and instability
134172
pub async fn calculate_coupling_metrics(&self, node_id: &str) -> Result<CouplingMetricsResult> {
135-
debug!("Calling fn::calculate_coupling_metrics({})", node_id);
173+
debug!(
174+
"Calling fn::calculate_coupling_metrics({}, project={})",
175+
node_id, self.project_id
176+
);
136177

137178
let results: Vec<CouplingMetricsResult> = self
138179
.db
139-
.query("RETURN fn::calculate_coupling_metrics($node_id)")
180+
.query("RETURN fn::calculate_coupling_metrics($project_id, $node_id)")
181+
.bind(("project_id", self.project_id.clone()))
140182
.bind(("node_id", node_id.to_string()))
141183
.await
142184
.map_err(|e| {
@@ -163,11 +205,15 @@ impl GraphFunctions {
163205
/// # Returns
164206
/// Vector of highly connected hub nodes sorted by degree (descending)
165207
pub async fn get_hub_nodes(&self, min_degree: i32) -> Result<Vec<HubNode>> {
166-
debug!("Calling fn::get_hub_nodes({})", min_degree);
208+
debug!(
209+
"Calling fn::get_hub_nodes({}, project={})",
210+
min_degree, self.project_id
211+
);
167212

168213
let result: Vec<HubNode> = self
169214
.db
170-
.query("RETURN fn::get_hub_nodes($min_degree)")
215+
.query("RETURN fn::get_hub_nodes($project_id, $min_degree)")
216+
.bind(("project_id", self.project_id.clone()))
171217
.bind(("min_degree", min_degree))
172218
.await
173219
.map_err(|e| {
@@ -199,13 +245,14 @@ impl GraphFunctions {
199245
depth: i32,
200246
) -> Result<Vec<DependencyNode>> {
201247
debug!(
202-
"Calling fn::get_reverse_dependencies({}, {}, {})",
203-
node_id, edge_type, depth
248+
"Calling fn::get_reverse_dependencies({}, {}, {}, project={})",
249+
node_id, edge_type, depth, self.project_id
204250
);
205251

206252
let result: Vec<DependencyNode> = self
207253
.db
208-
.query("RETURN fn::get_reverse_dependencies($node_id, $edge_type, $depth)")
254+
.query("RETURN fn::get_reverse_dependencies($project_id, $node_id, $edge_type, $depth)")
255+
.bind(("project_id", self.project_id.clone()))
209256
.bind(("node_id", node_id.to_string()))
210257
.bind(("edge_type", edge_type.to_string()))
211258
.bind(("depth", depth))

‎crates/codegraph-mcp/src/graph_tool_executor.rs‎

Lines changed: 29 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -101,10 +101,10 @@ impl GraphToolExecutor {
101101
stats.current_size = 0;
102102
}
103103

104-
/// Generate a cache key from tool name and parameters
105-
fn cache_key(tool_name: &str, parameters: &JsonValue) -> String {
106-
// Create deterministic key from function name + serialized params
107-
format!("{}:{}", tool_name, parameters.to_string())
104+
/// Generate a cache key from project, tool name, and parameters
105+
fn cache_key(project_id: &str, tool_name: &str, parameters: &JsonValue) -> String {
106+
// Create deterministic key from project + function name + serialized params
107+
format!("{}:{}:{}", project_id, tool_name, parameters.to_string())
108108
}
109109

110110
/// Execute a tool call from LLM
@@ -122,9 +122,11 @@ impl GraphToolExecutor {
122122
let _schema = GraphToolSchemas::get_by_name(tool_name)
123123
.ok_or_else(|| McpError::Protocol(format!("Unknown tool: {}", tool_name)))?;
124124

125+
let project_id = self.graph_functions.project_id();
126+
125127
// Check cache if enabled
126128
if self.cache_enabled {
127-
let cache_key = Self::cache_key(tool_name, &parameters);
129+
let cache_key = Self::cache_key(project_id, tool_name, &parameters);
128130

129131
// Try cache lookup
130132
{
@@ -177,7 +179,7 @@ impl GraphToolExecutor {
177179

178180
// Cache the result if enabled
179181
if self.cache_enabled {
180-
let cache_key = Self::cache_key(tool_name, &parameters);
182+
let cache_key = Self::cache_key(project_id, tool_name, &parameters);
181183
let mut cache = self.cache.lock();
182184
let was_evicted = cache.len() >= cache.cap().get();
183185
cache.put(cache_key, result.clone());
@@ -430,16 +432,34 @@ mod tests {
430432
"depth": 3
431433
});
432434

433-
let key1 = GraphToolExecutor::cache_key("get_transitive_dependencies", &params1);
434-
let key2 = GraphToolExecutor::cache_key("get_transitive_dependencies", &params2);
435-
let key3 = GraphToolExecutor::cache_key("get_transitive_dependencies", &params3);
435+
let project = "proj-a";
436+
437+
let key1 = GraphToolExecutor::cache_key(project, "get_transitive_dependencies", &params1);
438+
let key2 = GraphToolExecutor::cache_key(project, "get_transitive_dependencies", &params2);
439+
let key3 = GraphToolExecutor::cache_key(project, "get_transitive_dependencies", &params3);
436440

437441
// Same params should generate same key
438442
assert_eq!(key1, key2);
439443
// Different params should generate different key
440444
assert_ne!(key1, key3);
441445
}
442446

447+
#[test]
448+
fn test_cache_key_includes_project_scope() {
449+
let params = json!({
450+
"node_id": "nodes:123"
451+
});
452+
453+
let key_a = GraphToolExecutor::cache_key("proj-a", "get_hub_nodes", &params);
454+
let key_b = GraphToolExecutor::cache_key("proj-b", "get_hub_nodes", &params);
455+
456+
assert_ne!(key_a, key_b);
457+
assert!(
458+
key_a.starts_with("proj-a:"),
459+
"Project scope should prefix cache key"
460+
);
461+
}
462+
443463
#[test]
444464
fn test_cache_stats_initialization() {
445465
let stats = CacheStats {

0 commit comments

Comments
 (0)