Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
61 changes: 32 additions & 29 deletions apps/arbiter/lib/arbiter/quota.ex
Original file line number Diff line number Diff line change
Expand Up @@ -1032,28 +1032,21 @@ defmodule Arbiter.Quota do
# Each view carries its *own* account's spend and workspace breakdown — two
# accounts in one list are two separate budgets and must not be summed.
#
# The headline `cost_usd` is the sum of the breakdown rather than a second
# pass over the ledger, so `arb quota`'s account total always equals the
# per-workspace line printed under it, and one view costs one ledger read
# per workspace instead of two.
# The headline `cost_usd` is `provider_spend/1` read straight off the
# account's `usage_events`, NOT a sum of the per-workspace breakdown below:
# `account_fields/3`'s `workspaces` list is built from `workspace_spend/1`,
# which only sums rows carrying a `workspace_id`, so a probe/pre-flight row
# (`workspace_id: nil`, always a `provider_account_id`) would silently drop
# out of a summed total — the exact under-reporting bias bd-adyhvn measured.
# The total can therefore be *larger* than the sum of the breakdown lines
# printed under it; that gap is exactly the account's workspace-less spend.
defp decorate_view(view, cache) do
fields = account_fields(view.provider_account_id, view.provider, cache)
total = cost_for(view.provider, provider_spend(view.provider_account_id))

view
|> Map.merge(fields)
|> Map.put(:cost_usd, total_cost(fields.workspaces))
end

# `nil` (not `0.0`) when no workspace on the account has attributable spend,
# matching `cost_for/2` — the UI shows "—" rather than a misleading "$0.00".
defp total_cost(workspaces) do
workspaces
|> Enum.map(& &1.cost_usd)
|> Enum.filter(&is_number/1)
|> case do
[] -> nil
costs -> costs |> Enum.sum() |> Float.round(6)
end
|> Map.put(:cost_usd, total)
end

@doc """
Expand All @@ -1063,22 +1056,32 @@ defmodule Arbiter.Quota do
("claude" / "gemini" / "openai"); `cost_for/2` maps quota codes onto it.
Returns `%{}` on any error so cost is a best-effort add-on, never a failure.

"How much of this plan did I spend?" is an account question (§8), so this
is the sum over every workspace metered under the account. `usage_events`
does not carry `provider_account_id` until P9, so the account total is
reached through the workspace link rather than read off the ledger row —
`workspace_spend/1` is the per-workspace term §6's breakdown line prints.
"How much of this plan did I spend?" is an account question (§8), read
straight off `usage_events.provider_account_id` (P9) rather than summed
through the workspace link — `workspace_spend/1` sums *only* rows that
carry a `workspace_id`, so a probe/pre-flight row (`workspace_id: nil`, but
always a `provider_account_id` — §8's seam with bd-adyhvn) would silently
drop out of the total, reintroducing the exact under-reporting bias
bd-adyhvn measured. `workspace_spend/1` is still the per-workspace term
§6's breakdown line prints, since that one *is* workspace-scoped by
definition.
"""
@spec provider_spend(String.t() | nil) :: %{optional(String.t()) => float()}
def provider_spend(account_id) do
account_id
|> Resolver.workspace_ids()
|> Enum.map(&workspace_spend/1)
|> Enum.reduce(%{}, fn spend, acc ->
Map.merge(acc, spend, fn _provider, a, b -> a + b end)
end)
def provider_spend(nil), do: %{}

def provider_spend(account_id) when is_binary(account_id) do
since = DateTime.utc_now() |> DateTime.add(-@cost_window_days * 86_400, :second)

case Arbiter.Usage.summarize(by: :provider, since: since, provider_account_id: account_id) do
{:ok, rows} -> Map.new(rows, &{&1.group, &1.total_cost_usd})
_ -> %{}
end
rescue
_ -> %{}
end

def provider_spend(_), do: %{}

@doc """
Build the `t:spend_cache/0` memo for `accounts`: one `workspace_spend/1`
scan per distinct workspace metered under them.
Expand Down
108 changes: 22 additions & 86 deletions apps/arbiter/lib/arbiter/usage.ex
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,6 @@ defmodule Arbiter.Usage do

use Ash.Domain

alias Arbiter.Accounts.Resolver
alias Arbiter.Tasks.Dependency
alias Arbiter.Usage.Estimate
alias Arbiter.Usage.Event
Expand Down Expand Up @@ -107,24 +106,29 @@ defmodule Arbiter.Usage do
# `campaign` was the old name for the `epic` grouping. Accepted as a
# deprecated alias for one release; normalized to `:epic` before validation
# so every caller (CLI, REST, MCP) gets the same grouping.
@deprecated_by %{campaign: :epic}
#
# `account` is not deprecated — it's the spelling the design doc (§8) and
# task use for this grouping — but it rides the same alias mechanism as
# `campaign` so `--by account` normalizes to `:provider_account` before
# validation, same as every other caller.
@deprecated_by %{campaign: :epic, account: :provider_account}

@doc """
Roll up usage events into a list of summary rows.

