Skip to content

Add loop_cap as a new API parameter to help to lock loop expansion #49

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
hungpham10 merged 2 commits into from
Oct 24, 2026
Merged

Add loop_cap as a new API parameter to help to lock loop expansion #49

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
5cbb307
Add loop_cap as a new API parameter to help to lock loop expansion
hungpham10 Oct 24, 2026
d152e5a
style: apply rustfmt
hungpham10 Oct 24, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion crates/codegraph-api/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -466,10 +466,11 @@ impl GraphApi {
target_id: u64,
max_depth: usize,
max_paths: usize,
max_loop_iter: usize,
) -> Result<Vec<FullFlowPath>> {
self.index()
.await
.trace_full_flow(target_id, max_depth, max_paths)
.trace_full_flow(target_id, max_depth, max_paths, max_loop_iter)
.await
}

Expand Down
20 changes: 14 additions & 6 deletions crates/codegraph-graph/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2269,10 +2269,11 @@ impl GraphIndex {
target_id: u64,
max_depth: usize,
max_paths: usize,
max_loop_iter: usize,
) -> Result<Vec<FullFlowPath>> {
// Loop body chỉ trace 1 lần rồi exit (max_iter): unroll thêm chỉ nhân
// bản y hệt body — nhiễu thuần, không thêm thông tin flow nào.
let max_loop_iter = 1;
// max_loop_iter: số lần tối đa trace loop body (mặc định 1 — body chạy
// 1 lần rồi exit qua max_iter; unroll thêm chỉ nhân bản y hệt body).
let max_loop_iter = max_loop_iter.max(1);

// ── Kiểm tra target tồn tại ──
if !self.symbols.contains_key(&target_id) {
Expand Down Expand Up @@ -3022,9 +3023,16 @@ impl GraphIndex {
}
MARKER_CONTINUE => {
let f = st.frames.last_mut().unwrap();
if let Some(&(lp, _)) = f.loops.last() {
*st.loop_iters.entry((f.func_id, lp)).or_insert(0) += 1;
f.pos = lp + 1;
if let Some(&(lp, back)) = f.loops.last() {
let cnt = st.loop_iters.entry((f.func_id, lp)).or_insert(0);
if *cnt >= max_loop_iter {
// Vượt cap: thoát loop luôn (chống continue-chain).
f.loops.pop();
f.pos = back + 1;
} else {
*cnt += 1;
f.pos = lp + 1;
}
} else {
f.pos += 1;
}
Expand Down
68 changes: 54 additions & 14 deletions crates/codegraph-graph/tests/trace_flow_tests.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,10 +4,10 @@

use codegraph_core::{
CallRecord, EffectType, MARKER_BRANCH_END, MARKER_BREAK, MARKER_CALL_ENTER, MARKER_CALL_EXIT,
MARKER_CATCH_START, MARKER_IF_FALSE, MARKER_IF_TRUE, MARKER_LOOP, MARKER_LOOP_BACK,
MARKER_RETURN, MARKER_SWITCH_CASE, MARKER_SWITCH_CLOSE, MARKER_SWITCH_END, MARKER_SWITCH_START,
MARKER_THROW, MARKER_TRY_END, MARKER_TRY_START, PathEnding, SYMBOL_BASE, ScopeLevel, Symbol,
SymbolKind,
MARKER_CATCH_START, MARKER_CONTINUE, MARKER_IF_FALSE, MARKER_IF_TRUE, MARKER_LOOP,
MARKER_LOOP_BACK, MARKER_RETURN, MARKER_SWITCH_CASE, MARKER_SWITCH_CLOSE, MARKER_SWITCH_END,
MARKER_SWITCH_START, MARKER_THROW, MARKER_TRY_END, MARKER_TRY_START, PathEnding, SYMBOL_BASE,
ScopeLevel, Symbol, SymbolKind,
};
use codegraph_graph::{GraphIndex, ParseResult};
use std::collections::HashMap;
Expand Down Expand Up @@ -123,7 +123,7 @@ async fn if_else_paths_continue_after_block() {
(C, vec![C, MARKER_RETURN]),
]);
let idx = build_idx(chains, vec![]).await;
let paths = idx.trace_full_flow(A, 5, 50).await.unwrap();
let paths = idx.trace_full_flow(A, 5, 50, 1).await.unwrap();
assert!(paths.len() >= 2, "cần ≥2 path (true/false), got {paths:?}");
eprintln!("DBG-TEST {paths:?}");
for p in &paths {
Expand Down Expand Up @@ -165,7 +165,7 @@ async fn switch_paths_per_case() {
(C, vec![C, MARKER_RETURN]),
]);
let idx = build_idx(chains, vec![]).await;
let paths = idx.trace_full_flow(A, 5, 50).await.unwrap();
let paths = idx.trace_full_flow(A, 5, 50, 1).await.unwrap();
assert!(paths.len() >= 2, "mỗi case một path, got {paths:?}");
assert!(paths.iter().any(|p| names(p).contains(&"b".to_string())));
assert!(paths.iter().any(|p| names(p).contains(&"c".to_string())));
Expand Down Expand Up @@ -194,7 +194,7 @@ async fn loop_break_exits_then_continues() {
(B, vec![B, MARKER_RETURN]),
]);
let idx = build_idx(chains, vec![]).await;
let paths = idx.trace_full_flow(A, 5, 50).await.unwrap();
let paths = idx.trace_full_flow(A, 5, 50, 1).await.unwrap();
assert!(
paths.iter().any(|p| names(p).contains(&"b".to_string())
&& names(p).last() == Some(&"a".to_string())
Expand All @@ -214,7 +214,7 @@ async fn loop_max_iter_exits_and_continues() {
(B, vec![B, MARKER_RETURN]),
]);
let idx = build_idx(chains, vec![]).await;
let paths = idx.trace_full_flow(A, 5, 50).await.unwrap();
let paths = idx.trace_full_flow(A, 5, 50, 1).await.unwrap();
assert!(
paths
.iter()
Expand All @@ -239,7 +239,7 @@ async fn call_descends_and_returns_with_stack_depths() {
(C, vec![C, MARKER_RETURN]),
]);
let idx = build_idx(chains, vec![]).await;
let paths = idx.trace_full_flow(B, 5, 50).await.unwrap();
let paths = idx.trace_full_flow(B, 5, 50, 1).await.unwrap();
assert!(
paths.iter().any(|p| names(p).contains(&"c".to_string())),
"phải descend vào c: {paths:?}"
Expand Down Expand Up @@ -300,7 +300,7 @@ async fn throw_unwinds_to_matching_catch() {
label(A, 6, "catch ValueError"),
];
let idx = build_idx(chains, calls).await;
let paths = idx.trace_full_flow(B, 5, 50).await.unwrap();
let paths = idx.trace_full_flow(B, 5, 50, 1).await.unwrap();
assert!(
paths.iter().any(|p| {
p.ended_with == PathEnding::Returned
Expand Down Expand Up @@ -345,7 +345,7 @@ async fn throw_wrong_type_unwinds_through() {
label(A, 6, "catch ValueError"),
];
let idx = build_idx(chains, calls).await;
let paths = idx.trace_full_flow(B, 5, 50).await.unwrap();
let paths = idx.trace_full_flow(B, 5, 50, 1).await.unwrap();
assert!(
paths.iter().any(|p| p.ended_with == PathEnding::Threw),
"type không khớp phải unwind hết: {paths:?}"
Expand Down Expand Up @@ -392,7 +392,7 @@ async fn throw_crosses_two_frames_to_catch() {
label(A, 6, "catch ValueError"),
];
let idx = build_idx(chains, calls).await;
let paths = idx.trace_full_flow(C, 5, 50).await.unwrap();
let paths = idx.trace_full_flow(C, 5, 50, 1).await.unwrap();
assert!(
paths.iter().any(|p| {
p.ended_with == PathEnding::Returned && names(p).contains(&"handler".to_string())
Expand All @@ -414,7 +414,7 @@ async fn no_double_return_after_leaf_call() {
]);
let calls = vec![call(A, "c", 2)];
let idx = build_idx(chains, calls).await;
let paths = idx.trace_full_flow(A, 5, 50).await.unwrap();
let paths = idx.trace_full_flow(A, 5, 50, 1).await.unwrap();
assert!(!paths.is_empty());
for p in &paths {
let a_returns = p
Expand All @@ -437,9 +437,49 @@ async fn max_depth_one_still_returns_direct_caller_path() {
(B, vec![B, MARKER_RETURN]),
]);
let idx = build_idx(chains, vec![]).await;
let paths = idx.trace_full_flow(B, 1, 50).await.unwrap();
let paths = idx.trace_full_flow(B, 1, 50, 1).await.unwrap();
assert!(
paths.iter().any(|p| names(p).contains(&"a".to_string())),
"max_depth=1 phải trả path qua direct caller: {paths:?}"
);
}

#[tokio::test]
async fn continue_heavy_loop_bounded_by_loop_cap() {
// while { if (skip) { i++; continue; } i++; } — continue không được phép
// mở rộng loop vô hạn: cap (mặc định 1) phải kết thúc loop.
let chains = HashMap::from([(
A,
vec![
A,
MARKER_LOOP,
MARKER_IF_TRUE,
MARKER_CONTINUE,
MARKER_BRANCH_END,
A,
MARKER_LOOP_BACK,
A,
MARKER_RETURN,
],
)]);
let idx = build_idx(chains, vec![]).await;
let paths = idx.trace_full_flow(A, 5, 50, 1).await.unwrap();
assert!(
paths
.iter()
.any(|p| p.ended_with == PathEnding::Returned
&& names(p).last() == Some(&"a".to_string())),
"continue phải thoát loop theo cap rồi return: {paths:?}"
);
for p in &paths {
let loop_steps = p
.steps
.iter()
.filter(|s| s.action == codegraph_core::TraceAction::Loop)
.count();
assert!(
loop_steps <= 4,
"loop không được mở rộng quá cap: {loop_steps}"
);
}
}
4 changes: 3 additions & 1 deletion crates/codegraph-graphql/src/query.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -336,12 +336,14 @@ impl Query {
target_id: ID,
max_depth: Option<i32>,
max_paths: Option<i32>,
max_loop_iter: Option<i32>,
) -> GqlResult<Vec<FullFlowPath>> {
let target_id: u64 = parse_id(&target_id)?;
let max_depth = max_depth.unwrap_or(5).max(1) as usize;
let max_paths = max_paths.unwrap_or(50).max(1) as usize;
let max_loop_iter = max_loop_iter.unwrap_or(1).max(1) as usize;
let api = api_for(ctx).await?;
api.trace_full_flow(target_id, max_depth, max_paths)
api.trace_full_flow(target_id, max_depth, max_paths, max_loop_iter)
.await
.map_err(|e| async_graphql::Error::new(e.to_string()))
}
Expand Down
12 changes: 10 additions & 2 deletions crates/codegraph-mcp/src/tools.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -126,7 +126,8 @@ fn tool_defs() -> Vec<ToolDef> {
"Trace full execution flow to target function with stacktrace from callers. Returns all possible paths from callers through to target with call/return/throw/branch information.",
json!({ "type": "object", "properties": {
"target_id": { "type": "integer", "description": "Symbol ID of the target function" },
"max_depth": { "type": "integer", "default": 5, "description": "Maximum depth to search up callers" },
"max_depth": { "type": "integer", "default": 5, "description": "Số caller hops tối đa khi dò ngược caller→target (path dài hops+1). KHÔNG giới hạn loop nội bộ — dùng max_loop_iter cho loop." },
"max_loop_iter": { "type": "integer", "default": 1, "description": "Số lần tối đa trace loop body (mặc định 1: body chạy 1 lần rồi exit — chống phình output; tăng nếu cần trace nhiều iteration)." },
"max_paths": { "type": "integer", "default": 50, "description": "Maximum number of paths to return" }
}, "required": ["target_id"] }),
),
Expand Down Expand Up @@ -622,7 +623,14 @@ pub async fn dispatch_with_api(
.unwrap_or(5)
.max(1) as usize; // caller hops — path caller→target dài hops+1
let max_paths = args.get("max_paths").and_then(|v| v.as_u64()).unwrap_or(50) as usize;
let paths = api.trace_full_flow(target_id, max_depth, max_paths).await?;
let max_loop_iter = args
.get("max_loop_iter")
.and_then(|v| v.as_u64())
.unwrap_or(1)
.max(1) as usize;
let paths = api
.trace_full_flow(target_id, max_depth, max_paths, max_loop_iter)
.await?;
emit_value(
root.as_str(),
serde_json::to_value(&paths).unwrap_or(Value::Null),
Expand Down
Loading