Billing counters: UPDATE first, INSERT ON CONFLICT only on first use

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.

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},
     )
 
 

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(

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)