feat(admin): assign existing upstreams to pools

lib/codex_pooler/admin/upstream_assignment_workflow.ex

@@ -0,0 +1,131 @@
+defmodule CodexPooler.Admin.UpstreamAssignmentWorkflow do
+  @moduledoc """
+  Coordinates the operator workflow for attaching an existing upstream identity to a Pool.
+  """
+
+  require Logger
+
+  alias CodexPooler.Accounts.{Scope, User}
+  alias CodexPooler.Audit
+  alias CodexPooler.Events
+  alias CodexPooler.Jobs
+  alias CodexPooler.Pools
+  alias CodexPooler.Pools.Pool
+  alias CodexPooler.Upstreams
+  alias CodexPooler.Upstreams.Assignments, as: UpstreamAssignments
+  alias CodexPooler.Upstreams.Schemas.{PoolUpstreamAssignment, UpstreamIdentity}
+
+  @audit_action "upstream_account.assign_pool"
+  @catalog_trigger_kind "manual"
+
+  @type workflow_error :: %{required(:code) => atom(), required(:message) => String.t()}
+  @type workflow_result ::
+          {:ok, PoolUpstreamAssignment.t()}
+          | {:error, Ecto.Changeset.t() | workflow_error()}
+
+  @spec assign_to_pool(Scope.t(), Pool.t(), UpstreamIdentity.t() | Ecto.UUID.t()) ::
+          workflow_result()
+  def assign_to_pool(
+        %Scope{user: %User{}} = scope,
+        %Pool{} = pool,
+        identity_or_id
+      ) do
+    with {:ok, _decision} <-
+           Pools.require_capability(scope, Pools.capability(:pool_operate), pool_id: pool.id),
+         {:ok, identity} <- authorize_identity(scope, identity_or_id),
+         {:ok, assignment} <- UpstreamAssignments.assign_pool_assignment(pool, identity) do
+      record_audit(scope, pool, identity, assignment)
+      broadcast_assignment(pool, identity, assignment)
+      enqueue_catalog_sync(pool)
+
+      {:ok, assignment}
+    end
+  end
+
+  def assign_to_pool(_scope, _pool, _identity_or_id),
+    do: {:error, workflow_error(:invalid_request, "user scope and Pool are required")}
+
+  defp authorize_identity(%Scope{} = scope, identity_or_id) do
+    case normalize_identity(identity_or_id) do
+      %UpstreamIdentity{} = identity ->
+        authorize_loaded_identity(scope, identity)
+
+      nil ->
+        {:error, workflow_error(:upstream_identity_not_found, "upstream identity was not found")}
+    end
+  end
+
+  defp authorize_loaded_identity(scope, identity) do
+    if Pools.owner?(scope) or visible_identity?(scope, identity.id) do
+      {:ok, identity}
+    else
+      {:error,
+       workflow_error(
+         :capability_denied,
+         "the upstream identity is not available in the operator's Pool scope"
+       )}
+    end
+  end
+
+  defp visible_identity?(scope, identity_id) do
+    scope
+    |> Upstreams.list_visible_upstream_identities()
+    |> Enum.any?(&(&1.id == identity_id))
+  end
+
+  defp normalize_identity(%UpstreamIdentity{id: identity_id}),
+    do: Upstreams.get_upstream_identity(identity_id)
+
+  defp normalize_identity(identity_id) when is_binary(identity_id),
+    do: Upstreams.get_upstream_identity(identity_id)
+
+  defp normalize_identity(_identity_or_id), do: nil
+
+  defp record_audit(%Scope{user: %User{} = user}, pool, identity, assignment) do
+    Audit.record_user_event(user, %{
+      pool_id: pool.id,
+      action: @audit_action,
+      target_type: "upstream_identity",
+      target_id: identity.id,
+      details: %{
+        upstream_identity_id: identity.id,
+        pool_assignment_ids: [assignment.id],
+        assignment_status: assignment.status
+      }
+    })
+
+    :ok
+  end
+
+  defp broadcast_assignment(pool, identity, assignment) do
+    Events.broadcast_upstreams(pool, "upstream_assignment_assigned", %{
+      assignment_id: assignment.id,
+      upstream_identity_id: identity.id,
+      assignment_status: assignment.status
+    })
+
+    :ok
+  end
+
+  defp enqueue_catalog_sync(pool) do
+    case Jobs.enqueue_catalog_sync(pool, trigger_kind: @catalog_trigger_kind) do
+      {:ok, _job} ->
+        :ok
+
+      {:error, reason} ->
+        Logger.warning(fn ->
+          "upstream assignment catalog sync enqueue failed pool_id=#{pool.id} " <>
+            "trigger_kind=#{@catalog_trigger_kind} reason=#{catalog_enqueue_error_code(reason)}"
+        end)
+
+        :ok
+    end
+  end
+
+  defp catalog_enqueue_error_code(%Ecto.Changeset{}), do: "invalid_job"
+  defp catalog_enqueue_error_code(%{code: code}) when is_atom(code), do: Atom.to_string(code)
+  defp catalog_enqueue_error_code(reason) when is_atom(reason), do: Atom.to_string(reason)
+  defp catalog_enqueue_error_code(_reason), do: "unknown"
+
+  defp workflow_error(code, message), do: %{code: code, message: message}
+end

