feat: add PruneTree helper and depth parameter to get_rag_tree

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: tae2089

internal/mcp/handlers.go

@@ -1383,8 +1383,9 @@ func (h *handlers) buildRagIndex(ctx context.Context, request mcp.CallToolReques
 
 func (h *handlers) getRagTree(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
 	communityID := request.GetString("community_id", "")
+	depth := int(request.GetFloat("depth", 0))
 
-	key := "get_rag_tree:" + mustJSON(map[string]any{"community_id": communityID})
+	key := "get_rag_tree:" + mustJSON(map[string]any{"community_id": communityID, "depth": depth})
 	if h.cache != nil {
 		if cached, ok := h.cache.Get(key); ok {
 			return mcp.NewToolResultText(cached), nil
@@ -1406,6 +1407,10 @@ func (h *handlers) getRagTree(ctx context.Context, request mcp.CallToolRequest)
 		}
 	}
 
+	if depth > 0 {
+		node = ragindex.PruneTree(node, depth)
+	}
+
 	b, err := json.Marshal(node)
 	if err != nil {
 		return mcp.NewToolResultError(fmt.Sprintf("marshal tree: %v", err)), nil

internal/mcp/handlers_test.go

@@ -27,6 +27,7 @@ import (
 	"github.com/imtaebin/code-context-graph/internal/analysis/incremental"
 	"github.com/imtaebin/code-context-graph/internal/analysis/query"
 	"github.com/imtaebin/code-context-graph/internal/model"
+	"github.com/imtaebin/code-context-graph/internal/ragindex"
 	"github.com/imtaebin/code-context-graph/internal/store/gormstore"
 	"github.com/imtaebin/code-context-graph/internal/store/search"
 )
@@ -2110,3 +2111,69 @@ func TestGetRagTree_InvalidCommunityID(t *testing.T) {
 		t.Fatal("expected error for nonexistent community_id")
 	}
 }
+
+func TestGetRagTree_DepthLimitsChildren(t *testing.T) {
+	deps := setupTestDeps(t)
+
+	// 임시 인덱스 디렉토리 설정
+	tmpDir := t.TempDir()
+	deps.RagIndexDir = filepath.Join(tmpDir, ".ccg")
+
+	// DB에 community + node + CommunityMembership 생성
+	community := model.Community{Key: "auth", Label: "Auth Community", Strategy: "auto"}
+	if err := deps.DB.Create(&community).Error; err != nil {
+		t.Fatalf("create community: %v", err)
+	}
+
+	node := model.Node{
+		QualifiedName: "auth.Login",
+		Kind:          model.NodeKindFunction,
+		Name:          "Login",
+		FilePath:      "internal/auth/login.go",
+		StartLine:     1,
+		EndLine:       10,
+		Language:      "go",
+	}
+	if err := deps.DB.Create(&node).Error; err != nil {
+		t.Fatalf("create node: %v", err)
+	}
+
+	membership := model.CommunityMembership{
+		CommunityID: community.ID,
+		NodeID:      node.ID,
+	}
+	if err := deps.DB.Create(&membership).Error; err != nil {
+		t.Fatalf("create membership: %v", err)
+	}
+
+	// ragindex.Builder로 인덱스 빌드
+	b := &ragindex.Builder{
+		DB:       deps.DB,
+		OutDir:   filepath.Join(tmpDir, "docs"),
+		IndexDir: deps.RagIndexDir,
+	}
+	if _, _, err := b.Build(); err != nil {
+		t.Fatalf("Build: %v", err)
+	}
+
+	// depth=1로 get_rag_tree 호출: community 노드는 있지만 파일 노드는 없어야 함
+	result := callTool(t, deps, "get_rag_tree", map[string]any{
+		"depth": float64(1),
+	})
+	if result.IsError {
+		t.Fatalf("get_rag_tree error: %v", getTextContent(result))
+	}
+
+	var treeNode ragindex.TreeNode
+	if err := json.Unmarshal([]byte(getTextContent(result)), &treeNode); err != nil {
+		t.Fatalf("unmarshal tree: %v", err)
+	}
+
+	if len(treeNode.Children) == 0 {
+		t.Fatal("expected community nodes at depth=1, got none")
+	}
+	communityNode := treeNode.Children[0]
+	if len(communityNode.Children) != 0 {
+		t.Fatalf("expected 0 file children at depth=1, got %d", len(communityNode.Children))
+	}
+}

internal/mcp/server.go

@@ -266,6 +266,7 @@ func NewServer(deps *Deps) *server.MCPServer {
 			Tool: mcp.NewTool("get_rag_tree",
 				mcp.WithDescription("Get the RAG document tree for navigation. Call without arguments first to see all communities, then pass community_id to drill into a specific one."),
 				mcp.WithString("community_id", mcp.Description("Community node ID as shown in the tree (e.g. 'community:auth'). Omit to get the full tree.")),
+				mcp.WithNumber("depth", mcp.Description("Maximum tree depth to return (1=communities only, 2=communities+files, 3=including symbols). Default: 0 (unlimited).")),
 			),
 			Handler: h.getRagTree,
 		},

internal/ragindex/builder.go

@@ -303,3 +303,31 @@ func FindNode(root *TreeNode, id string) *TreeNode {
 	}
 	return nil
 }
+
+// PruneTree는 root 트리를 maxDepth 깊이까지만 포함한 새 트리를 반환한다.
+// maxDepth <= 0이면 전체 트리를 반환한다. 원본 트리는 변경하지 않는다.
+// depth 계산: root는 depth 0, root의 직계 자식은 depth 1.
+func PruneTree(root *TreeNode, maxDepth int) *TreeNode {
+	if root == nil {
+		return nil
+	}
+	return pruneNode(root, 0, maxDepth)
+}
+
+func pruneNode(n *TreeNode, currentDepth, maxDepth int) *TreeNode {
+	copied := &TreeNode{
+		ID:      n.ID,
+		Label:   n.Label,
+		Summary: n.Summary,
+		DocPath: n.DocPath,
+	}
+	if maxDepth <= 0 || currentDepth < maxDepth {
+		copied.Children = make([]*TreeNode, 0, len(n.Children))
+		for _, child := range n.Children {
+			copied.Children = append(copied.Children, pruneNode(child, currentDepth+1, maxDepth))
+		}
+	} else {
+		copied.Children = []*TreeNode{}
+	}
+	return copied
+}

internal/ragindex/builder_test.go

@@ -328,6 +328,61 @@ func TestBuilder_ProjectDesc(t *testing.T) {
 	}
 }
 
