From ae99b785796d28cd42a90bc0f4b5543378f3c441 Mon Sep 17 00:00:00 2001
From: hailin <hailin.zhao@gdzx.xyz>
Date: Sun, 25 Jan 2026 17:08:23 -0800
Subject: [PATCH] fix(conversation): track token usage and message count in
 conversation entity

Problem:
- Token usage was recorded to token_usage table but not to conversation entity
- Message count was not being incremented
- Dashboard showed 0 tokens for all conversations

Solution:
- Add inputTokens/outputTokens fields to StreamChunk interface
- Return token usage in 'end' chunk from ClaudeAgentServiceV2
- Capture token usage in conversation.service.ts sendMessage
- Call conversation.addTokens() and incrementMessageCount() after each exchange
- Consolidate conversation updates into single repo.update() call

Files changed:
- claude-agent-v2.service.ts: Add token fields to StreamChunk, return in 'end'
- conversation.service.ts: Track tokens and message counts properly

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
---
 .../services/conversation.service.ts          | 19 ++++++++++++++++---
 .../claude/claude-agent-v2.service.ts         |  7 +++++--
 2 files changed, 21 insertions(+), 5 deletions(-)

diff --git a/packages/services/conversation-service/src/application/services/conversation.service.ts b/packages/services/conversation-service/src/application/services/conversation.service.ts
index 6d490f5..13a2bd3 100644
--- a/packages/services/conversation-service/src/application/services/conversation.service.ts
+++ b/packages/services/conversation-service/src/application/services/conversation.service.ts
@@ -158,6 +158,8 @@ export class ConversationService {
     let fullResponse = '';
     const toolCalls: Array<{ name: string; input: unknown; result: unknown }> = [];
     let updatedState: ConversationContext['consultingState'] | undefined;
+    let inputTokens = 0;
+    let outputTokens = 0;
 
     // Stream response from Claude (with attachments for multimodal support)
     for await (const chunk of this.claudeAgentService.sendMessage(
@@ -181,6 +183,10 @@ export class ConversationService {
       } else if (chunk.type === 'state_update' && chunk.newState) {
         // V2: Capture updated consulting state
         updatedState = chunk.newState;
+      } else if (chunk.type === 'end') {
+        // Capture token usage from end chunk
+        inputTokens = chunk.inputTokens || 0;
+        outputTokens = chunk.outputTokens || 0;
       }
 
       yield chunk;
@@ -191,7 +197,6 @@ export class ConversationService {
       // Convert state to JSON-compatible format for database storage
       const stateForDb = JSON.parse(JSON.stringify(updatedState));
       conversation.updateConsultingState(stateForDb);
-      await this.conversationRepo.update(conversation);
     }
 
     // Save assistant response
@@ -205,12 +210,20 @@ export class ConversationService {
     });
     await this.messageRepo.save(assistantMessage);
 
+    // Update conversation statistics
+    conversation.incrementMessageCount('user');
+    conversation.incrementMessageCount('assistant');
+    conversation.addTokens(inputTokens, outputTokens);
+
     // Update conversation title if first message
-    if (conversation.messageCount === 0) {
+    if (conversation.messageCount === 2) {
+      // Now 2 after incrementing (user + assistant)
       const title = await this.generateTitle(params.content);
       conversation.title = title;
-      await this.conversationRepo.update(conversation);
     }
+
+    // Save all updates to conversation
+    await this.conversationRepo.update(conversation);
   }
 
   /**
diff --git a/packages/services/conversation-service/src/infrastructure/claude/claude-agent-v2.service.ts b/packages/services/conversation-service/src/infrastructure/claude/claude-agent-v2.service.ts
index 76efc2e..46c13dc 100644
--- a/packages/services/conversation-service/src/infrastructure/claude/claude-agent-v2.service.ts
+++ b/packages/services/conversation-service/src/infrastructure/claude/claude-agent-v2.service.ts
@@ -66,6 +66,9 @@ export interface StreamChunk {
   // V2新增
   stageName?: string;
   newState?: ConsultingState;
+  // Token usage (returned with 'end' chunk)
+  inputTokens?: number;
+  outputTokens?: number;
 }
 
 @Injectable()
@@ -342,7 +345,7 @@ export class ClaudeAgentServiceV2 implements OnModuleInit {
 
           // 7.6 返回更新后的状态
           yield { type: 'state_update', newState: state };
-          yield { type: 'end' };
+          yield { type: 'end', inputTokens: totalInputTokens, outputTokens: totalOutputTokens };
           return;
         }
 
@@ -403,7 +406,7 @@ export class ClaudeAgentServiceV2 implements OnModuleInit {
     }
 
     console.error('[ClaudeAgentV2] Tool loop exceeded maximum iterations');
-    yield { type: 'end' };
+    yield { type: 'end', inputTokens: totalInputTokens, outputTokens: totalOutputTokens };
   }
 
   /**