fix: address review feedback on PR-TRACE-1a

Must-fix issues:
- parent_edge_id now stores edge IDs (format: from_id:to_id:edge_type:hop)
  instead of node IDs, enabling proper tree reconstruction
- Budget counting now only counts newly discovered nodes; moved increment
  after visited check

Suggestions addressed:
- test_trace_direction_required: now validates pydantic ValidationError
- test_trace_inbound_callers_depth_2: removed vacuous >= 0 assertion
- test_trace_parent_edge_id_chain: now verifies parent_edge_id format
  and validates it references an edge that reaches the current node
- Tautological hop check: simplified to assert 0 in hops and hops <= {0, 1}
- Removed hardcoded row_kind from _edge_attrs_for_row
- Documented include_unresolved as 1a no-op (like prune_roles)

Tests: 22 passed, 1 skipped

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

Author: HumanBean17

mcp_trace.py

@@ -108,7 +108,6 @@ def _edge_attrs_for_row(row: dict[str, Any]) -> dict[str, Any]:
         if k not in {"source_id", "other_id", "edge_type", "stored_edge_type"}
         and v not in (None, "")
     }
-    attrs.setdefault("row_kind", "resolved")
     return attrs
 
 
@@ -372,8 +371,9 @@ def trace_v2(
 ) -> TraceOutput:
     """Multi-hop BFS traversal with pruning.
 
-    This is PR-TRACE-1a (core BFS engine). Pruning features are wired but
-    treated as no-ops until PR-TRACE-1b implementation.
+    This is PR-TRACE-1a (core BFS engine). Pruning features (prune_roles, fan_out_cap,
+    collapse_trivial, include_unresolved) are wired but treated as no-ops until
+    PR-TRACE-1b implementation.
     """
     # Validate required parameters.
     if not direction:
