diff --git a/plugins/lsp/conform-nvim.nix b/plugins/lsp/conform-nvim.nix
index a98a63e3..224cc51e 100644
--- a/plugins/lsp/conform-nvim.nix
+++ b/plugins/lsp/conform-nvim.nix
@@ -40,9 +40,9 @@ in {
       formatOnSave =
         helpers.defaultNullOpts.mkNullable
         (
-          with types;
+          with helpers.nixvimTypes;
             either
-            str
+            strLuaFn
             (
               submodule {
                 options = {
@@ -69,15 +69,23 @@ in {
         '';
 
       formatAfterSave =
-        helpers.defaultNullOpts.mkNullable (types.submodule {
-          options = {
-            lspFallback = mkOption {
-              type = types.bool;
-              default = true;
-              description = "See :help conform.format for details.";
-            };
-          };
-        })
+        helpers.defaultNullOpts.mkNullable
+        (
+          with helpers.nixvimTypes;
+            either
+            strLuaFn
+            (
+              submodule {
+                options = {
+                  lspFallback = mkOption {
+                    type = types.bool;
+                    default = true;
+                    description = "See :help conform.format for details.";
+                  };
+                };
+              }
+            )
+        )
         "see documentation"
         ''
           If this is set, Conform will run the formatter asynchronously after save.
@@ -110,9 +118,12 @@ in {
             timeout_ms = formatOnSave.timeoutMs;
           }
           else helpers.mkRaw formatOnSave;
-        format_after_save = helpers.ifNonNull' formatAfterSave {
-          lsp_fallback = formatAfterSave.lspFallback;
-        };
+        format_after_save =
+          if builtins.isAttrs formatAfterSave
+          then {
+            lsp_fallback = formatOnSave.lspFallback;
+          }
+          else helpers.mkRaw formatAfterSave;
         log_level = logLevel;
         notify_on_error = notifyOnError;
         inherit formatters;
diff --git a/tests/test-sources/plugins/lsp/conform-nvim.nix b/tests/test-sources/plugins/lsp/conform-nvim.nix
index ad733fa6..2a1b1fa8 100644
--- a/tests/test-sources/plugins/lsp/conform-nvim.nix
+++ b/tests/test-sources/plugins/lsp/conform-nvim.nix
@@ -94,4 +94,30 @@
       '';
     };
   };
+
+  custom_format_after_save_function = {
+    plugins.conform-nvim = {
+      enable = true;
+
+      formattersByFt = {
+        lua = ["stylua"];
+        python = ["isort" "black"];
+        javascript = [["prettierd" "prettier"]];
+        "*" = ["codespell"];
+        "_" = ["trimWhitespace"];
+      };
+
+      formatAfterSave = ''
+        function(bufnr)
+          if vim.g.disable_autoformat or vim.b[bufnr].disable_autoformat then
+            return
+          end
+          if not _conform_slow_format_filetypes[vim.bo[bufnr].filetype] then
+            return
+          end
+          return { lsp_fallback = true }
+        end
+      '';
+    };
+  };
 }