feat(balance): extend atomic balance reserve to async+sync+product paths

- balance.py: module-level redis.asyncio singleton (works in all
  processes without env.redis injection), centralized tpac/self-org
  skip policy, userid param, extend_reserve, invalidate_balance_cache,
  finalize updates max_cost on cold start (no reserve)
- llmclient.py: unified reserve injection in _inference_generator
  (covers all 22 dspy entries); fix luid mismatch (reserve key ==
  llmusage id == finalize key); stream exception refund
- syncinference.py / asyncinference.py: luid reuse + failure refunds
  (submit failure, outer exception, query_task_status FAILED)
- product_interface.py: reserve + refund on product stream/non-stream
- accounting.py: finalize after llm_accounting (was missing for
  direct llmage accounting path -> reserve leaked until TTL)
- init.py: lambda signatures synced
- v1/chat/completions/index.dspy: pass userid to reserve

llmage/accounting.py

@@ -195,6 +195,13 @@ async def llm_accounting(llmusage):
 			'accounting_status': 'accounted'
 		}
 		await sor.U('llmusage', ns)
+	# Finalize Redis balance reservation (non-critical)
+	try:
+		from .balance import finalize_balance
+		await finalize_balance(env, llmusage.id, trans_amount,
+						llmid=llmusage.llmid)
+	except Exception as e:
+		exception(f'finalize_balance failed (non-critical): {e}')
 
 async def get_accounting_llmusages(luid=None):
 	env = ServerEnv()

llmage/asyncinference.py

@@ -13,6 +13,7 @@ from appPublic.base64_to_file import base64_to_file, getFilenameFromBase64
 from ahserver.serverenv import get_serverenv, ServerEnv
 from ahserver.filestorage import FileStorage
 from .accounting import llm_accounting, llm_charging
+from .balance import refund_balance
 from .utils import *
 
 # Global set to keep references to background tasks
