edit: write an efficient text search-replace binary

This prevents streaming all the contents of a file from dagger.

Signed-off-by: Tibor Vass <teabee89@gmail.com>

Author: tiborvass

cmd/cu/stdio.go

@@ -1,10 +1,12 @@
 package main
 
 import (
+	"context"
 	"log/slog"
 	"os"
 
 	"dagger.io/dagger"
+	"github.com/dagger/container-use/environment"
 	"github.com/dagger/container-use/mcpserver"
 	"github.com/spf13/cobra"
 )
@@ -30,10 +32,15 @@ var stdioCmd = &cobra.Command{
 		}
 		defer dag.Close()
 
+		go warmCache(ctx, dag)
 		return mcpserver.RunStdioServer(ctx, dag)
 	},
 }
 
+func warmCache(ctx context.Context, dag *dagger.Client) {
+	environment.EditUtil(dag).Sync(ctx)
+}
+
 func init() {
 	rootCmd.AddCommand(stdioCmd)
 }

edit/cmd/edit/main.go

@@ -0,0 +1,30 @@
+package main
+
+import (
+	"fmt"
+	"io"
+	"os"
+	"strconv"
+
+	"github.com/tiborvass/replace"
+	"golang.org/x/text/transform"
+)
+
+func main() {
+	if len(os.Args)%3 != 0 || len(os.Args) < 6 {
+		fmt.Fprintf(os.Stderr, "usage: %s <source> <destination> <old_string1> <new_string1> <replace_count1> [...<old_stringN> <new_stringN> <replace_countN>]\n", os.Args[0])
+		fmt.Fprintln(os.Stderr, "       Reads stream from source and replaces in it replace_count times, old_string with new_string and writes to destination.")
+		fmt.Fprintln(os.Stderr, "       If replace_count is -1, it replaces all occurrences.")
+		os.Exit(1)
+	}
+	n := len(os.Args)/3 - 1
+	t := make([]transform.Transformer, n)
+	for i := range t {
+		replaceCount, err := strconv.Atoi(os.Args[5+i])
+		if err != nil {
+			fmt.Fprintf(os.Stderr, "replace_count must be an integer, received: %q\n", replaceCount)
+		}
+		t[i] = replace.StringN(os.Args[3+i], os.Args[4+i], replaceCount)
+	}
+	io.Copy(os.Stdout, replace.Chain(os.Stdin, t...))
+}

edit/edit.go

@@ -0,0 +1,6 @@
+package edit
+
+import _ "embed"
+
+//go:embed cmd/edit/main.go
+var Src string

environment/filesystem.go

@@ -2,15 +2,17 @@ package environment
 
 import (
 	"context"
+	_ "embed"
 	"fmt"
 	"strings"
 
 	"dagger.io/dagger"
 	"dagger.io/dagger/dag"
+	"github.com/dagger/container-use/edit"
 )
 
 // FIXME: See hack where it's used
