wip: opt environment in to tracking a branch

Signed-off-by: Alex Suraci <suraci.alex@gmail.com>

Author: vito

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) {

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",

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)
 	}