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.
+      """
     ]
   ]
 
@@ -2052,23 +2065,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

@@ -1791,9 +1791,11 @@ defmodule Ash.ActionInput do
   end
 
   def run_after_transaction_hooks(result, input) do
+    {result_for_hooks, pending_notifications} = peel_after_transaction_notifications(result, input)
+
     input.after_transaction
     |> Enum.reduce(
-      result,
+      result_for_hooks,
       fn after_transaction, result ->
         tracer = input.context[:private][:tracer]
 
@@ -1827,6 +1829,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 +1843,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,15 +1856,55 @@ 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:
 
         #{inspect(other)}
         """
     end
+    |> attach_after_transaction_notifications(pending_notifications, input)
+  end
+
+  defp peel_after_transaction_notifications({:ok, result, notifications}, input)
+       when has_return?(input) and is_list(notifications) do
+    {{:ok, result}, notifications}
+  end
+
+  defp peel_after_transaction_notifications({:ok, notifications}, input)
+       when has_no_return?(input) and is_list(notifications) do
+    # Hooks for untyped actions expect `:ok`; notifications are reattached after.
+    {:ok, notifications}
+  end
+
+  defp peel_after_transaction_notifications(result, _input), do: {result, :none}
+
+  defp attach_after_transaction_notifications(result, :none, _input), do: result
+  defp attach_after_transaction_notifications({:error, _} = error, _notifications, _input), do: error
+
+  defp attach_after_transaction_notifications({:ok, result}, notifications, input)
+       when has_return?(input) and is_list(notifications) do
+    {:ok, result, notifications}
   end
 
+  defp attach_after_transaction_notifications({:ok, result, hook_notifications}, notifications, input)
+       when has_return?(input) and is_list(notifications) and is_list(hook_notifications) do
+    {:ok, result, notifications ++ hook_notifications}
+  end
+
+  defp attach_after_transaction_notifications(:ok, notifications, input)
+       when has_no_return?(input) and is_list(notifications) do
+    {:ok, notifications}
+  end
+
+  defp attach_after_transaction_notifications({:ok, hook_notifications}, notifications, input)
+       when has_no_return?(input) and is_list(notifications) and is_list(hook_notifications) do
+    {:ok, notifications ++ hook_notifications}
+  end
+
+  defp attach_after_transaction_notifications(result, _notifications, _input), do: result
+
   @doc false
   def run_around_transaction_hooks(%{around_transaction: []} = input, func) do
     case func.(input) do
@@ -1865,6 +1914,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 +1928,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 +1941,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:
@@ -1895,8 +1952,62 @@ defmodule Ash.ActionInput do
   end
 
   def run_around_transaction_hooks(%{around_transaction: [around | rest]} = input, func) do
-    around.(input, fn input ->
-      run_around_transaction_hooks(%{input | around_transaction: rest}, func)
-    end)
+    pending_ref = make_ref()
+
+    result =
+      around.(input, fn input ->
+        case run_around_transaction_hooks(%{input | around_transaction: rest}, func) do
+          {:ok, term, notifications} when has_return?(input) and is_list(notifications) ->
+            Process.put(pending_ref, notifications)
+            {:ok, term}
+
+          {:ok, notifications} when has_no_return?(input) and is_list(notifications) ->
+            Process.put(pending_ref, notifications)
+            :ok
+
+          other ->
+            other
+        end
+      end)
+
+    pending = Process.get(pending_ref, [])
+    Process.delete(pending_ref)
+
+    attach_around_notifications(result, pending, input)
+  end
+
+  defp attach_around_notifications({:error, _} = error, _, _), do: error
+
+  defp attach_around_notifications(result, pending, input) do
+    {base, extra} =
+      case result do
+        {:ok, term, notifications} when is_list(notifications) -> {{:ok, term}, notifications}
+        other -> {other, []}
+      end
+
+    notifications = pending ++ extra
+
+    cond do
+      has_return?(input) ->
+        case base do
+          {:ok, term} ->
+            if notifications == [], do: {:ok, term}, else: {:ok, term, notifications}
+
+          other ->
+            other
+        end
+
+      has_no_return?(input) ->
+        case base do
+          :ok ->
+            if notifications == [], do: :ok, else: {:ok, notifications}
+
+          {:ok, more} when is_list(more) ->
+            {:ok, notifications ++ more}
+
+          other ->
+            other
+        end
+    end
   end
 end

lib/ash/actions/action.ex

@@ -133,6 +133,18 @@ 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, _opts)
+       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
@@ -199,31 +211,9 @@ 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
-
-              # Run after_transaction hooks
-              Ash.ActionInput.run_after_transaction_hooks(final_result, input)
+              finalize_notifications(notifications, input, opts, notify?)
+              |> build_result(result, input, opts)
+              |> Ash.ActionInput.run_after_transaction_hooks(input)
 
             {:error, error} ->
               error_result = {:error, Ash.Error.to_ash_error(error)}
@@ -250,17 +240,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}
@@ -270,7 +251,6 @@ defmodule Ash.Actions.Action do
               {:error, error}
           end
 
-        # Run after_transaction hooks
         Ash.ActionInput.run_after_transaction_hooks(result, input)
 
       {:error, error} ->
@@ -403,6 +383,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