From e668169ed9136014b87721548c587c7939e2dafc Mon Sep 17 00:00:00 2001
From: MFPires <muriloferreira2@gmail.com>
Date: Mon, 27 Jan 2025 23:02:25 -0300
Subject: [PATCH] feat: Add conversation context token counter

- Add contextTokens to ApiMetrics interface
- Calculate context size using last API request's tokens
- Display context token count in TaskHeader below total tokens
- Use exact token counts instead of character estimation

This helps users track the total size of their conversation context,
which is useful for managing context window limits.
---
 src/shared/getApiMetrics.ts                   | 12 ++++++++++++
 webview-ui/src/components/chat/ChatView.tsx   |  1 +
 webview-ui/src/components/chat/TaskHeader.tsx |  9 +++++++++
 3 files changed, 22 insertions(+)

diff --git a/src/shared/getApiMetrics.ts b/src/shared/getApiMetrics.ts
index bd7b1bb..b4caadc 100644
--- a/src/shared/getApiMetrics.ts
+++ b/src/shared/getApiMetrics.ts
@@ -6,6 +6,7 @@ interface ApiMetrics {
 	totalCacheWrites?: number
 	totalCacheReads?: number
 	totalCost: number
+	contextTokens: number // Total tokens in conversation (last message's tokensIn + tokensOut)
 }
 
 /**
@@ -32,8 +33,14 @@ export function getApiMetrics(messages: ClineMessage[]): ApiMetrics {
 		totalCacheWrites: undefined,
 		totalCacheReads: undefined,
 		totalCost: 0,
+		contextTokens: 0,
 	}
 
+	// Find the last api_req_started message to get the context size
+	const lastApiReq = [...messages]
+		.reverse()
+		.find((message) => message.type === "say" && message.say === "api_req_started" && message.text)
+
 	messages.forEach((message) => {
 		if (message.type === "say" && message.say === "api_req_started" && message.text) {
 			try {
@@ -55,6 +62,11 @@ export function getApiMetrics(messages: ClineMessage[]): ApiMetrics {
 				if (typeof cost === "number") {
 					result.totalCost += cost
 				}
+
+				// If this is the last api request, use its tokens for context size
+				if (message === lastApiReq) {
+					result.contextTokens = (tokensIn || 0) + (tokensOut || 0)
+				}
 			} catch (error) {
 				console.error("Error parsing JSON:", error)
 			}
diff --git a/webview-ui/src/components/chat/ChatView.tsx b/webview-ui/src/components/chat/ChatView.tsx
index da15b4e..7769f63 100644
--- a/webview-ui/src/components/chat/ChatView.tsx
+++ b/webview-ui/src/components/chat/ChatView.tsx
@@ -915,6 +915,7 @@ const ChatView = ({ isHidden, showAnnouncement, hideAnnouncement, showHistoryVie
 					cacheWrites={apiMetrics.totalCacheWrites}
 					cacheReads={apiMetrics.totalCacheReads}
 					totalCost={apiMetrics.totalCost}
+					contextTokens={apiMetrics.contextTokens}
 					onClose={handleTaskCloseButtonClick}
 				/>
 			) : (
diff --git a/webview-ui/src/components/chat/TaskHeader.tsx b/webview-ui/src/components/chat/TaskHeader.tsx
index 0b3f494..646dce3 100644
--- a/webview-ui/src/components/chat/TaskHeader.tsx
+++ b/webview-ui/src/components/chat/TaskHeader.tsx
@@ -16,6 +16,7 @@ interface TaskHeaderProps {
 	cacheWrites?: number
 	cacheReads?: number
 	totalCost: number
+	contextTokens: number
 	onClose: () => void
 }
 
@@ -27,6 +28,7 @@ const TaskHeader: React.FC<TaskHeaderProps> = ({
 	cacheWrites,
 	cacheReads,
 	totalCost,
+	contextTokens,
 	onClose,
 }) => {
 	const { apiConfiguration } = useExtensionState()
@@ -272,6 +274,13 @@ const TaskHeader: React.FC<TaskHeaderProps> = ({
 								{!isCostAvailable && <ExportButton />}
 							</div>
 
+							<div style={{ display: "flex", alignItems: "center", gap: "4px", flexWrap: "wrap" }}>
+								<span style={{ fontWeight: "bold" }}>Context:</span>
+								<span style={{ display: "flex", alignItems: "center", gap: "3px" }}>
+									{formatLargeNumber(contextTokens || 0)}
+								</span>
+							</div>
+
 							{shouldShowPromptCacheInfo && (cacheReads !== undefined || cacheWrites !== undefined) && (
 								<div style={{ display: "flex", alignItems: "center", gap: "4px", flexWrap: "wrap" }}>
 									<span style={{ fontWeight: "bold" }}>Cache:</span>