coder/coderd/database/queries/aicostcontrol.sql

-- name: UpsertAIModelPrices :exec
-- Upsert a batch of (provider, model) rows from a JSON array. Each element
-- must have provider, model, and the four price fields; null prices are
-- written as SQL NULL.
INSERT INTO ai_model_prices (
	provider, model, input_price, output_price, cache_read_price, cache_write_price
)
SELECT
	elem->>'provider',
	elem->>'model',
	(elem->>'input_price')::bigint,
	(elem->>'output_price')::bigint,
	(elem->>'cache_read_price')::bigint,
	(elem->>'cache_write_price')::bigint
FROM jsonb_array_elements(@seed::jsonb) AS elem
ON CONFLICT (provider, model) DO UPDATE SET
	input_price       = EXCLUDED.input_price,
	output_price      = EXCLUDED.output_price,
	cache_read_price  = EXCLUDED.cache_read_price,
	cache_write_price = EXCLUDED.cache_write_price,
	updated_at        = NOW();

-- name: GetAIModelPriceByProviderModel :one
SELECT *
FROM ai_model_prices
WHERE provider = @provider AND model = @model;

-- name: GetGroupAIBudget :one
SELECT *
FROM group_ai_budgets
WHERE group_id = @group_id;

-- name: UpsertGroupAIBudget :one
INSERT INTO group_ai_budgets (group_id, spend_limit_micros)
VALUES (@group_id, @spend_limit_micros)
ON CONFLICT (group_id) DO UPDATE SET
	spend_limit_micros = EXCLUDED.spend_limit_micros,
	updated_at  = NOW()
RETURNING *;

-- name: DeleteGroupAIBudget :one
DELETE FROM group_ai_budgets WHERE group_id = @group_id RETURNING *;

-- name: GetUserAIBudgetOverride :one
SELECT *
FROM user_ai_budget_overrides
WHERE user_id = @user_id;

-- name: UpsertUserAIBudgetOverride :one
INSERT INTO user_ai_budget_overrides (user_id, group_id, spend_limit_micros)
VALUES (@user_id, @group_id, @spend_limit_micros)
ON CONFLICT (user_id) DO UPDATE SET
	group_id           = EXCLUDED.group_id,
	spend_limit_micros = EXCLUDED.spend_limit_micros,
	updated_at         = NOW()
RETURNING *;

-- name: DeleteUserAIBudgetOverride :one
DELETE FROM user_ai_budget_overrides WHERE user_id = @user_id RETURNING *;

-- name: GetHighestGroupAIBudgetByUser :one
-- Returns the highest group AI budget across the groups the user belongs to,
-- breaking ties by the earliest organization membership. Implements the
-- "highest" budget policy. group_members_expanded is a UNION of group_members
-- and organization_members, so the implicit "Everyone" group
-- (group_id == organization_id) is included. Returns no rows when the user has
-- no budgeted groups. Callers should treat sql.ErrNoRows as "no group budget".
SELECT
	budget.group_id,
	budget.spend_limit_micros
FROM group_ai_budgets budget
JOIN group_members_expanded member ON member.group_id = budget.group_id
JOIN organizations ON organizations.id = member.organization_id
JOIN organization_members
	ON organization_members.user_id = member.user_id
	AND organization_members.organization_id = member.organization_id
WHERE member.user_id = @user_id
	AND organizations.deleted = false
ORDER BY
	budget.spend_limit_micros DESC,       -- highest wins
	organization_members.created_at ASC,  -- earliest organization membership
	budget.group_id ASC                   -- deterministic tiebreak
LIMIT 1;

-- name: GetUserEveryoneFallbackGroup :one
-- Returns the "Everyone" group (id == organization_id) to attribute a user's
-- spend to when no override or budgeted group applies. Prefers the default org,
-- then the earliest organization membership. Returns no rows when the user has
-- no organization membership.
SELECT organizations.id AS group_id
FROM organization_members
JOIN organizations ON organizations.id = organization_members.organization_id
WHERE organization_members.user_id = @user_id
	AND organizations.deleted = false
ORDER BY
	organizations.is_default DESC,        -- prefer the default org
	organization_members.created_at ASC,  -- earliest organization membership
	organizations.id ASC                  -- deterministic tiebreak
LIMIT 1;

-- name: IncrementUserAIDailySpend :one
-- Adds cost_micros to the spend for (user_id, effective_group_id, day).
-- The day parameter is normalized to its UTC calendar day before storage.
INSERT INTO ai_user_daily_spend (user_id, effective_group_id, day, spend_micros)
VALUES (@user_id, @effective_group_id, ((@day::timestamptz) AT TIME ZONE 'UTC')::date, @cost_micros)
ON CONFLICT (user_id, effective_group_id, day) DO UPDATE SET
	spend_micros = ai_user_daily_spend.spend_micros + EXCLUDED.spend_micros
RETURNING *;

-- name: GetUserAISpendSince :one
-- Total spend for (user_id, effective_group_id) on or after period_start until NOW.
-- The period_start parameter is normalized to its UTC calendar day.
SELECT
	@user_id::uuid AS user_id,
	@effective_group_id::uuid AS effective_group_id,
	((@period_start::timestamptz) AT TIME ZONE 'UTC')::date AS period_start,
	COALESCE(SUM(spend_micros), 0)::BIGINT AS spend_micros
FROM ai_user_daily_spend
WHERE user_id = @user_id
	AND effective_group_id = @effective_group_id
	AND day >= ((@period_start::timestamptz) AT TIME ZONE 'UTC')::date;