## Options

* `:by` — one of `#{inspect(@valid_by)}` (`:campaign` also accepted as a
deprecated alias for `:epic`). Required.
deprecated alias for `:epic`; `:account` accepted as an alias for
`:provider_account`, the spelling `docs/provider-account-design.md`
§8 and the CLI/REST docs use). Required.
* `:since` — `%DateTime{}` filter on `occurred_at`. Optional.
* `:workspace_id` — restrict to one workspace. Optional.
* `:provider_account_id` — restrict to one provider account
(`docs/provider-account-design.md` §3.3). `usage_events` carries no
account column until P9, so this narrows the query to the account's
*workspaces* — which over-selects, because a workspace is metered under
a different account per provider. With `by: :provider_account` the
result is then exact (every other account's group is dropped); with any
other grouping it is the workspace-set approximation. Optional.
(`docs/provider-account-design.md` §3.3, §8). Filters
`usage_events.provider_account_id` directly (P9), so it is exact with
every grouping — including rows with no `workspace_id` at all (a probe
or pre-flight row has no workspace but always has an account). Optional.
* `:session_ids` — restrict to a list of `session_id` values, pushed into
the query as `session_id in ^ids` rather than filtered after the read.
`Event` indexes `:session_id`, so this keeps a `:by :session` rollup for
Expand Down Expand Up @@ -166,7 +170,6 @@ defmodule Arbiter.Usage do
{:ok,
events
|> group_events(by)
|> exact_account_groups(by, opts)
|> Enum.map(&aggregate_group(by, &1))
|> sort_rollups(by)
|> maybe_limit(opts)}
Expand Down Expand Up @@ -308,23 +311,6 @@ defmodule Arbiter.Usage do
end
end

# `:provider_account_id` narrows the *query* to the account's workspaces,
# which is as exact as SQL can be while `usage_events` carries no account
# column. It is not exact enough on its own: a workspace is metered under a
# different account per provider, so its Claude rows come back for a Codex
# account too — measured on the live install, where a Codex account's 5h
# overage figure was $125 of Claude spend. The `:provider_account`
# grouping *is* per-event exact, so when both are given, drop every group
# but the one asked for.
defp exact_account_groups(groups, :provider_account, opts) do
case Keyword.get(opts, :provider_account_id) do
id when is_binary(id) and id != "" -> Map.take(groups, [id])
_ -> groups
end
end

defp exact_account_groups(groups, _by, _opts), do: groups

defp base_filter(query, opts) do
query
|> filter_since(Keyword.get(opts, :since))
Expand All @@ -344,15 +330,11 @@ defmodule Arbiter.Usage do
defp filter_workspace_id(query, ""), do: query
defp filter_workspace_id(query, ws), do: Ash.Query.filter(query, workspace_id == ^ws)

# `usage_events` carries no `provider_account_id` column until P9, so an
# account filter is the set of workspaces metered under it. An account with
# no workspaces matches nothing, which is the honest answer — not
# "everything".
defp filter_provider_account_id(query, nil), do: query
defp filter_provider_account_id(query, ""), do: query

defp filter_provider_account_id(query, account_id),
do: Ash.Query.filter(query, workspace_id in ^Resolver.workspace_ids(account_id))
do: Ash.Query.filter(query, provider_account_id == ^account_id)

defp filter_session_ids(query, nil), do: query
defp filter_session_ids(query, []), do: query
Expand Down Expand Up @@ -391,25 +373,15 @@ defmodule Arbiter.Usage do
defp group_events(events, :workspace),
do: Enum.group_by(events, &(&1.workspace_id || "(none)"))

# §5 row 14: the account rollup. Resolved through the workspace → account
# join rather than read off the row, because `usage_events` does not carry
# `provider_account_id` until P9. An event is attributed to the account its
# workspace is metered under **for that event's provider** — a workspace can
# hold a different account per provider — and rows with no resolvable
# account fall into the `(none)` sentinel rather than vanishing.
# §5 row 14, §8: the account rollup, read straight off
# `usage_events.provider_account_id` (P9). This is what makes probe/
# pre-flight rows count — they carry no `workspace_id` but always carry
# `provider_account_id` (§8's seam with bd-adyhvn) — dropping them here
# would reintroduce the exact under-reporting bias bd-adyhvn measured.
# Rows with no resolvable account fall into the `(none)` sentinel rather
# than vanishing.
defp group_events(events, :provider_account) do
index = account_index(events)
codes = provider_codes(events)
fallbacks = default_provider_codes(events, codes)

Enum.group_by(events, fn ev ->
code = Map.get(codes, ev.provider) || Map.get(fallbacks, ev.workspace_id)

index
|> Map.get(ev.workspace_id, %{})
|> Map.get(code)
|> Kernel.||("(none)")
end)
Enum.group_by(events, &(&1.provider_account_id || "(none)"))
end

defp group_events(events, :repo), do: Enum.group_by(events, &(&1.repo || "(none)"))
Expand All @@ -430,42 +402,6 @@ defmodule Arbiter.Usage do
end)
end

