improve listTree output

Author: brendan-kellam

packages/web/src/features/chat/agent.ts

@@ -22,6 +22,7 @@ import { grepDefinition } from "@/features/tools/grep";
 import { Source } from "./types";
 import { addLineNumbers, fileReferenceToString } from "./utils";
 import { tools } from "./tools";
+import { listTreeDefinition } from "../tools";
 
 const dedent = _dedent.withOptions({ alignValues: true });
 
@@ -223,27 +224,43 @@ const createAgentStream = async ({
                 if (toolName === readFileDefinition.name) {
                     onWriteSource({
                         type: 'file',
-                        language: output.metadata.language,
                         repo: output.metadata.repo,
                         path: output.metadata.path,
                         revision: output.metadata.revision,
                         name: output.metadata.path.split('/').pop() ?? output.metadata.path,
                     });
                 } else if (toolName === grepDefinition.name) {
                     output.metadata.files.forEach((file) => {
-                        onWriteSource(file);
+                        onWriteSource({
+                            type: 'file',
+                            repo: file.repo,
+                            path: file.path,
+                            revision: file.revision,
+                            name: file.path.split('/').pop() ?? file.path,
+                        });
                     });
                 } else if (toolName === findSymbolDefinitionsDefinition.name || toolName === findSymbolReferencesDefinition.name) {
                     output.metadata.files.forEach((file) => {
                         onWriteSource({
                             type: 'file',
-                            language: file.language,
                             repo: file.repo,
                             path: file.fileName,
                             revision: file.revision,
                             name: file.fileName.split('/').pop() ?? file.fileName,
                         });
                     });
+                } else if (toolName === listTreeDefinition.name) {
+                    output.metadata.entries
+                        .filter((entry) => entry.type === 'blob')
+                        .forEach((entry) => {
+                            onWriteSource({
+                                type: 'file',
+                                repo: output.metadata.repo,
+                                path: entry.path,
+                                revision: output.metadata.ref,
+                                name: entry.name,
+                            });
+                        });
                 }
             });
         },

packages/web/src/features/chat/types.ts

@@ -11,7 +11,6 @@ const fileSourceSchema = z.object({
     repo: z.string(),
     path: z.string(),
     name: z.string(),
-    language: z.string(),
     revision: z.string(),
 });
 export type FileSource = z.infer<typeof fileSourceSchema>;

packages/web/src/features/chat/utils.ts

@@ -187,7 +187,6 @@ export const createUIMessage = (text: string, mentions: MentionData[], selectedS
                     path: mention.path,
                     repo: mention.repo,
                     name: mention.name,
-                    language: mention.language,
                     revision: mention.revision,
                 }
                 return fileSource;

packages/web/src/features/tools/grep.ts

@@ -6,7 +6,6 @@ import escapeStringRegexp from "escape-string-regexp";
 import { ToolDefinition } from "./types";
 import { logger } from "./logger";
 import description from "./grep.txt";
-import { FileSource } from "../chat/types";
 
 const DEFAULT_SEARCH_LIMIT = 100;
 const MAX_LINE_LENGTH = 2000;
@@ -44,8 +43,15 @@ const grepShape = {
         .optional(),
 };
 
+export type GrepFile = {
+    path: string;
+    name: string;
+    repo: string;
+    revision: string;
+};
+
 export type GrepMetadata = {
-    files: FileSource[];
+    files: GrepFile[];
     query: string;
 };
 
@@ -102,13 +108,11 @@ export const grepDefinition: ToolDefinition<'grep', typeof grepShape, GrepMetada
 
         const metadata: GrepMetadata = {
             files: response.files.map((file) => ({
-                type: 'file',
                 path: file.fileName.text,
                 name: file.fileName.text.split('/').pop() ?? file.fileName.text,
-                language: file.language,
                 repo: file.repository,
                 revision: ref ?? 'HEAD',
-            })),
+            } satisfies GrepFile)),
             query,
         };
 

packages/web/src/features/tools/listTree.ts

@@ -52,12 +52,14 @@ export const listTreeDefinition: ToolDefinition<'list_tree', typeof listTreeShap
 
         if (!includeFiles && !includeDirectories) {
             const metadata: ListTreeMetadata = {
-                repo, ref, path: normalizedPath,
+                repo,
+                ref,
+                path: normalizedPath,
                 entries: [],
                 totalReturned: 0,
                 truncated: false,
             };
-            return { output: JSON.stringify(metadata), metadata };
+            return { output: 'No entries found', metadata };
         }
 
         const queue: Array<{ path: string; depth: number }> = [{ path: normalizedPath, depth: 0 }];
@@ -135,6 +137,38 @@ export const listTreeDefinition: ToolDefinition<'list_tree', typeof listTreeShap
             truncated,
         };
 
-        return { output: JSON.stringify(metadata), metadata };
+        const outputLines = [normalizedPath || '/'];
+
+        const childrenByPath = new Map<string, ListTreeEntry[]>();
+        for (const entry of sortedEntries) {
+            const siblings = childrenByPath.get(entry.parentPath) ?? [];
+            siblings.push(entry);
+            childrenByPath.set(entry.parentPath, siblings);
+        }
+
+        function renderEntries(parentPath: string) {
+            const children = childrenByPath.get(parentPath) ?? [];
+            for (const entry of children) {
+                const indent = '  '.repeat(entry.depth);
+                const label = entry.type === 'tree' ? `${entry.name}/` : entry.name;
+                outputLines.push(`${indent}${label}`);
+                if (entry.type === 'tree') {
+                    renderEntries(entry.path);
+                }
+            }
+        }
+
+        renderEntries(normalizedPath);
+
+        if (sortedEntries.length === 0) {
+            outputLines.push('  (no entries found)');
+        }
+
+        if (truncated) {
+            outputLines.push('');
+            outputLines.push(`(truncated — showing first ${normalizedMaxEntries} entries)`);
+        }
+
+        return { output: outputLines.join('\n'), metadata };
     },
 };