fix(responses): bill a background response's project only while the key stays on its team

This commit is contained in:
mateo-berri 2026-10-20 00:53:19 -07:00
parent a8ac9678e2
commit ded4f084c5
2 changed files with 53 additions and 2 deletions

View file

@ -110,7 +110,9 @@ class ManagedObjectCreatorAttribution:
A background response also carries its end user, the creating key's project, and the key,
team, and user budgets and per-model budgets, the way the read that billed it before did,
and never names a key-scoped creator as its user. A key that sets no budget of its own
takes its budget tier's, the same rule auth applies.
takes its budget tier's, the same rule auth applies. The project is billed only while the
key is still on the team the object was created under, so one charge never splits across
two teams.
"""
def __init__(self, prisma_client: PrismaClient, poller_name: str) -> None:
@ -128,7 +130,7 @@ class ManagedObjectCreatorAttribution:
"user_api_key_user_id": _creator_user_id(job.created_by),
**_shared_metadata(job, rows),
**({"user_api_key_end_user_id": end_user_id} if end_user_id else {}),
"user_api_key_project_id": getattr(rows.key, "project_id", None),
"user_api_key_project_id": _project_on_creating_team(rows.key, getattr(job, "team_id", None)),
**_key_budget_fields(rows.key),
**_budget_fields(
rows.team,
@ -227,6 +229,10 @@ def _organization_id(row: object) -> str | None:
return cast("str | None", getattr(row, "organization_id", None)) if row is not None else None
def _project_on_creating_team(key: object, creating_team_id: str | None) -> object:
return getattr(key, "project_id", None) if getattr(key, "team_id", None) == creating_team_id else None
def _budget_fields(row: object, spend_key: str, max_budget_key: str, model_max_budget_key: str) -> Mapping[str, object]:
if row is None:
return {}

View file

@ -1517,6 +1517,7 @@ async def test_the_billing_event_carries_the_creating_keys_project_and_every_cre
organization_id=None,
spend=0.0,
max_budget=None,
team_id="team-alpha",
project_id="project-1",
model_max_budget=key_budget,
)
@ -1551,6 +1552,50 @@ async def test_the_billing_event_carries_the_creating_keys_project_and_every_cre
}
@pytest.mark.asyncio
@pytest.mark.parametrize(
("key_team_id", "billed_project_id"),
[("team-alpha", "project-1"), ("team-beta", None)],
ids=["key_still_on_the_creating_team", "key_moved_to_another_team"],
)
async def test_the_billing_event_bills_the_keys_project_only_while_the_key_stays_on_the_creating_team(
key_team_id: str, billed_project_id: str | None
):
table: Final = _InMemoryManagedObjectTable(
rows=(
_StoredResponse(
id="job-moved",
unified_object_id="resp_moved",
status="queued",
created_by="alice",
api_key="hash-alice",
team_id="team-alpha",
),
)
)
billed_logger: Final = _RecordingBilledLogger()
poller: Final = _poller_over(table, billed_logger)
poller.prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
return_value=SimpleNamespace(
key_alias=None,
organization_id=None,
spend=0.0,
max_budget=None,
team_id=key_team_id,
project_id="project-1",
model_max_budget=None,
)
)
await poller.bill_if_finished(response_id="resp_moved", response=_terminal_response())
(billed,) = billed_logger.billed
assert (billed.metadata["user_api_key_team_id"], billed.metadata["user_api_key_project_id"]) == (
"team-alpha",
billed_project_id,
)
_TIER_MODEL_BUDGET: Final = {"gpt-5": {"max_budget": 4.0, "budget_duration": "1d"}}
_KEY_MODEL_BUDGET: Final = {"gpt-5": {"max_budget": 1.0, "budget_duration": "1d"}}