From 21bf30d2a653284ed52f2904bc69b13453188d40 Mon Sep 17 00:00:00 2001
From: yumoqing <yumoqing@gmail.com>
Date: Thu, 30 Jul 2026 10:49:06 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20add=20finally=20block=20to=20write=5Fllm?=
 =?UTF-8?q?usage=20=E2=80=94=20recover=20from=20GeneratorExit=20on=20clien?=
 =?UTF-8?q?t=20disconnect?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

When client disconnects mid-stream, Python GC kills the generator with
GeneratorExit (BaseException subclass), which is NOT caught by except Exception.
write_llmusage at line 93 was never reached, causing ~16% usage records lost.
finally block ensures llmusage is written even on GeneratorExit.
---
 llmage/llmclient.py | 13 +++++++++++--
 1 file changed, 11 insertions(+), 2 deletions(-)

diff --git a/llmage/llmclient.py b/llmage/llmclient.py
index 79db6c4..9bda9df 100644
--- a/llmage/llmclient.py
+++ b/llmage/llmclient.py
@@ -27,6 +27,7 @@ async def uapi_request(request, llm, callerid, callerorgid, params_kw=None):
 	outlines = []
 	txt = ''
 	luid = getID()
+	llmusage = None
 	try:
 		start_timestamp = time.time()
 		responsed_seconds = None
@@ -96,10 +97,18 @@ async def uapi_request(request, llm, callerid, callerorgid, params_kw=None):
 		estr = erase_apikey(e)
 		ed = {"error": f"ERROR:{estr}", "status": "FAILED" ,"llmusageid": luid}
 		s = json.dumps(ed, ensure_ascii=False)
-		s = ''.join(s.split('\n'))
+		s = ''.join(s.split('\\n'))
 		outlines.append(ed)
-		yield f'{s}\n'
+		yield f'{s}\\n'
 		return
+	finally:
+		# GeneratorExit / client disconnect — flush partial usage
+		if llmusage and llmusage.get('id') == luid:
+			try:
+				llmusage.status = llmusage.status or 'UNKNOWN'
+				await write_llmusage(llmusage)
+			except:
+				pass
 
 async def inference_generator(request, *args, params_kw=None, **kw):
 	env = request._run_ns.copy()