From f198255be2ec686993faa1e6f38d4e087a223822 Mon Sep 17 00:00:00 2001
From: Jacob Chung <chungjac@amazon.com>
Date: Fri, 18 Apr 2025 15:43:08 -0700
Subject: [PATCH] fix

---
 .../controllers/chat/telemetryHelper.ts       | 22 +++++++++
 packages/core/src/shared/db/chatDb/chatDb.ts  | 47 +++++++++++--------
 packages/core/src/shared/db/chatDb/util.ts    |  1 +
 3 files changed, 50 insertions(+), 20 deletions(-)

diff --git a/packages/core/src/codewhispererChat/controllers/chat/telemetryHelper.ts b/packages/core/src/codewhispererChat/controllers/chat/telemetryHelper.ts
index 862d8f1207e..25c3b370779 100644
--- a/packages/core/src/codewhispererChat/controllers/chat/telemetryHelper.ts
+++ b/packages/core/src/codewhispererChat/controllers/chat/telemetryHelper.ts
@@ -48,6 +48,8 @@ import { AdditionalContextPrompt } from '../../../amazonq/lsp/types'
 import { getUserPromptsDirectory, promptFileExtension } from '../../constants'
 import { isInDirectory } from '../../../shared/filesystemUtilities'
 import { CustomFormActionMessage } from '../../view/connector/connector'
+import { Database } from '../../../shared/db/chatDb/chatDb'
+import { Message } from '../../../shared/db/chatDb/util'
 
 export function logSendTelemetryEventFailure(error: any) {
     let requestId: string | undefined
@@ -217,6 +219,26 @@ export class CWCTelemetryHelper {
         telemetry.feedback_result.emit({ result: feedbackResult })
     }
 
+    public record_TODO(tabId: string, conversationId: string, message: Message) {
+        try {
+            telemetry.amazonq_addMessageHistory.run((span) => {
+                span.record({
+                    result: 'Succeeded',
+                    cwsprChatConversationId: conversationId,
+                    cwsprChatConversationType: 'todo',
+                    cwsprChatHistoryMessageCharacterCount: message.characterCount,
+                    cwsprChatMessageId: message.messageId,
+                    cwsprToolName: 'todo',
+                    cwsprToolUseId: 'todo',
+                    // message.userInputMessageContext?.toolResults
+                    // message.toolUses
+                })
+            })
+        } catch (e: any) {
+            getLogger().error('Unable to record ______ telemetry')
+        }
+    }
+
     public recordToolUseSuggested(toolUse: ToolUse, messageId: string) {
         try {
             telemetry.amazonq_toolUseSuggested.run((span) => {
diff --git a/packages/core/src/shared/db/chatDb/chatDb.ts b/packages/core/src/shared/db/chatDb/chatDb.ts
index 1cafa5990c9..2f1fe572f5d 100644
--- a/packages/core/src/shared/db/chatDb/chatDb.ts
+++ b/packages/core/src/shared/db/chatDb/chatDb.ts
@@ -21,6 +21,7 @@ import path from 'path'
 import { fs } from '../../fs/fs'
 import { getLogger } from '../../logger/logger'
 import { ChatMessage, ToolResultStatus } from '@amzn/codewhisperer-streaming'
+import { CWCTelemetryHelper } from '../../../codewhispererChat/controllers/chat/telemetryHelper'
 
 // Maximum number of characters to keep in history
 const MaxConversationHistoryCharacters = 600_000
@@ -275,6 +276,7 @@ export class Database {
                     ? message.body
                     : tabData?.title || 'Amazon Q Chat'
             message = this.formatChatHistoryMessage(message)
+            message.characterCount = this.calculateMessageCharacterCount(message)
             if (tabData) {
                 this.logger.info(`Found existing tab data, updating conversations`)
                 tabData.conversations = updateOrCreateConversation(tabData.conversations, conversationId, message)
@@ -292,6 +294,7 @@ export class Database {
                     conversations: [{ conversationId, clientType: ClientType.VSCode, messages: [message] }],
                 })
             }
+            CWCTelemetryHelper.instance.record_TODO(tabId, conversationId, message)
         }
     }
 
@@ -419,31 +422,35 @@ export class Database {
     private calculateCharacterCount(allMessages: Message[]): number {
         let count = 0
         for (const message of allMessages) {
-            // Count characters of all message text
-            count += message.body.length
-
-            // Count characters in tool uses
-            if (message.toolUses) {
-                try {
-                    for (const toolUse of message.toolUses) {
-                        count += JSON.stringify(toolUse).length
-                    }
-                } catch (e) {
-                    this.logger.error(`Error counting toolUses: ${String(e)}`)
+            count += message.characterCount ?? 0
+        }
+        this.logger.debug(`Current history characters: ${count}`)
+        return count
+    }
+
+    private calculateMessageCharacterCount(message: Message): number {
+        let count = message.body.length
+
+        // Count characters in tool uses
+        if (message.toolUses) {
+            try {
+                for (const toolUse of message.toolUses) {
+                    count += JSON.stringify(toolUse).length
                 }
+            } catch (e) {
+                this.logger.error(`Error counting toolUses: ${String(e)}`)
             }
-            // Count characters in tool results
-            if (message.userInputMessageContext?.toolResults) {
-                try {
-                    for (const toolResul of message.userInputMessageContext.toolResults) {
-                        count += JSON.stringify(toolResul).length
-                    }
-                } catch (e) {
-                    this.logger.error(`Error counting toolResults: ${String(e)}`)
+        }
+        // Count characters in tool results
+        if (message.userInputMessageContext?.toolResults) {
+            try {
+                for (const toolResul of message.userInputMessageContext.toolResults) {
+                    count += JSON.stringify(toolResul).length
                 }
+            } catch (e) {
+                this.logger.error(`Error counting toolResults: ${String(e)}`)
             }
         }
-        this.logger.debug(`Current history characters: ${count}`)
         return count
     }
 
diff --git a/packages/core/src/shared/db/chatDb/util.ts b/packages/core/src/shared/db/chatDb/util.ts
index 8995db90e8a..b884045e07a 100644
--- a/packages/core/src/shared/db/chatDb/util.ts
+++ b/packages/core/src/shared/db/chatDb/util.ts
@@ -57,6 +57,7 @@ export type Message = {
     userInputMessageContext?: UserInputMessageContext
     toolUses?: ToolUse[]
     shouldDisplayMessage?: boolean
+    characterCount?: number
 }
 
 /**