From 6f9bff1f1f0a3add1c134a699e1534565859a3ac Mon Sep 17 00:00:00 2001
From: ViperEkura <3081035982@qq.com>
Date: Mon, 13 Apr 2026 23:04:41 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E5=A2=9E=E5=8A=A0token=20=E7=BB=9F?=
 =?UTF-8?q?=E8=AE=A1?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 dashboard/src/components/MessageBubble.vue    | 27 ++++++++++++-
 dashboard/src/components/ProcessBlock.vue     |  5 ++-
 .../src/views/ConversationDetailView.vue      |  7 +++-
 luxx/models.py                                |  8 ++++
 luxx/routes/conversations.py                  | 18 +++++++++
 luxx/services/chat.py                         | 40 ++++++++++++++++---
 6 files changed, 95 insertions(+), 10 deletions(-)

diff --git a/dashboard/src/components/MessageBubble.vue b/dashboard/src/components/MessageBubble.vue
index 95880fc..7cd5db8 100644
--- a/dashboard/src/components/MessageBubble.vue
+++ b/dashboard/src/components/MessageBubble.vue
@@ -26,7 +26,12 @@
         </template>
       </div>
       <div class="message-footer">
-        <span class="token-count" v-if="tokenCount">{{ tokenCount }} tokens</span>
+        <span class="token-info" v-if="usage">
+          <span class="token-item" v-if="usage.prompt_tokens">输入: {{ usage.prompt_tokens }}</span>
+          <span class="token-item" v-if="usage.completion_tokens">输出: {{ usage.completion_tokens }}</span>
+          <span class="token-item total" v-if="usage.total_tokens">总计: {{ usage.total_tokens }}</span>
+        </span>
+        <span class="token-count" v-else-if="tokenCount">{{ tokenCount }} tokens</span>
         <span class="message-time">{{ formatTime(createdAt) }}</span>
         <button v-if="role === 'assistant'" class="ghost-btn accent" @click="copyContent" title="复制">
           <span v-html="copyIcon"></span>