@@ -479,8 +479,8 @@ def trace_v2(
     actual_depth = 0
     budget_hit = False
 
-    # Track incoming edge for each node (for parent_edge_id).
-    node_to_incoming_edge: dict[str, TraceEdge] = {}
+    # Track incoming edge ID for each node (for parent_edge_id).
+    node_to_incoming_edge_id: dict[str, str] = {}
 
     # For seed nodes, record them in nodes dict (always include seeds, filter doesn't apply).
     for sid in seed_ids:
@@ -524,24 +524,23 @@ def trace_v2(
         edges_by_parent: dict[str, list[TraceEdge]] = defaultdict(list)
 
         for src_id, src_rows in by_source.items():
-            parent_edge = node_to_incoming_edge.get(src_id)
-            parent_edge_id = parent_edge.to_id if parent_edge else None
+            parent_edge_id = node_to_incoming_edge_id.get(src_id)
 
             for row in src_rows:
                 other_id = str(row.get("other_id") or "")
                 if not other_id:
                     continue
 
-                # Check budget BEFORE counting.
+                if other_id in visited:
+                    continue
+
+                # Check budget BEFORE counting (only counts newly discovered nodes).
                 if total_discovered >= max_nodes_discovered:
                     budget_hit = True
                     break
 
                 total_discovered += 1
 
-                if other_id in visited:
-                    continue
-
                 # Load target node.
                 try:
                     other_kind = _resolve_node_kind(g, other_id)
@@ -561,6 +560,7 @@ def trace_v2(
 
                 # Record edge.
                 edge_type = str(row.get("edge_type") or "")
+                edge_id = f"{src_id}:{other_id}:{edge_type}:{hop}"
                 edge = TraceEdge(
                     from_id=src_id,
                     to_id=other_id,
@@ -570,13 +570,12 @@ def trace_v2(
                     attrs=_edge_attrs_for_row(row),
                 )
                 edges.append(edge)
-                edge_id = f"{src_id}:{other_id}:{edge_type}:{hop}"
                 edge_id_map[edge_id] = edge
                 edges_by_parent[src_id].append(edge)
 
-                # Track incoming edge for this node.
-                if other_id not in node_to_incoming_edge:
-                    node_to_incoming_edge[other_id] = edge
+                # Track incoming edge ID for this node (for parent_edge_id of children).
+                if other_id not in node_to_incoming_edge_id:
+                    node_to_incoming_edge_id[other_id] = edge_id
 
                 # PR-TRACE-1b: prune_roles, fan_out_cap, cross_service are no-ops here.
                 # For now, add all discovered nodes to frontier.

tests/test_mcp_trace.py

@@ -65,9 +65,9 @@ def test_trace_outbound_calls_depth_2(kuzu_graph: KuzuGraph) -> None:
     assert out.seed_ids == [seed_id]
     assert out.direction == "out"
     assert out.edge_types == ["CALLS"]
-    # Check that we have edges at hop 0 and hop 1.
+    # Check that we have edges at hop 0 and possibly hop 1.
     hops = {e.hop for e in out.edges}
-    assert hops == {0, 1} or hops == {0} or hops == {0, 1}
+    assert 0 in hops and hops <= {0, 1}
 
 
 def test_trace_inbound_callers_depth_2(kuzu_graph: KuzuGraph) -> None:
@@ -85,7 +85,6 @@ def test_trace_inbound_callers_depth_2(kuzu_graph: KuzuGraph) -> None:
     assert out.success is True
     assert out.seed_ids == [seed_id]
     assert out.direction == "in"
-    assert len(out.edges) >= 0
 
 
 def test_trace_max_paths_cap(kuzu_graph: KuzuGraph) -> None:
@@ -249,10 +248,20 @@ def test_trace_invalid_edge_type(kuzu_graph: KuzuGraph) -> None:
 
 
 def test_trace_direction_required(kuzu_graph: KuzuGraph) -> None:
-    """Missing direction returns success=False."""
-    # direction is required with Field(...) so this test validates the contract.
-    # We can't actually call without direction due to pydantic validation.
-    pass
+    """Missing direction is caught by pydantic validation (literal error)."""
+    from pydantic import ValidationError
+
+    seed_id = _find_method_with_outbound_calls(kuzu_graph)
+    if seed_id is None:
+        pytest.skip("No method with outbound calls in fixture")
+    # Pydantic validation rejects empty string before our code runs.
+    with pytest.raises(ValidationError, match="direction"):
+        trace_v2(
+            ids=seed_id,
+            direction="",
+            edge_types=["CALLS"],
+            graph=kuzu_graph,
+        )
 
 
 def test_trace_edge_types_required(kuzu_graph: KuzuGraph) -> None:
@@ -573,10 +582,20 @@ def test_trace_parent_edge_id_chain(kuzu_graph: KuzuGraph) -> None:
 
     for e in out.edges:
         if e.hop > 0:
-            # parent_edge_id should point to a valid edge.
-            # We check that the parent edge's to_id == e.from_id.
+            # parent_edge_id should be a valid edge identifier that matches an edge in the result.
             if e.parent_edge_id:
-                # Find the parent edge.
-                parent_edges = [p for p in out.edges if p.to_id == e.from_id]
-                # At least one parent should exist.
-                assert len(parent_edges) >= 1
+                # Parse the edge_id format: from_id:to_id:edge_type:hop
+                parts = e.parent_edge_id.split(":")
+                assert len(parts) == 4, f"Invalid parent_edge_id format: {e.parent_edge_id}"
+                parent_from_id, parent_to_id, parent_edge_type, parent_hop = parts
+                # Verify parent edge exists in result and parent.to_id == e.from_id
+                parent_exists = any(
+                    p.from_id == parent_from_id
+                    and p.to_id == parent_to_id
+                    and p.edge_type == parent_edge_type
+                    and p.hop == int(parent_hop)
+                    for p in out.edges
+                )
+                assert parent_exists, f"Parent edge {e.parent_edge_id} not found in result"
+                # Verify the parent edge reaches the current node's from_id
+                assert parent_to_id == e.from_id, f"Parent edge {e.parent_edge_id} to_id != {e.from_id}"