Vectrum
Sign In
80 lines · 2.3 KB · Go
package tools

import (
	"context"
	"fmt"
	"os"
	"path/filepath"
	"strings"
)

// WriteArgs is the write_file input.
type WriteArgs struct {
	Path    string `json:"path" jsonschema:"description=Path to write, absolute or relative to the working root"`
	Content string `json:"content" jsonschema:"description=Full contents of the file"`
}

// NewWrite builds the write_file tool. It is kept separate from edit_file so
// the two do not blur together in the metrics: a whole-file rewrite succeeding
// says nothing about whether targeted edits work.
func NewWrite(root string, track FileTracker) Tool {
	return New("write_file", strings.TrimSpace(`
Write a complete file, creating it or replacing it entirely. Use this for new
files. To change part of an existing file use edit_file instead — a full rewrite
discards anything you did not think to include.`),
		func(_ context.Context, in WriteArgs) (string, error) {
			path, err := resolve(root, in.Path)
			if err != nil {
				return "", err
			}

			info, statErr := os.Stat(path)
			existed := statErr == nil
			if existed && info.IsDir() {
				return "", fmt.Errorf("%s is a directory", in.Path)
			}

			content := in.Content
			perm := os.FileMode(0o644)
			if existed {
				// Overwriting is a destructive act on content the model may
				// only be guessing at, so the same read-before-write rule as
				// edit_file applies.
				if err := checkRead(track, path, in.Path); err != nil {
					return "", fmt.Errorf("%s (write_file replaces the whole file)", err)
				}
				perm = info.Mode().Perm()
				if prev, readErr := os.ReadFile(path); readErr == nil {
					if _, crlf := normalizeEndings(string(prev)); crlf {
						content = restoreEndings(normalizeInput(content), true)
					}
				}
			} else if dir := filepath.Dir(path); dir != "." {
				if err := os.MkdirAll(dir, 0o755); err != nil {
					return "", err
				}
			}

			if err := os.WriteFile(path, []byte(content), perm); err != nil {
				return "", err
			}
			if track != nil {
				if info, err := os.Stat(path); err == nil {
					track.MarkRead(path, info.ModTime())
				}
			}

			verb := "Created"
			if existed {
				verb = "Replaced"
			}
			return fmt.Sprintf("%s %s (%d bytes, %d lines).", verb, in.Path, len(content), countLines(content)), nil
		})
}

func countLines(s string) int {
	if s == "" {
		return 0
	}
	return len(splitLines(s))
}