@@ -50,6 +55,7 @@ const props = defineProps({
   toolCalls: { type: Array, default: () => [] },
   processSteps: { type: Array, default: () => [] },
   tokenCount: { type: Number, default: 0 },
+  usage: { type: Object, default: null },
   createdAt: { type: String, default: '' },
   deletable: { type: Boolean, default: false },
   attachments: { type: Array, default: () => [] },
@@ -131,6 +137,25 @@ const trashIcon = `<svg viewBox="0 0 24 24" width="14" height="14" fill="none" s
   font-size: 12px;
 }
 
+.token-info {
+  display: flex;
+  align-items: center;
+  gap: 6px;
+  font-size: 11px;
+  color: var(--text-tertiary);
+}
+
+.token-item {
+  padding: 2px 6px;
+  background: var(--bg-code);
+  border-radius: 4px;
+}
+
+.token-item.total {
+  font-weight: 600;
+  color: var(--accent-primary);
+}
+
 .token-count,
 .message-time {
   font-size: 12px;
diff --git a/dashboard/src/components/ProcessBlock.vue b/dashboard/src/components/ProcessBlock.vue
index 810a8ca..530ad37 100644
--- a/dashboard/src/components/ProcessBlock.vue
+++ b/dashboard/src/components/ProcessBlock.vue
@@ -276,6 +276,8 @@ const alertIcon = `<svg viewBox="0 0 24 24" width="16" height="16" fill="none" s
   cursor: pointer;
   font-size: 13px;
   transition: background 0.15s;
+  width: 100%;
+  box-sizing: border-box;
 }
 
 .thinking .step-header:hover,
@@ -370,8 +372,9 @@ const alertIcon = `<svg viewBox="0 0 24 24" width="16" height="16" fill="none" s
   overflow: hidden;
   text-overflow: ellipsis;
   white-space: nowrap;
-  flex: 1;
+  flex: 1 1 auto;
   min-width: 0;
+  max-width: 300px;
 }
 
 .arrow.open {
diff --git a/dashboard/src/views/ConversationDetailView.vue b/dashboard/src/views/ConversationDetailView.vue
index e012994..cfdedfb 100644
--- a/dashboard/src/views/ConversationDetailView.vue
+++ b/dashboard/src/views/ConversationDetailView.vue
@@ -41,6 +41,7 @@
               :tool-calls="msg.tool_calls"
               :process-steps="msg.process_steps"
               :token-count="msg.token_count"
+              :usage="msg.usage"
               :created-at="msg.created_at"
               :deletable="msg.role === 'user'"
               :attachments="msg.attachments"
@@ -233,13 +234,15 @@ const sendMessage = async () => {
           streamingMessage.value.process_steps.push(step)
         }
       },
-      onDone: () => {
+      onDone: (data) => {
         // 完成，添加到消息列表
         autoScroll.value = true
         if (streamingMessage.value) {
           messages.value.push({
             ...streamingMessage.value,
-            created_at: new Date().toISOString()
+            created_at: new Date().toISOString(),
+            token_count: data.token_count,
+            usage: data.usage
           })
           
           // 如果标题为空，自动用第一条用户消息作为标题
diff --git a/luxx/models.py b/luxx/models.py
index 3542dc8..fdf4c19 100644
--- a/luxx/models.py
+++ b/luxx/models.py
@@ -167,6 +167,7 @@ class Message(Base):
     role: Mapped[str] = mapped_column(String(16), nullable=False)  # user, assistant, system, tool
     content: Mapped[str] = mapped_column(Text, nullable=False, default="")
     token_count: Mapped[int] = mapped_column(Integer, default=0)
+    usage: Mapped[Optional[str]] = mapped_column(Text, nullable=True)  # JSON string for usage info
     created_at: Mapped[datetime] = mapped_column(DateTime, default=local_now)
     
     # Relationships
@@ -184,6 +185,13 @@ class Message(Base):
             "created_at": self.created_at.isoformat() if self.created_at else None
         }
         
+        # Parse usage JSON
+        if self.usage:
+            try:
+                result["usage"] = json.loads(self.usage)
+            except json.JSONDecodeError:
+                result["usage"] = None
+        
         # Parse content JSON
         try:
             content_obj = json.loads(self.content) if self.content else {}
diff --git a/luxx/routes/conversations.py b/luxx/routes/conversations.py
index 9722a05..d4c0b11 100644
--- a/luxx/routes/conversations.py
+++ b/luxx/routes/conversations.py
@@ -43,6 +43,7 @@ def list_conversations(
     db: Session = Depends(get_db)
 ):
     """Get conversation list"""
+    import json
     query = db.query(Conversation).filter(Conversation.user_id == current_user.id)
     result = paginate(query.order_by(Conversation.updated_at.desc()), page, page_size)
     
@@ -56,6 +57,23 @@ def list_conversations(
         ).order_by(Message.created_at).first()
         if first_msg:
             conv_dict['first_message'] = first_msg.content[:50] + ('...' if len(first_msg.content) > 50 else '')
+        
+        # Calculate total tokens from all assistant messages in this conversation
+        assistant_messages = db.query(Message).filter(
+            Message.conversation_id == c.id,
+            Message.role == 'assistant'
+        ).all()
+        total_tokens = 0
+        for msg in assistant_messages:
+            total_tokens += msg.token_count or 0
+            # Also try to get usage from the usage field
+            if msg.usage:
+                try:
+                    usage_obj = json.loads(msg.usage)
+                    total_tokens = usage_obj.get("total_tokens", total_tokens)
+                except:
+                    pass
+        conv_dict['token_count'] = total_tokens
         items.append(conv_dict)
     
     return success_response(data={
diff --git a/luxx/services/chat.py b/luxx/services/chat.py
index c53ab9b..bfcac63 100644
--- a/luxx/services/chat.py
+++ b/luxx/services/chat.py
@@ -1,6 +1,7 @@
 """Chat service module"""
 import json
 import uuid
+import logging
 from typing import List, Dict, Any, AsyncGenerator, Optional
 
 from luxx.models import Conversation, Message
@@ -9,7 +10,7 @@ from luxx.tools.core import registry
 from luxx.services.llm_client import LLMClient
 from luxx.config import config
 
-
+logger = logging.getLogger(__name__)
 # Maximum iterations to prevent infinite loops
 MAX_ITERATIONS = 10
 
@@ -130,6 +131,13 @@ class ChatService:
             all_tool_results = []
             step_index = 0
             
+            # Token usage tracking
+            total_usage = {
+                "prompt_tokens": 0,
+                "completion_tokens": 0,
+                "total_tokens": 0
+            }
+            
             # Global step IDs for thinking and text (persist across iterations)
             thinking_step_id = None
             thinking_step_idx = None
@@ -186,6 +194,13 @@ class ChatService:
                         yield _sse_event("error", {"content": f"Failed to parse response: {data_str}"})
                         return
                     
+                    # 提取 API 返回的 usage 信息
+                    if "usage" in chunk:
+                        usage = chunk["usage"]
+                        total_usage["prompt_tokens"] = usage.get("prompt_tokens", 0)
+                        total_usage["completion_tokens"] = usage.get("completion_tokens", 0)
+                        total_usage["total_tokens"] = usage.get("total_tokens", 0)
+                    
                     # Check for error in response
                     if "error" in chunk:
                         error_msg = chunk["error"].get("message", str(chunk["error"]))
@@ -362,18 +377,26 @@ class ChatService:
                 
                 # No tool calls - final iteration, save message
                 msg_id = str(uuid.uuid4())
+                
+                # 使用 API 返回的真实 completion_tokens，如果 API 没返回则降级使用估算值
+                actual_token_count = total_usage.get("completion_tokens", 0) or len(full_content) // 4
+                logger.info(f"[TOKEN] total_usage: {total_usage}, actual_token_count: {actual_token_count}")
+                
                 self._save_message(
                     conversation.id,
                     msg_id,
                     full_content,
                     all_tool_calls,
                     all_tool_results,
-                    all_steps
+                    all_steps,
+                    actual_token_count,
+                    total_usage
                 )
                 
                 yield _sse_event("done", {
                     "message_id": msg_id,
-                    "token_count": len(full_content) // 4
+                    "token_count": actual_token_count,
+                    "usage": total_usage
                 })
                 return
             
@@ -386,7 +409,9 @@ class ChatService:
                     full_content,
                     all_tool_calls,
                     all_tool_results,
-                    all_steps
+                    all_steps,
+                    actual_token_count,
+                    total_usage
                 )
             yield _sse_event("error", {"content": "Exceeded maximum tool call iterations"})
             
@@ -400,7 +425,9 @@ class ChatService:
         full_content: str,
         all_tool_calls: list,
         all_tool_results: list,
-        all_steps: list
+        all_steps: list,
+        token_count: int = 0,
+        usage: dict = None
     ):
         """Save the assistant message to database."""
         from luxx.database import SessionLocal
@@ -420,7 +447,8 @@ class ChatService:
                 conversation_id=conversation_id,
                 role="assistant",
                 content=json.dumps(content_json, ensure_ascii=False),
-                token_count=len(full_content) // 4
+                token_count=token_count,
+                usage=json.dumps(usage) if usage else None
             )
             db.add(msg)
             db.commit()