diff --git a/Justfile b/Justfile
index e75b202..81db32e 100644
--- a/Justfile
+++ b/Justfile
@@ -4,5 +4,5 @@ help:
 ci:
   nix-unit ./tests.nix
 
-fmt:
-  nixfmt .
+fmt *args:
+  treefmt {{args}}
diff --git a/shell.nix b/shell.nix
index 11a214d..3dfb02a 100644
--- a/shell.nix
+++ b/shell.nix
@@ -6,6 +6,6 @@ pkgs.mkShell {
     pkgs.just
     pkgs.nix-unit
     pkgs.treefmt
-    pkgs.nixfmt-tree
+    pkgs.nixfmt
   ];
 }
diff --git a/treefmt.toml b/treefmt.toml
new file mode 100644
index 0000000..d50e8eb
--- /dev/null
+++ b/treefmt.toml
@@ -0,0 +1,3 @@
+[formatter.nixfmt]
+command = "nixfmt"
+includes = [ "*.nix" ]
\ No newline at end of file