# `%{workspace_id => %{provider_code => account_id}}` for every workspace the
# window touches — one read, rather than one per event.
defp account_index(events) do
events
|> Enum.map(& &1.workspace_id)
|> Enum.reject(&(is_nil(&1) or &1 == ""))
|> Enum.uniq()
|> Map.new(&{&1, Resolver.account_ids(&1)})
end

# Ledger providers ("claude" / "openai" / "gemini") are not the quota
# provider codes the join rows use, and resolving "gemini" probes the PATH
# (`Arbiter.Quota.provider_code/1`), so map each distinct value once.
defp provider_codes(events) do
events
|> Enum.map(& &1.provider)
|> Enum.uniq()
|> Map.new(&{&1, Arbiter.Quota.provider_code(&1)})
end

# `usage_events.provider` is nullable, and a row that records no provider
# (or one with no tracked quota) would otherwise fall out of every account —
# silently under-reporting the overage figure `Arbiter.Quota.Overage` alerts
# on, which pre-P7 counted every row the workspace had. Attribute it to the
# account the workspace actually dispatches on, the same provider the gate
# reads its snapshot for. Only workspaces that have such a row pay for the
# lookup.
defp default_provider_codes(events, codes) do
events
|> Enum.filter(&is_nil(Map.get(codes, &1.provider)))
|> Enum.map(& &1.workspace_id)
|> Enum.reject(&(is_nil(&1) or &1 == ""))
|> Enum.uniq()
|> Map.new(&{&1, Arbiter.Quota.provider_code(Arbiter.Quota.default_provider(&1))})
end

defp task_attributed?(ev), do: is_binary(ev.task_id) and ev.task_id != ""
defp session_attributed?(ev), do: is_binary(ev.session_id) and ev.session_id != ""

Expand Down
59 changes: 32 additions & 27 deletions apps/arbiter/test/arbiter/quota/account_wide_hold_test.exs
Original file line number Diff line number Diff line change
Expand Up @@ -65,12 +65,19 @@ defmodule Arbiter.Quota.AccountWideHoldTest do
})
end

defp usage_event!(ws_id, cost, provider \\ "claude") do
# P10 (bd-icwk2k): `usage_events.provider_account_id` is a real column
# since P9, populated at write time by `Arbiter.Worker`'s own
# `AccountResolver.account_id(workspace_id, provider)` call — this mirrors
# that so these fixtures land exactly where the production write path
# would put them, rather than relying on `Usage.summarize/1` to infer it
# from the workspace join at read time (the pre-P9 approximation).
defp usage_event!(ws_id, cost, provider \\ "claude", account_id \\ nil) do
Ash.create!(Event, %{
workspace_id: ws_id,
task_id: "bd-p7-#{System.unique_integer([:positive])}",
step: :work,
provider: provider,
provider_account_id: account_id,
cost_usd: cost,
occurred_at: DateTime.utc_now()
})
Expand Down Expand Up @@ -243,8 +250,8 @@ defmodule Arbiter.Quota.AccountWideHoldTest do
link!(a, account)
link!(b, account)

usage_event!(a.id, 1.25)
usage_event!(b.id, 2.75)
usage_event!(a.id, 1.25, "claude", account.id)
usage_event!(b.id, 2.75, "claude", account.id)