@@ -67,7 +68,7 @@ async def async_uapi_request(request, llm,
 	uapi = env.UpAppApi(request)
 	userid = await env.uapi_data.get_calluserid(llm.upappid, orgid=llm.ownerid)
 	b = None
-	luid = getID()
+	luid = params_kw.get('_luid') or getID()
 	try:
 		start_timestamp = time.time()
 		if llm.callbackurl:
@@ -77,6 +78,12 @@ async def async_uapi_request(request, llm,
 		try:
 			b = await uapi.call(llm.upappid, llm.apiname, userid, params=params_kw)
 		except Exception as e:
+			# Refund balance reservation on submission failure
+			try:
+				if luid:
+					await refund_balance(ServerEnv(), luid)
+			except Exception:
+				pass
 			estr = erase_apikey(e)
 			ed = {"error": f"ERROR:{estr}", "status": "FAILED"}
 			exception(f'{ed}')
@@ -123,6 +130,12 @@ async def async_uapi_request(request, llm,
 		task.add_done_callback(_background_tasks.discard)
 
 	except Exception as e:
+		# Refund balance reservation on outer failure
+		try:
+			if luid:
+				await refund_balance(ServerEnv(), luid)
+		except Exception:
+			pass
 		ed = {"error": f"ERROR:{e}", "status": "FAILED"}
 		s = json.dumps(ed, ensure_ascii=False)
 		s = ''.join(s.split('\n'))
@@ -229,6 +242,12 @@ async def query_task_status(request, luid, onetime=False):
 						ns['usages'] = json.dumps(new_output['usage'])
 					await append_new_llmoutput(llmusage.ioinfo, new_output)
 					await modify_llmusage(ns)
+					if llmusage.status == 'FAILED':
+						# Async task failed — refund the balance reservation
+						try:
+							await refund_balance(ServerEnv(), luid)
+						except Exception:
+							pass
 				if  llmusage.status in ['UNKNOWN', 'FAILED', 'SUCCEEDED']:
 					critical(f'finished .. {llmusage.status=}')
 					return

llmage/balance.py

@@ -2,14 +2,54 @@
 Redis atomic balance reservation for llmage.
 Prevents concurrent overspend: pre-deduct before inference, settle after accounting.
 
+Uses a module-level redis.asyncio singleton so the atomic reserve works in
+EVERY process (web workers + backend_accounting) without needing env.redis
+to be injected. This fixes the previous fatal issue where env.redis was never
+set, so reserve always silently fell back to the non-atomic DB check.
+
 Key patterns:
   balance:{userorgid}           — current pre-deducted balance (int, cents)
-  reserve:{luid}                — {userorgid}|{llmid}|{max_cost}  (TTL 600s)
+  reserve:{luid}                — {userorgid}|{llmid}|{max_cost}  (TTL)
   model:max_cost:{llmid}        — historical max actual customer charge (int, cents)
+
+Reserve TTL: 600s for fast (stream/sync) calls, 3600s for async tasks that
+may run for minutes before finalize/refund.
 """
 import asyncio
 from appPublic.log import debug, exception
 
+# ── Module-level async Redis singleton ───────────────────────
+_redis = None
+_redis_lock = asyncio.Lock()
+
+
+def _redis_url():
+    try:
+        from appPublic.jsonConfig import getConfig
+        config = getConfig()
+        url = getattr(config.website, 'session_redis', None)
+        if url:
+            u = getattr(url, 'url', None)
+            if u:
+                return u
+    except Exception:
+        pass
+    return "redis://127.0.0.1:6379"
+
+
+async def _get_redis():
+    """Return a shared redis.asyncio client (lazy singleton)."""
+    global _redis
+    if _redis is not None:
+        return _redis
+    async with _redis_lock:
+        if _redis is None:
+            import redis.asyncio as aioredis
+            _redis = await aioredis.from_url(
+                _redis_url(), decode_responses=True)
+    return _redis
+
+
 # ── Lua scripts ──────────────────────────────────────────────
 
 RESERVE_LUA = """
@@ -19,6 +59,7 @@ local reserve_key = KEYS[3]
 local max_cost = tonumber(ARGV[1])
 local db_balance = tonumber(ARGV[2])    -- 0 means no DB fallback
 local reserve_val = ARGV[3]
+local ttl = tonumber(ARGV[4])
 
 -- Get or init max_cost: use stored if higher, else set from arg
 local stored_max = redis.call('GET', cost_key)
@@ -48,7 +89,7 @@ if balance < max_cost then
     return {0, 'insufficient balance'}
 end
 redis.call('DECRBY', bal_key, max_cost)
-redis.call('SETEX', reserve_key, 600, reserve_val)
+redis.call('SETEX', reserve_key, ttl, reserve_val)
 return {1, max_cost}
 """
 
@@ -123,24 +164,69 @@ def _from_cents(c):
     return round(int(c) / 100, 2)
 
 
-async def _redis_eval(env, script, nkeys, *keys_and_args):
+async def _update_max_cost(llmid, actual_cost):
-    """Eval Lua script via env's Redis connection."""
+    """Bump model:max_cost if actual exceeds stored value; persist to DB.
-    if not hasattr(env, 'redis') or env.redis is None:
-        raise RuntimeError('Redis unavailable')
-    return await asyncio.to_thread(env.redis.eval, script, nkeys, *keys_and_args)
 
-
+    Keeps the pre-deduct baseline growing even when the original reserve was
-async def reserve_balance(env, llmid, userorgid, luid):
+    skipped (no_history cold start) — otherwise reserve would never activate.
-    """Atomically check balance and deduct max_cost via Redis Lua."""
+    """
     try:
+        redis = await _get_redis()
+        cost_key = f'model:max_cost:{llmid}'
+        cost_cents = _cents(actual_cost)
+        stored = await redis.get(cost_key)
+        if stored and int(stored) >= cost_cents:
+            return
+        await redis.set(cost_key, cost_cents)
+        from .utils import update_model_max_cost
+        await update_model_max_cost(llmid, actual_cost)
+    except Exception as e:
+        debug(f'_update_max_cost failed: {e}')
+
+
+async def reserve_balance(env, llmid, userorgid, luid, ttl=600, userid=None):
+    """Atomically check balance and deduct max_cost via Redis Lua.
+
+    Policy (centralized so every entry behaves the same):
+      - Self-owned org (llm.ownerid == userorgid): skip reserve.
+      - tpac user (external balance system): skip reserve.
+      - No pricing (ppid empty): skip reserve (availability handled elsewhere).
+
+    Returns:
+      {'ok': True, 'max_cost': X}            — reserved X
+      {'ok': True, 'max_cost': 0, 'skip': ...} — reserve skipped by policy
+      {'ok': True, 'max_cost': 0, 'no_history': True} — no max_cost data yet
+      {'ok': False, 'reason': ...}           — insufficient balance
+      {'ok': True, 'max_cost': 0, 'no_redis': True} — Redis down, DB fallback
+    """
+    try:
+        # ── Policy checks (need llm info) ──
+        try:
+            from .utils import get_llmage_llm, get_user_tpac
+            llm = await get_llmage_llm(llmid)
+            if llm and llm.ownerid == userorgid:
+                return {'ok': True, 'max_cost': 0, 'skip': 'self_org'}
+            if not llm or not llm.ppid:
+                return {'ok': True, 'max_cost': 0, 'skip': 'no_ppid'}
+            # tpac user: balance lives in external system, skip redis reserve
+            if userid:
+                try:
+                    tpac = await get_user_tpac(userid)
+                    if tpac:
+                        return {'ok': True, 'max_cost': 0, 'skip': 'tpac'}
+                except Exception as e:
+                    debug(f'reserve_balance: tpac check failed: {e}')
+        except Exception as e:
+            debug(f'reserve_balance: policy check failed: {e}')
+
+        redis = await _get_redis()
         bal_key = f'balance:{userorgid}'
         cost_key = f'model:max_cost:{llmid}'
         reserve_key = f'reserve:{luid}'
 
-        # Load max_cost from DB if not cached
+        # Load max_cost from Redis if cached, else from DB history
-        redis = env.redis
         max_cost = 0
-        stored = await asyncio.to_thread(redis.get, cost_key)
+        stored = await redis.get(cost_key)
         if stored:
             max_cost = int(stored)
         else:
@@ -155,9 +241,9 @@ async def reserve_balance(env, llmid, userorgid, luid):
         if max_cost <= 0:
             return {'ok': True, 'max_cost': 0, 'no_history': True}
 
-        # Load balance from DB if not cached
+        # Load balance from DB if not cached in Redis
         db_balance = 0
-        stored_bal = await asyncio.to_thread(redis.get, bal_key)
+        stored_bal = await redis.get(bal_key)
         if not stored_bal:
             try:
                 from accounting.getaccount import getCustomerBalance
@@ -170,48 +256,46 @@ async def reserve_balance(env, llmid, userorgid, luid):
                 debug(f'reserve_balance: getCustomerBalance failed: {e}')
 
         reserve_val = f'{userorgid}|{llmid}|{max_cost}'
-        result = await _redis_eval(env, RESERVE_LUA, 3,
+        result = await redis.eval(RESERVE_LUA, 3,
             bal_key, cost_key, reserve_key,
-            str(max_cost), str(db_balance), reserve_val)
+            str(max_cost), str(db_balance), reserve_val, str(ttl))
 
         if result[0] == 0:
             return {'ok': False, 'reason': result[1]}
         return {'ok': True, 'max_cost': _from_cents(result[1])}
 
-    except RuntimeError:
-        debug('reserve_balance: Redis not configured, skip')
-        return {'ok': True, 'max_cost': 0, 'no_redis': True}
     except Exception as e:
         # Redis down or unreachable — fall back to DB check, don't block
         debug(f'reserve_balance: Redis error, falling back to DB: {e}')
         return {'ok': True, 'max_cost': 0, 'no_redis': True}
 
 
-async def finalize_balance(env, luid, actual_cost):
+async def finalize_balance(env, luid, actual_cost, llmid=None):
-    """After accounting: adjust balance, update max_cost."""
+    """After accounting: adjust balance, update max_cost.
+
+    llmid: needed for cold-start max_cost updates when no reserve existed
+    (reserve skipped due to no_history) — keeps the pre-deduct baseline
+    growing so future calls can reserve.
+    """
     try:
-        cost_key = 'model:max_cost:dummy'
+        redis = await _get_redis()
         reserve_key = f'reserve:{luid}'
 
-        # We need the llmid to get the correct cost_key, but we won't know it
+        reserve_val = await redis.get(reserve_key)
-        # until parsing the reserve. So use a placeholder and let Lua handle it.
-        # Actually FINALIZE_LUA only uses cost_key for updating, so pass a dummy.
-        # The real cost_key needs the llmid which comes from reserve parsing.
-        
-        redis = env.redis
-        reserve_val = await asyncio.to_thread(redis.get, reserve_key)
         if not reserve_val:
             debug(f'finalize_balance: no reserve for luid={luid}')
+            if llmid:
+                await _update_max_cost(llmid, actual_cost)
             return
 
-        # Parse to get llmid
+        # Parse to get llmid (reserve_val: userorgid|llmid|max_cost)
         parts = reserve_val.split('|')
         if len(parts) < 3:
             return
         llmid = parts[1]
         cost_key = f'model:max_cost:{llmid}'
 
-        result = await _redis_eval(env, FINALIZE_LUA, 2,
+        result = await redis.eval(FINALIZE_LUA, 2,
             cost_key, reserve_key, str(_cents(actual_cost)))
 
         if result[0] == 0:
@@ -224,14 +308,8 @@ async def finalize_balance(env, luid, actual_cost):
 
         # Persist max_cost to DB if actual exceeded previous max
         if actual_cost > max_c:
-            try:
+            await _update_max_cost(llmid, actual_cost)
-                from .utils import update_model_max_cost
-                await update_model_max_cost(llmid, actual_cost)
-            except Exception as e:
-                debug(f'finalize_balance: DB persist failed: {e}')
 
-    except RuntimeError:
-        debug('finalize_balance: Redis unavailable, skip')
     except Exception as e:
         exception(f'finalize_balance error: {e}')
 
@@ -239,13 +317,39 @@ async def finalize_balance(env, luid, actual_cost):
 async def refund_balance(env, luid):
     """Full refund on API failure."""
     try:
+        redis = await _get_redis()
         reserve_key = f'reserve:{luid}'
-        result = await _redis_eval(env, REFUND_LUA, 1, reserve_key)
+        result = await redis.eval(REFUND_LUA, 1, reserve_key)
         if result[0] == 1:
             debug(f'refund_balance: luid={luid} refund={_from_cents(result[1])}')
         else:
             debug(f'refund_balance: luid={luid} skipped ({result[1]})')
-    except RuntimeError:
-        pass
     except Exception as e:
         exception(f'refund_balance error: {e}')
+
+
+async def extend_reserve(env, luid, ttl):
+    """Extend TTL of an existing reserve.
+
+    Used when an entry pre-reserved with the short (600s) TTL but the
+    request is dispatched to async mode, where the task may run longer.
+    """
+    try:
+        redis = await _get_redis()
+        await redis.expire(f'reserve:{luid}', ttl)
+    except Exception as e:
+        debug(f'extend_reserve failed: {e}')
+
+
+async def invalidate_balance_cache(userorgid):
+    """Delete the cached Redis balance so the next reserve reloads from DB.
+
+    Must be called after recharge / recharge reversal, otherwise the
+    pre-deduct baseline stays stale and valid requests get rejected.
+    """
+    try:
+        redis = await _get_redis()
+        await redis.delete(f'balance:{userorgid}')
+        debug(f'invalidate_balance_cache: balance:{userorgid} deleted')
+    except Exception as e:
+        debug(f'invalidate_balance_cache failed: {e}')

llmage/init.py

@@ -173,8 +173,8 @@ def load_llmage():
 	env._reserve_balance = reserve_balance
 	env._finalize_balance = finalize_balance
 	env._refund_balance = refund_balance
-	env.reserve_balance = lambda llmid, userorgid, luid: reserve_balance(env, llmid, userorgid, luid)
+	env.reserve_balance = lambda llmid, userorgid, luid, ttl=600, userid=None: reserve_balance(env, llmid, userorgid, luid, ttl=ttl, userid=userid)
-	env.finalize_balance = lambda luid, actual_cost: finalize_balance(env, luid, actual_cost)
+	env.finalize_balance = lambda luid, actual_cost, llmid=None: finalize_balance(env, luid, actual_cost, llmid=llmid)
 	env.refund_balance = lambda luid: refund_balance(env, luid)
 	env.backup_accounted_llmusage = backup_accounted_llmusage
 	env.read_ioinfo_content = read_ioinfo_content

llmage/llmclient.py

@@ -14,7 +14,7 @@ from ahserver.filestorage import FileStorage
 from .asyncinference import async_uapi_request
 from .syncinference import sync_uapi_request
 from .accounting import llm_accounting, llm_charging
-from .balance import refund_balance
+from .balance import refund_balance, reserve_balance, extend_reserve
 from .utils import *
 
 async def uapi_request(request, llm, callerid, callerorgid, params_kw=None):
@@ -27,7 +27,7 @@ async def uapi_request(request, llm, callerid, callerorgid, params_kw=None):
 	userid = await env.uapi_data.get_calluserid(llm.upappid, orgid=llm.ownerid)
 	outlines = []
 	txt = ''
-	luid = getID()
+	luid = params_kw.get('_luid') or getID()
 	llmusage = None
 	try:
 		start_timestamp = time.time()
@@ -161,6 +161,29 @@ async def _inference_generator(request, callerid, callerorgid,
 		yield errmsg
 		return
 	params_kw.model = llm.model
+	# ── Unified balance reserve ────────────────────────────────────────
+	# Covers every entry that did not pre-reserve. v1/chat/completions
+	# reserves at dspy level and passes _luid; everyone else reserves here.
+	# Reusing _luid downstream also fixes the old mismatch where the
+	# reserve key could never equal the llmusage id used by finalize.
+	if params_kw.get('_luid'):
+		if llm.stream == 'async':
+			# Entry reserved with the short TTL; async tasks run longer
+			await extend_reserve(env, params_kw._luid, 3600)
+	else:
+		_luid = getID()
+		_ttl = 3600 if llm.stream == 'async' else 600
+		reserved = await reserve_balance(env, llm.id, callerorgid, _luid,
+							ttl=_ttl, userid=callerid)
+		if not reserved.get('ok'):
+			debug(f'balance reserve rejected: {reserved}')
+			errmsg = json.dumps({'status': 'FAILED',
+				'error': f'余额不足(balance reserve rejected): {reserved.get("reason")}'},
+				ensure_ascii=False) + '\n'
+			yield errmsg
+			return
+		params_kw._luid = _luid
+		params_kw._reserved = reserved
 	if llm.stream == 'async':
 		if llm.callbackurl:
 			cb_url = env.entire_url(llm.callbackurl)

llmage/product_interface.py

@@ -22,6 +22,7 @@ from .utils import (
     write_llmusage,
 )
 from .accounting import llm_charging
+from .balance import reserve_balance, refund_balance
 
 
 async def _resolve_llm(resource_ref_id):
@@ -219,24 +220,46 @@ async def execute_product_service(resource_ref_id, user_id, user_org_id, request
     if not params_kw.get('transno'):
         params_kw.transno = luid
 
+    # ── Balance reservation (product path) ─────────────────────────
+    # Same atomic pre-deduct as the direct dspy path; policy checks
+    # (self-org / tpac / no-ppid) live inside reserve_balance.
+    ttl = 3600 if full_llm.stream == 'async' else 600
+    reserved = await reserve_balance(env, full_llm.id, user_org_id, luid,
+                                     ttl=ttl, userid=user_id)
+    if not reserved.get('ok'):
+        return {'success': False,
+                'message': f'余额不足(balance reserve rejected): {reserved.get("reason")}',
+                'status': 'FAILED', 'task_id': luid}
+    params_kw._luid = luid
+
     try:
         if full_llm.stream == 'async':
             from .asyncinference import async_uapi_request_product
             result = await async_uapi_request_product(
                 full_llm, userid, user_id, user_org_id, params_kw, luid)
-            return result
-
         elif not full_llm.stream:
             from .syncinference import sync_uapi_request_product
             result = await sync_uapi_request_product(
                 full_llm, userid, user_id, user_org_id, params_kw, luid)
-            return result
         else:
             result = await _collect_stream(full_llm, userid, user_id,
                                            user_org_id, params_kw, luid)
+
+        # Sub-callers swallow exceptions and return success=False —
+        # refund here too (refund is idempotent via GETDEL)
+        if not result.get('success'):
+            try:
+                await refund_balance(env, luid)
+            except Exception:
+                pass
         return result
 
     except Exception as e:
+        # Refund balance reservation on failure
+        try:
+            await refund_balance(env, luid)
+        except Exception:
+            pass
         exception(f'execute_product_service error: {e}')
         return {'success': False, 'message': str(e), 'status': 'FAILED',
                 'task_id': luid}
@@ -349,6 +372,16 @@ async def execute_product_service_stream(resource_ref_id, user_id, user_org_id,
     if not params_kw.get('transno'):
         params_kw.transno = luid
 
+    # ── Balance reservation (product stream path) ──────────────────
+    reserved = await reserve_balance(env, full_llm.id, user_org_id, luid,
+                                     ttl=600, userid=user_id)
+    if not reserved.get('ok'):
+        yield {'chunk': None, 'usage_data': None, 'done': True,
+               'error': f'余额不足(balance reserve rejected): {reserved.get("reason")}',
+               'task_id': luid, 'status': 'FAILED'}
+        return
+    params_kw._luid = luid
+
     outlines = []
     txt = ''
     usage = None
@@ -416,6 +449,11 @@ async def execute_product_service_stream(resource_ref_id, user_id, user_org_id,
         }
 
     except Exception as e:
+        # Refund balance reservation on failure
+        try:
+            await refund_balance(env, luid)
+        except Exception:
+            pass
         exception(f'stream error: {e}')
         yield {'chunk': None, 'usage_data': None, 'done': True,
                'error': str(e), 'task_id': luid, 'status': 'FAILED'}

llmage/syncinference.py

@@ -14,6 +14,7 @@ from appPublic.base64_to_file import base64_to_file, getFilenameFromBase64
 from ahserver.serverenv import get_serverenv, ServerEnv
 from ahserver.filestorage import FileStorage
 from .accounting import llm_accounting, llm_charging
+from .balance import refund_balance
 from .utils import *
 
 async def sync_uapi_request(request, llm, callerid, callerorgid, params_kw=None):
@@ -28,7 +29,7 @@ async def sync_uapi_request(request, llm, callerid, callerorgid, params_kw=None)
 	outlines = []
 	b = None
 	d = None
-	luid = getID()
+	luid = params_kw.get('_luid') or getID()
 	try:
 		start_timestamp = time.time()
 		responsed_seconds = None
@@ -69,6 +70,12 @@ async def sync_uapi_request(request, llm, callerid, callerorgid, params_kw=None)
 		yield b
 		# await write_llmusage(llmusage)
 	except Exception as e:
+		# Refund balance reservation on failure
+		try:
+			if luid:
+				await refund_balance(ServerEnv(), luid)
+		except Exception:
+			pass
 		exception(f'{e=},{format_exc()}, {b=}')
 		estr = erase_apikey(e)
 		ed = {"error": f"ERROR:{estr}", "status": "FAILED" ,"llmusageid": luid}

wwwroot/v1/chat/completions/index.dspy

@@ -49,7 +49,7 @@ params_kw.llmcatelogid = catelogid
 
 debug(f'{params_kw.llmid=}')
 luid = getID()
-reserved = await env.reserve_balance(params_kw.llmid, userorgid, luid)
+reserved = await env.reserve_balance(params_kw.llmid, userorgid, luid, userid=userid)
 if not reserved.get('ok'):
 	debug(f'{userid=} balance not enough: {reserved}')
 	return openai_429()