agentu/internal/tools/code.go

package tools

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"go/ast"
	"go/parser"
	"go/printer"
	"go/token"
	"io/fs"
	"os"
	"path/filepath"
	"sort"
	"strings"

	"agentu/pkg/llm"
)

const defaultCodeSymbolFileLimit = 50

type CodeSymbolsTool struct {
	workingDir string
}

func NewCodeSymbolsTool(workingDir string) *CodeSymbolsTool {
	return &CodeSymbolsTool{workingDir: workingDir}
}

func (t *CodeSymbolsTool) Definition() llm.Tool {
	return llm.Tool{
		Type: "function",
		Function: llm.ToolFunction{
			Name:        "code_symbols",
			Description: "List Go code symbols (packages, imports, types, functions, methods) for a file or directory without reading full file contents.",
			Parameters: JSONSchema(`{
  "type": "object",
  "properties": {
    "path": {"type": "string", "description": "Go file or directory to inspect. Defaults to the configured agent working directory."},
    "max_files": {"type": "integer", "description": "Maximum number of Go files to inspect when path is a directory. Defaults to 50."}
  },
  "additionalProperties": false
}`),
		},
	}
}

type codeSymbolsArgs struct {
	Path     string `json:"path"`
	MaxFiles int    `json:"max_files"`
}

func (t *CodeSymbolsTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
	var args codeSymbolsArgs
	if err := json.Unmarshal(raw, &args); err != nil {
		return "", fmt.Errorf("parse arguments: %w", err)
	}
	path, err := resolvePath(t.workingDir, args.Path)
	if err != nil {
		return "", err
	}
	limit := args.MaxFiles
	if limit <= 0 {
		limit = defaultCodeSymbolFileLimit
	}

	files, err := goFiles(ctx, path, limit)
	if err != nil {
		return "", err
	}
	if len(files) == 0 {
		return "no Go files found", nil
	}

	fset := token.NewFileSet()
	var out strings.Builder
	for i, file := range files {
		if ctxErr := ctx.Err(); ctxErr != nil {
			return FormatResult(out.String()), ctxErr
		}
		if i > 0 {
			out.WriteString("\n")
		}
		if err := appendGoSymbols(&out, fset, file); err != nil {
			fmt.Fprintf(&out, "%s\nerror: %v\n", file, err)
		}
		if out.Len() > MaxToolOutputBytes {
			break
		}
	}
	return FormatResult(out.String()), nil
}

func goFiles(ctx context.Context, path string, limit int) ([]string, error) {
	info, err := os.Stat(path)
	if err != nil {
		return nil, err
	}
	if !info.IsDir() {
		if strings.HasSuffix(path, ".go") {
			return []string{path}, nil
		}
		return nil, fmt.Errorf("path is not a Go file: %s", path)
	}

	var files []string
	stop := errors.New("file limit reached")
	err = filepath.WalkDir(path, func(current string, d fs.DirEntry, err error) error {
		if err != nil {
			return nil
		}
		if ctxErr := ctx.Err(); ctxErr != nil {
			return ctxErr
		}
		if d.IsDir() {
			switch d.Name() {
			case ".git", "vendor", "node_modules":
				return filepath.SkipDir
			}
			return nil
		}
		if strings.HasSuffix(current, ".go") {
			files = append(files, current)
			if len(files) >= limit {
				return stop
			}
		}
		return nil
	})
	if err != nil && !errors.Is(err, stop) {
		return nil, err
	}
	sort.Strings(files)
	return files, nil
}

func appendGoSymbols(out *strings.Builder, fset *token.FileSet, path string) error {
	file, err := parser.ParseFile(fset, path, nil, parser.ImportsOnly|parser.ParseComments)
	if err != nil {
		return err
	}
	full, err := parser.ParseFile(fset, path, nil, 0)
	if err != nil {
		return err
	}

	fmt.Fprintf(out, "%s\n", path)
	fmt.Fprintf(out, "  package %s\n", full.Name.Name)
	if len(file.Imports) > 0 {
		imports := make([]string, 0, len(file.Imports))
		for _, imp := range file.Imports {
			imports = append(imports, strings.Trim(imp.Path.Value, "\""))
		}
		fmt.Fprintf(out, "  imports: %s\n", strings.Join(imports, ", "))
	}

	var types, funcs, methods []string
	for _, decl := range full.Decls {
		switch d := decl.(type) {
		case *ast.GenDecl:
			if d.Tok != token.TYPE {
				continue
			}
			for _, spec := range d.Specs {
				if ts, ok := spec.(*ast.TypeSpec); ok {
					kind := typeKind(ts.Type)
					types = append(types, fmt.Sprintf("%d: type %s %s", fset.Position(ts.Pos()).Line, ts.Name.Name, kind))
				}
			}
		case *ast.FuncDecl:
			line := fset.Position(d.Pos()).Line
			if d.Recv == nil {
				funcs = append(funcs, fmt.Sprintf("%d: func %s", line, d.Name.Name))
				continue
			}
			methods = append(methods, fmt.Sprintf("%d: method %s.%s", line, receiverName(d.Recv), d.Name.Name))
		}
	}
	appendSymbolGroup(out, "types", types)
	appendSymbolGroup(out, "funcs", funcs)
	appendSymbolGroup(out, "methods", methods)
	return nil
}

func appendSymbolGroup(out *strings.Builder, title string, values []string) {
	if len(values) == 0 {
		return
	}
	fmt.Fprintf(out, "  %s:\n", title)
	for _, value := range values {
		fmt.Fprintf(out, "    - %s\n", value)
	}
}

func typeKind(expr ast.Expr) string {
	switch expr.(type) {
	case *ast.StructType:
		return "struct"
	case *ast.InterfaceType:
		return "interface"
	case *ast.FuncType:
		return "func"
	default:
		return "alias"
	}
}

func receiverName(recv *ast.FieldList) string {
	if recv == nil || len(recv.List) == 0 {
		return "?"
	}
	var b bytes.Buffer
	if err := printer.Fprint(&b, token.NewFileSet(), recv.List[0].Type); err != nil {
		return "?"
	}
	return b.String()
}