bidirectional syncing to/from user worktree

environment/filesystem.go

@@ -2,8 +2,12 @@ package environment
 
 import (
 	"context"
+	"crypto/sha256"
 	"fmt"
 	"strings"
+
+	"dagger.io/dagger"
+	godiffpatch "github.com/sourcegraph/go-diff-patch"
 )
 
 func (env *Environment) FileRead(ctx context.Context, targetFile string, shouldReadEntireFile bool, startLineOneIndexedInclusive int, endLineOneIndexedInclusive int) (string, error) {
@@ -45,6 +49,86 @@ func (env *Environment) FileWrite(ctx context.Context, explanation, targetFile,
 	return nil
 }
 
+func (env *Environment) FileSearchReplace(ctx context.Context, explanation, targetFile, search, replace, matchID string) error {
+	contents, err := env.container().File(targetFile).Contents(ctx)
+	if err != nil {
+		return err
+	}
+
+	// Find all matches of the search text
+	matches := []int{}
+	searchIndex := 0
+	for {
+		index := strings.Index(contents[searchIndex:], search)
+		if index == -1 {
+			break
+		}
+		actualIndex := searchIndex + index
+		matches = append(matches, actualIndex)
+		searchIndex = actualIndex + 1
+	}
+
+	if len(matches) == 0 {
+		return fmt.Errorf("search text not found in file %s", targetFile)
+	}
+
+	// If there are multiple matches and no matchID is provided, return an error with all matches
+	if len(matches) > 1 && matchID == "" {
+		var matchDescriptions []string
+		for i, matchIndex := range matches {
+			// Generate a unique ID for each match
+			id := generateMatchID(targetFile, search, replace, i)
+
+			// Get context around the match (3 lines before and after)
+			context := getMatchContext(contents, matchIndex, len(search))
+
+			matchDescriptions = append(matchDescriptions, fmt.Sprintf("Match %d (ID: %s):\n%s", i+1, id, context))
+		}
+
+		return fmt.Errorf("multiple matches found for search text in %s. Please specify which_match parameter with one of the following IDs:\n\n%s",
+			targetFile, strings.Join(matchDescriptions, "\n\n"))
+	}
+
+	// Determine which match to replace
+	var targetMatchIndex int
+	if len(matches) == 1 {
+		targetMatchIndex = matches[0]
+	} else {
+		// Find the match with the specified ID
+		found := false
+		for i, matchIndex := range matches {
+			id := generateMatchID(targetFile, search, replace, i)
+			if id == matchID {
+				targetMatchIndex = matchIndex
+				found = true
+				break
+			}
+		}
+		if !found {
+			return fmt.Errorf("match ID %s not found", matchID)
+		}
+	}
+
+	// Replace the specific match
+	newContents := contents[:targetMatchIndex] + replace + contents[targetMatchIndex+len(search):]
+
+	// Apply the changes using `patch` so we don't have to spit out the entire
+	// contents
+	return env.ApplyPatch(ctx, godiffpatch.GeneratePatch(targetFile, contents, newContents))
+}
+
+func (env *Environment) ApplyPatch(ctx context.Context, patch string) error {
+	err := env.apply(ctx, env.container().
+		WithExec([]string{"patch", "-p1"}, dagger.ContainerWithExecOpts{
+			Stdin: patch,
+		}))
+	if err != nil {
+		return fmt.Errorf("failed applying file edit, skipping git propagation: %w", err)
+	}
+	env.Notes.Add("Apply patch")
+	return nil
+}
+
 func (env *Environment) FileDelete(ctx context.Context, explanation, targetFile string) error {
 	err := env.apply(ctx, env.container().WithoutFile(targetFile))
 	if err != nil {
@@ -65,3 +149,41 @@ func (env *Environment) FileList(ctx context.Context, path string) (string, erro
 	}
 	return out.String(), nil
 }
+
+// generateMatchID creates a unique ID for a match based on file, search, replace, and index
+func generateMatchID(targetFile, search, replace string, index int) string {
+	data := fmt.Sprintf("%s:%s:%s:%d", targetFile, search, replace, index)
+	hash := sha256.Sum256([]byte(data))
+	return fmt.Sprintf("%x", hash)[:8] // Use first 8 characters of hash
+}
+
+// getMatchContext returns the context around a match (3 lines before and after)
+func getMatchContext(contents string, matchIndex, matchLength int) string {
+	lines := strings.Split(contents, "\n")
+
+	// Find which line contains the match
+	currentPos := 0
+	matchLine := 0
+	for i, line := range lines {
+		if currentPos+len(line) >= matchIndex {
+			matchLine = i
+			break
+		}
+		currentPos += len(line) + 1 // +1 for newline
+	}
+
+	// Get context lines (3 before, match line, 3 after)
+	start := max(0, matchLine-3)
+	end := min(len(lines), matchLine+4)
+
+	contextLines := make([]string, 0, end-start)
+	for i := start; i < end; i++ {
+		prefix := "  "
+		if i == matchLine {
+			prefix = "> " // Mark the line containing the match
+		}
+		contextLines = append(contextLines, fmt.Sprintf("%s%s", prefix, lines[i]))
+	}
+
+	return strings.Join(contextLines, "\n")
+}

environment/state.go

@@ -10,9 +10,10 @@ type State struct {
 	CreatedAt time.Time `json:"created_at,omitempty"`
 	UpdatedAt time.Time `json:"updated_at,omitempty"`
 
-	Config    *EnvironmentConfig `json:"config,omitempty"`
-	Container string             `json:"container,omitempty"`
-	Title     string             `json:"title,omitempty"`
+	Config         *EnvironmentConfig `json:"config,omitempty"`
+	Container      string             `json:"container,omitempty"`
+	Title          string             `json:"title,omitempty"`
+	TrackingBranch string             `json:"tracking_branch,omitempty"`
 }
 
 func (s *State) Marshal() ([]byte, error) {

go.mod

@@ -56,6 +56,7 @@ require (
 	github.com/rogpeppe/go-internal v1.14.1 // indirect
 	github.com/sergi/go-diff v1.3.2-0.20230802210424-5b0b94c5c0d3 // indirect
 	github.com/sosodev/duration v1.3.1 // indirect
+	github.com/sourcegraph/go-diff-patch v0.0.0-20240223163233-798fd1e94a8e // indirect
 	github.com/spf13/cast v1.7.1 // indirect
 	github.com/spf13/pflag v1.0.6 // indirect
 	github.com/vektah/gqlparser/v2 v2.5.28 // indirect

go.sum

@@ -112,6 +112,8 @@ github.com/sergi/go-diff v1.3.2-0.20230802210424-5b0b94c5c0d3 h1:n661drycOFuPLCN
 github.com/sergi/go-diff v1.3.2-0.20230802210424-5b0b94c5c0d3/go.mod h1:A0bzQcvG0E7Rwjx0REVgAGH58e96+X0MeOfepqsbeW4=
 github.com/sosodev/duration v1.3.1 h1:qtHBDMQ6lvMQsL15g4aopM4HEfOaYuhWBw3NPTtlqq4=
 github.com/sosodev/duration v1.3.1/go.mod h1:RQIBBX0+fMLc/D9+Jb/fwvVmo0eZvDDEERAikUR6SDg=
+github.com/sourcegraph/go-diff-patch v0.0.0-20240223163233-798fd1e94a8e h1:H+jDTUeF+SVd4ApwnSFoew8ZwGNRfgb9EsZc7LcocAg=
+github.com/sourcegraph/go-diff-patch v0.0.0-20240223163233-798fd1e94a8e/go.mod h1:VsUklG6OQo7Ctunu0gS3AtEOCEc2kMB6r5rKzxAes58=
 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=

mcpserver/tools.go

@@ -10,6 +10,7 @@ import (
 	"log/slog"
 	"os"
 	"os/signal"
+	"strings"
 	"syscall"
 
 	"dagger.io/dagger"
@@ -130,18 +131,22 @@ func init() {
 		EnvironmentOpenTool,
 		EnvironmentCreateTool,
 		EnvironmentUpdateMetadataTool,
+		EnvironmentEnableTrackingTool,
 		EnvironmentConfigTool,
 
 		EnvironmentRunCmdTool,
 
 		EnvironmentFileReadTool,
 		EnvironmentFileListTool,
 		EnvironmentFileWriteTool,
+		EnvironmentFilePatchTool,
 		EnvironmentFileDeleteTool,
 
 		EnvironmentAddServiceTool,
 
 		EnvironmentCheckpointTool,
+
+		EnvironmentSyncFromUserTool,
 	)
 }
 
@@ -613,6 +618,73 @@ var EnvironmentFileWriteTool = &Tool{
 	},
 }
 
+var EnvironmentFilePatchTool = &Tool{
+	Definition: mcp.NewTool("environment_file_patch",
+		mcp.WithDescription("Find and replace text in a file."),
+		mcp.WithString("explanation",
+			mcp.Description("One sentence explanation for why this file is being edited."),
+		),
+		mcp.WithString("environment_source",
+			mcp.Description("Absolute path to the source git repository for the environment."),
+			mcp.Required(),
+		),
+		mcp.WithString("environment_id",
+			mcp.Description("The ID of the environment for this command. Must call `environment_create` first."),
+			mcp.Required(),
+		),
+		mcp.WithString("target_file",
+			mcp.Description("Path of the file to write, absolute or relative to the workdir."),
+			mcp.Required(),
+		),
+		mcp.WithString("search_text",
+			mcp.Description("The text to find and replace."),
+			mcp.Required(),
+		),
+		mcp.WithString("replace_text",
+			mcp.Description("The text to insert."),
+			mcp.Required(),
+		),
+		mcp.WithString("which_match",
+			mcp.Description("The ID of the match to replace, if there were multiple matches."),
+		),
+	),
+	Handler: func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+		repo, env, err := openEnvironment(ctx, request)
+		if err != nil {
+			return mcp.NewToolResultErrorFromErr("unable to open the environment", err), nil
+		}
+
+		targetFile, err := request.RequireString("target_file")
+		if err != nil {
+			return nil, err
+		}
+		search, err := request.RequireString("search_text")
+		if err != nil {
+			return nil, err
+		}
+		replace, err := request.RequireString("replace_text")
+		if err != nil {
+			return nil, err
+		}
+
+		if err := env.FileSearchReplace(ctx,
+			request.GetString("explanation", ""),
+			targetFile,
+			search,
+			replace,
+			request.GetString("which_match", ""),
+		); err != nil {
+			return mcp.NewToolResultErrorFromErr("failed to write 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
+	},
+}
+
 var EnvironmentFileDeleteTool = &Tool{
 	Definition: newEnvironmentTool(
 		"environment_file_delete",
@@ -756,3 +828,90 @@ Supported schemas are:
 		return mcp.NewToolResultText(fmt.Sprintf("Service added and started successfully: %s", string(output))), nil
 	},
 }
+
+var EnvironmentEnableTrackingTool = &Tool{
+	Definition: newEnvironmentTool(
+		"environment_enable_tracking",
+		"Enable branch tracking for an environment. When enabled, environment changes will be automatically synced to the user's working tree when on the tracked branch. CRITICAL: This is an opt-in feature that can only be enabled by explicit user request.",
+	),
+	Handler: func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+		repo, env, err := openEnvironment(ctx, request)
+		if err != nil {
+			return nil, err
+		}
+
+		// Get the current branch and tie it to an environment
+		currentBranch, err := repo.CurrentUserBranch(ctx)
+		if err != nil {
+			return nil, fmt.Errorf("failed to get current branch: %w", err)
+		}
+		if err := repo.TrackEnvironment(ctx, currentBranch, env.ID); err != nil {
+			return nil, fmt.Errorf("unable to set current branch tracking environment: %w", err)
+		}
+
+		// Set the tracking branch to the current branch
+		env.State.TrackingBranch = currentBranch
+
+		if err := repo.Update(ctx, env, request.GetString("explanation", "")); err != nil {
+			return nil, fmt.Errorf("unable to update the environment: %w", err)
+		}
+
+		out, err := marshalEnvironment(env)
+		if err != nil {
+			return nil, fmt.Errorf("failed to marshal environment: %w", err)
+		}
+		return mcp.NewToolResultText(fmt.Sprintf("Branch tracking enabled for branch '%s'. Environment changes will now be synced to the working tree when on this branch.\n%s", currentBranch, out)), nil
+	},
+}
+
+var EnvironmentSyncFromUserTool = &Tool{
+	Definition: newEnvironmentTool(
+		"environment_sync_from_user",
+		"Apply the user's unstaged changes to the environment and apply the environment's to the user's local worktree. ONLY RUN WHEN EXPLICITLY REQUESTED BY THE USER.",
+	),
+	Handler: func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+		repo, env, err := openEnvironment(ctx, request)
+		if err != nil {
+			return nil, err
+		}
+
+		currentBranch, err := repo.CurrentUserBranch(ctx)
+		if err != nil {
+			return nil, fmt.Errorf("failed to get current branch: %w", err)
+		}
+		branchEnv, err := repo.TrackedEnvironment(ctx, currentBranch)
+		if err != nil {
+			return nil, err
+		}
+		if branchEnv != env.ID {
+			return nil, fmt.Errorf("branch is tracking %s, not %s", branchEnv, env.ID)
+		}
+
+		patch, err := repo.DiffUserLocalChanges(ctx)
+		if err != nil {
+			return nil, fmt.Errorf("failed to generate patch: %w", err)
+		}
+		if len(patch) == 0 {
+			return mcp.NewToolResultText("No unstaged changes to pull."), nil
+		}
+
+		if err := env.ApplyPatch(ctx, patch); err != nil {
+			return nil, fmt.Errorf("failed to pull changes to environment: %w", err)
+		}
+
+		if err := repo.Update(ctx, env, request.GetString("explanation", "")); err != nil {
+			return nil, fmt.Errorf("unable to update the environment: %w", err)
+		}
+
+		if err := repo.ResetUserLocalChanges(ctx); err != nil {
+			return nil, fmt.Errorf("unable to reset user's worktree: %w", err)
+		}
+
+		var buf strings.Builder
+		if err := repo.Apply(ctx, env.ID, &buf); err != nil {
+			return nil, fmt.Errorf("unable to apply changes to user's worktree: %w\n\nlogs:\n%s", err, buf.String())
+		}
+
+		return mcp.NewToolResultText("Patch applied successfully to the environment:\n\n```patch\n" + string(patch) + "\n```"), nil
+	},
+}

repository/git.go

@@ -282,16 +282,25 @@ func (r *Repository) addGitNote(ctx context.Context, env *environment.Environmen
 	return r.propagateGitNotes(ctx, gitNotesLogRef)
 }
 
-func (r *Repository) currentUserBranch(ctx context.Context) (string, error) {
-	return RunGitCommand(ctx, r.userRepoPath, "branch", "--show-current")
+func (r *Repository) CurrentUserBranch(ctx context.Context) (string, error) {
+	currentBranch, err := RunGitCommand(ctx, r.userRepoPath, "branch", "--show-current")
+	if err != nil {
+		return "", err
+	}
+	// TODO(vito): pretty sure this is redundant, but consolidating from other
+	// places
+	branch := strings.TrimSpace(currentBranch)
+	if branch == "" {
+		return "", fmt.Errorf("no current branch (detached HEAD?)")
+	}
+	return branch, nil
 }
 
 func (r *Repository) mergeBase(ctx context.Context, env *environment.EnvironmentInfo) (string, error) {
-	currentBranch, err := r.currentUserBranch(ctx)
+	currentBranch, err := r.CurrentUserBranch(ctx)
 	if err != nil {
 		return "", err
 	}
-	currentBranch = strings.TrimSpace(currentBranch)
 	if currentBranch == "" {
 		currentBranch = "HEAD"
 	}