From 8ba4c6c22959c348d1c2860ed86999d27838fa8f Mon Sep 17 00:00:00 2001
From: yumoqing <yumoqing@gmail.com>
Date: Sun, 26 Apr 2026 13:26:15 +0800
Subject: [PATCH] bugfix

---
 llmage/asyncinference.py | 16 ++++++++++------
 1 file changed, 10 insertions(+), 6 deletions(-)

diff --git a/llmage/asyncinference.py b/llmage/asyncinference.py
index e5f221d..b0b6a7b 100644
--- a/llmage/asyncinference.py
+++ b/llmage/asyncinference.py
@@ -117,13 +117,10 @@ async def async_uapi_request(request, llm,
 		yield f'{s}\n'
 		return
 
-async def modify_llmusage_status(llmusage):
+async def modify_llmusage(ns):
 	env = ServerEnv()
 	async with get_sor_context(env, 'llmage') as sor:
-		await sor.U('llmusage', {
-			'id': llmusage.id,
-			'status': llmusage.status
-		})
+		await sor.U('llmusage', ns.copy())
 
 async def get_llm_llmusage(luid):
 	env = ServerEnv()
@@ -176,14 +173,21 @@ async def query_task_status(request, luid, onetime=False):
 				}
 			if lastoutout['status'] != new_output['status']:
 				llmusage.status = new_output['status']
+				ns = {
+					'id': llmusage.id,
+					'status': llmusage.status
+				}
+				if 'usage' in new_output.keys():
+					ns['usages'] = new_output['usage']
 				await append_new_llmoutput(llmusage.ioinfo, new_output)
-				await modify_llmusage_status(llmusage)
+				await modify_llmusage(ns)
 			if  llmusage.status in ['FAILED', 'SUCCEEDED']:
 				critical(f'finished .. {llmusage.status=}')
 				return
 
 			if onetime:
 				critical(f'onetime is true, returned')
+				return
 			await asyncio.sleep(llm.query_period or 30)
 			critical(f'{llm.query_period=} seconds will retry, {new_output["status"]=}')