allow list files to exit gracefully when there are too many files and add a non recursive version to limit return files

Author: m-y-mo

src/seclab_taskflows/mcp_servers/local_file_viewer.py

@@ -26,6 +26,8 @@
 
 LINE_LIMIT_FOR_FETCHING_FILE_CONTENT = int(os.getenv('LINE_LIMIT_FOR_FETCHING_FILE_CONTENT', default=1000))
 
+FILE_LIMIT_FOR_LIST_FILES = int(os.getenv('FILE_LIMIT_FOR_LIST_FILES', default=100))
+
 def is_subdirectory(directory, potential_subdirectory):
     directory_path = Path(directory)
     potential_subdirectory_path = Path(potential_subdirectory)
@@ -69,16 +71,21 @@ def search_zipfile(database_path, term, search_dir = None):
                             results[filename].append(i+1)
     return results
 
-def _list_files(database_path, root_dir = None):
+def _list_files(database_path, root_dir = None, recursive=True):
     results = []
     root_dir = strip_leading_dash(root_dir)
     with zipfile.ZipFile(database_path) as z:
         for entry in z.infolist():
             if entry.is_dir():
-                continue
+                if not recursive:
+                    dirname = remove_root_dir(entry.filename)
+                    if Path(dirname).parent == Path(root_dir):
+                        results.append(dirname)
             filename = remove_root_dir(entry.filename)
             if root_dir and not is_subdirectory(root_dir, filename):
                 continue
+            if not recursive and Path(filename).parent != Path(root_dir):
+                    continue
             results.append(filename)
     return results
 
@@ -152,8 +159,26 @@ async def list_files(
     if not source_path or not source_path.exists():
         return f"Invalid {owner} and {repo}. Check that the input is correct or try to fetch the repo from gh first."
     content = _list_files(source_path, path)
+    if len(content) > FILE_LIMIT_FOR_LIST_FILES:
+        return f"Too many files to display in {owner}/{repo} at path {path} ({len(content)} files). Try use `list_files_non_recursive` instead."
     return json.dumps(content, indent=2)
 
+@mcp.tool()
+async def list_files_non_recursive(
+    owner: str = Field(description="The owner of the repository"),
+    repo: str = Field(description="The name of the repository"),
+    path: str = Field(description="The path to the directory in the repository")) -> str:
+    """
+    List the files of a directory from a local GitHub repository non recursively.
+    """
+    source_path = Path(f"{LOCAL_GH_DIR}/{owner}/{repo}.zip")
+    source_path = sanitize_file_path(source_path, [LOCAL_GH_DIR])
+    if not source_path or not source_path.exists():
+        return f"Invalid {owner} and {repo}. Check that the input is correct or try to fetch the repo from gh first."
+    content = _list_files(source_path, path, recursive=False)
+    return json.dumps(content, indent=2)
+
+
 @mcp.tool()
 async def search_repo(
     owner: str = Field(description="The owner of the repository"),

src/seclab_taskflows/toolboxes/local_file_viewer.yaml

@@ -12,3 +12,4 @@ server_params:
   env:
     LOCAL_GH_DIR: "{{ env DATA_DIR }}"
     LINE_LIMIT_FOR_FETCHING_FILE_CONTENT: "{{ env LINE_LIMIT_FOR_FETCHING_FILE_CONTENT }}"
+    FILE_LIMIT_FOR_LIST_FILES: "{{ env FILE_LIMIT_FOR_LIST_FILES }}"