From b99f119e375f59f3df059e7a8bb1c623f6ea61a8 Mon Sep 17 00:00:00 2001
From: viettranx <viettranx@gmail.com>
Date: Tue, 31 Mar 2026 17:55:27 +0700
Subject: [PATCH] feat(kg): bidirectional multi-hop traversal
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Add reverse-edge UNION ALL to recursive CTE so traversal follows both
source→target and target→source edges. Reverse edges prefixed with ~
in via field (e.g. ~manages). Tool output shows directional arrows.
---
 .../store/pg/knowledge_graph_traversal.go     | 32 +++++++++++++++++++
 internal/tools/knowledge_graph.go             |  6 +++-
 2 files changed, 37 insertions(+), 1 deletion(-)

diff --git a/internal/store/pg/knowledge_graph_traversal.go b/internal/store/pg/knowledge_graph_traversal.go
index b1e7de9d..48e834f7 100644
--- a/internal/store/pg/knowledge_graph_traversal.go
+++ b/internal/store/pg/knowledge_graph_traversal.go
@@ -68,6 +68,22 @@ func (s *PGKnowledgeGraphStore) Traverse(ctx context.Context, agentID, userID, s
 			JOIN kg_entities  e ON r.target_entity_id = e.id AND e.agent_id = $2
 			WHERE p.depth < $%d
 			  AND NOT e.id::text = ANY(p.path)
+
+			UNION ALL
+
+			SELECT
+				e.id, e.agent_id, e.user_id, e.external_id,
+				e.name, e.entity_type, e.description,
+				e.properties, e.source_id, e.confidence,
+				e.created_at, e.updated_at,
+				p.depth + 1,
+				p.path || e.id::text,
+				'~' || r.relation_type
+			FROM paths p
+			JOIN kg_relations r ON p.id = r.target_entity_id AND r.agent_id = $2
+			JOIN kg_entities  e ON r.source_entity_id = e.id AND e.agent_id = $2
+			WHERE p.depth < $%[2]d
+			  AND NOT e.id::text = ANY(p.path)
 		)
 		SELECT
 			id, agent_id, user_id, external_id,
@@ -113,6 +129,22 @@ func (s *PGKnowledgeGraphStore) Traverse(ctx context.Context, agentID, userID, s
 			JOIN kg_entities  e ON r.target_entity_id = e.id AND e.user_id = $3
 			WHERE p.depth < $%d
 			  AND NOT e.id::text = ANY(p.path)
+
+			UNION ALL
+
+			SELECT
+				e.id, e.agent_id, e.user_id, e.external_id,
+				e.name, e.entity_type, e.description,
+				e.properties, e.source_id, e.confidence,
+				e.created_at, e.updated_at,
+				p.depth + 1,
+				p.path || e.id::text,
+				'~' || r.relation_type
+			FROM paths p
+			JOIN kg_relations r ON p.id = r.target_entity_id AND r.user_id = $3
+			JOIN kg_entities  e ON r.source_entity_id = e.id AND e.user_id = $3
+			WHERE p.depth < $%[2]d
+			  AND NOT e.id::text = ANY(p.path)
 		)
 		SELECT
 			id, agent_id, user_id, external_id,
diff --git a/internal/tools/knowledge_graph.go b/internal/tools/knowledge_graph.go
index ab749b99..00f62322 100644
--- a/internal/tools/knowledge_graph.go
+++ b/internal/tools/knowledge_graph.go
@@ -110,7 +110,11 @@ func (t *KnowledgeGraphSearchTool) executeTraversal(ctx context.Context, agentID
 		for _, r := range results {
 			sb.WriteString(fmt.Sprintf("- [depth %d] %s (%s)", r.Depth, r.Entity.Name, r.Entity.EntityType))
 			if r.Via != "" {
-				sb.WriteString(fmt.Sprintf(" via %q", r.Via))
+				if strings.HasPrefix(r.Via, "~") {
+					sb.WriteString(fmt.Sprintf(" ←[%s]—", r.Via[1:]))
+				} else {
+					sb.WriteString(fmt.Sprintf(" —[%s]→", r.Via))
+				}
 			}
 			if r.Entity.Description != "" {
 				sb.WriteString(fmt.Sprintf("\n  %s", r.Entity.Description))