From e5d712db5947de3805061688b7062552295e4f2a Mon Sep 17 00:00:00 2001
From: Saoud Rizwan <7799382+saoudrizwan@users.noreply.github.com>
Date: Wed, 29 Jan 2025 13:53:15 -0800
Subject: [PATCH] Add context window info to task header

---
 webview-ui/src/components/chat/ChatView.tsx   | 16 ++++
 webview-ui/src/components/chat/TaskHeader.tsx | 78 +++++++++++++++++++
 2 files changed, 94 insertions(+)

diff --git a/webview-ui/src/components/chat/ChatView.tsx b/webview-ui/src/components/chat/ChatView.tsx
index aec4e544a9..7a7e6b93ae 100644
--- a/webview-ui/src/components/chat/ChatView.tsx
+++ b/webview-ui/src/components/chat/ChatView.tsx
@@ -5,6 +5,7 @@ import { useDeepCompareEffect, useEvent, useMount } from "react-use"
 import { Virtuoso, type VirtuosoHandle } from "react-virtuoso"
 import styled from "styled-components"
 import {
+	ClineApiReqInfo,
 	ClineAsk,
 	ClineMessage,
 	ClineSayBrowserAction,
@@ -44,6 +45,20 @@ const ChatView = ({ isHidden, showAnnouncement, hideAnnouncement, showHistoryVie
 	// has to be after api_req_finished are all reduced into api_req_started messages
 	const apiMetrics = useMemo(() => getApiMetrics(modifiedMessages), [modifiedMessages])
 
+	const lastApiReqTotalTokens = useMemo(() => {
+		const getTotalTokensFromApiReqMessage = (msg: ClineMessage) => {
+			if (!msg.text) return 0
+			const { tokensIn, tokensOut, cacheWrites, cacheReads }: ClineApiReqInfo = JSON.parse(msg.text)
+			return (tokensIn || 0) + (tokensOut || 0) + (cacheWrites || 0) + (cacheReads || 0)
+		}
+		const lastApiReqMessage = findLast(modifiedMessages, (msg) => {
+			if (msg.say !== "api_req_started") return false
+			return getTotalTokensFromApiReqMessage(msg) > 0
+		})
+		if (!lastApiReqMessage) return undefined
+		return getTotalTokensFromApiReqMessage(lastApiReqMessage)
+	}, [modifiedMessages])
+
 	const [inputValue, setInputValue] = useState("")
 	const textAreaRef = useRef<HTMLTextAreaElement>(null)
 	const [textAreaDisabled, setTextAreaDisabled] = useState(false)
@@ -729,6 +744,7 @@ const ChatView = ({ isHidden, showAnnouncement, hideAnnouncement, showHistoryVie
 					cacheWrites={apiMetrics.totalCacheWrites}
 					cacheReads={apiMetrics.totalCacheReads}
 					totalCost={apiMetrics.totalCost}
+					lastApiReqTotalTokens={lastApiReqTotalTokens}
 					onClose={handleTaskCloseButtonClick}
 				/>
 			) : (
diff --git a/webview-ui/src/components/chat/TaskHeader.tsx b/webview-ui/src/components/chat/TaskHeader.tsx
index d04f0aecf4..aeede40047 100644
--- a/webview-ui/src/components/chat/TaskHeader.tsx
+++ b/webview-ui/src/components/chat/TaskHeader.tsx
@@ -8,6 +8,7 @@ import { formatLargeNumber } from "../../utils/format"
 import { formatSize } from "../../utils/size"
 import { vscode } from "../../utils/vscode"
 import Thumbnails from "../common/Thumbnails"
+import { normalizeApiConfiguration } from "../settings/ApiOptions"
 
 interface TaskHeaderProps {
 	task: ClineMessage
@@ -17,9 +18,39 @@ interface TaskHeaderProps {
 	cacheWrites?: number
 	cacheReads?: number
 	totalCost: number
+	lastApiReqTotalTokens?: number
 	onClose: () => void
 }
 
+const LinearProgress: React.FC<{ percentage: number }> = ({ percentage }) => (
+	<div
+		style={{
+			display: "flex",
+			alignItems: "center",
+			gap: "8px",
+			fontSize: "11px",
+		}}>
+		<div
+			style={{
+				width: "100px",
+				height: "4px",
+				backgroundColor: "color-mix(in srgb, var(--vscode-badge-foreground) 20%, transparent)",
+				borderRadius: "2px",
+				overflow: "hidden",
+			}}>
+			<div
+				style={{
+					width: `${percentage}%`,
+					height: "100%",
+					backgroundColor: "var(--vscode-badge-foreground)",
+					borderRadius: "2px",
+				}}
+			/>
+		</div>
+		<span>{Math.round(percentage)}%</span>
+	</div>
+)
+
 const TaskHeader: React.FC<TaskHeaderProps> = ({
 	task,
 	tokensIn,
@@ -28,6 +59,7 @@ const TaskHeader: React.FC<TaskHeaderProps> = ({
 	cacheWrites,
 	cacheReads,
 	totalCost,
+	lastApiReqTotalTokens,
 	onClose,
 }) => {
 	const { apiConfiguration, currentTaskItem, checkpointTrackerErrorMessage } = useExtensionState()
@@ -37,6 +69,9 @@ const TaskHeader: React.FC<TaskHeaderProps> = ({
 	const textContainerRef = useRef<HTMLDivElement>(null)
 	const textRef = useRef<HTMLDivElement>(null)
 
+	const { selectedModelInfo } = useMemo(() => normalizeApiConfiguration(apiConfiguration), [apiConfiguration])
+	const contextWindow = selectedModelInfo?.contextWindow
+
 	/*
 	When dealing with event listeners in React components that depend on state variables, we face a challenge. We want our listener to always use the most up-to-date version of a callback function that relies on current state, but we don't want to constantly add and remove event listeners as that function updates. This scenario often arises with resize listeners or other window events. Simply adding the listener in a useEffect with an empty dependency array risks using stale state, while including the callback in the dependencies can lead to unnecessary re-registrations of the listener. There are react hook libraries that provide a elegant solution to this problem by utilizing the useRef hook to maintain a reference to the latest callback function without triggering re-renders or effect re-runs. This approach ensures that our event listener always has access to the most current state while minimizing performance overhead and potential memory leaks from multiple listener registrations. 
 	Sources
@@ -105,6 +140,48 @@ const TaskHeader: React.FC<TaskHeaderProps> = ({
 
 	const shouldShowPromptCacheInfo = doesModelSupportPromptCache && apiConfiguration?.apiProvider !== "openrouter"
 
+	const ContextWindowComponent = (
+		<>
+			{isTaskExpanded && contextWindow && lastApiReqTotalTokens && (
+				<div
+					style={{
+						display: "flex",
+						justifyContent: "space-between",
+						alignItems: "center",
+						height: 17,
+					}}>
+					<div style={{ display: "flex", alignItems: "center", gap: "4px" }}>
+						<span style={{ fontWeight: "bold" }}>Context Window:</span>
+						<span>
+							{formatLargeNumber(lastApiReqTotalTokens)} (
+							{Math.round((lastApiReqTotalTokens / contextWindow) * 100)}%)
+						</span>
+					</div>
+					<div style={{ display: "flex", alignItems: "center", gap: "8px", flex: 1, marginLeft: 8 }}>
+						<div
+							style={{
+								flex: 1,
+								height: "4px",
+								backgroundColor: "color-mix(in srgb, var(--vscode-badge-foreground) 20%, transparent)",
+								borderRadius: "2px",
+								overflow: "hidden",
+							}}>
+							<div
+								style={{
+									width: `${(lastApiReqTotalTokens / contextWindow) * 100}%`,
+									height: "100%",
+									backgroundColor: "var(--vscode-badge-foreground)",
+									borderRadius: "2px",
+								}}
+							/>
+						</div>
+						<span>{formatLargeNumber(contextWindow)}</span>
+					</div>
+				</div>
+			)}
+		</>
+	)
+
 	return (
 		<div style={{ padding: "10px 13px 10px 13px" }}>
 			<div
@@ -354,6 +431,7 @@ const TaskHeader: React.FC<TaskHeaderProps> = ({
 									</span>
 								</div>
 							)}
+							{ContextWindowComponent}
 							{isCostAvailable && (
 								<div
 									style={{