assert_in_delta Overage.windowed_spend(account, nil), 4.0, 0.0001
assert_in_delta Overage.windowed_spend(account.id, nil), 4.0, 0.0001
Expand All @@ -258,8 +265,8 @@ defmodule Arbiter.Quota.AccountWideHoldTest do
link!(mine, account)
link!(theirs, other)

usage_event!(mine.id, 1.0)
usage_event!(theirs.id, 9.0)
usage_event!(mine.id, 1.0, "claude", account.id)
usage_event!(theirs.id, 9.0, "claude", other.id)

assert_in_delta Overage.windowed_spend(account, nil), 1.0, 0.0001
end
Expand All @@ -276,8 +283,8 @@ defmodule Arbiter.Quota.AccountWideHoldTest do
link!(ws, claude)
link_codex!(ws, codex)

usage_event!(ws.id, 100.0, "claude")
usage_event!(ws.id, 7.0, "openai")
usage_event!(ws.id, 100.0, "claude", claude.id)
usage_event!(ws.id, 7.0, "openai", codex.id)

assert_in_delta Overage.windowed_spend(codex, nil), 7.0, 0.0001
assert_in_delta Overage.windowed_spend(claude, nil), 100.0, 0.0001
Expand All @@ -297,8 +304,8 @@ defmodule Arbiter.Quota.AccountWideHoldTest do
link!(a, account)
link!(b, account)

usage_event!(a.id, 1.0)
usage_event!(b.id, 2.0)
usage_event!(a.id, 1.0, "claude", account.id)
usage_event!(b.id, 2.0, "claude", account.id)

assert {:ok, rows} = Usage.summarize(by: :provider_account)
row = Enum.find(rows, &(&1.group == account.id))
Expand All @@ -314,8 +321,8 @@ defmodule Arbiter.Quota.AccountWideHoldTest do
link!(mine, account)
link!(theirs, other)

usage_event!(mine.id, 1.0)
usage_event!(theirs.id, 9.0)
usage_event!(mine.id, 1.0, "claude", account.id)
usage_event!(theirs.id, 9.0, "claude", other.id)

assert {:ok, rows} = Usage.summarize(by: :provider_account, provider_account_id: account.id)
assert [%{group: group, rows: 1}] = rows
Expand All @@ -329,8 +336,8 @@ defmodule Arbiter.Quota.AccountWideHoldTest do
link!(ws, claude)
link_codex!(ws, codex)

usage_event!(ws.id, 100.0, "claude")
usage_event!(ws.id, 7.0, "openai")
usage_event!(ws.id, 100.0, "claude", claude.id)
usage_event!(ws.id, 7.0, "openai", codex.id)

assert {:ok, [%{group: group, rows: 1, total_cost_usd: cost}]} =
Usage.summarize(by: :provider_account, provider_account_id: codex.id)
Expand All @@ -339,22 +346,20 @@ defmodule Arbiter.Quota.AccountWideHoldTest do
assert_in_delta cost, 7.0, 0.0001
end

test "a row with no recorded provider falls back to the workspace's default provider" do
# `usage_events.provider` is nullable, and pre-P7 `windowed_spend/2`
# counted every row the workspace had regardless of provider. Dropping
# such a row from the account rollup would silently under-report
# overage and suppress the alert, so it is attributed to the account
# the workspace actually dispatches on — the same provider the gate
# reads the snapshot for.
test "a probe row with no workspace_id still lands in its account's group (§8)" do
# §8's seam with bd-adyhvn: a probe/pre-flight row has no workspace but
# always has an account, because it is issued *as* a credential.
# `usage_events.provider_account_id` (P9) is what makes this possible —
# there is no workspace to join through at all.
account = account!()
ws = workspace!()
link!(ws, account)

Ash.create!(Event, %{
workspace_id: ws.id,
task_id: "bd-p7-noprov",
step: :work,
provider: nil,
workspace_id: nil,
task_id: nil,
source: :preflight,
step: :other,
provider: "claude",
provider_account_id: account.id,
cost_usd: 3.0,
occurred_at: DateTime.utc_now()
})
Expand All @@ -366,7 +371,7 @@ defmodule Arbiter.Quota.AccountWideHoldTest do
assert_in_delta Overage.windowed_spend(account, nil), 3.0, 0.0001
end

test "a row whose workspace has no account lands in the (none) sentinel" do
test "a row with no provider_account_id lands in the (none) sentinel" do
usage_event!(workspace!().id, 1.0)

assert {:ok, rows} = Usage.summarize(by: :provider_account)
Expand Down
Loading
Loading