From 14be8ef8a0fdc7a62909b22f76fa37df3e51101b Mon Sep 17 00:00:00 2001 From: Alex Suraci Date: Mon, 7 Jul 2025 13:28:49 -0400 Subject: [PATCH 01/13] wip: file search/replace (needs tests) --- environment/filesystem.go | 111 ++++++++++++++++++++++++++++++++++++++ mcpserver/tools.go | 68 +++++++++++++++++++++++ 2 files changed, 179 insertions(+) diff --git a/environment/filesystem.go b/environment/filesystem.go index 84ead062..292748fb 100644 --- a/environment/filesystem.go +++ b/environment/filesystem.go @@ -2,6 +2,7 @@ package environment import ( "context" + "crypto/sha256" "fmt" "strings" ) @@ -45,6 +46,78 @@ 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 + err = env.apply(ctx, env.container().WithNewFile(targetFile, newContents)) + if err != nil { + return fmt.Errorf("failed applying file edit, skipping git propagation: %w", err) + } + env.Notes.Add("Edit %s", targetFile) + 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 +138,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") +} diff --git a/mcpserver/tools.go b/mcpserver/tools.go index 1433ab24..d9b5e45d 100644 --- a/mcpserver/tools.go +++ b/mcpserver/tools.go @@ -137,6 +137,7 @@ func init() { EnvironmentFileReadTool, EnvironmentFileListTool, EnvironmentFileWriteTool, + EnvironmentFilePatchTool, EnvironmentFileDeleteTool, EnvironmentAddServiceTool, @@ -613,6 +614,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", From 6d5e9ab67578472730e831ebea097d649c37b761 Mon Sep 17 00:00:00 2001 From: Alex Suraci Date: Mon, 7 Jul 2025 19:06:13 -0400 Subject: [PATCH 02/13] use patch command instead of writing entire file IMO it's OK to depend on the environment's `patch` command; it's just as universal as `sh` which we already depend on, and it's heavily used in various Linux distro packaging systems. Signed-off-by: Alex Suraci --- environment/filesystem.go | 11 +++++++++-- go.mod | 1 + go.sum | 2 ++ 3 files changed, 12 insertions(+), 2 deletions(-) diff --git a/environment/filesystem.go b/environment/filesystem.go index 292748fb..fd9a41a5 100644 --- a/environment/filesystem.go +++ b/environment/filesystem.go @@ -5,6 +5,9 @@ import ( "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) { @@ -109,8 +112,12 @@ func (env *Environment) FileSearchReplace(ctx context.Context, explanation, targ // Replace the specific match newContents := contents[:targetMatchIndex] + replace + contents[targetMatchIndex+len(search):] - // Apply the changes - err = env.apply(ctx, env.container().WithNewFile(targetFile, newContents)) + // Apply the changes using `patch` so we don't have to spit out the entire + // contents + err = env.apply(ctx, env.container(). + WithExec([]string{"patch", "-p1"}, dagger.ContainerWithExecOpts{ + Stdin: godiffpatch.GeneratePatch(targetFile, contents, newContents), + })) if err != nil { return fmt.Errorf("failed applying file edit, skipping git propagation: %w", err) } diff --git a/go.mod b/go.mod index 461f6838..2b822be1 100644 --- a/go.mod +++ b/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 diff --git a/go.sum b/go.sum index 44d4f368..97133155 100644 --- a/go.sum +++ b/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= From 45037f6ca7401aeb98153a5b671ff4a39918413c Mon Sep 17 00:00:00 2001 From: Alex Suraci Date: Tue, 15 Jul 2025 16:59:10 -0400 Subject: [PATCH 03/13] apply: preserve local (unstaged) changes This change allows you to run `cu apply` continuously, by doing a somewhat delicate git dance: 1. `git diff` to save the unstaged changes to a .patch file 2. `git reset --hard` to get back to a pristine state 3. `git merge --squash` (no `--autostash`) to pull in the env changes 4. `git commit -m "temporary commit"` 5. `git apply` the patch from 1 (but this stages them) 6. `git reset` to move them back to unstaged 7. `git reset --soft HEAD~1` to move the temp commit into staging The spookiest part is probably 2, since that'll nuke any non-agent changes the user had staged. Maybe there's a safer way? Signed-off-by: Alex Suraci --- repository/repository.go | 108 ++++++++++++++++++++++++++++++++++++++- 1 file changed, 107 insertions(+), 1 deletion(-) diff --git a/repository/repository.go b/repository/repository.go index 81803f53..34ec6e4e 100644 --- a/repository/repository.go +++ b/repository/repository.go @@ -11,6 +11,7 @@ import ( "path/filepath" "sort" "strings" + "time" "dagger.io/dagger" "github.com/dagger/container-use/environment" @@ -420,5 +421,110 @@ func (r *Repository) Apply(ctx context.Context, id string, w io.Writer) error { return err } - return RunInteractiveGitCommand(ctx, r.userRepoPath, w, "merge", "--autostash", "--squash", "--", "container-use/"+envInfo.ID) + // Create patch directory if it doesn't exist + configPath := os.ExpandEnv("$HOME/.config/container-use") + patchDir := filepath.Join(configPath, "patches") + if err := os.MkdirAll(patchDir, 0755); err != nil { + return fmt.Errorf("failed to create patch directory: %w", err) + } + + // Create a unique patch filename using timestamp and environment ID + patchFile := filepath.Join(patchDir, fmt.Sprintf("user-changes-%s-%d.patch", envInfo.ID, time.Now().Unix())) + + // Check if there are any unstaged changes + diffCmd := exec.CommandContext(ctx, "git", "diff") + diffCmd.Dir = r.userRepoPath + diffOutput, err := diffCmd.Output() + if err != nil { + return fmt.Errorf("failed to check for unstaged changes: %w", err) + } + + hasUnstagedChanges := len(diffOutput) > 0 + + if hasUnstagedChanges { + // Create a patch of only unstaged changes + fmt.Fprintf(w, "Saving unstaged user changes to %s...\n", patchFile) + + // Create the patch from unstaged changes only + patchCmd := exec.CommandContext(ctx, "git", "diff") + patchCmd.Dir = r.userRepoPath + patchOutput, err := patchCmd.Output() + if err != nil { + return fmt.Errorf("failed to create patch: %w", err) + } + + // Write patch to file + if err := os.WriteFile(patchFile, patchOutput, 0644); err != nil { + return fmt.Errorf("failed to write patch file: %w", err) + } + + // Reset to clean state + fmt.Fprintf(w, "Resetting to clean state...\n") + if err := RunInteractiveGitCommand(ctx, r.userRepoPath, w, "reset", "--hard", "HEAD"); err != nil { + return fmt.Errorf("failed to reset: %w", err) + } + } + + // Apply the merge without autostash + fmt.Fprintf(w, "Applying environment changes...\n") + if err := RunInteractiveGitCommand(ctx, r.userRepoPath, w, "merge", "--squash", "--", "container-use/"+envInfo.ID); err != nil { + // If merge fails, try to restore user changes + if hasUnstagedChanges { + fmt.Fprintf(w, "Merge failed, restoring user changes...\n") + applyCmd := exec.CommandContext(ctx, "git", "apply", patchFile) + applyCmd.Dir = r.userRepoPath + applyCmd.Stdout = w + applyCmd.Stderr = w + applyCmd.Run() // Ignore error as patch might partially apply + } + return fmt.Errorf("failed to merge: %w", err) + } + + // Apply user changes back + if hasUnstagedChanges { + fmt.Fprintf(w, "Restoring user changes...\n") + + // 1. Temporarily commit the agent's changes + commitCmd := exec.CommandContext(ctx, "git", "commit", "-m", "temp: agent changes") + commitCmd.Dir = r.userRepoPath + if err := commitCmd.Run(); err != nil { + fmt.Fprintf(w, "Warning: Failed to commit agent changes: %v\n", err) + return nil + } + + // 2. Apply the user's patch + applyCmd := exec.CommandContext(ctx, "git", "apply", patchFile) + applyCmd.Dir = r.userRepoPath + applyCmd.Stdout = w + applyCmd.Stderr = w + if err := applyCmd.Run(); err != nil { + fmt.Fprintf(w, "Warning: Failed to apply some user changes. Patch saved at: %s\n", patchFile) + fmt.Fprintf(w, "You can manually apply it with: git apply %s\n", patchFile) + // Try to recover by doing soft reset + resetCmd := exec.CommandContext(ctx, "git", "reset", "--soft", "HEAD~1") + resetCmd.Dir = r.userRepoPath + resetCmd.Run() + return nil + } + + // 3. Reset to unstage everything + resetCmd := exec.CommandContext(ctx, "git", "reset") + resetCmd.Dir = r.userRepoPath + if err := resetCmd.Run(); err != nil { + fmt.Fprintf(w, "Warning: Failed to reset: %v\n", err) + } + + // 4. Soft reset to bring agent changes back to staging + softResetCmd := exec.CommandContext(ctx, "git", "reset", "--soft", "HEAD~1") + softResetCmd.Dir = r.userRepoPath + if err := softResetCmd.Run(); err != nil { + fmt.Fprintf(w, "Warning: Failed to restore agent changes to staging: %v\n", err) + } + + // Clean up patch file on successful application + os.Remove(patchFile) + fmt.Fprintf(w, "User changes successfully restored as unstaged changes.\n") + } + + return nil } From 962825b75b0c8c7126d552180142f75c5613cd92 Mon Sep 17 00:00:00 2001 From: Alex Suraci Date: Tue, 15 Jul 2025 19:04:14 -0400 Subject: [PATCH 04/13] support syncing changes from user * new `environment_sync_from_user` tool: applies user's unstaged changes to the env, commits, syncs the env back to the user (bidirectional) * `git reset --hard` is now done after creating a refless stash, so the user can bring everything back if something goes wrong * add `Environment.ApplyPatch`, which currently depends on `patch` being available in the environment * clean up all the git code that wasn't using helpers Signed-off-by: Alex Suraci --- environment/filesystem.go | 10 ++-- mcpserver/tools.go | 44 ++++++++++++++++ repository/repository.go | 108 ++++++++++++++------------------------ 3 files changed, 91 insertions(+), 71 deletions(-) diff --git a/environment/filesystem.go b/environment/filesystem.go index fd9a41a5..2a3cb4f6 100644 --- a/environment/filesystem.go +++ b/environment/filesystem.go @@ -114,14 +114,18 @@ func (env *Environment) FileSearchReplace(ctx context.Context, explanation, targ // Apply the changes using `patch` so we don't have to spit out the entire // contents - err = env.apply(ctx, env.container(). + 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: godiffpatch.GeneratePatch(targetFile, contents, newContents), + Stdin: patch, })) if err != nil { return fmt.Errorf("failed applying file edit, skipping git propagation: %w", err) } - env.Notes.Add("Edit %s", targetFile) + env.Notes.Add("Apply patch") return nil } diff --git a/mcpserver/tools.go b/mcpserver/tools.go index d9b5e45d..2f72b240 100644 --- a/mcpserver/tools.go +++ b/mcpserver/tools.go @@ -10,6 +10,7 @@ import ( "log/slog" "os" "os/signal" + "strings" "syscall" "dagger.io/dagger" @@ -143,6 +144,8 @@ func init() { EnvironmentAddServiceTool, EnvironmentCheckpointTool, + + EnvironmentSyncFromUserTool, ) } @@ -824,3 +827,44 @@ Supported schemas are: return mcp.NewToolResultText(fmt.Sprintf("Service added and started successfully: %s", string(output))), 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 + } + + // Use a string builder to capture output + 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 + }, +} diff --git a/repository/repository.go b/repository/repository.go index 34ec6e4e..a348caf3 100644 --- a/repository/repository.go +++ b/repository/repository.go @@ -11,7 +11,6 @@ import ( "path/filepath" "sort" "strings" - "time" "dagger.io/dagger" "github.com/dagger/container-use/environment" @@ -415,68 +414,39 @@ func (r *Repository) Merge(ctx context.Context, id string, w io.Writer) error { return RunInteractiveGitCommand(ctx, r.userRepoPath, w, "merge", "--no-ff", "--autostash", "-m", "Merge environment "+envInfo.ID, "--", "container-use/"+envInfo.ID) } -func (r *Repository) Apply(ctx context.Context, id string, w io.Writer) error { +func (r *Repository) Apply(ctx context.Context, id string, w io.Writer) (rerr error) { envInfo, err := r.Info(ctx, id) if err != nil { return err } - // Create patch directory if it doesn't exist - configPath := os.ExpandEnv("$HOME/.config/container-use") - patchDir := filepath.Join(configPath, "patches") - if err := os.MkdirAll(patchDir, 0755); err != nil { - return fmt.Errorf("failed to create patch directory: %w", err) - } - - // Create a unique patch filename using timestamp and environment ID - patchFile := filepath.Join(patchDir, fmt.Sprintf("user-changes-%s-%d.patch", envInfo.ID, time.Now().Unix())) - - // Check if there are any unstaged changes - diffCmd := exec.CommandContext(ctx, "git", "diff") - diffCmd.Dir = r.userRepoPath - diffOutput, err := diffCmd.Output() + diffOutput, err := r.DiffUserLocalChanges(ctx) if err != nil { return fmt.Errorf("failed to check for unstaged changes: %w", err) } hasUnstagedChanges := len(diffOutput) > 0 - if hasUnstagedChanges { - // Create a patch of only unstaged changes - fmt.Fprintf(w, "Saving unstaged user changes to %s...\n", patchFile) - - // Create the patch from unstaged changes only - patchCmd := exec.CommandContext(ctx, "git", "diff") - patchCmd.Dir = r.userRepoPath - patchOutput, err := patchCmd.Output() - if err != nil { - return fmt.Errorf("failed to create patch: %w", err) - } - - // Write patch to file - if err := os.WriteFile(patchFile, patchOutput, 0644); err != nil { - return fmt.Errorf("failed to write patch file: %w", err) + fmt.Fprintf(w, "Creating virtual stash as backup...\n") + stashID, err := RunGitCommand(ctx, r.userRepoPath, "stash", "create") + if err != nil { + return fmt.Errorf("failed to stash changes: %w", err) + } + defer func() { + if rerr != nil { + fmt.Fprintf(w, "ERROR: %s\n", rerr) + fmt.Fprintf(w, "Your prior changes can be restored with `git stash apply %s`\n", stashID) } + }() - // Reset to clean state - fmt.Fprintf(w, "Resetting to clean state...\n") - if err := RunInteractiveGitCommand(ctx, r.userRepoPath, w, "reset", "--hard", "HEAD"); err != nil { - return fmt.Errorf("failed to reset: %w", err) - } + // Reset to clean state + if err := RunInteractiveGitCommand(ctx, r.userRepoPath, w, "reset", "--hard", "HEAD"); err != nil { + return fmt.Errorf("failed to reset: %w", err) } // Apply the merge without autostash fmt.Fprintf(w, "Applying environment changes...\n") if err := RunInteractiveGitCommand(ctx, r.userRepoPath, w, "merge", "--squash", "--", "container-use/"+envInfo.ID); err != nil { - // If merge fails, try to restore user changes - if hasUnstagedChanges { - fmt.Fprintf(w, "Merge failed, restoring user changes...\n") - applyCmd := exec.CommandContext(ctx, "git", "apply", patchFile) - applyCmd.Dir = r.userRepoPath - applyCmd.Stdout = w - applyCmd.Stderr = w - applyCmd.Run() // Ignore error as patch might partially apply - } return fmt.Errorf("failed to merge: %w", err) } @@ -485,46 +455,48 @@ func (r *Repository) Apply(ctx context.Context, id string, w io.Writer) error { fmt.Fprintf(w, "Restoring user changes...\n") // 1. Temporarily commit the agent's changes - commitCmd := exec.CommandContext(ctx, "git", "commit", "-m", "temp: agent changes") - commitCmd.Dir = r.userRepoPath - if err := commitCmd.Run(); err != nil { - fmt.Fprintf(w, "Warning: Failed to commit agent changes: %v\n", err) - return nil + if err := RunInteractiveGitCommand(ctx, r.userRepoPath, w, "commit", "-m", "temp: agent changes"); err != nil { + return fmt.Errorf("failed to commit agent changes: %w", err) } // 2. Apply the user's patch - applyCmd := exec.CommandContext(ctx, "git", "apply", patchFile) + applyCmd := exec.CommandContext(ctx, "git", "apply", "-") applyCmd.Dir = r.userRepoPath + applyCmd.Stdin = strings.NewReader(diffOutput) applyCmd.Stdout = w applyCmd.Stderr = w if err := applyCmd.Run(); err != nil { - fmt.Fprintf(w, "Warning: Failed to apply some user changes. Patch saved at: %s\n", patchFile) - fmt.Fprintf(w, "You can manually apply it with: git apply %s\n", patchFile) - // Try to recover by doing soft reset - resetCmd := exec.CommandContext(ctx, "git", "reset", "--soft", "HEAD~1") - resetCmd.Dir = r.userRepoPath - resetCmd.Run() - return nil + return fmt.Errorf("failed to apply user changes: %w", err) } - // 3. Reset to unstage everything - resetCmd := exec.CommandContext(ctx, "git", "reset") - resetCmd.Dir = r.userRepoPath - if err := resetCmd.Run(); err != nil { - fmt.Fprintf(w, "Warning: Failed to reset: %v\n", err) + // 3. Reset to unstage the user's changes + if err := RunInteractiveGitCommand(ctx, r.userRepoPath, w, "reset"); err != nil { + return fmt.Errorf("failed to reset user changes: %w", err) } // 4. Soft reset to bring agent changes back to staging - softResetCmd := exec.CommandContext(ctx, "git", "reset", "--soft", "HEAD~1") - softResetCmd.Dir = r.userRepoPath - if err := softResetCmd.Run(); err != nil { - fmt.Fprintf(w, "Warning: Failed to restore agent changes to staging: %v\n", err) + if err := RunInteractiveGitCommand(ctx, r.userRepoPath, w, "reset", "--soft", "HEAD~1"); err != nil { + return fmt.Errorf("failed to restore agent changes to staging: %w", err) } // Clean up patch file on successful application - os.Remove(patchFile) fmt.Fprintf(w, "User changes successfully restored as unstaged changes.\n") } return nil } + +func (r *Repository) DiffUserLocalChanges(ctx context.Context) (string, error) { + diff, err := RunGitCommand(ctx, r.userRepoPath, "diff") + if err != nil { + return "", fmt.Errorf("failed to get user diff: %w", err) + } + return diff, nil +} + +func (r *Repository) ResetUserLocalChanges(ctx context.Context) error { + if _, err := RunGitCommand(ctx, r.userRepoPath, "restore", "."); err != nil { + return fmt.Errorf("failed to reset unstaged changes: %w", err) + } + return nil +} From bbecb79384692f4b0f5d8c4420fd1294b7624c06 Mon Sep 17 00:00:00 2001 From: Alex Suraci Date: Tue, 15 Jul 2025 23:10:24 -0400 Subject: [PATCH 05/13] wip: opt environment in to tracking a branch Signed-off-by: Alex Suraci --- environment/state.go | 7 ++++--- mcpserver/tools.go | 37 +++++++++++++++++++++++++++++++++++++ repository/repository.go | 19 +++++++++++++++++++ 3 files changed, 60 insertions(+), 3 deletions(-) diff --git a/environment/state.go b/environment/state.go index ed7e2d4a..85cba15a 100644 --- a/environment/state.go +++ b/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) { diff --git a/mcpserver/tools.go b/mcpserver/tools.go index 2f72b240..917ada35 100644 --- a/mcpserver/tools.go +++ b/mcpserver/tools.go @@ -131,6 +131,7 @@ func init() { EnvironmentOpenTool, EnvironmentCreateTool, EnvironmentUpdateMetadataTool, + EnvironmentEnableTrackingTool, EnvironmentConfigTool, EnvironmentRunCmdTool, @@ -312,6 +313,42 @@ var EnvironmentUpdateMetadataTool = &Tool{ }, } +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 from the user's repository + currentBranch, err := repository.RunGitCommand(ctx, repo.SourcePath(), "branch", "--show-current") + if err != nil { + return nil, fmt.Errorf("unable to determine current branch: %w", err) + } + currentBranch = strings.TrimSpace(currentBranch) + if currentBranch == "" { + return nil, fmt.Errorf("not on a branch (detached HEAD state) - cannot enable tracking") + } + + // 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 EnvironmentConfigTool = &Tool{ Definition: newEnvironmentTool( "environment_config", diff --git a/repository/repository.go b/repository/repository.go index a348caf3..9bb96940 100644 --- a/repository/repository.go +++ b/repository/repository.go @@ -283,6 +283,25 @@ func (r *Repository) Update(ctx context.Context, env *environment.Environment, e if err := r.propagateToWorktree(ctx, env, explanation); err != nil { return err } + + // Check if branch tracking is enabled and we're on the tracked branch + if env.State.TrackingBranch != "" { + currentBranch, err := RunGitCommand(ctx, r.userRepoPath, "branch", "--show-current") + if err != nil { + // Log the error but don't fail the update + slog.Warn("Failed to check current branch for tracking", "error", err) + } else { + currentBranch = strings.TrimSpace(currentBranch) + if currentBranch == env.State.TrackingBranch { + // Apply environment changes to the user's working tree + if err := r.Apply(ctx, env.ID, io.Discard); err != nil { + // Log the error but don't fail the update to avoid breaking the environment + slog.Warn("Failed to apply tracking changes to working tree", "error", err, "branch", currentBranch) + } + } + } + } + if note := env.Notes.Pop(); note != "" { return r.addGitNote(ctx, env, note) } From bcab059906ac26d234e4de39f73ac8a6215c92e6 Mon Sep 17 00:00:00 2001 From: Alex Suraci Date: Wed, 16 Jul 2025 13:54:07 -0400 Subject: [PATCH 06/13] prevent syncing branch to wrong environment I hit a scenario where Claude compacted the message history and then created a new environment. It tried to sync my local changes to 'recover' the progress, which is unfortunate since it's only meant to do that when the user explicitly requests it. So, this adds a layer of defensiveness: now when you enable tracking for a branch, we also track which environment the branch is tracking, and error early if there's a mismatch. Signed-off-by: Alex Suraci --- mcpserver/tools.go | 83 +++++++++++++++++++++++----------------- repository/git.go | 17 ++++++-- repository/repository.go | 20 ++++++++++ 3 files changed, 80 insertions(+), 40 deletions(-) diff --git a/mcpserver/tools.go b/mcpserver/tools.go index 917ada35..5767737d 100644 --- a/mcpserver/tools.go +++ b/mcpserver/tools.go @@ -313,42 +313,6 @@ var EnvironmentUpdateMetadataTool = &Tool{ }, } -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 from the user's repository - currentBranch, err := repository.RunGitCommand(ctx, repo.SourcePath(), "branch", "--show-current") - if err != nil { - return nil, fmt.Errorf("unable to determine current branch: %w", err) - } - currentBranch = strings.TrimSpace(currentBranch) - if currentBranch == "" { - return nil, fmt.Errorf("not on a branch (detached HEAD state) - cannot enable tracking") - } - - // 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 EnvironmentConfigTool = &Tool{ Definition: newEnvironmentTool( "environment_config", @@ -865,6 +829,41 @@ Supported schemas are: }, } +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", @@ -876,6 +875,18 @@ var EnvironmentSyncFromUserTool = &Tool{ 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) + } + // Use a string builder to capture output patch, err := repo.DiffUserLocalChanges(ctx) if err != nil { diff --git a/repository/git.go b/repository/git.go index 5bae7ced..1f9695ab 100644 --- a/repository/git.go +++ b/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" } diff --git a/repository/repository.go b/repository/repository.go index 9bb96940..9b32e65d 100644 --- a/repository/repository.go +++ b/repository/repository.go @@ -513,6 +513,26 @@ func (r *Repository) DiffUserLocalChanges(ctx context.Context) (string, error) { return diff, nil } +func (r *Repository) TrackEnvironment(ctx context.Context, branch, envID string) error { + _, err := RunGitCommand(ctx, r.userRepoPath, "config", "branch."+branch+".environment", envID) + if err != nil { + return fmt.Errorf("failed to set branch tracking env: %w", err) + } + return nil +} + +func (r *Repository) TrackedEnvironment(ctx context.Context, branch string) (string, error) { + envID, err := RunGitCommand(ctx, r.userRepoPath, "config", "get", "--default=", "branch."+branch+".environment") + if err != nil { + return "", fmt.Errorf("failed to get branch tracking env: %w", err) + } + envID = strings.TrimSpace(envID) + if envID == "" { + return "", fmt.Errorf("branch %s is not tracking an environment", branch) + } + return envID, nil +} + func (r *Repository) ResetUserLocalChanges(ctx context.Context) error { if _, err := RunGitCommand(ctx, r.userRepoPath, "restore", "."); err != nil { return fmt.Errorf("failed to reset unstaged changes: %w", err) From 8db1008ae4c1c392e207242bf8608b1ee312e6c1 Mon Sep 17 00:00:00 2001 From: Alex Suraci Date: Wed, 16 Jul 2025 13:56:16 -0400 Subject: [PATCH 07/13] don't choke on files marked binary Signed-off-by: Alex Suraci --- repository/repository.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/repository/repository.go b/repository/repository.go index 9b32e65d..4a0a6292 100644 --- a/repository/repository.go +++ b/repository/repository.go @@ -506,7 +506,7 @@ func (r *Repository) Apply(ctx context.Context, id string, w io.Writer) (rerr er } func (r *Repository) DiffUserLocalChanges(ctx context.Context) (string, error) { - diff, err := RunGitCommand(ctx, r.userRepoPath, "diff") + diff, err := RunGitCommand(ctx, r.userRepoPath, "diff", "--binary") if err != nil { return "", fmt.Errorf("failed to get user diff: %w", err) } From 207fc1929e61f59730bc2085b90eceb2a814b729 Mon Sep 17 00:00:00 2001 From: Alex Suraci Date: Wed, 16 Jul 2025 13:56:47 -0400 Subject: [PATCH 08/13] error when env->user sync fails Signed-off-by: Alex Suraci --- mcpserver/tools.go | 1 - repository/repository.go | 12 +++++------- 2 files changed, 5 insertions(+), 8 deletions(-) diff --git a/mcpserver/tools.go b/mcpserver/tools.go index 5767737d..7bf4add5 100644 --- a/mcpserver/tools.go +++ b/mcpserver/tools.go @@ -887,7 +887,6 @@ var EnvironmentSyncFromUserTool = &Tool{ return nil, fmt.Errorf("branch is tracking %s, not %s", branchEnv, env.ID) } - // Use a string builder to capture output patch, err := repo.DiffUserLocalChanges(ctx) if err != nil { return nil, fmt.Errorf("failed to generate patch: %w", err) diff --git a/repository/repository.go b/repository/repository.go index 4a0a6292..86f473d0 100644 --- a/repository/repository.go +++ b/repository/repository.go @@ -286,17 +286,15 @@ func (r *Repository) Update(ctx context.Context, env *environment.Environment, e // Check if branch tracking is enabled and we're on the tracked branch if env.State.TrackingBranch != "" { - currentBranch, err := RunGitCommand(ctx, r.userRepoPath, "branch", "--show-current") + currentBranch, err := r.CurrentUserBranch(ctx) if err != nil { - // Log the error but don't fail the update - slog.Warn("Failed to check current branch for tracking", "error", err) + return fmt.Errorf("failed to check current branch for tracking: %w", err) } else { - currentBranch = strings.TrimSpace(currentBranch) if currentBranch == env.State.TrackingBranch { // Apply environment changes to the user's working tree - if err := r.Apply(ctx, env.ID, io.Discard); err != nil { - // Log the error but don't fail the update to avoid breaking the environment - slog.Warn("Failed to apply tracking changes to working tree", "error", err, "branch", currentBranch) + var logs strings.Builder + if err := r.Apply(ctx, env.ID, &logs); err != nil { + return fmt.Errorf("failed to apply tracking changes to working tree: %w\n\nlogs:\n%s\n", err, logs.String()) } } } From 856d418004c452a5d69cd9eb65489d1f7c81ff98 Mon Sep 17 00:00:00 2001 From: Alex Suraci Date: Wed, 16 Jul 2025 13:56:52 -0400 Subject: [PATCH 09/13] use 3-way merge when applying user changes this supports the following flow: 1. made local edits, e.g. to write a failing test 2. create a tracking environment and tell it to sync my changes previously this would fail even though the patch is a no-op; now we'll use a 3-way merge, which accepts no-ops, and we'll --check first to avoid leaving conflict markers Signed-off-by: Alex Suraci --- repository/repository.go | 25 +++++++++++++++++++++---- 1 file changed, 21 insertions(+), 4 deletions(-) diff --git a/repository/repository.go b/repository/repository.go index 86f473d0..873e7933 100644 --- a/repository/repository.go +++ b/repository/repository.go @@ -471,13 +471,30 @@ func (r *Repository) Apply(ctx context.Context, id string, w io.Writer) (rerr er if hasUnstagedChanges { fmt.Fprintf(w, "Restoring user changes...\n") - // 1. Temporarily commit the agent's changes - if err := RunInteractiveGitCommand(ctx, r.userRepoPath, w, "commit", "-m", "temp: agent changes"); err != nil { + // 1. Temporarily commit the agent's changes (if any) + if err := RunInteractiveGitCommand(ctx, r.userRepoPath, w, + "commit", + // it's simpler/safer to just keep this path unconditional than check if there + // were staged stuff before, so just allow an empty commit; we'll be + // getting rid of it immediately anyway + "--allow-empty", + "-m", "temp: agent changes (you should not see this)", + ); err != nil { return fmt.Errorf("failed to commit agent changes: %w", err) } - // 2. Apply the user's patch - applyCmd := exec.CommandContext(ctx, "git", "apply", "-") + // 2. Apply the user's patch, using 3-way merge with --check beforehand to + // avoid leaving conflict markers + var checkOut strings.Builder + checkCmd := exec.CommandContext(ctx, "git", "apply", "--3way", "--check", "-") + checkCmd.Dir = r.userRepoPath + checkCmd.Stdin = strings.NewReader(diffOutput) + checkCmd.Stdout = &checkOut + checkCmd.Stderr = &checkOut + if err := checkCmd.Run(); err != nil { + return fmt.Errorf("conflict detected when re-applying user changes:\n%s", checkOut.String()) + } + applyCmd := exec.CommandContext(ctx, "git", "apply", "--3way", "-") applyCmd.Dir = r.userRepoPath applyCmd.Stdin = strings.NewReader(diffOutput) applyCmd.Stdout = w From b362d18e7f1e448dc9579f8b5e7d4e18a5a27fd9 Mon Sep 17 00:00:00 2001 From: Alex Suraci Date: Wed, 16 Jul 2025 20:32:30 -0400 Subject: [PATCH 10/13] Revert "use 3-way merge when applying user changes" this is still not robust enough to cover the scenario where the user deleted a file. we should just do a hard reset to the environment's changes instead, and not attempt to restore their changes, since we just synced them over. This reverts commit 856d418004c452a5d69cd9eb65489d1f7c81ff98. --- repository/repository.go | 25 ++++--------------------- 1 file changed, 4 insertions(+), 21 deletions(-) diff --git a/repository/repository.go b/repository/repository.go index 873e7933..86f473d0 100644 --- a/repository/repository.go +++ b/repository/repository.go @@ -471,30 +471,13 @@ func (r *Repository) Apply(ctx context.Context, id string, w io.Writer) (rerr er if hasUnstagedChanges { fmt.Fprintf(w, "Restoring user changes...\n") - // 1. Temporarily commit the agent's changes (if any) - if err := RunInteractiveGitCommand(ctx, r.userRepoPath, w, - "commit", - // it's simpler/safer to just keep this path unconditional than check if there - // were staged stuff before, so just allow an empty commit; we'll be - // getting rid of it immediately anyway - "--allow-empty", - "-m", "temp: agent changes (you should not see this)", - ); err != nil { + // 1. Temporarily commit the agent's changes + if err := RunInteractiveGitCommand(ctx, r.userRepoPath, w, "commit", "-m", "temp: agent changes"); err != nil { return fmt.Errorf("failed to commit agent changes: %w", err) } - // 2. Apply the user's patch, using 3-way merge with --check beforehand to - // avoid leaving conflict markers - var checkOut strings.Builder - checkCmd := exec.CommandContext(ctx, "git", "apply", "--3way", "--check", "-") - checkCmd.Dir = r.userRepoPath - checkCmd.Stdin = strings.NewReader(diffOutput) - checkCmd.Stdout = &checkOut - checkCmd.Stderr = &checkOut - if err := checkCmd.Run(); err != nil { - return fmt.Errorf("conflict detected when re-applying user changes:\n%s", checkOut.String()) - } - applyCmd := exec.CommandContext(ctx, "git", "apply", "--3way", "-") + // 2. Apply the user's patch + applyCmd := exec.CommandContext(ctx, "git", "apply", "-") applyCmd.Dir = r.userRepoPath applyCmd.Stdin = strings.NewReader(diffOutput) applyCmd.Stdout = w From 77a2330d8a30b5e0efbbef9c03ff45c5074834e2 Mon Sep 17 00:00:00 2001 From: Alex Suraci Date: Wed, 16 Jul 2025 20:35:15 -0400 Subject: [PATCH 11/13] discard local changes when re-syncing to env we literally just migrated their unstaged changes to the env, so this should be equivalent and much more foolproof Signed-off-by: Alex Suraci --- cmd/container-use/apply.go | 2 +- environment/integration/merge_test.go | 6 +++--- mcpserver/tools.go | 11 +++-------- repository/repository.go | 6 +++--- 4 files changed, 10 insertions(+), 15 deletions(-) diff --git a/cmd/container-use/apply.go b/cmd/container-use/apply.go index bac9ae7b..cf440eba 100644 --- a/cmd/container-use/apply.go +++ b/cmd/container-use/apply.go @@ -43,7 +43,7 @@ git commit -m "Add backend API implementation"`, env := args[0] - if err := repo.Apply(ctx, env, os.Stdout); err != nil { + if err := repo.Apply(ctx, env, os.Stdout, false); err != nil { return fmt.Errorf("failed to apply environment: %w", err) } diff --git a/environment/integration/merge_test.go b/environment/integration/merge_test.go index f69347a0..18c90224 100644 --- a/environment/integration/merge_test.go +++ b/environment/integration/merge_test.go @@ -79,7 +79,7 @@ func TestRepositoryApply(t *testing.T) { // Apply the environment (squash merge) var applyOutput bytes.Buffer - err = repo.Apply(ctx, env.ID, &applyOutput) + err = repo.Apply(ctx, env.ID, &applyOutput, true) require.NoError(t, err, "Apply should succeed: %s", applyOutput.String()) // Verify we're still on the initial branch @@ -146,7 +146,7 @@ func TestRepositoryApplyNonExistent(t *testing.T) { // Try to apply non-existent environment var applyOutput bytes.Buffer - err := repo.Apply(ctx, "non-existent-env", &applyOutput) + err := repo.Apply(ctx, "non-existent-env", &applyOutput, true) assert.Error(t, err, "Applying non-existent environment should fail") assert.Contains(t, err.Error(), "not found") }) @@ -203,7 +203,7 @@ func TestRepositoryApplyWithConflicts(t *testing.T) { // Try to apply - this should fail due to conflict var applyOutput bytes.Buffer - err = repo.Apply(ctx, env.ID, &applyOutput) + err = repo.Apply(ctx, env.ID, &applyOutput, true) // The apply should fail due to conflict assert.Error(t, err, "Apply should fail due to conflict") diff --git a/mcpserver/tools.go b/mcpserver/tools.go index 7bf4add5..a3f23fb7 100644 --- a/mcpserver/tools.go +++ b/mcpserver/tools.go @@ -867,7 +867,7 @@ var EnvironmentEnableTrackingTool = &Tool{ 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.", + "Apply the user's unstaged changes to the environment. 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) @@ -899,17 +899,12 @@ var EnvironmentSyncFromUserTool = &Tool{ 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()) + if err := repo.Update(ctx, env, request.GetString("explanation", "")); err != nil { + return nil, fmt.Errorf("unable to update the environment: %w", err) } return mcp.NewToolResultText("Patch applied successfully to the environment:\n\n```patch\n" + string(patch) + "\n```"), nil diff --git a/repository/repository.go b/repository/repository.go index 86f473d0..298bb8ef 100644 --- a/repository/repository.go +++ b/repository/repository.go @@ -293,7 +293,7 @@ func (r *Repository) Update(ctx context.Context, env *environment.Environment, e if currentBranch == env.State.TrackingBranch { // Apply environment changes to the user's working tree var logs strings.Builder - if err := r.Apply(ctx, env.ID, &logs); err != nil { + if err := r.Apply(ctx, env.ID, &logs, false); err != nil { return fmt.Errorf("failed to apply tracking changes to working tree: %w\n\nlogs:\n%s\n", err, logs.String()) } } @@ -431,7 +431,7 @@ func (r *Repository) Merge(ctx context.Context, id string, w io.Writer) error { return RunInteractiveGitCommand(ctx, r.userRepoPath, w, "merge", "--no-ff", "--autostash", "-m", "Merge environment "+envInfo.ID, "--", "container-use/"+envInfo.ID) } -func (r *Repository) Apply(ctx context.Context, id string, w io.Writer) (rerr error) { +func (r *Repository) Apply(ctx context.Context, id string, w io.Writer, discard bool) (rerr error) { envInfo, err := r.Info(ctx, id) if err != nil { return err @@ -468,7 +468,7 @@ func (r *Repository) Apply(ctx context.Context, id string, w io.Writer) (rerr er } // Apply user changes back - if hasUnstagedChanges { + if hasUnstagedChanges && !discard { fmt.Fprintf(w, "Restoring user changes...\n") // 1. Temporarily commit the agent's changes From 6868137c420fe79031a82680ea5629857ace1a0c Mon Sep 17 00:00:00 2001 From: Alex Suraci Date: Fri, 18 Jul 2025 20:03:17 -0400 Subject: [PATCH 12/13] track by default, except for ephemeral envs Signed-off-by: Alex Suraci --- cmd/container-use/list.go | 4 +-- environment/environment.go | 12 ++++---- environment/filesystem.go | 1 - environment/state.go | 1 + mcpserver/tools.go | 60 ++++++++++++++------------------------ repository/repository.go | 4 +-- 6 files changed, 34 insertions(+), 48 deletions(-) diff --git a/cmd/container-use/list.go b/cmd/container-use/list.go index 896cbbb5..b5583b85 100644 --- a/cmd/container-use/list.go +++ b/cmd/container-use/list.go @@ -33,11 +33,11 @@ Use -q for environment IDs only, useful for scripting.`, } tw := tabwriter.NewWriter(os.Stdout, 0, 0, 2, ' ', 0) - fmt.Fprintln(tw, "ID\tTITLE\tCREATED\tUPDATED") + fmt.Fprintln(tw, "ID\tBRANCH\tTITLE\tCREATED\tUPDATED\tEPHEMERAL") defer tw.Flush() for _, envInfo := range envInfos { - fmt.Fprintf(tw, "%s\t%s\t%s\t%s\n", envInfo.ID, truncate(app, envInfo.State.Title, 40), humanize.Time(envInfo.State.CreatedAt), humanize.Time(envInfo.State.UpdatedAt)) + fmt.Fprintf(tw, "%s\t%s\t%s\t%s\t%s\t%v\n", envInfo.ID, envInfo.State.TrackingBranch, truncate(app, envInfo.State.Title, 40), humanize.Time(envInfo.State.CreatedAt), humanize.Time(envInfo.State.UpdatedAt), envInfo.State.Ephemeral) } return nil }, diff --git a/environment/environment.go b/environment/environment.go index 7c184b9a..da2556f9 100644 --- a/environment/environment.go +++ b/environment/environment.go @@ -31,15 +31,17 @@ type Environment struct { mu sync.RWMutex } -func New(ctx context.Context, dag *dagger.Client, id, title string, config *EnvironmentConfig, initialSourceDir *dagger.Directory) (*Environment, error) { +func New(ctx context.Context, dag *dagger.Client, id, branch, title string, ephemeral bool, config *EnvironmentConfig, initialSourceDir *dagger.Directory) (*Environment, error) { env := &Environment{ EnvironmentInfo: &EnvironmentInfo{ ID: id, State: &State{ - Config: config, - Title: title, - CreatedAt: time.Now(), - UpdatedAt: time.Now(), + Config: config, + Title: title, + TrackingBranch: branch, + Ephemeral: ephemeral, + CreatedAt: time.Now(), + UpdatedAt: time.Now(), }, }, dag: dag, diff --git a/environment/filesystem.go b/environment/filesystem.go index 2a3cb4f6..e77a5503 100644 --- a/environment/filesystem.go +++ b/environment/filesystem.go @@ -6,7 +6,6 @@ import ( "fmt" "strings" - "dagger.io/dagger" godiffpatch "github.com/sourcegraph/go-diff-patch" ) diff --git a/environment/state.go b/environment/state.go index 85cba15a..a0a47152 100644 --- a/environment/state.go +++ b/environment/state.go @@ -14,6 +14,7 @@ type State struct { Container string `json:"container,omitempty"` Title string `json:"title,omitempty"` TrackingBranch string `json:"tracking_branch,omitempty"` + Ephemeral bool `json:"ephemeral,omitempty"` } func (s *State) Marshal() ([]byte, error) { diff --git a/mcpserver/tools.go b/mcpserver/tools.go index a3f23fb7..8cf0d89c 100644 --- a/mcpserver/tools.go +++ b/mcpserver/tools.go @@ -10,7 +10,6 @@ import ( "log/slog" "os" "os/signal" - "strings" "syscall" "dagger.io/dagger" @@ -131,7 +130,6 @@ func init() { EnvironmentOpenTool, EnvironmentCreateTool, EnvironmentUpdateMetadataTool, - EnvironmentEnableTrackingTool, EnvironmentConfigTool, EnvironmentRunCmdTool, @@ -236,6 +234,13 @@ Environment configuration is managed by the user via cu config commands.`, mcp.Description("Short description of the work that is happening in this environment."), mcp.Required(), ), + mcp.WithBoolean("ephemeral", + mcp.Description("Whether this environment is for a sub-task of a larger task."), + mcp.Required(), + ), + mcp.WithString("background_branch", + mcp.Description("A user-supplied branch name to create and track, instead of the current branch."), + ), ), Handler: func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) { repo, err := openRepository(ctx, request) @@ -246,17 +251,31 @@ Environment configuration is managed by the user via cu config commands.`, if err != nil { return nil, err } + branch := request.GetString("background_branch", "") + ephemeral := request.GetBool("ephemeral", false) + if branch == "" && !ephemeral { + branch, err = repo.CurrentUserBranch(ctx) + if err != nil { + return nil, fmt.Errorf("failed to get current branch: %w", err) + } + } dag, ok := ctx.Value(daggerClientKey{}).(*dagger.Client) if !ok { return nil, fmt.Errorf("dagger client not found in context") } - env, err := repo.Create(ctx, dag, title, request.GetString("explanation", "")) + env, err := repo.Create(ctx, dag, branch, title, request.GetString("explanation", ""), ephemeral) if err != nil { return nil, fmt.Errorf("failed to create environment: %w", err) } + if !ephemeral { + if err := repo.TrackEnvironment(ctx, branch, env.ID); err != nil { + return nil, fmt.Errorf("unable to set current branch tracking environment: %w", err) + } + } + out, err := marshalEnvironment(env) if err != nil { return nil, fmt.Errorf("failed to marshal environment: %w", err) @@ -829,41 +848,6 @@ Supported schemas are: }, } -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", diff --git a/repository/repository.go b/repository/repository.go index 298bb8ef..90c9bc94 100644 --- a/repository/repository.go +++ b/repository/repository.go @@ -144,7 +144,7 @@ func (r *Repository) exists(ctx context.Context, id string) error { // Create creates a new environment with the given description and explanation. // Requires a dagger client for container operations during environment initialization. -func (r *Repository) Create(ctx context.Context, dag *dagger.Client, description, explanation string) (*environment.Environment, error) { +func (r *Repository) Create(ctx context.Context, dag *dagger.Client, branch, description, explanation string, ephemeral bool) (*environment.Environment, error) { id := petname.Generate(2, "-") worktree, err := r.initializeWorktree(ctx, id) if err != nil { @@ -173,7 +173,7 @@ func (r *Repository) Create(ctx context.Context, dag *dagger.Client, description return nil, err } - env, err := environment.New(ctx, dag, id, description, config, baseSourceDir) + env, err := environment.New(ctx, dag, id, branch, description, ephemeral, config, baseSourceDir) if err != nil { return nil, err } From 8f50d05c56956af1cf287b08c8512e08c597565c Mon Sep 17 00:00:00 2001 From: Alex Suraci Date: Fri, 18 Jul 2025 20:03:30 -0400 Subject: [PATCH 13/13] use Directory.withPatch Signed-off-by: Alex Suraci --- environment/filesystem.go | 7 +++---- go.mod | 4 ++-- go.sum | 6 ++---- 3 files changed, 7 insertions(+), 10 deletions(-) diff --git a/environment/filesystem.go b/environment/filesystem.go index e77a5503..96b39010 100644 --- a/environment/filesystem.go +++ b/environment/filesystem.go @@ -117,10 +117,9 @@ func (env *Environment) FileSearchReplace(ctx context.Context, explanation, targ } func (env *Environment) ApplyPatch(ctx context.Context, patch string) error { - err := env.apply(ctx, env.container(). - WithExec([]string{"patch", "-p1"}, dagger.ContainerWithExecOpts{ - Stdin: patch, - })) + ctr := env.container() + err := env.apply(ctx, ctr. + WithDirectory(".", ctr.Directory(".").WithPatch(patch))) if err != nil { return fmt.Errorf("failed applying file edit, skipping git propagation: %w", err) } diff --git a/go.mod b/go.mod index 2b822be1..4a42d9c9 100644 --- a/go.mod +++ b/go.mod @@ -5,7 +5,7 @@ go 1.24.3 toolchain go1.24.4 require ( - dagger.io/dagger v0.18.12 + dagger.io/dagger v0.18.14 github.com/charmbracelet/bubbletea v1.3.5 github.com/charmbracelet/fang v0.3.0 github.com/charmbracelet/lipgloss v1.1.0 @@ -14,6 +14,7 @@ require ( github.com/mark3labs/mcp-go v0.29.0 github.com/mitchellh/go-homedir v1.1.0 github.com/pelletier/go-toml/v2 v2.2.4 + github.com/sourcegraph/go-diff-patch v0.0.0-20240223163233-798fd1e94a8e github.com/spf13/cobra v1.9.1 github.com/stretchr/testify v1.10.0 github.com/tiborvass/go-watch v0.0.0-20250607214558-08999a83bf8b @@ -56,7 +57,6 @@ 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 diff --git a/go.sum b/go.sum index 97133155..021f5c84 100644 --- a/go.sum +++ b/go.sum @@ -1,7 +1,5 @@ -dagger.io/dagger v0.18.11 h1:6lSfemlbGM2HmdOjhgevrX2+orMDGKU/xTaBMZ+otyY= -dagger.io/dagger v0.18.11/go.mod h1:azlZ24m2br95t0jQHUBpL5SiafeqtVDLl1Itlq6GO+4= -dagger.io/dagger v0.18.12 h1:s7v8aHlzDUogZ/jW92lHC+gljCNRML+0mosfh13R4vs= -dagger.io/dagger v0.18.12/go.mod h1:azlZ24m2br95t0jQHUBpL5SiafeqtVDLl1Itlq6GO+4= +dagger.io/dagger v0.18.14 h1:7+VFqNJffm6Qa8ckNRMfsM64sI5dXbRnZswCQ1jnDF0= +dagger.io/dagger v0.18.14/go.mod h1:azlZ24m2br95t0jQHUBpL5SiafeqtVDLl1Itlq6GO+4= github.com/99designs/gqlgen v0.17.75 h1:GwHJsptXWLHeY7JO8b7YueUI4w9Pom6wJTICosDtQuI= github.com/99designs/gqlgen v0.17.75/go.mod h1:p7gbTpdnHyl70hmSpM8XG8GiKwmCv+T5zkdY8U8bLog= github.com/Khan/genqlient v0.8.1 h1:wtOCc8N9rNynRLXN3k3CnfzheCUNKBcvXmVv5zt6WCs=