Vectrum
Sign In
91 lines · 3.0 KB · Go
package tools

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

// 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, opts ...Option) Tool {
	o := applyOptions(opts)
	// Anything modified after this moment came into being during the
	// session — a scaffold's output, a file made with bash — and is not the
	// user's work in progress. Replacing it unread is not a destructive act
	// on content the model is guessing at; the model made it.
	sessionStart := time.Now()
	return New("write_file", o.describe(`
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. A file that existed before this
session must be read before it is replaced; one created during the session
(by a generator, or by you) may be replaced directly.`),
		func(ctx 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 — to files that predate the session.
				if !info.ModTime().After(sessionStart) {
					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"
			}
			summary := fmt.Sprintf("%s %s (%d bytes, %d lines).", verb, in.Path, len(content), countLines(content))
			return o.withDiagnostics(ctx, root, path, summary), nil
		})
}

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