feat: add multi-hop bidirectional context traversal

Implement depth-aware graph traversal for GetContext using a
bidirectional recursive CTE in PostgreSQL. Previously, GetContext
only returned 1-hop neighbors and ignored the depth parameter.

- Add GetBidirectional to EdgeRepository interface
- Implement bidirectional recursive CTE with cycle detection via
  path array, DISTINCT dedup, and LIMIT 1000 safety cap
- Default depth=1 for backward compatibility, cap at 5
- Add 5 unit tests covering multi-hop, cycles, depth capping

Closes #236

Author: ravisuhag

core/entity/edge.go

@@ -34,6 +34,7 @@ type EdgeRepository interface {
 	GetByTarget(ctx context.Context, ns *namespace.Namespace, urn string, filter EdgeFilter) ([]Edge, error)
 	GetDownstream(ctx context.Context, ns *namespace.Namespace, urn string, depth int) ([]Edge, error)
 	GetUpstream(ctx context.Context, ns *namespace.Namespace, urn string, depth int) ([]Edge, error)
+	GetBidirectional(ctx context.Context, ns *namespace.Namespace, urn string, depth int) ([]Edge, error)
 	Delete(ctx context.Context, ns *namespace.Namespace, sourceURN, targetURN, edgeType string) error
 	DeleteByURN(ctx context.Context, ns *namespace.Namespace, urn string) error
 }

core/entity/service.go

@@ -124,6 +124,9 @@ func (s *Service) Suggest(ctx context.Context, ns *namespace.Namespace, text str
 	return nil, nil
 }
 
