From 5c1248184b96357e03462c7d0db1a0fc7435bcd4 Mon Sep 17 00:00:00 2001
From: Matt Low <matt@mlow.ca>
Date: Tue, 21 May 2024 00:08:17 +0000
Subject: [PATCH] Update dir_tree to have maximum depth of 5

Until we have some mechanism in place for confirming tool calls with the
user before executing, it's dangerous to allow unlimited depth
---
 pkg/lmcli/tools/dir_tree.go | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/pkg/lmcli/tools/dir_tree.go b/pkg/lmcli/tools/dir_tree.go
index fb0e2c4..7d7d8d9 100644
--- a/pkg/lmcli/tools/dir_tree.go
+++ b/pkg/lmcli/tools/dir_tree.go
@@ -39,7 +39,7 @@ var DirTreeTool = model.Tool{
 		{
 			Name:        "depth",
 			Type:        "integer",
-			Description: "Depth of directory recursion. Default 0. Use -1 for unlimited.",
+			Description: "Depth of directory recursion. Defaults to 0 (no recursion), maximum of 5.",
 		},
 	},
 	Impl: func(tool *model.Tool, args map[string]interface{}) (string, error) {
@@ -61,6 +61,7 @@ var DirTreeTool = model.Tool{
 				if depth, err = strconv.Atoi(v); err != nil {
 					return "", fmt.Errorf("invalid `depth` value, expected integer but got string that cannot convert: %v", tmp)
 				}
+				depth = max(0, min(5, depth))
 			default:
 				return "", fmt.Errorf("expected int or string for max_depth, got %T", tmp)
 			}
@@ -123,7 +124,7 @@ func buildTree(output *strings.Builder, path string, prefix string, depth int) e
 		output.WriteString(prefix + branch + file.Name())
 		if file.IsDir() {
 			output.WriteString("/\n")
-			if depth != 0 {
+			if depth > 0 {
 				var nextPrefix string
 				if isLast {
 					nextPrefix = prefix + "    "