Add usage limits
- Dominant language
- Python
- Stars
- 183
- Forks
- 41
- Avg merge
- 8h 6m
- Merged PRs (30d)
- 2
Description
## Overview
Loom currently has no concept of usage limits — model access is either globally enabled/disabled (`enabled_model_ids`) or allow-listed per agent (`Agent.allowed_model_ids`), but nothing constrains *how much* a user or group can consume. Alongside the more granular cost reporting work (issue #21), add usage limits: administrators define limits scoped to a user or group, for a specific model or model family, measured in token count or budget ($), each with one of three enforcement modes — warn (non-blocking), throttle (slow down), or block (stop entirely).
## Context
### Current State
**No usage-limit, rate-limit, or quota mechanism exists anywhere in the backend**
- Grep for `rate_limit|throttle|quota` across `backend/app/` only turns up AWS SDK error-code translations (`backend/app/routers/registry.py` L68, `backend/app/routers/memories.py` L175/177 — `ServiceQuotaExceededException`/`ThrottledException` mapped to HTTP 429) — these just pass through upstream AWS throttling, they are not a limiting mechanism Loom implements itself. This feature introduces the concept from scratch.
**Model identity is not persisted per invocation today**
- `Invocation` (`backend/app/models/invocation.py` L1-101) has token/cost columns (`input_tokens`, `output_tokens`, `estimated_cost`, etc.) but **no `model_id` column at all**.
- `Agent.allowed_model_ids` (`backend/app/models/agent.py` L53, helpers L106-117) is a JSON array of model IDs allowed for that agent — a per-agent allow-list, not a per-invocation record.
- The invoke endpoint validates `request_body.model_id` against `agent.get_allowed_model_ids()` (`backend/app/routers/invocations.py` L1420-1429) but the resolved model ID is **never written onto the `Invocation` row** created shortly after (L1472-1478). A usage-limit feature scoped "per model family or specific model" (R2) requires adding a durable `model_id` column to `Invocation` first — there is currently no way to answer "how many tokens has this user spent on model X" at all.
- There is no persisted "model family" taxonomy. The closest analog is provider grouping in `backend/app/services/model_catalog.py` (`_normalize_model_id()` L60-68, `provider` field `"bedrock"`/`"litellm"` L230/419) — this is presentation-layer grouping, not a stored family concept. "Model family" scoping (R2) needs its own definition (e.g. by `litellm_provider`, or a model-ID prefix pattern) since none exists today.
- Global model enable/disable lives in a single flat `SiteSetting` key `enabled_model_ids` (`backend/app/routers/settings.py` L29, read via `get_enabled_model_ids()` L528-530) — a deployment-wide allow/deny set with no group/user scoping and no numeric threshold, structurally unrelated to what this issue needs.
**Single choke point exists for enforcement, but real spend is only known after the fact**
- `POST /{agent_id}/invoke` (`backend/app/routers/invocations.py` L1364-1481) is the one and only HTTP endpoint that creates `InvocationSession`/`Invocation` rows for agent invocations — there is no separate invoke path per agent runtime/harness/MCP tool call at the Loom-backend layer. By the time this handler validates `model_id` (L1420-1429) and checks group-based access using `user.groups` (L1393-1418), all the context a limit check needs (user, groups, agent, model) is already resolved in one place — this is the natural pre-flight point for `block`/`warn`/`throttle` checks, before `db.add(invocation)` and before the runtime call.
- The backend does not proxy model tokens itself — it invokes the deployed agent runtime via boto3 (`invoke_agent`/`invoke_agent_ws`, `backend/app/services/agentcore.py` L79/163/203); the agent runtime itself calls the model, typically through the LiteLLM proxy using a per-agent virtual key (`backend/app/services/litellm.py` L1-20).
- `input_tokens`/`estimated_cost` are computed only after the invocation completes (around `invocations.py` L1340-1361), and even then they're estimates — actual cost is reconciled later, out of band, by `backend/app/services/usage_poller.py` (`_poll_once()` L32-152, a plain `asyncio` loop polling every `POLL_INTERVAL_SECONDS = 600`, L159-170; note its `start_usage_poller()` entry point could not be confirmed as wired into `backend/app/main.py`'s lifespan — verify this actually runs before relying on it for reconciliation). **Implication: precise, real-time mid-stream throttling against exact spend is not achievable** — enforcement must work off the running total from completed (estimated, then reconciled) invocations, not a live token counter.
- No APScheduler/Celery exists — the `usage_poller.py` `asyncio`-loop pattern is the only precedent for periodic aggregation (e.g. recomputing a user's or group's trailing-window spend), and a usage-limit aggregation job should follow the same shape.
**Existing policy-row pattern to follow for admin configuration**
- `ApprovalPolicy` (`backend/app/models/approval_policy.py` L8-45) is a strong structural precedent: a named policy with a `policy_type` discriminator, an `approval_mode` action enum (`require_approval`/`notify_only`), a numeric `timeout_seconds` threshold, and — most relevant — a JSON `agent_scope` column with a type discriminator (`{"type": "all"}` / `{"type": "specific", "agent_ids": [...]}` / `{"type": "tag_filter", ...}`, L16). This scope-by-discriminator shape maps directly onto what a usage-limit policy needs: scope (user vs. group), measure (tokens vs. budget), model target (specific model vs. family), and enforcement action (warn/throttle/block).
- Frontend precedent: `frontend/src/components/ApprovalPolicyPanel.tsx`, rendered as a tab inside `frontend/src/pages/SecurityAdminPage.tsx` (L6, L55), using shadcn `Select` dropdowns for enum fields (L211-223, L227-237) and numeric inputs for thresholds (L246, L266). `frontend/src/pages/TaggingPage.tsx` (L14-142) shows the CRUD-list-with-inline-edit pattern (`listTagPolicies`/`createTagPolicy`/`updateTagPolicy`/`deleteTagPolicy`) a Usage Limits admin panel should mirror.
- Identity for scoping: no persisted `User`/`Group` model exists — `UserInfo` (`backend/app/dependencies/auth.py` L129-149: `sub`, `username`, `groups`, `scopes`, `idp_type`) is resolved fresh per request from JWT claims. Any user/group-scoped limit must reference `username`/`groups` values as they appear in `UserInfo`, not a foreign key into a persisted identity table.
### Key Files
- `backend/app/models/invocation.py` — needs a `model_id` column added; the substrate all usage aggregation reads from
- `backend/app/routers/invocations.py` L1364-1481 — the single invoke choke point where limit checks (block/warn) and throttling must be inserted
- `backend/app/models/agent.py` L53, L106-117 — existing `allowed_model_ids` pattern for per-agent model scoping
- `backend/app/services/model_catalog.py` — provider/model metadata; needs a "family" grouping definition
- `backend/app/services/usage_poller.py` — the only existing periodic-aggregation precedent (`asyncio` loop, 10-min interval)
- `backend/app/models/approval_policy.py`, `backend/app/routers/security.py` (or wherever `ApprovalPolicy` CRUD lives) — the policy-row/scope-discriminator pattern to mirror
- `frontend/src/components/ApprovalPolicyPanel.tsx`, `frontend/src/pages/SecurityAdminPage.tsx`, `frontend/src/pages/TaggingPage.tsx` — frontend CRUD/admin-panel precedent
- `backend/app/dependencies/auth.py` — `UserInfo` (`username`, `groups`) as the identity a limit's scope binds to
## Requirements
### R1: Administrators can set limits at the group or user level
- Add a `UsageLimit` (or similarly named) policy model, following the `ApprovalPolicy` scope-discriminator pattern (`approval_policy.py` L16): a `scope` JSON column of the shape `{"type": "user", "username": ...}` or `{"type": "group", "group": ...}`, since there is no persisted `User`/`Group` table to foreign-key against — scope values must match `UserInfo.username`/`UserInfo.groups` (`auth.py` L129-149) as asserted by the IdP at request time
- Add admin CRUD endpoints and a frontend panel (new tab on `SecurityAdminPage.tsx`, alongside `ApprovalPolicyPanel.tsx`) for creating/editing/deleting usage limits
- Decide and document limit precedence when multiple limits apply to the same request (e.g. a user-level limit and a group-level limit both matching) — most-restrictive-wins is the safer default, but this must be explicit, not incidental
- Decide how a limit is evaluated when a user belongs to multiple groups with different limits — same precedence question applies
### R2: Administrators can set limits per model family or specific model
- Add a `model_id` column to `Invocation` (currently absent per Context) and populate it from the already-validated `request_body.model_id` in the invoke endpoint (`invocations.py` L1420-1429/1472-1478) — without this, there is no way to attribute token/cost usage to a specific model at all
- Define "model family" explicitly for this feature (e.g. grouping by `litellm_provider` or a model-ID prefix scheme from `model_catalog.py`), since no such taxonomy is persisted today — do not assume one exists
- The `UsageLimit` model target should be a discriminated field mirroring the scope pattern in R1: `{"type": "model", "model_id": ...}` or `{"type": "family", "family": ...}`
- A limit with no model target specified should apply across all models for the given user/group scope — make this default explicit in the policy schema
### R3: Administrators can set limits by token count or by budget
- The `UsageLimit` model needs a `measure` field (`"tokens"` or `"budget"`) and a numeric `threshold`, plus a time window (e.g. daily/weekly/monthly/rolling) — none of this exists today; define the window semantics explicitly since usage resets/rolling windows affect both enforcement and how the aggregation job (R4) computes running totals
- Token-count limits should sum `input_tokens`/`output_tokens` from `Invocation` (existing columns); budget limits should sum `estimated_cost` (existing, estimate-only per Context) and be reconciled against actual cost once `usage_poller.py` updates a row's `cost_source` — document that near-real-time budget enforcement necessarily operates on estimates, with actuals only available after asynchronous reconciliation
- Support per-user and per-group limits independently choosing token vs. budget measurement (a group could be budget-capped while a specific user within it is token-capped)
### R4: Enforcement actions — warn, throttle, and block
- Add a pre-flight usage-limit check in the invoke endpoint (`invocations.py`, before `db.add(invocation)` around L1472), evaluated against the running total for the matching `UsageLimit`(s) resolved from the authenticated `UserInfo` and requested `model_id`
- **Warn**: non-blocking — allow the invocation to proceed, but surface a warning (e.g. in the SSE response payload or via a dashboard notification) that the user/group is approaching or has exceeded a warn-tier threshold
- **Throttle**: since real-time token metering mid-invocation isn't available (Context), implement throttling as an added delay before dispatching the runtime call (or a reduced concurrency ceiling) once a user/group crosses the throttle threshold, rather than attempting to slow an in-flight model stream
- **Block**: reject the invoke request with an HTTP 4xx (429, consistent with the existing `ServiceQuotaExceededException`/`ThrottledException` → 429 convention in `registry.py`/`memories.py`) before any runtime call or `Invocation` row is created
- Running-total computation needs a periodic aggregation job following the `usage_poller.py` `asyncio`-loop pattern (confirm first whether `usage_poller.py` is actually wired into `main.py`'s lifespan today, per Context — if not, this is also the moment to fix that) rather than recomputing full sums synchronously on every invoke call, to keep the pre-flight check cheap
- Document exactly which enforcement tier takes precedence when a request matches limits at multiple tiers simultaneously (e.g. within warn range for tokens but over the block threshold for budget) — block should win over throttle, which should win over warn, regardless of which measure triggered it
## Testing
- Run backend tests: `cd backend && make test`
- Add unit tests for `UsageLimit` scope/model-target/measure resolution, including precedence when multiple limits match (R1)
- Add tests verifying `model_id` is correctly persisted on new `Invocation` rows and that model-family grouping resolves as defined in R2
- Add tests for each enforcement tier: warn returns success with a warning indicator, throttle adds the expected delay without blocking, block returns 429 and creates no `Invocation` row
- Add a test for the aggregation job's running-total computation across a rolling/reset window boundary
- Manually verify: configure a low token limit for a test user, exhaust it, and confirm warn → throttle → block behavior transitions as expected in the invoke flow and is reflected in the admin usage-limits panel
## Out of Scope
- Real-time, mid-stream token metering or exact live spend tracking — enforcement operates on the most recent completed/reconciled invocation totals, not a live counter (per Context)
- A persisted `User`/`Group` table — this issue scopes limits against `UserInfo.username`/`UserInfo.groups` as asserted by the IdP per request, consistent with how group-based access is already checked in the invoke endpoint
- Per-agent usage limits (distinct from per-user/per-group) — `Agent.allowed_model_ids` already provides agent-level model restriction; combining that with usage-limit scoping is a possible future extension, not required here
- Automatic remediation beyond the three defined actions (e.g. auto-upgrading a user's tier, auto-notifying billing systems) — warn/throttle/block are the only enforcement behaviors this issue implements
Contributor guide
Research direction
Start by reading backend/app/routers/invocations.py, backend/app/models/invocation.py, backend/app/models/approval_policy.py, and backend/app/services/usage_poller.py, then inspect the related frontend panels. Map the unresolved decisions around model families, time windows, precedence, estimates, and enforcement before implementation. Done means the policy model, admin CRUD, invocation attribution, aggregation, and warn/throttle/block behavior satisfy R1-R4.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- aws, python, typescript
- Domain
- api, backend, databases, frontend
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 25/100