diff --git a/main.go b/main.go
index 4811475..337d72d 100644
--- a/main.go
+++ b/main.go
@@ -270,9 +270,15 @@ func cleanupGitFiles(files []string) error {
 			logger.Printf("Successfully committed file: %s with message: 'Track %s'", filename, filename)
 		} else {
 			logger.Printf("File %s is already tracked. Restoring it to the working tree.", file)
-			worktree.Restore(&git.RestoreOptions{
-				Files: []string{file},
+			err := worktree.Restore(&git.RestoreOptions{
+				Files:    []string{file},
+				Staged:   true,
+				Worktree: true,
 			})
+			if err != nil {
+				fmt.Fprintf(os.Stderr, "Error restoring file: %v\n", err)
+				return fmt.Errorf("error restoring file: %w", err)
+			}
 			logger.Printf("File %s restored successfully.", file)
 		}
 	}