pkg/codesearch: add read-file command

Just provides full file contents as last resort.

7 files changed, 66 insertions, 0 deletions

diff --git a/pkg/aflow/tool/codesearcher/codesearcher.go b/pkg/aflow/tool/codesearcher/codesearcher.go
index c336a0ca4..34db81b80 100644
--- a/pkg/aflow/tool/codesearcher/codesearcher.go
+++ b/pkg/aflow/tool/codesearcher/codesearcher.go
@@ -17,6 +17,11 @@ var Tools = []aflow.Tool{
 	aflow.NewFuncTool("codesearch-dir-index", dirIndex, `
 Tool provides list of source files and subdirectories in the given directory in the source tree.
 `),
+	aflow.NewFuncTool("read-file", readFile, `
+Tool provides full contents of a single source file as is. Avoid using this tool if there are better
+and more specialized tools for the job, because source files may be large and contain lots
+of unrelated information.
+`),
 	aflow.NewFuncTool("codesearch-file-index", fileIndex, `
 Tool provides list of entities defined in the given source file.
 Entity can be function, struct, or global variable.
@@ -68,6 +73,15 @@ type dirIndexResult struct {
 	Files   []string `jsonschema:"List of source files."`
 }
 
+type readFileArgs struct {
+	File string `jsonschema:"Source file path."`
+}
+
+type readFileResult struct {
+	Missing  bool   `jsonschema:"Set to true if the requested file does not exist."`
+	Contents string `jsonschema:"File contents."`
+}
+
 type fileIndexArgs struct {
 	SourceFile string `jsonschema:"Source file path."`
 }
@@ -154,6 +168,15 @@ func dirIndex(ctx *aflow.Context, state prepareResult, args dirIndexArgs) (dirIn
 	return res, err
 }
 
+func readFile(ctx *aflow.Context, state prepareResult, args readFileArgs) (readFileResult, error) {
+	ok, contents, err := state.Index.ReadFile(args.File)
+	res := readFileResult{
+		Missing:  !ok,
+		Contents: contents,
+	}
+	return res, err
+}
+
 func fileIndex(ctx *aflow.Context, state prepareResult, args fileIndexArgs) (fileIndexResult, error) {
 	ok, entities, err := state.Index.FileIndex(args.SourceFile)
 	res := fileIndexResult{
diff --git a/pkg/codesearch/codesearch.go b/pkg/codesearch/codesearch.go
index 8e0259af7..96ee1c696 100644
--- a/pkg/codesearch/codesearch.go
+++ b/pkg/codesearch/codesearch.go
@@ -45,6 +45,13 @@ var Commands = []Command{
 		}
 		return b.String(), nil
 	}},
+	{"read-file", 1, func(index *Index, args []string) (string, error) {
+		ok, contents, err := index.ReadFile(args[0])
+		if err != nil || !ok {
+			return notFound, err
+		}
+		return contents, nil
+	}},
 	{"file-index", 1, func(index *Index, args []string) (string, error) {
 		ok, entities, err := index.FileIndex(args[0])
 		if err != nil || !ok {
@@ -137,6 +144,27 @@ func (index *Index) DirIndex(dir string) (bool, []string, []string, error) {
 	return exists, subdirs, files, nil
 }
 
+func (index *Index) ReadFile(file string) (bool, string, error) {
+	if err := escaping(file); err != nil {
+		return false, "", nil
+	}
+	for _, dir := range index.srcDirs {
+		data, err := os.ReadFile(filepath.Join(dir, file))
+		if err != nil {
+			if os.IsNotExist(err) {
+				continue
+			}
+			var errno syscall.Errno
+			if errors.As(err, &errno) && errno == syscall.EISDIR {
+				return false, "", nil
+			}
+			return false, "", err
+		}
+		return true, string(data), nil
+	}
+	return false, "", nil
+}
+
 func (index *Index) FileIndex(file string) (bool, []Entity, error) {
 	var entities []Entity
 	for _, def := range index.db.Definitions {
diff --git a/pkg/codesearch/testdata/mm/slub.c b/pkg/codesearch/testdata/mm/slub.c
index e69de29bb..4b6741a7f 100644
--- a/pkg/codesearch/testdata/mm/slub.c
+++ b/pkg/codesearch/testdata/mm/slub.c
@@ -0,0 +1,2 @@
+// slub.c contents.
+// This file is used in read-file test.
diff --git a/pkg/codesearch/testdata/query-read-file-dir b/pkg/codesearch/testdata/query-read-file-dir
new file mode 100644
index 000000000..210a326cd
--- /dev/null
+++ b/pkg/codesearch/testdata/query-read-file-dir
@@ -0,0 +1,3 @@
+read-file mm
+
+not found
diff --git a/pkg/codesearch/testdata/query-read-file-escaping b/pkg/codesearch/testdata/query-read-file-escaping
new file mode 100644
index 000000000..fca2abf6a
--- /dev/null
+++ b/pkg/codesearch/testdata/query-read-file-escaping
@@ -0,0 +1,3 @@
+read-file mm/../../codesearch.go
+
+not found
diff --git a/pkg/codesearch/testdata/query-read-file-existing b/pkg/codesearch/testdata/query-read-file-existing
new file mode 100644
index 000000000..9fc30a164
--- /dev/null
+++ b/pkg/codesearch/testdata/query-read-file-existing
@@ -0,0 +1,4 @@
+read-file mm/slub.c
+
+// slub.c contents.
+// This file is used in read-file test.
diff --git a/pkg/codesearch/testdata/query-read-file-missing b/pkg/codesearch/testdata/query-read-file-missing
new file mode 100644
index 000000000..ac7bead8d
--- /dev/null
+++ b/pkg/codesearch/testdata/query-read-file-missing
@@ -0,0 +1,3 @@
+read-file file-that-does-not-exist.c
+
+not found