lib/codex_pooler/audit.ex

@@ -36,6 +36,7 @@ defmodule CodexPooler.Audit do
     {"Pool invite created", "invite.create"},
     {"Pool invite revoked", "invite.revoke"},
     {"Upstream account imported", "upstream_account.import"},
+    {"Upstream account assigned to Pool", "upstream_account.assign_pool"},
     {"Upstream account paused", "upstream_account.pause"},
     {"Upstream account reactivated", "upstream_account.reactivate"},
     {"Upstream account token refresh queued", "upstream_account.refresh_enqueue"},

lib/codex_pooler/upstreams.ex

@@ -49,9 +49,11 @@ defmodule CodexPooler.Upstreams do
   @spec list_upstream_identities(keyword()) :: [UpstreamIdentity.t()]
   def list_upstream_identities(opts \\ []) do
     status = Keyword.get(opts, :status)
+    excluded_status = Keyword.get(opts, :exclude_status)
 
     UpstreamIdentity
     |> maybe_where_status(status)
+    |> maybe_exclude_status(excluded_status)
     |> order_by([identity], asc: identity.account_label)
     |> Repo.all()
   end
@@ -248,4 +250,9 @@ defmodule CodexPooler.Upstreams do
 
   defp maybe_where_status(query, status),
     do: from(identity in query, where: identity.status == ^status)
+
+  defp maybe_exclude_status(query, nil), do: query
+
+  defp maybe_exclude_status(query, status),
+    do: from(identity in query, where: identity.status != ^status)
 end

lib/codex_pooler/upstreams/assignments.ex

@@ -23,6 +23,9 @@ defmodule CodexPooler.Upstreams.Assignments do
   defdelegate sync_pool_assignments_for_pool_edit(pool, selected_ids, opts \\ []),
     to: PoolAssignments
 
+  @spec assign_pool_assignment(Pool.t(), identity_ref(), map()) :: assignment_result()
+  defdelegate assign_pool_assignment(pool, identity_or_id, attrs \\ %{}), to: PoolAssignments
+
   @spec put_assignment_cooldown(assignment_ref(), DateTime.t(), map()) :: assignment_result()
   defdelegate put_assignment_cooldown(assignment_or_id, cooldown_until, attrs \\ %{}),
     to: PoolAssignments

lib/codex_pooler/upstreams/assignments/pool_assignments.ex

@@ -19,7 +19,7 @@ defmodule CodexPooler.Upstreams.Assignments.PoolAssignments do
   @health_active PoolUpstreamAssignment.active_health_status()
   @health_cooldown PoolUpstreamAssignment.cooldown_health_status()
   @health_disabled PoolUpstreamAssignment.disabled_health_status()
-  @pending UpstreamIdentity.pending_status()
+  @pending PoolUpstreamAssignment.pending_status()
 
   @type lifecycle_error :: %{required(:code) => atom(), required(:message) => String.t()}
   @type lifecycle_result :: {:ok, map()} | {:error, lifecycle_error()}
@@ -60,6 +60,31 @@ defmodule CodexPooler.Upstreams.Assignments.PoolAssignments do
   def create_pool_assignment(_pool, _identity_or_id, _attrs),
     do: {:error, lifecycle_error(:pool_not_found, "pool was not found")}
 
+  @doc """
+  Assigns an already-linked upstream identity to a Pool as an operational
+  assignment.
+
+  `create_pool_assignment/3` intentionally defaults to `pending` for
+  onboarding. The admin "Assign to Pool" workflow must instead create an
+  active assignment and restore an older pending/deleted row when one already
+  occupies the Pool/identity slot.
+  """
+  @spec assign_pool_assignment(Pool.t(), identity_ref(), map()) :: assignment_result()
+  def assign_pool_assignment(pool, identity_or_id, attrs \\ %{})
+
+  def assign_pool_assignment(%Pool{} = pool, identity_or_id, attrs) when is_map(attrs) do
+    case identity_id(identity_or_id) do
+      identity_id when is_binary(identity_id) ->
+        assign_pool_assignment_transaction(pool, identity_id, attrs)
+
+      nil ->
+        {:error, lifecycle_error(:upstream_identity_not_found, "upstream identity was not found")}
+    end
+  end
+
+  def assign_pool_assignment(_pool, _identity_or_id, _attrs),
+    do: {:error, lifecycle_error(:pool_not_found, "pool was not found")}
+
   @spec sync_pool_assignments_for_pool_edit(Pool.t(), [Ecto.UUID.t()], keyword()) ::
           :ok | {:error, term()}
   def sync_pool_assignments_for_pool_edit(pool, selected_ids, opts \\ [])
