feat: Cost Display in Task Header - Suppress Zero Cost Values and Ensure Visibility for Gemini, OpenAI, LM Studio, and Ollama (#2662)

test: Add unit tests for TaskHeader component cost display logic

Author: d-oit

webview-ui/src/components/chat/TaskHeader.tsx

@@ -122,13 +122,8 @@ const TaskHeader: React.FC<TaskHeaderProps> = ({
 	}, [task.text, windowWidth])
 
 	const isCostAvailable = useMemo(() => {
-		return (
-			apiConfiguration?.apiProvider !== "openai" &&
-			apiConfiguration?.apiProvider !== "ollama" &&
-			apiConfiguration?.apiProvider !== "lmstudio" &&
-			apiConfiguration?.apiProvider !== "gemini"
-		)
-	}, [apiConfiguration?.apiProvider])
+		return totalCost !== null && totalCost !== undefined && totalCost > 0 && !isNaN(totalCost)
+	}, [totalCost])
 
 	const shouldShowPromptCacheInfo = doesModelSupportPromptCache && apiConfiguration?.apiProvider !== "openrouter"
 

webview-ui/src/components/chat/__tests__/TaskHeader.test.tsx

@@ -0,0 +1,114 @@
+import React from "react"
+import { render, screen } from "@testing-library/react"
+import TaskHeader from "../TaskHeader"
+import { ApiConfiguration } from "../../../../../src/shared/api"
+
+// Mock the vscode API
+jest.mock("@/utils/vscode", () => ({
+	vscode: {
+		postMessage: jest.fn(),
+	},
+}))
+
+// Mock the ExtensionStateContext
+jest.mock("../../../context/ExtensionStateContext", () => ({
+	useExtensionState: () => ({
+		apiConfiguration: {
+			apiProvider: "anthropic",
+			apiKey: "test-api-key", // Add relevant fields
+			apiModelId: "claude-3-opus-20240229", // Add relevant fields
+		} as ApiConfiguration, // Optional: Add type assertion if ApiConfiguration is imported
+		currentTaskItem: null,
+	}),
+}))
+
+describe("TaskHeader", () => {
+	const defaultProps = {
+		task: { text: "Test task", images: [] },
+		tokensIn: 100,
+		tokensOut: 50,
+		doesModelSupportPromptCache: true,
+		totalCost: 0.05,
+		contextTokens: 200,
+		onClose: jest.fn(),
+	}
+
+	it("should display cost when totalCost is greater than 0", () => {
+		render(
+			<TaskHeader
+				{...defaultProps}
+				task={{
+					type: "say",
+					ts: Date.now(),
+					text: "Test task",
+					images: [],
+				}}
+			/>,
+		)
+		expect(screen.getByText("$0.0500")).toBeInTheDocument()
+	})
+
+	it("should not display cost when totalCost is 0", () => {
+		render(
+			<TaskHeader
+				{...defaultProps}
+				totalCost={0}
+				task={{
+					type: "say",
+					ts: Date.now(),
+					text: "Test task",
+					images: [],
+				}}
+			/>,
+		)
+		expect(screen.queryByText("$0.0000")).not.toBeInTheDocument()
+	})
+
+	it("should not display cost when totalCost is null", () => {
+		render(
+			<TaskHeader
+				{...defaultProps}
+				totalCost={null as any}
+				task={{
+					type: "say",
+					ts: Date.now(),
+					text: "Test task",
+					images: [],
+				}}
+			/>,
+		)
+		expect(screen.queryByText(/\$/)).not.toBeInTheDocument()
+	})
+
+	it("should not display cost when totalCost is undefined", () => {
+		render(
+			<TaskHeader
+				{...defaultProps}
+				totalCost={undefined as any}
+				task={{
+					type: "say",
+					ts: Date.now(),
+					text: "Test task",
+					images: [],
+				}}
+			/>,
+		)
+		expect(screen.queryByText(/\$/)).not.toBeInTheDocument()
+	})
+
+	it("should not display cost when totalCost is NaN", () => {
+		render(
+			<TaskHeader
+				{...defaultProps}
+				totalCost={NaN}
+				task={{
+					type: "say",
+					ts: Date.now(),
+					text: "Test task",
+					images: [],
+				}}
+			/>,
+		)
+		expect(screen.queryByText(/\$/)).not.toBeInTheDocument()
+	})
+})