+// maxContextDepth caps the maximum traversal depth for context queries.
+const maxContextDepth = 5
+
 // GetContext assembles a context subgraph around an entity.
 func (s *Service) GetContext(ctx context.Context, ns *namespace.Namespace, urn string, depth int) (*ContextGraph, error) {
 	ent, err := s.repo.GetByURN(ctx, ns, urn)
@@ -134,23 +137,28 @@ func (s *Service) GetContext(ctx context.Context, ns *namespace.Namespace, urn s
 	cg := &ContextGraph{Entity: ent}
 
 	if s.edges != nil {
-		_ = depth // TODO: use depth for multi-hop traversal
-		outgoing, _ := s.edges.GetBySource(ctx, ns, urn, EdgeFilter{Current: true})
-		incoming, _ := s.edges.GetByTarget(ctx, ns, urn, EdgeFilter{Current: true})
-		cg.Edges = append(outgoing, incoming...)
+		if depth <= 0 {
+			depth = 1
+		}
+		if depth > maxContextDepth {
+			depth = maxContextDepth
+		}
+
+		cg.Edges, err = s.edges.GetBidirectional(ctx, ns, urn, depth)
+		if err != nil {
+			return nil, fmt.Errorf("get context edges: %w", err)
+		}
 
 		seen := map[string]bool{urn: true}
 		for _, e := range cg.Edges {
-			relURN := e.TargetURN
-			if relURN == urn {
-				relURN = e.SourceURN
-			}
-			if seen[relURN] {
-				continue
-			}
-			seen[relURN] = true
-			if rel, err := s.repo.GetByURN(ctx, ns, relURN); err == nil {
-				cg.Related = append(cg.Related, rel)
+			for _, candidate := range []string{e.SourceURN, e.TargetURN} {
+				if seen[candidate] {
+					continue
+				}
+				seen[candidate] = true
+				if rel, err := s.repo.GetByURN(ctx, ns, candidate); err == nil {
+					cg.Related = append(cg.Related, rel)
+				}
 			}
 		}
 	}

core/entity/service_test.go

@@ -166,6 +166,221 @@ func TestService_GetTypes(t *testing.T) {
 	}
 }
 
+// mockEdgeRepo is a simple in-memory edge repository for testing.
+type mockEdgeRepo struct {
+	edges []Edge
+}
+
+func (m *mockEdgeRepo) Upsert(_ context.Context, _ *namespace.Namespace, e *Edge) error {
+	m.edges = append(m.edges, *e)
+	return nil
+}
+
+func (m *mockEdgeRepo) GetBySource(_ context.Context, _ *namespace.Namespace, urn string, _ EdgeFilter) ([]Edge, error) {
+	var result []Edge
+	for _, e := range m.edges {
+		if e.SourceURN == urn {
+			result = append(result, e)
+		}
+	}
+	return result, nil
+}
+
+func (m *mockEdgeRepo) GetByTarget(_ context.Context, _ *namespace.Namespace, urn string, _ EdgeFilter) ([]Edge, error) {
+	var result []Edge
+	for _, e := range m.edges {
+		if e.TargetURN == urn {
+			result = append(result, e)
+		}
+	}
+	return result, nil
+}
+
+func (m *mockEdgeRepo) GetDownstream(_ context.Context, _ *namespace.Namespace, _ string, _ int) ([]Edge, error) {
+	return nil, nil
+}
+
+func (m *mockEdgeRepo) GetUpstream(_ context.Context, _ *namespace.Namespace, _ string, _ int) ([]Edge, error) {
+	return nil, nil
+}
+
+func (m *mockEdgeRepo) GetBidirectional(_ context.Context, _ *namespace.Namespace, urn string, depth int) ([]Edge, error) {
+	// BFS traversal up to depth hops in both directions.
+	type frontier struct {
+		urn   string
+		level int
+	}
+	visited := map[string]bool{urn: true}
+	queue := []frontier{{urn: urn, level: 0}}
+	var result []Edge
+	seen := map[string]bool{} // dedup edges by source+target+type
+
+	for len(queue) > 0 {
+		cur := queue[0]
+		queue = queue[1:]
+		if cur.level >= depth {
+			continue
+		}
+		for _, e := range m.edges {
+			key := e.SourceURN + "|" + e.TargetURN + "|" + e.Type
+			var neighbor string
+			if e.SourceURN == cur.urn {
+				neighbor = e.TargetURN
+			} else if e.TargetURN == cur.urn {
+				neighbor = e.SourceURN
+			} else {
+				continue
+			}
+			if !seen[key] {
+				seen[key] = true
+				result = append(result, e)
+			}
+			if !visited[neighbor] {
+				visited[neighbor] = true
+				queue = append(queue, frontier{urn: neighbor, level: cur.level + 1})
+			}
+		}
+	}
+	return result, nil
+}
+
+func (m *mockEdgeRepo) Delete(_ context.Context, _ *namespace.Namespace, _, _, _ string) error {
+	return nil
+}
+
+func (m *mockEdgeRepo) DeleteByURN(_ context.Context, _ *namespace.Namespace, _ string) error {
+	return nil
+}
+
+func TestService_GetContext_DefaultDepth(t *testing.T) {
+	repo := newMockRepo()
+	edges := &mockEdgeRepo{}
+	svc := NewService(repo, edges, nil)
+	ctx := context.Background()
+	ns := namespace.DefaultNamespace
+
+	// A -> B -> C (linear chain)
+	_, _ = svc.Upsert(ctx, ns, &Entity{URN: "urn:a", Type: TypeTable, Name: "a"})
+	_, _ = svc.Upsert(ctx, ns, &Entity{URN: "urn:b", Type: TypeTable, Name: "b"})
+	_, _ = svc.Upsert(ctx, ns, &Entity{URN: "urn:c", Type: TypeTable, Name: "c"})
+	edges.edges = []Edge{
+		{SourceURN: "urn:a", TargetURN: "urn:b", Type: "lineage"},
+		{SourceURN: "urn:b", TargetURN: "urn:c", Type: "lineage"},
+	}
+
+	// depth=0 should default to 1 (only direct neighbors of B)
+	cg, err := svc.GetContext(ctx, ns, "urn:b", 0)
+	if err != nil {
+		t.Fatalf("GetContext failed: %v", err)
+	}
+	if len(cg.Edges) != 2 {
+		t.Errorf("expected 2 edges at depth 1, got %d", len(cg.Edges))
+	}
+	if len(cg.Related) != 2 {
+		t.Errorf("expected 2 related entities, got %d", len(cg.Related))
+	}
+}
+
+func TestService_GetContext_MultiHop(t *testing.T) {
+	repo := newMockRepo()
+	edges := &mockEdgeRepo{}
+	svc := NewService(repo, edges, nil)
+	ctx := context.Background()
+	ns := namespace.DefaultNamespace
+
+	// A -> B -> C -> D
+	_, _ = svc.Upsert(ctx, ns, &Entity{URN: "urn:a", Type: TypeTable, Name: "a"})
+	_, _ = svc.Upsert(ctx, ns, &Entity{URN: "urn:b", Type: TypeTable, Name: "b"})
+	_, _ = svc.Upsert(ctx, ns, &Entity{URN: "urn:c", Type: TypeTable, Name: "c"})
+	_, _ = svc.Upsert(ctx, ns, &Entity{URN: "urn:d", Type: TypeTable, Name: "d"})
+	edges.edges = []Edge{
+		{SourceURN: "urn:a", TargetURN: "urn:b", Type: "lineage"},
+		{SourceURN: "urn:b", TargetURN: "urn:c", Type: "lineage"},
+		{SourceURN: "urn:c", TargetURN: "urn:d", Type: "lineage"},
+	}
+
+	// depth=2 from B should reach A, C, and D
+	cg, err := svc.GetContext(ctx, ns, "urn:b", 2)
+	if err != nil {
+		t.Fatalf("GetContext failed: %v", err)
+	}
+	if len(cg.Edges) != 3 {
+		t.Errorf("expected 3 edges at depth 2, got %d", len(cg.Edges))
+	}
+	if len(cg.Related) != 3 {
+		t.Errorf("expected 3 related entities (a, c, d), got %d", len(cg.Related))
+	}
+}
+
+func TestService_GetContext_MaxDepthCap(t *testing.T) {
+	repo := newMockRepo()
+	edges := &mockEdgeRepo{}
+	svc := NewService(repo, edges, nil)
+	ctx := context.Background()
+	ns := namespace.DefaultNamespace
+
+	_, _ = svc.Upsert(ctx, ns, &Entity{URN: "urn:a", Type: TypeTable, Name: "a"})
+
+	// depth=10 should be capped to maxContextDepth (5), not error
+	cg, err := svc.GetContext(ctx, ns, "urn:a", 10)
+	if err != nil {
+		t.Fatalf("GetContext with large depth should not error: %v", err)
+	}
+	if cg.Entity.URN != "urn:a" {
+		t.Errorf("expected entity urn:a, got %s", cg.Entity.URN)
+	}
+}
+
+func TestService_GetContext_NilEdges(t *testing.T) {
+	repo := newMockRepo()
+	svc := NewService(repo, nil, nil) // no edge repo
+	ctx := context.Background()
+	ns := namespace.DefaultNamespace
+
+	_, _ = svc.Upsert(ctx, ns, &Entity{URN: "urn:a", Type: TypeTable, Name: "a"})
+
+	cg, err := svc.GetContext(ctx, ns, "urn:a", 2)
+	if err != nil {
+		t.Fatalf("GetContext with nil edges should not error: %v", err)
+	}
+	if len(cg.Edges) != 0 {
+		t.Errorf("expected 0 edges, got %d", len(cg.Edges))
+	}
+	if len(cg.Related) != 0 {
+		t.Errorf("expected 0 related, got %d", len(cg.Related))
+	}
+}
+
+func TestService_GetContext_CycleHandling(t *testing.T) {
+	repo := newMockRepo()
+	edges := &mockEdgeRepo{}
+	svc := NewService(repo, edges, nil)
+	ctx := context.Background()
+	ns := namespace.DefaultNamespace
+
+	// Cycle: A -> B -> C -> A
+	_, _ = svc.Upsert(ctx, ns, &Entity{URN: "urn:a", Type: TypeTable, Name: "a"})
+	_, _ = svc.Upsert(ctx, ns, &Entity{URN: "urn:b", Type: TypeTable, Name: "b"})
+	_, _ = svc.Upsert(ctx, ns, &Entity{URN: "urn:c", Type: TypeTable, Name: "c"})
+	edges.edges = []Edge{
+		{SourceURN: "urn:a", TargetURN: "urn:b", Type: "lineage"},
+		{SourceURN: "urn:b", TargetURN: "urn:c", Type: "lineage"},
+		{SourceURN: "urn:c", TargetURN: "urn:a", Type: "lineage"},
+	}
+
+	// depth=3 should not infinite loop
+	cg, err := svc.GetContext(ctx, ns, "urn:a", 3)
+	if err != nil {
+		t.Fatalf("GetContext with cycle should not error: %v", err)
+	}
+	if len(cg.Edges) != 3 {
+		t.Errorf("expected 3 edges in cycle, got %d", len(cg.Edges))
+	}
+	if len(cg.Related) != 2 {
+		t.Errorf("expected 2 related entities (b, c), got %d", len(cg.Related))
+	}
+}
+
 func TestService_Search_NilRepos(t *testing.T) {
 	svc := NewService(newMockRepo(), nil, nil)
 	results, err := svc.Search(context.Background(), SearchConfig{Text: "test"})

store/edge_repository.go

@@ -77,6 +77,41 @@ func (r *EdgeRepository) GetUpstream(ctx context.Context, ns *namespace.Namespac
 	return r.traverse(ctx, ns, urn, depth, "upstream")
 }
 
+func (r *EdgeRepository) GetBidirectional(ctx context.Context, ns *namespace.Namespace, urn string, depth int) ([]entity.Edge, error) {
+	if depth <= 0 {
+		depth = 1
+	}
+
+	query := `
+		WITH RECURSIVE graph(source_urn, target_urn, type, properties, depth, path, frontier) AS (
+			SELECT source_urn, target_urn, type, properties, 1, ARRAY[source_urn], target_urn
+			FROM edges
+			WHERE namespace_id = $1 AND source_urn = $2 AND valid_to IS NULL
+		UNION ALL
+			SELECT source_urn, target_urn, type, properties, 1, ARRAY[target_urn], source_urn
+			FROM edges
+			WHERE namespace_id = $1 AND target_urn = $2 AND valid_to IS NULL
+		UNION ALL
+			SELECT e.source_urn, e.target_urn, e.type, e.properties, g.depth + 1, g.path || g.frontier, e.target_urn
+			FROM edges e
+			JOIN graph g ON e.source_urn = g.frontier
+			WHERE e.target_urn <> ALL(g.path) AND e.valid_to IS NULL AND g.depth < $3
+		UNION ALL
+			SELECT e.source_urn, e.target_urn, e.type, e.properties, g.depth + 1, g.path || g.frontier, e.source_urn
+			FROM edges e
+			JOIN graph g ON e.target_urn = g.frontier
+			WHERE e.source_urn <> ALL(g.path) AND e.valid_to IS NULL AND g.depth < $3
+		)
+		SELECT DISTINCT source_urn, target_urn, type, properties FROM graph
+		LIMIT 1000`
+
+	var models []edgeModel
+	if err := r.client.SelectContext(ctx, &models, query, ns.ID, urn, depth); err != nil {
+		return nil, fmt.Errorf("traverse bidirectional: %w", err)
+	}
+	return toEdgeList(models), nil
+}
+
 func (r *EdgeRepository) Delete(ctx context.Context, ns *namespace.Namespace, sourceURN, targetURN, edgeType string) error {
 	_, err := r.client.ExecContext(ctx,
 		`UPDATE edges SET valid_to = now() WHERE namespace_id = $1 AND source_urn = $2 AND target_urn = $3 AND type = $4 AND valid_to IS NULL`,