@@ -570,6 +595,73 @@ defmodule CodexPooler.Upstreams.Assignments.PoolAssignments do
   defp pool_id(id) when is_binary(id), do: id
   defp pool_id(_id), do: nil
 
+  defp assignment_for_pool_identity(pool_id, identity_id) do
+    Repo.one(
+      from assignment in PoolUpstreamAssignment,
+        where:
+          assignment.pool_id == ^pool_id and
+            assignment.upstream_identity_id == ^identity_id,
+        lock: "FOR UPDATE"
+    )
+  end
+
+  defp assign_pool_assignment_transaction(pool, identity_id, attrs) do
+    Repo.transaction(fn ->
+      identity_id
+      |> lock_upstream_identity()
+      |> assign_locked_identity(pool, attrs)
+      |> case do
+        {:ok, assignment} -> assignment
+        {:error, reason} -> Repo.rollback(reason)
+      end
+    end)
+  end
+
+  defp lock_upstream_identity(identity_id) do
+    Repo.one(
+      from identity in UpstreamIdentity,
+        where: identity.id == ^identity_id,
+        lock: "FOR UPDATE"
+    )
+  end
+
+  defp assign_locked_identity(nil, _pool, _attrs) do
+    {:error, lifecycle_error(:upstream_identity_not_found, "upstream identity was not found")}
+  end
+
+  defp assign_locked_identity(%UpstreamIdentity{status: @deleted}, _pool, _attrs) do
+    {:error,
+     lifecycle_error(
+       :upstream_identity_not_assignable,
+       "deleted upstream identities cannot be assigned"
+     )}
+  end
+
+  defp assign_locked_identity(%UpstreamIdentity{} = identity, pool, attrs) do
+    assignment_attrs =
+      attrs
+      |> atomize_attrs()
+      |> Map.merge(%{
+        status: @assignment_active,
+        health_status: @health_active,
+        eligibility_status: @eligible,
+        cooldown_until: nil,
+        disabled_at: nil
+      })
+
+    case assignment_for_pool_identity(pool.id, identity.id) do
+      %PoolUpstreamAssignment{status: status} = assignment
+      when status in [@pending, @assignment_deleted] ->
+        update_pool_assignment(assignment, assignment_attrs)
+
+      %PoolUpstreamAssignment{} = assignment ->
+        {:ok, assignment}
+
+      nil ->
+        create_pool_assignment(pool, identity, assignment_attrs)
+    end
+  end
+
   defp atomize_attrs(attrs) when is_map(attrs) do
     Map.new(attrs, fn
       {key, value} when is_binary(key) -> {String.to_existing_atom(key), value}

lib/codex_pooler/upstreams/lifecycle/account_lifecycle.ex

@@ -20,6 +20,7 @@ defmodule CodexPooler.Upstreams.Lifecycle.AccountLifecycle do
   @refresh_failed UpstreamIdentity.refresh_failed_status()
   @reauth_required UpstreamIdentity.reauth_required_status()
   @deleted UpstreamIdentity.deleted_status()
+  @assignment_pending PoolUpstreamAssignment.pending_status()
   @assignment_active PoolUpstreamAssignment.active_status()
   @assignment_paused PoolUpstreamAssignment.paused_status()
   @assignment_refresh_due PoolUpstreamAssignment.refresh_due_status()
@@ -31,6 +32,7 @@ defmodule CodexPooler.Upstreams.Lifecycle.AccountLifecycle do
   @health_disabled PoolUpstreamAssignment.disabled_health_status()
   @reactivatable_statuses [@active, @paused, @refresh_due, @refresh_failed]
   @reactivatable_assignment_statuses [
+    @assignment_pending,
     @assignment_active,
     @assignment_paused,
     @assignment_refresh_due,

lib/codex_pooler_web/live/admin/components/pages/upstreams/account_card.ex

@@ -523,6 +523,16 @@ defmodule CodexPoolerWeb.Admin.UpstreamPageComponents.AccountCard do
         tabindex="0"
         class="menu dropdown-content z-20 mt-2 w-60 rounded-box border border-base-300 bg-base-100 p-2 text-left shadow-xl"
       >
+        <li>
+          <AdminComponents.dropdown_action_item
+            id={"assign-pool-upstream-account-#{@account.identity.id}"}
+            icon="hero-server-stack"
+            label="Assign to Pool"
+            phx-click="open_assign_pool"
+            phx-value-id={@account.identity.id}
+            disabled={@account.identity.status == "deleted"}
+          />
+        </li>
         <li>
           <AdminComponents.dropdown_action_item
             id={"rename-upstream-account-#{@account.identity.id}"}