fix(admin): deduplicate upstream stats rows

Author: masterkain

lib/codex_pooler/admin/stats.ex

@@ -495,27 +495,103 @@ defmodule CodexPooler.Admin.Stats do
     entries_by_identity = Enum.group_by(settlements, & &1.upstream_identity_id)
 
     quota_accounts
-    |> Enum.map(fn account ->
-      entries = Map.get(entries_by_identity, account.upstream_identity_id, [])
+    |> Enum.group_by(& &1.upstream_identity_id)
+    |> Enum.map(fn {upstream_identity_id, accounts} ->
+      entries = Map.get(entries_by_identity, upstream_identity_id, [])
+      canonical_account = canonical_upstream_account(accounts)
 
       %{
-        pool_upstream_assignment_id: account.pool_upstream_assignment_id,
-        upstream_identity_id: account.upstream_identity_id,
-        assignment_label: account.assignment_label,
-        upstream_label: account.upstream_label,
-        status: account.assignment_status,
-        health_status: account.health_status,
-        quota_state: account.state,
+        pool_upstream_assignment_id: single_assignment_id(accounts),
+        upstream_identity_id: upstream_identity_id,
+        assignment_label: shared_account_value(accounts, :assignment_label),
+        upstream_label:
+          shared_account_value(accounts, :upstream_label) || canonical_account.upstream_label,
+        status: aggregate_account_value(accounts, :assignment_status),
+        health_status: aggregate_account_value(accounts, :health_status),
+        quota_state: aggregate_account_value(accounts, :state, :mixed),
+        assignment_count: length(accounts),
         requests: sum_integer(entries, :request_count),
         total_tokens: sum_integer(entries, :total_tokens),
         settled_cost_micros: sum_decimal_integer(entries, :settled_cost_micros)
       }
     end)
     |> Enum.sort_by(fn row ->
-      {-row.total_tokens, -row.requests, row.assignment_label || row.upstream_label || ""}
+      {-row.total_tokens, -row.requests, upstream_table_label(row),
+       row.upstream_identity_id || ""}
     end)
   end
 
+  @spec canonical_upstream_account([map()]) :: map()
+  defp canonical_upstream_account(accounts) do
+    Enum.min_by(accounts, &upstream_account_sort_key/1, fn -> %{} end)
+  end
+
+  @spec single_assignment_id([map()]) :: Ecto.UUID.t() | nil
+  defp single_assignment_id([account]), do: Map.get(account, :pool_upstream_assignment_id)
+  defp single_assignment_id(_accounts), do: nil
+
+  @spec shared_account_value([map()], atom()) :: term() | nil
+  defp shared_account_value(accounts, field) do
+    case distinct_account_values(accounts, field) do
+      [value] -> value
+      _values -> nil
+    end
+  end
+
+  @spec aggregate_account_value([map()], atom(), term()) :: term() | nil
+  defp aggregate_account_value(accounts, field, mixed_value \\ "mixed") do
+    case distinct_account_values(accounts, field) do
+      [value] -> value
+      [] -> nil
+      _values -> mixed_value
+    end
+  end
+
+  @spec distinct_account_values([map()], atom()) :: [term()]
+  defp distinct_account_values(accounts, field) do
+    accounts
+    |> Enum.map(&Map.get(&1, field))
+    |> Enum.reject(&blank_value?/1)
+    |> Enum.uniq()
+  end
+
+  @spec upstream_account_sort_key(map()) :: {integer(), String.t(), String.t()}
+  defp upstream_account_sort_key(account) do
+    {
+      assignment_status_rank(Map.get(account, :assignment_status)),
+      safe_string(Map.get(account, :assignment_label) || Map.get(account, :upstream_label)),
+      safe_string(Map.get(account, :pool_upstream_assignment_id))
+    }
+  end
+
+  @spec assignment_status_rank(term()) :: non_neg_integer()
+  defp assignment_status_rank("active"), do: 0
+  defp assignment_status_rank("pending"), do: 1
+  defp assignment_status_rank("refresh_due"), do: 2
+  defp assignment_status_rank("refreshing"), do: 3
+  defp assignment_status_rank("paused"), do: 4
+  defp assignment_status_rank("refresh_failed"), do: 5
+  defp assignment_status_rank("reauth_required"), do: 6
+  defp assignment_status_rank("disabled"), do: 7
+  defp assignment_status_rank("errored"), do: 8
+  defp assignment_status_rank("deleted"), do: 9
+  defp assignment_status_rank(_status), do: 10
+
+  @spec upstream_table_label(map()) :: String.t()
+  defp upstream_table_label(row) do
+    safe_string(row.assignment_label || row.upstream_label)
+  end
+
+  @spec blank_value?(term()) :: boolean()
+  defp blank_value?(nil), do: true
+  defp blank_value?(value) when is_binary(value), do: String.trim(value) == ""
+  defp blank_value?(_value), do: false
+
+  @spec safe_string(term()) :: String.t()
+  defp safe_string(nil), do: ""
+  defp safe_string(value) when is_binary(value), do: value
+  defp safe_string(value), do: to_string(value)
+
   defp usage_pool_name(entries, pool_names_by_id) do
     entries
     |> Enum.map(& &1.pool_id)

test/codex_pooler/admin/stats_test.exs

@@ -14,6 +14,7 @@ defmodule CodexPooler.Admin.StatsTest do
   alias CodexPooler.Jobs
   alias CodexPooler.Jobs.RuntimeStateCleanupWorker
   alias CodexPooler.Repo
+  alias CodexPooler.Upstreams.Assignments.PoolAssignments
 
   test "build_dashboard/2 returns pool-scoped KPI, table, chart, session, and quota aggregates" do
     scope = owner_scope()
@@ -187,6 +188,91 @@ defmodule CodexPooler.Admin.StatsTest do
            ] = dashboard.tables.upstreams
   end
 