-const fileEditBaseImage = "busybox"
+const fileUtilsBaseImage = "busybox"
 
 func (env *Environment) FileRead(ctx context.Context, targetFile string, shouldReadEntireFile bool, startLineOneIndexed int, endLineOneIndexedInclusive int) (string, error) {
 	file, err := env.container().File(targetFile).Contents(ctx)
@@ -82,18 +84,39 @@ func (env *Environment) FileGrep(ctx context.Context, path, pattern, include str
 	args := []string{"/bin/grep", "-E", "--", pattern, include}
 
 	dir := env.container().Rootfs().Directory(path)
-	out, err := dag.Container().From(fileEditBaseImage).WithMountedDirectory("/mnt", dir).WithWorkdir("/mnt").WithExec(args).Stdout(ctx)
+	out, err := dag.Container().From(fileUtilsBaseImage).WithMountedDirectory("/mnt", dir).WithWorkdir("/mnt").WithExec(args).Stdout(ctx)
 	if err != nil {
 		return "", err
 	}
 	return out, nil
 }
 
-func (env *Environment) FileEdit(ctx context.Context, targetFile string, edits []string) error {
+type FileEdit struct {
+	Old        string
+	New        string
+	ReplaceAll bool
+}
+
+func EditUtil(dag *dagger.Client) *dagger.Container {
+	editBin := dag.Container().From("golang:alpine").
+		WithNewFile("/edit.go", edit.Src).
+		WithEnvVariable("CGO_ENABLED", "0").
+		WithExec([]string{"go", "build", "-o", "/edit", "-ldflags", "-w -s", "/edit.go"}).File("/edit")
+	return dag.Container().From("scratch").WithFile("/edit", editBin).WithEntrypoint([]string{"/edit"})
+}
+
+func (env *Environment) FileEdit(ctx context.Context, dag *dagger.Client, targetFile string, edits []FileEdit) error {
 	// Hack: use busybox to run `sed` since dagger doesn't have native file editing primitives.
-	args := []string{"/bin/sh", "-c", fmt.Sprintf("sed -ri'' -- %s /target && cp /target /new", strings.Join(edits, " "))}
+	args := []string{"/edit", "/target", "/new"}
+	for _, edit := range edits {
+		replaceCount := "1"
+		if edit.ReplaceAll {
+			replaceCount = "-1"
+		}
+		args = append(args, edit.Old, edit.New, replaceCount)
+	}
 
-	newFile := dag.Container().From(fileEditBaseImage).WithMountedFile("/target", env.container().File(targetFile)).WithExec(args).File("/new")
+	newFile := EditUtil(dag).WithMountedFile("/target", env.container().File(targetFile)).WithExec(args).File("/new")
 	err := env.apply(ctx, env.container().WithFile(targetFile, newFile))
 	if err != nil {
 		return fmt.Errorf("failed applying file edit, skipping git propagation: %w", err)

go.mod

@@ -13,6 +13,8 @@ require (
 	github.com/spf13/cobra v1.9.1
 	github.com/stretchr/testify v1.10.0
 	github.com/tiborvass/go-watch v0.0.0-20250607214558-08999a83bf8b
+	github.com/tiborvass/replace v0.0.0-20250708165616-d642c0f9c3ff
+	golang.org/x/text v0.26.0
 )
 
 require (
@@ -55,7 +57,6 @@ require (
 	golang.org/x/sync v0.15.0 // indirect
 	golang.org/x/sys v0.33.0 // indirect
 	golang.org/x/term v0.32.0 // indirect
-	golang.org/x/text v0.26.0 // indirect
 	google.golang.org/genproto/googleapis/api v0.0.0-20250603155806-513f23925822 // indirect
 	google.golang.org/genproto/googleapis/rpc v0.0.0-20250603155806-513f23925822 // indirect
 	google.golang.org/grpc v1.73.0 // indirect

go.sum

@@ -27,6 +27,7 @@ github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag=
 github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE=
 github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek=
 github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps=
+github.com/google/go-cmp v0.3.0/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU=
 github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
 github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
 github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
@@ -46,6 +47,8 @@ github.com/mark3labs/mcp-go v0.29.0 h1:sH1NBcumKskhxqYzhXfGc201D7P76TVXiT0fGVhab
 github.com/mark3labs/mcp-go v0.29.0/go.mod h1:rXqOudj/djTORU/ThxYx8fqEVj/5pvTuuebQ2RC7uk4=
 github.com/mitchellh/go-homedir v1.1.0 h1:lukF9ziXFxDFPkA1vsr5zpc1XuPDn/wFntq5mG+4E0Y=
 github.com/mitchellh/go-homedir v1.1.0/go.mod h1:SfyaCUpYCn1Vlf4IUYiD9fPX4A5wJrkLzIz1N1q0pr0=
+github.com/pkg/errors v0.8.1 h1:iURUrRGxPUNPdy5/HRSm+Yj6okJ6UtLINN0Q9M4+h3I=
+github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pkg/term v1.1.0 h1:xIAAdCMh3QIAy+5FrE8Ad8XoDhEU4ufwbaSozViP9kk=
 github.com/pkg/term v1.1.0/go.mod h1:E25nymQcrSllhX42Ok8MRm1+hyBdHY0dCeiKZ9jpNGw=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
@@ -61,6 +64,7 @@ github.com/spf13/cast v1.7.1 h1:cuNEagBQEHWN1FnbGEjCXL2szYEXqfJPbP2HNUaca9Y=
 github.com/spf13/cast v1.7.1/go.mod h1:ancEpBxwJDODSW/UG4rDrAqiKolqNNh2DX3mk86cAdo=
 github.com/spf13/cobra v1.9.1 h1:CXSaggrXdbHK9CF+8ywj8Amf7PBRmPCOJugH954Nnlo=
 github.com/spf13/cobra v1.9.1/go.mod h1:nDyEzZ8ogv936Cinf6g1RU9MRY64Ir93oCnqb9wxYW0=
+github.com/spf13/pflag v1.0.3/go.mod h1:DYY7MBk1bdzusC3SYhjObp+wFpr4gzcvqqNjLnInEg4=
 github.com/spf13/pflag v1.0.6 h1:jFzHGLGAlb3ruxLB8MhbI6A8+AQX/2eW4qeyNZXNp2o=
 github.com/spf13/pflag v1.0.6/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
@@ -69,6 +73,8 @@ github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOf
 github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
 github.com/tiborvass/go-watch v0.0.0-20250607214558-08999a83bf8b h1:W24fsALOtQ9v3b0mK4yR8wrmhPx4lqJAMMJ+d338fqM=
 github.com/tiborvass/go-watch v0.0.0-20250607214558-08999a83bf8b/go.mod h1:oAWYkECp9mFVuJQQzHtoHhepQKbme1gLM4fYH0KWvzk=
+github.com/tiborvass/replace v0.0.0-20250708165616-d642c0f9c3ff h1:zpP7bpKTsqLLgsUngIXlmd1X9PpJD1Bca1dQfYacNf8=
+github.com/tiborvass/replace v0.0.0-20250708165616-d642c0f9c3ff/go.mod h1:9+jQ4zDLeiANhfwMbvl6qKw/sW+aN1m6AjhxHZKN40s=
 github.com/vektah/gqlparser/v2 v2.5.28 h1:bIulcl3LF69ba6EiZVGD88y4MkM+Jxrf3P2MX8xLRkY=
 github.com/vektah/gqlparser/v2 v2.5.28/go.mod h1:D1/VCZtV3LPnQrcPBeR/q5jkSQIPti0uYCP/RI0gIeo=
 github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4=
@@ -109,17 +115,25 @@ go.opentelemetry.io/proto/otlp v1.7.0 h1:jX1VolD6nHuFzOYso2E73H85i92Mv8JQYk0K9vz
 go.opentelemetry.io/proto/otlp v1.7.0/go.mod h1:fSKjH6YJ7HDlwzltzyMj036AJ3ejJLCgCSHGj4efDDo=
 go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto=
 go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE=
+golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
+golang.org/x/net v0.0.0-20190311183353-d8887717615a/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
 golang.org/x/net v0.41.0 h1:vBTly1HeNPEn3wtREYfy4GZ/NECgw2Cnl+nK6Nz3uvw=
 golang.org/x/net v0.41.0/go.mod h1:B/K4NNqkfmg07DQYrbwvSluqCJOOXwUjeb/5lOisjbA=
+golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.15.0 h1:KWH3jNZsfyT6xfAfKiz6MRNmd46ByHDYaZ7KSkCtdW8=
 golang.org/x/sync v0.15.0/go.mod h1:1dzgHSNfp02xaA81J2MS99Qcpr2w7fw1gpm99rleRqA=
+golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20200909081042-eff7692f9009/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.33.0 h1:q3i8TbbEz+JRD9ywIRlyRAQbM0qF7hu24q3teo2hbuw=
 golang.org/x/sys v0.33.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k=
 golang.org/x/term v0.32.0 h1:DR4lr0TjUs3epypdhTOkMmuF5CDFJ/8pOnbzMZPQ7bg=
 golang.org/x/term v0.32.0/go.mod h1:uZG1FhGx848Sqfsq4/DlJr3xGGsYMu/L5GW4abiaEPQ=
+golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
+golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
 golang.org/x/text v0.26.0 h1:P42AVeLghgTYr4+xUnTRKDMqpar+PtX7KWuNQL21L8M=
 golang.org/x/text v0.26.0/go.mod h1:QK15LZJUUQVJxhz7wXgxSy/CJaTFjd0G+YLonydOVQA=
+golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
+golang.org/x/tools v0.0.0-20190624222133-a101b041ded4/go.mod h1:/rFqwRUd4F7ZHNgwSSTFct+R/Kf4OFW1sUzUTQQTgfc=
 google.golang.org/genproto/googleapis/api v0.0.0-20250603155806-513f23925822 h1:oWVWY3NzT7KJppx2UKhKmzPq4SRe0LdCijVRwvGeikY=
 google.golang.org/genproto/googleapis/api v0.0.0-20250603155806-513f23925822/go.mod h1:h3c4v36UTKzUiuaOKQ6gr3S+0hovBtUrXzTG/i3+XEc=
 google.golang.org/genproto/googleapis/rpc v0.0.0-20250603155806-513f23925822 h1:fc6jSaCT0vBduLYZHYrBBNY4dsWuvgyff9noRNDdBeE=
@@ -136,3 +150,5 @@ gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
 gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
 gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
+gotest.tools/v3 v3.0.2 h1:kG1BFyqVHuQoVQiR1bWGnfz/fmHvvuiSPIV7rvl360E=
+gotest.tools/v3 v3.0.2/go.mod h1:3SzNCllyD9/Y+b5r9JIKQ474KzkZyqLqEfYqMsX94Bk=

mcpserver/tools.go

@@ -761,30 +761,25 @@ var EnvironmentFileEditTool = &Tool{
 			return mcp.NewToolResultErrorFromErr("unable to open the environment", err), nil
 		}
 
-		targetFile, err := request.RequireString("target_file")
-		if err != nil {
-			return nil, err
-		}
+		dag, _ := ctx.Value(daggerClientKey{}).(*dagger.Client)
 
-		args := request.GetArguments()
-		v, ok := args["edits"]
-		if !ok {
-			return nil, fmt.Errorf("could not find `edits` argument")
+		var args struct {
+			TargetFile string
+			Edits      []environment.FileEdit
 		}
-		edits, ok := v.([]string)
-		if !ok {
-			return nil, fmt.Errorf("`edits` argument is expected to be a []string")
+		if err := request.BindArguments(&args); err != nil {
+			return nil, fmt.Errorf("could not bind arguments")
 		}
 
-		if err := env.FileEdit(ctx, targetFile, edits); err != nil {
+		if err := env.FileEdit(ctx, dag, args.TargetFile, args.Edits); err != nil {
 			return mcp.NewToolResultErrorFromErr("failed to edit file", err), nil
 		}
 
 		if err := repo.Update(ctx, env, request.GetString("explanation", "")); err != nil {
 			return mcp.NewToolResultErrorFromErr("unable to update the environment", err), nil
 		}
 
-		return mcp.NewToolResultText(fmt.Sprintf("file %s edited successfully and committed to container-use/ remote", targetFile)), nil
+		return mcp.NewToolResultText(fmt.Sprintf("file %s edited successfully and committed to container-use/ remote", args.TargetFile)), nil
 	},
 }