From 02bee776abb61e82bb9f05e5c2d12dbe8a01c7dc Mon Sep 17 00:00:00 2001
From: jayhack <jayhack.0@gmail.com>
Date: Wed, 12 Feb 2025 18:03:15 -0800
Subject: [PATCH 1/2] .

---
 .../extensions/tools/file_operations.py       | 72 +++++++++++--------
 1 file changed, 41 insertions(+), 31 deletions(-)

diff --git a/src/codegen/extensions/tools/file_operations.py b/src/codegen/extensions/tools/file_operations.py
index d21d8ffac..42a15bbfe 100644
--- a/src/codegen/extensions/tools/file_operations.py
+++ b/src/codegen/extensions/tools/file_operations.py
@@ -1,6 +1,5 @@
 """File operations for manipulating the codebase."""
 
-import os
 from typing import Any, Literal
 
 from codegen import Codebase
@@ -36,9 +35,6 @@ def view_file(codebase: Codebase, filepath: str) -> dict[str, Any]:
 def list_directory(codebase: Codebase, dirpath: str = "./", depth: int = 1) -> dict[str, Any]:
     """List contents of a directory.
 
-    TODO(CG-10729): add support for directories that only including non-SourceFiles (code files). At the moment,
-     only files and directories that have SourceFile objects are included.
-
     Args:
         codebase: The codebase to operate on
         dirpath: Path to directory relative to workspace root
@@ -46,7 +42,21 @@ def list_directory(codebase: Codebase, dirpath: str = "./", depth: int = 1) -> d
                Use -1 for unlimited depth.
 
     Returns:
-        Dict containing directory contents and metadata, or error information if directory not found
+        Dict containing directory contents and metadata in a nested structure:
+        {
+            "path": str,
+            "name": str,
+            "files": list[str],
+            "subdirectories": [
+                {
+                    "path": str,
+                    "name": str,
+                    "files": list[str],
+                    "subdirectories": [...],
+                },
+                ...
+            ]
+        }
     """
     try:
         directory = codebase.get_directory(dirpath)
@@ -56,33 +66,33 @@ def list_directory(codebase: Codebase, dirpath: str = "./", depth: int = 1) -> d
     if not directory:
         return {"error": f"Directory not found: {dirpath}"}
 
-    # Get immediate files
-    files = []
-    subdirs = []
-
-    for item in directory.items.values():
-        if isinstance(item, Directory):
-            subdirs.append(item.name)
-        else:
-            # Get full filename with extension from filepath
-            files.append(os.path.basename(item.filepath))
-
-    # If depth > 1 or unlimited (-1), recursively get subdirectories
-    if depth != 1:
-        new_depth = depth - 1 if depth > 1 else -1
-        for item in directory.items.values():
-            if isinstance(item, Directory):
-                subdir_result = list_directory(codebase, os.path.join(dirpath, item.name), depth=new_depth)
-                if "error" not in subdir_result:
-                    files.extend(subdir_result["files"])
-                    subdirs.extend(subdir_result["subdirectories"])
+    def get_directory_info(dir_obj: Directory, current_depth: int) -> dict[str, Any]:
+        """Helper function to get directory info recursively."""
+        # Get direct files
+        all_files = []
+        for file in dir_obj.files:
+            if file.directory == dir_obj:
+                all_files.append(file.filepath.split("/")[-1])
+
+        # Get direct subdirectories
+        subdirs = []
+        for subdir in dir_obj.subdirectories:
+            # Only include direct descendants
+            if subdir.parent == dir_obj:
+                if current_depth != 1:
+                    new_depth = current_depth - 1 if current_depth > 1 else -1
+                    subdirs.append(get_directory_info(subdir, new_depth))
+                else:
+                    # At max depth, just include name
+                    subdirs.append(subdir.name)
+        return {
+            "name": dir_obj.name,
+            "path": dir_obj.dirpath,
+            "files": all_files,
+            "subdirectories": subdirs,
+        }
 
-    return {
-        "path": str(directory.path),  # Convert PosixPath to string
-        "name": directory.name,
-        "files": files,
-        "subdirectories": subdirs,
-    }
+    return get_directory_info(directory, depth)
 
 
 def edit_file(codebase: Codebase, filepath: str, content: str) -> dict[str, Any]:

From 507671fa5a804711a66b17f6871a870093d1a248 Mon Sep 17 00:00:00 2001
From: jayhack <jayhack.0@gmail.com>
Date: Wed, 12 Feb 2025 18:05:35 -0800
Subject: [PATCH 2/2] .

---
 src/codegen/extensions/tools/file_operations.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/src/codegen/extensions/tools/file_operations.py b/src/codegen/extensions/tools/file_operations.py
index 42a15bbfe..55794cb67 100644
--- a/src/codegen/extensions/tools/file_operations.py
+++ b/src/codegen/extensions/tools/file_operations.py
@@ -41,6 +41,9 @@ def list_directory(codebase: Codebase, dirpath: str = "./", depth: int = 1) -> d
         depth: How deep to traverse the directory tree. Default is 1 (immediate children only).
                Use -1 for unlimited depth.
 
+    TODO(CG-10729): add support for directories that only including non-SourceFiles (code files). At the moment,
+     only files and directories that have SourceFile objects are included.
+
     Returns:
         Dict containing directory contents and metadata in a nested structure:
         {