+// TestPruneTree_Depth1: depth=1이면 root와 직계 자식만 반환, 손자 노드 없음.
+func TestPruneTree_Depth1(t *testing.T) {
+	root := &ragindex.TreeNode{
+		ID:    "root",
+		Label: "Root",
+		Children: []*ragindex.TreeNode{
+			{
+				ID:    "c1",
+				Label: "Community 1",
+				Children: []*ragindex.TreeNode{
+					{ID: "f1", Label: "file.go"},
+				},
+			},
+		},
+	}
+
+	result := ragindex.PruneTree(root, 1)
+	if len(result.Children) != 1 {
+		t.Fatalf("expected 1 child, got %d", len(result.Children))
+	}
+	if len(result.Children[0].Children) != 0 {
+		t.Fatalf("expected 0 grandchildren at depth=1, got %d", len(result.Children[0].Children))
+	}
+	// 원본 트리는 변경되지 않아야 함
+	if len(root.Children[0].Children) != 1 {
+		t.Fatal("PruneTree must not modify the original tree")
+	}
+}
+
+// TestPruneTree_NegativeDepth: depth <= 0이면 트리 전체를 반환한다.
+func TestPruneTree_NegativeDepth(t *testing.T) {
+	root := &ragindex.TreeNode{
+		ID: "root",
+		Children: []*ragindex.TreeNode{
+			{ID: "c1", Children: []*ragindex.TreeNode{{ID: "f1"}}},
+		},
+	}
+
+	result := ragindex.PruneTree(root, 0)
+	if len(result.Children) != 1 {
+		t.Fatalf("expected 1 child, got %d", len(result.Children))
+	}
+	if len(result.Children[0].Children) != 1 {
+		t.Fatalf("expected 1 grandchild at depth=0 (unlimited), got %d", len(result.Children[0].Children))
+	}
+}
+
+// TestPruneTree_NilRoot: nil 입력 → nil 반환.
+func TestPruneTree_NilRoot(t *testing.T) {
+	result := ragindex.PruneTree(nil, 2)
+	if result != nil {
+		t.Fatal("expected nil for nil root")
+	}
+}
+
 // TestFindNode: FindNode가 재귀적으로 트리에서 노드를 찾는지 검증한다.
 func TestFindNode(t *testing.T) {
 	root := &ragindex.TreeNode{