From c5c9c2aa0295cc6fd2dbc5fc30fbdf20e3aa1620 Mon Sep 17 00:00:00 2001 From: Rob Kooper Date: Wed, 28 Oct 2026 01:17:46 -0500 Subject: [PATCH] Billing counters: UPDATE first, INSERT ON CONFLICT only on first use update_stats, subtract_coins and the conversation counter no longer SELECT the counter row before bumping it. A new _bump helper runs the atomic UPDATE and, only when it matches no row, falls back to the dialect INSERT ... ON CONFLICT DO UPDATE with the same increments, which also covers concurrent first requests. The savepoints and swallowed IntegrityErrors go away. A steady-state bill is one UPDATE per counter and draws no serial id. Part of #171 --- lumen/blueprints/api/routes.py | 2 +- lumen/blueprints/chat/routes.py | 20 +- lumen/services/llm.py | 137 +++++++------ .../test_billing_counters_postgres.py | 146 ++++++++++++++ tests/unit/test_billing_counters.py | 184 ++++++++++++++++++ 5 files changed, 403 insertions(+), 86 deletions(-) create mode 100644 tests/integration/test_billing_counters_postgres.py create mode 100644 tests/unit/test_billing_counters.py diff --git a/lumen/blueprints/api/routes.py b/lumen/blueprints/api/routes.py index 0ff5d5ae..9b71403a 100644 --- a/lumen/blueprints/api/routes.py +++ b/lumen/blueprints/api/routes.py @@ -495,7 +495,7 @@ def fill_missing(r): if row is None: row = EntityModelConsent(entity_id=entity_id, model_config_id=config.id) fill_missing(row) - # Race-safe first insert (same pattern as subtract_coins): two + # Race-safe first insert (savepoint + IntegrityError): two # concurrent first acks — overlapping retries, or two keys of one # entity — would otherwise make the loser's commit raise # IntegrityError and 500 instead of the promised idempotent 200. diff --git a/lumen/blueprints/chat/routes.py b/lumen/blueprints/chat/routes.py index 4fed35ae..ec3a399f 100644 --- a/lumen/blueprints/chat/routes.py +++ b/lumen/blueprints/chat/routes.py @@ -8,8 +8,6 @@ import pypdf from flask import Blueprint, Response, current_app, jsonify, render_template, request, session from sqlalchemy import and_, func, or_, select -from sqlalchemy import update as sa_update -from sqlalchemy.exc import IntegrityError from werkzeug.utils import secure_filename from lumen.decorators import login_required @@ -23,6 +21,7 @@ from lumen.services import model_cache from lumen.services.live_state import get_live_state from lumen.services.llm import ( + _bump, bulk_model_access_info, check_coin_budget, get_pool_limit, @@ -74,18 +73,13 @@ def _count_conversation_started(entity_id): """Bump the persistent conversation counter in entity_stats. The counter survives conversation deletion and disabled storage. Uses the - same race-safe ensure-row + atomic-increment pattern as llm.py's stats. + same UPDATE-first, ON CONFLICT-on-first-use helper as llm.py's stats. """ - if db.session.execute(select(EntityStat).filter_by(entity_id=entity_id)).scalar_one_or_none() is None: - try: - with db.session.begin_nested(): - db.session.add(EntityStat(entity_id=entity_id, requests=0, input_tokens=0, output_tokens=0, cost=0)) - except IntegrityError: - pass - db.session.execute( - sa_update(EntityStat) - .where(EntityStat.entity_id == entity_id) - .values(conversations=EntityStat.conversations + 1) + _bump( + EntityStat, + {"entity_id": entity_id}, + {"conversations": EntityStat.conversations + 1}, + {"requests": 0, "input_tokens": 0, "output_tokens": 0, "cost": 0, "conversations": 1}, ) diff --git a/lumen/services/llm.py b/lumen/services/llm.py index d814dce9..828cb94e 100644 --- a/lumen/services/llm.py +++ b/lumen/services/llm.py @@ -11,7 +11,7 @@ from flask import current_app, has_request_context, request from sqlalchemy import select from sqlalchemy import update as sa_update -from sqlalchemy.exc import IntegrityError +from sqlalchemy.dialects import postgresql, sqlite from sqlalchemy.ext.compiler import compiles from sqlalchemy.sql.expression import FunctionElement @@ -573,6 +573,33 @@ def get_effective_limit(entity_id: int, model_config_id: int, model_config=None) return get_pool_limit(entity_id) +def _bump(model, key: dict, increments: dict, insert_values: dict): + """Apply ``increments`` to the ``model`` row matching ``key``, creating it on first use. + + The row exists on every call after an entity's first, so the atomic UPDATE + runs alone: one statement, no existence SELECT. Only when it matches no row + does the dialect ``INSERT … ON CONFLICT (key) DO UPDATE`` run, inserting + ``key`` plus ``insert_values``; if a concurrent first request inserted the + row in between, the conflict applies the same ``increments`` to it, so both + requests are counted and nothing raises. Not an upsert every time: Postgres + draws the serial ``id`` (model_stats, entity_balances) before it detects + the conflict, so each steady-state call would burn a sequence value. + + ``increments`` reference the existing row through the model's columns, + which ON CONFLICT DO UPDATE resolves against the conflicting row. + """ + where = [getattr(model, col) == value for col, value in key.items()] + result = db.session.execute(sa_update(model).where(*where).values(**increments)) + if result.rowcount: + return + dialect = postgresql if db.engine.dialect.name == "postgresql" else sqlite + db.session.execute( + dialect.insert(model) + .values(**key, **insert_values) + .on_conflict_do_update(index_elements=list(key), set_=increments) + ) + + def subtract_coins(entity_id: int, model_config_id: int, coin_cost: float, effective=_UNSET): """Deduct coin_cost from the entity's pool balance (no-op for unlimited or blocked). @@ -588,32 +615,16 @@ def subtract_coins(entity_id: int, model_config_id: int, coin_cost: float, effec if max_coins == -2: return - # Ensure balance row exists (first API use before login creates it). - # Check first: the row exists for every entity after its first request, so an - # unconditional INSERT would fail — and log a Postgres duplicate-key ERROR — - # on every subsequent request. Same race-safe pattern as update_stats below: - # if two concurrent first-requests both see None, the loser's IntegrityError - # is swallowed and the atomic UPDATE succeeds for both. - if db.session.execute( - select(EntityBalance).filter_by(entity_id=entity_id) - ).scalar_one_or_none() is None: - try: - with db.session.begin_nested(): - db.session.add(EntityBalance( - entity_id=entity_id, - coins_left=starting, - last_refill_at=utcnow(), - )) - except IntegrityError: - pass - # Single atomic deduction, floored at 0: deduct when affordable, otherwise zero # (the budget is a soft limit — see check_coin_budget). One statement so a # concurrent refill/credit can never be clobbered back to 0 by a separate zeroing. - db.session.execute( - sa_update(EntityBalance) - .where(EntityBalance.entity_id == entity_id) - .values(coins_left=_greatest(0, EntityBalance.coins_left - coin_cost)) + # The row is created on first use (first API use before login creates it), + # already deducted from the starting balance. + _bump( + EntityBalance, + {"entity_id": entity_id}, + {"coins_left": _greatest(0, EntityBalance.coins_left - coin_cost)}, + {"coins_left": max(0.0, float(starting) - float(coin_cost)), "last_refill_at": utcnow()}, ) db.session.flush() @@ -841,55 +852,37 @@ def update_stats( """ now = utcnow() - # Ensure ModelStat row exists before the atomic increment. - # Use try/except to handle the race where two concurrent first-requests both - # see None and both attempt the INSERT — the loser gets IntegrityError which - # we swallow; the atomic UPDATE below then succeeds for both. - if db.session.execute( - select(ModelStat).filter_by(entity_id=entity_id, model_config_id=model_config_id, source=source) - ).scalar_one_or_none() is None: - try: - with db.session.begin_nested(): - db.session.add(ModelStat( - entity_id=entity_id, model_config_id=model_config_id, source=source, - requests=0, input_tokens=0, output_tokens=0, cost=0, - )) - except IntegrityError: - pass - - db.session.execute( - sa_update(ModelStat) - .where(ModelStat.entity_id == entity_id, ModelStat.model_config_id == model_config_id, ModelStat.source == source) - .values( - requests=ModelStat.requests + 1, - input_tokens=ModelStat.input_tokens + input_tokens, - output_tokens=ModelStat.output_tokens + output_tokens, - audio_seconds=ModelStat.audio_seconds + audio_seconds, - cost=ModelStat.cost + cost, - last_used_at=now, - ) + _bump( + ModelStat, + {"entity_id": entity_id, "model_config_id": model_config_id, "source": source}, + { + "requests": ModelStat.requests + 1, + "input_tokens": ModelStat.input_tokens + input_tokens, + "output_tokens": ModelStat.output_tokens + output_tokens, + "audio_seconds": ModelStat.audio_seconds + audio_seconds, + "cost": ModelStat.cost + cost, + "last_used_at": now, + }, + { + "requests": 1, "input_tokens": input_tokens, "output_tokens": output_tokens, + "audio_seconds": audio_seconds, "cost": cost, "last_used_at": now, + }, ) - - # Ensure EntityStat row exists before the atomic increment. - # Same race-safe pattern as ModelStat above. - if db.session.execute(select(EntityStat).filter_by(entity_id=entity_id)).scalar_one_or_none() is None: - try: - with db.session.begin_nested(): - db.session.add(EntityStat(entity_id=entity_id, requests=0, input_tokens=0, output_tokens=0, cost=0)) - except IntegrityError: - pass - - db.session.execute( - sa_update(EntityStat) - .where(EntityStat.entity_id == entity_id) - .values( - requests=EntityStat.requests + 1, - input_tokens=EntityStat.input_tokens + input_tokens, - output_tokens=EntityStat.output_tokens + output_tokens, - audio_seconds=EntityStat.audio_seconds + audio_seconds, - cost=EntityStat.cost + cost, - last_used_at=now, - ) + _bump( + EntityStat, + {"entity_id": entity_id}, + { + "requests": EntityStat.requests + 1, + "input_tokens": EntityStat.input_tokens + input_tokens, + "output_tokens": EntityStat.output_tokens + output_tokens, + "audio_seconds": EntityStat.audio_seconds + audio_seconds, + "cost": EntityStat.cost + cost, + "last_used_at": now, + }, + { + "requests": 1, "input_tokens": input_tokens, "output_tokens": output_tokens, + "audio_seconds": audio_seconds, "cost": cost, "last_used_at": now, + }, ) log = RequestLog( diff --git a/tests/integration/test_billing_counters_postgres.py b/tests/integration/test_billing_counters_postgres.py new file mode 100644 index 00000000..a5d2833e --- /dev/null +++ b/tests/integration/test_billing_counters_postgres.py @@ -0,0 +1,146 @@ +"""Billing counters under real PostgreSQL concurrency. + +SQLite serialises writers, so the first-use race the ON CONFLICT fallback +exists for only happens here: several requests for a fresh entity all find no +row with their UPDATE and all reach the INSERT. Each has to be counted once, +without an error. The steady-state check pins why the fallback is not a plain +upsert every time: Postgres draws the serial ``id`` before it detects the +conflict, so an upsert per request would burn a sequence value per request. +""" + +import os +import threading +import uuid + +import pytest +from sqlalchemy import select, text + +from tests.conftest import no_snapshot_refresher +from tests.integration.conftest import TEST_CONFIG + +pytestmark = pytest.mark.postgres + +THREADS = 8 +STARTING = 100 +EFFECTIVE = (float(STARTING), 0.0, float(STARTING)) + + +@pytest.fixture(scope="module") +def pg_app(pg_migrated_isolated): + """A Flask app bound to this module's migrated PostgreSQL database. + + ``config.Config`` reads ``DATABASE_URL`` at class-definition time and the root + conftest's SQLite app has already imported it, so the attribute is patched too. + """ + url = pg_migrated_isolated[0] + previous_env = os.environ.get("DATABASE_URL") + os.environ["DATABASE_URL"] = url + os.environ["CONFIG_YAML"] = TEST_CONFIG + os.environ["BACKGROUND_WORKER"] = "false" + from config import Config + previous_uri = Config.SQLALCHEMY_DATABASE_URI + Config.SQLALCHEMY_DATABASE_URI = url + from lumen import create_app + with pytest.MonkeyPatch.context() as mp: + no_snapshot_refresher(mp) + application = create_app() + Config.SQLALCHEMY_DATABASE_URI = previous_uri + application.config["TESTING"] = True + with application.app_context(): + from lumen.extensions import db + assert db.engine.dialect.name == "postgresql" + yield application + if previous_env is None: + os.environ.pop("DATABASE_URL", None) + else: + os.environ["DATABASE_URL"] = previous_env + + +@pytest.fixture +def fresh_entity(pg_app): + """A new entity with a coin limit but no balance or stats rows yet, and a model.""" + from lumen.extensions import db + from lumen.models.entity import Entity + from lumen.models.entity_limit import EntityLimit + from lumen.models.model_config import ModelConfig + with pg_app.app_context(): + tag = uuid.uuid4().hex[:8] + entity = Entity(entity_type="user", email=f"billing-{tag}@example.com", + name=f"Billing {tag}", initials="BI", active=True) + model = ModelConfig(model_name=f"billing-{tag}", input_cost_per_million=1.0, + output_cost_per_million=1.0) + db.session.add_all([entity, model]) + db.session.flush() + db.session.add(EntityLimit(entity_id=entity.id, max_coins=STARTING, refresh_coins=0, + starting_coins=STARTING)) + db.session.commit() + return entity.id, model.id + + +def _bill(entity_id, model_id, cost): + from lumen.extensions import db + from lumen.services.llm import subtract_coins, update_stats + update_stats(entity_id, model_id, "api", 10, 5, cost) + subtract_coins(entity_id, model_id, cost, effective=EFFECTIVE) + db.session.commit() + + +def _counters(entity_id, model_id): + from lumen.extensions import db + from lumen.models.entity_balance import EntityBalance + from lumen.models.entity_stat import EntityStat + from lumen.models.model_stat import ModelStat + db.session.expire_all() + ms = db.session.execute(select(ModelStat).filter_by(entity_id=entity_id, model_config_id=model_id)).scalar_one() + es = db.session.get(EntityStat, entity_id) + bal = db.session.execute(select(EntityBalance).filter_by(entity_id=entity_id)).scalar_one() + return ms, es, bal + + +def test_concurrent_first_requests_are_each_counted_once(pg_app, fresh_entity): + entity_id, model_id = fresh_entity + costs = [float(i) for i in range(1, THREADS + 1)] + barrier = threading.Barrier(THREADS) + errors = [] + + def worker(cost): + try: + with pg_app.app_context(): + barrier.wait(timeout=30) + _bill(entity_id, model_id, cost) + except Exception as exc: # collected and asserted on below + errors.append(exc) + + threads = [threading.Thread(target=worker, args=(c,)) for c in costs] + for t in threads: + t.start() + for t in threads: + t.join(timeout=60) + + assert errors == [] + with pg_app.app_context(): + ms, es, bal = _counters(entity_id, model_id) + assert ms.requests == THREADS + assert es.requests == THREADS + assert ms.input_tokens == es.input_tokens == 10 * THREADS + assert float(ms.cost) == float(es.cost) == pytest.approx(sum(costs)) + assert float(bal.coins_left) == pytest.approx(STARTING - sum(costs)) + + +def test_steady_state_billing_draws_no_sequence_values(pg_app, fresh_entity): + entity_id, model_id = fresh_entity + seq = text("SELECT last_value, is_called FROM model_stats_id_seq " + "UNION ALL SELECT last_value, is_called FROM entity_balances_id_seq") + with pg_app.app_context(): + from lumen.extensions import db + _bill(entity_id, model_id, 1.0) # first use creates the rows + before = db.session.execute(seq).all() + db.session.commit() + + for _ in range(10): + _bill(entity_id, model_id, 1.0) + + assert db.session.execute(seq).all() == before + ms, es, bal = _counters(entity_id, model_id) + assert ms.requests == es.requests == 11 + assert float(bal.coins_left) == pytest.approx(STARTING - 11) diff --git a/tests/unit/test_billing_counters.py b/tests/unit/test_billing_counters.py new file mode 100644 index 00000000..9f16677f --- /dev/null +++ b/tests/unit/test_billing_counters.py @@ -0,0 +1,184 @@ +"""Billing counters: UPDATE first, INSERT … ON CONFLICT only on an entity's first use. + +The steady-state test records every statement the way +``tests/routes/test_preflight_queries.py`` does: once the counter rows exist, +``update_stats`` + ``subtract_coins`` run only their UPDATEs and the RequestLog +INSERT — no existence SELECT and no SAVEPOINT. + +The fallback tests force the UPDATE to match nothing while the row does exist, +which is what a concurrent first request looks like from the loser's side: the +row appears between its UPDATE and its INSERT. The ON CONFLICT branch then has +to apply the increment to the existing row, exactly once. +""" +from contextlib import contextmanager +from decimal import Decimal + +import pytest +from sqlalchemy import event +from sqlalchemy.sql.dml import Update + +STARTING = 100 +EFFECTIVE = (100.0, 0.0, float(STARTING)) + + +@contextmanager +def _statements(): + from lumen.extensions import db + log = [] + + def _record(conn, cursor, statement, parameters, context, executemany): + log.append(" ".join(statement.split())) + + event.listen(db.engine, "before_cursor_execute", _record) + try: + yield log + finally: + event.remove(db.engine, "before_cursor_execute", _record) + + +@contextmanager +def _updates_match_nothing(monkeypatch): + """Make every UPDATE report ``rowcount == 0`` without touching the database.""" + from lumen.extensions import db + real_execute = db.session.execute + + class _NoRows: + rowcount = 0 + + def execute(statement, *args, **kwargs): + if isinstance(statement, Update): + return _NoRows() + return real_execute(statement, *args, **kwargs) + + with monkeypatch.context() as mp: + mp.setattr(db.session, "execute", execute) + yield + + +def _seed(entity_id): + from lumen.extensions import db + from lumen.models.entity_limit import EntityLimit + db.session.add(EntityLimit(entity_id=entity_id, max_coins=100, refresh_coins=0, starting_coins=STARTING)) + db.session.commit() + + +def _bill(entity_id, model_id, cost=2.0): + from lumen.extensions import db + from lumen.services.llm import subtract_coins, update_stats + update_stats(entity_id, model_id, "api", 10, 5, cost) + subtract_coins(entity_id, model_id, cost, effective=EFFECTIVE) + db.session.commit() + + +def _counters(entity_id, model_id): + from sqlalchemy import select + + from lumen.extensions import db + from lumen.models.entity_balance import EntityBalance + from lumen.models.entity_stat import EntityStat + from lumen.models.model_stat import ModelStat + db.session.expire_all() + ms = db.session.execute(select(ModelStat).filter_by(entity_id=entity_id, model_config_id=model_id, source="api")).scalar_one() + es = db.session.get(EntityStat, entity_id) + bal = db.session.execute(select(EntityBalance).filter_by(entity_id=entity_id)).scalar_one() + return ms, es, bal + + +def test_steady_state_billing_runs_no_select_and_no_savepoint(app, test_user, test_model): + entity_id, model_id = test_user["id"], test_model["id"] + with app.app_context(): + _seed(entity_id) + _bill(entity_id, model_id) # first use creates the rows + + with _statements() as log: + _bill(entity_id, model_id) + + verbs = [s.split()[0].upper() for s in log] + assert "SELECT" not in verbs, log + assert not [s for s in log if "SAVEPOINT" in s.upper()], log + assert verbs == ["UPDATE", "UPDATE", "INSERT", "UPDATE"], log + assert "INTO request_logs" in log[2] + + +def test_first_use_inserts_each_counter_row_once(app, test_user, test_model): + entity_id, model_id = test_user["id"], test_model["id"] + with app.app_context(): + _seed(entity_id) + with _statements() as log: + _bill(entity_id, model_id, cost=3.0) + + assert not [s for s in log if s.upper().startswith("SELECT")], log + upserts = [s for s in log if "ON CONFLICT" in s] + assert [s.split()[2] for s in upserts] == ["model_stats", "entity_stats", "entity_balances"] + + ms, es, bal = _counters(entity_id, model_id) + assert (ms.requests, ms.input_tokens, ms.output_tokens) == (1, 10, 5) + assert (es.requests, es.input_tokens, es.output_tokens) == (1, 10, 5) + assert float(bal.coins_left) == STARTING - 3.0 + assert bal.last_refill_at is not None + + +def test_first_use_balance_is_floored_at_zero(app, test_user, test_model): + entity_id, model_id = test_user["id"], test_model["id"] + with app.app_context(): + _seed(entity_id) + _bill(entity_id, model_id, cost=STARTING + 50) + _, _, bal = _counters(entity_id, model_id) + assert float(bal.coins_left) == 0 + + +def test_on_conflict_fallback_adds_exactly_once(app, test_user, test_model, monkeypatch): + """Rows exist but the UPDATE matched none: ON CONFLICT bumps the existing rows once.""" + entity_id, model_id = test_user["id"], test_model["id"] + with app.app_context(): + _seed(entity_id) + _bill(entity_id, model_id, cost=2.0) + + with _updates_match_nothing(monkeypatch), _statements() as log: + _bill(entity_id, model_id, cost=5.0) + + assert len([s for s in log if "ON CONFLICT" in s]) == 3, log + ms, es, bal = _counters(entity_id, model_id) + assert (ms.requests, ms.input_tokens, ms.output_tokens) == (2, 20, 10) + assert float(ms.cost) == pytest.approx(7.0) + assert (es.requests, es.input_tokens, es.output_tokens) == (2, 20, 10) + assert float(es.cost) == pytest.approx(7.0) + assert float(bal.coins_left) == STARTING - 2.0 - 5.0 + + +def test_on_conflict_fallback_floors_the_existing_balance(app, test_user, test_model, monkeypatch): + entity_id, model_id = test_user["id"], test_model["id"] + with app.app_context(): + from lumen.extensions import db + from lumen.models.entity_balance import EntityBalance + _seed(entity_id) + db.session.add(EntityBalance(entity_id=entity_id, coins_left=Decimal("4"))) + db.session.commit() + + with _updates_match_nothing(monkeypatch): + _bill(entity_id, model_id, cost=10.0) + + _, _, bal = _counters(entity_id, model_id) + assert float(bal.coins_left) == 0 + + +def test_conversation_counter_first_use_and_fallback(app, test_user, monkeypatch): + entity_id = test_user["id"] + with app.app_context(): + from lumen.blueprints.chat.routes import _count_conversation_started + from lumen.extensions import db + from lumen.models.entity_stat import EntityStat + + with _statements() as log: + _count_conversation_started(entity_id) + db.session.commit() + assert not [s for s in log if s.upper().startswith("SELECT")], log + db.session.expire_all() + es = db.session.get(EntityStat, entity_id) + assert (es.conversations, es.requests) == (1, 0) + + with _updates_match_nothing(monkeypatch): + _count_conversation_started(entity_id) + db.session.commit() + db.session.expire_all() + assert db.session.get(EntityStat, entity_id).conversations == 2