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