From 92c156e8ba5ce33eb8d1dd2026df1a1f16efb6c1 Mon Sep 17 00:00:00 2001 From: Gang Tao Date: Mon, 17 Aug 2026 15:49:30 -0700 Subject: [PATCH 1/7] =?UTF-8?q?docs(k8s):=20app-only=20=E2=80=94=20account?= =?UTF-8?q?=20for=20the=20dedicated=20`tpk`=20database=20(#58/#59)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The k8s manifests (#56/#57) predate the change that moved all tpk streams under a dedicated `tpk` database (#59). For app-only — which connects to an externally-managed Timeplus Enterprise — the app now runs `CREATE DATABASE IF NOT EXISTS tpk` on startup, so the pre-existing DB user needs CREATE DATABASE (first run) plus read/write on that database. A restricted user would otherwise fail at startup/ingest. - app-only.yaml: header now states the app creates its own database on the existing server; add a `TIMEPLUS_DATABASE` env (default `tpk`) with the privilege note and the pre-create fallback for restricted users. - deploy/k8s/README.md: add a "tpk database" note to the App-only section (CREATE DATABASE grant + pre-create SQL), and a TIMEPLUS_DATABASE row to the Configuration table. Notes that enterprise/allinone are unaffected (they provision a full-privilege tpk user). Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01XPzYZXxyTdj5G25KoujpHb --- deploy/k8s/README.md | 17 +++++++++++++++++ deploy/k8s/app-only.yaml | 13 ++++++++++++- 2 files changed, 29 insertions(+), 1 deletion(-) 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" From 2b28d17b2316058d06808adff46d6dd069b134b5 Mon Sep 17 00:00:00 2001 From: Gang Tao Date: Mon, 17 Aug 2026 16:08:34 -0700 Subject: [PATCH 2/7] feat(chat): per-user daily token budget for non-admin users (#62) Operators can cap daily LLM token consumption per non-admin user so a single user can't run up unbounded cost against the shared agent. admin (the reserved super-role) is exempt. (Token budget only; the optional question cap from #62 is deferred.) - usage.py: append-only `chat_usage` stream (ts, username, tokens) written per turn independent of TPK_CHAT_AUDIT, plus DbUsage (used_today/record). Reads fail OPEN (unreachable store never blocks a user); writes are best-effort. - db: create chat_usage in ensure_schema (+ drop_schema); add kg_roles.daily_token_limit (uint32) with the same _add_column_if_missing lazy migration as capabilities (existing roles -> 0 = unlimited). - auth.Role + kg_roles carry daily_token_limit; api /roles create/update/list carry it (validated >= 0). - config.daily_token_limit(): global fallback (env TPK_DAILY_TOKEN_LIMIT > [server].daily_token_limit > 0) applied when a role sets no limit of its own. - server /chat: sum usage_metadata tokens across the turn's model calls; before running the agent, block a non-admin who is at/over budget with HTTP 429 + reset time (next UTC midnight). Enforcement is next-turn (a turn's cost is only known once it runs); the crossing turn completes. Metering + enforcement only active when a usage store is present (production), so unit tests that inject a fake agent are untouched. - web: role editor gains a "Daily token budget per user" field (add + edit); chat surfaces the 429 with the server's friendly reset message. - tests: usage helpers + fail-open, config precedence, enforcement (over/under/ admin-bypass/no-store), and backend round-trips (role limit, daily count). Verified end-to-end on live timeplusd. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01XPzYZXxyTdj5G25KoujpHb --- .env.example | 1 + README.md | 1 + deploy/docker/repos.container.toml | 1 + repos.toml | 1 + src/tpk/api.py | 10 +++- src/tpk/auth.py | 26 ++++++---- src/tpk/config.py | 8 +++ src/tpk/db.py | 18 ++++++- src/tpk/server.py | 82 ++++++++++++++++++++++++++++-- src/tpk/usage.py | 79 ++++++++++++++++++++++++++++ tests/test_backend.py | 32 ++++++++++++ tests/test_config.py | 11 ++++ tests/test_server.py | 67 ++++++++++++++++++++++++ tests/test_usage.py | 71 ++++++++++++++++++++++++++ web/src/Chat.tsx | 10 +++- web/src/Users.tsx | 42 ++++++++++++--- 16 files changed, 434 insertions(+), 26 deletions(-) create mode 100644 src/tpk/usage.py create mode 100644 tests/test_usage.py diff --git a/.env.example b/.env.example index 027f114..1bc8ac9 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=0 # [server].daily_token_limit 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..461ef39 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` | `0` | Per-user daily token budget for non-admins (`0` = unlimited; a role's own `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..d7ac6c9 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 = 0 # env: TPK_DAILY_TOKEN_LIMIT (non-admin per-user/day; 0 = unlimited) [repos.proton-enterprise] github = "timeplus-io/proton-enterprise" diff --git a/repos.toml b/repos.toml index edb1743..ea514cd 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 = 0 # env: TPK_DAILY_TOKEN_LIMIT (per-user/day, non-admin; 0 = unlimited; a role's own limit wins) [repos.proton-enterprise] github = "timeplus-io/proton-enterprise" diff --git a/src/tpk/api.py b/src/tpk/api.py index 4f76624..5da0c8c 100644 --- a/src/tpk/api.py +++ b/src/tpk/api.py @@ -84,6 +84,8 @@ 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 = unlimited, #62). + daily_token_limit: int = 0 class DeleteRole(BaseModel): @@ -388,7 +390,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,6 +407,8 @@ 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 = unlimited)") 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 @@ -414,7 +419,8 @@ def api_upsert_role(body: UpsertRole, actor: User = Depends(auth.require_cap(aut _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..3ec327f 100644 --- a/src/tpk/auth.py +++ b/src/tpk/auth.py @@ -66,7 +66,8 @@ def expand_capabilities(caps) -> set[str]: _USER_COLUMNS = ["username", "password_hash", "role", "must_change_password", "disabled", "created_at", "updated_at"] -_ROLE_COLUMNS = ["name", "entry_keys", "capabilities", "description", "updated_at"] +_ROLE_COLUMNS = ["name", "entry_keys", "capabilities", "description", + "daily_token_limit", "updated_at"] _SESSION_COLUMNS = ["token_hash", "username", "expires_at", "created_at"] @@ -85,6 +86,9 @@ 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 = unlimited). See + # #62; a global fallback (config.daily_token_limit) applies when this is 0. + daily_token_limit: int = 0 def _parse_capabilities(raw: str) -> list[str]: @@ -194,30 +198,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..7ea3565 100644 --- a/src/tpk/config.py +++ b/src/tpk/config.py @@ -110,6 +110,14 @@ 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 > 0 (unlimited). + Applies to users whose role does not set its own `daily_token_limit`.""" + val = setting("TPK_DAILY_TOKEN_LIMIT", "server", "daily_token_limit", 0, 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..a6231b8 100644 --- a/src/tpk/db.py +++ b/src/tpk/db.py @@ -173,13 +173,17 @@ def ensure_schema(client, prefix: str = "") -> None: ], pk="username")) 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 = unlimited); backfilled onto + # roles created before #62 as 0 (unlimited). + _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 +212,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 +229,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..54cd215 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,12 +152,14 @@ 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 auth = auth or AuthLayer(stream_prefix) @@ -162,6 +188,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} @@ -201,6 +236,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 +251,33 @@ 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: the role's own limit, else the + # global fallback (config). admin is never limited (branch skipped). + role_limit = role.daily_token_limit if role else 0 + turn_limit = role_limit if role_limit > 0 else 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: + from tpk.usage import day_window + + _, 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 +292,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 +321,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 +446,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..46166fb --- /dev/null +++ b/src/tpk/usage.py @@ -0,0 +1,79 @@ +"""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 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_config.py b/tests/test_config.py index b55c288..0cf5d20 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -137,6 +137,17 @@ 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() == 0 # default: unlimited + 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 + + 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..37ebbf3 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -86,11 +86,78 @@ 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_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..b1267a7 --- /dev/null +++ b/tests/test_usage.py @@ -0,0 +1,71 @@ +"""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_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..1a90eee 100644 --- a/web/src/Chat.tsx +++ b/web/src/Chat.tsx @@ -311,7 +311,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; } diff --git a/web/src/Users.tsx b/web/src/Users.tsx index 19c7190..46f9060 100644 --- a/web/src/Users.tsx +++ b/web/src/Users.tsx @@ -3,15 +3,17 @@ import { apiFetch } from "./api"; import { CAP, CAPABILITY_OPTIONS, type Capability, expandCaps, hasCap } from "./capabilities"; type ApiUser = { username: string; role: string; must_change_password: boolean; disabled: boolean }; -type ApiRole = { name: string; entry_keys: string[]; description: string; capabilities: string[] }; +type ApiRole = { name: string; entry_keys: string[]; description: string; + capabilities: string[]; daily_token_limit: number }; type Repo = { entry_key: string }; -type RoleEdit = { entry_keys: string[]; description: string; capabilities: string[] }; +type RoleEdit = { entry_keys: string[]; description: string; capabilities: string[]; + daily_token_limit: number }; // Least-privilege: force an explicit role pick rather than defaulting new // users to admin. A new role starts chat-only. const EMPTY_USER = { username: "", password: "", role: "", must_change_password: true }; const EMPTY_ROLE = { name: "", entry_keys: [] as string[], description: "", - capabilities: [CAP.chat] as string[] }; + capabilities: [CAP.chat] as string[], daily_token_limit: 0 }; // Checking a `:manage` capability implies its `:view` sibling; unchecking is // free. Returns the next selected set. @@ -88,7 +90,8 @@ export default function Users({ capabilities, isAdmin }: setEntryKeys([...new Set(repos.map((x) => x.entry_key))].sort()); setRoleEdits(Object.fromEntries( r.map((role) => [role.name, { entry_keys: role.entry_keys, - description: role.description, capabilities: role.capabilities ?? [] }]) + description: role.description, capabilities: role.capabilities ?? [], + daily_token_limit: role.daily_token_limit ?? 0 }]) )); setError(""); } catch (e) { @@ -104,7 +107,7 @@ export default function Users({ capabilities, isAdmin }: function toggleEntryKey(name: string, key: string) { setRoleEdits((prev) => { - const cur = prev[name] ?? { entry_keys: [], description: "", capabilities: [] }; + const cur = prev[name] ?? { entry_keys: [], description: "", capabilities: [], daily_token_limit: 0 }; const has = cur.entry_keys.includes(key); return { ...prev, @@ -115,11 +118,18 @@ export default function Users({ capabilities, isAdmin }: function toggleRoleCap(name: string, key: Capability) { setRoleEdits((prev) => { - const cur = prev[name] ?? { entry_keys: [], description: "", capabilities: [] }; + const cur = prev[name] ?? { entry_keys: [], description: "", capabilities: [], daily_token_limit: 0 }; return { ...prev, [name]: { ...cur, capabilities: toggleCap(cur.capabilities, key) } }; }); } + function setRoleLimit(name: string, value: number) { + setRoleEdits((prev) => { + const cur = prev[name] ?? { entry_keys: [], description: "", capabilities: [], daily_token_limit: 0 }; + return { ...prev, [name]: { ...cur, daily_token_limit: Math.max(0, value || 0) } }; + }); + } + function memberCount(roleName: string): number { return users.filter((u) => u.role === roleName).length; } @@ -342,7 +352,8 @@ export default function Users({ capabilities, isAdmin }:
{roles.map((r) => { const edit = roleEdits[r.name] - ?? { entry_keys: r.entry_keys, description: r.description, capabilities: r.capabilities ?? [] }; + ?? { entry_keys: r.entry_keys, description: r.description, + capabilities: r.capabilities ?? [], daily_token_limit: r.daily_token_limit ?? 0 }; const count = memberCount(r.name); return (
@@ -355,7 +366,8 @@ export default function Users({ capabilities, isAdmin }: {canManage && } {canManage && (confirmingRole === r.name ? ( @@ -398,6 +410,13 @@ export default function Users({ capabilities, isAdmin }: ...prev, [r.name]: { ...edit, description: e.target.value }, }))} />
+ +
+ + setRoleLimit(r.name, parseInt(e.target.value, 10))} /> +
); })} @@ -522,6 +541,13 @@ export default function Users({ capabilities, isAdmin }: setNewRole({ ...newRole, description: e.target.value })} /> +
+ + setNewRole({ ...newRole, + daily_token_limit: Math.max(0, parseInt(e.target.value, 10) || 0) })} /> +
Capabilities
From c05ca9483c7b3d0e2b92df859ca19f9d1f4b5027 Mon Sep 17 00:00:00 2001 From: Gang Tao Date: Mon, 17 Aug 2026 16:27:02 -0700 Subject: [PATCH 3/7] feat(chat): show the user their daily token budget (used / left / reset) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Surface the budget proactively, not only on a 429. Adds GET /chat/usage (chat-cap gated) returning {limited, used, limit, remaining, reset} for the current user — limited:false for admins, unlimited roles, or when no budget is enforced, so the UI shows no indicator. Extracts the effective-limit rule into usage.effective_daily_limit (role's own limit, else the global fallback), reused by /chat enforcement and the new endpoint. Chat UI fetches /chat/usage on load and after each turn (a turn spends tokens; the append-only read lags slightly so it refreshes again shortly after), and shows "used / limit", tokens left, and the reset time in the header and empty state for users who have a budget. Tests: /chat/usage for a limited user, and unlimited for admin / no-store. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01XPzYZXxyTdj5G25KoujpHb --- src/tpk/server.py | 32 ++++++++++++++++++++++++++++---- src/tpk/usage.py | 8 ++++++++ tests/test_server.py | 24 ++++++++++++++++++++++++ web/src/Chat.tsx | 43 +++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 103 insertions(+), 4 deletions(-) diff --git a/src/tpk/server.py b/src/tpk/server.py index 54cd215..7c0bd1d 100644 --- a/src/tpk/server.py +++ b/src/tpk/server.py @@ -161,6 +161,7 @@ def create_app( 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") @@ -221,6 +222,32 @@ 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, + unlimited roles, 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: any failure degrades to unlimited.""" + 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(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)] @@ -253,8 +280,7 @@ async def chat(req: ChatRequest, user: User = Depends(auth.require_cap(auth_mod. scope = frozenset(role.entry_keys) if role else frozenset() # Effective daily token budget: the role's own limit, else the # global fallback (config). admin is never limited (branch skipped). - role_limit = role.daily_token_limit if role else 0 - turn_limit = role_limit if role_limit > 0 else daily_token_limit() + turn_limit = effective_daily_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 @@ -263,8 +289,6 @@ async def chat(req: ChatRequest, user: User = Depends(auth.require_cap(auth_mod. if usage is not None and turn_limit > 0: used = await run_in_threadpool(lambda: usage.used_today(user.username)) if used >= turn_limit: - from tpk.usage import day_window - _, reset = day_window() raise HTTPException( status_code=429, diff --git a/src/tpk/usage.py b/src/tpk/usage.py index 46166fb..bb0f297 100644 --- a/src/tpk/usage.py +++ b/src/tpk/usage.py @@ -21,6 +21,14 @@ USAGE_COLUMNS = ["ts", "username", "tokens"] +def effective_daily_limit(role, global_default: int) -> int: + """The daily token budget that applies to a user: their role's own limit + if it sets one, else the global fallback. 0 means unlimited. `role` may be + a Role or None (unreadable/absent role -> falls back to the global).""" + 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.""" diff --git a/tests/test_server.py b/tests/test_server.py index 37ebbf3..bd3f772 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -158,6 +158,30 @@ def test_chat_no_usage_store_skips_enforcement(monkeypatch): assert client.post("/chat", json={"message": "q"}).status_code == 200 +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/web/src/Chat.tsx b/web/src/Chat.tsx index 1a90eee..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; @@ -402,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); } } @@ -526,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" })} +
+ + )}
)} @@ -547,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) => (