-- name: GetOrganizationGroupsAISpend :many
-- Returns AI spend limits and aggregate spend for groups in @group_ids that
-- belong to @organization_id, on or after period_start until NOW. The spend
-- limit is null when the group has no configured budget.
-- The period_start parameter is normalized to its UTC calendar day.
SELECT
	groups.id AS group_id,
	groups.organization_id AS organization_id,
	budget.spend_limit_micros AS spend_limit_micros,
	COALESCE(SUM(spend.spend_micros), 0)::BIGINT AS current_spend_micros
FROM groups
LEFT JOIN group_ai_budgets budget ON budget.group_id = groups.id
LEFT JOIN ai_user_daily_spend spend
	ON spend.effective_group_id = groups.id
	AND spend.day >= ((@period_start::timestamptz) AT TIME ZONE 'UTC')::date
WHERE groups.organization_id = @organization_id
	AND groups.id = ANY(@group_ids::uuid[])
GROUP BY groups.id, budget.spend_limit_micros
ORDER BY groups.id;

-- name: GetGroupMembersAISpend :many
-- Returns each user's AI spend attributed to the queried group, on or after
-- period_start until NOW. Only current members of the queried group are
-- returned. spend_limit_micros and limit_source are populated only when the
-- queried group is the user's effective budget source. The effective group
-- falls back to the Everyone group, and effective_group_id is null only when
-- that group belongs to a different organization than the queried group.
-- The period_start parameter is normalized to its UTC calendar day.
-- TODO(AIGOV-527): unify effective group resolution in a single place.
WITH queried_group AS (
	-- The queried group's org, used to detect cross-org effective groups.
	SELECT organization_id
	FROM groups
	WHERE id = @group_id
),
filtered_users AS (
	-- Users from @user_ids that are members of the queried group. Uses
	-- group_members_expanded so the implicit Everyone group counts.
	SELECT DISTINCT user_id
	FROM group_members_expanded
	WHERE group_id = @group_id
		AND user_id = ANY(@user_ids::uuid[])
),
user_highest_group AS (
	-- Per user, the highest-limit group they belong to. Uses
	-- group_members_expanded so the implicit Everyone group counts.
	SELECT DISTINCT ON (member.user_id)
		member.user_id,
		budget.group_id,
		budget.spend_limit_micros
	FROM group_ai_budgets budget
	JOIN group_members_expanded member ON member.group_id = budget.group_id
	JOIN organizations ON organizations.id = member.organization_id
	JOIN organization_members
		ON organization_members.user_id = member.user_id
		AND organization_members.organization_id = member.organization_id
	WHERE member.user_id IN (SELECT user_id FROM filtered_users)
		AND organizations.deleted = false
	ORDER BY member.user_id, budget.spend_limit_micros DESC, organization_members.created_at ASC, budget.group_id ASC
),
user_fallback_group AS (
	-- Per user, the Everyone group to fall back to when no override or budgeted
	-- group applies. The Everyone group has id == organization_id. Prefers the
	-- default org, then the earliest organization membership.
	SELECT DISTINCT ON (organization_members.user_id)
		organization_members.user_id,
		organizations.id AS group_id
	FROM organization_members
	JOIN organizations ON organizations.id = organization_members.organization_id
	WHERE organization_members.user_id IN (SELECT user_id FROM filtered_users)
		AND organizations.deleted = false
	ORDER BY organization_members.user_id, organizations.is_default DESC, organization_members.created_at ASC, organizations.id ASC
),
effective AS (
	-- Effective budget per user: a per-user override wins over the highest-limit
	-- group, which wins over the Everyone group fallback.
	SELECT
		filtered_users.user_id,
		COALESCE(override.group_id, user_highest_group.group_id, user_fallback_group.group_id) AS raw_effective_group_id,
		COALESCE(override.spend_limit_micros, user_highest_group.spend_limit_micros) AS spend_limit_micros,
		(CASE
			WHEN override.group_id IS NOT NULL THEN 'user_override'
			WHEN user_highest_group.group_id IS NOT NULL THEN 'group'
		END)::text AS limit_source
	FROM filtered_users
	LEFT JOIN user_ai_budget_overrides override ON override.user_id = filtered_users.user_id
	LEFT JOIN user_highest_group ON user_highest_group.user_id = filtered_users.user_id
	LEFT JOIN user_fallback_group ON user_fallback_group.user_id = filtered_users.user_id
),
applied_budget AS (
	-- The limit and source only for users whose effective budget source is the
	-- queried group.
	SELECT user_id, spend_limit_micros, limit_source
	FROM effective
	WHERE raw_effective_group_id = @group_id
)
-- Spend is aggregated for the queried group, not the user's effective group.
SELECT
	effective.user_id,
	queried_group.organization_id,
	effective_group.id AS effective_group_id,
	applied_budget.spend_limit_micros,
	applied_budget.limit_source,
	COALESCE(SUM(spend.spend_micros), 0)::BIGINT AS group_spend_micros
FROM effective
CROSS JOIN queried_group
LEFT JOIN groups effective_group
	ON effective_group.id = effective.raw_effective_group_id
	AND effective_group.organization_id = queried_group.organization_id
-- A LEFT JOIN leaves spend_limit_micros and limit_source null for users
-- whose effective budget source is not the queried group.
LEFT JOIN applied_budget ON applied_budget.user_id = effective.user_id
LEFT JOIN ai_user_daily_spend spend
	ON spend.user_id = effective.user_id
	AND spend.effective_group_id = @group_id
	AND spend.day >= ((@period_start::timestamptz) AT TIME ZONE 'UTC')::date
GROUP BY
	effective.user_id,
	queried_group.organization_id,
	effective_group.id,
	applied_budget.spend_limit_micros,
	applied_budget.limit_source
ORDER BY effective.user_id;