diff --git a/crates/codegraph-api/src/lib.rs b/crates/codegraph-api/src/lib.rs index ad0b8d37f4..d45192fad9 100644 --- a/crates/codegraph-api/src/lib.rs +++ b/crates/codegraph-api/src/lib.rs @@ -466,10 +466,11 @@ impl GraphApi { target_id: u64, max_depth: usize, max_paths: usize, + max_loop_iter: usize, ) -> Result> { 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 } diff --git a/crates/codegraph-graph/src/lib.rs b/crates/codegraph-graph/src/lib.rs index 9da11d89c6..ab62b09e2f 100644 --- a/crates/codegraph-graph/src/lib.rs +++ b/crates/codegraph-graph/src/lib.rs @@ -2269,10 +2269,11 @@ impl GraphIndex { target_id: u64, max_depth: usize, max_paths: usize, + max_loop_iter: usize, ) -> Result> { - // 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) { @@ -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; } diff --git a/crates/codegraph-graph/tests/trace_flow_tests.rs b/crates/codegraph-graph/tests/trace_flow_tests.rs index 02333c85aa..afe6dbd388 100644 --- a/crates/codegraph-graph/tests/trace_flow_tests.rs +++ b/crates/codegraph-graph/tests/trace_flow_tests.rs @@ -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; @@ -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 { @@ -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()))); @@ -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()) @@ -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() @@ -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:?}" @@ -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 @@ -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:?}" @@ -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()) @@ -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 @@ -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}" + ); + } +} diff --git a/crates/codegraph-graphql/src/query.rs b/crates/codegraph-graphql/src/query.rs index a3aaa0ad89..a1e5584b20 100644 --- a/crates/codegraph-graphql/src/query.rs +++ b/crates/codegraph-graphql/src/query.rs @@ -336,12 +336,14 @@ impl Query { target_id: ID, max_depth: Option, max_paths: Option, + max_loop_iter: Option, ) -> GqlResult> { 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())) } diff --git a/crates/codegraph-mcp/src/tools.rs b/crates/codegraph-mcp/src/tools.rs index 48caf94665..b2ef601c57 100644 --- a/crates/codegraph-mcp/src/tools.rs +++ b/crates/codegraph-mcp/src/tools.rs @@ -126,7 +126,8 @@ fn tool_defs() -> Vec { "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"] }), ), @@ -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),