From 670325e8ebd3c50afe766c8cdaab0b58450cec7f Mon Sep 17 00:00:00 2001
From: yumoqing <yumoqing@gmail.com>
Date: Fri, 31 Jul 2026 16:39:43 +0800
Subject: [PATCH] fix: move refund to uapi_request except, remove gen()
 try/except that re-raised causing 500
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

gen() try/except re-raised internal errors from uapi_request → 500 response.
Moved refund_balance to uapi_request's except block where errors are properly
handled (yield error, normal generator exit).
---
 llmage/llmclient.py                    |  9 +++++++++
 wwwroot/v1/chat/completions/index.dspy | 25 ++++++++-----------------
 2 files changed, 17 insertions(+), 17 deletions(-)

diff --git a/llmage/llmclient.py b/llmage/llmclient.py
index 9bda9df..2b493ee 100644
--- a/llmage/llmclient.py
+++ b/llmage/llmclient.py
@@ -93,6 +93,15 @@ async def uapi_request(request, llm, callerid, callerorgid, params_kw=None):
 		llmusage.accounting_status = 'created'
 		await write_llmusage(llmusage)
 	except Exception as e:
+		# Refund balance reservation on failure
+		try:
+			from .balance import refund_balance
+			from ahserver.serverenv import ServerEnv
+			reserve_luid = params_kw.get('_luid') if params_kw else None
+			if reserve_luid:
+				await refund_balance(ServerEnv(), reserve_luid)
+		except:
+			pass
 		exception(f'{e=},{format_exc()}')
 		estr = erase_apikey(e)
 		ed = {"error": f"ERROR:{estr}", "status": "FAILED" ,"llmusageid": luid}
diff --git a/wwwroot/v1/chat/completions/index.dspy b/wwwroot/v1/chat/completions/index.dspy
index 75ad5bc..6b726bf 100644
--- a/wwwroot/v1/chat/completions/index.dspy
+++ b/wwwroot/v1/chat/completions/index.dspy
@@ -1,19 +1,13 @@
 async def gen():
 	env = request._run_ns.copy()
 	f = partial(inference_generator, request, params_kw=params_kw)
-	try:
-		if params_kw.stream:
-			async for l in f():
-				yield f'data: {l}\n'
-			yield 'data: [DONE]\n\n'
-		else:
-			async for l in f():
-				yield l
-	except Exception as e:
-		luid = params_kw.get('_luid')
-		if luid:
-			await env.refund_balance(luid)
-		raise
+	if params_kw.stream:
+		async for l in f():
+			yield f'data: {l}\\n'
+		yield 'data: [DONE]\\n\\n'
+	else:
+		async for l in f():
+			yield l
 
 debug_params('params_kw', params_kw)
 catelogid = params_kw.catelogid or 't2t'
@@ -43,8 +37,7 @@ params_kw.llmid = llmid
 params_kw.llmcatelogid = catelogid
 
 debug(f'{params_kw.llmid=}')
-import uuid
-luid = str(uuid.uuid4())
+luid = getID()
 reserved = await env.reserve_balance(params_kw.llmid, userorgid, luid)
 if not reserved.get('ok'):
 	debug(f'{userid=} balance not enough: {reserved}')
@@ -55,6 +48,4 @@ f = await checkCustomerBalance(params_kw.llmid, userid, userorgid)
 if not f and reserved.get('no_redis'):
 	debug(f'{userid=} balance not enough (DB fallback)')
 	return openai_429()
-# debug(f'{tools=}, {request._run_ns.tools=}')
 return await env.stream_response(request, gen, content_type='application/json')
-