diff --git a/.env.example b/.env.example index 027f114..476d487 100644 --- a/.env.example +++ b/.env.example @@ -68,6 +68,7 @@ OPENAI_API_KEY= # TPK_DB_WAIT_SECONDS=60 # [db].wait_seconds serve: DB connect retry budget # TPK_SESSION_TTL=86400 # [server].session_ttl login session lifetime (s) # TPK_CHAT_AUDIT=1 # [server].chat_audit false/0 disables chat auditing +# TPK_DAILY_TOKEN_LIMIT=500000 # [server].daily_token_limit global fallback non-admin per-user/day cap; 0 = unlimited # TPK_CHECKOUT_DIR=~/.tpk/checkouts # [server].checkout_dir github checkout cache # TPK_EXTRACTION_BACKEND= # [llm].backend auto | claude | openai # TPK_EXTRACTION_MODEL= # [llm].model semantic-extraction model diff --git a/README.md b/README.md index a04f332..ceea7bd 100644 --- a/README.md +++ b/README.md @@ -76,6 +76,7 @@ export the matching env var — whichever suits your deployment. Secrets are | `TPK_CHECKOUT_DIR` | `[server].checkout_dir` | `~/.tpk/checkouts` | GitHub checkout cache | | `TPK_SESSION_TTL` | `[server].session_ttl` | `86400` | Login session lifetime (seconds) | | `TPK_CHAT_AUDIT` | `[server].chat_audit` | `true` | Chat Q&A auditing (`false`/`0` disables) | +| `TPK_DAILY_TOKEN_LIMIT` | `[server].daily_token_limit` | `500000` | Global fallback per-user daily token budget for non-admins (`0` = unlimited; a per-user or role `daily_token_limit` wins) | | `TPK_EXTRACTION_BACKEND` | `[llm].backend` | `auto` | Semantic-extraction backend (`auto`\|`claude`\|`openai`) | | `TPK_EXTRACTION_MODEL` | `[llm].model` | backend default | Semantic-extraction model | diff --git a/deploy/docker/repos.container.toml b/deploy/docker/repos.container.toml index 319208a..3980802 100644 --- a/deploy/docker/repos.container.toml +++ b/deploy/docker/repos.container.toml @@ -45,6 +45,7 @@ token_budget = 16000 # checkout_dir = "/opt/tpk/.checkouts" # env: TPK_CHECKOUT_DIR # session_ttl = 86400 # env: TPK_SESSION_TTL # chat_audit = true # env: TPK_CHAT_AUDIT +# daily_token_limit = 500000 # env: TPK_DAILY_TOKEN_LIMIT (global fallback, non-admin per-user/day; 0 = unlimited) [repos.proton-enterprise] github = "timeplus-io/proton-enterprise" diff --git a/deploy/k8s/README.md b/deploy/k8s/README.md index 5023fc5..8a9d8b2 100644 --- a/deploy/k8s/README.md +++ b/deploy/k8s/README.md @@ -99,6 +99,22 @@ Or, to remap the port, a headless Service + an `EndpointSlice` pointing at the DB's IP on its real port, exposed as `port: 8123`. The DB user in `TIMEPLUS_USER` must already exist on your timeplusd — `app-only.yaml` does not provision users. +**The `tpk` database.** All tpk streams live under a dedicated database (`tpk` +by default, `TIMEPLUS_DATABASE`), which the app creates on startup (`CREATE +DATABASE IF NOT EXISTS`). Against a shared, externally-managed Timeplus this is +the one extra grant to check: the `TIMEPLUS_USER` needs **CREATE DATABASE** (the +first time) plus read/write on that database. If that user isn't allowed to +create databases, pre-create it and grant access, then keep `TIMEPLUS_DATABASE` +pointed at it: + +```sql +CREATE DATABASE IF NOT EXISTS tpk; +-- grant your tpk user read/write on tpk (per your Timeplus access model) +``` + +This only applies to app-only: `enterprise.yaml` and `allinone.yaml` provision a +`tpk` user with full privileges, so database creation just works there. + ## 3. Build the knowledge graph (ingest) The corpus is defined in the `repos.toml` baked into the image. Run ingest once @@ -169,6 +185,7 @@ manifests, others stubbed as commented-out examples): | `TPK_AGENT_PROVIDER` / `TPK_AGENT_MODEL` | chat-agent backend + model | | `TPK_EXTRACTION_BACKEND` | `tpk ingest` semantic backend (`openai`\|`claude`\|`auto`) | | `TPK_DB_BACKEND` | `timeplusd` (Enterprise, mutable streams) or `proton` | +| `TIMEPLUS_DATABASE` | database all tpk streams live under (default `tpk`; app-only: user needs CREATE DATABASE) | | `TPK_DB_WAIT_SECONDS` | how long the app waits for the DB on boot | ### Custom corpus (`repos.toml`) via ConfigMap diff --git a/deploy/k8s/app-only.yaml b/deploy/k8s/app-only.yaml index 3d109e8..d2440b0 100644 --- a/deploy/k8s/app-only.yaml +++ b/deploy/k8s/app-only.yaml @@ -1,7 +1,11 @@ # App-only deployment: the tpk `app` (chat agent + web UI + ingest + MCP) # ALONE, connecting to a Timeplus Enterprise (timeplusd) that is ALREADY # running — deployed and managed separately (its own Helm chart / operator / -# existing cluster). Nothing here provisions or runs a database. +# existing cluster). This does not run a database *server*; but on startup the +# app DOES create its own database (`tpk` by default) on that server and puts +# all its streams there. So the DB user you point it at needs CREATE DATABASE +# (first run only) plus read/write on that database. If the user is restricted, +# pre-create the database and grant it, then set TIMEPLUS_DATABASE (see below). # # Use this when you already operate Timeplus Enterprise and just want to add the # knowledge agent against it. If you want tpk to bring up its own timeplusd too, @@ -84,6 +88,13 @@ spec: # Enterprise timeplusd -> `timeplusd` backend (mutable streams). - name: TPK_DB_BACKEND value: "timeplusd" + # All tpk streams live under this database, which the app creates on + # startup (CREATE DATABASE IF NOT EXISTS) — so the DB user needs + # CREATE DATABASE the first time, plus read/write on it. If the user + # can't create databases, pre-create this one and grant access, or + # point at an existing database you control. + - name: TIMEPLUS_DATABASE + value: "tpk" # The DB is already up; a short connect-retry covers transient blips. - name: TPK_DB_WAIT_SECONDS value: "120" diff --git a/docker-compose.allinone.yml b/docker-compose.allinone.yml index 37f1ee3..c5f1f7e 100644 --- a/docker-compose.allinone.yml +++ b/docker-compose.allinone.yml @@ -46,6 +46,9 @@ services: TPK_AGENT_MODEL: ${TPK_AGENT_MODEL:-} TPK_AGENT_REASONING_EFFORT: ${TPK_AGENT_REASONING_EFFORT:-} TPK_CHAT_AUDIT: ${TPK_CHAT_AUDIT:-} + # Global fallback daily token budget for non-admin chat (0 = unlimited); + # a role's own daily_token_limit (Users -> Roles) overrides it. + TPK_DAILY_TOKEN_LIMIT: ${TPK_DAILY_TOKEN_LIMIT:-} # Semantic-extraction (graphify) backend. `auto` is ambiguous when both # API keys are set (picks claude), so set this to openai|claude to choose. # The model comes from OPENAI_MODEL / ANTHROPIC_MODEL (graphify reads them). diff --git a/docker-compose.yml b/docker-compose.yml index cf8245f..17cc9cb 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -91,6 +91,9 @@ services: TPK_AGENT_MODEL: ${TPK_AGENT_MODEL:-} TPK_AGENT_REASONING_EFFORT: ${TPK_AGENT_REASONING_EFFORT:-} TPK_CHAT_AUDIT: ${TPK_CHAT_AUDIT:-} + # Global fallback daily token budget for non-admin chat (0 = unlimited); + # a role's own daily_token_limit (Users -> Roles) overrides it. + TPK_DAILY_TOKEN_LIMIT: ${TPK_DAILY_TOKEN_LIMIT:-} # Semantic-extraction (graphify) backend: openai|claude|auto. `auto` is # ambiguous when both API keys are set; model comes from OPENAI_MODEL / # ANTHROPIC_MODEL. diff --git a/repos.toml b/repos.toml index edb1743..94232d3 100644 --- a/repos.toml +++ b/repos.toml @@ -51,6 +51,7 @@ token_budget = 16000 # checkout_dir = "~/.tpk/checkouts" # env: TPK_CHECKOUT_DIR (github checkout cache) # session_ttl = 86400 # env: TPK_SESSION_TTL (login session lifetime, seconds) # chat_audit = true # env: TPK_CHAT_AUDIT (false/0 disables chat auditing) +# daily_token_limit = 500000 # env: TPK_DAILY_TOKEN_LIMIT (global fallback per-user/day, non-admin; 0 = unlimited; a per-user/role limit wins) [repos.proton-enterprise] github = "timeplus-io/proton-enterprise" diff --git a/src/tpk/api.py b/src/tpk/api.py index 4f76624..c113266 100644 --- a/src/tpk/api.py +++ b/src/tpk/api.py @@ -64,6 +64,8 @@ class AddUser(BaseModel): password: str role: str must_change_password: bool = True + # Per-user daily token budget override (0 = inherit role/global, #62). + daily_token_limit: int = 0 class UpdateUser(BaseModel): @@ -71,6 +73,7 @@ class UpdateUser(BaseModel): role: str | None = None disabled: bool | None = None password: str | None = None # admin reset; sets must_change_password + daily_token_limit: int | None = None # omitted -> unchanged class DeleteUser(BaseModel): @@ -84,6 +87,9 @@ class UpsertRole(BaseModel): # Omitted -> the chat-only default (friendly for API-only callers); an # explicit [] is honored as a zero-capability role. capabilities: list[str] | None = None + # Daily per-user token budget for this role's members (0 = inherit the + # global fallback, #62). + daily_token_limit: int = 0 class DeleteRole(BaseModel): @@ -200,7 +206,8 @@ def _client(): def _user_json(u): return {"username": u.username, "role": u.role, "must_change_password": u.must_change_password, - "disabled": u.disabled} + "disabled": u.disabled, + "daily_token_limit": u.daily_token_limit} def _check_role_exists(client, name: str): if name != auth_mod.ROLE_ADMIN and auth_mod.get_role(client, name, prefix=prefix) is None: @@ -234,6 +241,17 @@ def _guard_admin_target(actor, target: "auth_mod.User"): if actor.role != auth_mod.ROLE_ADMIN and target.role == auth_mod.ROLE_ADMIN: raise HTTPException(403, "cannot manage admin users") + def _guard_token_limit(actor, requested, current): + """A daily token budget is a cost-governance lever: only admins may + change one (#62/#24 bounded delegation). A non-admin users:manage + delegate may still manage users/roles as long as it leaves the budget + unchanged, but may not set or raise it — else a limited manager could + hand out unlimited/huge LLM spend through these endpoints.""" + if requested is None: + return + if actor.role != auth_mod.ROLE_ADMIN and int(requested) != int(current or 0): + raise HTTPException(403, "only an admin can change the daily token budget") + def _guard_grant(client, actor, caps, entry_keys): """Reject a grant (role capabilities/entry_keys) exceeding the actor's own. No-op for admins.""" @@ -339,11 +357,15 @@ def api_add_user(body: AddUser, actor: User = Depends(auth.require_cap(auth_mod. client = _client() if auth_mod.get_user(client, body.username, prefix=prefix) is not None: raise HTTPException(409, "user already exists") + if body.daily_token_limit < 0: + raise HTTPException(400, "daily_token_limit must be >= 0 (0 = inherit)") + _guard_token_limit(actor, body.daily_token_limit, 0) _check_role_exists(client, body.role) _guard_role_assignment(client, actor, body.role) auth_mod.upsert_user(client, auth_mod.User( body.username, auth_mod.hash_password(body.password), body.role, - must_change_password=body.must_change_password), prefix=prefix) + must_change_password=body.must_change_password, + daily_token_limit=body.daily_token_limit), prefix=prefix) return {"ok": True} @router.post("/users/update") @@ -359,6 +381,11 @@ def api_update_user(body: UpdateUser, actor: User = Depends(auth.require_cap(aut _check_role_exists(client, role) _guard_role_assignment(client, actor, role) disabled = body.disabled if body.disabled is not None else u.disabled + if body.daily_token_limit is not None and body.daily_token_limit < 0: + raise HTTPException(400, "daily_token_limit must be >= 0 (0 = inherit)") + _guard_token_limit(actor, body.daily_token_limit, u.daily_token_limit) + token_limit = body.daily_token_limit if body.daily_token_limit is not None \ + else u.daily_token_limit password_hash, must_change = u.password_hash, u.must_change_password if body.password is not None: err = auth_mod.validate_new_password(body.password) @@ -366,7 +393,8 @@ def api_update_user(body: UpdateUser, actor: User = Depends(auth.require_cap(aut raise HTTPException(400, err) password_hash, must_change = auth_mod.hash_password(body.password), True auth_mod.upsert_user(client, auth_mod.User( - u.username, password_hash, role, must_change, disabled), prefix=prefix) + u.username, password_hash, role, must_change, disabled, + daily_token_limit=token_limit), prefix=prefix) if disabled or body.password is not None: auth_mod.delete_user_sessions(client, u.username, prefix=prefix) return {"ok": True} @@ -388,7 +416,8 @@ def api_delete_user(body: DeleteUser, actor: User = Depends(auth.require_cap(aut @router.get("/roles") def api_list_roles(actor: User = Depends(auth.require_cap(auth_mod.CAP_USERS_VIEW))): return [{"name": r.name, "entry_keys": r.entry_keys, - "capabilities": r.capabilities, "description": r.description} + "capabilities": r.capabilities, "description": r.description, + "daily_token_limit": r.daily_token_limit} for r in auth_mod.list_roles(_client(), prefix=prefix)] @router.post("/roles") @@ -404,17 +433,22 @@ def api_upsert_role(body: UpsertRole, actor: User = Depends(auth.require_cap(aut unknown = [c for c in caps if c not in auth_mod.ALL_CAPABILITIES] if unknown: raise HTTPException(400, f"unknown capabilities: {', '.join(unknown)}") + if body.daily_token_limit < 0: + raise HTTPException(400, "daily_token_limit must be >= 0 (0 = inherit)") client = _client() # Guard both the new values AND (when overwriting) the role's current # privileges -- else a non-admin manager could neuter a role more # privileged than their own grant, stripping its members (mirrors the # delete-path guard). existing = auth_mod.get_role(client, body.name, prefix=prefix) + _guard_token_limit(actor, body.daily_token_limit, + existing.daily_token_limit if existing else 0) if existing is not None: _guard_grant(client, actor, existing.capabilities, existing.entry_keys) _guard_grant(client, actor, caps, body.entry_keys) auth_mod.upsert_role(client, auth_mod.Role( - body.name, body.entry_keys, body.description, caps), prefix=prefix) + body.name, body.entry_keys, body.description, caps, + daily_token_limit=body.daily_token_limit), prefix=prefix) return {"ok": True} @router.post("/roles/delete") diff --git a/src/tpk/auth.py b/src/tpk/auth.py index d3000df..250c5d2 100644 --- a/src/tpk/auth.py +++ b/src/tpk/auth.py @@ -65,8 +65,9 @@ def expand_capabilities(caps) -> set[str]: _hasher = PasswordHasher() # argon2id defaults _USER_COLUMNS = ["username", "password_hash", "role", "must_change_password", - "disabled", "created_at", "updated_at"] -_ROLE_COLUMNS = ["name", "entry_keys", "capabilities", "description", "updated_at"] + "disabled", "daily_token_limit", "created_at", "updated_at"] +_ROLE_COLUMNS = ["name", "entry_keys", "capabilities", "description", + "daily_token_limit", "updated_at"] _SESSION_COLUMNS = ["token_hash", "username", "expires_at", "created_at"] @@ -77,6 +78,9 @@ class User: role: str must_change_password: bool = False disabled: bool = False + # Per-user daily token budget override (0 = inherit the role/global limit). + # Takes precedence over the role's daily_token_limit when > 0. See #62. + daily_token_limit: int = 0 @dataclass @@ -85,6 +89,10 @@ class Role: entry_keys: list[str] = field(default_factory=list) description: str = "" capabilities: list[str] = field(default_factory=lambda: list(DEFAULT_CAPABILITIES)) + # Daily per-user token budget for members of this role (0 = inherit the + # global fallback, config.daily_token_limit; NOT unlimited unless the global + # itself is 0). A per-user override still wins over this. See #62. + daily_token_limit: int = 0 def _parse_capabilities(raw: str) -> list[str]: @@ -137,29 +145,31 @@ def upsert_user(client, user: User, prefix: str = "") -> None: client.insert( db.qualified("kg_users", prefix), [[user.username, user.password_hash, user.role, - user.must_change_password, user.disabled, _now(), _now()]], + user.must_change_password, user.disabled, + max(int(user.daily_token_limit), 0), _now(), _now()]], column_names=_USER_COLUMNS, ) def get_user(client, username: str, prefix: str = "") -> User | None: rows = client.query( - f"SELECT username, password_hash, role, must_change_password, disabled" + f"SELECT username, password_hash, role, must_change_password, disabled, daily_token_limit" f" FROM {db.latest(db.qualified('kg_users', prefix))} WHERE username = %(u)s", parameters={"u": username}, ).result_rows if not rows: return None - u, h, r, mc, dis = rows[0] - return User(u, h, r, bool(mc), bool(dis)) + u, h, r, mc, dis, lim = rows[0] + return User(u, h, r, bool(mc), bool(dis), daily_token_limit=int(lim or 0)) def list_users(client, prefix: str = "") -> list[User]: rows = client.query( - f"SELECT username, password_hash, role, must_change_password, disabled" + f"SELECT username, password_hash, role, must_change_password, disabled, daily_token_limit" f" FROM {db.latest(db.qualified('kg_users', prefix))} ORDER BY username" ).result_rows - return [User(u, h, r, bool(mc), bool(dis)) for u, h, r, mc, dis in rows] + return [User(u, h, r, bool(mc), bool(dis), daily_token_limit=int(lim or 0)) + for u, h, r, mc, dis, lim in rows] def delete_user(client, username: str, prefix: str = "") -> None: @@ -194,30 +204,32 @@ def upsert_role(client, role: Role, prefix: str = "") -> None: client.insert( db.qualified("kg_roles", prefix), [[role.name, json.dumps(role.entry_keys), json.dumps(role.capabilities), - role.description, _now()]], + role.description, max(int(role.daily_token_limit), 0), _now()]], column_names=_ROLE_COLUMNS, ) def get_role(client, name: str, prefix: str = "") -> Role | None: rows = client.query( - f"SELECT name, entry_keys, capabilities, description FROM {db.latest(db.qualified('kg_roles', prefix))}" - f" WHERE name = %(n)s", + f"SELECT name, entry_keys, capabilities, description, daily_token_limit" + f" FROM {db.latest(db.qualified('kg_roles', prefix))} WHERE name = %(n)s", parameters={"n": name}, ).result_rows if not rows: return None - n, keys, caps, desc = rows[0] - return Role(n, json.loads(keys) if keys else [], desc, _parse_capabilities(caps)) + n, keys, caps, desc, limit = rows[0] + return Role(n, json.loads(keys) if keys else [], desc, _parse_capabilities(caps), + daily_token_limit=int(limit or 0)) def list_roles(client, prefix: str = "") -> list[Role]: rows = client.query( - f"SELECT name, entry_keys, capabilities, description FROM {db.latest(db.qualified('kg_roles', prefix))}" - f" ORDER BY name" + f"SELECT name, entry_keys, capabilities, description, daily_token_limit" + f" FROM {db.latest(db.qualified('kg_roles', prefix))} ORDER BY name" ).result_rows - return [Role(n, json.loads(k) if k else [], d, _parse_capabilities(c)) - for n, k, c, d in rows] + return [Role(n, json.loads(k) if k else [], d, _parse_capabilities(c), + daily_token_limit=int(lim or 0)) + for n, k, c, d, lim in rows] def delete_role(client, name: str, prefix: str = "") -> None: diff --git a/src/tpk/config.py b/src/tpk/config.py index 4051050..1ab024d 100644 --- a/src/tpk/config.py +++ b/src/tpk/config.py @@ -110,6 +110,15 @@ def database() -> str: return name +def daily_token_limit() -> int: + """Global fallback daily per-user token budget for non-admin chat (#62): + env TPK_DAILY_TOKEN_LIMIT > [server].daily_token_limit > 500000. Applies to + users whose role/user override does not set its own `daily_token_limit`. + Set to 0 for unlimited-by-default.""" + val = setting("TPK_DAILY_TOKEN_LIMIT", "server", "daily_token_limit", 500000, cast=int) + return max(int(val), 0) + + @dataclass(frozen=True) class RepoConfig: name: str diff --git a/src/tpk/db.py b/src/tpk/db.py index 9a84634..67a1c7b 100644 --- a/src/tpk/db.py +++ b/src/tpk/db.py @@ -169,17 +169,25 @@ def ensure_schema(client, prefix: str = "") -> None: client.command(_keyed_stream(prefix, "kg_users", [ "username string", "password_hash string", "role string", "must_change_password bool", "disabled bool", + "daily_token_limit uint32", "created_at datetime64(3, 'UTC')", "updated_at datetime64(3, 'UTC')", ], pk="username")) + # Per-user daily token budget override (0 = inherit role/global), #62; + # backfilled onto users created before it as 0. + _add_column_if_missing(client, qualified("kg_users", prefix), "daily_token_limit", "uint32") client.command(_keyed_stream(prefix, "kg_roles", [ "name string", "entry_keys string", "capabilities string", - "description string", "updated_at datetime64(3, 'UTC')", + "description string", "daily_token_limit uint32", + "updated_at datetime64(3, 'UTC')", ], pk="name")) # Backfill for deployments whose kg_roles predates the capabilities # column (CREATE ... IF NOT EXISTS won't add it to an existing stream). # Existing role rows keep an empty cell, which auth._parse_capabilities # reads as the chat-only migration default. _add_column_if_missing(client, qualified("kg_roles", prefix), "capabilities", "string") + # Daily per-user token budget per role (0 = inherit the global fallback); + # backfilled onto roles created before #62 as 0 (inherit). + _add_column_if_missing(client, qualified("kg_roles", prefix), "daily_token_limit", "uint32") client.command(_keyed_stream(prefix, "kg_sessions", [ "token_hash string", "username string", "expires_at datetime64(3, 'UTC')", "created_at datetime64(3, 'UTC')", @@ -208,6 +216,16 @@ def ensure_schema(client, prefix: str = "") -> None: error string ) """) + # Append-only per-turn token usage, for daily-budget enforcement (#62). + # Written independent of TPK_CHAT_AUDIT so a token budget can't be silently + # disabled by turning auditing off. + client.command(f""" + CREATE STREAM IF NOT EXISTS {qualified('chat_usage', prefix)} ( + ts datetime64(3, 'UTC'), + username string, + tokens uint32 + ) + """) def drop_schema(client, prefix: str) -> None: @@ -215,6 +233,6 @@ def drop_schema(client, prefix: str) -> None: raise ValueError("refusing to drop unprefixed (production) streams") for name in ( "kg_nodes", "kg_edges", "kg_ingest_log", "kg_repos", - "kg_users", "kg_roles", "kg_sessions", "chat_audit_log", + "kg_users", "kg_roles", "kg_sessions", "chat_audit_log", "chat_usage", ): client.command(f"DROP STREAM IF EXISTS {qualified(name, prefix)}") diff --git a/src/tpk/server.py b/src/tpk/server.py index 640b30d..445c075 100644 --- a/src/tpk/server.py +++ b/src/tpk/server.py @@ -9,7 +9,7 @@ from pathlib import Path from typing import Literal -from fastapi import Depends, FastAPI +from fastapi import Depends, FastAPI, HTTPException from fastapi.concurrency import run_in_threadpool from fastapi.responses import StreamingResponse from fastapi.staticfiles import StaticFiles @@ -76,6 +76,30 @@ def _chunk_thinking(chunk) -> str: return "" +def _usage_tokens(msg) -> int: + """Total tokens for one model call, read from an `on_chat_model_end` + output message. Prefers LangChain's normalized `usage_metadata`, falling + back to raw `response_metadata` token_usage (OpenAI-style) so gateways that + only pass the raw shape still count. Returns 0 when usage is unavailable + (some gateways omit it) — the caller treats a 0-cost turn as free.""" + if msg is None: + return 0 + um = getattr(msg, "usage_metadata", None) + if isinstance(um, dict): + total = um.get("total_tokens") + if total is None: + total = (um.get("input_tokens") or 0) + (um.get("output_tokens") or 0) + return max(int(total or 0), 0) + rm = getattr(msg, "response_metadata", None) or {} + tu = rm.get("token_usage") or rm.get("usage") or {} + if isinstance(tu, dict): + total = tu.get("total_tokens") + if total is None: + total = (tu.get("prompt_tokens") or 0) + (tu.get("completion_tokens") or 0) + return max(int(total or 0), 0) + return 0 + + def _build_kg_and_repos(): """Production path: one shared `KnowledgeGraph` (+ the parsed repo config it needs) built up front in `create_app`, reused by both the @@ -128,13 +152,16 @@ def _live_corpus(): def create_app( - agent=None, stream_prefix: str = "", auth=None, kg=None, audit_sink=None + agent=None, stream_prefix: str = "", auth=None, kg=None, audit_sink=None, + usage=None, ) -> FastAPI: import tpk.auth as auth_mod from tpk.agent import RECURSION_LIMIT from tpk.api import create_api_router from tpk.auth import AuthLayer, User, create_auth_router + from tpk.config import daily_token_limit from tpk.tools import ROLE_SCOPE + from tpk.usage import day_window, effective_daily_limit auth = auth or AuthLayer(stream_prefix) app = FastAPI(title="timeplus-knowledge") @@ -162,6 +189,15 @@ def create_app( # `kg` built over the test stream prefix -- so this branch never # runs there and never touches the network in unit tests. kg, repos_for_agent = _build_kg_and_repos() + # Per-user daily token budget (#62). Always-on in production (not gated + # by the audit toggle); a fresh client per read/write. Tests inject + # their own `usage` (or leave it None to disable enforcement). + if usage is None: + from tpk import db + from tpk.config import Settings + from tpk.usage import DbUsage + + usage = DbUsage(stream_prefix, lambda: db.get_client(Settings.from_env())) state = {"agent": agent} @@ -186,6 +222,33 @@ def chat_model(user: User = Depends(auth.require_cap(auth_mod.CAP_CHAT))): except Exception: return {"provider": None, "model": None} + @app.get("/chat/usage") + def chat_usage_status(user: User = Depends(auth.require_cap(auth_mod.CAP_CHAT))): + """This user's daily token budget for the UI (#62): how much is used, + how much is left, and when it resets. `limited: false` for admins, when + the effective limit is 0 (unlimited), or when no usage store is active — + the UI then shows no budget indicator. Sync def -> FastAPI runs the DB + reads in a threadpool. Best-effort: a role-lookup failure falls back to + the global default limit, and the usage read fails open (0 used).""" + if user.role == auth_mod.ROLE_ADMIN or usage is None: + return {"limited": False} + try: + role = auth_mod.get_role(auth._client(), user.role, prefix=stream_prefix) + except Exception: + role = None + limit = effective_daily_limit(user.daily_token_limit, role, daily_token_limit()) + if limit <= 0: + return {"limited": False} + used = usage.used_today(user.username) + _, reset = day_window() + return { + "limited": True, + "used": used, + "limit": limit, + "remaining": max(0, limit - used), + "reset": reset.isoformat(), + } + @app.post("/chat") async def chat(req: ChatRequest, user: User = Depends(auth.require_cap(auth_mod.CAP_CHAT))): messages = [(t.role, t.content) for t in req.history] + [("user", req.message)] @@ -201,6 +264,7 @@ async def chat(req: ChatRequest, user: User = Depends(auth.require_cap(auth_mod. audit_provider, audit_model = "", "" scope = None + turn_limit = 0 # effective daily token budget for this user (0 = unlimited) if user.role != auth_mod.ROLE_ADMIN: try: # Off the event loop: against an unreachable-but-not-refusing @@ -215,6 +279,31 @@ async def chat(req: ChatRequest, user: User = Depends(auth.require_cap(auth_mod. # scope = tools see nothing, rather than falling through to # unrestricted (None) access. scope = frozenset(role.entry_keys) if role else frozenset() + # Effective daily token budget by precedence: the user's own + # override, else the role's limit, else the global fallback. admin + # is never limited (this branch is skipped for admins). + turn_limit = effective_daily_limit(user.daily_token_limit, role, daily_token_limit()) + + # Enforce the budget BEFORE running the agent (#62). Enforcement is + # necessarily next-turn: a turn's cost is only known once it runs, so + # the turn that crosses the line completes and the NEXT one is + # blocked. Reads fail open (usage.used_today swallows + returns 0). + if usage is not None and turn_limit > 0: + used = await run_in_threadpool(lambda: usage.used_today(user.username)) + if used >= turn_limit: + _, reset = day_window() + raise HTTPException( + status_code=429, + detail={ + "message": ( + f"Daily token budget reached ({used}/{turn_limit}). " + f"Access resets at {reset.isoformat()}." + ), + "used": used, + "limit": turn_limit, + "reset": reset.isoformat(), + }, + ) async def stream(): # Set inside stream(), not the handler body: the generator runs @@ -229,6 +318,7 @@ async def stream(): answer_text = "" audit_status = "ok" audit_error = "" + turn_tokens = 0 # summed across the turn's model calls (#62) try: full: list[str] = [] final_text = "" @@ -257,9 +347,13 @@ async def stream(): # The last model turn's message is the authoritative # answer — token deltas can miss it entirely for # models that stream on a reasoning channel (gpt-oss). - end_text = _chunk_text(event.get("data", {}).get("output")) + output = event.get("data", {}).get("output") + end_text = _chunk_text(output) if end_text: final_text = end_text + # Sum token cost across every model call in the turn + # (the tool loop makes several) for the daily budget. + turn_tokens += _usage_tokens(output) elif kind == "on_tool_start": yield _sse( { @@ -378,6 +472,14 @@ async def stream(): await run_in_threadpool(audit_sink, record) except Exception: logger.exception("chat audit failed; skipping") + # Daily-budget metering (#62): record this turn's token cost, + # off the event loop and best-effort. Written even on a partial + # (errored) turn — tokens spent still count against the budget. + if usage is not None and turn_tokens > 0: + try: + await run_in_threadpool(usage.record, user.username, turn_tokens) + except Exception: + logger.exception("usage record failed; skipping") return StreamingResponse(stream(), media_type="text/event-stream") diff --git a/src/tpk/usage.py b/src/tpk/usage.py new file mode 100644 index 0000000..456e429 --- /dev/null +++ b/src/tpk/usage.py @@ -0,0 +1,93 @@ +"""Per-user LLM token metering + daily-budget accounting (issue #62). + +One row per chat turn is written to the append-only ``chat_usage`` stream, +independent of the chat-audit toggle (``TPK_CHAT_AUDIT``), so a daily token +budget stays enforceable even when auditing is off. Enforcement reads the +day's running total from the same stream. + +Both read and write are best-effort at the call site: a token budget must never +be a single point of failure for chat. Reads fail *open* (an unreachable store +returns 0 used, so nobody is wrongly blocked); writes are swallowed. +""" + +import logging +from datetime import datetime, timedelta, timezone + +from tpk import db + +logger = logging.getLogger(__name__) + +# Insert column order for the append-only chat_usage stream (see db.ensure_schema). +USAGE_COLUMNS = ["ts", "username", "tokens"] + + +def effective_daily_limit(user_limit: int, role, global_default: int) -> int: + """The daily token budget that applies to a user, by precedence (#62): + the user's own override, else their role's limit, else the global fallback. + At the user and role levels, 0 means "inherit the next level down" (not + unlimited); only a 0 that reaches the global fallback is truly unlimited + (i.e. when TPK_DAILY_TOKEN_LIMIT is 0). `role` may be a Role or None (an + unreadable/absent role falls through to the global).""" + ul = int(user_limit or 0) + if ul > 0: + return ul + role_limit = getattr(role, "daily_token_limit", 0) or 0 + return int(role_limit) if role_limit > 0 else max(int(global_default), 0) + + +def day_window(now: datetime | None = None) -> tuple[datetime, datetime]: + """The current budgeting window: [UTC-midnight-today, UTC-midnight-tomorrow). + A per-UTC-calendar-day window — predictable and simple to explain.""" + now = now or datetime.now(timezone.utc) + start = now.astimezone(timezone.utc).replace(hour=0, minute=0, second=0, microsecond=0) + return start, start + timedelta(days=1) + + +def record_usage(client, username: str, tokens: int, prefix: str = "", + ts: datetime | None = None) -> None: + """Append one usage row (a chat turn's total token cost).""" + client.insert( + db.qualified("chat_usage", prefix), + [[ts or datetime.now(timezone.utc), username, int(tokens)]], + column_names=USAGE_COLUMNS, + ) + + +def tokens_used_today(client, username: str, prefix: str = "", + now: datetime | None = None) -> int: + """Sum of tokens a user has spent in the current day window.""" + start, _ = day_window(now) + rows = client.query( + f"SELECT sum(tokens) FROM table({db.qualified('chat_usage', prefix)})" + " WHERE username = %(u)s AND ts >= %(start)s", + parameters={"u": username, "start": start}, + ).result_rows + return int(rows[0][0]) if rows and rows[0][0] is not None else 0 + + +class DbUsage: + """DB-backed usage store the /chat handler uses to enforce + record. A + fresh single-use client per call avoids timeplus_connect's concurrent-query + restriction (same reasoning as the audit sink).""" + + def __init__(self, prefix: str, client_factory): + self.prefix = prefix + self._client_factory = client_factory + + def used_today(self, username: str, now: datetime | None = None) -> int: + try: + return tokens_used_today(self._client_factory(), username, + prefix=self.prefix, now=now) + except Exception: + # Fail OPEN: never block a user because the usage store is + # unreachable. Log and treat as zero used. + logger.exception("usage read failed; treating as 0 used (fail-open)") + return 0 + + def record(self, username: str, tokens: int) -> None: + if tokens <= 0: + return + try: + record_usage(self._client_factory(), username, tokens, prefix=self.prefix) + except Exception: + logger.exception("usage write failed; skipping") diff --git a/tests/test_backend.py b/tests/test_backend.py index 6102888..9e2a8eb 100644 --- a/tests/test_backend.py +++ b/tests/test_backend.py @@ -164,3 +164,35 @@ def test_backend_session_lifecycle(backend_tp): auth.delete_session(client, token, prefix=prefix) _eventually(lambda: True if auth.get_session(client, token, prefix=prefix) is None else None) assert auth.get_session(client, token, prefix=prefix) is None + + +def test_backend_role_daily_token_limit_roundtrip(backend_tp): + from tpk import auth + client, prefix, backend = backend_tp + + auth.upsert_role(client, auth.Role("member", ["r1@v1"], "desc", ["chat"], + daily_token_limit=12345), prefix=prefix) + got = _eventually(lambda: auth.get_role(client, "member", prefix=prefix)) + assert got is not None and got.daily_token_limit == 12345 + # a role that doesn't set a limit reads back as 0 (unlimited) + auth.upsert_role(client, auth.Role("open", [], "", ["chat"]), prefix=prefix) + got2 = _eventually(lambda: auth.get_role(client, "open", prefix=prefix)) + assert got2 is not None and got2.daily_token_limit == 0 + + +def test_backend_chat_usage_daily_count(backend_tp): + from tpk import usage + client, prefix, backend = backend_tp + + usage.record_usage(client, "bob", 1000, prefix=prefix) + usage.record_usage(client, "bob", 500, prefix=prefix) + usage.record_usage(client, "alice", 999, prefix=prefix) + got = None + for _ in range(20): + got = usage.tokens_used_today(client, "bob", prefix=prefix) + if got == 1500: + break + time.sleep(0.3) + assert got == 1500 + # scoping is per-user + assert usage.tokens_used_today(client, "alice", prefix=prefix) == 999 diff --git a/tests/test_capabilities.py b/tests/test_capabilities.py index f742d82..c6b4537 100644 --- a/tests/test_capabilities.py +++ b/tests/test_capabilities.py @@ -215,6 +215,23 @@ def test_manager_cannot_grant_corpus_beyond_own(c, manager): assert r.status_code == 403 +def test_manager_cannot_set_token_budget(c, client, prefix, manager): + # A non-admin manager may create/manage users and roles, but not set or + # change a daily token budget (admin-only cost lever, #62). + r = c.post("/api/roles", json={"name": "sub-budget", "entry_keys": ["alpha@v1"], + "capabilities": [CAP_CHAT], "daily_token_limit": 999999}, + headers=_hdr(manager)) + assert r.status_code == 403 + r = c.post("/api/users", json={"username": "richuser", "password": "password-2", + "role": "manager", "daily_token_limit": 999999}, + headers=_hdr(manager)) + assert r.status_code == 403 + # ...but managing without touching the budget (limit 0 = inherit) is fine. + assert c.post("/api/roles", json={"name": "sub-ok", "entry_keys": ["alpha@v1"], + "capabilities": [CAP_CHAT]}, + headers=_hdr(manager)).status_code == 200 + + def test_manager_cannot_assign_admin_role(c, client, prefix, manager): r = c.post("/api/users", json={"username": "eviladmin", "password": "password-2", "role": auth.ROLE_ADMIN}, headers=_hdr(manager)) diff --git a/tests/test_config.py b/tests/test_config.py index b55c288..84b3879 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -137,6 +137,19 @@ def test_database_rejects_bad_value(config_file, monkeypatch): database() +def test_daily_token_limit_precedence(config_file, monkeypatch): + from tpk.config import daily_token_limit + + monkeypatch.delenv("TPK_DAILY_TOKEN_LIMIT", raising=False) + assert daily_token_limit() == 500000 # built-in default budget + config_file.write_text("[server]\ndaily_token_limit = 50000\n") + assert daily_token_limit() == 50000 # file over default + monkeypatch.setenv("TPK_DAILY_TOKEN_LIMIT", "1000") + assert daily_token_limit() == 1000 # env over file + monkeypatch.setenv("TPK_DAILY_TOKEN_LIMIT", "0") + assert daily_token_limit() == 0 # explicit 0 = unlimited + + def test_settings_reads_db_section_from_file(config_file, monkeypatch): for var in ("TIMEPLUS_HOST", "TIMEPLUS_USER", "TPK_STREAM_PREFIX", "TPK_DB_BACKEND"): monkeypatch.delenv(var, raising=False) diff --git a/tests/test_server.py b/tests/test_server.py index 16964cc..553a8c3 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -86,11 +86,125 @@ def _parse_sse(body: str) -> list[dict]: return [json.loads(line[len("data: "):]) for line in body.splitlines() if line.startswith("data: ")] +def _end_usage(total_tokens=0, text=""): + """An on_chat_model_end event whose output carries usage_metadata (#62).""" + class Msg: + content = text + usage_metadata = {"total_tokens": total_tokens} + return {"event": "on_chat_model_end", "data": {"output": Msg()}} + + +class _Usage: + """In-memory usage store stub for enforcement/metering tests.""" + + def __init__(self, used=0): + self._used = used + self.recorded = [] + + def used_today(self, username, now=None): + return self._used + + def record(self, username, tokens): + self.recorded.append((username, tokens)) + + def test_healthz(): client = TestClient(create_app(agent=FakeAgent([]), auth=_StubAuth())) assert client.get("/healthz").json() == {"status": "ok"} +def test_chat_over_daily_budget_returns_429(monkeypatch): + # No per-role limit reachable in unit tests (stub._client raises), so the + # global fallback applies; used == limit -> blocked. + monkeypatch.setenv("TPK_DAILY_TOKEN_LIMIT", "1000") + monkeypatch.delenv("TPK_CONFIG", raising=False) + usage = _Usage(used=1000) + auth = _StubAuth(User("bob", "", "member"), caps=["chat"]) + client = TestClient(create_app(agent=FakeAgent([_tok("hi")]), auth=auth, usage=usage)) + resp = client.post("/chat", json={"message": "q"}) + assert resp.status_code == 429 + detail = resp.json()["detail"] + assert detail["limit"] == 1000 and detail["used"] == 1000 + assert "reset" in detail and "budget" in detail["message"].lower() + + +def test_chat_under_budget_proceeds_and_records_usage(monkeypatch): + monkeypatch.setenv("TPK_DAILY_TOKEN_LIMIT", "10000") + monkeypatch.delenv("TPK_CONFIG", raising=False) + usage = _Usage(used=500) + auth = _StubAuth(User("bob", "", "member"), caps=["chat"]) + agent = FakeAgent([_tok("hi "), _end_usage(1234, "hi there")]) + client = TestClient(create_app(agent=agent, auth=auth, usage=usage)) + resp = client.post("/chat", json={"message": "q"}) + assert resp.status_code == 200 + # the turn's token cost is metered to the store + assert usage.recorded == [("bob", 1234)] + + +def test_chat_admin_never_limited(monkeypatch): + monkeypatch.setenv("TPK_DAILY_TOKEN_LIMIT", "10") + usage = _Usage(used=10_000_000) # way over, but admin is exempt + client = TestClient(create_app(agent=FakeAgent([_tok("hi")]), + auth=_StubAuth(User("root", "", "admin")), usage=usage)) + resp = client.post("/chat", json={"message": "q"}) + assert resp.status_code == 200 + + +def test_chat_no_usage_store_skips_enforcement(monkeypatch): + # Unit tests that don't inject a usage store must not enforce (or touch a DB). + monkeypatch.setenv("TPK_DAILY_TOKEN_LIMIT", "1") + auth = _StubAuth(User("bob", "", "member"), caps=["chat"]) + client = TestClient(create_app(agent=FakeAgent([_tok("hi")]), auth=auth)) + assert client.post("/chat", json={"message": "q"}).status_code == 200 + + +def test_chat_user_override_wins_over_global(monkeypatch): + # A tight per-user override beats a generous global default. + monkeypatch.setenv("TPK_DAILY_TOKEN_LIMIT", "1000000") + monkeypatch.delenv("TPK_CONFIG", raising=False) + usage = _Usage(used=100) + user = User("bob", "", "member", daily_token_limit=100) + auth = _StubAuth(user, caps=["chat"]) + client = TestClient(create_app(agent=FakeAgent([_tok("hi")]), auth=auth, usage=usage)) + resp = client.post("/chat", json={"message": "q"}) + assert resp.status_code == 429 + assert resp.json()["detail"]["limit"] == 100 + + +def test_chat_usage_status_reflects_user_override(monkeypatch): + monkeypatch.setenv("TPK_DAILY_TOKEN_LIMIT", "1000000") + monkeypatch.delenv("TPK_CONFIG", raising=False) + usage = _Usage(used=40) + user = User("bob", "", "member", daily_token_limit=100) + client = TestClient(create_app(agent=FakeAgent([]), auth=_StubAuth(user, caps=["chat"]), usage=usage)) + body = client.get("/chat/usage").json() + assert body["limit"] == 100 and body["remaining"] == 60 + + +def test_chat_usage_status_for_limited_user(monkeypatch): + monkeypatch.setenv("TPK_DAILY_TOKEN_LIMIT", "50000") + monkeypatch.delenv("TPK_CONFIG", raising=False) + usage = _Usage(used=1234) + auth = _StubAuth(User("bob", "", "member"), caps=["chat"]) + client = TestClient(create_app(agent=FakeAgent([]), auth=auth, usage=usage)) + body = client.get("/chat/usage").json() + assert body["limited"] is True + assert body["used"] == 1234 and body["limit"] == 50000 + assert body["remaining"] == 48766 and "reset" in body + + +def test_chat_usage_status_unlimited_for_admin_and_no_store(monkeypatch): + monkeypatch.setenv("TPK_DAILY_TOKEN_LIMIT", "50000") + # admin -> unlimited even with a store + admin = TestClient(create_app(agent=FakeAgent([]), auth=_StubAuth(User("root", "", "admin")), + usage=_Usage(used=9))) + assert admin.get("/chat/usage").json() == {"limited": False} + # non-admin but no usage store -> unlimited (unit path) + auth = _StubAuth(User("bob", "", "member"), caps=["chat"]) + nostore = TestClient(create_app(agent=FakeAgent([]), auth=auth)) + assert nostore.get("/chat/usage").json() == {"limited": False} + + def test_chat_model_endpoint(monkeypatch): monkeypatch.setenv("TPK_AGENT_PROVIDER", "anthropic") monkeypatch.setenv("TPK_AGENT_MODEL", "anthropic.claude-opus-4-8") diff --git a/tests/test_usage.py b/tests/test_usage.py new file mode 100644 index 0000000..d073108 --- /dev/null +++ b/tests/test_usage.py @@ -0,0 +1,90 @@ +"""Token metering + daily-budget accounting (issue #62).""" + +from datetime import datetime, timezone + +from tpk import usage + + +class _FakeClient: + def __init__(self, sum_val=0): + self.inserts = [] + self.queries = [] + self._sum = sum_val + + def insert(self, stream, rows, column_names): + self.inserts.append((stream, rows, column_names)) + + def query(self, sql, parameters=None): + self.queries.append((sql, parameters)) + return type("R", (), {"result_rows": [(self._sum,)]})() + + +def test_day_window_is_utc_calendar_day(): + now = datetime(2026, 8, 17, 15, 30, 0, tzinfo=timezone.utc) + start, reset = usage.day_window(now) + assert start == datetime(2026, 8, 17, 0, 0, 0, tzinfo=timezone.utc) + assert reset == datetime(2026, 8, 18, 0, 0, 0, tzinfo=timezone.utc) + + +def test_record_usage_inserts_qualified_row(monkeypatch): + monkeypatch.delenv("TIMEPLUS_DATABASE", raising=False) + c = _FakeClient() + ts = datetime(2026, 8, 17, 12, 0, 0, tzinfo=timezone.utc) + usage.record_usage(c, "bob", 1234, prefix="p_", ts=ts) + stream, rows, cols = c.inserts[0] + assert stream == "tpk.p_chat_usage" + assert rows == [[ts, "bob", 1234]] + assert cols == ["ts", "username", "tokens"] + + +def test_tokens_used_today_sums_and_scopes_query(monkeypatch): + monkeypatch.delenv("TIMEPLUS_DATABASE", raising=False) + c = _FakeClient(sum_val=5000) + now = datetime(2026, 8, 17, 9, 0, 0, tzinfo=timezone.utc) + assert usage.tokens_used_today(c, "bob", prefix="p_", now=now) == 5000 + sql, params = c.queries[0] + assert "table(tpk.p_chat_usage)" in sql + assert params["u"] == "bob" + assert params["start"] == datetime(2026, 8, 17, 0, 0, 0, tzinfo=timezone.utc) + + +def test_tokens_used_today_handles_null_sum(monkeypatch): + monkeypatch.delenv("TIMEPLUS_DATABASE", raising=False) + c = _FakeClient(sum_val=None) # sum() over no rows -> NULL + assert usage.tokens_used_today(c, "bob", now=datetime.now(timezone.utc)) == 0 + + +def test_effective_daily_limit_precedence(): + class _Role: + daily_token_limit = 5000 + # user override wins over role and global + assert usage.effective_daily_limit(100, _Role(), 9999) == 100 + # no user override -> role limit + assert usage.effective_daily_limit(0, _Role(), 9999) == 5000 + # no user or role limit -> global fallback + assert usage.effective_daily_limit(0, None, 9999) == 9999 + # a role limit of 0 means INHERIT (not unlimited): falls through to global + assert usage.effective_daily_limit(0, _Role0(), 9999) == 9999 + # nothing set anywhere (global also 0) -> 0 (truly unlimited) + assert usage.effective_daily_limit(0, None, 0) == 0 + + +class _Role0: + daily_token_limit = 0 + + +def test_db_usage_read_fails_open(): + def boom(): + raise RuntimeError("store down") + store = usage.DbUsage("", boom) + # Never block a user because the store is unreachable. + assert store.used_today("bob") == 0 + + +def test_db_usage_record_swallows_and_skips_zero(): + c = _FakeClient() + store = usage.DbUsage("", lambda: c) + store.record("bob", 0) # non-positive -> no write + assert c.inserts == [] + store.record("bob", 42) + assert c.inserts and c.inserts[0][1] == [[c.inserts[0][1][0][0], "bob", 42]] diff --git a/web/src/Chat.tsx b/web/src/Chat.tsx index d14f9a8..c128e42 100644 --- a/web/src/Chat.tsx +++ b/web/src/Chat.tsx @@ -190,6 +190,13 @@ function SourceCard({ n, source }: { n?: number; source: SourceEventPayload }) { // Chat // -------------------------------------------------------------------- +// Daily token budget for the current user (#62). `limited: false` for admins, +// unlimited roles, or when no budget is enforced -> no indicator shown. +type UsageInfo = + | { limited: true; used: number; limit: number; remaining: number; reset: string } + | { limited: false } + | null; + export default function Chat({ initialInput, onConsumeInitial, @@ -204,6 +211,7 @@ export default function Chat({ const [busy, setBusy] = useState(false); const [corpusTags, setCorpusTags] = useState([]); const [agentModel, setAgentModel] = useState(null); + const [usage, setUsage] = useState(null); // Citations live on the read_source trace rows (issue #40): clicking a row // opens that source fragment in the side panel (like the previous Sources // panel), rather than inline. `activeSource` is the one being shown, or null. @@ -280,6 +288,16 @@ export default function Chat({ return () => { cancelled = true; }; }, []); + // This user's daily token budget (#62): fetched on load and refreshed after + // each turn (a turn spends tokens). Silent on failure — just no indicator. + async function refreshUsage() { + try { + const resp = await apiFetch("/chat/usage"); + if (resp.ok) setUsage(await resp.json()); + } catch { /* leave the last-known usage in place */ } + } + useEffect(() => { void refreshUsage(); }, []); + async function send(overrideMessage?: string) { const message = (overrideMessage ?? input).trim(); if (!message || busy) return; @@ -311,7 +329,15 @@ export default function Chat({ body: JSON.stringify({ message, history }), }); if (!resp.ok || !resp.body) { - update((t) => ({ ...t, content: t.content + `\n\n[error] HTTP ${resp.status}` })); + // Daily token budget reached (#62): show the server's friendly message + // (with the reset time) rather than a bare HTTP code. + let msg = `\n\n[error] HTTP ${resp.status}`; + if (resp.status === 429) { + const detail = await resp.json().catch(() => null); + const m = detail?.detail?.message ?? detail?.message; + msg = m ? `\n\n${m}` : "\n\n[error] Daily usage limit reached."; + } + update((t) => ({ ...t, content: t.content + msg })); finish("error"); return; } @@ -394,6 +420,11 @@ export default function Chat({ finish("error"); } finally { setBusy(false); + // The turn's token cost is recorded server-side as the stream ends; the + // append-only usage read lags slightly, so refresh now and once shortly + // after so the "used/left" figures reflect the turn just completed. + void refreshUsage(); + setTimeout(() => void refreshUsage(), 1500); } } @@ -518,6 +549,19 @@ export default function Chat({ {agentModel} )} + {usage?.limited && ( + <> +
daily tokens
+ + {usage.used.toLocaleString()} / {usage.limit.toLocaleString()} + +
+ {usage.remaining.toLocaleString()} left · resets{" "} + {new Date(usage.reset).toLocaleTimeString([], { hour: "2-digit", minute: "2-digit" })} +
+ + )} )} @@ -539,6 +583,13 @@ export default function Chat({ model {agentModel} )} + {usage?.limited && ( +
+ {usage.used.toLocaleString()} / {usage.limit.toLocaleString()} + {" "}tokens used today · {usage.remaining.toLocaleString()} left · resets{" "} + {new Date(usage.reset).toLocaleTimeString([], { hour: "2-digit", minute: "2-digit" })} +
+ )}
{SUGGESTED_QUESTIONS.map((q) => (