+  test "build_dashboard/2 collapses repeated assignments for the same upstream identity" do
+    scope = owner_scope()
+    first_pool = pool_fixture(%{slug: "stats-shared-upstream-a", name: "Stats Shared A"})
+    second_pool = pool_fixture(%{slug: "stats-shared-upstream-b", name: "Stats Shared B"})
+    third_pool = pool_fixture(%{slug: "stats-shared-label", name: "Stats Shared Label"})
+    %{api_key: first_api_key} = active_api_key_fixture(first_pool)
+    %{api_key: second_api_key} = active_api_key_fixture(second_pool)
+    %{api_key: third_api_key} = active_api_key_fixture(third_pool)
+
+    %{identity: shared_identity, assignment: first_assignment} =
+      upstream_assignment_fixture(first_pool, %{
+        account_label: "Shared account",
+        assignment_label: "Pool A custom label"
+      })
+
+    assert {:ok, second_assignment} =
+             PoolAssignments.create_pool_assignment(second_pool, shared_identity, %{
+               assignment_label: "Pool B custom label",
+               status: "active",
+               health_status: "active",
+               eligibility_status: "eligible"
+             })
+
+    %{identity: same_label_identity, assignment: same_label_assignment} =
+      upstream_assignment_fixture(third_pool, %{
+        account_label: "Shared account",
+        assignment_label: "Pool A custom label"
+      })
+
+    as_of = ~U[2026-01-10 12:00:00.000000Z]
+    occurred_at = ~U[2026-01-10 11:30:00.000000Z]
+
+    insert_timed_usage!(
+      first_pool,
+      first_api_key,
+      first_assignment,
+      shared_identity,
+      occurred_at,
+      40
+    )
+
+    insert_timed_usage!(
+      second_pool,
+      second_api_key,
+      second_assignment,
+      shared_identity,
+      occurred_at,
+      60
+    )
+
+    insert_timed_usage!(
+      third_pool,
+      third_api_key,
+      same_label_assignment,
+      same_label_identity,
+      occurred_at,
+      25
+    )
+
+    assert {:ok, dashboard} = Stats.build_dashboard(scope, %{window: "1h", as_of: as_of})
+
+    assert [
+             %{
+               upstream_identity_id: shared_identity_id,
+               assignment_label: nil,
+               upstream_label: "Shared account",
+               status: "active",
+               assignment_count: 2,
+               requests: 2,
+               total_tokens: 100
+             },
+             %{
+               upstream_identity_id: same_label_identity_id,
+               assignment_label: "Pool A custom label",
+               upstream_label: "Shared account",
+               assignment_count: 1,
+               requests: 1,
+               total_tokens: 25
+             }
+           ] = dashboard.tables.upstreams
+
+    assert shared_identity_id == shared_identity.id
+    assert same_label_identity_id == same_label_identity.id
+  end
+
   test "build_dashboard/2 returns hourly model usage top five plus Other for sub-day windows" do
     %{user: owner} = bootstrap_owner_fixture()
     %{user: admin} = operator_fixture(owner, %{"email" => unique_user_email()})