Union Forms and NoSuchInput Error

I’m building an application with a plugin system where each component can have different configuration parameters based on the plugin type selected.

I’ve implemented the Union Forms approach from the documentation for my plugin parameters, but I’m encountering an issue.

My Setup

I have a resource (Plugin) with a params attribute that uses a union type to accommodate different plugin parameter structures:

defmodule MyApp.Plugins.Plugin do
  use Ash.Resource, otp_app: :my_app, domain: MyApp.Plugins, data_layer: AshPostgres.DataLayer

  # Other configuration...

  attributes do
    integer_primary_key :id
    attribute :name, :string
    attribute :active, :boolean, default: false
    attribute :plugin_type, :module
    # This is the union attribute that holds different parameter configurations
    attribute :params, :plugin_params
    # Other attributes...
  end

  # Actions, relationships, etc.
end

The union type for plugin parameters:

defmodule MyApp.Types.PluginParams do
  defmodule TypeAParams do
    use Ash.Resource, data_layer: :embedded

    attributes do
      attribute :window_size, :integer, default: 3000
      attribute :threshold, :float, default: 0.70
      # Other parameters...
    end

    actions do
      defaults [:read, :destroy, create: :*, update: :*]
    end
  end

  use Ash.Type.NewType,
    subtype_of: :union,
    constraints: [
      types: [
        type_a: [
          type: TypeAParams,
          tag: :type,
          tag_value: :type_a
        ]
        # Other plugin types...
      ]
    ]
end

In my form component:

<.inputs_for :let={fc} field={@form[:params]}>
  <.input
    field={fc[:_union_type]}
    phx-change="type-changed"
    type="select"
    label="Plugin Type"
    options={["Type A": "type_a"]}
  />

  <%= case fc.params["_union_type"] do %>
    <% "type_a" -> %>
      <.input type="number" label="Window Size" field={fc[:window_size]} />
      <% # Other type_a inputs... %>
  <% end %>
</.inputs_for>

The Issue

When rendering the form, the current values are correctly displayed in the inputs. However, when I try to change any input in the nested form, I get this error:

%Ash.Error.Invalid.NoSuchInput{
  calculation: nil,
  resource: MyApp.Types.PluginParams.TypeAParams,
  action: :update,
  input: "window_size",
  inputs: MapSet.new([]),
  did_you_mean: [],
  # Other error details...
}

I’m following the type-changed handler from the documentation:

def handle_event("type-changed", %{"_target" => path} = params, socket) do
  new_type = get_in(params, path)
  path = :lists.droplast(path)

  form =
    socket.assigns.form
    |> AshPhoenix.Form.remove_form(path)
    |> AshPhoenix.Form.add_form(path, params: %{"_union_type" => new_type})

  {:noreply, assign(socket, :form, form)}
end

Any insights on what I might be doing wrong? The form renders correctly with the initial values but fails when trying to update any field.