fix: handle notifications when running nested actions

lib/ash.ex

@@ -890,6 +890,19 @@ defmodule Ash do
     load: [
       type: :any,
       doc: "A load statement to apply on the resulting records after the action is invoked."
+    ],
+    return_notifications?: [
+      type: :boolean,
+      default: false,
+      doc: """
+      Use this if you're running ash actions in your own transaction and you want to manually handle sending notifications.
+
+      If a transaction is ongoing, and this is false, notifications will be aggregated in the process dictionary,
+      otherwise the return value is `{:ok, result, notifications}` (or `{:ok, notifications}` for actions with no return type)
+
+      To send notifications later, use `Ash.Notifier.notify(notifications)`. It sends any notifications
+      that can be sent, and returns the rest.
+      """
     ]
   ]
 
@@ -1944,23 +1957,32 @@ defmodule Ash do
   """
   @doc spark_opts: [{1, @run_action_opts}]
   @spec run_action(input :: Ash.ActionInput.t(), opts :: Keyword.t()) ::
-          :ok | {:ok, term} | {:error, Ash.Error.t()}
+          :ok
+          | {:ok, term}
+          | {:ok, list(Ash.Notifier.Notification.t())}
+          | {:ok, term, list(Ash.Notifier.Notification.t())}
+          | {:error, Ash.Error.t()}
   def run_action(input, opts \\ []) do
     Ash.Helpers.expect_options!(opts)
     domain = Ash.Helpers.domain!(input, opts)
 
     with {:ok, opts} <- RunActionOpts.validate(opts),
          opts <- RunActionOpts.to_options(opts),
          input = %{input | domain: domain},
-         {:ok, _resource} <- Ash.Domain.Info.resource(domain, input.resource),
-         {:ok, result} <- Ash.Actions.Action.run(domain, input, opts) do
-      {:ok, result}
-    else
-      :ok ->
-        :ok
+         {:ok, _resource} <- Ash.Domain.Info.resource(domain, input.resource) do
+      case Ash.Actions.Action.run(domain, input, opts) do
+        {:ok, result, notifications} ->
+          {:ok, result, notifications}
 
-      {:error, error} ->
-        {:error, Ash.Error.to_error_class(error)}
+        {:ok, result} ->
+          {:ok, result}
+
+        :ok ->
+          :ok
+
+        {:error, error} ->
+          {:error, Ash.Error.to_error_class(error)}
+      end
     end
   end
 

lib/ash/action_input.ex

@@ -1827,6 +1827,12 @@ defmodule Ash.ActionInput do
       {:ok, new_result} when has_return?(input) ->
         {:ok, new_result}
 
+      {:ok, new_result, notifications} when has_return?(input) and is_list(notifications) ->
+        {:ok, new_result, notifications}
+
+      {:ok, notifications} when has_no_return?(input) and is_list(notifications) ->
+        {:ok, notifications}
+
       {:error, error} ->
         {:error, error}
 
@@ -1835,6 +1841,7 @@ defmodule Ash.ActionInput do
         Invalid return value from after_transaction hook. Because this action has a return type I expected one of:
 
         * {:ok, term}
+        * {:ok, term, notifications}
         * {:error, error}
 
         Got:
@@ -1847,6 +1854,7 @@ defmodule Ash.ActionInput do
         Invalid return value from after_transaction hook. Because this action has no return type I expected one of:
 
         * :ok
+        * {:ok, notifications}
         * {:error, error}
 
         Got:
@@ -1865,6 +1873,12 @@ defmodule Ash.ActionInput do
       {:ok, term} when has_return?(input) ->
         {:ok, term}
 
+      {:ok, term, notifications} when has_return?(input) and is_list(notifications) ->
+        {:ok, term, notifications}
+
+      {:ok, notifications} when has_no_return?(input) and is_list(notifications) ->
+        {:ok, notifications}
+
       {:error, error} ->
         {:error, error}
 
@@ -1873,6 +1887,7 @@ defmodule Ash.ActionInput do
         Invalid return value from around_transaction hook. Because this action has no return type, I expected one of:
 
         * :ok
+        * {:ok, notifications}
         * {:error, error}
 
         Got:
@@ -1885,6 +1900,7 @@ defmodule Ash.ActionInput do
         Invalid return value from around_transaction hook. Because this action has a return type, I expected one of:
 
         * {:ok, term}
+        * {:ok, term, notifications}
         * {:error, error}
 
         Got:

lib/ash/actions/action.ex

@@ -77,7 +77,9 @@ defmodule Ash.Actions.Action do
             else
               run_without_transaction(domain, input, module, run_opts, context, opts)
             end
+            |> strip_return_notifications_for_hooks(input, opts)
           end)
+          |> restore_return_notifications_after_hooks(input, opts)
 
         result = maybe_load(result, input, domain, opts)
 
@@ -133,6 +135,20 @@ defmodule Ash.Actions.Action do
   defp maybe_load(:ok, _input, _domain, _opts), do: :ok
   defp maybe_load({:ok, nil}, _input, _domain, _opts), do: {:ok, nil}
 
+  defp maybe_load({:ok, result, notifications}, input, domain, opts) do
+    case maybe_load({:ok, result}, input, domain, opts) do
+      {:ok, loaded} -> {:ok, loaded, notifications}
+      other -> other
+    end
+  end
+
+  defp maybe_load({:ok, notifications}, %{action: %{returns: nil}} = _input, _domain, %{
+         return_notifications?: true
+       })
+       when is_list(notifications) do
+    {:ok, notifications}
+  end
+
   defp maybe_load({:ok, result}, input, domain, opts) do
     constraints = input.action.constraints || []
     returns = input.action.returns
@@ -154,6 +170,41 @@ defmodule Ash.Actions.Action do
 
   defp maybe_load(other, _input, _domain, _opts), do: other
 
+  defp strip_return_notifications_for_hooks({:ok, result, notifications}, input, %{
+         return_notifications?: true
+       })
+       when not is_nil(input.action.returns) do
+    Process.put(:ash_return_notifications, notifications)
+    {:ok, result}
+  end
+
+  defp strip_return_notifications_for_hooks({:ok, notifications}, %{action: %{returns: nil}}, %{
+         return_notifications?: true
+       }) do
+    Process.put(:ash_return_notifications, notifications)
+    :ok
+  end
+
+  defp strip_return_notifications_for_hooks(result, _input, _opts), do: result
+
+  defp restore_return_notifications_after_hooks(result, input, %{return_notifications?: true}) do
+    case Process.delete(:ash_return_notifications) do
+      nil ->
+        result
+
+      notifications when not is_nil(input.action.returns) ->
+        case result do
+          {:ok, result} -> {:ok, result, notifications}
+          other -> other
+        end
+
+      notifications ->
+        {:ok, notifications}
+    end
+  end
+
+  defp restore_return_notifications_after_hooks(result, _input, _opts), do: result
+
   defp run_with_transaction(domain, input, module, run_opts, context, opts) do
     # Run before_transaction hooks first
     case Ash.ActionInput.run_before_transaction_hooks(input) do
@@ -199,30 +250,15 @@ defmodule Ash.Actions.Action do
           )
           |> case do
             {:ok, {:ok, result, notifications}} ->
-              notifications =
-                if notify? && !opts[:return_notifications?] do
-                  Enum.concat(
-                    notifications || [],
-                    Process.delete(:ash_notifications) || []
-                  )
-                else
-                  notifications || []
-                end
-
-              remaining = Ash.Notifier.notify(notifications)
-
-              Ash.Actions.Helpers.warn_missed!(input.resource, input.action, %{
-                resource_notifications: remaining
-              })
-
               final_result =
-                if input.action.returns do
-                  {:ok, result}
-                else
-                  :ok
-                end
+                finalize_notifications(
+                  notifications,
+                  input,
+                  opts,
+                  notify?
+                )
+                |> build_result(result, input, opts)
 
-              # Run after_transaction hooks
               Ash.ActionInput.run_after_transaction_hooks(final_result, input)
 
             {:error, error} ->
@@ -250,17 +286,8 @@ defmodule Ash.Actions.Action do
             :ok ->
               case run_with_hooks(module, input, run_opts, context, false) do
                 {:ok, result, notifications} ->
-                  remaining = Ash.Notifier.notify(notifications)
-
-                  Ash.Actions.Helpers.warn_missed!(input.resource, input.action, %{
-                    resource_notifications: remaining
-                  })
-
-                  if input.action.returns do
-                    {:ok, result}
-                  else
-                    :ok
-                  end
+                  finalize_notifications(notifications, input, opts, false)
+                  |> build_result(result, input, opts)
 
                 {:error, error} ->
                   {:error, error}
@@ -403,6 +430,67 @@ defmodule Ash.Actions.Action do
     )
   end
 
+  defp finalize_notifications(notifications, input, opts, notify?) do
+    notifications = List.wrap(notifications)
+
+    notifications =
+      if notify? && !opts[:return_notifications?] do
+        Enum.concat(notifications, Process.delete(:ash_notifications) || [])
+      else
+        notifications
+      end
+
+    if opts[:return_notifications?] do
+      notifications
+    else
+      if Process.get(:ash_started_transaction?) && !notify? do
+        current_notifications = List.wrap(Process.get(:ash_notifications, []))
+
+        Process.put(:ash_notifications, current_notifications ++ notifications)
+      else
+        remaining = Ash.Notifier.notify(notifications)
+
+        Ash.Actions.Helpers.warn_missed!(input.resource, input.action, %{
+          resource_notifications: remaining
+        })
+      end
+
+      :ok
+    end
+  end
+
+  defp build_result(:ok, result, input, opts) do
+    if opts[:return_notifications?] do
+      if input.action.returns do
+        {:ok, result, []}
+      else
+        {:ok, []}
+      end
+    else
+      if input.action.returns do
+        {:ok, result}
+      else
+        :ok
+      end
+    end
+  end
+
+  defp build_result(notifications, result, input, opts) when is_list(notifications) do
+    if opts[:return_notifications?] do
+      if input.action.returns do
+        {:ok, result, notifications}
+      else
+        {:ok, notifications}
+      end
+    else
+      if input.action.returns do
+        {:ok, result}
+      else
+        :ok
+      end
+    end
+  end
+
   defp run_with_hooks(module, input, run_opts, context, in_transaction?) do
     # Run before_action hooks
     case Ash.ActionInput.run_before_actions(input) do

test/actions/generic_actions_test.exs

@@ -775,6 +775,16 @@ defmodule Ash.Test.Actions.GenericActionsTest do
 
       assert result == "Processed: test"
     end
+
+    test "return_notifications?: true returns notifications" do
+      assert {:ok, "Processed: test", notifications} =
+               Post
+               |> Ash.ActionInput.for_action(:with_notifications, %{message: "test"})
+               |> Ash.run_action(return_notifications?: true)
+
+      assert length(notifications) == 1
+      assert hd(notifications).data == %{message: "test"}
+    end
   end
 
   describe "action-level preparations and validations" do